Operand-driven diagonal algebra for the graded path

Project.toml

@@ -26,6 +26,10 @@ VectorInterface = "409d34a3-91d5-4945-b6ec-7529ddf182d8"
 [weakdeps]
 SUNRepresentations = "1a50b95c-7aac-476d-a9ce-2bfc675fc617"
 
+[sources.TensorAlgebra]
+rev = "main"
+url = "https://github.com/ITensor/TensorAlgebra.jl"
+
 [extensions]
 GradedArraysSUNRepresentationsExt = "SUNRepresentations"
 
@@ -42,7 +46,7 @@ SUNRepresentations = "0.3, 0.4"
 ScopedValues = "1"
 SplitApplyCombine = "1.2.3"
 StridedViews = "0.5"
-TensorAlgebra = "0.19.2"
+TensorAlgebra = "0.20"
 TensorKit = "0.17"
 TensorKitSectors = "0.3"
 VectorInterface = "0.4.8, 0.5"

src/GradedArrays.jl

@@ -59,13 +59,13 @@ include("fusedgradedvector.jl")
 # The graded broadcast-style lattice is defined here, before `FusedGradedDiagonal` (which adds its
 # own `FusedGradedDiagonalStyle` under the lattice) and `GradedArray` (which adds `GradedStyle`).
 include("broadcast.jl")
-include("fusedgradeddiagonal.jl")
 include("adjointfusedgradedarray.jl")
 include("fusedgradedblocks.jl")
 
 include("sectorproduct.jl")
 
 include("fusion.jl")
+include("fusedgradeddiagonal.jl")
 include("tensoralgebra.jl")
 include("cat.jl")
 

src/abstractfusedgradedarray.jl

@@ -60,7 +60,7 @@ end
 # generics live in `tensoralgebra.jl` (overloaded on `AbstractArray`, since `GradedArray` shares them
 # but is not an `AbstractFusedGradedArray`).
 
-function isblockdiagonal(A::AbstractFusedGradedMatrix)
+function isblockdiag(A::AbstractFusedGradedMatrix)
     for bI in eachblockstoredindex(A)
         row, col = Tuple(bI)
         row == col || return false
@@ -78,6 +78,23 @@ function LinearAlgebra.isdiag(A::AbstractFusedGradedMatrix)
     return true
 end
 
+# Square in the graded sense: equal codomain and domain axes, so every coupled block is square and sits
+# on the block diagonal. Stronger than equal overall dimensions, since a graded matrix can be
+# dimension-square with rectangular blocks. This is the meaningful "square" for a graded operator (its
+# main diagonal, an eigendecomposition, an endomorphism).
+issquare(A::AbstractFusedGradedMatrix) = axis_codomain(A) == axis_domain(A)
+
+# Throwing form of `issquare`, for use as a precondition guard. Our own (not `LinearAlgebra.checksquare`,
+# whose contract returns the matrix size); this one just checks and returns nothing.
+function checksquare(A::AbstractFusedGradedMatrix)
+    issquare(A) || throw(
+        DimensionMismatch(
+            "graded matrix is not square: its codomain and domain axes differ"
+        )
+    )
+    return nothing
+end
+
 # ---------------------------------------------------------------------------
 #  axes / size — derived from the per-variant `biaxes` core
 # ---------------------------------------------------------------------------

src/blocksparseinterface.jl

@@ -1,15 +1,17 @@
 # Block-sparse interface functions owned by GradedArrays.
 #
 # GradedArrays implements a block-sparse interface on its own graded array and axis types.
-# These names are duplicated with BlockSparseArrays by design: GradedArrays owns them here so
+# These names are mostly duplicated with BlockSparseArrays by design: GradedArrays owns them here so
 # it does not depend on BlockSparseArrays. They are internal (not exported); downstream reaches
-# them by qualified import, e.g. `using GradedArrays: eachblockstoredindex`.
+# them by qualified import, e.g. `using GradedArrays: eachblockstoredindex`. `isblockdiag` is a
+# deliberate exception: BlockSparseArrays spells it `isblockdiagonal`, but we match the name of
+# `LinearAlgebra.isdiag` instead.
 
 function eachblockstoredindex end
 function eachblockaxis end
 function mortar_axis end
 function blocktype end
-function isblockdiagonal end
+function isblockdiag end
 
 # The number of stored (symmetry-allowed) blocks. Counted from the stored block indices directly
 # rather than via `storedlength(blocks(a))`, whose generic `length(storedvalues(...))` would

src/fusedgradeddiagonal.jl

@@ -45,6 +45,15 @@ entries of the block at `sectors[i]`. The axis is derived from the blocks, as fo
 """
 fusedgradeddiagonal(sectordata) = FusedGradedDiagonal(fusedgradedvector(sectordata))
 
+# Densify to a full `FusedGradedMatrix`, forwarding the diagonal's block backend `V` through the
+# `{T,S,V}` undef constructor so the result stays on the same device. Each `Diagonal` block becomes a
+# dense block and `copyto!` zeros the off-diagonal. Used where a diagonal result is not representable
+# as a `FusedGradedDiagonal` (a non-`{1,1}` matricize / bond-split unmatricize).
+function FusedGradedMatrix(d::FusedGradedDiagonal{T, S, V}) where {T, S, V}
+    m = FusedGradedMatrix{T, S, V}(undef, axis_codomain(d), axis_domain(d))
+    return copyto!(m, d)
+end
+
 sectordata(d::FusedGradedDiagonal) = map(Diagonal, sectordata(MAK.diagview(d)))
 
 # ---- accessors ----
@@ -77,6 +86,31 @@ function TensorAlgebra.permuteddims(d::FusedGradedDiagonal, perm)
     return d
 end
 
+# ---- matricize ----
+
+# A `{1,1}` matricization of the diagonal is the identity (a diagonal is already a matrix). Any other
+# codomain rank bends a leg, which matrix-level fused storage cannot represent.
+TensorAlgebra.matricize(::GradedMatricize, d::FusedGradedDiagonal, ::Val{1}) = d
+function TensorAlgebra.matricize(
+        style::GradedMatricize, d::FusedGradedDiagonal, ndims_codomain::Val
+    )
+    throw(
+        ArgumentError(
+            "a matrix-level fused array matricizes only with a single codomain leg"
+        )
+    )
+end
+
+# The product of two diagonal fused matrices over a single contracted leg is again diagonal, so
+# allocate a `FusedGradedDiagonal` (the block-wise `mul!` fills it via `Diagonal * Diagonal`). Mixed
+# diagonal/dense products fall through to the general `AbstractFusedGradedMatrix` method (dense).
+function allocate_output(
+        ::typeof(*), A::FusedGradedDiagonal, B::FusedGradedDiagonal
+    )
+    Tout = Base.promote_op(*, eltype(A), eltype(B))
+    return FusedGradedDiagonal{Tout}(undef, axis_codomain(A))
+end
+
 # ---- broadcasting ----
 
 struct FusedGradedDiagonalStyle <: AbstractFusedGradedStyle{2} end

src/fusedgradedmatrix.jl

@@ -4,6 +4,14 @@
 
 using MatrixAlgebraKit: MatrixAlgebraKit as MAK
 
+# Length of the contiguous stored buffer: the sum of codomain-block times domain-block sizes over
+# the coupled sectors the two axes share.
+function fusedbufferlength(codomain::FusedGradedOneTo, domain::FusedGradedOneTo)
+    codl, doml = sectordatalengths(codomain), sectordatalengths(domain)
+    coupled = intersect(keys(codl), keys(doml))
+    return sum(c -> codl[c] * doml[c], coupled; init = 0)
+end
+
 """
     FusedGradedMatrix{T,S<:SectorRange,V<:DenseVector{T}}
 
@@ -33,9 +41,7 @@ struct FusedGradedMatrix{T, S <: SectorRange, V <: DenseVector{T}} <:
             )
         )
         # Validate the buffer length against the block total (SectorData does the same check on access).
-        codl, doml = sectordatalengths(cod), sectordatalengths(dom)
-        coupled = intersect(keys(codl), keys(doml))
-        total = sum(c -> codl[c] * doml[c], coupled; init = 0)
+        total = fusedbufferlength(cod, dom)
         length(buffer) == total ||
             throw(
             DimensionMismatch(
@@ -68,12 +74,21 @@ function FusedGradedMatrix{T}(
     ) where {T}
     cod = FusedGradedOneTo(codomain)
     dom = FusedGradedOneTo(domain)
-    codl, doml = sectordatalengths(cod), sectordatalengths(dom)
-    coupled = intersect(keys(codl), keys(doml))
-    buffer = Vector{T}(undef, sum(c -> codl[c] * doml[c], coupled; init = 0))
+    buffer = Vector{T}(undef, fusedbufferlength(cod, dom))
     return FusedGradedMatrix(buffer, cod, dom)
 end
 
+# Same as the `{T}` method but allocates the buffer as the given `V`, so a caller can forward its own
+# block backend (for example a GPU buffer) instead of defaulting to `Vector{T}`.
+function FusedGradedMatrix{T, S, V}(
+        ::UndefInitializer, codomain::AbstractGradedOneTo, domain::AbstractGradedOneTo
+    ) where {T, S, V}
+    cod = FusedGradedOneTo(codomain)
+    dom = FusedGradedOneTo(domain)
+    buffer = V(undef, fusedbufferlength(cod, dom))
+    return FusedGradedMatrix{T, S, V}(buffer, cod, dom)
+end
+
 """
     fusedgradedmatrix(sectors .=> data, codomain, domain)
     fusedgradedmatrix(sectordata::Dictionary, codomain, domain)
@@ -161,9 +176,13 @@ axes_codomain(m::FusedGradedMatrix) = (m.axis_codomain,)
 axes_domain(m::FusedGradedMatrix) = (m.axis_domain,)
 
 # The main diagonal as an owned `FusedGradedVector` whose block at each coupled sector is that block's
-# diagonal; the fresh buffer means writing it does not touch `m`. A write-through `diagview` of a
-# `FusedGradedMatrix` is not yet supported. Off-diagonals are unsupported.
+# diagonal; the fresh buffer means writing it does not touch `m`. Restricted to equal codomain and
+# domain axes (square blocks): only then do the per-block diagonals coincide with the matrix's main
+# diagonal. With rectangular blocks the dense diagonal drifts off the blocks into off-diagonal bands,
+# so concatenating per-block diagonals is a different operation; iterate blocks explicitly for that. A
+# write-through `diagview` of a `FusedGradedMatrix` is not yet supported. Off-diagonals are unsupported.
 function LinearAlgebra.diag(m::FusedGradedMatrix)
+    checksquare(m)
     return fusedgradedvector(map(MAK.diagview, sectordata(m)))
 end
 function LinearAlgebra.diag(m::FusedGradedMatrix, k::Integer)

src/gradedarray.jl

@@ -431,6 +431,14 @@ end
 
 TensorAlgebra.zero!(fa::GradedArray) = (zero!(matricize(fa)); fa)
 TensorAlgebra.scale!(fa::GradedArray, α::Number) = (scale!(matricize(fa), α); fa)
+function LinearAlgebra.rmul!(fa::GradedArray, α::Number)
+    LinearAlgebra.rmul!(matricize(fa), α)
+    return fa
+end
+function LinearAlgebra.lmul!(α::Number, fa::GradedArray)
+    LinearAlgebra.lmul!(α, matricize(fa))
+    return fa
+end
 LinearAlgebra.norm(fa::GradedArray, p::Real = 2) = LinearAlgebra.norm(matricize(fa), p)
 Base.fill!(fa::GradedArray, v) = (fill!(matricize(fa), v); fa)
 Base.iszero(fa::GradedArray) = iszero(matricize(fa))
@@ -638,6 +646,9 @@ function TensorAlgebra.unmatricize(
     return GradedArray(m, axes_codomain, axes_domain)
 end
 
+# A `{1,1}` unmatricize (one codomain axis, one domain axis) reproduces the diagonal's own square
+# bond and is the endomorphism identity: the result stays diagonal, wrapped up to the tensor-level
+# `GradedArray`. `check_input` rejects a mismatched single-axis split rather than densifying it.
 function TensorAlgebra.check_input(
         ::typeof(unmatricize), d::FusedGradedDiagonal, axes_codomain::Tuple, axes_domain::Tuple
     )
@@ -650,13 +661,21 @@ function TensorAlgebra.check_input(
     return nothing
 end
 
-# Keep the diagonal factor (`S`/`D`) bare so a spectrum is a true diagonal matrix.
 function TensorAlgebra.unmatricize(
-        ::GradedMatricize, d::FusedGradedDiagonal, axes_codomain::Tuple,
-        axes_domain::Tuple
+        ::GradedMatricize, d::FusedGradedDiagonal,
+        axes_codomain::Tuple{<:AbstractGradedOneTo}, axes_domain::Tuple{<:AbstractGradedOneTo}
     )
     check_input(unmatricize, d, axes_codomain, axes_domain)
-    return d
+    return GradedArray(d, axes_codomain, axes_domain)
+end
+
+# Any other split is a genuine bond-split (for example a rank-4 generalized-diagonal tensor), not
+# representable as a `FusedGradedDiagonal`, so densify to a `FusedGradedMatrix` and reconstruct
+# through its own `unmatricize`.
+function TensorAlgebra.unmatricize(
+        style::GradedMatricize, d::FusedGradedDiagonal, axes_codomain::Tuple, axes_domain::Tuple
+    )
+    return unmatricize(style, FusedGradedMatrix(d), axes_codomain, axes_domain)
 end
 
 # ============================  contraction  ============================
@@ -688,7 +707,8 @@ function TensorAlgebra.matricize(
 end
 
 function TensorAlgebra.unmatricizeperm!(
-        ::GradedMatricize, a_dest::GradedArray{<:Any, <:Any, N}, m::FusedGradedMatrix,
+        ::GradedMatricize, a_dest::GradedArray{<:Any, <:Any, N},
+        m::AbstractFusedGradedMatrix,
         invperm_codomain::Tuple{Vararg{Int}}, invperm_domain::Tuple{Vararg{Int}}
     ) where {N}
     # Permute `a_dest` into the matricized leg order to get the matricized-order axes with correct

src/matrixalgebrakit.jl

@@ -1,4 +1,5 @@
 using MatrixAlgebraKit: MatrixAlgebraKit as MAK
+using TensorAlgebra: MatrixAlgebra as MA
 
 # Length of the main diagonal of a matrix (e.g. the number of singular values a block produces).
 diaglength(a::AbstractArray) = minimum(size(a))
@@ -100,7 +101,7 @@ for f! in (
         :left_polar!, :right_polar!,
     )
     @eval function MAK.$f!(A::FusedGradedMatrix, F, alg::FusedGradedMatrixAlgorithm)
-        $(f! in (:eig_full!, :eigh_full!) && :(LinearAlgebra.checksquare(A)))
+        $(f! in (:eig_full!, :eigh_full!) && :(checksquare(A)))
         for c in eachsector(A, F...)
             Ac = getsectordata(A, c)
             Fc = map(x -> getsectordata(x, c), F)
@@ -118,7 +119,7 @@ for f! in (
         :project_isometric!,
     )
     @eval function MAK.$f!(A::FusedGradedMatrix, N, alg::FusedGradedMatrixAlgorithm)
-        $(f! in (:eig_vals!, :eigh_vals!) && :(LinearAlgebra.checksquare(A)))
+        $(f! in (:eig_vals!, :eigh_vals!) && :(checksquare(A)))
         for c in eachsector(A, N)
             Ac = getsectordata(A, c)
             Nc = getsectordata(N, c)
@@ -133,7 +134,7 @@ end
 # with the same block structure as the input, so they iterate the stored blocks directly.
 for f! in (:project_hermitian!, :project_antihermitian!)
     @eval function MAK.$f!(A::FusedGradedMatrix, out, alg::FusedGradedMatrixAlgorithm)
-        LinearAlgebra.checksquare(A)
+        checksquare(A)
         for I in eachblockstoredindex(A)
             MAK.$f!(view(A, I), view(out, I), FusedSectorMatrixAlgorithm(alg.alg))
         end
@@ -390,27 +391,22 @@ MAK.diagview(d::FusedGradedDiagonal) = d.diag
 # falling through to LinearAlgebra's scalar-indexing `Diagonal*Matrix` impl.
 MAK.diagonal(v::FusedGradedVector) = FusedGradedDiagonal(v)
 
-# `pow_diag_safe!` for a block-diagonal graded matrix: clamp-power each reduced diagonal
-# block. Only the reduced (degeneracy) data is touched, and that is correct even in the
-# non-abelian case: a diagonal factor is `Diagonal(λ) ⊗ I` per sector, and `f(A ⊗ I) =
-# f(A) ⊗ I`, so the power passes straight to the reduced eigenvalues. This is why the
-# diagonal power is well defined here whereas a general element-wise `map!` on a graded
-# array is not.
-function TensorAlgebra.MatrixAlgebra.pow_diag_safe!(
-        Dp::FusedGradedDiagonal, D::FusedGradedDiagonal, p, tol
+# `pow_diag_safe!` for a graded matrix that is diagonal: a `FusedGradedDiagonal`, or a
+# `FusedGradedMatrix` that happens to be runtime-diagonal (the `isdiag` fast path in
+# `sqrth_invsqrth_safe` powers such a matrix directly, without an eigendecomposition).
+# Delegating per reduced block reuses the generic diagonal-only kernel, which is correct even
+# in the non-abelian case: a diagonal factor is `Diagonal(λ) ⊗ I` per sector, and `f(A ⊗ I) =
+# f(A) ⊗ I`, so the power passes straight to the reduced eigenvalues. This is why the diagonal
+# power is well defined here whereas a general element-wise `map!` on a graded array is not.
+function MA.pow_diag_safe!(
+        Dp::AbstractFusedGradedMatrix, D::AbstractFusedGradedMatrix, p, tol
     )
-    dp, d = MAK.diagview(Dp), MAK.diagview(D)
-    for c in eachsector(d)
-        map!(x -> _clamped_pow(x, p, tol), sectordata(dp, c), sectordata(d, c))
+    for c in eachsector(D)
+        MA.pow_diag_safe!(sectordata(Dp, c), sectordata(D, c), p, tol)
     end
     return Dp
 end
 
-# Vendored from `TensorAlgebra.MatrixAlgebra` (not part of its public API): clamp entries
-# below `tol` to zero, then raise to `p`; a negative entry above `tol` lets `real(d)^p`
-# error for fractional `p`, enforcing the PSD precondition per-power.
-_clamped_pow(d, p, tol) = abs(d) < tol ? zero(d) : real(d)^p
-
 # Count how many elements are kept for a given index specification and block size
 _count_kept(::Colon, n) = n
 _count_kept(ind::AbstractVector{Bool}, _) = count(ind)

test/Project.toml

@@ -34,7 +34,7 @@ SUNRepresentations = "0.3, 0.4"
 SafeTestsets = "0.1"
 StableRNGs = "1"
 Suppressor = "0.2.8"
-TensorAlgebra = "0.19.2"
+TensorAlgebra = "0.20"
 TensorKit = "0.17"
 TensorKitSectors = "0.3"
 Test = "1.10"

test/test_factorizations.jl

@@ -603,11 +603,24 @@ end
 
             P = sqrth_safe(S, (1,), (2,))
             Pinv = invsqrth_safe(S, (1,), (2,))
-            @test P isa FusedGradedDiagonal
-            @test Pinv isa FusedGradedDiagonal
+            # The pattern-taking form ends in `unmatricize`, which wraps a diagonal `{1,1}`
+            # result up to the tensor-level `GradedArray`. The fast
+            # `pow_diag_safe` path still runs, so the backing stays a `FusedGradedDiagonal`
+            # rather than densifying to the eigenvalue-power path.
+            @test matricize(P) isa FusedGradedDiagonal
+            @test matricize(Pinv) isa FusedGradedDiagonal
             for (i, s) in enumerate(sects)
-                @test diag(sectordata(P)[s]) ≈ sqrt.(svals[i])
-                @test diag(sectordata(Pinv)[s]) ≈ inv.(sqrt.(svals[i]))
+                @test diag(sectordata(matricize(P))[s]) ≈ sqrt.(svals[i])
+                @test diag(sectordata(matricize(Pinv))[s]) ≈ inv.(sqrt.(svals[i]))
+            end
+
+            # A dense-stored runtime-diagonal matrix takes the same fast path, with
+            # `pow_diag_safe!` delegating per block.
+            M = FusedGradedMatrix(S)
+            Pm = sqrth_safe(M, (1,), (2,))
+            @test matricize(Pm) isa FusedGradedMatrix
+            for (i, s) in enumerate(sects)
+                @test diag(sectordata(matricize(Pm))[s]) ≈ sqrt.(svals[i])
             end
         end
     end