From feb74e1662be353053961285672069217ecb5e82 Mon Sep 17 00:00:00 2001 From: mtfishman Date: Mon, 9 Mar 2026 10:48:07 -0400 Subject: [PATCH 01/13] [WIP] Fix matricize --- src/tensoralgebra.jl | 127 +++++++++++++++++++++++++++++----- test/test_tensoralgebraext.jl | 7 +- 2 files changed, 111 insertions(+), 23 deletions(-) diff --git a/src/tensoralgebra.jl b/src/tensoralgebra.jl index 8e51f76b..5a55f7a9 100644 --- a/src/tensoralgebra.jl +++ b/src/tensoralgebra.jl @@ -1,5 +1,5 @@ using BlockArrays: blocks, eachblockaxes1 -using BlockSparseArrays: BlockSparseArray, blockrange, blockreshape +using BlockSparseArrays: BlockSparseArray, blockrange, blockreshape, eachblockstoredindex using GradedArrays: GradedArray, GradedUnitRange, SectorRange, flip, invblockperm, sectormergesortperm, sectorsortperm, trivial, unmerged_tensor_product, × using TensorAlgebra: TensorAlgebra, AbstractBlockPermutation, BlockedTuple, FusionStyle, @@ -49,7 +49,7 @@ end function TensorAlgebra.tensor_product_axis( ::SectorFusion, ::Val{:domain}, r1::SectorUnitRange, r2::SectorUnitRange ) - return dual(r1 ⊗ r2) + return flip(r1 ⊗ r2) end function TensorAlgebra.tensor_product_axis( style::BlockReshapeFusion, side::Val{:codomain}, @@ -90,11 +90,20 @@ function TensorAlgebra.matricize( ) where {Ncodomain} biperm = trivialbiperm(ndims_codomain, Val(ndims(a))) ax_codomain, ax_domain = blocks(axes(a)[biperm]) - ax_codomain = tensor_product(ax_codomain...) - ax_domain = flip(tensor_product(ax_domain...)) + ax_codomain = + isempty(ax_codomain) ? trivial(sector_type(a)) : tensor_product(ax_codomain...) + ax_domain = + isempty(ax_domain) ? trivial(sector_type(a)) : flip(tensor_product(ax_domain...)) return SectorDelta{eltype(a)}((ax_codomain, ax_domain)) end +function TensorAlgebra.unmatricize( + ::SectorFusion, m::SectorDelta, + codomain_axes::Tuple{Vararg{SectorRange}}, + domain_axes::Tuple{Vararg{SectorRange}} + ) + return SectorDelta{eltype(m)}((codomain_axes..., domain_axes...)) +end function TensorAlgebra.unmatricize( ::SectorFusion, m::AbstractMatrix, codomain_axes::Tuple{Vararg{AbstractUnitRange}}, @@ -109,27 +118,107 @@ function TensorAlgebra.unmatricize( return a end - # First, fuse axes to get `sectormergesortperm`. - # Then unpermute the blocks. + # Compute the unmerged fused axes (block structure before sectormergesort). fused_axes = matricize_axes(BlockReshapeFusion(), m, codomain_axes, domain_axes) - blockperms = sectorsortperm.(fused_axes) - sorted_axes = map((r, I) -> only(axes(r[I])), fused_axes, blockperms) + # Split merged blocks back to unmerged block structure, then unmatricize. + m_unmerged = unsectormergesort(m, fused_axes) + return unmatricize(FusionStyle(BlockSparseArray), m_unmerged, blocked_axes) +end + +# Split merged blocks back to unmerged block structure (inverse of sectormergesort for arrays). +function unsectormergesort(m::AbstractArray, original_axes) + N = ndims(m) + perms = sectormergesortperm.(original_axes) + + result = zeros(eltype(m), original_axes) - # TODO: This is doing extra copies of the blocks, - # use `@view a[axes_prod...]` instead. - # That will require implementing some reindexing logic - # for this combination of slicing. - m_unblocked = m[sorted_axes...] - m_blockpermed = m_unblocked[invblockperm.(blockperms)...] - return unmatricize(FusionStyle(BlockSparseArray), m_blockpermed, blocked_axes) + for I in eachblockstoredindex(m) + dest_idx = map(Int, Tuple(I)) + merged_block_data = m[I] + + # For each dimension, find which source blocks map to this merged block + # and their offsets within it. + src_groups = ntuple(N) do dim + grp = blocks(perms[dim])[dest_idx[dim]] + cumoff = 0 + map(eachindex(grp)) do k + off = cumoff + cumoff += length(original_axes[dim][grp[k]]) + return (grp[k], off) + end + end + + # Iterate over all combinations of source blocks. + for src_combo in Iterators.product(src_groups...) + src_blocks = ntuple(dim -> src_combo[dim][1], N) # Block{1, Int64} + src_offsets = ntuple(dim -> src_combo[dim][2], N) + src_sizes = ntuple(dim -> length(original_axes[dim][src_blocks[dim]]), N) + src_region = + ntuple(dim -> (src_offsets[dim] + 1):(src_offsets[dim] + src_sizes[dim]), N) + src_data = merged_block_data[src_region...] + if !all(iszero, src_data) + result[src_blocks...] = src_data + end + end + end + + return result end # Sort the blocks by sector and then merge the common sectors. function sectormergesort(a::AbstractArray) - # TODO: fix this, no clue why broken and no clue how to fix - return a + axs = axes(a) + perms = sectormergesortperm.(axs) + merged_axes = sectormergesort.(axs) + + result = zeros(eltype(a), merged_axes) + + N = ndims(a) + + function find_dest_and_pos(perm, src_block::Block{1}) + for (igroup, grp) in enumerate(blocks(perm)) + pos = findfirst(==(src_block), grp) + isnothing(pos) || return (igroup, pos) + end + return error("Block $src_block not found in permutation") + end + + # Group stored source blocks by their destination block + dest_groups = Dict{NTuple{N, Int}, Vector{Tuple{Block{N, Int}, NTuple{N, Int}}}}() + + for I in eachblockstoredindex(a) + src_block_indices = Tuple(I) + + dest_and_pos = ntuple(N) do dim + return find_dest_and_pos(perms[dim], src_block_indices[dim]) + end + dest_idx = ntuple(dim -> dest_and_pos[dim][1], N) + + offsets = ntuple(N) do dim + igroup, pos = dest_and_pos[dim] + grp = blocks(perms[dim])[igroup] + return sum(length(axs[dim][grp[k]]) for k in 1:(pos - 1); init = 0) + end + + push!(get!(dest_groups, dest_idx, []), (I, offsets)) + end + + # Build each destination block from its contributing source blocks + for (dest_idx, src_list) in dest_groups + dest_size = ntuple(dim -> length(merged_axes[dim][Block(dest_idx[dim])]), N) + dest_data = zeros(eltype(a), dest_size) + + for (I, offsets) in src_list + src_data = a[I] + src_sizes = size(src_data) + dest_region = + ntuple(dim -> (offsets[dim] + 1):(offsets[dim] + src_sizes[dim]), N) + dest_data[dest_region...] = src_data + end + + result[Block.(dest_idx)...] = dest_data + end - I = sectormergesortperm.(axes(a)) - return a[I...] + return result end diff --git a/test/test_tensoralgebraext.jl b/test/test_tensoralgebraext.jl index a106bffb..3c76063d 100644 --- a/test/test_tensoralgebraext.jl +++ b/test/test_tensoralgebraext.jl @@ -32,10 +32,8 @@ end @test space_isequal(trivial_gradedrange((gN, gS)), gNS) end -broken = true - const elts = (Float32, Float64, Complex{Float32}, Complex{Float64}) -broken || @testset "`contract` `GradedArray` (eltype=$elt)" for elt in elts +@testset "`contract` `GradedArray` (eltype=$elt)" for elt in elts @testset "matricize" begin d1 = gradedrange([U1(0) => 1, U1(1) => 1]) d2 = gradedrange([U1(0) => 1, U1(1) => 1]) @@ -82,7 +80,8 @@ broken || @testset "`contract` `GradedArray` (eltype=$elt)" for elt in elts @test a == unmatricize(m, (), (d1, d2, dual(d1), dual(d2))) end - @testset "contract with U(1)" begin + contract_broken = true + contract_broken || @testset "contract with U(1)" begin d = gradedrange([U1(0) => 2, U1(1) => 3]) a1 = randn_blockdiagonal(elt, (d, d, dual(d), dual(d))) a2 = randn_blockdiagonal(elt, (d, d, dual(d), dual(d))) From d95d4c88733b83f35db9c539365afd62c6829c4b Mon Sep 17 00:00:00 2001 From: mtfishman Date: Mon, 9 Mar 2026 13:18:32 -0400 Subject: [PATCH 02/13] Better fix --- src/tensoralgebra.jl | 37 +++++++++++++++++++++---------------- 1 file changed, 21 insertions(+), 16 deletions(-) diff --git a/src/tensoralgebra.jl b/src/tensoralgebra.jl index 5a55f7a9..fba6a9a3 100644 --- a/src/tensoralgebra.jl +++ b/src/tensoralgebra.jl @@ -1,4 +1,4 @@ -using BlockArrays: blocks, eachblockaxes1 +using BlockArrays: Block, BlockVector, blocks, eachblockaxes1 using BlockSparseArrays: BlockSparseArray, blockrange, blockreshape, eachblockstoredindex using GradedArrays: GradedArray, GradedUnitRange, SectorRange, flip, invblockperm, sectormergesortperm, sectorsortperm, trivial, unmerged_tensor_product, × @@ -127,6 +127,9 @@ function TensorAlgebra.unmatricize( end # Split merged blocks back to unmerged block structure (inverse of sectormergesort for arrays). +_asblock(i::Block{1}) = i +_asblock(i::Integer) = Block(i) + function unsectormergesort(m::AbstractArray, original_axes) N = ndims(m) perms = sectormergesortperm.(original_axes) @@ -144,8 +147,9 @@ function unsectormergesort(m::AbstractArray, original_axes) cumoff = 0 map(eachindex(grp)) do k off = cumoff - cumoff += length(original_axes[dim][grp[k]]) - return (grp[k], off) + src_block = _asblock(grp[k]) + cumoff += length(original_axes[dim][src_block]) + return (src_block, off) end end @@ -166,30 +170,27 @@ function unsectormergesort(m::AbstractArray, original_axes) return result end -# Sort the blocks by sector and then merge the common sectors. -function sectormergesort(a::AbstractArray) +# Materialize grouped block indexing used by `sectormergesort`. +function Base.getindex( + a::BlockSparseArray, + perms::Vararg{<:BlockVector{<:Union{Int, Block{1}}}, N} + ) where {N} axs = axes(a) - perms = sectormergesortperm.(axs) merged_axes = sectormergesort.(axs) - result = zeros(eltype(a), merged_axes) - N = ndims(a) - - function find_dest_and_pos(perm, src_block::Block{1}) + function find_dest_and_pos(perm::BlockVector, src_block::Block{1}) for (igroup, grp) in enumerate(blocks(perm)) - pos = findfirst(==(src_block), grp) + pos = findfirst(i -> _asblock(i) == src_block, grp) isnothing(pos) || return (igroup, pos) end return error("Block $src_block not found in permutation") end - # Group stored source blocks by their destination block + # Group stored source blocks by destination grouped block and local offsets. dest_groups = Dict{NTuple{N, Int}, Vector{Tuple{Block{N, Int}, NTuple{N, Int}}}}() - for I in eachblockstoredindex(a) src_block_indices = Tuple(I) - dest_and_pos = ntuple(N) do dim return find_dest_and_pos(perms[dim], src_block_indices[dim]) end @@ -198,13 +199,12 @@ function sectormergesort(a::AbstractArray) offsets = ntuple(N) do dim igroup, pos = dest_and_pos[dim] grp = blocks(perms[dim])[igroup] - return sum(length(axs[dim][grp[k]]) for k in 1:(pos - 1); init = 0) + return sum(length(axs[dim][_asblock(grp[k])]) for k in 1:(pos - 1); init = 0) end push!(get!(dest_groups, dest_idx, []), (I, offsets)) end - # Build each destination block from its contributing source blocks for (dest_idx, src_list) in dest_groups dest_size = ntuple(dim -> length(merged_axes[dim][Block(dest_idx[dim])]), N) dest_data = zeros(eltype(a), dest_size) @@ -222,3 +222,8 @@ function sectormergesort(a::AbstractArray) return result end + +function sectormergesort(a::AbstractArray) + I = sectormergesortperm.(axes(a)) + return a[I...] +end From 34edc832eb6d4d668b18ae7d39399645323fcca6 Mon Sep 17 00:00:00 2001 From: mtfishman Date: Mon, 9 Mar 2026 13:57:21 -0400 Subject: [PATCH 03/13] Cleanup --- src/tensoralgebra.jl | 113 ++++--------------------------------------- 1 file changed, 9 insertions(+), 104 deletions(-) diff --git a/src/tensoralgebra.jl b/src/tensoralgebra.jl index fba6a9a3..5e0f445b 100644 --- a/src/tensoralgebra.jl +++ b/src/tensoralgebra.jl @@ -1,5 +1,5 @@ -using BlockArrays: Block, BlockVector, blocks, eachblockaxes1 -using BlockSparseArrays: BlockSparseArray, blockrange, blockreshape, eachblockstoredindex +using BlockArrays: blocks, eachblockaxes1 +using BlockSparseArrays: BlockSparseArray, blockrange, blockreshape using GradedArrays: GradedArray, GradedUnitRange, SectorRange, flip, invblockperm, sectormergesortperm, sectorsortperm, trivial, unmerged_tensor_product, × using TensorAlgebra: TensorAlgebra, AbstractBlockPermutation, BlockedTuple, FusionStyle, @@ -118,109 +118,14 @@ function TensorAlgebra.unmatricize( return a end - # Compute the unmerged fused axes (block structure before sectormergesort). + # First, fuse axes to get `sectormergesortperm`. + # Then unpermute the blocks. fused_axes = matricize_axes(BlockReshapeFusion(), m, codomain_axes, domain_axes) - - # Split merged blocks back to unmerged block structure, then unmatricize. - m_unmerged = unsectormergesort(m, fused_axes) - return unmatricize(FusionStyle(BlockSparseArray), m_unmerged, blocked_axes) -end - -# Split merged blocks back to unmerged block structure (inverse of sectormergesort for arrays). -_asblock(i::Block{1}) = i -_asblock(i::Integer) = Block(i) - -function unsectormergesort(m::AbstractArray, original_axes) - N = ndims(m) - perms = sectormergesortperm.(original_axes) - - result = zeros(eltype(m), original_axes) - - for I in eachblockstoredindex(m) - dest_idx = map(Int, Tuple(I)) - merged_block_data = m[I] - - # For each dimension, find which source blocks map to this merged block - # and their offsets within it. - src_groups = ntuple(N) do dim - grp = blocks(perms[dim])[dest_idx[dim]] - cumoff = 0 - map(eachindex(grp)) do k - off = cumoff - src_block = _asblock(grp[k]) - cumoff += length(original_axes[dim][src_block]) - return (src_block, off) - end - end - - # Iterate over all combinations of source blocks. - for src_combo in Iterators.product(src_groups...) - src_blocks = ntuple(dim -> src_combo[dim][1], N) # Block{1, Int64} - src_offsets = ntuple(dim -> src_combo[dim][2], N) - src_sizes = ntuple(dim -> length(original_axes[dim][src_blocks[dim]]), N) - src_region = - ntuple(dim -> (src_offsets[dim] + 1):(src_offsets[dim] + src_sizes[dim]), N) - src_data = merged_block_data[src_region...] - if !all(iszero, src_data) - result[src_blocks...] = src_data - end - end - end - - return result -end - -# Materialize grouped block indexing used by `sectormergesort`. -function Base.getindex( - a::BlockSparseArray, - perms::Vararg{<:BlockVector{<:Union{Int, Block{1}}}, N} - ) where {N} - axs = axes(a) - merged_axes = sectormergesort.(axs) - result = zeros(eltype(a), merged_axes) - - function find_dest_and_pos(perm::BlockVector, src_block::Block{1}) - for (igroup, grp) in enumerate(blocks(perm)) - pos = findfirst(i -> _asblock(i) == src_block, grp) - isnothing(pos) || return (igroup, pos) - end - return error("Block $src_block not found in permutation") - end - - # Group stored source blocks by destination grouped block and local offsets. - dest_groups = Dict{NTuple{N, Int}, Vector{Tuple{Block{N, Int}, NTuple{N, Int}}}}() - for I in eachblockstoredindex(a) - src_block_indices = Tuple(I) - dest_and_pos = ntuple(N) do dim - return find_dest_and_pos(perms[dim], src_block_indices[dim]) - end - dest_idx = ntuple(dim -> dest_and_pos[dim][1], N) - - offsets = ntuple(N) do dim - igroup, pos = dest_and_pos[dim] - grp = blocks(perms[dim])[igroup] - return sum(length(axs[dim][_asblock(grp[k])]) for k in 1:(pos - 1); init = 0) - end - - push!(get!(dest_groups, dest_idx, []), (I, offsets)) - end - - for (dest_idx, src_list) in dest_groups - dest_size = ntuple(dim -> length(merged_axes[dim][Block(dest_idx[dim])]), N) - dest_data = zeros(eltype(a), dest_size) - - for (I, offsets) in src_list - src_data = a[I] - src_sizes = size(src_data) - dest_region = - ntuple(dim -> (offsets[dim] + 1):(offsets[dim] + src_sizes[dim]), N) - dest_data[dest_region...] = src_data - end - - result[Block.(dest_idx)...] = dest_data - end - - return result + blockperms = sectorsortperm.(fused_axes) + sorted_axes = map((r, I) -> only(axes(r[I])), fused_axes, blockperms) + m_unblocked = m[sorted_axes...] + m_blockpermed = m_unblocked[invblockperm.(blockperms)...] + return unmatricize(FusionStyle(BlockSparseArray), m_blockpermed, blocked_axes) end function sectormergesort(a::AbstractArray) From a1c248b5f66360f8ce49df5ef10af739186e850e Mon Sep 17 00:00:00 2001 From: mtfishman Date: Mon, 9 Mar 2026 14:48:37 -0400 Subject: [PATCH 04/13] GradedUnitRange slicing fixes --- src/gradedarray.jl | 29 +++++++++++++++++++++++++++++ test/test_gradedunitrange.jl | 2 +- 2 files changed, 30 insertions(+), 1 deletion(-) diff --git a/src/gradedarray.jl b/src/gradedarray.jl index a5a0c299..06312c7d 100644 --- a/src/gradedarray.jl +++ b/src/gradedarray.jl @@ -74,6 +74,35 @@ function BlockSparseArrays.mortar_axis(geachblockaxis::AbstractVector{<:SectorUn return blockrange(geachblockaxis) end +# Keep graded labels when slicing by block vectors. +function BlockSparseArrays.blockedunitrange_getindices( + g::GradedUnitRange, indices::Vector{<:BlockArrays.Block{1}} + ) + gblocks = map(index -> g[index], indices) + new_axis = mortar_axis(sectorrange.(sector.(gblocks), Base.OneTo.(length.(gblocks)))) + return BlockArrays.mortar(gblocks, (new_axis,)) +end + +function BlockSparseArrays.blockedunitrange_getindices( + g::GradedUnitRange, indices::BlockArrays.AbstractBlockVector{<:BlockArrays.Block{1}} + ) + blks = map(bs -> BlockArrays.mortar(map(b -> g[b], bs)), BlockArrays.blocks(indices)) + new_sectors = map(bs -> sectors(g)[Int.(bs)], BlockArrays.blocks(indices)) + @assert all(allequal.(new_sectors)) + new_lengths = length.(blks) + new_axis = mortar_axis(sectorrange.(first.(new_sectors), Base.OneTo.(new_lengths))) + return BlockArrays.mortar(blks, (new_axis,)) +end +function Base.getindex(g::GradedUnitRange, indices::Vector{<:BlockArrays.Block{1}}) + return BlockSparseArrays.blockedunitrange_getindices(g, indices) +end +function Base.getindex( + g::GradedUnitRange, + indices::BlockArrays.AbstractBlockVector{<:BlockArrays.Block{1}} + ) + return BlockSparseArrays.blockedunitrange_getindices(g, indices) +end + to_gradedrange(g::GradedUnitRange) = g function Base.show(io::IO, g::GradedUnitRange) diff --git a/test/test_gradedunitrange.jl b/test/test_gradedunitrange.jl index 730701be..92c8c782 100644 --- a/test/test_gradedunitrange.jl +++ b/test/test_gradedunitrange.jl @@ -214,7 +214,7 @@ using Test: @test, @test_throws, @testset v = mortar([[Block(2), Block(2)], [Block(1)]]) a = g[v] @test a isa BlockVector - @test_broken only(axes(a)) isa GradedOneTo + @test only(axes(a)) isa GradedOneTo @test_broken space_isequal( only(axes(a)), gradedrange(["y" => 6, "x" => 2]; isdual = isdual(g)) From 0fa7b2bd54eb9a9a8c308bb6ab9e0c0c57a837aa Mon Sep 17 00:00:00 2001 From: mtfishman Date: Mon, 9 Mar 2026 15:43:53 -0400 Subject: [PATCH 05/13] More progress --- src/GradedArrays.jl | 1 + src/gradedarray.jl | 11 ++++++++--- src/sectorarray.jl | 16 ++++++++++++++++ src/tensoralgebra.jl | 29 ++++++++++++++++++++++++++++- 4 files changed, 53 insertions(+), 4 deletions(-) diff --git a/src/GradedArrays.jl b/src/GradedArrays.jl index a54040ef..6e6f47a6 100644 --- a/src/GradedArrays.jl +++ b/src/GradedArrays.jl @@ -15,6 +15,7 @@ export dual, flip, gradedrange, isdual, # imports # ------- +import BlockSparseArrays import KroneckerArrays: × using BlockArrays: BlockArrays, Block, blocksize using BlockSparseArrays: @view!, AbstractBlockSparseArray, BlockOneTo, blockrange diff --git a/src/gradedarray.jl b/src/gradedarray.jl index 06312c7d..3d656828 100644 --- a/src/gradedarray.jl +++ b/src/gradedarray.jl @@ -79,7 +79,8 @@ function BlockSparseArrays.blockedunitrange_getindices( g::GradedUnitRange, indices::Vector{<:BlockArrays.Block{1}} ) gblocks = map(index -> g[index], indices) - new_axis = mortar_axis(sectorrange.(sector.(gblocks), Base.OneTo.(length.(gblocks)))) + new_multiplicities = sector_multiplicity.(gblocks) + new_axis = mortar_axis(sectorrange.(sector.(gblocks), Base.OneTo.(new_multiplicities))) return BlockArrays.mortar(gblocks, (new_axis,)) end @@ -89,8 +90,12 @@ function BlockSparseArrays.blockedunitrange_getindices( blks = map(bs -> BlockArrays.mortar(map(b -> g[b], bs)), BlockArrays.blocks(indices)) new_sectors = map(bs -> sectors(g)[Int.(bs)], BlockArrays.blocks(indices)) @assert all(allequal.(new_sectors)) - new_lengths = length.(blks) - new_axis = mortar_axis(sectorrange.(first.(new_sectors), Base.OneTo.(new_lengths))) + new_multiplicities = map(BlockArrays.blocks(indices)) do bs + return sum(b -> sector_multiplicity(g[b]), bs; init = 0) + end + new_axis = mortar_axis( + sectorrange.(first.(new_sectors), Base.OneTo.(new_multiplicities)) + ) return BlockArrays.mortar(blks, (new_axis,)) end function Base.getindex(g::GradedUnitRange, indices::Vector{<:BlockArrays.Block{1}}) diff --git a/src/sectorarray.jl b/src/sectorarray.jl index 716986ab..da132737 100644 --- a/src/sectorarray.jl +++ b/src/sectorarray.jl @@ -297,6 +297,22 @@ function Base.similar( similar(A, elt, kroneckerfactors.(axs, 2)) ) end +# disambiguate for unblocked views of BlockSparseArrays with graded axes +function Base.similar( + A::BlockSparseArrays.UnblockedSubArray, + elt::Type, + axs::Tuple{SectorUnitRange, Vararg{SectorUnitRange}} + ) + data_axes = kroneckerfactors.(axs, 2) + data = similar(Array{elt, length(axs)}, data_axes) + for I in CartesianIndices(data) + @inbounds data[I] = A[I] + end + return SectorArray( + kroneckerfactors.(axs, 1), + data + ) +end function Base.similar( ::Type{A}, axs::Tuple{SectorUnitRange, Vararg{SectorUnitRange}} diff --git a/src/tensoralgebra.jl b/src/tensoralgebra.jl index 5e0f445b..21987f6d 100644 --- a/src/tensoralgebra.jl +++ b/src/tensoralgebra.jl @@ -1,7 +1,9 @@ -using BlockArrays: blocks, eachblockaxes1 +using ArrayLayouts: ArrayLayouts +using BlockArrays: Block, BlockLayout, blocks, eachblockaxes1 using BlockSparseArrays: BlockSparseArray, blockrange, blockreshape using GradedArrays: GradedArray, GradedUnitRange, SectorRange, flip, invblockperm, sectormergesortperm, sectorsortperm, trivial, unmerged_tensor_product, × +using SparseArraysBase: SparseLayout using TensorAlgebra: TensorAlgebra, AbstractBlockPermutation, BlockedTuple, FusionStyle, ReshapeFusion, matricize, matricize_axes, tensor_product_axis, trivialbiperm, tuplemortar, unmatricize @@ -12,6 +14,31 @@ TensorAlgebra.FusionStyle(::Type{<:SectorDelta}) = SectorFusion() TensorAlgebra.FusionStyle(::Type{<:GradedArray}) = SectorFusion() TensorAlgebra.FusionStyle(::Type{<:SectorUnitRange}) = SectorFusion() +function ArrayLayouts.sub_materialize( + layout::BlockLayout{<:SparseLayout}, + a::SubArray{ + <:Any, + <:Any, + <:BlockSparseArrays.BlockSparseArray{<:Any, <:Any, <:SectorArray}, + }, + axes + ) + blocktype_a = BlockSparseArrays.blocktype(parent(a)) + a_dest = BlockSparseArray{eltype(a), length(axes), blocktype_a}(undef, axes) + for I in CartesianIndices(blocks(a_dest)) + b = Block(Tuple(I)) + block_a = @view a[b] + if !iszero(block_a) + block_dest = similar(a_dest[b]) + for J in CartesianIndices(block_dest) + @inbounds block_dest[J] = block_a[J] + end + a_dest[b] = block_dest + end + end + return a_dest +end + using BlockArrays: AbstractBlockArray const BlockReshapeFusion = typeof(FusionStyle(AbstractBlockArray)) From bf57f68a4811289079f3aa498f0227f5fdc2a441 Mon Sep 17 00:00:00 2001 From: mtfishman Date: Mon, 9 Mar 2026 17:32:41 -0400 Subject: [PATCH 06/13] Bump versions --- Project.toml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/Project.toml b/Project.toml index bc4b8d66..2005021e 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "GradedArrays" uuid = "bc96ca6e-b7c8-4bb6-888e-c93f838762c2" -version = "0.6.17" +version = "0.6.18" authors = ["ITensor developers and contributors"] [workspace] @@ -34,7 +34,7 @@ GradedArraysSUNRepresentationsExt = "SUNRepresentations" [compat] ArrayLayouts = "1" BlockArrays = "1.6" -BlockSparseArrays = "0.10" +BlockSparseArrays = "0.10.35" Compat = "4.16" FillArrays = "1.13" HalfIntegers = "1.6" From ec0e37696852d57d64415f4d2779b67b2544911e Mon Sep 17 00:00:00 2001 From: mtfishman Date: Mon, 9 Mar 2026 17:35:48 -0400 Subject: [PATCH 07/13] Import style --- src/GradedArrays.jl | 7 +++---- src/sectorrange.jl | 2 +- 2 files changed, 4 insertions(+), 5 deletions(-) diff --git a/src/GradedArrays.jl b/src/GradedArrays.jl index 6e6f47a6..002f629e 100644 --- a/src/GradedArrays.jl +++ b/src/GradedArrays.jl @@ -15,12 +15,11 @@ export dual, flip, gradedrange, isdual, # imports # ------- -import BlockSparseArrays -import KroneckerArrays: × using BlockArrays: BlockArrays, Block, blocksize -using BlockSparseArrays: @view!, AbstractBlockSparseArray, BlockOneTo, blockrange +using BlockSparseArrays: + BlockSparseArrays, @view!, AbstractBlockSparseArray, BlockOneTo, blockrange using KroneckerArrays -using KroneckerArrays: AbstractKroneckerArray, CartesianProductUnitRange +using KroneckerArrays: KroneckerArrays, AbstractKroneckerArray, CartesianProductUnitRange, × using LinearAlgebra: LinearAlgebra, Adjoint using SparseArraysBase: isstored using TypeParameterAccessors: type_parameters, unspecify_type_parameters diff --git a/src/sectorrange.jl b/src/sectorrange.jl index a7dfb2f3..7d12c74f 100644 --- a/src/sectorrange.jl +++ b/src/sectorrange.jl @@ -1,6 +1,6 @@ # This file defines the interface for type Sector # all fusion categories (Z{2}, SU2, Ising...) are subtypes of Sector -import TensorKitSectors as TKS +using TensorKitSectors: TensorKitSectors as TKS """ SectorRange(sector::TKS.Sector, isdual::Bool) From 3bdc65122af44d7b177d2d744943eb39690437b7 Mon Sep 17 00:00:00 2001 From: mtfishman Date: Mon, 9 Mar 2026 17:45:57 -0400 Subject: [PATCH 08/13] Bump BlockSparseArrays compat --- Project.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Project.toml b/Project.toml index 2005021e..90d2b536 100644 --- a/Project.toml +++ b/Project.toml @@ -34,7 +34,7 @@ GradedArraysSUNRepresentationsExt = "SUNRepresentations" [compat] ArrayLayouts = "1" BlockArrays = "1.6" -BlockSparseArrays = "0.10.35" +BlockSparseArrays = "0.10.36" Compat = "4.16" FillArrays = "1.13" HalfIntegers = "1.6" From 03e7e009a77779a035512a3fd57d1e5b1bad9aa1 Mon Sep 17 00:00:00 2001 From: mtfishman Date: Mon, 9 Mar 2026 17:49:03 -0400 Subject: [PATCH 09/13] Fix BlockSparseArrays compat --- Project.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Project.toml b/Project.toml index 90d2b536..2005021e 100644 --- a/Project.toml +++ b/Project.toml @@ -34,7 +34,7 @@ GradedArraysSUNRepresentationsExt = "SUNRepresentations" [compat] ArrayLayouts = "1" BlockArrays = "1.6" -BlockSparseArrays = "0.10.36" +BlockSparseArrays = "0.10.35" Compat = "4.16" FillArrays = "1.13" HalfIntegers = "1.6" From fa79e4981093f81c5eda5d190ee1b6def0a682ab Mon Sep 17 00:00:00 2001 From: mtfishman Date: Mon, 9 Mar 2026 17:58:14 -0400 Subject: [PATCH 10/13] Improve import style --- src/GradedArrays.jl | 2 +- src/gradedarray.jl | 12 ++++++------ src/sectorarray.jl | 4 ++-- src/sectorproduct.jl | 29 ++++++++++++++++------------- 4 files changed, 25 insertions(+), 22 deletions(-) diff --git a/src/GradedArrays.jl b/src/GradedArrays.jl index 002f629e..6d589187 100644 --- a/src/GradedArrays.jl +++ b/src/GradedArrays.jl @@ -19,7 +19,7 @@ using BlockArrays: BlockArrays, Block, blocksize using BlockSparseArrays: BlockSparseArrays, @view!, AbstractBlockSparseArray, BlockOneTo, blockrange using KroneckerArrays -using KroneckerArrays: KroneckerArrays, AbstractKroneckerArray, CartesianProductUnitRange, × +using KroneckerArrays: KroneckerArrays, AbstractKroneckerArray, CartesianProductUnitRange using LinearAlgebra: LinearAlgebra, Adjoint using SparseArraysBase: isstored using TypeParameterAccessors: type_parameters, unspecify_type_parameters diff --git a/src/gradedarray.jl b/src/gradedarray.jl index 3d656828..96ac77b6 100644 --- a/src/gradedarray.jl +++ b/src/gradedarray.jl @@ -38,16 +38,16 @@ isdual(g::GradedUnitRange) = isdual(first(eachblockaxis(g))) # crash for empty. flux(a::AbstractBlockedUnitRange, I::Block{1}) = flux(a[I]) -function ×(g1::GradedOneTo, g2::GradedOneTo) +function KroneckerArrays.:×(g1::GradedOneTo, g2::GradedOneTo) v = vec([a × b for a in eachblockaxis(g1), b in eachblockaxis(g2)]) return mortar_axis(v) end -×(g::GradedUnitRange, a::AbstractUnitRange) = ×(g, to_gradedrange(a)) -×(a::AbstractUnitRange, g::GradedUnitRange) = ×(to_gradedrange(a), g) -×(g::GradedUnitRange, a::SectorRange) = ×(g, to_gradedrange(a)) -×(a::SectorRange, g::GradedUnitRange) = ×(to_gradedrange(a), g) -function ×(g1::GradedUnitRange, g2::GradedUnitRange) +KroneckerArrays.:×(g::GradedUnitRange, a::AbstractUnitRange) = ×(g, to_gradedrange(a)) +KroneckerArrays.:×(a::AbstractUnitRange, g::GradedUnitRange) = ×(to_gradedrange(a), g) +KroneckerArrays.:×(g::GradedUnitRange, a::SectorRange) = ×(g, to_gradedrange(a)) +KroneckerArrays.:×(a::SectorRange, g::GradedUnitRange) = ×(to_gradedrange(a), g) +function KroneckerArrays.:×(g1::GradedUnitRange, g2::GradedUnitRange) v = vec([a × b for a in eachblockaxis(g1), b in eachblockaxis(g2)]) return mortar_axis(v) end diff --git a/src/sectorarray.jl b/src/sectorarray.jl index da132737..b54c33a0 100644 --- a/src/sectorarray.jl +++ b/src/sectorarray.jl @@ -33,8 +33,8 @@ function sectorrange( return sectorrange(to_sector(sector), args...) end -×(a::SectorRange, g::AbstractUnitRange) = cartesianrange(a, g) -×(g::AbstractUnitRange, a::SectorRange) = cartesianrange(a, g) +KroneckerArrays.:×(a::SectorRange, g::AbstractUnitRange) = cartesianrange(a, g) +KroneckerArrays.:×(g::AbstractUnitRange, a::SectorRange) = cartesianrange(a, g) to_gradedrange(g::SectorUnitRange) = mortar_axis([g]) diff --git a/src/sectorproduct.jl b/src/sectorproduct.jl index 7734b8e6..5c1afb89 100644 --- a/src/sectorproduct.jl +++ b/src/sectorproduct.jl @@ -161,38 +161,41 @@ end # ================================= Cartesian Product ==================================== -const sectorproduct = × +const sectorproduct = KroneckerArrays.:× -×(c::SectorRange) = SectorRange(SectorProduct(label(c))) -×(c1::SectorRange, c2::SectorRange) = SectorRange(×(label(c1), label(c2))) -×(c1::TKS.Sector, c2::TKS.Sector) = ×(SectorProduct(c1), SectorProduct(c2)) +KroneckerArrays.:×(c::SectorRange) = SectorRange(SectorProduct(label(c))) +KroneckerArrays.:×(c1::SectorRange, c2::SectorRange) = SectorRange(×(label(c1), label(c2))) +KroneckerArrays.:×(c1::TKS.Sector, c2::TKS.Sector) = ×(SectorProduct(c1), SectorProduct(c2)) -function ×(p1::SectorProduct{<:Tuple}, p2::SectorProduct{<:Tuple}) +function KroneckerArrays.:×(p1::SectorProduct{<:Tuple}, p2::SectorProduct{<:Tuple}) return SectorProduct(arguments(p1)..., arguments(p2)...) end -function ×(p1::SectorProduct{<:NamedTuple}, p2::SectorProduct{<:NamedTuple}) +function KroneckerArrays.:×( + p1::SectorProduct{<:NamedTuple}, + p2::SectorProduct{<:NamedTuple} + ) isdisjoint(keys(arguments(p1)), keys(arguments(p2))) || throw(ArgumentError("keys of SectorProducts must be distinct")) return SectorProduct(merge(arguments(p1), arguments(p2))) end -function ×(a::SectorProduct, b::SectorProduct) +function KroneckerArrays.:×(a::SectorProduct, b::SectorProduct) isempty(arguments(a)) && return b isempty(arguments(b)) && return a throw(MethodError(×, typeof.((a, b)))) end -×(nt1::NamedTuple) = to_sector(nt1) -×(nt1::NamedTuple, nt2::NamedTuple) = ×(to_sector(nt1), to_sector(nt2)) -×(c1::NamedTuple, c2::SectorRange) = ×(to_sector(c1), c2) -×(c1::SectorRange, c2::NamedTuple) = ×(c1, to_sector(c2)) +KroneckerArrays.:×(nt1::NamedTuple) = to_sector(nt1) +KroneckerArrays.:×(nt1::NamedTuple, nt2::NamedTuple) = ×(to_sector(nt1), to_sector(nt2)) +KroneckerArrays.:×(c1::NamedTuple, c2::SectorRange) = ×(to_sector(c1), c2) +KroneckerArrays.:×(c1::SectorRange, c2::NamedTuple) = ×(c1, to_sector(c2)) -function ×(pairs::Pair...) +function KroneckerArrays.:×(pairs::Pair...) keys = Symbol.(first.(pairs)) vals = last.(pairs) return ×(NamedTuple{keys}(vals)) end -function ×(sr1::SectorOneTo, sr2::SectorOneTo) +function KroneckerArrays.:×(sr1::SectorOneTo, sr2::SectorOneTo) isdual(sr1) == isdual(sr2) || throw(ArgumentError("SectorProduct duality must match")) sr = sectorrange( sector(sr1) × sector(sr2), From fde2d716913f53d09fbb08b225a250e7b9628053 Mon Sep 17 00:00:00 2001 From: mtfishman Date: Mon, 9 Mar 2026 19:21:46 -0400 Subject: [PATCH 11/13] Simplify --- src/sectorarray.jl | 11 +++++++++++ src/tensoralgebra.jl | 29 +---------------------------- 2 files changed, 12 insertions(+), 28 deletions(-) diff --git a/src/sectorarray.jl b/src/sectorarray.jl index b54c33a0..ed7f2a50 100644 --- a/src/sectorarray.jl +++ b/src/sectorarray.jl @@ -329,6 +329,17 @@ function Base.copy!(C::SectorArray, A::SectorArray) copy!(arg2(C), arg2(A)) return C end +function Base.copyto!(dst::SectorArray, src::BlockSparseArrays.UnblockedSubArray) + @invoke copyto!(dst.data::AbstractArray, src::AbstractArray) + return dst +end +function Base.copyto!( + dst::SectorArray, + src::SubArray{<:Any, <:Any, <:BlockSparseArrays.AbstractBlockSparseArray} + ) + @invoke copyto!(dst.data::AbstractArray, src::AbstractArray) + return dst +end function Base.convert( ::Type{SectorArray{T₁, N, I, A}}, diff --git a/src/tensoralgebra.jl b/src/tensoralgebra.jl index 21987f6d..5e0f445b 100644 --- a/src/tensoralgebra.jl +++ b/src/tensoralgebra.jl @@ -1,9 +1,7 @@ -using ArrayLayouts: ArrayLayouts -using BlockArrays: Block, BlockLayout, blocks, eachblockaxes1 +using BlockArrays: blocks, eachblockaxes1 using BlockSparseArrays: BlockSparseArray, blockrange, blockreshape using GradedArrays: GradedArray, GradedUnitRange, SectorRange, flip, invblockperm, sectormergesortperm, sectorsortperm, trivial, unmerged_tensor_product, × -using SparseArraysBase: SparseLayout using TensorAlgebra: TensorAlgebra, AbstractBlockPermutation, BlockedTuple, FusionStyle, ReshapeFusion, matricize, matricize_axes, tensor_product_axis, trivialbiperm, tuplemortar, unmatricize @@ -14,31 +12,6 @@ TensorAlgebra.FusionStyle(::Type{<:SectorDelta}) = SectorFusion() TensorAlgebra.FusionStyle(::Type{<:GradedArray}) = SectorFusion() TensorAlgebra.FusionStyle(::Type{<:SectorUnitRange}) = SectorFusion() -function ArrayLayouts.sub_materialize( - layout::BlockLayout{<:SparseLayout}, - a::SubArray{ - <:Any, - <:Any, - <:BlockSparseArrays.BlockSparseArray{<:Any, <:Any, <:SectorArray}, - }, - axes - ) - blocktype_a = BlockSparseArrays.blocktype(parent(a)) - a_dest = BlockSparseArray{eltype(a), length(axes), blocktype_a}(undef, axes) - for I in CartesianIndices(blocks(a_dest)) - b = Block(Tuple(I)) - block_a = @view a[b] - if !iszero(block_a) - block_dest = similar(a_dest[b]) - for J in CartesianIndices(block_dest) - @inbounds block_dest[J] = block_a[J] - end - a_dest[b] = block_dest - end - end - return a_dest -end - using BlockArrays: AbstractBlockArray const BlockReshapeFusion = typeof(FusionStyle(AbstractBlockArray)) From 695696795ed07ee3e61c0e3c4be02c3955e77d56 Mon Sep 17 00:00:00 2001 From: mtfishman Date: Tue, 10 Mar 2026 10:24:06 -0400 Subject: [PATCH 12/13] Bring back comments --- src/tensoralgebra.jl | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/src/tensoralgebra.jl b/src/tensoralgebra.jl index 5e0f445b..c454078a 100644 --- a/src/tensoralgebra.jl +++ b/src/tensoralgebra.jl @@ -121,8 +121,14 @@ function TensorAlgebra.unmatricize( # First, fuse axes to get `sectormergesortperm`. # Then unpermute the blocks. fused_axes = matricize_axes(BlockReshapeFusion(), m, codomain_axes, domain_axes) + blockperms = sectorsortperm.(fused_axes) sorted_axes = map((r, I) -> only(axes(r[I])), fused_axes, blockperms) + + # TODO: This is doing extra copies of the blocks, + # use `@view a[axes_prod...]` instead. + # That will require implementing some reindexing logic + # for this combination of slicing. m_unblocked = m[sorted_axes...] m_blockpermed = m_unblocked[invblockperm.(blockperms)...] return unmatricize(FusionStyle(BlockSparseArray), m_blockpermed, blocked_axes) From b6b4c4cbe5ae6597dcf769e9466af9bae0790e13 Mon Sep 17 00:00:00 2001 From: mtfishman Date: Tue, 10 Mar 2026 10:48:49 -0400 Subject: [PATCH 13/13] Add back comment on sectormergesort --- src/tensoralgebra.jl | 1 + 1 file changed, 1 insertion(+) diff --git a/src/tensoralgebra.jl b/src/tensoralgebra.jl index c454078a..650d87fd 100644 --- a/src/tensoralgebra.jl +++ b/src/tensoralgebra.jl @@ -134,6 +134,7 @@ function TensorAlgebra.unmatricize( return unmatricize(FusionStyle(BlockSparseArray), m_blockpermed, blocked_axes) end +# Sort the blocks by sector and then merge the common sectors. function sectormergesort(a::AbstractArray) I = sectormergesortperm.(axes(a)) return a[I...]