Block-aware methods on graded storage

Project.toml

@@ -1,6 +1,6 @@
 name = "GradedArrays"
 uuid = "bc96ca6e-b7c8-4bb6-888e-c93f838762c2"
-version = "0.9.1"
+version = "0.9.2"
 authors = ["ITensor developers <support@itensor.org> and contributors"]
 
 [workspace]
@@ -50,7 +50,7 @@ Random = "1.10"
 SUNRepresentations = "0.3, 0.4"
 SparseArraysBase = "0.9"
 SplitApplyCombine = "1.2.3"
-TensorAlgebra = "0.8, 0.9"
+TensorAlgebra = "0.9.5"
 TensorKitSectors = "0.3"
 TypeParameterAccessors = "0.4"
 julia = "1.10"

src/GradedArrays.jl

@@ -22,15 +22,16 @@ export dual, flip, gradedrange, isdual,
 
 # imports
 # -------
-import FunctionImplementations as FI
 using BlockArrays: BlockArrays, AbstractBlockArray, AbstractBlockVector,
     AbstractBlockedUnitRange, Block, BlockIndexRange, BlockVector, BlockedOneTo,
     blockedrange, blocklasts, blocklength, blocklengths, blocks, eachblockaxes1
 using BlockSparseArrays: BlockSparseArrays, blockdiagindices, blockstoredlength,
     eachblockaxis, eachblockstoredindex, mortar_axis
 using Dictionaries: Dictionaries, Dictionary, dictionary, gettoken, gettokenvalue
+using FunctionImplementations: FunctionImplementations as FI
 using KroneckerArrays: KroneckerArrays, kroneckerfactors, ×, ⊗
 using LinearAlgebra: LinearAlgebra, Adjoint, Diagonal, mul!
+using Random: Random, AbstractRNG
 using SparseArraysBase: SparseArraysBase
 using TensorAlgebra: TensorAlgebra, BlockedTuple, FusionStyle, bipermutedimsopadd!,
     check_input, matricize, matricize_axes, permutedimsadd!, tensor_product_axis,

src/abeliangradedarray.jl

@@ -18,6 +18,8 @@ struct AbelianGradedArray{T, N, D <: AbstractArray{T, N}, S <: SectorRange} <:
     axes::NTuple{N, GradedOneTo{S}}
 end
 
+const AbelianGradedMatrix{T, D, S} = AbelianGradedArray{T, 2, D, S}
+
 # ---------------------------------------------------------------------------
 #  Constructors
 # ---------------------------------------------------------------------------
@@ -83,13 +85,20 @@ datatype(::Type{<:AbelianGradedArray{T, N, D, S}}) where {T, N, D, S} = D
 #  view (primitive): returns AbelianSectorArray sharing data with blockdata
 # ---------------------------------------------------------------------------
 
-function Base.view(a::AbelianGradedArray{T, N}, I::Block{N}) where {T, N}
+# Shared implementation: build the `AbelianSectorArray` view for a stored block.
+function view_abelian(a::AbelianGradedArray{T, N}, I::Block{N}) where {T, N}
     bk = Int.(Tuple(I))
     haskey(a.blockdata, bk) || error("Block $bk is not stored.")
     sects = ntuple(d -> sectors(axes(a, d))[bk[d]], Val(N))
     return AbelianSectorArray(sects, a.blockdata[bk])
 end
 
+Base.view(a::AbelianGradedArray{T, N}, I::Block{N}) where {T, N} = view_abelian(a, I)
+
+# Disambiguate against `view(::AbstractGradedArray{T, N}, ::Vararg{Block{1}, N})` for
+# N=1, where the splatted form collapses to a single Block{1} argument.
+Base.view(a::AbelianGradedArray{T, 1}, I::Block{1}) where {T} = view_abelian(a, I)
+
 # ---------------------------------------------------------------------------
 #  blocks — lazy view delegating to view (following BlockArrays convention)
 # ---------------------------------------------------------------------------
@@ -248,6 +257,23 @@ function Base.similar(
     D_N′ = isconcretetype(D_N) ? D_N : Array{T, N}
     return AbelianGradedArray{T, N, D_N′, S}(undef, axes)
 end
+
+# Allocate a graded array from any prototype when the requested axes are
+# `GradedOneTo`. The axes carry the structure; the prototype only fixes a
+# fallback datatype. Two overloads to resolve ambiguity with `BlockArrays`'
+# `StridedArray`-specific methods.
+function Base.similar(
+        ::AbstractArray, ::Type{T},
+        axes::Tuple{GradedOneTo{S}, Vararg{GradedOneTo{S}}}
+    ) where {T, S}
+    return AbelianGradedArray{T}(undef, axes)
+end
+function Base.similar(
+        ::StridedArray, ::Type{T},
+        axes::Tuple{GradedOneTo{S}, Vararg{GradedOneTo{S}}}
+    ) where {T, S}
+    return AbelianGradedArray{T}(undef, axes)
+end
 function Base.similar(
         a::AbstractGradedArray{T}, axes::Tuple{Vararg{GradedOneTo}}
     ) where {T}
@@ -260,6 +286,43 @@ function Base.similar(a::AbelianGradedArray{T}) where {T}
     return similar(a, T)
 end
 
+# Block-wise copy; the default falls through to scalar `copyto!`.
+function Base.copy(a::AbelianGradedArray{T, N, D, S}) where {T, N, D, S}
+    return AbelianGradedArray{T, N, D, S}(
+        Dict{NTuple{N, Int}, D}(k => copy(v) for (k, v) in a.blockdata),
+        a.axes
+    )
+end
+
+function Base.copyto!(
+        dest::AbelianGradedArray{<:Any, N},
+        src::AbelianGradedArray{<:Any, N}
+    ) where {N}
+    axes(dest) == axes(src) ||
+        throw(
+        DimensionMismatch(
+            "copyto! axes mismatch: dest $(axes(dest)), src $(axes(src))"
+        )
+    )
+    # Matching axes mean matching allowed-block keys (every allowed block is
+    # allocated), so copy each block into the existing destination buffer.
+    for (k, v) in src.blockdata
+        copyto!(dest.blockdata[k], v)
+    end
+    return dest
+end
+
+# Conjugate element-wise *and* flip axis duality, mirroring `Base.conj` on the
+# axis types (`SectorRange`/`GradedOneTo`/`SectorOneTo`). Without the axis flip,
+# `conj(t)` would leave bra-layer tensors with the same duality as the ket, and
+# any contraction between them would silently pair non-dual against non-dual.
+function Base.conj(a::AbelianGradedArray{T, N, D, S}) where {T, N, D, S}
+    return AbelianGradedArray{T, N, D, S}(
+        Dict{NTuple{N, Int}, D}(k => conj(v) for (k, v) in a.blockdata),
+        map(conj, a.axes)
+    )
+end
+
 # ---------------------------------------------------------------------------
 #  sectortype
 # ---------------------------------------------------------------------------
@@ -285,30 +348,143 @@ function Base.permutedims!(
     return y
 end
 
-# ---------------------------------------------------------------------------
-#  fill! / zero! / scale!
-# ---------------------------------------------------------------------------
+# Block-aware iszero: non-stored blocks are implicitly zero, so an
+# `AbelianGradedArray` is zero iff every stored block is zero. The generic
+# `Base.iszero(::AbstractArray) = all(iszero, x)` path iterates elements,
+# which throws on the no-scalar-indexing guard.
+function Base.iszero(a::AbelianGradedArray)
+    return all(iszero, values(a.blockdata))
+end
 
-scale!(a::AbstractArray, β::Number) = (a .*= β; a)
-function scale!(a::AbstractGradedArray, β::Number)
-    for bI in eachblockstoredindex(a)
-        scale!(view(a, bI), β)
+# Block-aware random fills: dispatch to the underlying block's `rand!`/`randn!`,
+# bypassing the generic `AbstractArray` fallbacks that go through scalar indexing.
+# The 3-arg `Random.rand!(rng, A, sp::Sampler)` form is what Random's `rand!(A)`
+# / `rand!(A, X)` / `rand!(rng, A, X)` shims ultimately call, so overriding it
+# covers every entry point.
+function Random.rand!(rng::AbstractRNG, a::AbelianGradedArray, sp::Random.Sampler)
+    for b in values(a.blockdata)
+        Random.rand!(rng, b, sp)
     end
     return a
 end
-
-function FI.zero!(a::AbstractGradedArray)
-    for bI in eachblockstoredindex(a)
-        FI.zero!(view(a, bI))
+function Random.randn!(rng::AbstractRNG, a::AbelianGradedArray)
+    for b in values(a.blockdata)
+        Random.randn!(rng, b)
     end
     return a
 end
 
-function Base.fill!(a::AbelianGradedArray, v)
-    iszero(v) || throw(
-        ArgumentError("fill! with nonzero value is not supported for AbelianGradedArray")
+# Constructors `rand(rng, T, axes)` / `randn(rng, T, axes)` for graded axes:
+# allocate an `AbelianGradedArray` from the axes, then fill via the block-aware
+# in-place methods above. The generic `Base.rand`/`randn` fallbacks build a
+# `Matrix` from `length.(axes)`, which loses the graded structure.
+function Base.rand(
+        rng::AbstractRNG, ::Type{T},
+        ax::Tuple{GradedOneTo, Vararg{GradedOneTo}}
+    ) where {T}
+    return Random.rand!(rng, AbelianGradedArray{T}(undef, ax))
+end
+function Base.randn(
+        rng::AbstractRNG, ::Type{T},
+        ax::Tuple{GradedOneTo, Vararg{GradedOneTo}}
+    ) where {T}
+    return Random.randn!(rng, AbelianGradedArray{T}(undef, ax))
+end
+
+# Block-aware diagonal check: block-diagonal (no off-diagonal stored blocks), and each
+# stored diagonal block is itself diagonal. Bypasses the generic scalar-indexing path.
+function LinearAlgebra.isdiag(A::AbelianGradedMatrix)
+    BlockSparseArrays.isblockdiagonal(A) || return false
+    for bI in eachblockstoredindex(A)
+        LinearAlgebra.isdiag(view(A, bI)) || return false
+    end
+    return true
+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.
+function TensorAlgebra.projectto!(dest::AbelianGradedArray, src::AbstractArray)
+    size(dest) == size(src) || throw(
+        DimensionMismatch(
+            "projectto!: dest has size $(size(dest)), src has size $(size(src))"
+        )
     )
-    return FI.zero!(a)
+    FI.zero!(dest)
+    for b in allowedblocks(axes(dest))
+        block_ranges = ntuple(d -> axes(dest, d)[Block(Int(Tuple(b)[d]))], ndims(dest))
+        view(dest, b) .= view(src, block_ranges...)
+    end
+    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...
+    )
+    TensorAlgebra.projectto!(dest, src)
+    isapprox(src, Array(dest); kwargs...) ||
+        throw(InexactError(:checked_projectto!, typeof(dest), src))
+    return dest
+end
+
+# Materialize the graded array into a dense `Array` for the default
+# `checked_projectto!`/`isapprox`-after path.
+function Base.Array(a::AbelianGradedArray{T, N}) where {T, N}
+    dest = zeros(T, size(a))
+    for bI in eachblockstoredindex(a)
+        block_ranges = ntuple(d -> axes(a, d)[Block(Int(Tuple(bI)[d]))], N)
+        copyto!(view(dest, block_ranges...), view(a, bI))
+    end
+    return dest
+end
+
+function LinearAlgebra.norm(a::AbelianGradedArray, p::Real = 2)
+    if p == Inf
+        isempty(a.blockdata) && return zero(float(real(eltype(a))))
+        return maximum(Base.Fix2(LinearAlgebra.norm, p), values(a.blockdata))
+    elseif p > 0
+        s = zero(float(real(eltype(a))))
+        for b in values(a.blockdata)
+            s += LinearAlgebra.norm(b, p)^p
+        end
+        return s^inv(p)
+    else
+        throw(ArgumentError("Norm with non-positive p ($p) is not defined"))
+    end
+end
+
+function LinearAlgebra.dot(a::AbelianGradedArray, b::AbelianGradedArray)
+    axes(a) == axes(b) ||
+        throw(DimensionMismatch("dot axes mismatch: a $(axes(a)), b $(axes(b))"))
+    # Matching axes mean matching allowed-block keys, so each `a` block has a
+    # counterpart in `b`.
+    s = zero(LinearAlgebra.dot(zero(eltype(a)), zero(eltype(b))))
+    for (k, ablk) in pairs(a.blockdata)
+        s += LinearAlgebra.dot(ablk, b.blockdata[k])
+    end
+    return s
+end
+
+# Forbidden blocks are zero, so the total is the sum over the stored blocks.
+function Base.sum(a::AbelianGradedArray)
+    s = zero(eltype(a))
+    for b in values(a.blockdata)
+        s += sum(b)
+    end
+    return s
+end
+
+# Scalar `*` / `/` are inherited from Base's `AbstractArray`-scalar methods, which
+# forward to broadcasting (`a .* x` / `a ./ x`). `AbelianGradedArray` supports the
+# linear-broadcast path, so no dedicated overrides are needed here.
+
+# `LinearAlgebra.normalize` infers its result eltype via `typeof(first(a)/nrm)`,
+# which scalar-indexes opaque storage.
+function LinearAlgebra.normalize(a::AbelianGradedArray, p::Real = 2)
+    return a / LinearAlgebra.norm(a, p)
 end
 
 # ---------------------------------------------------------------------------

src/abeliansectorarray.jl

@@ -104,15 +104,6 @@ function Base.permutedims!(y::AbelianSectorArray, x::AbelianSectorArray, perm)
     TensorAlgebra.permutedimsopadd!(y, identity, x, perm, true, false)
     return y
 end
-# Eager fallback for non-bosonic braiding: a lazy view of a fermionic array
-# would skip the permutation phase under scalar indexing.
-function FI.permuteddims(x::AbelianSectorArray, perm)
-    return if TKS.BraidingStyle(sectortype(x)) isa TKS.Bosonic
-        AbelianSectorArray(permutedims(sector(x), perm), FI.permuteddims(data(x), perm))
-    else
-        permutedims(x, perm)
-    end
-end
 
 # ========================  mul!  ========================
 

src/abeliansectordelta.jl

@@ -38,9 +38,6 @@ function Base.permutedims(x::AbelianSectorDelta, perm)
     new_sectors = ntuple(n -> x.sectors[perm[n]], Val(ndims(x)))
     return AbelianSectorDelta{eltype(x)}(new_sectors)
 end
-function FI.permuteddims(x::AbelianSectorDelta, perm)
-    return permutedims(x, perm)
-end
 
 # ========================  adjoint / broadcasting  ========================