[ITensors] Fix checkflux for QN ITensor with no blocks (#1210)

* Allow no blocks in checkflux

* Fix typo breaking test

* Add regression test

* Improvements to flux and checkflux functions

src/qn/flux.jl

@@ -1,4 +1,3 @@
-flux(T::ITensor, args...) = flux(tensor(T), args...)
 
 """
     flux(T::ITensor)
@@ -7,47 +6,75 @@ Returns the flux of the ITensor.
 
 If the ITensor is empty or it has no QNs, returns `nothing`.
 """
-function flux(T::ITensor)
-  return flux(tensor(T))
-end
+flux(T::ITensor, args...) = flux(tensor(T), args...)
 
-function checkflux(T::ITensor, flux_check)
-  return checkflux(tensor(T), flux_check)
-end
+"""
+    checkflux(T::ITensor)
 
-function checkflux(T::ITensor)
-  return checkflux(tensor(T))
-end
+Check that fluxes of all non-zero blocks of a blocked or symmetric ITensor
+are equal. Throws an error if one or more blocks have a different flux.
+"""
+checkflux(T::ITensor, flux_check) = checkflux(tensor(T), flux_check)
+
+"""
+    checkflux(T::ITensor, flux)
+
+Check that fluxes of all non-zero blocks of a blocked or symmetric Tensor
+equal the value `flux`. Throws an error if one or more blocks does not have this flux.
+"""
+checkflux(T::ITensor) = checkflux(tensor(T))
 
 #
 # Tensor versions
 # TODO: Move to NDTensors when QN functionality
 # is moved there.
 #
 
-flux(T::Tensor, args...) = flux(inds(T), args...)
+"""
+    flux(T::Tensor, block::Block)
+
+Compute the flux of a specific block of a Tensor, 
+regardless of whether this block is present or not in the storage.
+"""
+flux(T::Tensor, block::Block) = flux(inds(T), block)
+
+"""
+    flux(T::Tensor, i::Integer, is::Integer...)
+
+Compute the flux of a specific element of a Tensor, 
+regardless of whether this element is zero or non-zero.
+"""
+flux(T::Tensor, i::Integer, is::Integer...) = flux(inds(T), i, is...)
+
+"""
+    flux(T::Tensor)
 
+Return the flux of a Tensor, based on what non-zero blocks it
+has. If the Tensor is not blocked or has no non-zero blocks,
+this function returns `nothing`. 
+"""
 function flux(T::Tensor)
   (!hasqns(T) || isempty(T)) && return nothing
   @debug_check checkflux(T)
   block1 = first(eachnzblock(T))
   return flux(T, block1)
 end
 
-function checkflux(T::Tensor, flux_check)
-  for b in nzblocks(T)
-    fluxTb = flux(T, b)
-    if fluxTb != flux_check
-      error(
-        "Block $b has flux $fluxTb that is inconsistent with the desired flux $flux_check"
-      )
-    end
-  end
-  return nothing
-end
+allfluxequal(T::Tensor, flux_to_check) = all(b -> flux(T, b) == flux_to_check, nzblocks(T))
+allfluxequal(T::Tensor) = allequal(flux(T, b) for b in nzblocks(T))
 
-function checkflux(T::Tensor)
-  b1 = first(nzblocks(T))
-  fluxTb1 = flux(T, b1)
-  return checkflux(T, fluxTb1)
-end
+"""
+    checkflux(T::Tensor)
+
+Check that fluxes of all non-zero blocks of a blocked or symmetric Tensor
+are equal. Throws an error if one or more blocks have a different flux.
+"""
+checkflux(T::Tensor) = allfluxequal(T) ? nothing : error("Fluxes not all equal")
+
+"""
+    checkflux(T::Tensor, flux)
+
+Check that fluxes of all non-zero blocks of a blocked or symmetric Tensor
+equal the value `flux`. Throws an error if one or more blocks does not have this flux.
+"""
+checkflux(T::Tensor, flux) = allfluxequal(T, flux) ? nothing : error("Fluxes not all equal")

test/base/test_qnitensor.jl

@@ -97,6 +97,10 @@ Random.seed!(1234)
     ]
     @test_throws ErrorException ITensor(A, i', dag(i); tol=1e-8)
     @test ITensor(A, i', dag(i); tol=1e-8, checkflux=false) isa ITensor
+
+    # Construct from zero matrix. Flux check should still pass
+    # (Regression test for issue #1209)
+    @test ITensor(zeros(3, 3), i', dag(i)) isa ITensor
   end
 
   @testset "similartype regression test" begin
@@ -124,12 +128,14 @@ Random.seed!(1234)
     @test T[2] == 0
     @test T[3] == 1
     @test T[4] == 2
+    # Test fluxes of specific elements:
     @test flux(T, 1) == QN(0)
     @test flux(T, 2) == QN(0)
     @test flux(T, 3) == QN(1)
     @test flux(T, 4) == QN(1)
     @test_throws BoundsError flux(T, 5)
     @test_throws BoundsError flux(T, 0)
+    # Test fluxes of specific Blocks
     @test flux(T, Block(1)) == QN(0)
     @test flux(T, Block(2)) == QN(1)
     @test_throws BoundsError flux(T, Block(0))
@@ -548,18 +554,21 @@ Random.seed!(1234)
     @testset "Combine set direction" begin
       i1 = Index([QN(0) => 2, QN(1) => 3], "i1")
       A = randomITensor(i1', dag(i1))
+      # Test that checkflux does not throw an error:
       @test isnothing(ITensors.checkflux(A))
       C = combiner(dag(i1); dir=ITensors.Out)
       c = combinedind(C)
       @test dir(c) == ITensors.Out
       AC = A * C
       @test nnz(AC) == nnz(A)
       @test nnzblocks(AC) == nnzblocks(A)
+      # Test that checkflux does not throw an error:
       @test isnothing(ITensors.checkflux(AC))
       Ap = AC * dag(C)
       @test nnz(Ap) == nnz(A)
       @test nnzblocks(Ap) == nnzblocks(A)
       @test hassameinds(Ap, A)
+      # Test that checkflux does not throw an error:
       @test isnothing(ITensors.checkflux(AC))
       @test A ≈ Ap
     end