Add allow_denser as option to the algorithm

Author: gdalle

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

@@ -189,15 +189,25 @@ 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).
 
+# 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 full (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
@@ -212,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]
@@ -227,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
@@ -236,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
@@ -260,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).
 
@@ -356,27 +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; allow_superset=true)
+    check_same_pattern(A, S; allow_denser)
     nzA = nonzeros(A)
     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)))
+    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,27 +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; allow_superset=true)
+    check_same_pattern(A, S; allow_denser)
     nzA = nonzeros(A)
     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
+                    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
@@ -520,40 +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; allow_superset=true)
+        check_same_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)
-            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
+                        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
@@ -796,7 +838,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_same_pattern(A, S)
     m, n = size(A)
     Br_and_Bc = _join_compressed!(result, Br, Bc)
     # pretend A is larger

src/graph.jl

@@ -99,8 +99,10 @@ end
 
 Return a [`SparsityPatternCSC`](@ref) corresponding to the matrix `[0 Aᵀ; A 0]`, with a minimum of allocations.
 """
-bidirectional_pattern(A::AbstractMatrix; symmetric_pattern::Bool) =
-    bidirectional_pattern(SparsityPatternCSC(SparseMatrixCSC(A)); symmetric_pattern)
+function bidirectional_pattern(A::AbstractMatrix; symmetric_pattern::Bool)
+    S = SparsityPatternCSC(SparseMatrixCSC(A))
+    return bidirectional_pattern(S; symmetric_pattern)
+end
 
 function bidirectional_pattern(S::SparsityPatternCSC{T}; symmetric_pattern::Bool) where {T}
     m, n = size(S)
@@ -171,7 +173,7 @@ function bidirectional_pattern(S::SparsityPatternCSC{T}; symmetric_pattern::Bool
 
     # Create the SparsityPatternCSC of the augmented adjacency matrix
     S_and_Sᵀ = SparsityPatternCSC{T}(p, p, colptr, rowval)
-    return S_and_Sᵀ, edge_to_index
+    return S, S_and_Sᵀ, edge_to_index
 end
 
 function build_edge_to_index(S::SparsityPatternCSC{T}) where {T}