A few more specializations for truncation and diagonal

Author: kshyatt

src/linalg/factorizations.jl

@@ -188,3 +188,30 @@ for f! in (
     @eval MAK.$f!(::AbstractBlockTensorMap, x, ::DiagonalAlgorithm) =
         error("Blocktensors are incompatible with diagonal algorithm")
 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)
+        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)
+        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)
+        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)
+        return nothing
+    end
+    return tdst
+end

src/linalg/linalg.jl

@@ -239,3 +239,23 @@ function LinearAlgebra.isposdef!(t::AbstractBlockTensorMap)
     end
     return true
 end
+
+function LinearAlgebra.lmul!(D::DiagonalTensorMap, t::AbstractBlockTensorMap)
+    domain(D) == codomain(t) || throw(SpaceMismatch())
+    TensorKit.foreachblock(t, D) do c, (tblock, bs...)
+        tblock′ = lmul!(bs..., copy_dense!(similar_dense(tblock), tblock))
+        tblock === tblock′ || copyto!(tblock, tblock′)
+        return tblock
+    end
+    return t
+end
+
+function LinearAlgebra.rmul!(t::AbstractBlockTensorMap, D::DiagonalTensorMap)
+    codomain(D) == domain(t) || throw(SpaceMismatch())
+    TensorKit.foreachblock(t, D) do c, (tblock, bs...)
+        tblock′ = rmul!(copy_dense!(similar_dense(tblock), tblock), bs...)
+        tblock === tblock′ || copyto!(tblock, tblock′)
+        return tblock
+    end
+    return t
+end