From acc046f26d932fbc43296136dd71b4eb7dab6426 Mon Sep 17 00:00:00 2001
From: Jishnu Bhattacharya <jishnub.github@gmail.com>
Date: Wed, 2 Oct 2024 18:25:52 +0530
Subject: [PATCH 1/3] Diagonal-sandwiched triple product for SparseMatrixCSC

---
 src/linalg.jl  |  3 +++
 test/linalg.jl | 16 ++++++++++++++++
 2 files changed, 19 insertions(+)

diff --git a/src/linalg.jl b/src/linalg.jl
index 131a21bc..1ba5e351 100644
--- a/src/linalg.jl
+++ b/src/linalg.jl
@@ -205,6 +205,9 @@ const SparseOrTri{Tv,Ti} = Union{SparseMatrixCSCUnion{Tv,Ti},SparseTriangular{Tv
 *(A::AdjOrTrans{<:Any,<:AbstractSparseMatrixCSC}, B::SparseOrTri) = spmatmul(copy(A), B)
 *(A::AdjOrTrans{<:Any,<:AbstractSparseMatrixCSC}, B::AdjOrTrans{<:Any,<:AbstractSparseMatrixCSC}) = spmatmul(copy(A), copy(B))
 
+(*)(Da::Diagonal, A::Union{SparseMatrixCSCUnion, AdjOrTrans{<:Any,<:AbstractSparseMatrixCSC}}, Db::Diagonal) =
+    Da * (A * Db)
+
 # Gustavson's matrix multiplication algorithm revisited.
 # The result rowval vector is already sorted by construction.
 # The auxiliary Vector{Ti} xb is replaced by a Vector{Bool} of same length.
diff --git a/test/linalg.jl b/test/linalg.jl
index 45d42d9f..9024f1c8 100644
--- a/test/linalg.jl
+++ b/test/linalg.jl
@@ -912,4 +912,20 @@ end
     @test sparse(3I, 4, 5) == sparse(1:4, 1:4, 3, 4, 5)
     @test sparse(3I, 5, 4) == sparse(1:4, 1:4, 3, 5, 4)
 end
+
+@testset "diagonal-sandwiched triple multiplication" begin
+    D = Diagonal(1:4)
+    S = sprand(Int, 4, 4, 0.2)
+    A = Array(S)
+    C = D * S * D
+    @test C isa SparseMatrixCSC
+    @test C ≈ D * A * D
+    C = D * S' * D
+    @test C isa SparseMatrixCSC
+    @test C ≈ D * A' * D
+    C = D * view(S, :, :) * D
+    @test C isa SparseMatrixCSC
+    @test C ≈ D * A * D
+end
+
 end

From 2a367284225dd44c94f57abf9498f8b3071f328e Mon Sep 17 00:00:00 2001
From: Jishnu Bhattacharya <jishnub.github@gmail.com>
Date: Fri, 4 Oct 2024 19:35:29 +0530
Subject: [PATCH 2/3] Specialize for SparseMatrixCSC

---
 src/linalg.jl  | 20 ++++++++++++++++++--
 test/linalg.jl | 17 +++++++++--------
 2 files changed, 27 insertions(+), 10 deletions(-)

diff --git a/src/linalg.jl b/src/linalg.jl
index 1ba5e351..282d44b4 100644
--- a/src/linalg.jl
+++ b/src/linalg.jl
@@ -205,8 +205,24 @@ const SparseOrTri{Tv,Ti} = Union{SparseMatrixCSCUnion{Tv,Ti},SparseTriangular{Tv
 *(A::AdjOrTrans{<:Any,<:AbstractSparseMatrixCSC}, B::SparseOrTri) = spmatmul(copy(A), B)
 *(A::AdjOrTrans{<:Any,<:AbstractSparseMatrixCSC}, B::AdjOrTrans{<:Any,<:AbstractSparseMatrixCSC}) = spmatmul(copy(A), copy(B))
 
-(*)(Da::Diagonal, A::Union{SparseMatrixCSCUnion, AdjOrTrans{<:Any,<:AbstractSparseMatrixCSC}}, Db::Diagonal) =
-    Da * (A * Db)
+(*)(Da::Diagonal, A::Union{SparseMatrixCSCUnion, AdjOrTrans{<:Any,<:AbstractSparseMatrixCSC}}, Db::Diagonal) = Da * (A * Db)
+function (*)(Da::Diagonal, A::SparseMatrixCSC, Db::Diagonal)
+    T = promote_op(matprod, eltype(Da), promote_op(matprod, eltype(A), eltype(Db)))
+    dest = similar(A, T)
+    vals_dest = nonzeros(dest)
+    rows = rowvals(A)
+    vals = nonzeros(A)
+    da, db = map(parent, (Da, Db))
+    for col in axes(A,2)
+        dbcol = db[col]
+        for i in nzrange(A, col)
+            row = rows[i]
+            val = vals[i]
+            vals_dest[i] = da[row] * val * dbcol
+        end
+    end
+    dest
+end
 
 # Gustavson's matrix multiplication algorithm revisited.
 # The result rowval vector is already sorted by construction.
diff --git a/test/linalg.jl b/test/linalg.jl
index 9024f1c8..7c8ae0a4 100644
--- a/test/linalg.jl
+++ b/test/linalg.jl
@@ -914,18 +914,19 @@ end
 end
 
 @testset "diagonal-sandwiched triple multiplication" begin
-    D = Diagonal(1:4)
-    S = sprand(Int, 4, 4, 0.2)
+    D1 = Diagonal(1:4)
+    D2 = Diagonal(2:2:8)
+    S = sprand(4, 4, 0.2)
     A = Array(S)
-    C = D * S * D
+    C = D1 * S * D2
     @test C isa SparseMatrixCSC
-    @test C ≈ D * A * D
-    C = D * S' * D
+    @test C ≈ D1 * A * D2
+    C = D1 * S' * D2
     @test C isa SparseMatrixCSC
-    @test C ≈ D * A' * D
-    C = D * view(S, :, :) * D
+    @test C ≈ D1 * A' * D2
+    C = D1 * view(S, :, :) * D2
     @test C isa SparseMatrixCSC
-    @test C ≈ D * A * D
+    @test C ≈ D1 * A * D2
 end
 
 end

From dbe4321a7950f84f808b7bf37fce4e844b9c26ef Mon Sep 17 00:00:00 2001
From: Jishnu Bhattacharya <jishnub.github@gmail.com>
Date: Fri, 4 Oct 2024 20:14:58 +0530
Subject: [PATCH 3/3] Test for non-square matrix

---
 src/linalg.jl  |  2 ++
 test/linalg.jl | 13 ++++++++-----
 2 files changed, 10 insertions(+), 5 deletions(-)

diff --git a/src/linalg.jl b/src/linalg.jl
index 282d44b4..9cf91d29 100644
--- a/src/linalg.jl
+++ b/src/linalg.jl
@@ -207,6 +207,8 @@ const SparseOrTri{Tv,Ti} = Union{SparseMatrixCSCUnion{Tv,Ti},SparseTriangular{Tv
 
 (*)(Da::Diagonal, A::Union{SparseMatrixCSCUnion, AdjOrTrans{<:Any,<:AbstractSparseMatrixCSC}}, Db::Diagonal) = Da * (A * Db)
 function (*)(Da::Diagonal, A::SparseMatrixCSC, Db::Diagonal)
+    (size(Da, 2) == size(A,1) && size(A,2) == size(Db,1)) ||
+        throw(DimensionMismatch("incompatible sizes"))
     T = promote_op(matprod, eltype(Da), promote_op(matprod, eltype(A), eltype(Db)))
     dest = similar(A, T)
     vals_dest = nonzeros(dest)
diff --git a/test/linalg.jl b/test/linalg.jl
index 7c8ae0a4..d3f004ca 100644
--- a/test/linalg.jl
+++ b/test/linalg.jl
@@ -914,19 +914,22 @@ end
 end
 
 @testset "diagonal-sandwiched triple multiplication" begin
-    D1 = Diagonal(1:4)
-    D2 = Diagonal(2:2:8)
-    S = sprand(4, 4, 0.2)
+    S = sprand(4, 6, 0.2)
+    D1 = Diagonal(axes(S,1))
+    D2 = Diagonal(axes(S,2) .+ 4)
     A = Array(S)
     C = D1 * S * D2
     @test C isa SparseMatrixCSC
     @test C ≈ D1 * A * D2
-    C = D1 * S' * D2
+    C = D2 * S' * D1
     @test C isa SparseMatrixCSC
-    @test C ≈ D1 * A' * D2
+    @test C ≈ D2 * A' * D1
     C = D1 * view(S, :, :) * D2
     @test C isa SparseMatrixCSC
     @test C ≈ D1 * A * D2
+
+    @test_throws DimensionMismatch D2 * S * D2
+    @test_throws DimensionMismatch D1 * S * D1
 end
 
 end