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..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 @@ -168,33 +177,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/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 diff --git a/test/abstracttensor/indexing.jl b/test/abstracttensor/indexing.jl index b85d454..ef23edf 100644 --- a/test/abstracttensor/indexing.jl +++ b/test/abstracttensor/indexing.jl @@ -2,61 +2,215 @@ using BlockTensorKit using BlockTensorKit: sprand using Test using TensorKit +using LinearAlgebra V = SumSpace(ℂ^2, ℂ^3, ℂ^2) -blockt = rand(V ⊗ V ⊗ V) - -# 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 -for I in eachindex(blockt) - @test blockt[I] === blockt[I] -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 - 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 - -# invalid indexing -@test_throws ArgumentError blockt[[1]] -@test_throws ArgumentError blockt[:] +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 + + # 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 + + @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 + +# 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 + +# the parent array has its own slicing implementation +@testset "parent array slicing" begin + st = sprand(V ⊗ V ⊗ V, 0.5) + 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 +@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 + +# 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