Add Data{N} type, undef constructors, and data/block indexing (#150)

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.8.1"
+version = "0.8.2"
 authors = ["ITensor developers <support@itensor.org> and contributors"]
 
 [workspace]

src/GradedArrays.jl

@@ -11,19 +11,22 @@ export AbstractSectorArray,
 export AbstractGradedArray, AbstractGradedMatrix
 export AbelianGradedArray, AbelianGradedVector, AbelianGradedMatrix
 export FusedGradedMatrix
-export gradedrange
 
 export dual, flip, gradedrange, isdual,
     data, dataaxes, dataaxes1, datalength, datalengths,
+    eachdataaxis, eachsectoraxis,
     sector, sectoraxes, sectoraxes1, sectorlength, sectorlengths,
-    sectors, sectortype
+    sectors, sectortype,
+    Data
 
 # imports
 # -------
 import FunctionImplementations as FI
 using BlockArrays: BlockArrays, AbstractBlockArray, AbstractBlockVector, Block,
-    BlockIndexRange, BlockVector, blocklength, blocklengths, blocks, eachblockaxes1
-using BlockSparseArrays: BlockSparseArrays, eachblockaxis, eachblockstoredindex, mortar_axis
+    BlockIndexRange, BlockVector, BlockedOneTo, blockedrange, blocklength, blocklengths,
+    blocks, eachblockaxes1
+using BlockSparseArrays:
+    BlockSparseArrays, blockdiagindices, eachblockaxis, eachblockstoredindex, mortar_axis
 using KroneckerArrays: KroneckerArrays, kroneckerfactors, ×, ⊗
 using LinearAlgebra: LinearAlgebra, Adjoint, mul!
 using TensorAlgebra: TensorAlgebra, BlockedTuple, FusionStyle, matricize, matricize_axes,
@@ -33,6 +36,7 @@ using TensorKitSectors: TensorKitSectors as TKS
 using TypeParameterAccessors: type_parameters, unspecify_type_parameters
 
 include("sectorrange.jl")
+include("data.jl")
 include("sectoroneto.jl")
 include("gradedoneto.jl")
 include("abstractsectordelta.jl")

src/abeliangradedarray.jl

@@ -27,11 +27,14 @@ end
 function AbelianGradedArray{T, N, D, S}(
         ::UndefInitializer, axs::NTuple{N, GradedOneTo{S}}
     ) where {T, N, D <: AbstractArray{T, N}, S <: SectorRange}
+    block_axes = map(eachdataaxis, axs)
+    function allocate_block(bk)
+        bk_inds = Int.(Tuple(bk))
+        return similar(D, ntuple(d -> block_axes[d][bk_inds[d]], Val(N)))
+    end
     bks = allowedblocks(axs)
     blockdata = Dict{NTuple{N, Int}, D}(
-        Int.(Tuple(bk)) =>
-            D(undef, ntuple(d -> blocklengths(axs[d])[Int(Tuple(bk)[d])], Val(N)))
-            for bk in bks
+        Int.(Tuple(bk)) => allocate_block(bk) for bk in bks
     )
     return AbelianGradedArray{T, N, D, S}(blockdata, axs)
 end
@@ -66,16 +69,12 @@ BlockSparseArrays.blocktype(a::AbelianGradedArray) = BlockSparseArrays.blocktype
 #  view (primitive): returns AbelianSectorArray sharing data with blockdata
 # ---------------------------------------------------------------------------
 
-# view: returns a AbelianSectorArray sharing data (errors for unstored blocks)
-function Base.view(a::AbelianGradedArray{T, N}, I::Vararg{Block{1}, N}) where {T, N}
-    bk = ntuple(d -> Int(I[d]), Val(N))
+function Base.view(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
-function Base.view(a::AbelianGradedArray{T, N}, I::Block{N}) where {T, N}
-    return view(a, Tuple(I)...)
-end
 
 # ---------------------------------------------------------------------------
 #  blocks — lazy view delegating to view (following BlockArrays convention)
@@ -108,32 +107,6 @@ function Base.setindex!(
     return b
 end
 
-# ---------------------------------------------------------------------------
-#  getindex / setindex! on AbelianGradedArray with Block — convenience wrappers
-# ---------------------------------------------------------------------------
-
-# getindex: returns a copy (errors for unstored blocks)
-function Base.getindex(a::AbelianGradedArray{T, N}, I::Vararg{Block{1}, N}) where {T, N}
-    return copy(view(a, I...))
-end
-function Base.getindex(a::AbelianGradedArray{T, N}, I::Block{N}) where {T, N}
-    return a[Tuple(I)...]
-end
-
-# setindex!: Block{N} unpacks to Vararg{Block{1}, N} (following BlockArrays convention)
-function Base.setindex!(a::AbelianGradedArray{<:Any, N}, value, I::Block{N}) where {N}
-    return setindex!(a, value, Tuple(I)...)
-end
-
-# Primitive: view existing block, then broadcast in.
-# Handles both AbelianSectorArray and raw data values.
-function Base.setindex!(
-        a::AbelianGradedArray{<:Any, N}, value::AbstractArray{<:Any, N}, I::Vararg{Block{1}, N}
-    ) where {N}
-    view(a, I...) .= value
-    return a
-end
-
 # ---------------------------------------------------------------------------
 #  Splitting getindex: each I[d][k] = Block(b)[r] means dest block k comes
 #  from source block b at subrange r. Inverse of the merging getindex.

src/abeliansectorarray.jl

@@ -36,6 +36,12 @@ function AbelianSectorArray(
     ) where {N}
     return AbelianSectorArray(delta.sectors, data)
 end
+function AbelianSectorArray{T, N, A, S}(
+        delta::AbelianSectorDelta{<:Any, N, S},
+        data::A
+    ) where {T, N, A <: AbstractArray{T, N}, S <: SectorRange}
+    return AbelianSectorArray{T, N, A, S}(delta.sectors, data)
+end
 
 const AbelianSectorVector{T, A <: AbstractVector{T}, S <: SectorRange} =
     AbelianSectorArray{T, 1, A, S}
@@ -55,23 +61,19 @@ datatype(::Type{AbelianSectorArray{T, N, A, S}}) where {T, N, A, S} = A
 # AbstractArray interface
 # -----------------------
 function Base.axes(sa::AbelianSectorArray)
-    sa_axes = sectoraxes(sa)
-    da_axes = dataaxes(sa)
-    return ntuple(d -> SectorOneTo(sa_axes[d], length(da_axes[d])), Val(ndims(sa)))
+    return map(SectorOneTo, sectoraxes(sa), dataaxes(sa))
 end
 
 Base.copy(A::AbelianSectorArray) = AbelianSectorArray(sector(A), copy(data(A)))
 
 # similar for AbelianSectorArray with SectorOneTo axes.
 # Delegates to similar on the data array for the data dimensions.
 function Base.similar(
-        a::AbelianSectorArray,
+        ::AbelianSectorArray,
         ::Type{T},
         axes::Tuple{SectorOneTo, Vararg{SectorOneTo}}
     ) where {T}
-    sects = sector.(axes)
-    data_ax = data.(axes)
-    return AbelianSectorArray(sects, similar(data(a), T, data_ax))
+    return AbelianSectorArray{T}(undef, axes)
 end
 
 function Base.fill!(A::AbelianSectorArray, v)
@@ -88,7 +90,7 @@ function Base.convert(
         x::AbelianSectorArray{T₂, N, B, S}
     )::AbelianSectorArray{T₁, N, A, S} where {T₁, T₂, N, A, B, S}
     A === B && return x
-    return AbelianSectorArray(sector(x), convert(A, data(x)))
+    return AbelianSectorArray{T₁, N, A, S}(sector(x), convert(A, data(x)))
 end
 
 # ========================  permutedims  ========================

src/abstractgradedarray.jl

@@ -26,3 +26,54 @@ function Base.setindex!(::AbstractGradedArray, _, ::Vararg{Int})
         "Scalar indexing is not supported for AbstractGradedArray. Use block indexing."
     )
 end
+
+# ---------------------------------------------------------------------------
+#  Block indexing interface
+#
+#  Concrete subtypes must implement:
+#    view(a::ConcreteType, ::Block{N})  → sector-wrapped view (e.g. SectorMatrix)
+#
+#  Everything else is derived here.
+# ---------------------------------------------------------------------------
+
+function Base.view(a::AbstractGradedArray{T, N}, I::Vararg{Block{1}, N}) where {T, N}
+    return view(a, Block(Int.(I)))
+end
+
+function Base.getindex(a::AbstractGradedArray{T, N}, I::Block{N}) where {T, N}
+    return copy(view(a, I))
+end
+function Base.getindex(a::AbstractGradedArray{T, N}, I::Vararg{Block{1}, N}) where {T, N}
+    return a[Block(Int.(I))]
+end
+
+function Base.setindex!(a::AbstractGradedArray{<:Any, N}, value, I::Block{N}) where {N}
+    return setindex!(a, value, Tuple(I)...)
+end
+function Base.setindex!(
+        a::AbstractGradedArray{<:Any, N}, value, I::Vararg{Block{1}, N}
+    ) where {N}
+    view(a, I...) .= value
+    return a
+end
+
+# ---------------------------------------------------------------------------
+#  Data indexing — raw block data without sector wrappers
+#
+#  Built on top of Block view: view(a, Data(I)) = data(view(a, Block(I)))
+# ---------------------------------------------------------------------------
+
+function Base.view(a::AbstractGradedArray{T, N}, I::Data{N}) where {T, N}
+    return data(view(a, Block(I)))
+end
+
+function Base.getindex(a::AbstractGradedArray{T, N}, I::Data{N}) where {T, N}
+    return copy(view(a, I))
+end
+
+function Base.setindex!(
+        a::AbstractGradedArray{<:Any, N}, value::AbstractArray{<:Any, N}, I::Data{N}
+    ) where {N}
+    view(a, I) .= value
+    return a
+end

src/data.jl

@@ -0,0 +1,14 @@
+"""
+    Data{N}
+
+Block-data indexing type analogous to `BlockArrays.Block{N}`. Indexing a graded
+array with `Data(i, j, ...)` accesses the raw data array for that block, without
+sector metadata wrappers.
+"""
+struct Data{N}
+    n::NTuple{N, Int}
+end
+Data(n::Vararg{Int, N}) where {N} = Data{N}(n)
+BlockArrays.Block(I::Data) = Block(I.n)
+Data(I::Block) = Data(Int.(Tuple(I)))
+Base.Tuple(I::Data) = I.n