From 86a847ff1157d186b289d826c3e8af5e54c06788 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Thu, 10 Sep 2026 12:03:43 -0400 Subject: [PATCH] fix: mixed dense/sparse `add!`, `scale!` and `inner` (#74) These were only defined for `BlockTensorMap`/`BlockTensorMap` and `SparseBlockTensorMap`/`SparseBlockTensorMap` pairs, so a mixed pair fell through to TensorKit's generic `AbstractTensorMap` methods, which operate on the fused per-sector `BlockMatrix`es. That path throws a `BoundsError` whenever a coupled sector has an interior zero-length block, which is the normal situation for a `SumSpace` with heterogeneous summands. Replace the specializations by a single `AbstractBlockTensorMap` method each, branching on `issparse` only where the cases differ. Also drop the dense-only `axpy!`, `axpby!`, `dot` and `mul!`, which reimplement TensorKit generics that forward to this layer (the deleted `mul!` referenced `SparseArrayKit._zero!`, which is not a dependency of this package). Co-Authored-By: Claude Opus 5 (1M context) --- src/linalg/linalg.jl | 42 --------------------- src/tensors/vectorinterface.jl | 69 ++++++++++------------------------ test/linalg/vectorinterface.jl | 57 +++++++++++++++++++++++----- 3 files changed, 66 insertions(+), 102 deletions(-) diff --git a/src/linalg/linalg.jl b/src/linalg/linalg.jl index 7f59324..aedc2bb 100644 --- a/src/linalg/linalg.jl +++ b/src/linalg/linalg.jl @@ -9,48 +9,6 @@ Base.:(*)(α::Number, t::AbstractBlockTensorMap) = scale(t, α) Base.:(/)(t::AbstractBlockTensorMap, α::Number) = scale(t, inv(α)) Base.:(\)(α::Number, t::AbstractBlockTensorMap) = scale(t, inv(α)) -function LinearAlgebra.axpy!(α::Number, t1::BlockTensorMap, t2::BlockTensorMap) - space(t1) == space(t2) || throw(SpaceMismatch()) - for (i, v) in nonzero_pairs(t1) - t2[i] = axpy!(α, v, t2[i]) - end - return t2 -end - -function LinearAlgebra.axpby!(α::Number, t1::BlockTensorMap, β::Number, t2::BlockTensorMap) - space(t1) == space(t2) || throw(SpaceMismatch()) - rmul!(t2, β) - for (i, v) in nonzero_pairs(t1) - t2[i] = axpy!(α, v, t2[i]) - end - return t2 -end - -function LinearAlgebra.dot(t1::BlockTensorMap, t2::BlockTensorMap) - size(t1) == size(t2) || throw(DimensionMismatch("dot arguments have different size")) - - s = zero(promote_type(scalartype(t1), scalartype(t2))) - if nonzero_length(t1) >= nonzero_length(t2) - @inbounds for (I, v) in nonzero_pairs(t1) - s += dot(v, t2[I]) - end - else - @inbounds for (I, v) in nonzero_pairs(t2) - s += dot(t1[I], v) - end - end - return s -end - -function LinearAlgebra.mul!(C::BlockTensorMap, α::Number, A::BlockTensorMap) - space(C) == space(A) || throw(SpaceMismatch()) - SparseArrayKit._zero!(parent(C)) - for (i, v) in nonzero_pairs(A) - C[i] = mul!(C[i], α, v) - end - return C -end - for TA in (:AbstractBlockTensorMap, :(TK.DiagonalTensorMap), :(TK.AdjointTensorMap), :TensorMap), TB in (:AbstractBlockTensorMap, :(TK.DiagonalTensorMap), :(TK.AdjointTensorMap), :TensorMap) (TA === :AbstractBlockTensorMap || TB === :AbstractBlockTensorMap) || continue diff --git a/src/tensors/vectorinterface.jl b/src/tensors/vectorinterface.jl index d9092dd..f7fbe66 100644 --- a/src/tensors/vectorinterface.jl +++ b/src/tensors/vectorinterface.jl @@ -10,30 +10,13 @@ function VI.scale!(t::AbstractBlockTensorMap, α::Number) return t end -function VI.scale!(ty::BlockTensorMap, tx::BlockTensorMap, α::Number) +function VI.scale!(ty::AbstractBlockTensorMap, tx::AbstractBlockTensorMap, α::Number) space(ty) == space(tx) || throw(SpaceMismatch("$(space(ty)) ≠ $(space(tx))")) - scale!(parent(ty), parent(tx), α) - return ty -end -function VI.scale!(ty::SparseBlockTensorMap, tx::SparseBlockTensorMap, α::Number) - space(ty) == space(tx) || throw(SpaceMismatch("$(space(ty)) ≠ $(space(tx))")) - y_notin_x = setdiff(nonzero_keys(ty), nonzero_keys(tx)) - x_notin_y = setdiff(nonzero_keys(tx), nonzero_keys(ty)) - inboth = intersect(nonzero_keys(ty), nonzero_keys(tx)) - - # remove elements that are not in tx - for k in y_notin_x - delete!(ty.data, k) - end - # in-place scale elements that are in both - for k in inboth - ty[k] = scale!(ty[k], tx[k], α) + # entries of ty that are structurally zero in tx have to be zeroed out + issparse(tx) && zerovector!(ty) + for (I, v) in nonzero_pairs(tx) + ty[I] = scale!!(ty[I], v, α) end - # new scale for elements in x that are not in y - for k in x_notin_y - ty[k] = scale(tx[k], α) - end - return ty end @@ -61,41 +44,27 @@ function VI.add(ty::AbstractBlockTensorMap, tx::AbstractBlockTensorMap, α::Numb return add!(scale!(tdst, ty, β), tx, α) end -function VI.add!(ty::BlockTensorMap, tx::BlockTensorMap, α::Number, β::Number) +function VI.add!(ty::AbstractBlockTensorMap, tx::AbstractBlockTensorMap, α::Number, β::Number) space(ty) == space(tx) || throw(SpaceMismatch("$(space(ty)) ≠ $(space(tx))")) - add!(parent(ty), parent(tx), α, β) - return ty -end -function VI.add!(ty::SparseBlockTensorMap, tx::SparseBlockTensorMap, α::Number, β::Number) - space(ty) == space(tx) || throw(SpaceMismatch("$(space(ty)) ≠ $(space(tx))")) - y_notin_x = setdiff(nonzero_keys(ty), nonzero_keys(tx)) - x_notin_y = setdiff(nonzero_keys(tx), nonzero_keys(ty)) - inboth = intersect(nonzero_keys(ty), nonzero_keys(tx)) - - for k in y_notin_x - ty[k] = scale!!(ty[k], β) - end - for k in x_notin_y - ty[k] = scale(tx[k], α) - end - for k in inboth - ty[k] = add!!(ty[k], tx[k], α, β) + isone(β) || scale!(ty, β) + for (I, v) in nonzero_pairs(tx) + ty[I] = add!!(ty[I], v, α, One()) end - return ty end # inner # ----- -function VI.inner(x::BlockTensorMap, y::BlockTensorMap) - space(y) == space(x) || throw(SpaceMismatch()) - return inner(parent(x), parent(y)) -end -function VI.inner(x::SparseBlockTensorMap, y::SparseBlockTensorMap) - space(x) == space(y) || throw(SpaceMismatch()) - both_nonzero = intersect(nonzero_keys(x), nonzero_keys(y)) +function VI.inner(x::AbstractBlockTensorMap, y::AbstractBlockTensorMap) + space(x) == space(y) || throw(SpaceMismatch("$(space(x)) ≠ $(space(y))")) T = VI.promote_inner(x, y) - return sum(both_nonzero; init = zero(T)) do k - inner(x[k], y[k]) + # only entries that are nonzero in both contribute + ks = if issparse(x) && issparse(y) + intersect(nonzero_keys(x), nonzero_keys(y)) + else + nonzero_keys(issparse(y) ? y : x) + end + return sum(ks; init = zero(T)) do I + inner(x[I], y[I]) end end diff --git a/test/linalg/vectorinterface.jl b/test/linalg/vectorinterface.jl index 5692f50..3d5a57a 100644 --- a/test/linalg/vectorinterface.jl +++ b/test/linalg/vectorinterface.jl @@ -14,17 +14,17 @@ Vtr = ( V = Vtr -@testset "VectorInterface $(issparse ? "SparseBlockTensorMap" : "BlockTensorMap")" for issparse in - ( - false, true, +storagename(sparse) = sparse ? "SparseBlockTensorMap" : "BlockTensorMap" +makestorage(sparse, W) = sparse ? sprand(Float64, W, 0.5) : rand(W) + +@testset "VectorInterface $(storagename(sparse_y)) ← $(storagename(sparse_x))" for ( + sparse_y, sparse_x, + ) in ( + (false, false), (true, true), (false, true), (true, false), ) - if issparse - t = sprand(Float64, *(V[1:3]...) ← *(V[4], V[5]), 0.5) - t′ = sprand(Float64, *(V[1:3]...) ← *(V[4], V[5]), 0.5) - else - t = rand(*(V[1:3]...) ← *(V[4], V[5])) - t′ = rand(*(V[1:3]...) ← *(V[4], V[5])) - end + W = *(V[1:3]...) ← *(V[4], V[5]) + t = makestorage(sparse_y, W) + t′ = makestorage(sparse_x, W) @testset "scalartype" begin @test Float64 === @inferred scalartype(t) @@ -132,6 +132,15 @@ V = Vtr @test @inferred inner(t″, t‴) ≈ conj(α) * β * inner(t, t′) end + @testset "mixed-storage in-place" begin + α, β = rand(2) + @test @inferred scale!(zerovector(t), t′, α) ≈ scale(t′, α) + @test @inferred axpy!(α, t′, deepcopy(t)) ≈ add(t, t′, α) + @test @inferred axpby!(α, t′, β, deepcopy(t)) ≈ add(t, t′, α, β) + @test @inferred mul!(zerovector(t), α, t′) ≈ scale(t′, α) + @test @inferred dot(t, t′) ≈ inner(t, t′) + end + @testset "general linalg" begin α, β = rand(ComplexF64, 2) @test (α * α) * t ≈ α * (α * t) @@ -143,3 +152,31 @@ V = Vtr @test inner(t, t′) ≈ conj(inner(t′, t)) end end + +# regression test for #74: heterogeneous summands give interior zero-length coupled blocks +@testset "mixed storage with interior zero blocks" begin + S = Vect[U1Irrep] + P = S(0 => 1, 1 => 1, 2 => 1) + Vu = SumSpace(S(0 => 1), S(1 => 2), S(-1 => 2)) + W = (Vu ⊗ SumSpace(P)) ← SumSpace(P) + + tsp = SparseBlockTensorMap{TensorMap{Float64, S, 2, 1, Vector{Float64}}}(undef, W) + for (I, Wi) in zip(CartesianIndices(size(tsp)), BlockTensorKit.eachspace(tsp)) + dim(Wi) == 0 || (tsp[I] = randn(Float64, Wi)) + end + tdn = BlockTensorMap(tsp) + @test convert(TensorMap, tdn) ≈ convert(TensorMap, tsp) + + α, β = rand(2) + ref = add(tdn, tdn, α, β) + for (y, x) in ((tdn, tsp), (tsp, tdn)) + @test add!(deepcopy(y), x, α, β) ≈ ref + @test add(y, x, α, β) ≈ ref + @test y + x ≈ add(tdn, tdn) + @test norm(y - x) ≈ 0 atol = 1.0e-12 + @test scale!(zerovector(y), x, α) ≈ scale(tdn, α) + @test axpy!(α, x, deepcopy(y)) ≈ add(tdn, tdn, α) + @test inner(y, x) ≈ inner(tdn, tdn) + @test dot(y, x) ≈ inner(tdn, tdn) + end +end