Merge pull request #9 from numericalEFT/treedivBZ

Add BaryCheb

Project.toml

@@ -6,12 +6,13 @@ version = "0.1.0"
 [deps]
 AbstractTrees = "1520ce14-60c1-5f80-bbc7-55ef81b5835c"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 
 [compat]
-StaticArrays = "1"
 AbstractTrees = "0.3"
+StaticArrays = "1"
 julia = "1.6"
 
 [extras]

example/graphene_treetest.jl

@@ -5,9 +5,10 @@ using LinearAlgebra
 
 include("graphene.jl")
 
+T = 0.05
+μ = -2.0
+
 function density(k; band=1)
-    T = 0.1
-    μ = -2.5
 
     dim = la.dim
     n = la.numatoms
@@ -20,7 +21,7 @@ end
 
 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^14, maxdepth=7, mindepth=1, N=2)
+tg = treegridfromdensity(k -> density(k), latvec; atol=1 / 2^12, maxdepth=9, mindepth=1, N=4, type=:barycheb)
 
 X, Y = zeros(Float64, length(tg)), zeros(Float64, length(tg))
 for (pi, p) in enumerate(tg)
@@ -30,7 +31,7 @@ end
 
 println("size:$(length(tg)),\t length:$(size(tg)),\t efficiency:$(efficiency(tg))")
 
-smap = SymMap(tg, k -> density(k); atol=1e-4)
+smap = SymMap{Float64}(tg, k -> density(k); atol=1e-14, rtol=1e-10)
 # println(smap.map)
 # println(smap.inv_map)
 println("compress:$(smap.reduced_length/length(smap.map))")
@@ -40,3 +41,24 @@ k = 4π
 p = plot(legend=false, size=(1024, 1024), xlim=(-k, k), ylim=(-k, k))
 scatter!(p, X, Y, marker=:cross, markersize=2)
 savefig(p, "run/tg.pdf")
+
+# test integrate
+
+function green_real(k, n; band=1)
+    dim = la.dim
+    ln = la.numatoms
+    ϵ = dispersion(dim, ln, ham, k)[band] - μ
+
+    # return - ϵ / ((π * T * (2*n + 1))^2 + ϵ^2)
+    return 1 / (exp((ϵ) / T) + 1.0)
+end
+
+# data = zeros(Float64, length(tg))
+data = MappedData(smap)
+for i in 1:length(tg)
+    data[i] = green_real(tg[i], 1)
+end
+
+n0 = integrate(data, tg)
+println("n0=$n0")
+

src/SpaceGrid.jl

@@ -7,6 +7,9 @@ using LinearAlgebra
 
 # Write your package code here.
 
+include("barycheb.jl")
+using .BaryCheb
+
 include("basemesh.jl")
 using .BaseMesh
 

src/barycheb.jl

@@ -0,0 +1,260 @@
+module BaryCheb
+
+using ..StaticArrays
+
+export BaryCheb1D, interp1D, interpND, integrate1D, integrateND
+######################################################
+#---------------- 1D barycheb ------------------------
+######################################################
+
+"""
+barychebinit(n)
+Get Chebyshev nodes of first kind and corresponding barycentric Lagrange interpolation weights. 
+Reference: Berrut, J.P. and Trefethen, L.N., 2004. Barycentric lagrange interpolation. SIAM review, 46(3), pp.501-517.
+# Arguments
+- `n`: order of the Chebyshev interpolation
+# Returns
+- Chebyshev nodes
+- Barycentric Lagrange interpolation weights
+"""
+function barychebinit(n)
+    x = zeros(Float64, n)
+    w = similar(x)
+    for i in 1:n
+        c = (2i - 1)π / (2n)
+        x[n - i + 1] = cos(c)
+        w[n - i + 1] = (-1)^(i - 1) * sin(c)
+    end
+    return x, w
+end
+
+function vandermonde(x)
+    # generate vandermonde matrix for x
+    n = length(x)
+    vmat = zeros(Float64, (n, n))
+    for i in 1:n
+        for j in 1:n
+            vmat[i,j] = x[i]^(j-1)
+        end
+    end
+    return vmat
+end
+
+function invvandermonde(order)
+    # generate inverse of vandermonde matrix for cheb x of given order
+    n = order
+    vmat = zeros(Float64, (n, n))
+    for i in 1:n
+        for j in 1:n
+            c = (2n-2i+1)π/(2n)
+            x = cos(c)
+            vmat[i,j] = x^(j-1)
+        end
+    end
+    return inv(transpose(vmat))
+end
+
+@inline function divfactorial(j, i)
+    if i == 0
+        return 1
+    elseif i == 1
+        return 1/j
+    elseif i == -1
+        return j-1
+    else
+        return factorial(j-1) / factorial(j+i-1)
+    end
+end
+
+
+function weightcoef(a, i::Int, n)
+    # integrate when i=1; differentiate when i=-1
+    b = zeros(Float64, n)
+    for j in 1:n
+        if j+i-1 > 0
+            # b[j] = a^(j+i-1)/(j+i-1)
+            b[j] = a^(j+i-1) * divfactorial(j, i)
+        elseif j+i-1 == 0
+            b[j] = 1
+        else
+            b[j] = 0
+        end
+    end
+    return b
+end
+
+@inline function calcweight(invmat, b)
+    return invmat*b
+end
+
+"""
+function barycheb(n, x, f, wc, xc)
+Barycentric Lagrange interpolation at Chebyshev nodes
+Reference: Berrut, J.P. and Trefethen, L.N., 2004. Barycentric lagrange interpolation. SIAM review, 46(3), pp.501-517.
+# Arguments
+- `n`: order of the Chebyshev interpolation
+- `x`: coordinate to interpolate
+- `f`: array of size n, function at the Chebyshev nodes
+- `wc`: array of size n, Barycentric Lagrange interpolation weights
+- `xc`: array of size n, coordinates of Chebyshev nodes
+# Returns
+- Interpolation result
+"""
+function barycheb(n, x, f, wc, xc)
+    for j in 1:n
+        if x == xc[j]
+            return f[j]
+        end    
+    end
+
+    num, den = 0.0, 0.0
+    for j in 1:n
+        q = wc[j] / (x - xc[j])
+        num += q * f[j]
+        den += q
+    end
+    return num / den
+end
+
+function barychebND(n, xs, f, wc, xc, DIM)
+    haseq = false
+    eqinds = zeros(Int, DIM)
+    for i in 1:DIM
+        for j in 1:n
+            if xs[i] == xc[j]
+                eqinds[i] = j
+                haseq = true
+            end
+        end
+    end
+
+    if haseq
+        newxs = [xs[i] for i in 1:DIM if eqinds[i]==0]
+        newDIM = length(newxs)
+        if newDIM == 0
+            return f[CartesianIndex(eqinds...)]
+        else
+            newf = view(f, [(i==0) ? (1:n) : (i) for i in eqinds]...)
+            return _barychebND_noneq(n, newxs, newf, wc, xc, newDIM)
+        end
+    else
+        return _barychebND_noneq(n, xs, f, wc, xc, DIM)
+    end
+end
+
+function _barychebND_noneq(n, xs, f, wc, xc, DIM)
+    # deal with the case when there's no xs[i] = xc[j]
+    inds = CartesianIndices(NTuple{DIM, Int}(ones(Int, DIM) .* n))
+    num, den = 0.0, 0.0
+    for (indi, ind) in enumerate(inds)
+        q = 1.0
+        for i in 1:DIM
+            q *= wc[ind[i]] / (xs[i] - xc[ind[i]])
+        end
+        num += q * f[indi]
+        den += q
+    end
+    return num / den
+end
+
+function barycheb2(n, x, f, wc, xc)
+    for j in 1:n
+        if x == xc[j]
+            return f[j, :]
+        end    
+    end
+
+    den = 0.0
+    num = zeros(eltype(f), size(f)[2])
+    for j in 1:n
+        q = wc[j] / (x - xc[j])
+        num += q .* f[j, :]
+        den += q
+    end
+    return num ./ den
+end
+
+function chebint(n, a, b, f, invmat)
+    wc = weightcoef(b, 1, n) .- weightcoef(a, 1, n)
+    intw = calcweight(invmat, wc)
+    return sum(intw .* f)
+end
+
+function chebdiff(n, x, f, invmat)
+    wc = weightcoef(x, -1, n)
+    intw = calcweight(invmat, wc)
+    return sum(intw .* f)
+end
+
+struct BaryCheb1D{N}
+    # wrapped barycheb 1d grid, x in [-1, 1]
+    x::SVector{N, Float64}
+    w::SVector{N, Float64}
+    invmat::SMatrix{N, N, Float64}
+
+    function BaryCheb1D(N::Int)
+        x, w = barychebinit(N)
+        invmat = invvandermonde(N)
+
+        return new{N}(x, w, invmat)
+    end
+end
+
+Base.getindex(bc::BaryCheb1D, i) = bc.x[i]
+
+function interp1D(data, xgrid::BaryCheb1D{N}, x) where {N}
+    return barycheb(N, x, data, xgrid.w, xgrid.x)
+end
+
+function integrate1D(data, xgrid::BaryCheb1D{N}; x1=-1, x2=1) where {N}
+    return chebint(N, x1, x2, data, xgrid.invmat)
+end
+
+function interpND(data, xgrid::BaryCheb1D{N}, xs) where {N}
+    return barychebND(N, xs, data, xgrid.w, xgrid.x, length(xs))
+end
+
+function integrateND(data, xgrid::BaryCheb1D{N}, x1s, x2s) where {N}
+    DIM = length(x1s)
+    @assert DIM == length(x2s)
+
+    intws = zeros(Float64, (DIM, N))
+    for i in 1:DIM
+        wc = weightcoef(x2s[i], 1, N) .- weightcoef(x1s[i], 1, N)
+        intws[i, :] = calcweight(xgrid.invmat, wc)
+    end
+
+    result = 0.0
+    inds = CartesianIndices(NTuple{DIM, Int}(ones(Int, DIM) .* N))
+    for (indi, ind) in enumerate(inds)
+        w = 1.0
+        for i in 1:DIM
+            w *= intws[i, ind[i]]
+        end
+        result += data[indi] * w
+    end
+
+    return result
+end
+
+function integrateND(data, xgrid::BaryCheb1D{N}, DIM) where {N}
+    @assert N ^ DIM == length(data)
+
+    intws = zeros(Float64, (DIM, N))
+    wc = weightcoef(1.0, 1, N) .- weightcoef(-1.0, 1, N)
+    intw = calcweight(xgrid.invmat, wc)
+
+    result = 0.0
+    inds = CartesianIndices(NTuple{DIM, Int}(ones(Int, DIM) .* N))
+    for (indi, ind) in enumerate(inds)
+        w = 1.0
+        for i in 1:DIM
+            w *= intw[ind[i]]
+        end
+        result += data[indi] * w
+    end
+
+    return result
+end
+
+end