Slicing `getindex` is quadratic: replace the per-block `findfirst` scan with inverse index maps

[lkdvos]

Summary

Slicing getindex on an AbstractBlockTensorMap searches the source index grid once per stored block, making it O(nnz × prod(length.(indices))).
Inverting the index maps instead makes it O(nnz + Σ size(t, d)), which is a 2600× speedup on a 256-channel sparse block tensor and removes the quadratic scaling entirely.

Where

src/tensors/abstractblocktensor/abstractarray.jl, in both the SliceIndex method (L95-115) and its Strided.SliceIndex disambiguation twin (L118-137), which have identical bodies:

Rsrc = CartesianIndices(t)[indices′...]
Rdst = CartesianIndices(tdst)

for (I, v) in nonzero_pairs(t)
    j = findfirst(==(I), Rsrc)     # <-- linear scan of the whole slice region
    isnothing(j) && continue
    tdst[Rdst[j]] = v
end

Rsrc has prod(length.(indices′)) entries and is rescanned from scratch for every stored block.
For a dense BlockTensorMap it is worse again, since nonzero_pairs then covers every block.

Measurements

BlockTensorKit v0.3.15, Julia 1.12.6.
A banded sparse block tensor (nnz = n, one block per row) sliced to every other row and column, i.e. t[1:2:n, 1:1, 1:1, 1:2:n]:

n	nnz	slice	current	inverse maps	speedup
16	16	8×8	0.0031 ms	0.0004 ms	7.5×
32	32	16×16	0.0188 ms	0.0006 ms	34×
64	64	32×32	0.1415 ms	0.0009 ms	160×
128	128	64×64	1.1317 ms	0.0014 ms	802×
256	256	128×128	8.9381 ms	0.0034 ms	2595×

Dense BlockTensorMap, same slice:

n	blocks	current	inverse maps	speedup
16	256	0.0385 ms	0.0035 ms	11×
32	1024	0.5206 ms	0.0118 ms	44×
64	4096	7.8955 ms	0.0452 ms	175×

The current column grows ~8× per doubling of n (nnz × region = 2× × 4×); the inverse-map column grows linearly.

Proposed fix

Build, per dimension, the inverse of the index map once, then look up each stored block's destination:

function slice_invmap(t, indices::Vararg{Any, N}) where {N}
    Vsp = space(eachspace(t)[indices...])
    tdst = similar(t, Vsp)
    length(tdst) == 0 && return tdst

    indices′ = map(ind -> ind isa Int ? (ind:ind) : ind, indices)
    invmaps = ntuple(N) do d
        m = zeros(Int, size(t, d))
        for (j, i) in enumerate(indices′[d])
            m[i] == 0 && (m[i] = j)   # first occurrence, matching `findfirst`
        end
        return m
    end

    for (I, v) in nonzero_pairs(t)
        J = ntuple(d -> invmaps[d][I[d]], N)
        any(iszero, J) && continue
        tdst[CartesianIndex(J)] = v
    end
    return tdst
end

This is the prototype benchmarked above.
It is deliberately a drop-in: same output spaces, same stored blocks, same values (asserted in the MWE below).
Loose ends a real implementation would need to cover, none of which affect the asymptotics:

• Logical indexing. Bool vectors are part of the SliceIndex union, so they need normalizing (to_indices / findall) before the inverse map is built.
• Linear / single-index slicing. With one index argument, CartesianIndices(t)[ind] yields a vector rather than an N-d grid, so that path wants the LinearIndices analogue.
• Deduplication. The two method bodies are byte-identical; the fix is a good moment to have the disambiguation method forward to a shared implementation.

Note: repeated indices currently drop blocks

Not the point of this issue, but it surfaced while checking the prototype against current behaviour, and it constrains what "drop-in" means.

findfirst returns only the first match, so an index repeated in the slice silently loses blocks:

t = banded(4)                       # nonzeros at (1,1,1,2), (2,1,1,3), (3,1,1,4), (4,1,1,1)
s = t[[1, 1, 2], 1:1, 1:1, 1:4]

Source rows 1 and 2 carry 2 stored blocks, and row 1 is selected twice, so the result should have 3 stored blocks: s[1,1,1,2] and s[2,1,1,2] both from t[1,1,1,2], plus s[3,1,1,3] from t[2,1,1,3].
It has 2 — the second copy of row 1 comes out empty.

The prototype above reproduces this exactly (hence the m[i] == 0 && guard) so that the performance change is behaviour-preserving.
If the repeated-index semantics should instead follow Base, that is a separate fix and the inverse map would become Vector{Vector{Int}} (or the loop would push to all matching destinations).

MWE

using BlockTensorKit
using BlockTensorKit: nonzero_pairs, nonzero_length, eachspace
using TensorKit
using BenchmarkTools

const V = ℂ^2
S(n) = SumSpace(fill(V, n)...)
const TT = TensorMap{Float64, ComplexSpace, 2, 2, Vector{Float64}}

function banded(n)
    t = SparseBlockTensorMap{TT}(undef, S(n) ⊗ V ← V ⊗ S(n))
    for i in 1:n
        t[i, 1, 1, mod1(i + 1, n)] = randn(V ⊗ V ← V ⊗ V)
    end
    return t
end

# ... slice_invmap as above ...

for n in (16, 32, 64, 128, 256)
    t = banded(n)
    rows, cols = collect(1:2:n), collect(1:2:n)

    a = t[rows, 1:1, 1:1, cols]
    b = slice_invmap(t, rows, 1:1, 1:1, cols)
    @assert space(a) == space(b)
    @assert nonzero_length(a) == nonzero_length(b)
    for (I, v) in nonzero_pairs(a)
        @assert b[I] ≈ v
    end

    t_cur = minimum(@benchmark($t[$rows, 1:1, 1:1, $cols])).time / 1e6
    t_new = minimum(@benchmark(slice_invmap($t, $rows, 1:1, 1:1, $cols))).time / 1e6
    println(n, "\t", round(t_cur; digits = 4), " ms\t", round(t_new; digits = 4), " ms\t",
        round(t_cur / t_new; digits = 1), "x")
end

Where this bites

MPSKit's two-site Jordan-MPO derivative operator restricts the continuing–continuing block to the MPO channels that can actually carry a contribution across both bonds, which means four of these slices per bond.
For a chain with two-body terms out to range 16 (a 115×116 A-block with 104 stored blocks) the four slices cost 1.9 ms and 12.9 MB, about 16 % of the whole per-bond operator construction — enough that it is not obvious the restriction pays for itself.
With the asymptotics fixed it would be free.