From 311e60254d98a1e79bc4c24fafef897c317ff33b Mon Sep 17 00:00:00 2001 From: lkdvos Date: Tue, 25 Aug 2026 12:58:05 -0400 Subject: [PATCH 1/3] fix: use cartesian keys in `nonzero_pairs` `pairs(A::AbstractVector)` uses linear indices, so `nonzero_pairs` yielded `Int` keys for the parent of a one-dimensional dense `BlockTensorMap`, inconsistent with `nonzero_keys`. Consumers comparing against `CartesianIndex` then silently matched nothing, which made `t[:]` and `t[1:2]` return uninitialized blocks. Co-Authored-By: Claude Opus 5 (1M context) --- src/tensors/abstractblocktensor/sparsity.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/tensors/abstractblocktensor/sparsity.jl b/src/tensors/abstractblocktensor/sparsity.jl index 6c059c7..76ecfc5 100644 --- a/src/tensors/abstractblocktensor/sparsity.jl +++ b/src/tensors/abstractblocktensor/sparsity.jl @@ -7,7 +7,7 @@ nonzero_length(t::AbstractBlockTensorMap) = nonzero_length(parent(t)) nonzero_values(A::AbstractArray) = values(A) nonzero_keys(A::AbstractArray) = eachindex(IndexCartesian(), A) -nonzero_pairs(A::AbstractArray) = pairs(A) +nonzero_pairs(A::AbstractArray) = pairs(IndexCartesian(), A) nonzero_length(A::AbstractArray) = length(A) issparse(t::AbstractTensorMap) = false From 0ba68ee22ba5aea8d66013317b43e50063022227 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Tue, 25 Aug 2026 12:58:19 -0400 Subject: [PATCH 2/3] perf: invert index maps when slicing block tensors Slicing `getindex` scanned the whole source index grid once per stored block, making it `O(nnz * prod(length.(indices)))` and materializing the slice region as an `Array{CartesianIndex}`. Invert the index maps per dimension instead, so each stored block finds its destinations in `O(1)`: n nnz before after speedup 16 16 0.0045 ms 0.0015 ms 3x 64 64 0.1452 ms 0.0036 ms 40x 256 256 8.9455 ms 0.0107 ms 836x 512 512 71.1836 ms 0.0188 ms 3786x (banded sparse block tensor, sliced to every other channel; the dense `BlockTensorMap` variant goes from 7.94 ms to 0.087 ms at n = 64.) `SparseTensorArray` had a second copy of the same logic for slicing the parent array; both now share `_copyslice!`. Along the way this fixes four bugs in the old index handling: - Repeated indices dropped blocks instead of duplicating them, and for a dense destination the dropped slot was returned as uninitialized memory. They now duplicate, as they do for `AbstractArray`. - Single-index slicing of one-dimensional block tensors (`t[:]`, `t[1:2]`) returned uninitialized blocks. - `Integer` indices other than `Int` threw: `t[UInt(1), 1:2, 1]` a `BoundsError`, `t[Int8(1), Int8(1), Int8(1)]` a `MethodError`. - Logical masks threw `CanonicalIndexError` on the parent array path, `parent(t)[[true, false, true], :, :]`. The `Vararg{Strided.SliceIndex}` methods are kept: they disambiguate against TensorKit's own `getindex`/`setindex!` for `AbstractTensorMap`, so removing them turns `t[:, :, :]` into a `MethodError`. Co-Authored-By: Claude Opus 5 (1M context) --- docs/src/blocktensors.md | 1 + src/BlockTensorKit.jl | 1 + src/auxiliary/sliceindices.jl | 57 +++++++ src/auxiliary/sparsetensorarray.jl | 23 +-- .../abstractblocktensor/abstractarray.jl | 112 +++++--------- test/abstracttensor/indexing.jl | 146 ++++++++++++------ 6 files changed, 197 insertions(+), 143 deletions(-) create mode 100644 src/auxiliary/sliceindices.jl diff --git a/docs/src/blocktensors.md b/docs/src/blocktensors.md index bba2b9c..5c2e23a 100644 --- a/docs/src/blocktensors.md +++ b/docs/src/blocktensors.md @@ -71,6 +71,7 @@ s[1] += 2 * s[1] Slicing operations are also supported, and the `AbstractBlockTensorMap` can be sliced in the same way as an `AbstractArray{AbstractTensorMap}`. There is however one elementary difference: as the slices still contain tensors with the same amount of legs, there can be no reduction in the number of dimensions. In particular, in contrast to `AbstractArray`, scalar dimensions are not discarded, and as a result, linear index slicing is not allowed. +Repeated indices duplicate the selected tensors, as they would for an `AbstractArray`. ```@repl blocktensors ndims(t[1, 1, :]) == 3 diff --git a/src/BlockTensorKit.jl b/src/BlockTensorKit.jl index 4d354b8..09cbd7a 100644 --- a/src/BlockTensorKit.jl +++ b/src/BlockTensorKit.jl @@ -38,6 +38,7 @@ import TupleTools as TT import MatrixAlgebraKit as MAK include("auxiliary/blockarrays.jl") +include("auxiliary/sliceindices.jl") # Spaces include("vectorspaces/sumspace.jl") diff --git a/src/auxiliary/sliceindices.jl b/src/auxiliary/sliceindices.jl new file mode 100644 index 0000000..f471102 --- /dev/null +++ b/src/auxiliary/sliceindices.jl @@ -0,0 +1,57 @@ +# Slice indices +# ------------- +# inverting index maps, to copy sliced blocks in `O(nnz)` instead of `O(nnz * length(dst))` + +const SliceIndex = Union{Strided.SliceIndex, AbstractVector{<:Integer}} + +_key_tuple(I::CartesianIndex) = I.I +_key_tuple(i::Integer) = (Int(i),) + +""" + _invert_index(n::Int, ind) -> m + +Invert an index into a dimension of length `n`, such that `_dstrange(m, i)` yields all +destination coordinates selecting source coordinate `i`. +""" +_invert_index(::Int, i::Integer) = Int(i) +_invert_index(::Int, r::AbstractUnitRange{Int}) = r +function _invert_index(n::Int, ind) + ptr = zeros(Int, n + 1) + for i in ind + 1 ≤ i ≤ n || throw(BoundsError(Base.OneTo(n), i)) + ptr[i + 1] += 1 + end + cumsum!(ptr, ptr) + dsts = Vector{Int}(undef, length(ind)) + pos = copy(ptr) + for (j, i) in enumerate(ind) + dsts[pos[i] += 1] = j + end + return ptr, dsts +end + +_dstrange(m::Int, i::Int) = i == m ? (1:1) : (1:0) +function _dstrange(m::AbstractUnitRange{Int}, i::Int) + j = i - first(m) + 1 + return 1 ≤ j ≤ length(m) ? (j:j) : (1:0) +end +_dstrange((ptr, dsts)::Tuple{Vector{Int}, Vector{Int}}, i::Int) = + view(dsts, (ptr[i] + 1):ptr[i + 1]) + +""" + _copyslice!(tdst, tsrc, inds::NTuple{N,Any}) -> tdst + +Copy the nonzero blocks of `tsrc` selected by `inds` into `tdst`, where `inds` holds one +normalized index per dimension of `tsrc`. +""" +function _copyslice!(tdst, tsrc, inds::NTuple{N, Any}) where {N} + maps = map(_invert_index, size(tsrc), inds) + for (I, v) in nonzero_pairs(tsrc) + rs = map(_dstrange, maps, _key_tuple(I)) + any(isempty, rs) && continue + for J in Iterators.product(rs...) + tdst[J...] = v + end + end + return tdst +end diff --git a/src/auxiliary/sparsetensorarray.jl b/src/auxiliary/sparsetensorarray.jl index 2756abb..347e640 100644 --- a/src/auxiliary/sparsetensorarray.jl +++ b/src/auxiliary/sparsetensorarray.jl @@ -168,33 +168,12 @@ end # non-scalar indexing # ------------------- # specialisations to have non-scalar indexing behave as expected - -_newindex(i::Int, range::Int) = i == range ? (1,) : nothing -function _newindex(i::Int, range::AbstractVector{Int}) - k = findfirst(==(i), range) - return k === nothing ? nothing : (k,) -end -_newindices(::Tuple{}, ::Tuple{}) = () -function _newindices(I::Tuple, indices::Tuple) - i = _newindex(I[1], indices[1]) - Itail = _newindices(Base.tail(I), Base.tail(indices)) - (i === nothing || Itail === nothing) && return nothing - return (i..., Itail...) -end - function Base._unsafe_getindex( ::IndexCartesian, t::SparseTensorArray{S, N₁, N₂, T, N}, I::Vararg{Union{Real, AbstractArray}, N}, ) where {S, N₁, N₂, T, N} dest = similar(t, eltype(t), space(eachspace(t)[I...])) - indices = Base.to_indices(t, I) - for (k, v) in t.data - newI = _newindices(k.I, indices) - if newI !== nothing - dest[newI...] = v - end - end - return dest + return _copyslice!(dest, t, Base.to_indices(t, I)) end # Space checking diff --git a/src/tensors/abstractblocktensor/abstractarray.jl b/src/tensors/abstractblocktensor/abstractarray.jl index e099a83..186720d 100644 --- a/src/tensors/abstractblocktensor/abstractarray.jl +++ b/src/tensors/abstractblocktensor/abstractarray.jl @@ -90,51 +90,30 @@ end getindex!(parent(t), I) # slicing getindex needs to correctly allocate output blocktensor: -const SliceIndex = Union{Strided.SliceIndex, AbstractVector{<:Union{Integer, Bool}}} - -Base.@propagate_inbounds function Base.getindex( - t::AbstractBlockTensorMap, indices::Vararg{SliceIndex} - ) - V = space(eachspace(t)[indices...]) - tdst = similar(t, V) +@propagate_inbounds Base.getindex(t::AbstractBlockTensorMap, indices::Vararg{SliceIndex}) = + _slice_getindex(t, indices...) +# disambiguate: TensorKit/src/tensors/abstracttensor.jl:540 +@propagate_inbounds Base.getindex(t::AbstractBlockTensorMap, indices::Vararg{Strided.SliceIndex}) = + _slice_getindex(t, indices...) + +@propagate_inbounds function _slice_getindex( + t::AbstractBlockTensorMap, indices::Vararg{Any, M} + ) where {M} + M == numind(t) || return _slice_getindex_single(t, indices...) + inds = Base.to_indices(t, indices) + inds isa NTuple{M, Int} && return parent(t)[inds...] + tdst = similar(t, space(eachspace(t)[inds...])) length(tdst) == 0 && return tdst - - # prevent discarding of singleton dimensions - indices′ = map(indices) do ind - return ind isa Int ? (ind:ind) : ind - end - Rsrc = CartesianIndices(t)[indices′...] - Rdst = CartesianIndices(tdst) - - for (I, v) in nonzero_pairs(t) - j = findfirst(==(I), Rsrc) - isnothing(j) && continue - tdst[Rdst[j]] = v - end - return tdst + return _copyslice!(tdst, t, inds) end -# disambiguate: -@propagate_inbounds function Base.getindex( - t::AbstractBlockTensorMap, indices::Vararg{Strided.SliceIndex} - ) - V = space(eachspace(t)[indices...]) - tdst = similar(t, V) - length(tdst) == 0 && return tdst - - # prevent discarding of singleton dimensions - indices′ = map(indices) do ind - return ind isa Int ? (ind:ind) : ind - end - Rsrc = CartesianIndices(t)[indices′...] - Rdst = CartesianIndices(tdst) - - for (I, v) in nonzero_pairs(t) - j = findfirst(==(I), Rsrc) - isnothing(j) && continue - tdst[Rdst[j]] = v - end - return tdst +# a single index is only supported when it selects a single nontrivial dimension +@noinline function _slice_getindex_single( + t::AbstractBlockTensorMap, indices::Vararg{Any, M} + ) where {M} + space(eachspace(t)[indices...]) # errors as before if unsupported + d = something(findfirst(>(1), size(t)), 1) + return _slice_getindex(t, ntuple(i -> i == d ? only(indices) : 1, numind(t))...) end # TODO: check if this fallback is fair @@ -151,9 +130,21 @@ function Base.setindex!(::AbstractBlockTensorMap, ::AbstractTensorMap, ::FusionT end # setindex verifies structure is correct -@inline function Base.setindex!( - t::AbstractBlockTensorMap, v::AbstractTensorMap, indices::Vararg{SliceIndex} - ) +@propagate_inbounds Base.setindex!( + t::AbstractBlockTensorMap, v::AbstractTensorMap, indices::Vararg{SliceIndex} +) = _slice_setindex!(t, v, indices...) +@propagate_inbounds Base.setindex!( + t::AbstractBlockTensorMap, v::AbstractBlockTensorMap, indices::Vararg{SliceIndex} +) = _slice_setindex!(t, v, indices...) +# disambiguate: TensorKit/src/tensors/abstracttensor.jl:552 +@propagate_inbounds Base.setindex!( + t::AbstractBlockTensorMap, v::AbstractTensorMap, indices::Vararg{Strided.SliceIndex} +) = _slice_setindex!(t, v, indices...) +@propagate_inbounds Base.setindex!( + t::AbstractBlockTensorMap, v::AbstractBlockTensorMap, indices::Vararg{Strided.SliceIndex} +) = _slice_setindex!(t, v, indices...) + +@inline function _slice_setindex!(t::AbstractBlockTensorMap, v::AbstractTensorMap, indices...) @boundscheck begin checkbounds(t, indices...) checkspaces(t, v, indices...) @@ -161,39 +152,12 @@ end @inbounds parent(t)[indices...] = v return t end -# setindex with blocktensor needs to correctly slice-assign -@inline function Base.setindex!( - t::AbstractBlockTensorMap, v::AbstractBlockTensorMap, indices::Vararg{SliceIndex} - ) +# a blocktensor needs to be slice-assigned +@inline function _slice_setindex!(t::AbstractBlockTensorMap, v::AbstractBlockTensorMap, indices...) @boundscheck begin checkbounds(t, indices...) checkspaces(t, v, indices...) end - - @inbounds copyto!(view(parent(t), indices...), parent(v)) - return t -end - -# disambiguate -@inline function Base.setindex!( - t::AbstractBlockTensorMap, v::AbstractTensorMap, indices::Vararg{Strided.SliceIndex} - ) - @boundscheck begin - checkbounds(t, indices...) - checkspaces(t, v, indices...) - end - @inbounds parent(t)[indices...] = v - return t -end -# disambiguate -@inline function Base.setindex!( - t::AbstractBlockTensorMap, v::AbstractBlockTensorMap, indices::Vararg{Strided.SliceIndex}, - ) - @boundscheck begin - checkbounds(t, indices...) - checkspaces(t, v, indices...) - end - @inbounds copyto!(view(parent(t), indices...), parent(v)) return t end diff --git a/test/abstracttensor/indexing.jl b/test/abstracttensor/indexing.jl index b85d454..f61aeb1 100644 --- a/test/abstracttensor/indexing.jl +++ b/test/abstracttensor/indexing.jl @@ -2,61 +2,113 @@ using BlockTensorKit using BlockTensorKit: sprand using Test using TensorKit +using LinearAlgebra V = SumSpace(ℂ^2, ℂ^3, ℂ^2) -blockt = rand(V ⊗ V ⊗ V) +for (label, blockt, TT) in ( + ("dense", rand(V ⊗ V ⊗ V), BlockTensorMap), + ("sparse", sprand(V ⊗ V ⊗ V, 0.5), SparseBlockTensorMap), + ) + @testset "$label indexing" begin + # scalar indexing + @test @inferred(blockt[1]) isa TensorMap + @test @inferred(blockt[1, 1, 1]) isa TensorMap + @test @inferred(blockt[CartesianIndex(1, 1, 1)]) isa TensorMap -# scalar indexing -@test @inferred(blockt[1]) isa TensorMap -@test @inferred(blockt[1, 1, 1]) isa TensorMap -@test @inferred(blockt[CartesianIndex(1, 1, 1)]) isa TensorMap + # colon indexing + @test @inferred(blockt[:, :, :]) == blockt + nnz = nonzero_length(blockt) + for I in eachindex(blockt) + if I in nonzero_keys(blockt) + @test blockt[I] === blockt[I] + else + @test norm(blockt[I]) == 0 + @test nonzero_length(blockt) == nnz + end + end -# colon indexing -blockt2 = @inferred blockt[:, :, :] -@test blockt2 == blockt -for I in eachindex(blockt) - @test blockt[I] === blockt[I] + @test size(@inferred(blockt[1, :, 1])) == (1, 3, 1) + @test size(blockt[1, [1, 3], 1]) == (1, 2, 1) + blockt3 = @inferred blockt[[1], [1], 1] + @test blockt3 isa TT + @test length(blockt3) == 1 + + # repeated indices duplicate blocks, as for `AbstractArray` + b = @inferred blockt[[1, 1, 3], :, :] + @test size(b) == (3, 3, 3) + for i in 1:3, j in 1:3 + @test b[1, i, j] == blockt[1, i, j] + @test b[2, i, j] == blockt[1, i, j] + @test b[3, i, j] == blockt[3, i, j] + end + # a dropped duplicate used to surface as uninitialized memory + @test norm(b[2, 1, 1]) == norm(blockt[1, 1, 1]) + @test ndims(blockt[:, 1:2, [1, 1]]) == 3 + + # logical indexing + @test @inferred(blockt[[true, false, true], :, 1]) == blockt[[1, 3], :, 1] + @test blockt[BitVector([true, false, true]), :, :] == blockt[[1, 3], :, :] + @test size(blockt[falses(3), :, :]) == (0, 3, 3) + @test_throws BoundsError blockt[[true, false], :, :] + @test_throws BoundsError blockt[[true, false, true, false], :, :] + + # reversed and stepped ranges + @test blockt[3:-1:1, :, :] == blockt[[3, 2, 1], :, :] + @test blockt[1:2:3, :, :] == blockt[[1, 3], :, :] + + # empty slices + @test size(blockt[Int[], :, :]) == (0, 3, 3) + @test length(blockt[Int[], :, :]) == 0 + + # integer types other than `Int` + @test blockt[Int8(1), 1:2, 1] == blockt[1, 1:2, 1] + @test blockt[UInt(1), 1:2, 1] == blockt[1, 1:2, 1] + @test blockt[Int8(1), Int8(1), Int8(1)] == blockt[1, 1, 1] + + # invalid indexing + @test_throws ArgumentError blockt[:] + @test_throws ArgumentError blockt[[1]] + @test_throws BoundsError blockt[:, :] + @test_throws BoundsError blockt[4, :, :] + @test_throws BoundsError @inbounds blockt[4, :, :] + + # slice assignment (index 1 and 3 share their space) + t2 = copy(blockt) + t2[[1], :, :] = blockt[[3], :, :] + for i in 1:3, j in 1:3 + @test t2[1, i, j] == blockt[3, i, j] + end + end end -@test size(@inferred(blockt[1, :, 1])) == (1, 3, 1) -blockt2 = blockt[1, [1, 3], 1] -@test size(blockt2) == (1, 2, 1) -blockt3 = @inferred blockt[[1], [1], 1] -@test blockt3 isa BlockTensorMap -@test length(blockt3) == 1 - -# invalid indexing -@test_throws ArgumentError blockt[:] -@test_throws ArgumentError blockt[[1]] - -blockt = sprand(V ⊗ V ⊗ V, 0.5) - -# scalar indexing -@test @inferred(blockt[1]) isa TensorMap -@test @inferred(blockt[1, 1, 1]) isa TensorMap -@test @inferred(blockt[CartesianIndex(1, 1, 1)]) isa TensorMap - -# colon indexing -blockt2 = @inferred blockt[:, :, :] -@test blockt2 == blockt -nnz = nonzero_length(blockt) -for I in eachindex(blockt) - if I in nonzero_keys(blockt) - @test blockt[I] === blockt[I] - else - @test norm(blockt[I]) == 0 - @test nonzero_length(blockt) == nnz +# single-index slicing of one-dimensional block tensors +for (label, t1) in (("dense", rand(V ← one(V))), ("sparse", sprand(V ← one(V), 0.8))) + @testset "$label 1-dimensional indexing" begin + @test t1[:] == t1 + @test norm(t1[1:2])^2 ≈ norm(t1[1])^2 + norm(t1[2])^2 + @test size(t1[[1, 1]]) == (2,) + @test t1[[1, 1]][2] == t1[1] end end -@test size(@inferred(blockt[1, :, 1])) == (1, 3, 1) -blockt2 = blockt[1, [1, 3], 1] -@test size(blockt2) == (1, 2, 1) -blockt3 = @inferred blockt[[1], [1], 1] -@test blockt3 isa SparseBlockTensorMap -@test length(blockt3) == 1 +# the parent array has its own slicing implementation +@testset "parent array slicing" begin + st = sprand(V ⊗ V ⊗ V, 0.5) + @test parent(st)[[true, false, true], :, :] == parent(st[[1, 3], :, :]) + @test parent(st)[[1, 1, 3], :, :] == parent(st[[1, 1, 3], :, :]) +end -# invalid indexing -@test_throws ArgumentError blockt[[1]] -@test_throws ArgumentError blockt[:] +# guard against slicing work proportional to the destination region again +@testset "slicing scales with nnz" begin + D = 512 + tb = spzeros(Float64, SumSpace(fill(ℂ^1, D)...) ⊗ SumSpace(ℂ^2) ← SumSpace(ℂ^2) ⊗ SumSpace(fill(ℂ^1, D)...)) + for i in 1:D + tb[i, 1, 1, i] = rand(eachspace(tb)[i, 1, 1, i]) + i < D && (tb[i, 1, 1, i + 1] = rand(eachspace(tb)[i, 1, 1, i + 1])) + end + inds = (2:(D - 1), :, :, 2:(D - 1)) + s = tb[inds...] + @test nonzero_length(s) == 2 * (D - 2) - 1 + @test @allocated(tb[inds...]) < 2_000_000 +end From 8e1dfc17ffa450a1857f9a81b0e3f9073333180c Mon Sep 17 00:00:00 2001 From: lkdvos Date: Tue, 25 Aug 2026 13:25:17 -0400 Subject: [PATCH 3/3] perf: make sparse slice assignment scale with nnz MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The three `SparseTensorArray` `copyto!` methods all walked the assigned region rather than the stored blocks, so `t[inds...] = v` and `cat` cost `O(prod(length.(inds)))` regardless of how little was stored — the mirror image of the `getindex` problem in #67. Reuse the inverted index maps instead: shape region nnz before after speedup 2d n=64 3_844 64 0.150 ms 0.015 ms 10x 2d n=256 64_516 256 2.267 ms 0.036 ms 63x 2d n=512 260_100 512 13.752 ms 0.051 ms 269x 4d n=24 234_256 24 9.459 ms 0.009 ms 1040x 4d n=48 4_477_456 48 140.997 ms 0.013 ms 10600x narrow 4 512 0.005 ms 0.004 ms 1x `copyto!(dst, view(parent(t), inds...))` improves likewise, 14.6 ms to 0.041 ms for a 260_100-block region. The narrow row matters: a small assignment into a large sparse array must not become proportional to the destination's stored blocks. `_deletemissing!` therefore sweeps whichever of the two is smaller, making it `O(nnz(src) + min(region, nnz(dst)))` — never worse than before. `copyto!(t::SparseTensorArray, ::SubArray)` needs no sweep at all: the destination spans exactly the viewed region, so everything not copied is dropped. Also fix the `nonzero_*` accessors for `SparseTensorArray`, which fell through to the `AbstractArray` fallbacks and reported every index rather than the stored ones. `nonzero_keys` and `nonzero_length` were always wrong this way; `nonzero_pairs` became wrong when it started going through `pairs(IndexCartesian(), A)`, which bypasses `Base.pairs(::SparseTensorArray)`. The visible effects were that slicing the parent array densified its result and that `copyto!` over regions stored explicitly-zero blocks. Verified semantics-preserving: the stored-key sets after slice assignment, `copyto!` from a view, region-to-region `copyto!` and `cat` are identical to the previous implementation once the accessors are fixed. Co-Authored-By: Claude Opus 5 (1M context) --- src/auxiliary/sparsetensorarray.jl | 69 +++++++++++-------- test/abstracttensor/indexing.jl | 106 ++++++++++++++++++++++++++++- 2 files changed, 143 insertions(+), 32 deletions(-) diff --git a/src/auxiliary/sparsetensorarray.jl b/src/auxiliary/sparsetensorarray.jl index 347e640..7831407 100644 --- a/src/auxiliary/sparsetensorarray.jl +++ b/src/auxiliary/sparsetensorarray.jl @@ -34,6 +34,11 @@ Base.pairs(A::SparseTensorArray) = pairs(A.data) Base.keys(A::SparseTensorArray) = keys(A.data) Base.values(A::SparseTensorArray) = values(A.data) +nonzero_keys(A::SparseTensorArray) = keys(A.data) +nonzero_values(A::SparseTensorArray) = values(A.data) +nonzero_pairs(A::SparseTensorArray) = pairs(A.data) +nonzero_length(A::SparseTensorArray) = length(A.data) + TensorKit.space(A::SparseTensorArray) = A.space TensorKit.codomain(A::SparseTensorArray) = codomain(space(A)) TensorKit.domain(A::SparseTensorArray) = domain(space(A)) @@ -100,38 +105,42 @@ function Base.similar( return SparseTensorArray{S, N₁, N₂, T, N}(Dict{CartesianIndex{N}, T}(), spaces) end -Base.@propagate_inbounds function Base.copyto!( - t::SparseTensorArray, v::SubArray{T, N, A} - ) where {T, N, A <: SparseTensorArray} - undropped_parentindices = map(Base.parentindices(v)) do I - I isa Base.ScalarIndex ? (I:I) : I - end +_undropped(inds::Tuple) = map(I -> I isa Base.ScalarIndex ? (I:I) : I, inds) - for I in eachindex(IndexCartesian(), t) - parentI = CartesianIndex(Base.reindex(undropped_parentindices, I.I)) - if haskey(parent(v), parentI) - t[I] = parent(v)[parentI] - else - delete!(t, I) +# clear the entries of `A` selected by `inds` that `v` does not store +function _deletemissing!(A::SparseTensorArray, inds::Tuple, v) + # sweep whichever of the selected region and the stored entries is smaller + if length(v) ≤ nonzero_length(A) + for I in eachindex(IndexCartesian(), v) + haskey(v, I) || delete!(A, CartesianIndex(Base.reindex(inds, I.I))) + end + else + maps = map(_invert_index, size(A), inds) + for J in collect(nonzero_keys(A)) + rs = map(_dstrange, maps, J.I) + any(isempty, rs) && continue + any(P -> haskey(v, CartesianIndex(P)), Iterators.product(rs...)) && continue + delete!(A, J) end end - return t + return A end +# the destination spans exactly the viewed region, so everything not copied is dropped Base.@propagate_inbounds function Base.copyto!( - t::SubArray{T, N, A}, v::SparseTensorArray + t::SparseTensorArray, v::SubArray{T, N, A} ) where {T, N, A <: SparseTensorArray} - undropped_parentindices = map(Base.parentindices(t)) do I - I isa Base.ScalarIndex ? (I:I) : I - end + empty!(t) + return _copyslice!(t, parent(v), _undropped(Base.parentindices(v))) +end - for I in eachindex(IndexCartesian(), v) - if haskey(v, I) - t[I] = v[I] - else - parentI = CartesianIndex(Base.reindex(undropped_parentindices, I.I)) - delete!(parent(t), parentI) - end +Base.@propagate_inbounds function Base.copyto!( + t::SubArray{T, N, A}, v::SparseTensorArray + ) where {T, N, A <: SparseTensorArray} + inds = _undropped(Base.parentindices(t)) + _deletemissing!(parent(t), inds, v) + for (I, x) in nonzero_pairs(v) + parent(t)[Base.reindex(inds, I.I)...] = x end return t end @@ -154,12 +163,12 @@ Base.@propagate_inbounds function Base.copyto!( checkbounds(src, first(Rsrc)) checkbounds(src, last(Rsrc)) end - CRdest = CartesianIndices(Rdest) - CRsrc = CartesianIndices(Rsrc) - ΔI = first(CRdest) - first(CRsrc) - for I in CRsrc - if Rsrc[I] in nonzero_keys(src) - dest[Rdest[I + ΔI]] = src[Rsrc[I]] + maps = map(_invert_index, size(src), Rsrc.indices) + for (I, x) in nonzero_pairs(src) + rs = map(_dstrange, maps, I.I) + any(isempty, rs) && continue + for P in Iterators.product(rs...) + dest[Rdest[P...]] = x end end return dest diff --git a/test/abstracttensor/indexing.jl b/test/abstracttensor/indexing.jl index f61aeb1..ef23edf 100644 --- a/test/abstracttensor/indexing.jl +++ b/test/abstracttensor/indexing.jl @@ -95,8 +95,92 @@ end # the parent array has its own slicing implementation @testset "parent array slicing" begin st = sprand(V ⊗ V ⊗ V, 0.5) - @test parent(st)[[true, false, true], :, :] == parent(st[[1, 3], :, :]) - @test parent(st)[[1, 1, 3], :, :] == parent(st[[1, 1, 3], :, :]) + for inds in ([true, false, true], [1, 3], [1, 1, 3], 1:2) + a, b = parent(st)[inds, :, :], parent(st[inds, :, :]) + @test a == b + @test length(a.data) == length(b.data) # must not densify + end +end + +# `nonzero_*` on the parent array must report stored entries, not every entry +@testset "sparse parent accessors" begin + st = sprand(V ⊗ V ⊗ V, 0.5) + A = parent(st) + stored = length(A.data) + @test 0 < stored < length(A) + @test nonzero_length(A) == stored + @test length(collect(nonzero_keys(A))) == stored + @test length(collect(nonzero_pairs(A))) == stored + @test length(collect(nonzero_values(A))) == stored + @test length(parent(A[1:3, :, :]).data) == stored +end + +@testset "sparse slice assignment" begin + Vh = SumSpace(ℂ^2, ℂ^2, ℂ^2) # homogeneous, so any index can be assigned to any other + for _ in 1:10 + t = sprand(Vh ⊗ Vh ⊗ Vh, 0.4) + src = sprand(Vh ⊗ Vh ⊗ Vh, 0.4)[1:2, :, :] + expected = Set(I for I in nonzero_keys(t) if I[1] > 2) + union!(expected, nonzero_keys(src)) + blocks = Dict(I => src[I] for I in nonzero_keys(src)) + for I in nonzero_keys(t) + I[1] > 2 && (blocks[I] = t[I]) + end + t[1:2, :, :] = src + @test Set(nonzero_keys(t)) == expected + for (I, x) in blocks + @test t[I] === x + end + end + + # structural zeros of the source clear the destination + t = sprand(Vh ⊗ Vh ⊗ Vh, 1.0) + @test nonzero_length(t) == 27 + t[1:2, :, :] = spzeros(Float64, Vh ⊗ Vh ⊗ Vh)[1:2, :, :] + @test nonzero_length(t) == 9 + @test all(I -> I[1] == 3, nonzero_keys(t)) +end + +@testset "copyto! on the parent array" begin + Vh = SumSpace(ℂ^2, ℂ^2, ℂ^2) # homogeneous, so blocks can move between indices + st = sprand(Vh ⊗ Vh ⊗ Vh, 0.5) + A = parent(st) + + # into a fresh array from a view: the destination spans the view exactly + dst = parent(st[1:2, :, :]) + dst[1, 1, 1] = rand(eachspace(st)[1, 1, 1]) + copyto!(dst, view(A, 1:2, :, :)) + @test Set(nonzero_keys(dst)) == Set(I for I in nonzero_keys(A) if I[1] ≤ 2) + for I in nonzero_keys(dst) + @test dst[I] === A[I] + end + + # region-to-region, including a stepped region + for st_ in (1, 2) + a, b = parent(sprand(Vh ⊗ Vh ⊗ Vh, 0.5)), parent(sprand(Vh ⊗ Vh ⊗ Vh, 0.5)) + Rsrc = CartesianIndices((1:st_:3, 1:1, 1:1)) + Rdest = CartesianIndices((1:st_:3, 2:2, 3:3)) + before = Dict(I => a[I] for I in nonzero_keys(a)) + copyto!(a, Rdest, b, Rsrc) + for (Pd, Ps) in zip(Rdest, Rsrc) + if Ps in nonzero_keys(b) + @test a[Pd] === b[Ps] + else + @test (Pd in keys(before)) == (Pd in nonzero_keys(a)) + end + end + end +end + +@testset "cat" begin + st = sprand(V ⊗ V ⊗ V, 0.5) + c = cat(st, st; dims = 1) + @test size(c) == (6, 3, 3) + @test nonzero_length(c) == 2 * nonzero_length(st) + for I in nonzero_keys(st) + @test c[I] === st[I] + @test c[I + CartesianIndex(3, 0, 0)] === st[I] + end end # guard against slicing work proportional to the destination region again @@ -112,3 +196,21 @@ end @test nonzero_length(s) == 2 * (D - 2) - 1 @test @allocated(tb[inds...]) < 2_000_000 end + +# guard against assignment work proportional to the assigned region +@testset "assignment scales with nnz" begin + D = 48 + Vb = SumSpace(fill(ℂ^1, D)...) + tb = spzeros(Float64, Vb ⊗ Vb ← Vb ⊗ Vb) + for i in 1:D + tb[i, i, i, i] = rand(eachspace(tb)[i, i, i, i]) + end + inds = ntuple(_ -> 2:(D - 1), 4) + src = tb[inds...] + tb[inds...] = src # warm up + nnz = nonzero_length(tb) + # the assigned region holds (D - 2)^4 ~ 4.5e6 blocks but only D are stored, so anything + # proportional to the region takes >100 ms here while this takes microseconds + @test (@elapsed tb[inds...] = src) < 0.02 + @test nonzero_length(tb) == nnz +end