From 949de13b4831fac48fcb7f0eccce41a1fa5be607 Mon Sep 17 00:00:00 2001 From: Alberto Mercurio Date: Sun, 2 Jun 2024 20:25:15 +0200 Subject: [PATCH 1/3] Add support for sum and kronecker product of COO matrices --- .gitignore | 2 ++ src/coo_linalg.jl | 58 +++++++++++++++++++++++++++++++++++++++++++++++ test/runtests.jl | 15 ++++++++++++ 3 files changed, 75 insertions(+) diff --git a/.gitignore b/.gitignore index 20fe29d..b73746e 100644 --- a/.gitignore +++ b/.gitignore @@ -3,3 +3,5 @@ *.jl.mem /Manifest.toml /docs/build/ + +.vscode diff --git a/src/coo_linalg.jl b/src/coo_linalg.jl index 629e9a0..7f784fa 100644 --- a/src/coo_linalg.jl +++ b/src/coo_linalg.jl @@ -320,6 +320,64 @@ end +(D::Diagonal, A::SparseMatrixCOO) = A + D +function Base.:+(A::SparseMatrixCOO{T1}, B::SparseMatrixCOO{T2}) where {T1<:Number, T2<:Number} + A.n == B.n || throw(ArgumentError("A and B must have the same number of rows")) + A.m == B.m || throw(ArgumentError("A and B must have the same number of columns")) + + T = promote_type(T1, T2) + + rowval_colvalA = collect(zip(A.rows, A.cols)) + rowval_colvalB = collect(zip(B.rows, B.cols)) + + rowval_colval = union(rowval_colvalA, rowval_colvalB) + rowval = first.(rowval_colval) + colval = last.(rowval_colval) + + nzval = similar(rowval, T) + @inbounds for i in eachindex(rowval) + nzval[i] = zero(T) + for j in eachindex(rowval_colvalA) + if rowval_colvalA[j] == rowval_colval[i] + nzval[i] += A.vals[j] + break + end + end + for j in eachindex(rowval_colvalB) + if rowval_colvalB[j] == rowval_colval[i] + nzval[i] += B.vals[j] + break + end + end + end + + return SparseMatrixCOO(A.m, A.n, rowval, colval, nzval) +end + +function LinearAlgebra.kron(A::SparseMatrixCOO, B::SparseMatrixCOO) + mA, nA = size(A) + mB, nB = size(B) + out_shape = (mA * mB, nA * nB) + Annz = nnz(A) + Bnnz = nnz(B) + + if Annz == 0 || Bnnz == 0 + return SparseMatrixCOO(Int[], Int[], T[], out_shape...) + end + + row = (A.rows .- 1) .* mB + row = repeat(row, inner = Bnnz) + col = (A.cols .- 1) .* nB + col = repeat(col, inner = Bnnz) + data = repeat(A.vals, inner = Bnnz) + + row .+= repeat(B.rows, outer = Annz) + col .+= repeat(B.cols, outer = Annz) + + data .*= repeat(B.vals, outer = Annz) + + return SparseMatrixCOO(out_shape[1], out_shape[2], row, col, data) +end + # maximum! functions replace_if_minusinf(val::T, replacement::T) where {T} = (val == -T(Inf)) ? replacement : val function LinearAlgebra.maximum!(f::Function, v::AbstractVector{T}, A::SparseMatrixCOO{T}) where {T} diff --git a/test/runtests.jl b/test/runtests.jl index c2d6b09..eac87f8 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -293,6 +293,11 @@ end B_coo = D + A_coo @test norm(B_csc - B_coo) ≤ sqrt(eps()) * norm(B_csc) @test issorted(B_coo.cols) + + B = sprand(Float64, 20, 20, 0.1) + C_csc = A + B + C_coo = A_coo + SparseMatrixCOO(B) + @test norm(C_csc - C_coo) ≤ sqrt(eps()) * norm(C_csc) end @testset "row/col reduce" begin @@ -336,3 +341,13 @@ end maximum!(abs, v_coo', As_coo) @test norm(v - v_coo) ≤ sqrt(eps()) * norm(v) end + +@testset "Kronecker product" begin + A = sprand(Float64, 10, 15, 0.2) + B = sprand(Float64, 5, 7, 0.3) + A_coo = SparseMatrixCOO(A) + B_coo = SparseMatrixCOO(B) + C = kron(A, B) + C_coo = kron(A_coo, B_coo) + @test norm(C - C_coo) ≤ sqrt(eps()) * norm(C) +end From 728c9bae83fcbfb83d126a8d199cb203abeb37fa Mon Sep 17 00:00:00 2001 From: Alberto Mercurio <61953577+albertomercurio@users.noreply.github.com> Date: Wed, 12 Jun 2024 16:26:39 +0200 Subject: [PATCH 2/3] Update src/coo_linalg.jl Co-authored-by: Tangi Migot --- src/coo_linalg.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/coo_linalg.jl b/src/coo_linalg.jl index 7f784fa..9b031ce 100644 --- a/src/coo_linalg.jl +++ b/src/coo_linalg.jl @@ -321,7 +321,7 @@ end +(D::Diagonal, A::SparseMatrixCOO) = A + D function Base.:+(A::SparseMatrixCOO{T1}, B::SparseMatrixCOO{T2}) where {T1<:Number, T2<:Number} - A.n == B.n || throw(ArgumentError("A and B must have the same number of rows")) + A.n == B.n || throw(DimensionMismatch()) A.m == B.m || throw(ArgumentError("A and B must have the same number of columns")) T = promote_type(T1, T2) From 7d346ca720fe4583dfbf2e1ed2d73114a6479425 Mon Sep 17 00:00:00 2001 From: Alberto Mercurio <61953577+albertomercurio@users.noreply.github.com> Date: Wed, 12 Jun 2024 16:26:54 +0200 Subject: [PATCH 3/3] Update src/coo_linalg.jl Co-authored-by: Tangi Migot --- src/coo_linalg.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/coo_linalg.jl b/src/coo_linalg.jl index 9b031ce..b837d4d 100644 --- a/src/coo_linalg.jl +++ b/src/coo_linalg.jl @@ -322,7 +322,7 @@ end function Base.:+(A::SparseMatrixCOO{T1}, B::SparseMatrixCOO{T2}) where {T1<:Number, T2<:Number} A.n == B.n || throw(DimensionMismatch()) - A.m == B.m || throw(ArgumentError("A and B must have the same number of columns")) + A.m == B.m || throw(DimensionMismatch()) T = promote_type(T1, T2)