Merge pull request #8 from numericalEFT/treedivBZ

add floor

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))")

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

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)

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

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

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