Bug fixes, broadcasting, matricize, and display improvements

Project.toml

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

src/GradedArrays.jl

@@ -22,13 +22,14 @@ export dual, flip, gradedrange, isdual,
 # imports
 # -------
 import FunctionImplementations as FI
-using BlockArrays: BlockArrays, AbstractBlockArray, AbstractBlockVector, Block,
-    BlockIndexRange, BlockVector, BlockedOneTo, blockedrange, blocklength, blocklengths,
-    blocks, eachblockaxes1
-using BlockSparseArrays:
-    BlockSparseArrays, blockdiagindices, eachblockaxis, eachblockstoredindex, mortar_axis
+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 KroneckerArrays: KroneckerArrays, kroneckerfactors, ×, ⊗
 using LinearAlgebra: LinearAlgebra, Adjoint, mul!
+using SparseArraysBase: SparseArraysBase
 using TensorAlgebra: TensorAlgebra, BlockedTuple, FusionStyle, matricize, matricize_axes,
     permutedimsadd!, permutedimsopadd!, tensor_product_axis, trivial_axis, trivialbiperm,
     tryflattenlinear, unmatricize

src/abeliangradedarray.jl

@@ -52,6 +52,19 @@ function AbelianGradedArray{T}(
     return AbelianGradedArray{T}(init, axs)
 end
 
+# Convert any `AbstractGradedMatrix` (e.g. a `FusedGradedMatrix`) to an
+# `AbelianGradedArray` with the same axes and stored blocks.
+function AbelianGradedArray(m::AbstractGradedMatrix)
+    # Assumes each allowed block of the target is also stored in `m` — every
+    # `similar` allocation is overwritten by the loop below, so no `zero!`
+    # is needed.
+    a = similar(m, axes(m))
+    for I in eachblockstoredindex(m)
+        a[Data(I)] = view(m, Data(I))
+    end
+    return a
+end
+
 # ---------------------------------------------------------------------------
 #  AbstractArray interface
 # ---------------------------------------------------------------------------
@@ -117,6 +130,9 @@ function Base.getindex(
         a::AbelianGradedArray{T, N}, I::Vararg{AbstractVector{<:BlockIndexRange{1}}, N}
     ) where {T, N}
     ax_dest = ntuple(d -> axes(a, d)[I[d]], Val(N))
+    # `zero!` is needed: we only copy sub-ranges from stored source blocks,
+    # so destination block regions outside those sub-ranges (and destination
+    # blocks with no source counterpart) must start at 0.
     a_dest = FI.zero!(similar(a, ax_dest))
     # Map source block b → list of (dest BlockIndexRange, src subrange).
     src_to_dests = ntuple(Val(N)) do d
@@ -159,6 +175,9 @@ function Base.getindex(
         a::AbelianGradedArray{T, N}, I::Vararg{AbstractBlockVector{<:Block{1}}, N}
     ) where {T, N}
     ax_dest = ntuple(d -> axes(a, d)[I[d]], Val(N))
+    # `zero!` is needed: each source block writes into a sub-range of one
+    # destination block, so remaining sub-ranges (and destination blocks
+    # with no source counterpart) must start at 0.
     a_dest = FI.zero!(similar(a, ax_dest))
     ax = axes(a)
     # Map source Block → BlockIndexRange encoding dest block + subrange within it
@@ -198,6 +217,18 @@ function BlockSparseArrays.eachblockstoredindex(a::AbelianGradedArray{T, N}) whe
     return (Block(k) for k in keys(a.blockdata))
 end
 
+# Implement the `SparseArraysBase` interface on `AbelianBlocks` (the lazy
+# block view) so that `storedlength(blocks(a))` — and by extension
+# `blockstoredlength(a)` — reflects the dict-of-keys storage rather than
+# treating every slot as stored.
+function SparseArraysBase.eachstoredindex(b::AbelianBlocks{T, N}) where {T, N}
+    return (CartesianIndex(k) for k in keys(b.parent.blockdata))
+end
+SparseArraysBase.storedvalues(b::AbelianBlocks) = values(b.parent.blockdata)
+function SparseArraysBase.isstored(b::AbelianBlocks{T, N}, I::Vararg{Int, N}) where {T, N}
+    return haskey(b.parent.blockdata, I)
+end
+
 # ---------------------------------------------------------------------------
 #  similar
 # ---------------------------------------------------------------------------
@@ -240,7 +271,9 @@ sectortype(::Type{<:AbelianGradedArray{T, N, D, S}}) where {T, N, D, S} = S
 
 function Base.permutedims(a::AbelianGradedArray{<:Any, N}, perm) where {N}
     dest_axes = ntuple(i -> axes(a)[perm[i]], Val(N))
-    a_dest = FI.zero!(similar(a, dest_axes))
+    # No `zero!` here: `permutedims!` → `permutedimsopadd!(β=0)` already
+    # zeros the destination before writing.
+    a_dest = similar(a, dest_axes)
     return permutedims!(a_dest, a, perm)
 end
 
@@ -288,19 +321,32 @@ const AbelianGradedMatrix{T, D, S} = AbelianGradedArray{T, 2, D, S}
 #  show
 # ---------------------------------------------------------------------------
 
-function Base.show(io::IO, ::MIME"text/plain", a::AbelianGradedArray{T, N}) where {T, N}
+function Base.summary(io::IO, a::AbelianGradedArray)
     block_str = join(map(g -> string(blocklength(g)), axes(a)), "×")
     size_str = join(map(string, size(a)), "×")
-    nstored = length(collect(eachblockstoredindex(a)))
-    print(io, block_str, "-blocked ", size_str, " AbelianGradedArray{", T, "}")
+    nstored = blockstoredlength(a)
+    print(io, block_str, "-blocked ", size_str, " ", typeof(a))
     print(io, " with ", nstored, " stored block", nstored == 1 ? "" : "s")
     return nothing
 end
 
-function Base.show(io::IO, a::AbelianGradedArray{T, N}) where {T, N}
+function Base.show(io::IO, ::MIME"text/plain", a::AbelianGradedArray)
+    summary(io, a)
+    println(io, ":")
+    for (d, g) in pairs(axes(a))
+        print(io, "  Dim $d: ")
+        show(io, g)
+        println(io)
+    end
+    isempty(a) && return nothing
+    Base.print_array(io, a)
+    return nothing
+end
+
+function Base.show(io::IO, a::AbelianGradedArray)
     block_str = join(map(g -> string(blocklength(g)), axes(a)), "×")
     size_str = join(map(string, size(a)), "×")
-    print(io, block_str, "-blocked ", size_str, " AbelianGradedArray{", T, "}")
+    print(io, block_str, "-blocked ", size_str, " ", typeof(a))
     return nothing
 end
 

src/abeliansectordelta.jl

@@ -106,7 +106,7 @@ function fermion_contraction_phase(
     length_codomain <= ndims(x) ||
         throw(ArgumentError(lazy"Cannot contract more than ndim legs ($N > $(ndims(x))"))
 
-    parity = mapreduce(⊻, enumerate(axes(x))) do (n, ax)
+    parity = mapreduce(⊻, pairs(axes(x))) do (n, ax)
         return (n <= length_codomain) & isdual(ax) & fermionparity(ax)
     end
     return ifelse(parity, -1, 1)

src/abstractgradedarray.jl

@@ -77,3 +77,23 @@ function Base.setindex!(
     view(a, I) .= value
     return a
 end
+
+# ---------------------------------------------------------------------------
+#  Display — convert to BlockSparseArray for printing
+# ---------------------------------------------------------------------------
+
+using BlockSparseArrays: BlockSparseArray
+
+function _to_blocksparsearray(a::AbstractGradedArray{T, N}) where {T, N}
+    blocked_axes = map(g -> blockedrange(blocklengths(g)), axes(a))
+    bsa = BlockSparseArray{T}(undef, blocked_axes)
+    for bI in eachblockstoredindex(a)
+        blk = view(a, bI)
+        bsa[bI] = collect(Array(sector(blk)) ⊗ data(blk))
+    end
+    return bsa
+end
+
+function Base.print_array(io::IO, a::AbstractGradedArray)
+    return Base.print_array(io, _to_blocksparsearray(a))
+end

src/abstractsectorarray.jl

@@ -53,3 +53,26 @@ function FI.zero!(a::AbstractSectorArray)
     FI.zero!(data(a))
     return a
 end
+
+# ========================  display  ========================
+
+function Base.print_array(io::IO, sa::AbstractSectorArray)
+    Base.print_array(io, sector(sa))
+    println(io, "\n ⊗")
+    Base.print_array(io, data(sa))
+    return nothing
+end
+
+function Base.show(io::IO, sa::AbstractSectorArray)
+    show(io, sector(sa))
+    print(io, " ⊗ ")
+    show(io, data(sa))
+    return nothing
+end
+
+function Base.show(io::IO, ::MIME"text/plain", sa::AbstractSectorArray)
+    summary(io, sa)
+    println(io, ":")
+    Base.print_array(io, sa)
+    return nothing
+end

src/broadcast.jl

@@ -52,39 +52,39 @@ function Base.copyto!(dest::AbstractSectorArray, bc::BC.Broadcasted{<:SectorStyl
     return dest
 end
 
-# ========================  AbelianGradedArray broadcasting  ========================
+# ========================  GradedArray broadcasting  ========================
 
-struct AbelianGradedStyle{N} <: BC.AbstractArrayStyle{N} end
-AbelianGradedStyle{N}(::Val{M}) where {N, M} = AbelianGradedStyle{M}()
+struct GradedStyle{N} <: BC.AbstractArrayStyle{N} end
+GradedStyle{N}(::Val{M}) where {N, M} = GradedStyle{M}()
 
-function BC.BroadcastStyle(::Type{<:AbelianGradedArray{<:Any, N}}) where {N}
-    return AbelianGradedStyle{N}()
+function BC.BroadcastStyle(::Type{<:AbstractGradedArray{<:Any, N}}) where {N}
+    return GradedStyle{N}()
 end
 function BC.BroadcastStyle(
-        style::AbelianGradedStyle{N},
+        style::GradedStyle{N},
         ::BC.DefaultArrayStyle{0}
     ) where {N}
     return style
 end
 function BC.BroadcastStyle(
         ::BC.DefaultArrayStyle{0},
-        style::AbelianGradedStyle{N}
+        style::GradedStyle{N}
     ) where {N}
     return style
 end
-BC.BroadcastStyle(s1::AbelianGradedStyle{N}, ::AbelianGradedStyle{N}) where {N} = s1
+BC.BroadcastStyle(s1::GradedStyle{N}, ::GradedStyle{N}) where {N} = s1
 
 # TODO: Ideally this would incorporate information from all broadcast arguments
 # (or their blocktypes) when computing similar, rather than picking one argument.
-function Base.similar(bc::BC.Broadcasted{<:AbelianGradedStyle}, elt::Type)
+function Base.similar(bc::BC.Broadcasted{<:GradedStyle}, elt::Type)
     bc′ = BC.flatten(bc)
-    arg = bc′.args[findfirst(arg -> arg isa AbelianGradedArray, bc′.args)]
+    arg = bc′.args[findfirst(arg -> arg isa AbstractGradedArray, bc′.args)]
     return similar(arg, elt)
 end
 
-function Base.copyto!(dest::AbelianGradedArray, bc::BC.Broadcasted{<:AbelianGradedStyle})
+function Base.copyto!(dest::AbstractGradedArray, bc::BC.Broadcasted{<:GradedStyle})
     lb = tryflattenlinear(bc)
     isnothing(lb) &&
-        throw(ArgumentError("AbelianGradedArray broadcasting requires linear operations"))
+        throw(ArgumentError("AbstractGradedArray broadcasting requires linear operations"))
     return copyto!(dest, lb)
 end

src/fusedgradedmatrix.jl

@@ -40,6 +40,11 @@ function FusedGradedMatrix(
     ) where {T, S <: SectorRange, D <: AbstractMatrix{T}}
     return FusedGradedMatrix{T, D, S}(sectors, blocks)
 end
+function FusedGradedMatrix(pairs::AbstractVector{<:Pair})
+    sectors = first.(pairs)
+    blocks = last.(pairs)
+    return FusedGradedMatrix(sectors, blocks)
+end
 
 # ========================  undef constructors  ========================
 
@@ -138,11 +143,17 @@ end
 
 # ========================  mul!  ========================
 
-function check_mul_axes(
-        C::FusedGradedMatrix, A::FusedGradedMatrix, B::FusedGradedMatrix
-    )
+function check_input(::typeof(*), A::FusedGradedMatrix, B::FusedGradedMatrix)
     axes(A, 2) == dual(axes(B, 1)) ||
         throw(DimensionMismatch("sector mismatch in contracted dimension"))
+    return nothing
+end
+
+function check_input(
+        ::typeof(mul!),
+        C::FusedGradedMatrix, A::FusedGradedMatrix, B::FusedGradedMatrix
+    )
+    check_input(*, A, B)
     axes(C, 1) == axes(A, 1) || throw(DimensionMismatch())
     axes(C, 2) == axes(B, 2) || throw(DimensionMismatch())
     return nothing
@@ -152,7 +163,7 @@ function LinearAlgebra.mul!(
         C::FusedGradedMatrix, A::FusedGradedMatrix, B::FusedGradedMatrix,
         α::Number, β::Number
     )
-    check_mul_axes(C, A, B)
+    check_input(mul!, C, A, B)
     for I in blockdiagindices(C)
         mul!(view(C, Data(I)), view(A, Data(I)), view(B, Data(I)), α, β)
     end
@@ -172,6 +183,7 @@ function allocate_output(::typeof(*), A::FusedGradedMatrix, B::FusedGradedMatrix
 end
 
 function Base.:(*)(A::FusedGradedMatrix, B::FusedGradedMatrix)
+    check_input(*, A, B)
     C = allocate_output(*, A, B)
     return mul!(C, A, B)
 end
@@ -185,18 +197,30 @@ end
 
 # ========================  show  ========================
 
-function Base.show(io::IO, ::MIME"text/plain", m::FusedGradedMatrix{T}) where {T}
+function Base.summary(io::IO, m::FusedGradedMatrix)
     nblocks = length(m.sectors)
-    print(io, nblocks, "-block FusedGradedMatrix{", T, "} with sectors ")
-    print(io, "[")
-    join(io, label.(m.sectors), ", ")
+    print(io, nblocks, "-block ", typeof(m), " with sectors [")
+    join(io, m.sectors, ", ")
     print(io, "]")
     return nothing
 end
 
-function Base.show(io::IO, m::FusedGradedMatrix{T}) where {T}
+function Base.show(io::IO, ::MIME"text/plain", m::FusedGradedMatrix)
+    summary(io, m)
+    println(io, ":")
+    for (d, g) in pairs(axes(m))
+        print(io, "  Dim $d: ")
+        show(io, g)
+        println(io)
+    end
+    isempty(m.sectors) && return nothing
+    Base.print_array(io, m)
+    return nothing
+end
+
+function Base.show(io::IO, m::FusedGradedMatrix)
     nblocks = length(m.sectors)
-    print(io, nblocks, "-block FusedGradedMatrix{", T, "}")
+    print(io, nblocks, "-block ", typeof(m))
     return nothing
 end
 
@@ -225,17 +249,3 @@ function FusedGradedMatrix(a::AbelianGradedMatrix{T}) where {T}
     end
     return m
 end
-
-"""
-    AbelianGradedArray(m::FusedGradedMatrix)
-
-Convert a `FusedGradedMatrix` to a 2D `AbelianGradedArray`.
-Inverse of `FusedGradedMatrix(::AbelianGradedArray)`.
-"""
-function AbelianGradedArray(m::FusedGradedMatrix{T}) where {T}
-    a = similar(m, axes(m))
-    for I in blockdiagindices(m)
-        a[Data(I)] = view(m, Data(I))
-    end
-    return a
-end