feat: allow full direct decompression into larger SparseMatrixCSC

Author: gdalle

src/decompression.jl

@@ -189,6 +189,8 @@ The out-of-place alternative is [`decompress`](@ref).
 
 !!! note
     In-place decompression is faster when `A isa SparseMatrixCSC`.
+      - In general, this case requires the sparsity pattern of `A` to match the sparsity pattern `S` from which the coloring result was computed.
+      - For a coloring result with `decompression=:direct`, we also allow _full_ decompression into an `A` whose sparsity pattern is a strict superset of `S`.
 
 Compression means summing either the columns or the rows of `A` which share the same color.
 It is done by calling [`compress`](@ref).
@@ -356,10 +358,25 @@ end
 function decompress!(A::SparseMatrixCSC, B::AbstractMatrix, result::ColumnColoringResult)
     (; compressed_indices) = result
     S = result.bg.S2
-    check_same_pattern(A, S)
+    check_same_pattern(A, S; allow_superset=true)
     nzA = nonzeros(A)
-    for k in eachindex(nzA, compressed_indices)
-        nzA[k] = B[compressed_indices[k]]
+    if nnz(A) == nnz(S)
+        for k in eachindex(compressed_indices)
+            nzA[k] = B[compressed_indices[k]]
+        end
+    else  # nnz(A) > nnz(Z)
+        fill!(nonzeros(A), zero(eltype(A)))
+        rvA, rvS = rowvals(A), rowvals(S)
+        shift = 0
+        for j in axes(S, 2)
+            for k in nzrange(S, j)
+                i = rvS[k]
+                while (k + shift) < A.colptr[j] || rvA[k + shift] < i
+                    shift += 1
+                end
+                nzA[k + shift] = B[compressed_indices[k]]
+            end
+        end
     end
     return A
 end
@@ -418,10 +435,25 @@ end
 function decompress!(A::SparseMatrixCSC, B::AbstractMatrix, result::RowColoringResult)
     (; compressed_indices) = result
     S = result.bg.S2
-    check_same_pattern(A, S)
+    check_same_pattern(A, S; allow_superset=true)
     nzA = nonzeros(A)
-    for k in eachindex(nzA, compressed_indices)
-        nzA[k] = B[compressed_indices[k]]
+    if nnz(A) == nnz(S)
+        for k in eachindex(nzA, compressed_indices)
+            nzA[k] = B[compressed_indices[k]]
+        end
+    else  # nnz(A) > nnz(S)
+        fill!(nonzeros(A), zero(eltype(A)))
+        rvA, rvS = rowvals(A), rowvals(S)
+        shift = 0
+        for j in axes(S, 2)
+            for k in nzrange(S, j)
+                i = rvS[k]
+                while (k + shift) < A.colptr[j] || rvA[k + shift] < i
+                    shift += 1
+                end
+                nzA[k + shift] = B[compressed_indices[k]]
+            end
+        end
     end
     return A
 end
@@ -492,9 +524,24 @@ function decompress!(
     (; S) = ag
     nzA = nonzeros(A)
     if uplo == :F
-        check_same_pattern(A, S)
-        for k in eachindex(nzA, compressed_indices)
-            nzA[k] = B[compressed_indices[k]]
+        check_same_pattern(A, S; allow_superset=true)
+        if nnz(A) == nnz(S)
+            for k in eachindex(nzA, compressed_indices)
+                nzA[k] = B[compressed_indices[k]]
+            end
+        else  # nnz(A) > nnz(S)
+            fill!(nonzeros(A), zero(eltype(A)))
+            rvA, rvS = rowvals(A), rowvals(S)
+            shift = 0
+            for j in axes(S, 2)
+                for k in nzrange(S, j)
+                    i = rvS[k]
+                    while (k + shift) < A.colptr[j] || rvA[k + shift] < i
+                        shift += 1
+                    end
+                    nzA[k + shift] = B[compressed_indices[k]]
+                end
+            end
         end
     else
         rvS = rowvals(S)

src/matrices.jl

@@ -61,25 +61,24 @@ function respectful_similar(A::Union{Symmetric,Hermitian}, ::Type{T}) where {T}
     return respectful_similar(sparse(A), T)
 end
 
-"""
-    same_pattern(A, B)
-
-Perform a partial equality check on the sparsity patterns of `A` and `B`:
-
-- if the return is `true`, they might have the same sparsity pattern but we're not sure
-- if the return is `false`, they definitely don't have the same sparsity pattern
-"""
-same_pattern(A, B) = size(A) == size(B)
-
-function same_pattern(
-    A::Union{SparseMatrixCSC,SparsityPatternCSC},
-    B::Union{SparseMatrixCSC,SparsityPatternCSC},
-)
-    return size(A) == size(B) && nnz(A) == nnz(B)
+same_pattern(A::AbstractMatrix, S; allow_superset::Bool=false) = true
+function same_pattern(A::SparseMatrixCSC, S; allow_superset::Bool=false)
+    return allow_superset ? nnz(A) >= nnz(S) : nnz(A) == nnz(S)
 end
 
-function check_same_pattern(A, S)
-    if !same_pattern(A, S)
-        throw(DimensionMismatch("`A` and `S` must have the same sparsity pattern."))
+function check_same_pattern(A, S; allow_superset::Bool=false)
+    if size(A) != size(S)
+        throw(
+            DimensionMismatch(
+                "Decompression target must have the same size as sparsity  pattern"
+            ),
+        )
+    elseif !same_pattern(A, S; allow_superset)
+        throw(
+            DimensionMismatch(
+                """Decompression target must $(allow_superset ? "contain the nonzeros of the sparsity pattern" : "be equal to the sparsity pattern") used for coloring""",
+            ),
+        )
     end
+    return true
 end

test/matrices.jl

@@ -48,5 +48,8 @@ end
     @test !same_pattern(A2, S)
     @test same_pattern(Matrix(A2), S)
 
+    @test_throws DimensionMismatch check_same_pattern(vcat(A1, A1), S)
     @test_throws DimensionMismatch check_same_pattern(A2, S)
+    @test check_same_pattern(A1, S; allow_superset=true)
+    @test check_same_pattern(A2, S; allow_superset=true)
 end

test/utils.jl

@@ -86,6 +86,18 @@ function test_coloring_decompression(
             @test decompress(B, result) ≈ A0  # check result wasn't modified
             @test decompress!(respectful_similar(A, eltype(B)), B, result) ≈ A0
             @test decompress!(respectful_similar(A, eltype(B)), B, result) ≈ A0
+            if decompression == :direct && A isa SparseMatrixCSC
+                A_bigger = respectful_similar(A, eltype(B))
+                for _ in 1:10
+                    nb_coeffs_added = rand(1:minimum(size(A)))
+                    for _ in nb_coeffs_added
+                        i = rand(axes(A, 1))
+                        j = rand(axes(A, 2))
+                        A_bigger[i, j] = one(eltype(B))
+                    end
+                    @test decompress!(A_bigger, B, result) ≈ A0
+                end
+            end
         end
 
         @testset "Single-color decompression" begin