add droptol (#26)

* initial commit * update Project.toml * remove fastchebinterp warning * make symmetry test more readable * initial commit * initial commit * finish rebasing * bugfix in MassVelocityInterp * initial commit * update Project.toml * remove fastchebinterp warning * make symmetry test more readable * initial commit * finish rebasing * bugfix in MassVelocityInterp * test droptol
lxvm · Jan 31, 2024 · 5264124 · 5264124
1 parent ab53a7e
commit 5264124
Show file tree

Hide file tree

Showing 6 changed files with 119 additions and 58 deletions.
diff --git a/Project.toml b/Project.toml
@@ -10,34 +10,43 @@ EquiBaryInterp = "c5458111-f826-47c0-abe9-f36aa00852ec"
 FourierSeriesEvaluators = "2a892dea-6eef-4bb5-9d1c-de966c9f6db5"
 HChebInterp = "78faba9b-a54b-441f-8118-62407cbe4e59"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+Logging = "56ddb016-857b-54e1-b83d-db4d58db5568"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 
+[weakdeps]
+Brillouin = "23470ee3-d0df-4052-8b1a-8cbd6363e7f0"
+Makie = "ee78f7c6-11fb-53f2-987a-cfe4a2b5a57a"
+PlotlyJS = "f0f68f2c-4968-5e81-91da-67840de0976a"
+Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d"
+WannierIO = "cb1bc77f-5443-4951-af9f-05b616a3e422"
+
+[extensions]
+AutoBZUnitfulExt = "Unitful"
+AutoBZWannierIOExt = "WannierIO"
+BrillouinMakieExt = ["Brillouin", "Makie"]
+BrillouinPlotlyJSExt = ["Brillouin", "PlotlyJS"]
+
 [compat]
 Aqua = "0.8"
 AutoBZCore = "0.3"
 BaryRational = "0.1,1"
 Brillouin = "0.5"
 EquiBaryInterp = "0.1"
 FourierSeriesEvaluators = "1"
-HChebInterp = "0.1,1"
+HChebInterp = "1"
 LinearAlgebra = "1.9"
+Logging = "1.9"
 Makie = "0.19"
 OffsetArrays = "1"
 PlotlyJS = "0.18"
 Reexport = "1"
 StaticArrays = "1"
-WannierIO = "0.2"
 Test = "1.9"
 Unitful = "1"
+WannierIO = "0.2"
 julia = "1.9"
 
-[extensions]
-AutoBZUnitfulExt = "Unitful"
-BrillouinMakieExt = ["Brillouin", "Makie"]
-BrillouinPlotlyJSExt = ["Brillouin", "PlotlyJS"]
-AutoBZWannierIOExt = "WannierIO"
-
 [extras]
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 AutoBZCore = "66bd3e16-1600-45cf-8f55-0b550710682b"
@@ -52,10 +61,3 @@ WannierIO = "cb1bc77f-5443-4951-af9f-05b616a3e422"
 
 [targets]
 test = ["Aqua", "AutoBZCore", "LinearAlgebra", "StaticArrays", "OffsetArrays", "Test"]
-
-[weakdeps]
-Brillouin = "23470ee3-d0df-4052-8b1a-8cbd6363e7f0"
-Makie = "ee78f7c6-11fb-53f2-987a-cfe4a2b5a57a"
-PlotlyJS = "f0f68f2c-4968-5e81-91da-67840de0976a"
-WannierIO = "cb1bc77f-5443-4951-af9f-05b616a3e422"
-Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d"
diff --git a/src/AutoBZ.jl b/src/AutoBZ.jl
@@ -52,6 +52,7 @@ module AutoBZ
 
 using LinearAlgebra
 using LinearAlgebra: checksquare
+using Logging
 
 using StaticArrays
 using Reexport

diff --git a/src/interp.jl b/src/interp.jl
@@ -237,7 +237,7 @@ struct MassVelocityInterp{C,B,G,N,T,iip,F,DF,TA} <: AbstractVelocityInterp{C,B,G
     h::HessianSeries{N,T,iip,F,DF}
     A::TA
     function MassVelocityInterp{C,B,G}(h::HessianSeries{N,T,iip,F,DF}, A::TA) where {C,B,G,N,T,iip,F,DF,TA}
-        @assert gauge(parentseries(h)) == GaugeDefault(parentseries(h))
+        @assert gauge(_parentseries(h)) == GaugeDefault(_parentseries(h))
         return new{C,B,G,N,T,iip,F,DF,TA}(h, A)
     end
 end

diff --git a/src/self_energies_io.jl b/src/self_energies_io.jl
@@ -139,9 +139,5 @@ function construct_chebyshev(omegas, values::Vector{<:Number}, order; atol=1e-6,
 end
 function construct_chebyshev(omegas, values::Vector{T}, order; atol=1e-6, tol=1e-13, mmax=100) where {T<:SArray}
     interp = ntuple(n -> aaa(omegas, getindex.(values, n); tol=tol, mmax=mmax), length(T))
-    ndims(T) > 1 && @warn """
-    Matrix-valued interpolation may not work without the FastChebInterp version at
-    `import Pkg; Pkg.add(url="https://github.com/lxvm/FastChebInterp.jl.git#pr_sarray")`
-    """
     hchebinterp(x -> T(map(f -> f(x), interp)), extrema(omegas)...; order=order, atol=atol)
 end
diff --git a/src/wannier90io.jl b/src/wannier90io.jl
@@ -37,7 +37,7 @@ function parse_hamiltonian(file::IO, ::Type{F}) where {F<:AbstractFloat}
         end
         C[k] = SMatrix{num_wann,num_wann,Complex{F},num_wann^2}(c)
     end
-    return (; date_time, num_wann, nrpts, degen, irvec, C)
+    return (; date_time, num_wann, nrpts, degen, irvec, H=C)
 end
 
 """
@@ -54,13 +54,12 @@ keyword `compact` to specify:
 - `:U`: store the upper triangle of the coefficients
 - `:S`: store the lower triangle of the symmetrized coefficients, `(c+c')/2`
 """
-function load_interp(::Type{<:HamiltonianInterp}, seed; precision=Float64, gauge=GaugeDefault(HamiltonianInterp), compact=:N, soc=nothing)
+function load_interp(::Type{<:HamiltonianInterp}, seed; precision=Float64, gauge=GaugeDefault(HamiltonianInterp), compact=:N, soc=nothing, droptol=eps(precision))
     (; nkpt) = parse_wout(seed * ".wout", precision)
-    (; num_wann, degen, irvec, C) = parse_hamiltonian(seed * "_hr.dat", precision)
+    (; num_wann, degen, irvec, H) = parse_hamiltonian(seed * "_hr.dat", precision)
     check_degen(degen, nkpt)
-    C_ = load_coefficients(Val{compact}(), num_wann, irvec, degen, C)[1]
-    offset = map(s -> -div(s,2)-1, size(C_))
-    f = FourierSeries(C_; period=freq2rad(one(precision)), offset=offset)
+    (C, origin), = load_coefficients(Val{compact}(), droptol, num_wann, irvec, degen, H)
+    f = FourierSeries(C; period=freq2rad(one(precision)), offset=Tuple(-origin))
     if soc === nothing
         return HamiltonianInterp(Freq2RadSeries(f), gauge=gauge)
     else
@@ -76,29 +75,57 @@ load_coeff(T, ::Val{:L}, c) = T(c)
 load_coeff(T, ::Val{:U}, c) = T(c')
 load_coeff(T, ::Val{:S}, c) = T(0.5*(c+c'))
 
-function load_coefficients(compact, num_wann, irvec, degen, cs...)
-    return load_coefficients(compact, num_wann, irvec, degen, cs)
+function load_coefficients(compact, droptol, num_wann, irvec, degen, cs...)
+    return load_coefficients(compact, droptol, num_wann, irvec, degen, cs)
 end
-function load_coefficients(compact, num_wann, irvec, degen, cs::NTuple{N}) where N
+function load_coefficients(compact, reltol, num_wann, irvec, degen, cs::NTuple{N}) where N
     T = load_coeff_type(eltype(eltype(eltype(cs))), compact, num_wann)
-    nmodes = zeros(Int, 3)
-    for idx in irvec
-        @inbounds for i in 1:3
-            if (n = abs(idx[i])) > nmodes[i]
-                nmodes[i] = n
-            end
-        end
-    end
-    Cs = ntuple(_ -> zeros(T, (2nmodes .+ 1)...), Val(N))
+    bounds = ntuple(n -> extrema(idx -> idx[n], irvec), Val(3))
+    lbounds = map(first, bounds)
+    ubounds = map(last,  bounds)
+    nmodes = map((lb, ub) -> ub-lb+1, lbounds, ubounds)
+    Cs = ntuple(_ -> zeros(T, nmodes), Val(N))
+    origins = map(C -> CartesianIndex(ntuple(n -> firstindex(C,n)-lbounds[n], Val(3))), Cs)
     for (i,idx) in enumerate(irvec)
-        idx_ = CartesianIndex((idx .+ nmodes .+ 1)...)
-        for (j,c) in enumerate(cs)
-            Cs[j][idx_] = load_coeff(T, compact, c[i])/degen[i]
+        _idx_ = CartesianIndex(idx...)
+        for (c,C,o) in zip(cs, Cs, origins)
+            C[_idx_+o] = load_coeff(T, compact, c[i])/degen[i]
         end
     end
-    return Cs
+    return map((C, origin) -> droptol(C, origin, reltol), Cs, origins)
 end
 
+infnorm(x::Number) = abs(x)
+infnorm(x::AbstractArray) = maximum(infnorm, x)
+
+# based on FastChebInterp.jl and helpful for truncating zero-padded coefficients
+function droptol(C::AbstractArray, origin::CartesianIndex, reltol)
+    abstol = infnorm(C) * reltol
+    norm_geq_tol = >=(abstol)∘infnorm
+    # https://juliaarrays.github.io/OffsetArrays.jl/stable/internals/#Caveats
+    # compute the new size, dropping values below tol at both ends of axes
+    newlb = ntuple(ndims(C)) do dim
+        n = firstindex(C, dim)
+        while n < lastindex(C, dim)
+            r = let n=n; ntuple(i -> axes(C,i)[i == dim ? (n:n) : (begin:end)], ndims(C)); end
+            any(norm_geq_tol, @view C[CartesianIndices(r)]) && break
+            n += 1
+        end
+        n
+    end
+    newub = ntuple(ndims(C)) do dim
+        n = lastindex(C, dim)
+        while n > firstindex(C, dim)
+            r = let n=n; ntuple(i -> axes(C,i)[i == dim ? (n:n) : (begin:end)], ndims(C)); end
+            any(norm_geq_tol, @view C[CartesianIndices(r)]) && break
+            n -= 1
+        end
+        n
+    end
+    newC = C[CartesianIndices(map(:, newlb, newub))]
+    neworigin = origin + CartesianIndex(ntuple(n -> firstindex(newC,n)-newlb[n], ndims(C)))
+    return (newC, neworigin)
+end
 
 """
     parse_position_operator(filename, rot=I)
@@ -166,21 +193,20 @@ Fourier coefficients in compact form, use the keyword `compact` to specify:
 Note that in some cases the coefficients are not Hermitian even though the
 values of the series are.
 """
-function load_interp(::Type{<:BerryConnectionInterp}, seed; precision=Float64, coord=CoordDefault(BerryConnectionInterp), period=one(precision), compact=:N, soc=nothing)
+function load_interp(::Type{<:BerryConnectionInterp}, seed; precision=Float64, coord=CoordDefault(BerryConnectionInterp), period=one(precision), compact=:N, soc=nothing, droptol=eps(precision))
     (; A, nkpt) = parse_wout(seed * ".wout", precision)
     invA = inv(A) # compute inv(A) for map from Cartesian to lattice coordinates
     rot, irot = pick_rot(coord, A, invA)
     (; degen) = parse_hamiltonian(seed * "_hr.dat", precision)
     check_degen(degen, nkpt)
     (; num_wann, irvec, As) = parse_position_operator(seed * "_r.dat", precision, rot)
-    A1, A2, A3 = load_coefficients(Val{compact}(), num_wann, irvec, degen, As...)
-    F1_ = FourierSeries(A1; period=one(precision), offset=map(s -> -div(s,2)-1, size(A1)))
-    F1  = soc === nothing ? F1_ : WrapperFourierSeries(wrap_soc, F1_)
-    F2_ = FourierSeries(A2; period=one(precision), offset=map(s -> -div(s,2)-1, size(A2)))
-    F2  = soc === nothing ? F2_ : WrapperFourierSeries(wrap_soc, F2_)
-    F3_ = FourierSeries(A3; period=one(precision), offset=map(s -> -div(s,2)-1, size(A3)))
-    F3  = soc === nothing ? F3_ : WrapperFourierSeries(wrap_soc, F3_)
-    BerryConnectionInterp{coord}(ManyFourierSeries(F1, F2, F3; period=period), irot; coord=coord)
+    (A1,o1), (A2,o2), (A3,o3) = load_coefficients(Val{compact}(), droptol, num_wann, irvec, degen, As)
+    F1 = FourierSeries(A1; period=one(precision), offset=Tuple(-o1))
+    F2 = FourierSeries(A2; period=one(precision), offset=Tuple(-o2))
+    F3 = FourierSeries(A3; period=one(precision), offset=Tuple(-o3))
+    Fs = (F1, F2, F3)
+    Ms = soc === nothing ? Fs : map(f -> WrapperFourierSeries(wrap_soc, f), Fs)
+    BerryConnectionInterp{coord}(ManyFourierSeries(Ms...; period=period), irot; coord=coord)
 end
 
 """
@@ -348,6 +374,11 @@ function load_autobz(bz::IBZ, seedname::String; precision=Float64, kws...)
     return load_bz(convert(AbstractBZ{3}, bz), A, B, species, reduce(hcat, frac_lat); kws..., coordinates="lattice")
 end
 
+struct CompactDisplay{T}
+    obj::T
+end
+Base.show(io::IO, a::CompactDisplay) = show(io, a.obj)
+
 """
     load_wannier90_data(seedname::String; bz::AbstractBZ=FBZ(), interp::AbstractWannierInterp=HamiltonianInterp, kwargs...)
 
@@ -373,13 +404,10 @@ function load_wannier90_data(seedname::String; precision=Float64, load_interp=lo
             val = wi(S*k)
             return calc_interp_error(wi, val, ref)
         end
-        msg = """
-        Testing that the interpolant's Hamiltonian satisfies the symmetries at random k-point
-            k = $(k)
-        and found a maximum error $(err) for symmetry
-            bz.syms[$(idx)] = $(bz.syms[idx])
-        """
-        err > sqrt(eps(one(err)))*oneunit(err) ? @warn(msg) : @info(msg)
+        sym = CompactDisplay(bz.syms[idx])
+        kpt = CompactDisplay(k)
+        msg = "Symmetry test of the interpolant's Hamiltonian at a random k-point"
+        @logmsg (err > sqrt(eps(one(err)))*oneunit(err) ? Warn : Info) msg seedname kpt sym err
     end
     return (wi, bz)
 end

diff --git a/test/utils.jl b/test/utils.jl
@@ -29,4 +29,38 @@ using Test, AutoBZ, OffsetArrays
         shift!(f2, -o)
     end
 
+end
+
+@testset "droptol" begin
+
+    Cpad = [
+        0.0 0.0 0.0 0.0
+        0.0 0.1 0.0 0.0
+        0.2 0.4 0.2 0.0
+        0.3 0.5 0.3 0.0
+        0.2 0.4 0.2 0.0
+        0.0 0.1 0.0 0.0
+    ]
+    OCpad = OffsetMatrix(Array(Cpad), -3:2, -1:2)
+    C = Cpad[begin+1:end, begin:end-1]
+    OC = OffsetMatrix(C, -2:2, -1:1)
+
+    @test @inferred(AutoBZ.droptol(Cpad,  CartesianIndex(4,2), eps())) == (C, CartesianIndex(3,2))
+    @test @inferred(AutoBZ.droptol(OCpad, CartesianIndex(0,0), eps())) == (C, CartesianIndex(3,2))
+    @test @inferred(AutoBZ.droptol(C,  CartesianIndex(3,2), eps())) == (C, CartesianIndex(3,2))
+    @test @inferred(AutoBZ.droptol(OC, CartesianIndex(0,0), eps())) == (C, CartesianIndex(3,2))
+
+
+    Cpad2 = [
+        0.0 0.0 0.0 0.0
+        0.0 0.0 0.0 0.0
+        0.0 0.1 0.0 0.0
+        0.2 0.4 0.2 0.0
+        0.3 0.5 0.3 0.0
+        0.2 0.4 0.2 0.0
+        0.0 0.1 0.0 0.0
+    ]
+    OCpad2 = OffsetMatrix(Array(Cpad2), -4:2, -1:2)
+    @test @inferred(AutoBZ.droptol(Cpad2,  CartesianIndex(5,2), eps())) == (C, CartesianIndex(3,2))
+    @test @inferred(AutoBZ.droptol(OCpad2, CartesianIndex(0,0), eps())) == (C, CartesianIndex(3,2))
 end