Implement the TensorAlgebra 0.17 projection and matrix-function interfaces

Project.toml

@@ -1,6 +1,6 @@
 name = "GradedArrays"
 uuid = "bc96ca6e-b7c8-4bb6-888e-c93f838762c2"
-version = "0.13.7"
+version = "0.13.8"
 authors = ["ITensor developers <support@itensor.org> and contributors"]
 
 [workspace]
@@ -44,7 +44,7 @@ SUNRepresentations = "0.3, 0.4"
 SparseArraysBase = "0.10"
 SplitApplyCombine = "1.2.3"
 StridedViews = "0.5"
-TensorAlgebra = "0.16"
+TensorAlgebra = "0.17"
 TensorKit = "0.17"
 TensorKitSectors = "0.3"
 julia = "1.10"

src/abeliangradedarray.jl

@@ -485,15 +485,26 @@ function LinearAlgebra.isdiag(A::AbelianGradedMatrix)
     return true
 end
 
+# Throw unless `sz1` and `sz2` are equal ignoring trailing length-1 axes (an axis beyond one
+# size's rank counts as length 1, mirroring `Base.size(A, d)` for `d > ndims(A)`), guarding a
+# `reshape` against silently reinterpreting same-length data of a genuinely different shape.
+# Vendored from TensorAlgebra rather than reused because it is not part of its public API.
+function check_project_size(sz1::Dims, sz2::Dims)
+    all(i -> get(sz1, i, 1) == get(sz2, i, 1), 1:max(length(sz1), length(sz2))) || throw(
+        DimensionMismatch("sizes $sz1 and $sz2 differ beyond trailing length-1 axes")
+    )
+    return nothing
+end
+
 # Orthogonal projection of a dense source into the symmetry-allowed subspace.
 # Magnitude-blind: forbidden-block entries of `src` are dropped without inspection.
-# Use `TensorAlgebra.checked_projectto!` to verify the discarded weight is small.
+# The `TensorAlgebra.project` wrapper verifies the discarded weight is small.
 function TensorAlgebra.projectto!(dest::AbelianGradedArray, src::AbstractArray)
-    size(dest) == size(src) || throw(
-        DimensionMismatch(
-            "projectto!: dest has size $(size(dest)), src has size $(size(src))"
-        )
-    )
+    # Reshape `src` to `size(dest)` (a no-op when the ranks already match), so a lower-rank
+    # `src` may omit trailing length-1 axes (e.g. an auxiliary flux-canceling leg); a genuine
+    # shape mismatch errors rather than reinterpreting the data.
+    check_project_size(size(src), size(dest))
+    src = reshape(src, size(dest))
     zero!(dest)
     for b in allowedblocks(axes(dest))
         block_ranges = ntuple(d -> axes(dest, d)[Block(Int(Tuple(b)[d]))], ndims(dest))
@@ -502,20 +513,85 @@ function TensorAlgebra.projectto!(dest::AbelianGradedArray, src::AbstractArray)
     return dest
 end
 
-# Compare via `Array(dest)` so the generic `isapprox(::AbstractArray, ::AbelianGradedArray)`
-# path doesn't fall back to a `src - dest` broadcast that scalar-indexes the block storage.
-# `kwargs` (e.g. `atol`/`rtol`) are forwarded to `isapprox`, matching the generic verb.
-function TensorAlgebra.checked_projectto!(
-        dest::AbelianGradedArray, src::AbstractArray; kwargs...
+# `allocate_project` with graded axes routes both the codomain-led and the (empty-codomain)
+# domain-led cases to `allocate_project_graded`, taking the graded structure from whichever side
+# is non-empty, the same two-entry split `similar_map` uses.
+function TensorAlgebra.allocate_project(
+        src::AbstractArray,
+        codomain_axes::Tuple{GradedOneTo, Vararg{GradedOneTo}},
+        domain_axes::Tuple{Vararg{GradedOneTo}}
     )
-    TensorAlgebra.projectto!(dest, src)
-    isapprox(src, Array(dest); kwargs...) ||
-        throw(InexactError(:checked_projectto!, typeof(dest), src))
-    return dest
+    return allocate_project_graded(src, codomain_axes, domain_axes)
+end
+function TensorAlgebra.allocate_project(
+        src::AbstractArray,
+        codomain_axes::Tuple{},
+        domain_axes::Tuple{GradedOneTo, Vararg{GradedOneTo}}
+    )
+    return allocate_project_graded(src, codomain_axes, domain_axes)
+end
+
+# The destination allocation, and where a trailing surplus axis in `src` gets its space derived.
+# With no surplus this is plain `similar_map`; a single trailing surplus axis is an auxiliary leg
+# appended as the last domain axis, its space derived (see `infer_aux_space`) so the result is
+# symmetry-allowed. The trailing position matches how `stack` lays out an operator multiplet and
+# the Julia convention that trailing length-1 axes are implicit. For the matricized result `X` the
+# Gram matrix `X * X'` contracts the aux away (for a spin multiplet, `X * X' = S·S`, the Casimir).
+function allocate_project_graded(src, codomain_axes, domain_axes)
+    nphys = length(codomain_axes) + length(domain_axes)
+    ndims(src) <= nphys &&
+        return TensorAlgebra.similar_map(src, codomain_axes, domain_axes)
+    ndims(src) == nphys + 1 || throw(
+        ArgumentError(
+            "`project`: expected at most one trailing auxiliary axis beyond the $nphys \
+            given axes, got a rank-$(ndims(src)) input"
+        )
+    )
+    aux = infer_aux_space(src, codomain_axes, domain_axes)
+    return TensorAlgebra.similar_map(src, codomain_axes, (domain_axes..., aux))
+end
+
+# The space of `src`'s trailing auxiliary axis, derived so the projected result is symmetry-
+# allowed. Abelian sectors are one-dimensional, so each length-1 slice along the aux axis carries
+# a single charge (`projected_charge`); contiguous equal charges merge into one sector of that
+# multiplicity (matching the `TensorMap` backend), while non-contiguous repeats stay separate to
+# preserve slice order. A multi-slice aux comes out as a direct sum (e.g. stacking `[S⁺, S⁻]`
+# gives `[U1(2), U1(-2)]`, an MPO-virtual-leg structure).
+function infer_aux_space(src::AbstractArray, codomain_axes, domain_axes)
+    aux_dim = length(codomain_axes) + length(domain_axes) + 1
+    qs = map(eachslice(src; dims = aux_dim)) do slice
+        return projected_charge(slice, codomain_axes, domain_axes)
+    end
+    ps = Pair{eltype(qs), Int}[]
+    for q in qs
+        if isempty(ps) || first(ps[end]) != q
+            push!(ps, q => 1)
+        else
+            ps[end] = q => (last(ps[end]) + 1)
+        end
+    end
+    return gradedrange(ps)
+end
+
+# Net charge of a dense operator, read from its dominant-magnitude entry: find the block
+# holding that entry over the stored axes (domain dualized to match `zeros_map`/`similar_map`)
+# and fuse that block's per-axis sectors, each with its axis's arrow applied (the same fusion
+# `allowedblocks` is built on, so the charge lines up with which blocks `project` keeps). This is
+# the abelian per-slice primitive `infer_aux_space` builds the aux space from; the general
+# (possibly multi-sector) derivation lives in the `TensorMap` backend.
+function projected_charge(src::AbstractArray, codomain_axes, domain_axes)
+    stored = (codomain_axes..., conj.(domain_axes)...)
+    src = reshape(src, length.(stored))
+    I = Tuple(findmax(abs, src)[2])
+    secs = map(stored, I) do ax, i
+        return eachsectoraxis(ax)[Int(BlockArrays.findblock(ax, i))]
+    end
+    return reduce(tensor_product, secs)
 end
 
-# Materialize the graded array into a dense `Array` for the default
-# `checked_projectto!`/`isapprox`-after path.
+# Materialize the graded array into a dense `Array`. The checked projection verbs reach this
+# through `convert(Array, dest)` to compare elementwise against a dense source, which would
+# otherwise fall back to a `src - dest` broadcast that scalar-indexes the block storage.
 function Base.Array(a::AbelianGradedArray{T, <:Any, N}) where {T, N}
     dest = zeros(T, size(a))
     for bI in eachblockstoredindex(a)
@@ -727,17 +803,18 @@ end
 
 Construct an `AbelianGradedArray` by projecting the dense data of `a` onto the
 symmetry-allowed blocks of the graded axes `(ax1, axs...)`, via
-`TensorAlgebra.checked_projectto!` (which errors if `a` has weight outside
+`TensorAlgebra.project` (which errors if `a` has weight outside
 the allowed blocks). `a` is reshaped to `length.((ax1, axs...))` first, so a
 trailing size-1 bond can be supplied implicitly. Each axis carries its own arrow,
 so index with `dual`/`conj` axes to set duality.
 """
 function Base.getindex(a::AbstractArray, ax1::GradedOneTo, axs::GradedOneTo...)
     dest_axes = (ax1, axs...)
-    a_reshaped = reshape(a, length.(dest_axes))
-    return TensorAlgebra.checked_projectto!(
-        similar(a_reshaped, eltype(a), dest_axes), a_reshaped
-    )
+    # Match `a` to the requested axes up to trailing length-1 axes: indexing selects exactly these
+    # axes, so the surplus-axis derivation branch of `project` must not trigger, and a genuine
+    # shape mismatch errors rather than reinterpreting the data.
+    check_project_size(size(a), length.(dest_axes))
+    return TensorAlgebra.project(reshape(a, length.(dest_axes)), dest_axes)
 end
 # Disambiguate the single-axis case for a concrete `Array`: `Base.getindex(::Array,
 # ::AbstractUnitRange{<:Integer})` and the projection method above are otherwise equally

src/abstractgradedarray.jl

@@ -117,6 +117,11 @@ function TensorAlgebra.scale!(a::AbstractGradedArray, β::Number)
     return a
 end
 
+# The `LinearAlgebra` spelling of blockwise scaling (the generic fallback
+# scalar-indexes).
+LinearAlgebra.rmul!(a::AbstractGradedArray, β::Number) = TensorAlgebra.scale!(a, β)
+LinearAlgebra.lmul!(β::Number, a::AbstractGradedArray) = TensorAlgebra.scale!(a, β)
+
 function TensorAlgebra.zero!(a::AbstractGradedArray)
     for bI in eachblockstoredindex(a)
         zero!(view(a, bI))

src/fusedgradedmatrix.jl

@@ -69,6 +69,27 @@ end
 # Block-diagonal by construction (one block per sector), so just check each block.
 LinearAlgebra.isdiag(A::FusedGradedMatrix) = all(LinearAlgebra.isdiag, A.blocks)
 
+# Blockwise copy: the generic `AbstractArray` fallback copies elementwise, which
+# scalar-indexes (disallowed for graded arrays).
+function Base.copy(A::FusedGradedMatrix)
+    return FusedGradedMatrix(A.codomain, A.domain, map(copy, A.blocks))
+end
+
+# Materialize into a dense `Array` blockwise, like the `AbelianGradedArray` method
+# (the generic fallback copies elementwise, which scalar-indexes). Sectors present
+# in only one of the two axes have no stored block and stay zero.
+function Base.Array(a::FusedGradedMatrix{T}) where {T}
+    dest = zeros(T, size(a))
+    rowax, colax = axes(a)
+    rowsectors, colsectors = collect(keys(a.codomain)), collect(keys(a.domain))
+    for (s, b) in pairs(a.blocks)
+        r = rowax[Block(findfirst(==(s), rowsectors))]
+        c = colax[Block(findfirst(==(s), colsectors))]
+        copyto!(view(dest, r, c), b)
+    end
+    return dest
+end
+
 # Block-diagonal by construction, so any matrix function `f(A) = blkdiag(f(blk_i))` for
 # each stored block — covers `sqrt`, `exp`, `log`, etc. Routes around the generic
 # `LinearAlgebra` impls that scalar-index for triangular / Hermitian detection.

src/fusedgradedvector.jl

@@ -149,6 +149,18 @@ end
 
 # ========================  Accessors  ========================
 
+# Blockwise copyto!: the generic `AbstractArray` fallback copies elementwise, which
+# scalar-indexes (disallowed for graded arrays). Also the write path for mutating a
+# `MAK.diagview(::FusedGradedMatrix)` (whose blocks alias the matrix diagonals).
+function Base.copyto!(dest::FusedGradedVector, src::FusedGradedVector)
+    keys(dest.blocks) == keys(src.blocks) ||
+        throw(ArgumentError("`copyto!` requires matching sectors"))
+    for s in keys(src.blocks)
+        copyto!(dest.blocks[s], src.blocks[s])
+    end
+    return dest
+end
+
 BlockArrays.blocklength(v::FusedGradedVector) = length(v.axis)
 
 function blocktype(::Type{<:FusedGradedVector{T, S, D}}) where {T, S, D}
@@ -276,6 +288,18 @@ function Base.copy(v::FusedGradedVector)
     return FusedGradedVector(copy(v.axis), map(copy, v.blocks))
 end
 
+# Materialize into a dense `Array` blockwise, like the `AbelianGradedArray` method
+# (the generic fallback copies elementwise, which scalar-indexes).
+function Base.Array(v::FusedGradedVector{T}) where {T}
+    dest = zeros(T, size(v))
+    ax = only(axes(v))
+    sectors = collect(keys(v.axis))
+    for (s, b) in pairs(v.blocks)
+        copyto!(view(dest, ax[Block(findfirst(==(s), sectors))]), b)
+    end
+    return dest
+end
+
 # ======================== LinearAlgebra ======================
 
 function LinearAlgebra.norm(A::FusedGradedVector, p::Real = 2)

src/gradedoneto.jl

@@ -63,6 +63,9 @@ end
 trivial(g::GradedOneTo) = trivial(typeof(g))
 
 TensorAlgebra.trivialrange(R::Type{<:GradedOneTo}) = trivial(R)
+function TensorAlgebra.trivialrange(::Type{GradedOneTo{S}}, n::Integer) where {S}
+    return gradedrange([trivial(S) => n])
+end
 
 """
     gradedrange(xs::AbstractVector{<:Pair})

src/matrixalgebrakit.jl

@@ -12,6 +12,7 @@ for f in [
         :lq_compact, :lq_full, :lq_null,
         :eig_full, :eig_vals, :eigh_full, :eigh_vals,
         :left_polar, :right_polar,
+        :project_hermitian, :project_antihermitian, :project_isometric,
     ]
     f! = Symbol(f, :!)
     @eval function MAK.default_algorithm(
@@ -283,6 +284,32 @@ function MAK.initialize_output(
     return P, Wᴴ
 end
 
+# Projections
+# -----------
+# Same output conventions as the generic implementations: hermitian and
+# antihermitian project in place, isometric writes to a fresh output.
+function MAK.initialize_output(
+        ::typeof(MAK.project_hermitian!),
+        A::FusedGradedMatrix,
+        alg::GradedBlockAlgorithm
+    )
+    return A
+end
+function MAK.initialize_output(
+        ::typeof(MAK.project_antihermitian!),
+        A::FusedGradedMatrix,
+        alg::GradedBlockAlgorithm
+    )
+    return A
+end
+function MAK.initialize_output(
+        ::typeof(MAK.project_isometric!),
+        A::FusedGradedMatrix,
+        alg::GradedBlockAlgorithm
+    )
+    return similar(A)
+end
+
 # Truncation support
 # ------------------
 
@@ -303,6 +330,26 @@ function MAK.diagonal(v::FusedGradedVector)
     return FusedGradedMatrix(v.axis, v.axis, diag_blocks)
 end
 
+# `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::FusedGradedMatrix, D::FusedGradedMatrix, p, tol
+    )
+    foreachblock(MAK.diagview(Dp), MAK.diagview(D)) do _, (σp, σ)
+        return map!(d -> _clamped_pow(d, p, tol), σp, σ)
+    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

@@ -33,7 +33,7 @@ SUNRepresentations = "0.3, 0.4"
 SafeTestsets = "0.1"
 SparseArraysBase = "0.10"
 Suppressor = "0.2.8"
-TensorAlgebra = "0.16"
+TensorAlgebra = "0.17"
 TensorKit = "0.17"
 TensorKitSectors = "0.3"
 Test = "1.10"