Fix broadcasting over sector and non-sector arrays (#66)

Author: mtfishman

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.20"
+version = "0.4.21"
 
 [deps]
 ArrayLayouts = "4c555306-a7a7-4459-81d9-ec55ddd5c99a"
@@ -26,7 +26,7 @@ GradedArraysTensorAlgebraExt = "TensorAlgebra"
 [compat]
 ArrayLayouts = "1"
 BlockArrays = "1.6"
-BlockSparseArrays = "0.8, 0.9.3"
+BlockSparseArrays = "0.9.5"
 Compat = "4.16"
 FillArrays = "1.13"
 HalfIntegers = "1.6"

ext/GradedArraysTensorAlgebraExt/GradedArraysTensorAlgebraExt.jl

@@ -62,6 +62,14 @@ function TensorAlgebra.unmatricize(
   m::AbstractMatrix,
   blocked_axes::BlockedTuple{2,<:Any,<:Tuple{Vararg{AbstractUnitRange}}},
 )
+  if isempty(blocked_axes)
+    # Handle edge case of empty blocked_axes, which can occur
+    # when matricizing a 0-dimensional array (a scalar).
+    a = similar(m, ())
+    a[] = only(m)
+    return a
+  end
+
   # First, fuse axes to get `sectormergesortperm`.
   # Then unpermute the blocks.
   fused_axes = matricize_axes(blocked_axes)

src/sectorunitrange.jl

@@ -159,3 +159,10 @@ end
 function Base.reshape(A::AbstractArray, ax::Tuple{SectorOneTo,Vararg{SectorOneTo}})
   return reshape(A, ungrade.(ax))
 end
+
+# Fixes issues when broadcasting over mixtures of arrays
+# where some have SectorOneTo axes and some have OneTo axes,
+# which can show up in BlockSparseArrays blockwise broadcasting.
+# See https://github.com/ITensor/GradedArrays.jl/pull/65.
+Base.Broadcast.axistype(r1::SectorOneTo, ::Base.OneTo) = r1
+Base.Broadcast.axistype(::Base.OneTo, r2::SectorOneTo) = r2

test/test_sectorunitrange.jl

@@ -46,7 +46,7 @@ using TestExtras: @constinferred
   @test length(sr) == 6
   @test firstindex(sr) == 1
   @test lastindex(sr) == 6
-  @test eltype(sr) === Int
+  @test eltype(sr) ≡ Int
   @test step(sr) == 1
   @test eachindex(sr) == Base.oneto(6)
   @test only(axes(sr)) isa SectorOneTo
@@ -59,9 +59,9 @@ using TestExtras: @constinferred
   @test isnothing(iterate(sr, 6))
 
   # Base.Slice
-  @test axes(Base.Slice(sr)) === (sr,)
-  @test Base.axes1(Base.Slice(sr)) === sr
-  @test Base.unsafe_indices(Base.Slice(sr)) === (sr,)
+  @test axes(Base.Slice(sr)) ≡ (sr,)
+  @test Base.axes1(Base.Slice(sr)) ≡ sr
+  @test Base.unsafe_indices(Base.Slice(sr)) ≡ (sr,)
 
   @test sr == 1:6
   @test sr == sr
@@ -113,8 +113,8 @@ using TestExtras: @constinferred
   @test blockisequal(sr, sr)
 
   # GradedUnitRanges interface
-  @test sector_type(sr) === SU{3,2}
-  @test sector_type(typeof(sr)) === SU{3,2}
+  @test sector_type(sr) ≡ SU{3,2}
+  @test sector_type(typeof(sr)) ≡ SU{3,2}
   @test sectors(sr) == [SU((1, 0))]
   @test sector_multiplicity(sr) == 2
   @test sector_multiplicities(sr) == [2]
@@ -138,7 +138,7 @@ using TestExtras: @constinferred
   end
   @test sr[2:3] == 2:3
   @test (@constinferred getindex(sr, 2:3)) isa UnitRange
-  @test sr[Block(1)] === sr
+  @test sr[Block(1)] ≡ sr
   @test_throws BlockBoundsError sr[Block(2)]
 
   sr2 = (@constinferred getindex(sr, (:, 2)))
@@ -162,5 +162,10 @@ using TestExtras: @constinferred
   # Slice sector range with sector range
   sr1 = sectorrange(U1(1), 4)
   sr2 = sectorrange(U1(1), 3)
-  @test sr1[sr2] === sr2
+  @test sr1[sr2] ≡ sr2
+
+  sr = sectorrange(U1(1), 4)
+  @test Broadcast.axistype(sr, sr) ≡ sr
+  @test Broadcast.axistype(sr, Base.OneTo(4)) ≡ sr
+  @test Broadcast.axistype(Base.OneTo(4), sr) ≡ sr
 end