Reduce overhead: cached block access, SumSpace sector hashing, sparse missing entries

src/auxiliary/blockarrays.jl

@@ -16,6 +16,17 @@ function copy_dense!(Adense, A)
     return Adense
 end
 
+function copy_blocks!(A::BlockMatrix, Adense)
+    for block_index in Iterators.product(blockaxes(A)...)
+        a = view(A, block_index...)
+        isempty(a) && continue
+        indices = getindex.(axes(A), block_index)
+        a .= view(Adense, indices...)
+    end
+    return A
+end
+copy_blocks!(A, Adense) = copyto!(A, Adense)
+
 const BlockBlasMat{T <: MAK.BlasFloat} = BlockMatrix{T}
 
 function MAK.zero!(A::BlockBlasMat)

src/auxiliary/sparsetensorarray.jl

@@ -44,9 +44,9 @@ TensorKit.codomain(A::SparseTensorArray) = codomain(space(A))
 TensorKit.domain(A::SparseTensorArray) = domain(space(A))
 
 TensorKit.numout(A::SparseTensorArray) = numout(eltype(A))
-TensorKit.numout(::Type{T}) where {T <: SparseTensorArray} = numout(eltype(A))
+TensorKit.numout(::Type{T}) where {T <: SparseTensorArray} = numout(eltype(T))
 TensorKit.numin(A::SparseTensorArray) = numin(eltype(A))
-TensorKit.numin(::Type{T}) where {T <: SparseTensorArray} = numin(eltype(A))
+TensorKit.numin(::Type{T}) where {T <: SparseTensorArray} = numin(eltype(T))
 
 # AbstractArray interface
 # -----------------------
@@ -93,6 +93,7 @@ function Base.delete!(A::SparseTensorArray, I::Vararg{Int, N}) where {N}
 end
 Base.delete!(A::SparseTensorArray, I::CartesianIndex) = delete!(A.data, I)
 Base.empty!(A::SparseTensorArray) = empty!(A.data)
+Base.filter!(f, A::SparseTensorArray) = (filter!(f, A.data); A)
 function Base.haskey(A::SparseTensorArray, I::Vararg{Int, N}) where {N}
     return haskey(A.data, CartesianIndex(I))
 end

src/linalg/factorizations.jl

@@ -28,7 +28,7 @@ for f! in (
             Fblocks′ = MAK.$f!(copy_dense!(similar_dense(tblock), tblock), alg)
             # deal with the case where the output is not in-place
             for (b′, b) in zip(Fblocks′, Fblocks)
-                b === b′ || copyto!(b, b′)
+                b === b′ || copy_blocks!(b, b′)
             end
             return nothing
         end
@@ -46,7 +46,7 @@ for f! in (
         TensorKit.foreachblock(t, N) do _, (tblock, Nblock)
             Nblock′ = MAK.$f!(copy_dense!(similar_dense(tblock), tblock), alg)
             # deal with the case where the output is not the same as the input
-            Nblock === Nblock′ || copyto!(Nblock, Nblock′)
+            Nblock === Nblock′ || copy_blocks!(Nblock, Nblock′)
             return nothing
         end
         return N
@@ -186,25 +186,25 @@ end
 function TensorKit.Factorizations.truncate_domain!(tdst::AbstractBlockTensorMap, tsrc::AbstractBlockTensorMap, inds)
     TensorKit.foreachblock(tdst, tsrc) do c, (dst_block, src_block)
         I = get(inds, c, nothing)
-        dst_dense = copy_dense!(similar_dense(dst_block), dst_block)
+        dst_dense = similar_dense(dst_block)
         src_dense = copy_dense!(similar_dense(src_block), src_block)
         @assert !isnothing(I)
         @views dst_dense .= src_dense[:, I]
         # deal with the case where the output is not in-place
-        dst_dense === dst_block || copyto!(dst_block, dst_dense)
+        dst_dense === dst_block || copy_blocks!(dst_block, dst_dense)
         return nothing
     end
     return tdst
 end
 function TensorKit.Factorizations.truncate_codomain!(tdst::AbstractBlockTensorMap, tsrc::AbstractBlockTensorMap, inds)
     TensorKit.foreachblock(tdst, tsrc) do c, (dst_block, src_block)
         I = get(inds, c, nothing)
-        dst_dense = copy_dense!(similar_dense(dst_block), dst_block)
+        dst_dense = similar_dense(dst_block)
         src_dense = copy_dense!(similar_dense(src_block), src_block)
         @assert !isnothing(I)
         @views dst_dense .= src_dense[I, :]
         # deal with the case where the output is not in-place
-        dst_dense === dst_block || copyto!(dst_block, dst_dense)
+        dst_dense === dst_block || copy_blocks!(dst_block, dst_dense)
         return nothing
     end
     return tdst

src/linalg/linalg.jl

@@ -157,10 +157,34 @@ for (TC, TA, TB) in Iterators.product(
     end
 end
 
-function LinearAlgebra.norm(tA::BlockTensorMap, p::Real = 2)
+function LinearAlgebra.norm(tA::AbstractBlockTensorMap, p::Real = 2)
     vals = nonzero_values(tA)
     isempty(vals) && return norm(zero(scalartype(tA)), p)
-    return LinearAlgebra.norm(norm.(vals), p)
+    return LinearAlgebra.norm(Iterators.map(Base.Fix2(norm, p), vals), p)
+end
+
+function Base.:(==)(t1::AbstractBlockTensorMap, t2::AbstractBlockTensorMap)
+    space(t1) == space(t2) || return false
+    for I in nonzero_keys(t1)
+        v1 = t1[I]
+        (haskey(t2, I) ? v1 == t2[I] : iszero(norm(v1))) || return false
+    end
+    if issparse(t1)
+        for I in nonzero_keys(t2)
+            haskey(t1, I) || iszero(norm(t2[I])) || return false
+        end
+    end
+    return true
+end
+
+function LinearAlgebra.tr(t::AbstractBlockTensorMap)
+    domain(t) == codomain(t) ||
+        throw(SpaceMismatch("Trace of a tensor only exist when domain == codomain"))
+    s = zero(scalartype(t)) * zero(TK.dimscalartype(sectortype(t)))
+    for I in nonzero_keys(t)
+        TT.getindices(I.I, codomainind(t)) == TT.getindices(I.I, domainind(t)) && (s += tr(t[I]))
+    end
+    return s
 end
 
 for f in (:real, :imag)
@@ -173,7 +197,7 @@ for f in (:real, :imag)
             end
             return t′
         else
-            msg = "`$f` has not been implemented for `BlockTensorMap{$(S)}`."
+            msg = "`$f` has not been implemented for `BlockTensorMap{$(spacetype(t))}`."
             throw(ArgumentError(msg))
         end
     end
@@ -208,7 +232,7 @@ function LinearAlgebra.lmul!(D::DiagonalTensorMap, t::AbstractBlockTensorMap)
     domain(D) == codomain(t) || throw(SpaceMismatch())
     TensorKit.foreachblock(t, D) do c, (tblock, Dblock)
         tblock′ = lmul!(Dblock, copy_dense!(similar_dense(tblock), tblock))
-        tblock === tblock′ || copyto!(tblock, tblock′)
+        tblock === tblock′ || copy_blocks!(tblock, tblock′)
         return tblock
     end
     return t
@@ -218,7 +242,7 @@ function LinearAlgebra.rmul!(t::AbstractBlockTensorMap, D::DiagonalTensorMap)
     codomain(D) == domain(t) || throw(SpaceMismatch())
     TensorKit.foreachblock(t, D) do c, (tblock, Dblock)
         tblock′ = rmul!(copy_dense!(similar_dense(tblock), tblock), Dblock)
-        tblock === tblock′ || copyto!(tblock, tblock′)
+        tblock === tblock′ || copy_blocks!(tblock, tblock′)
         return tblock
     end
     return t

src/tensors/abstractblocktensor/abstracttensormap.jl

@@ -10,34 +10,14 @@ end
 
 eachspace(t::AbstractBlockTensorMap) = SumSpaceIndices(space(t))
 
-@inline function TensorKit.subblock(
-        t::AbstractBlockTensorMap, (f₁, f₂)::Tuple{FusionTree, FusionTree}
-    )
-    sectortype(t) === sectortype(f₁) === sectortype(f₂) ||
-        throw(SectorMismatch("Not a valid sectortype for this tensor"))
-    numout(t) == length(f₁) && numin(t) == length(f₂) ||
-        throw(DimensionMismatch("Invalid number of fusiontree legs for this tensor"))
-
-    subblocks = map(eachspace(t), parent(t)) do V, x
-        sz = (dims(codomain(V), f₁.uncoupled)..., dims(domain(V), f₂.uncoupled)...)
-        if prod(sz) == 0
-            data = storagetype(t)(undef, 0)
-            return sreshape(StridedView(data), sz)
-        else
-            return x[f₁, f₂]
-        end
-    end
-
-    return mortar(subblocks)
-end
 @inline function Base.setindex!(
         t::AbstractBlockTensorMap, v::AbstractBlockArray, f₁::FusionTree, f₂::FusionTree
     )
     for I in eachindex(t)
-        b = v[Block(I.I)]
-        if !isempty(b)
-            getindex!(t, I)[f₁, f₂] = v[Block(I.I)]
-        end
+        b = view(v, Block(I.I))
+        isempty(b) && continue
+        x = issparse(t) ? getindex!(t, I) : t[I]
+        x[f₁, f₂] = b
     end
     return t
 end
@@ -53,57 +33,122 @@ end
     return setindex!(t, v′, f₁, f₂)
 end
 
-function TensorKit.block(t::AbstractBlockTensorMap, c::Sector)
-    sectortype(t) == typeof(c) || throw(SectorMismatch())
+# Blocks and subblocks
+# --------------------
+_entrycache(f, t::AbstractBlockTensorMap) =
+    issparse(t) ? Dict(I => f(x) for (I, x) in nonzero_pairs(t)) : map(f, parent(t))
 
-    rows = prod(TT.getindices(size(t), codomainind(t)))
-    cols = prod(TT.getindices(size(t), domainind(t)))
+_blockcache(x) = x
+_blockcache(x::AbstractTensorMap) = (b = TK.blocks(x); b isa TK.BlockIterator ? b : x)
+_cachedblock(x, c::Sector) = block(x, c)
+_cachedblock(b::TK.BlockIterator, c::Sector) = b[c]
 
-    if rows == 0 || cols == 0
-        allblocks = Matrix{TK.blocktype(eltype(t))}(undef, rows, cols)
+_subblockcache(x) = x
+_subblockcache(x::AbstractTensorMap) = sectortype(x) === Trivial ? x : TK.subblocks(x)
+_cachedsubblock(x, f) = subblock(x, f)
+_cachedsubblock(s::TK.SubblockIterator, f) = s[f]
 
-        rowaxes = Int[]
-        if rows != 0
-            W′ = codomain(t) ← zero(spacetype(t))
-            for V in eachspace(W′)
-                push!(rowaxes, blockdim(codomain(V), c))
-            end
-        end
+TensorKit.blocks(t::AbstractBlockTensorMap) =
+    TK.BlockIterator(t, (blocksectors(t), _entrycache(_blockcache, t)))
+TensorKit.block(t::AbstractBlockTensorMap, c::Sector) = _block(t, parent(t), c)
+
+TensorKit.subblocks(t::AbstractBlockTensorMap) =
+    TK.SubblockIterator(t, (TK.fusiontrees(t), _entrycache(_subblockcache, t)))
+TensorKit.subblock(t::AbstractBlockTensorMap, f::Tuple{FusionTree, FusionTree}) = _subblock(t, parent(t), f)
+
+Base.length(iter::TK.BlockIterator{<:AbstractBlockTensorMap}) = length(first(iter.structure))
+Base.isdone(iter::TK.BlockIterator{<:AbstractBlockTensorMap}, state...) =
+    Base.isdone(first(iter.structure), state...)
+function Base.iterate(iter::TK.BlockIterator{<:AbstractBlockTensorMap}, state...)
+    next = iterate(first(iter.structure), state...)
+    isnothing(next) && return next
+    c, newstate = next
+    return c => iter[c], newstate
+end
+Base.getindex(iter::TK.BlockIterator{<:AbstractBlockTensorMap}, c::Sector) =
+    _block(iter.t, last(iter.structure), c)
+
+Base.length(iter::TK.SubblockIterator{<:AbstractBlockTensorMap}) = length(first(iter.structure))
+Base.isdone(iter::TK.SubblockIterator{<:AbstractBlockTensorMap}, i::Int = 1) = i > length(iter)
+function Base.iterate(iter::TK.SubblockIterator{<:AbstractBlockTensorMap}, i::Int = 1)
+    i > length(iter) && return nothing
+    f = TK.gettokenvalue(first(iter.structure), i)
+    return f => iter[f], i + 1
+end
+Base.getindex(iter::TK.SubblockIterator{<:AbstractBlockTensorMap}, i::Int) =
+    iter[TK.gettokenvalue(first(iter.structure), i)]
+Base.getindex(iter::TK.SubblockIterator{<:AbstractBlockTensorMap}, f::Tuple{FusionTree, FusionTree}) =
+    _subblock(iter.t, last(iter.structure), f)
 
-        colaxes = Int[]
-        if cols != 0
-            W′ = zero(spacetype(t)) ← domain(t)
-            for V in eachspace(W′)
-                push!(colaxes, blockdim(domain(V), c))
-            end
+function _block(t::AbstractBlockTensorMap, entries, c::Sector)
+    sectortype(t) == typeof(c) || throw(SectorMismatch())
+    rowdims = _blockdims(codomain(t), c)
+    coldims = _blockdims(domain(t), c)
+    allblocks = Matrix{_blockeltype(typeof(t))}(undef, length(rowdims), length(coldims))
+    for (k, I) in enumerate(eachindex(IndexCartesian(), t))
+        allblocks[k] = if !issparse(t) || haskey(entries, I)
+            _cachedblock(entries[I], c)
+        else
+            i, j = Tuple(CartesianIndices(allblocks)[k])
+            _zeroblock(t, I, rowdims[i], coldims[j], c)
         end
+    end
+    return mortar(allblocks, rowdims, coldims)
+end
 
-        return mortar(allblocks, rowaxes, colaxes)
+function _subblock(t::AbstractBlockTensorMap, entries, (f₁, f₂)::Tuple{FusionTree, FusionTree})
+    sectortype(t) === sectortype(f₁) === sectortype(f₂) ||
+        throw(SectorMismatch("Not a valid sectortype for this tensor"))
+    numout(t) == length(f₁) && numin(t) == length(f₂) ||
+        throw(DimensionMismatch("Invalid number of fusiontree legs for this tensor"))
+    W = eachspace(t)
+    subblocks = map(eachindex(IndexCartesian(), t)) do I
+        V = W[I]
+        sz = (dims(codomain(V), f₁.uncoupled)..., dims(domain(V), f₂.uncoupled)...)
+        if prod(sz) == 0 || (issparse(t) && !haskey(entries, I))
+            return sreshape(StridedView(zerovector!(similar(storagetype(t), prod(sz)))), sz)
+        else
+            return _cachedsubblock(entries[I], (f₁, f₂))
+        end
     end
+    return mortar(subblocks)
+end
+
+function _blockdims(P::ProductSumSpace{S, N}, c::Sector) where {S, N}
+    return vec(
+        map(CartesianIndices(map(length, P.spaces))) do I
+            return blockdim(ProductSpace{S, N}(map(getindex, P.spaces, Tuple(I))), c)
+        end
+    )
+end
 
-    allblocks = map(Base.Fix2(block, c), parent(t))
-    return mortar(reshape(allblocks, rows, cols))
+function _zeroblock(t::AbstractBlockTensorMap, I, d₁::Int, d₂::Int, c::Sector)
+    eltype(t) <: TensorMap || return block(zerovector!(similar(eltype(t), eachspace(t)[I])), c)
+    data = zerovector!(similar(storagetype(t), d₁ * d₂))
+    return reshape(view(data, 1:(d₁ * d₂)), (d₁, d₂))
 end
 
 # TODO: this might get fixed once new tensormap is implemented
 TensorKit.blocksectors(t::AbstractBlockTensorMap) = blocksectors(space(t))
 TensorKit.hasblock(t::AbstractBlockTensorMap, c::Sector) = c in blocksectors(t)
 
-TensorKit.blocks(t::AbstractBlockTensorMap) = TK.BlockIterator(t, blocksectors(t))
-Base.@assume_effects :foldable function TensorKit.blocktype(::Type{TT}) where {TT <: AbstractBlockTensorMap}
+Base.@assume_effects :foldable function _blockeltype(::Type{TT}) where {TT <: AbstractBlockTensorMap}
     T = scalartype(TT)
     B = TK.blocktype(eltype(TT))
-    (B <: AbstractMatrix{T}) || (B = AbstractMatrix{T}) # safeguard against type-instability
+    return B <: AbstractMatrix{T} ? B : AbstractMatrix{T} # safeguard against type-instability
+end
+function TensorKit.blocktype(::Type{TT}) where {TT <: AbstractBlockTensorMap}
     BS = NTuple{2, BlockedOneTo{Int, Vector{Int}}}
-    return BlockMatrix{T, Matrix{B}, BS}
+    return BlockMatrix{scalartype(TT), Matrix{_blockeltype(TT)}, BS}
 end
 
-function Base.iterate(iter::TK.BlockIterator{<:AbstractBlockTensorMap}, state...)
-    next = iterate(iter.structure, state...)
-    isnothing(next) && return next
-    c, newstate = next
-    return c => block(iter.t, c), newstate
+function TensorKit.foreachblock(f, t::AbstractBlockTensorMap, ts...; scheduler = nothing)
+    tensors = (t, ts...)
+    iters = map(_blockcache, tensors)
+    foreach(union(blocksectors.(tensors)...)) do c
+        return f(c, map(Base.Fix2(_cachedblock, c), iters))
+    end
+    return nothing
 end
-Base.getindex(iter::TK.BlockIterator{<:AbstractBlockTensorMap}, c::Sector) = block(iter.t, c)
 
 TK.storagetype(::Type{TT}) where {TT <: AbstractBlockTensorMap} = storagetype(eltype(TT))

src/tensors/abstractblocktensor/conversion.jl

@@ -1,16 +1,38 @@
 # Conversion
 # ----------
 
-function _copy_subblocks!(tdst, tsrc)
-    N₁, N₂ = numout(tsrc), numin(tsrc)
-    offsets = ntuple(i -> _sumspace_offsets(i <= N₁ ? codomain(tsrc)[i] : domain(tsrc)[i - N₁]), N₁ + N₂)
-    # a single `SubblockIterator` holds the structure: `tdst[f₁, f₂]` would look it up in TensorKit's locked cache every time
-    dstblocks = subblocks(tdst)
-    for (k, v) in nonzero_pairs(tsrc), ((f₁, f₂), src) in subblocks(v)
-        ranges = map(offsets, (f₁.uncoupled..., f₂.uncoupled...), Tuple(k)) do o, c, kᵢ
-            return (o[c][kᵢ] + 1):o[c][kᵢ + 1]
+function _subblock_ranges(offsets, (f₁, f₂), I::CartesianIndex)
+    return map(offsets, (f₁.uncoupled..., f₂.uncoupled...), Tuple(I)) do o, c, kᵢ
+        return (o[c][kᵢ] + 1):o[c][kᵢ + 1]
+    end
+end
+
+function _subblock_pairs(t::AbstractTensorMap)
+    sectortype(t) === Trivial || return subblocks(t)
+    f = TK.trivial_fusiontree(t)
+    return (f => subblock(t, f),)
+end
+
+function _copy_subblocks!(tdst::TensorMap, tsrc::AbstractBlockTensorMap)
+    offsets = map(_sumspace_offsets, eachspace(tsrc).sumspaces)
+    dst = Base.Fix1(_cachedsubblock, _subblockcache(tdst))
+    for (I, v) in nonzero_pairs(tsrc), (f, src) in _subblock_pairs(v)
+        copy!(view(dst(f), _subblock_ranges(offsets, f, I)...), src)
+    end
+    return tdst
+end
+
+function _copy_subblocks!(tdst::AbstractBlockTensorMap, tsrc::AbstractTensorMap)
+    offsets = map(_sumspace_offsets, eachspace(tdst).sumspaces)
+    src = Base.Fix1(_cachedsubblock, _subblockcache(tsrc))
+    for (I, V) in pairs(eachspace(tdst))
+        v = similar(eltype(tdst), V)
+        vblocks = _subblock_pairs(v)
+        issparse(tdst) && all(isempty ∘ last, vblocks) && continue
+        for (f, dst) in vblocks
+            copy!(dst, view(src(f), _subblock_ranges(offsets, f, I)...))
         end
-        copy!(dstblocks[(f₁, f₂)][ranges...], src)
+        @inbounds tdst[I] = v
     end
     return tdst
 end