test droptol

lxvm · Jan 31, 2024 · 66d96a6 · 66d96a6
1 parent 589778b
commit 66d96a6
Show file tree

Hide file tree

Showing 2 changed files with 40 additions and 8 deletions.
diff --git a/src/wannier90io.jl b/src/wannier90io.jl
@@ -102,30 +102,28 @@ infnorm(x::AbstractArray) = maximum(infnorm, x)
 function droptol(C::AbstractArray, origin::CartesianIndex, reltol)
     abstol = infnorm(C) * reltol
     norm_geq_tol = >=(abstol)∘infnorm
-    all(norm_geq_tol, C) && return (C, origin)
-
+    # https://juliaarrays.github.io/OffsetArrays.jl/stable/internals/#Caveats
     # compute the new size, dropping values below tol at both ends of axes
-    N = ndims(C)
-    newlb = ntuple(Val(N)) do dim
+    newlb = ntuple(ndims(C)) do dim
         n = firstindex(C, dim)
         while n < lastindex(C, dim)
-            r = ntuple(i -> axes(C,i)[i == dim ? (n:n) : (begin:end)], Val(N))
+            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(Val(N)) do dim
+    newub = ntuple(ndims(C)) do dim
         n = lastindex(C, dim)
         while n > firstindex(C, dim)
-            r = ntuple(i -> axes(C,i)[i == dim ? (n:n) : (begin:end)], Val(N))
+            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], Val(N)))
+    neworigin = origin + CartesianIndex(ntuple(n -> firstindex(newC,n)-newlb[n], ndims(C)))
     return (newC, neworigin)
 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