define checkspaces and findfirstblock_sector (#47)

Author: ogauthe

Project.toml

@@ -1,7 +1,7 @@
 name = "GradedArrays"
 uuid = "bc96ca6e-b7c8-4bb6-888e-c93f838762c2"
 authors = ["ITensor developers <support@itensor.org> and contributors"]
-version = "0.4.13"
+version = "0.4.14"
 
 [deps]
 BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e"

src/abstractsector.jl

@@ -76,3 +76,7 @@ end
 
 # ================================  GradedUnitRanges interface  ==================================
 sector_type(S::Type{<:AbstractSector}) = S
+
+function findfirstblock(g::AbstractGradedUnitRange, s::AbstractSector)
+  return findfirstblock_sector(g::AbstractGradedUnitRange, s)
+end

src/gradedunitrange.jl

@@ -149,6 +149,13 @@ function gradedunitrange_getindices(::NotAbelianStyle, g::AbstractUnitRange, ind
   return blockedunitrange_getindices(ungrade(g), indices)
 end
 
+# do not overload BlockArrays.findblock to avoid ambiguity for findblock(g, 1)
+function findfirstblock_sector(g::AbstractGradedUnitRange, s)
+  i = findfirst(==(s), sectors(g))
+  isnothing(i) && return nothing
+  return Block(i)
+end
+
 # ==================================  Base interface  ======================================
 
 # needed in BlockSparseArrays

src/gradedunitrange_interface.jl

@@ -54,3 +54,14 @@ sectors(v::AbstractBlockVector) = mapreduce(sectors, vcat, blocks(v))
 function space_isequal(a1::AbstractUnitRange, a2::AbstractUnitRange)
   return (isdual(a1) == isdual(a2)) && sectors(a1) == sectors(a2) && blockisequal(a1, a2)
 end
+
+function checkspaces(::Type{Bool}, axes1, axes2)
+  return length(axes1) == length(axes2) && all(space_isequal.(axes1, axes2))
+end
+
+function checkspaces(ax1, ax2)
+  return checkspaces(Bool, ax1, ax2) || throw(ArgumentError(lazy"$ax1 does not match $ax2"))
+end
+
+checkspaces_dual(::Type{Bool}, axes1, axes2) = checkspaces(Bool, axes1, dual.(axes2))
+checkspaces_dual(axes1, axes2) = checkspaces(axes1, dual.(axes2))

test/test_gradedunitrange.jl

@@ -21,7 +21,11 @@ using GradedArrays:
   SectorOneTo,
   SectorUnitRange,
   SU,
+  checkspaces,
+  checkspaces_dual,
   dual,
+  findfirstblock,
+  findfirstblock_sector,
   flip,
   gradedrange,
   isdual,
@@ -96,6 +100,8 @@ using Test: @test, @test_throws, @testset
     @test findblock(g, 2) == Block(1)
     @test findblock(g, 3) == Block(2)
     @test findblockindex(g, 3) == Block(2)[1]
+    @test findfirstblock_sector(g, "y") == Block(2)
+    @test isnothing(findfirstblock_sector(g, "a"))
 
     @test axes(Base.Slice(g)) isa Tuple{typeof(g)}
     @test AbstractUnitRange{Int}(g) == 1:7
@@ -205,8 +211,15 @@ using Test: @test, @test_throws, @testset
     @test space_isequal(only(axes(a)), gradedrange(["y" => 6, "x" => 2]; isdual=isdual(g)))
   end
 
+  @test checkspaces(Bool, (g1, g1d), (g1, g1d))
+
   @test space_isequal(g1d, dual(g1))
   @test space_isequal(dual(g1d), g1)
+  @test checkspaces((g1, g1d), (g1, g1d))
+  @test checkspaces_dual(Bool, (g1, g1d), (g1d, g1))
+  @test checkspaces_dual((g1, g1d), (g1d, g1))
+  @test_throws ArgumentError checkspaces((g1, g1), (g1, g1d))
+  @test_throws ArgumentError checkspaces_dual((g1, g1), (g1, g1d))
 
   for a in (
     combine_blockaxes(g1, b0),
@@ -274,6 +287,8 @@ end
   @test g[4] == 4
   @test g[Block(1)[1]] == 1
   @test g[Block(2)[1]] == 3
+  @test findfirstblock(g, SU((1, 0))) == Block(2)
+  @test isnothing(findfirstblock(g, SU((2, 0))))
 
   # Non-abelian slicing operations
   a = g[2:4]