perf: invert index maps when slicing and assigning block tensors

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

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")

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

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

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,49 +130,34 @@ 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...)
     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

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