Implement single-color decompression (#81)

* Add single-color decompression

* Encode undefined hubs as negative

* Typos

* Typo

* Fix test utils

Author: gdalle

docs/src/api.md

@@ -35,6 +35,7 @@ row_groups
 compress
 decompress
 decompress!
+decompress_single_color!
 ```
 
 ## Orders

docs/src/dev.md

@@ -26,15 +26,15 @@ SparseMatrixColorings.symmetric_coefficient
 SparseMatrixColorings.star_coloring
 SparseMatrixColorings.acyclic_coloring
 SparseMatrixColorings.group_by_color
-SparseMatrixColorings.get_matrix
 SparseMatrixColorings.StarSet
 SparseMatrixColorings.TreeSet
 ```
 
 ## Concrete coloring results
 
 ```@docs
-SparseMatrixColorings.NonSymmetricColoringResult
+SparseMatrixColorings.ColumnColoringResult
+SparseMatrixColorings.RowColoringResult
 SparseMatrixColorings.StarSetColoringResult
 SparseMatrixColorings.TreeSetColoringResult
 SparseMatrixColorings.LinearSystemColoringResult

src/SparseMatrixColorings.jl

@@ -54,6 +54,6 @@ export ColoringProblem, GreedyColoringAlgorithm, AbstractColoringResult
 export coloring
 export column_colors, row_colors
 export column_groups, row_groups
-export compress, decompress, decompress!
+export compress, decompress, decompress!, decompress_single_color!
 
 end

src/coloring.jl

@@ -131,8 +131,23 @@ $TYPEDFIELDS
 struct StarSet
     "a mapping from edges (pair of vertices) to their star index"
     star::Dict{Tuple{Int,Int},Int}
-    "a mapping from star indices to their hub (the hub is `0` if the star only contains one edge)"
+    "a mapping from star indices to their hub (undefined hubs for single-edge stars are the negative value of one of the vertices, picked arbitrarily)"
     hub::Vector{Int}
+    "a mapping from star indices to the vector of their spokes"
+    spokes::Vector{Vector{Int}}
+end
+
+function StarSet(star, hub)
+    spokes = [Int[] for s in eachindex(hub)]
+    for ((i, j), s) in pairs(star)
+        h = hub[s]
+        if i == abs(h)
+            push!(spokes[s], j)
+        elseif j == abs(h)
+            push!(spokes[s], i)
+        end
+    end
+    return StarSet(star, hub, spokes)
 end
 
 _sort(u, v) = (min(u, v), max(u, v))
@@ -185,7 +200,7 @@ function _update_stars!(
                 hub[star[vq]] = v  # this may already be true
                 star[vw] = star[vq]
             else  # vw forms a new star
-                push!(hub, 0)  # hub is yet undefined
+                push!(hub, -max(v, w))  # hub is undefined so we set it to a negative value, but it allows us to remember one of the two vertices
                 star[vw] = length(hub)
             end
         end
@@ -194,13 +209,6 @@ function _update_stars!(
 end
 
 """
-    symmetric_coefficient(
-        i::Integer, j::Integer,
-        color::AbstractVector{<:Integer},
-        group::AbstractVector{<:AbstractVector{<:Integer}},
-        S::AbstractMatrix{Bool}
-    )
-
     symmetric_coefficient(
         i::Integer, j::Integer,
         color::AbstractVector{<:Integer},
@@ -209,30 +217,12 @@ end
 
 Return the indices `(k, c)` such that `A[i, j] = B[k, c]`, where `A` is the uncompressed symmetric matrix and `B` is the column-compressed matrix.
 
-The first version corresponds to algorithm `DirectRecover1` in the paper, the second to `DirectRecover2`.
+This function corresponds to algorithm `DirectRecover2` in the paper.
 
 # References
 
-> [_Efficient Computation of Sparse Hessians Using Coloring and Automatic Differentiation_](https://pubsonline.informs.org/doi/abs/10.1287/ijoc.1080.0286), Gebremedhin et al. (2009), Figures 2 and 3
+> [_Efficient Computation of Sparse Hessians Using Coloring and Automatic Differentiation_](https://pubsonline.informs.org/doi/abs/10.1287/ijoc.1080.0286), Gebremedhin et al. (2009), Figure 3
 """
-function symmetric_coefficient end
-
-function symmetric_coefficient(
-    i::Integer,
-    j::Integer,
-    color::AbstractVector{<:Integer},
-    group::AbstractVector{<:AbstractVector{<:Integer}},
-    S::AbstractMatrix,
-)
-    for j2 in group[color[j]]
-        j2 == j && continue
-        if !iszero(S[i, j2])
-            return j, color[i]
-        end
-    end
-    return i, color[j]
-end
-
 function symmetric_coefficient(
     i::Integer, j::Integer, color::AbstractVector{<:Integer}, star_set::StarSet
 )
@@ -246,11 +236,7 @@ function symmetric_coefficient(
         i, j = j, i
     end
     star_id = star[i, j]
-    h = hub[star_id]
-    if h == 0
-        # pick arbitrary hub
-        h = i
-    end
+    h = abs(hub[star_id])
     if h == j
         # i is the spoke
         return i, color[h]

src/decompression.jl

@@ -116,7 +116,7 @@ true
 - [`AbstractColoringResult`](@ref)
 """
 function decompress(B::AbstractMatrix{R}, result::AbstractColoringResult) where {R<:Real}
-    S = get_matrix(result)
+    @compat (; S) = result
     A = respectful_similar(S, R)
     return decompress!(A, B, result)
 end
@@ -183,33 +183,84 @@ true
 """
 function decompress! end
 
-## NonSymmetricColoringResult
+"""
+    decompress_single_color!(
+        A::AbstractMatrix, b::AbstractVector, c::Integer,
+        result::AbstractColoringResult,
+    )
+
+Decompress the vector `b` corresponding to color `c` in-place into `A`, given a coloring `result` of the sparsity pattern of `A`.
+
+- If `result` comes from a `:nonsymmetric` structure with `:column` partition, this will update the columns of `A` that share color `c` (whose sum makes up `b`).
+- If `result` comes from a `:nonsymmetric` structure with `:row` partition, this will update the rows of `A` that share color `c` (whose sum makes up `b`).
+- If `result` comes from a `:symmetric` structure with `:column` partition, this will update the coefficients of `A` whose value is deduced from color `c`.
+
+!!! warning
+    This function will only update some coefficients of `A`, without resetting the rest to zero.
+
+# See also
+
+- [`ColoringProblem`](@ref)
+- [`AbstractColoringResult`](@ref)
+- [`decompress!`](@ref)
+"""
+function decompress_single_color! end
+
+## ColumnColoringResult
 
 function decompress!(
-    A::AbstractMatrix{R}, B::AbstractMatrix{R}, result::NonSymmetricColoringResult{:column}
+    A::AbstractMatrix{R}, B::AbstractMatrix{R}, result::ColumnColoringResult
 ) where {R<:Real}
-    S = get_matrix(result)
+    @compat (; S, color) = result
     check_same_pattern(A, S)
     A .= zero(R)
-    color = column_colors(result)
     rvS = rowvals(S)
     for j in axes(S, 2)
+        cj = color[j]
         for k in nzrange(S, j)
             i = rvS[k]
-            cj = color[j]
             A[i, j] = B[i, cj]
         end
     end
     return A
 end
 
+function decompress_single_color!(
+    A::AbstractMatrix{R}, b::AbstractVector{R}, c::Integer, result::ColumnColoringResult
+) where {R<:Real}
+    @compat (; S, group) = result
+    check_same_pattern(A, S)
+    view(A, :, group[c]) .= zero(R)
+    rvS = rowvals(S)
+    for j in group[c]
+        for k in nzrange(S, j)
+            i = rvS[k]
+            A[i, j] = b[i]
+        end
+    end
+    return A
+end
+
 function decompress!(
-    A::AbstractMatrix{R}, B::AbstractMatrix{R}, result::NonSymmetricColoringResult{:row}
+    A::SparseMatrixCSC{R}, B::AbstractMatrix{R}, result::ColumnColoringResult
 ) where {R<:Real}
-    S = get_matrix(result)
+    @compat (; S, compressed_indices) = result
+    check_same_pattern(A, S)
+    nzA = nonzeros(A)
+    for k in eachindex(nzA, compressed_indices)
+        nzA[k] = B[compressed_indices[k]]
+    end
+    return A
+end
+
+## RowColoringResult
+
+function decompress!(
+    A::AbstractMatrix{R}, B::AbstractMatrix{R}, result::RowColoringResult
+) where {R<:Real}
+    @compat (; S, color) = result
     check_same_pattern(A, S)
     A .= zero(R)
-    color = row_colors(result)
     rvS = rowvals(S)
     for j in axes(S, 2)
         for k in nzrange(S, j)
@@ -221,28 +272,30 @@ function decompress!(
     return A
 end
 
-function decompress!(
-    A::SparseMatrixCSC{R}, B::AbstractMatrix{R}, result::NonSymmetricColoringResult{:column}
+function decompress_single_color!(
+    A::AbstractMatrix{R}, b::AbstractVector{R}, c::Integer, result::RowColoringResult
 ) where {R<:Real}
-    S = get_matrix(result)
+    @compat (; S, Sᵀ, group) = result
     check_same_pattern(A, S)
-    nzA = nonzeros(A)
-    ind = result.compressed_indices
-    for i in eachindex(nzA, ind)
-        nzA[i] = B[ind[i]]
+    view(A, group[c], :) .= zero(R)
+    rvSᵀ = rowvals(Sᵀ)
+    for i in group[c]
+        for k in nzrange(Sᵀ, i)
+            j = rvSᵀ[k]
+            A[i, j] = b[j]
+        end
     end
     return A
 end
 
 function decompress!(
-    A::SparseMatrixCSC{R}, B::AbstractMatrix{R}, result::NonSymmetricColoringResult{:row}
+    A::SparseMatrixCSC{R}, B::AbstractMatrix{R}, result::RowColoringResult
 ) where {R<:Real}
-    S = get_matrix(result)
+    @compat (; S, compressed_indices) = result
     check_same_pattern(A, S)
     nzA = nonzeros(A)
-    ind = result.compressed_indices
-    for i in eachindex(nzA, ind)
-        nzA[i] = B[ind[i]]
+    for k in eachindex(nzA, compressed_indices)
+        nzA[k] = B[compressed_indices[k]]
     end
     return A
 end
@@ -252,16 +305,43 @@ end
 function decompress!(
     A::AbstractMatrix{R}, B::AbstractMatrix{R}, result::StarSetColoringResult
 ) where {R<:Real}
-    S = get_matrix(result)
+    @compat (; S, color, star_set) = result
+    @compat (; star, hub, spokes) = star_set
     check_same_pattern(A, S)
     A .= zero(R)
-    color = column_colors(result)
-    rvS = rowvals(S)
-    for j in axes(S, 2)
-        for k in nzrange(S, j)
-            i = rvS[k]
-            k, c = symmetric_coefficient(i, j, color, result.star_set)
-            A[i, j] = B[k, c]
+    for i in axes(A, 1)
+        if !iszero(S[i, i])
+            A[i, i] = B[i, color[i]]
+        end
+    end
+    for s in eachindex(hub, spokes)
+        j = abs(hub[s])
+        for i in spokes[s]
+            A[i, j] = B[i, color[j]]
+            A[j, i] = B[i, color[j]]
+        end
+    end
+    return A
+end
+
+function decompress_single_color!(
+    A::AbstractMatrix{R}, b::AbstractVector{R}, c::Integer, result::StarSetColoringResult
+) where {R<:Real}
+    @compat (; S, color, group, star_set) = result
+    @compat (; hub, spokes) = star_set
+    check_same_pattern(A, S)
+    for i in axes(A, 1)
+        if !iszero(S[i, i]) && color[i] == c
+            A[i, i] = b[i]
+        end
+    end
+    for s in eachindex(hub, spokes)
+        j = abs(hub[s])
+        if color[j] == c
+            for i in spokes[s]
+                A[i, j] = b[i]
+                A[j, i] = b[i]
+            end
         end
     end
     return A
@@ -270,12 +350,11 @@ end
 function decompress!(
     A::SparseMatrixCSC{R}, B::AbstractMatrix{R}, result::StarSetColoringResult
 ) where {R<:Real}
-    S = get_matrix(result)
+    @compat (; S, compressed_indices) = result
     check_same_pattern(A, S)
     nzA = nonzeros(A)
-    ind = result.compressed_indices
-    for i in eachindex(nzA, ind)
-        nzA[i] = B[ind[i]]
+    for k in eachindex(nzA, compressed_indices)
+        nzA[k] = B[compressed_indices[k]]
     end
     return A
 end
@@ -287,11 +366,9 @@ end
 function decompress!(
     A::AbstractMatrix{R}, B::AbstractMatrix{R}, result::TreeSetColoringResult
 ) where {R<:Real}
-    S = get_matrix(result)
+    @compat (; S, color, vertices_by_tree, reverse_bfs_orders, buffer) = result
     check_same_pattern(A, S)
     A .= zero(R)
-    color = column_colors(result)
-    @compat (; vertices_by_tree, reverse_bfs_orders, buffer) = result
 
     if eltype(buffer) == R
         buffer_right_type = buffer
@@ -327,10 +404,10 @@ end
 function decompress!(
     A::AbstractMatrix{R}, B::AbstractMatrix{R}, result::LinearSystemColoringResult
 ) where {R<:Real}
-    S = get_matrix(result)
+    @compat (;
+        S, color, strict_upper_nonzero_inds, T_factorization, strict_upper_nonzeros_A
+    ) = result
     check_same_pattern(A, S)
-    color = column_colors(result)
-    @compat (; strict_upper_nonzero_inds, T_factorization, strict_upper_nonzeros_A) = result
 
     # TODO: for some reason I cannot use ldiv! with a sparse QR
     strict_upper_nonzeros_A = T_factorization \ vec(B)

src/interface.jl

@@ -178,7 +178,7 @@ function coloring(
     S = sparse(A)
     bg = bipartite_graph(S)
     color = partial_distance2_coloring(bg, Val(2), algo.order)
-    return NonSymmetricColoringResult{:column}(S, color)
+    return ColumnColoringResult(S, color)
 end
 
 function coloring(
@@ -190,7 +190,7 @@ function coloring(
     S = sparse(A)
     bg = bipartite_graph(S)
     color = partial_distance2_coloring(bg, Val(1), algo.order)
-    return NonSymmetricColoringResult{:row}(S, color)
+    return RowColoringResult(S, color)
 end
 
 function coloring(