Add allow_denser kwarg for direct decompression into denser SparseMatrixCSC

docs/src/dev.md

@@ -58,7 +58,6 @@ SparseMatrixColorings.valid_dynamic_order
 ```@docs
 SparseMatrixColorings.respectful_similar
 SparseMatrixColorings.matrix_versions
-SparseMatrixColorings.same_pattern
 ```
 
 ## Visualization

src/adtypes.jl

@@ -6,7 +6,7 @@ function coloring(
 )
     bg = BipartiteGraph(A)
     color = convert(Vector{eltype(bg)}, ADTypes.column_coloring(A, algo))
-    return ColumnColoringResult(A, bg, color)
+    return ColumnColoringResult(A, bg, color; allow_denser=false)
 end
 
 function coloring(
@@ -17,5 +17,5 @@ function coloring(
 )
     bg = BipartiteGraph(A)
     color = convert(Vector{eltype(bg)}, ADTypes.row_coloring(A, algo))
-    return RowColoringResult(A, bg, color)
+    return RowColoringResult(A, bg, color; allow_denser=false)
 end

src/constant.jl

@@ -9,12 +9,22 @@ Indeed, for symmetric coloring problems, we need more than just the vector of co
 
 # Constructors
 
-    ConstantColoringAlgorithm{partition}(matrix_template, color)
-    ConstantColoringAlgorithm(matrix_template, color; partition=:column)
-
+    ConstantColoringAlgorithm{partition}(
+        matrix_template,
+        color;
+        allow_denser=false
+    )
+
+    ConstantColoringAlgorithm(
+        matrix_template, color;
+        partition=:column,
+        allow_denser=false
+    )
+
 - `partition::Symbol`: either `:row` or `:column`.
 - `matrix_template::AbstractMatrix`: matrix for which the vector of colors was precomputed (the algorithm will only accept matrices of the exact same size).
 - `color::Vector{<:Integer}`: vector of integer colors, one for each row or column (depending on `partition`).
+- `allow_denser::Bool`: whether or not to allow decompression into a `SparseMatrixCSC` which has more nonzeros than the original sparsity pattern (see [`decompress!`](@ref) to know when this is implemented).
 
 !!! warning
     The second constructor (based on keyword arguments) is type-unstable.
@@ -73,30 +83,38 @@ struct ConstantColoringAlgorithm{
     matrix_template::M
     color::Vector{T}
     result::R
+    allow_denser::Bool
 end
 
 function ConstantColoringAlgorithm{:column}(
-    matrix_template::AbstractMatrix, color::Vector{<:Integer}
+    matrix_template::AbstractMatrix, color::Vector{<:Integer}; allow_denser::Bool=false
 )
     bg = BipartiteGraph(matrix_template)
-    result = ColumnColoringResult(matrix_template, bg, color)
+    result = ColumnColoringResult(matrix_template, bg, color; allow_denser)
     T, M, R = eltype(bg), typeof(matrix_template), typeof(result)
-    return ConstantColoringAlgorithm{:column,M,T,R}(matrix_template, color, result)
+    return ConstantColoringAlgorithm{:column,M,T,R}(
+        matrix_template, color, result, allow_denser
+    )
 end
 
 function ConstantColoringAlgorithm{:row}(
-    matrix_template::AbstractMatrix, color::Vector{<:Integer}
+    matrix_template::AbstractMatrix, color::Vector{<:Integer}; allow_denser::Bool=false
 )
     bg = BipartiteGraph(matrix_template)
-    result = RowColoringResult(matrix_template, bg, color)
+    result = RowColoringResult(matrix_template, bg, color; allow_denser)
     T, M, R = eltype(bg), typeof(matrix_template), typeof(result)
-    return ConstantColoringAlgorithm{:row,M,T,R}(matrix_template, color, result)
+    return ConstantColoringAlgorithm{:row,M,T,R}(
+        matrix_template, color, result, allow_denser
+    )
 end
 
 function ConstantColoringAlgorithm(
-    matrix_template::AbstractMatrix, color::Vector{<:Integer}; partition::Symbol=:column
+    matrix_template::AbstractMatrix,
+    color::Vector{<:Integer};
+    partition::Symbol=:column,
+    allow_denser::Bool=false,
 )
-    return ConstantColoringAlgorithm{partition}(matrix_template, color)
+    return ConstantColoringAlgorithm{partition}(matrix_template, color; allow_denser)
 end
 
 function coloring(

src/decompression.jl

@@ -193,9 +193,21 @@ The out-of-place alternative is [`decompress`](@ref).
 Compression means summing either the columns or the rows of `A` which share the same color.
 It is done by calling [`compress`](@ref).
 
+# Details
+
 For `:symmetric` coloring results (and for those only), an optional positional argument `uplo in (:U, :L, :F)` can be passed to specify which part of the matrix `A` should be updated: the Upper triangle, the Lower triangle, or the Full matrix.
 When `A isa SparseMatrixCSC`, using the `uplo` argument requires a target matrix which only stores the relevant triangle(s).
 
+Some coloring algorithms ([`GreedyColoringAlgorithm`](@ref) and [`ConstantColoringAlgorithm`](@ref)) have an option called `allow_denser`, which enables in-place decompression into a `SparseMatrixCSC` containing more structural nonzeros than the sparsity pattern used for coloring.
+
+!!! warning
+    Decompression into a denser `SparseMatrixCSC` is only implemented when all the conditions below are satisfied simultaneously:
+      - the partition is `:row` or `:column` (not `:bidirectional`)
+      - the decompression is `:direct` (not `:substitution`)
+      - the decompression is for all colors (so [`decompress_single_color!`](@ref) will not work)
+
+    Outside of these cases, it is up to the user to make sure that the sparsity pattern of the decompression target is an exact match. 
+
 # Example
 
 ```jldoctest
@@ -210,6 +222,8 @@ julia> A = sparse([
 
 julia> result = coloring(A, ColoringProblem(), GreedyColoringAlgorithm());
 
+julia> result_tolerant = coloring(A, ColoringProblem(), GreedyColoringAlgorithm(; allow_denser=true));
+
 julia> collect.(column_groups(result))
 3-element Vector{Vector{Int64}}:
  [1, 2, 4]
@@ -225,6 +239,8 @@ julia> B = compress(A, result)
 
 julia> A2 = similar(A);
 
+julia> A3 = similar(A); A3[1, 1] = A3[end, end] = 1;
+
 julia> decompress!(A2, B, result)
 4×6 SparseMatrixCSC{Int64, Int64} with 9 stored entries:
  ⋅  ⋅  4  6  ⋅  9
@@ -234,6 +250,16 @@ julia> decompress!(A2, B, result)
 
 julia> A2 == A
 true
+
+julia> decompress!(A3, B, result_tolerant)
+4×6 SparseMatrixCSC{Int64, Int64} with 11 stored entries:
+ 0  ⋅  4  6  ⋅  9
+ 1  ⋅  ⋅  ⋅  7  ⋅
+ ⋅  2  ⋅  ⋅  8  ⋅
+ ⋅  3  5  ⋅  ⋅  0
+
+julia> A3 == A
+true
 ```
 
 # See also
@@ -258,6 +284,8 @@ Decompress the vector `b` corresponding to color `c` in-place into `A`, given a
 !!! warning
     This function will only update some coefficients of `A`, without resetting the rest to zero.
 
+# Details
+
 For `:symmetric` coloring results (and for those only), an optional positional argument `uplo in (:U, :L, :F)` can be passed to specify which part of the matrix `A` should be updated: the Upper triangle, the Lower triangle, or the Full matrix.
 When `A isa SparseMatrixCSC`, using the `uplo` argument requires a target matrix which only stores the relevant triangle(s).
 
@@ -324,7 +352,7 @@ end
 function decompress!(A::AbstractMatrix, B::AbstractMatrix, result::ColumnColoringResult)
     (; color) = result
     S = result.bg.S2
-    check_same_pattern(A, S)
+    @assert size(A) == size(S)
     fill!(A, zero(eltype(A)))
     rvS = rowvals(S)
     for j in axes(S, 2)
@@ -342,7 +370,7 @@ function decompress_single_color!(
 )
     (; group) = result
     S = result.bg.S2
-    check_same_pattern(A, S)
+    @assert size(A) == size(S)
     rvS = rowvals(S)
     for j in group[c]
         for k in nzrange(S, j)
@@ -354,12 +382,28 @@ function decompress_single_color!(
 end
 
 function decompress!(A::SparseMatrixCSC, B::AbstractMatrix, result::ColumnColoringResult)
-    (; compressed_indices) = result
+    (; compressed_indices, allow_denser) = result
     S = result.bg.S2
-    check_same_pattern(A, S)
+    check_compatible_pattern(A, S; allow_denser)
     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(S)
+        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
+                    nzA[k + shift] = zero(eltype(A))
+                    shift += 1
+                end
+                nzA[k + shift] = B[compressed_indices[k]]
+            end
+        end
+        nzA[(nnz(S) + shift + 1):end] .= zero(eltype(A))
     end
     return A
 end
@@ -369,7 +413,7 @@ function decompress_single_color!(
 )
     (; group) = result
     S = result.bg.S2
-    check_same_pattern(A, S)
+    check_compatible_pattern(A, S)
     rvS = rowvals(S)
     nzA = nonzeros(A)
     for j in group[c]
@@ -386,7 +430,7 @@ end
 function decompress!(A::AbstractMatrix, B::AbstractMatrix, result::RowColoringResult)
     (; color) = result
     S = result.bg.S2
-    check_same_pattern(A, S)
+    @assert size(A) == size(S)
     fill!(A, zero(eltype(A)))
     rvS = rowvals(S)
     for j in axes(S, 2)
@@ -404,7 +448,7 @@ function decompress_single_color!(
 )
     (; group) = result
     S, Sᵀ = result.bg.S2, result.bg.S1
-    check_same_pattern(A, S)
+    @assert size(A) == size(S)
     rvSᵀ = rowvals(Sᵀ)
     for i in group[c]
         for k in nzrange(Sᵀ, i)
@@ -416,12 +460,28 @@ function decompress_single_color!(
 end
 
 function decompress!(A::SparseMatrixCSC, B::AbstractMatrix, result::RowColoringResult)
-    (; compressed_indices) = result
+    (; compressed_indices, allow_denser) = result
     S = result.bg.S2
-    check_same_pattern(A, S)
+    check_compatible_pattern(A, S; allow_denser)
     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)
+        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
+                    nzA[k + shift] = zero(eltype(A))
+                    shift += 1
+                end
+                nzA[k + shift] = B[compressed_indices[k]]
+            end
+        end
+        nzA[(nnz(S) + shift + 1):end] .= zero(eltype(A))
     end
     return A
 end
@@ -433,7 +493,7 @@ function decompress!(
 )
     (; ag, compressed_indices) = result
     (; S) = ag
-    uplo == :F && check_same_pattern(A, S)
+    @assert size(A) == size(S)
     fill!(A, zero(eltype(A)))
 
     rvS = rowvals(S)
@@ -457,7 +517,7 @@ function decompress_single_color!(
 )
     (; ag, compressed_indices, group) = result
     (; S) = ag
-    uplo == :F && check_same_pattern(A, S)
+    @assert size(A) == size(S)
 
     lower_index = (c - 1) * S.n + 1
     upper_index = c * S.n
@@ -488,25 +548,54 @@ end
 function decompress!(
     A::SparseMatrixCSC, B::AbstractMatrix, result::StarSetColoringResult, uplo::Symbol=:F
 )
-    (; ag, compressed_indices) = result
+    (; ag, compressed_indices, allow_denser) = result
     (; 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_compatible_pattern(A, S; allow_denser)
+        if nnz(A) == nnz(S)
+            for k in eachindex(nzA, compressed_indices)
+                nzA[k] = B[compressed_indices[k]]
+            end
+        else  # nnz(A) > nnz(S)
+            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
+                        nzA[k + shift] = zero(eltype(A))
+                        shift += 1
+                    end
+                    nzA[k + shift] = B[compressed_indices[k]]
+                end
+            end
+            nzA[(nnz(S) + shift + 1):end] .= zero(eltype(A))
         end
     else
-        rvS = rowvals(S)
+        rvA, rvS = rowvals(A), rowvals(S)
         l = 0  # assume A has the same pattern as the triangle
+        shift = 0
         for j in axes(S, 2)
             for k in nzrange(S, j)
                 i = rvS[k]
                 if in_triangle(i, j, uplo)
                     l += 1
-                    nzA[l] = B[compressed_indices[k]]
+                    while (l + shift) < A.colptr[j] || rvA[l + shift] < i
+                        nzA[l + shift] = zero(eltype(A))
+                        shift += 1
+                    end
+                    nzA[l + shift] = B[compressed_indices[k]]
                 end
             end
+            nzA[(l + shift + 1):end] .= zero(eltype(A))
+            if !allow_denser && shift > 0
+                throw(
+                    DimensionMismatch(
+                        "Decompression target must have exactly as many nonzeros as the sparsity pattern used for coloring",
+                    ),
+                )
+            end
         end
     end
     return A
@@ -519,7 +608,7 @@ function decompress!(
 )
     (; ag, color, reverse_bfs_orders, buffer) = result
     (; S) = ag
-    uplo == :F && check_same_pattern(A, S)
+    @assert size(A) == size(S)
     R = eltype(A)
     fill!(A, zero(R))
 
@@ -582,7 +671,7 @@ function decompress!(
     (; S) = ag
     A_colptr = A.colptr
     nzA = nonzeros(A)
-    uplo == :F && check_same_pattern(A, S)
+    uplo == :F && check_compatible_pattern(A, S)
 
     if eltype(buffer) == R
         buffer_right_type = buffer
@@ -679,7 +768,7 @@ function decompress!(
 )
     (; color, strict_upper_nonzero_inds, M_factorization, strict_upper_nonzeros_A) = result
     S = result.ag.S
-    uplo == :F && check_same_pattern(A, S)
+    @assert size(A) == size(S)
 
     # TODO: for some reason I cannot use ldiv! with a sparse QR
     strict_upper_nonzeros_A = M_factorization \ vec(B)
@@ -740,6 +829,7 @@ function decompress!(
     A::AbstractMatrix, Br::AbstractMatrix, Bc::AbstractMatrix, result::BicoloringResult
 )
     m, n = size(A)
+    @assert size(A) == size(result.S)
     Br_and_Bc = _join_compressed!(result, Br, Bc)
     A_and_Aᵀ = decompress(Br_and_Bc, result.symmetric_result)
     copyto!(A, A_and_Aᵀ[(n + 1):(n + m), 1:n])  # original matrix in bottom left corner
@@ -749,7 +839,8 @@ end
 function decompress!(
     A::SparseMatrixCSC, Br::AbstractMatrix, Bc::AbstractMatrix, result::BicoloringResult
 )
-    (; large_colptr, large_rowval, symmetric_result) = result
+    (; S, large_colptr, large_rowval, symmetric_result) = result
+    check_compatible_pattern(A, S)
     m, n = size(A)
     Br_and_Bc = _join_compressed!(result, Br, Bc)
     # pretend A is larger