Strict axis equality, remove space_isequal (#138)

Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: mtfishman

Project.toml

@@ -1,6 +1,6 @@
 name = "GradedArrays"
 uuid = "bc96ca6e-b7c8-4bb6-888e-c93f838762c2"
-version = "0.6.23"
+version = "0.7.0"
 authors = ["ITensor developers <support@itensor.org> and contributors"]
 
 [workspace]

docs/Project.toml

@@ -9,6 +9,6 @@ path = ".."
 
 [compat]
 Documenter = "1"
-GradedArrays = "0.6"
+GradedArrays = "0.7"
 ITensorFormatter = "0.2.27"
 Literate = "2"

examples/Project.toml

@@ -5,4 +5,4 @@ GradedArrays = "bc96ca6e-b7c8-4bb6-888e-c93f838762c2"
 path = ".."
 
 [compat]
-GradedArrays = "0.6"
+GradedArrays = "0.7"

src/GradedArrays.jl

@@ -11,7 +11,7 @@ export gradedrange
 export dual, flip, gradedrange, isdual,
     sector, sector_multiplicities, sector_multiplicity,
     sectorrange, sectors, sector_type,
-    space_isequal, ungrade
+    ungrade
 
 # imports
 # -------

src/broadcast.jl

@@ -31,29 +31,11 @@ function Base.Broadcast.materialize(a::SectorArray)
     return ofsector(a, Base.Broadcast.materialize(a.data))
 end
 
-function check_sector_broadcast_axes(a::SectorArray, b::SectorArray)
-    axes(a) == axes(b) ||
-        throw(ArgumentError("SectorArray linear broadcasting requires matching axes"))
-    return nothing
-end
-
 function TensorAlgebra.:+ₗ(a::SectorArray, b::SectorArray)
-    check_sector_broadcast_axes(a, b)
+    _check_add_axes(a, b)
     return ofsector(a, a.data +ₗ b.data)
 end
 
-function TensorAlgebra.add!(dest::AbstractArray, src::SectorArray, α::Number, β::Number)
-    require_unique_fusion(src)
-    TensorAlgebra.add!(dest, src.data, α, β)
-    return dest
-end
-
-function TensorAlgebra.add!(dest::SectorArray, src::SectorArray, α::Number, β::Number)
-    check_sector_broadcast_axes(dest, src)
-    TensorAlgebra.add!(dest.data, src.data, α, β)
-    return dest
-end
-
 function TensorAlgebra.:*ₗ(α::Number, a::SectorArray)
     return ofsector(a, α *ₗ a.data)
 end
@@ -108,10 +90,10 @@ function Base.similar(bc::BC.Broadcasted{<:GradedStyle}, elt::Type, ax)
     return graded_similar(arg, elt, ax)
 end
 
-function check_graded_broadcast_axes(a::AbstractArray, b::AbstractArray)
-    all(dim -> space_isequal(axes(a, dim), axes(b, dim)), 1:ndims(a)) ||
+function _check_add_axes(a::AbstractArray, b::AbstractArray)
+    axes(a) == axes(b) ||
         throw(
-        ArgumentError("GradedArray linear broadcasting requires matching graded axes")
+        ArgumentError("linear broadcasting requires matching axes")
     )
     return nothing
 end
@@ -211,7 +193,7 @@ function copy_lazygraded(a::LazyGradedArray)
 end
 
 function TensorAlgebra.:+ₗ(a::LazyGradedArray, b::LazyGradedArray)
-    check_graded_broadcast_axes(a, b)
+    _check_add_axes(a, b)
     return AddGradedArray(a, b)
 end
 TensorAlgebra.:*ₗ(α::Number, a::GradedArray) = ScaledGradedArray(α, a)

src/gradedarray.jl

@@ -45,9 +45,16 @@ function KroneckerArrays.:×(g1::GradedUnitRange, g2::GradedUnitRange)
     return mortar_axis(v)
 end
 
-function space_isequal(a1::AbstractUnitRange, a2::AbstractUnitRange)
-    return (isdual(a1) == isdual(a2)) && sectors(a1) == sectors(a2) && blockisequal(a1, a2)
+function Base.isequal(a::GradedUnitRange, b::GradedUnitRange)
+    ea = eachblockaxis(a)
+    eb = eachblockaxis(b)
+    return length(ea) == length(eb) && all(splat(isequal), zip(ea, eb))
 end
+Base.:(==)(a::GradedUnitRange, b::GradedUnitRange) = isequal(a, b)
+Base.:(==)(::GradedUnitRange, ::AbstractUnitRange) = false
+Base.:(==)(::AbstractUnitRange, ::GradedUnitRange) = false
+Base.isequal(::GradedUnitRange, ::AbstractUnitRange) = false
+Base.isequal(::AbstractUnitRange, ::GradedUnitRange) = false
 
 function BlockSparseArrays.blockrange(xs::Vector{<:GradedUnitRange})
     baxis = mapreduce(eachblockaxis, vcat, xs)
@@ -123,6 +130,21 @@ const GradedArray{T, N, I, A, Blocks, Axes <: NTuple{N, GradedUnitRange{I}}} =
 const GradedMatrix{T, I, A, Blocks, Axes} = GradedArray{T, 2, A, Blocks, Axes}
 const GradedVector{T, I, A, Blocks, Axes} = GradedArray{T, 1, A, Blocks, Axes}
 
+# Override ArrayLayouts._check_mul_axes for GradedArray.
+# For graded matrices, the contracted axes satisfy axes(A,2) == dual(axes(B,1)),
+# not axes(A,2) == axes(B,1), so the default check needs to account for duality.
+const GradedMatrixOrAdj = Union{
+    GradedMatrix, LinearAlgebra.AdjOrTrans{<:Any, <:GradedMatrix},
+}
+function ArrayLayouts._check_mul_axes(A::GradedMatrixOrAdj, B::GradedMatrixOrAdj)
+    axes(A, 2) == dual(axes(B, 1)) || throw(
+        DimensionMismatch(
+            "second axis of A, $(axes(A, 2)), and first axis of B, $(axes(B, 1)), must match"
+        )
+    )
+    return nothing
+end
+
 # Specific overloads
 # ------------------
 # convert Array to SectorArray upon insertion

src/sectorarray.jl

@@ -36,6 +36,13 @@ end
 KroneckerArrays.:×(a::SectorRange, g::AbstractUnitRange) = cartesianrange(a, g)
 KroneckerArrays.:×(g::AbstractUnitRange, a::SectorRange) = cartesianrange(a, g)
 
+function Base.isequal(a::SectorUnitRange, b::SectorUnitRange)
+    return isequal(kroneckerfactors(a), kroneckerfactors(b))
+end
+function Base.:(==)(a::SectorUnitRange, b::SectorUnitRange)
+    return isequal(a, b)
+end
+
 to_gradedrange(g::SectorUnitRange) = mortar_axis([g])
 
 """
@@ -385,10 +392,10 @@ end
 # TODO: Define this as part of:
 # `check_input(::typeof(mul!), ::SectorMatrix, ::SectorMatrix, ::SectorMatrix)`
 function check_mul_axes(c::SectorMatrix, a::SectorMatrix, b::SectorMatrix)
-    space_isequal(axes(a, 2), dual(axes(b, 1))) ||
+    axes(a, 2) == dual(axes(b, 1)) ||
         throw(DimensionMismatch("$(axes(a, 2)) != dual($(axes(b, 1))))"))
-    space_isequal(axes(c, 1), axes(a, 1)) || throw(DimensionMismatch())
-    space_isequal(axes(c, 2), axes(b, 2)) || throw(DimensionMismatch())
+    axes(c, 1) == axes(a, 1) || throw(DimensionMismatch())
+    axes(c, 2) == axes(b, 2) || throw(DimensionMismatch())
     return nothing
 end
 
@@ -413,6 +420,18 @@ function KroneckerArrays.:(⊗)(
     return SectorArray(A.sectors, collect(T, data))
 end
 
+function TensorAlgebra.add!(dest::AbstractArray, src::SectorArray, α::Number, β::Number)
+    require_unique_fusion(src)
+    TensorAlgebra.add!(dest, src.data, α, β)
+    return dest
+end
+
+function TensorAlgebra.add!(dest::SectorArray, src::SectorArray, α::Number, β::Number)
+    _check_add_axes(dest, src)
+    TensorAlgebra.add!(dest.data, src.data, α, β)
+    return dest
+end
+
 # TODO: can we avoid this?
 function Base.materialize!(dst::SectorArray, src::KroneckerArrays.KroneckerBroadcasted)
     Base.materialize!(kroneckerfactors(dst, 1), kroneckerfactors(src, 1))

test/Project.toml

@@ -25,7 +25,7 @@ path = ".."
 Aqua = "0.8.11"
 BlockArrays = "1.6"
 BlockSparseArrays = "0.10"
-GradedArrays = "0.6"
+GradedArrays = "0.7"
 ITensorPkgSkeleton = "0.3.42"
 KroneckerArrays = "0.3"
 LinearAlgebra = "1.10"

test/test_exports.jl

@@ -30,7 +30,6 @@ using Test: @test, @testset
         :sectorrange,
         :sectors,
         :sector_type,
-        :space_isequal,
         :ungrade,
     ]
     @test issetequal(names(GradedArrays), exports)