diff --git a/example/graphene_treetest.jl b/example/graphene_treetest.jl index 35bae90..0b4f2c2 100644 --- a/example/graphene_treetest.jl +++ b/example/graphene_treetest.jl @@ -18,19 +18,19 @@ function density(k; band=1) # return (exp((ϵ) / T) / T)/(exp((ϵ) / T) + 1.0)^2 end -latvec = [2 0; 1 sqrt(3)] .* (2π) +latvec = [2 0; 1 sqrt(3)]' .* (2π) # println(SpaceGrid.GridTree._calc_subpoints(3, [3,3], latvec, 2)) -tg = treegridfromdensity(k -> density(k), latvec; atol=1 / 2^18, maxdepth=9, mindepth=1, N=2) +tg = treegridfromdensity(k -> density(k), latvec; atol=1 / 2^14, maxdepth=7, mindepth=1, N=2) -X, Y = zeros(Float64, size(tg)), zeros(Float64, size(tg)) +X, Y = zeros(Float64, length(tg)), zeros(Float64, length(tg)) for (pi, p) in enumerate(tg) X[pi], Y[pi] = p[1], p[2] # println(p) end -println("size:$(size(tg)),\t length:$(length(tg)),\t efficiency:$(efficiency(tg))") +println("size:$(length(tg)),\t length:$(size(tg)),\t efficiency:$(efficiency(tg))") -smap = SymMap(tg, k -> density(k); atol=1e-6) +smap = SymMap(tg, k -> density(k); atol=1e-4) # println(smap.map) # println(smap.inv_map) println("compress:$(smap.reduced_length/length(smap.map))") diff --git a/src/basemesh.jl b/src/basemesh.jl index 32fdb19..25ea536 100644 --- a/src/basemesh.jl +++ b/src/basemesh.jl @@ -9,6 +9,11 @@ abstract type AbstractMesh end struct UniformMesh{DIM, N} <: AbstractMesh origin::SVector{DIM, Float64} latvec::SMatrix{DIM, DIM, Float64} + invlatvec::SMatrix{DIM, DIM, Float64} + + function UniformMesh{DIM, N}(origin, latvec) where {DIM, N} + return new{DIM, N}(origin, latvec, inv(latvec)) + end end Base.length(mesh::UniformMesh{DIM, N}) where {DIM, N} = N @@ -26,15 +31,15 @@ end function Base.getindex(mesh::UniformMesh{DIM, N}, inds...) where {DIM, N} pos = Vector(mesh.origin) for (ni, n) in enumerate(inds) - pos = pos .+ mesh.latvec[ni, :] .* (n - 1) ./ (length(mesh)) + pos = pos .+ mesh.latvec[:, ni] .* (n - 1 + 0.5) ./ (length(mesh)) end return pos end function Base.getindex(mesh::UniformMesh{DIM, N}, i::Int) where {DIM, N} - inds = digits(i-1, base = N) + inds = digits(i-1, base = N, pad = DIM) pos = Vector(mesh.origin) for (ni, n) in enumerate(inds) - pos = pos .+ mesh.latvec[ni, :] .* n ./ (N) + pos = pos .+ mesh.latvec[:, ni] .* (n + 0.5) ./ (N) end return pos end @@ -44,4 +49,27 @@ Base.lastindex(mesh::UniformMesh) = size(mesh) Base.iterate(mesh::UniformMesh) = (mesh[1],1) Base.iterate(mesh::UniformMesh, state) = (state>=size(mesh)) ? nothing : (mesh[state+1],state+1) +function Base.floor(mesh::UniformMesh{DIM, N}, x) where {DIM, N} + # find index of nearest grid point to the point + displacement = SVector{DIM, Float64}(x) - mesh.origin + # println(displacement) + inds = (mesh.invlatvec * displacement) .* (N) .+ 1 + indexall = 1 + # println((mesh.invlatvec * displacement)) + # println(inds) + for i in 1:DIM + if inds[i] < 1 + indexi = 1 + elseif inds[i] >= N + indexi = N + else + indexi = floor(Int, inds[i]) + end + # println("$(i):$(indexi)") + indexall += (indexi - 1) * N ^ (i-1) + end + + return indexall +end + end diff --git a/src/tree.jl b/src/tree.jl index 314bf63..6d31514 100644 --- a/src/tree.jl +++ b/src/tree.jl @@ -8,6 +8,7 @@ using ..LinearAlgebra export GridNode, TreeGrid, uniformtreegrid, treegridfromdensity, efficiency, SymMap struct GridNode{DIM} + index::Vector{Int} # index in treegrid.subgrids, 0 if has children depth::Int pos::SVector{DIM,Int64} children::Vector{GridNode{DIM}} @@ -17,7 +18,7 @@ function GridNode{DIM}( isfine; depth=0, pos=SVector{DIM,Int64}(zeros(Int64, DIM)), maxdepth=10, mindepth=0) where {DIM} if (isfine(depth, pos) && depth >= mindepth) || depth >= maxdepth - return GridNode{DIM}(depth, pos, Vector{GridNode{DIM}}([])) + return GridNode{DIM}([0, ], depth, pos, Vector{GridNode{DIM}}([])) else children = Vector{GridNode{DIM}}([]) for i in 0:2^DIM-1 @@ -26,12 +27,30 @@ function GridNode{DIM}( childpos = pos .* 2 .+ bin push!(children, GridNode{DIM}(isfine; depth=childdepth, pos=childpos, maxdepth=maxdepth, mindepth=mindepth)) end - return GridNode{DIM}(depth, pos, children) + return GridNode{DIM}([0, ], depth, pos, children) end end AbstractTrees.children(node::GridNode) = node.children +function Base.floor(node::GridNode{DIM}, x) where {DIM} + # x is dimensionless pos that matches pos ./ 2^depth + if isempty(node.children) + return node.index[1] + else + midpos = (node.pos .* 2.0 .+ 1.0) ./ 2^(node.depth + 1) + + index = 1 + for i in 1:DIM + if x[i] > midpos[i] + index = index + 2^(DIM-i) + end + end + + return floor(node.children[index], x) + end +end + function efficiency(root::GridNode{DIM}) where {DIM} np, depth = 0, 0 for node in PostOrderDFS(root) @@ -48,11 +67,25 @@ end struct TreeGrid{DIM,SG} root::GridNode{DIM} latvec::SMatrix{DIM,DIM,Float64} + invlatvec::SMatrix{DIM, DIM, Float64} subgrids::Vector{SG} end efficiency(tg::TreeGrid) = efficiency(tg.root) +function Base.floor(tg::TreeGrid{DIM, SG}, x) where {DIM, SG} + dimlessx = tg.invlatvec * SVector{DIM, Float64}(x) .+ 0.5 + + tgi = floor(tg.root, dimlessx) + mesh = tg.subgrids[tgi] + + sgi = floor(mesh, x) + + sgsize = size(tg.subgrids[1]) + + return (tgi - 1) * sgsize + sgi +end + function _calc_area(latvec) return abs(det(latvec)) end @@ -63,7 +96,7 @@ function _calc_point(depth, pos, latvec) origin = zeros(size(ratio)) for i in 1:DIM - origin .+= ratio[i] .* latvec[i, :] + origin .+= ratio[i] .* latvec[:, i] end return origin @@ -97,19 +130,24 @@ end function uniformtreegrid(isfine, latvec; maxdepth=10, mindepth=0, DIM=2, N=2) root = GridNode{DIM}(isfine; maxdepth=maxdepth, mindepth=mindepth) subgrids = Vector{UniformMesh{DIM,N}}([]) + + i = 1 for node in PostOrderDFS(root) if isempty(node.children) depth = node.depth origin = _calc_origin(node, latvec) mesh = UniformMesh{DIM,N}(origin, latvec ./ 2^depth) push!(subgrids, mesh) + node.index[1] = i + i = i+1 end end - return TreeGrid{DIM,UniformMesh{DIM,N}}(root, latvec, subgrids) + + return TreeGrid{DIM,UniformMesh{DIM,N}}(root, latvec, inv(latvec), subgrids) end -Base.length(tg::TreeGrid{DIM,SG}) where {DIM,SG} = length(tg.subgrids) -Base.size(tg::TreeGrid{DIM,SG}) where {DIM,SG} = length(tg) * size(tg.subgrids[1]) +Base.length(tg::TreeGrid{DIM,SG}) where {DIM,SG} = length(tg.subgrids) * size(tg.subgrids[1]) +Base.size(tg::TreeGrid{DIM,SG}) where {DIM,SG} = (length(tg.subgrids), size(tg.subgrids[1])) # index and iterator function Base.getindex(tg::TreeGrid{DIM,SG}, i) where {DIM,SG} sgsize = size(tg.subgrids[1]) @@ -122,10 +160,10 @@ function Base.getindex(tg::TreeGrid{DIM,SG}, i, j) where {DIM,SG} return getindex(tg.subgrids[i], j) end Base.firstindex(tg::TreeGrid) = 1 -Base.lastindex(tg::TreeGrid) = size(tg) +Base.lastindex(tg::TreeGrid) = length(tg) Base.iterate(tg::TreeGrid) = (tg[1], 1) -Base.iterate(tg::TreeGrid, state) = (state >= size(tg)) ? nothing : (tg[state+1], state + 1) +Base.iterate(tg::TreeGrid, state) = (state >= length(tg)) ? nothing : (tg[state+1], state + 1) function densityisfine(density, latvec, depth, pos, atol, DIM; N=3) # compare results from subgrid of N+1 and N+3 @@ -168,10 +206,10 @@ struct SymMap{T} inv_map::Vector{Vector{Int}} function SymMap(tg::TreeGrid, density; atol=1e-6) - map = zeros(Int, size(tg)) + map = zeros(Int, length(tg)) reduced_vals = [] inv_map = [] - for pi in 1:size(tg) + for pi in 1:length(tg) # println(pi, " ", p) p = tg[pi] val = density(p) diff --git a/test/basemesh.jl b/test/basemesh.jl new file mode 100644 index 0000000..9f38250 --- /dev/null +++ b/test/basemesh.jl @@ -0,0 +1,54 @@ +@testset "Base Mesh" begin + + function dispersion(k) + me = 0.5 + return dot(k, k) / 2me + end + + function density(k) + T = 0.01 + μ = 1.0 + + ϵ = dispersion(k) - μ + + # return 1 / (exp((ϵ) / T) + 1.0) + return (π*T)^2/((π*T)^2 + ϵ^2) + # return (exp((ϵ) / T) / T)/(exp((ϵ) / T) + 1.0)^2 + end + + @testset "Uniform Mesh" begin + δ = 1e-6 + + # 2D + shift = [[0, δ], [0, -δ], [δ, 0], [-δ, 0]] + origin = [0.0 0.0] + latvec = [2 0;1 sqrt(3)]' + # latvec = [1 0; 0 1]' + + umesh = UniformMesh{2, 3}(origin, latvec) + + for i in 1:size(umesh) + for j in 1:length(shift) + # println("i=$i, point:$(umesh[i]), shift:$(shift[j])") + @test i == floor(umesh, umesh[i] + shift[j]) + end + end + + # 3D + shift = [[0, 0, δ], [0, 0, -δ], [δ, 0, 0], [-δ, 0, 0], [0, δ, 0], [0, -δ, 0]] + origin = [0.0 0.0 0.0] + latvec = [1.0 0 0; 0 1.0 0; 0 0 1.0]' + + + umesh = UniformMesh{3, 3}(origin, latvec) + + for i in 1:size(umesh) + for j in 1:length(shift) + # println("i=$i, point:$(umesh[i]), shift:$(shift[j])") + @test i == floor(umesh, umesh[i] + shift[j]) + end + end + + end + +end diff --git a/test/runtests.jl b/test/runtests.jl index 7cfb286..d2ec66d 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -6,6 +6,7 @@ using Test @testset "SpaceGrid.jl" begin if isempty(ARGS) include("tree.jl") + include("basemesh.jl") else include(ARGS[1]) end diff --git a/test/tree.jl b/test/tree.jl index 83a60c0..a8f8d41 100644 --- a/test/tree.jl +++ b/test/tree.jl @@ -15,11 +15,39 @@ # return (exp((ϵ) / T) / T)/(exp((ϵ) / T) + 1.0)^2 end - latvec = [2 0; 1 sqrt(3)] .* (2π) - tg = treegridfromdensity(k->density(k), latvec; atol = 1/2^12, maxdepth = 6, mindepth = 1, N = 2) + latvec = [2 0; 1 sqrt(3)]' .* (2π) + tg = treegridfromdensity(k->density(k), latvec; atol = 1/2^10, maxdepth = 5, mindepth = 1, N = 2) println("size:$(size(tg)),\t length:$(length(tg)),\t efficiency:$(efficiency(tg))") + @testset "TreeGrid" begin + + # test node.index + for node in PostOrderDFS(tg.root) + if node.index[1] != 0 + origin1 = GridTree._calc_origin(node, tg.latvec) + origin2 = tg.subgrids[node.index[1]].origin + @test origin1 == origin2 + end + end + + # test floor of GridNode + for node in PostOrderDFS(tg.root) + if node.index[1] != 0 + x = (node.pos .* 2.0 .+ 1.0) ./ 2^(node.depth + 1) + index = floor(tg.root, x) + @test index == node.index[1] + end + end + + # test floor of TreeGrid + for i in 1:length(tg) + p = tg[i] + @test i == floor(tg, p) + end + + end + @testset "SymMap" begin atol = 1e-6 smap = SymMap(tg, k->density(k); atol = atol) @@ -29,7 +57,7 @@ vals = smap._vals # test map - for i in 1:size(tg) + for i in 1:length(tg) @test isapprox(density(tg[i]), vals[smap.map[i]], atol = 2atol) end