Implement clever symmetric decompression with star set (#25)

* Implement clever symmetric decompression with star set

* Field typo

* Typo

* Spoke handling

* Fix

Author: gdalle

Project.toml

@@ -1,7 +1,7 @@
 name = "SparseMatrixColorings"
 uuid = "0a514795-09f3-496d-8182-132a7b665d35"
 authors = ["Guillaume Dalle <22795598+gdalle@users.noreply.github.com>"]
-version = "0.3.3"
+version = "0.3.4"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"

docs/src/api.md

@@ -13,6 +13,7 @@ GreedyColoringAlgorithm
 column_coloring
 row_coloring
 symmetric_coloring
+symmetric_coloring_detailed
 ```
 
 ## Public, not exported
@@ -36,6 +37,7 @@ decompress_rows
 decompress_rows!
 decompress_symmetric
 decompress_symmetric!
+StarSet
 ```
 
 ## Private
@@ -64,6 +66,7 @@ vertices
 ```@docs
 partial_distance2_coloring
 star_coloring
+star_coloring_detailed
 ```
 
 ### Testing

src/SparseMatrixColorings.jl

@@ -46,8 +46,9 @@ include("check.jl")
 @compat public decompress_columns, decompress_columns!
 @compat public decompress_rows, decompress_rows!
 @compat public decompress_symmetric, decompress_symmetric!
+@compat public StarSet
 
 export GreedyColoringAlgorithm
-export column_coloring, row_coloring, symmetric_coloring
+export column_coloring, row_coloring, symmetric_coloring, symmetric_coloring_detailed
 
 end

src/adtypes.jl

@@ -26,7 +26,7 @@ end
 """
     column_coloring(A::AbstractMatrix, algo::GreedyColoringAlgorithm)
 
-Compute a partial distance-2 coloring of the columns in the bipartite graph of the matrix `A`.
+Compute a partial distance-2 coloring of the columns in the bipartite graph of the matrix `A`, return a vector of integer colors.
 
 Function defined by ADTypes, re-exported by SparseMatrixColorings.
 
@@ -64,7 +64,7 @@ end
 """
     row_coloring(A::AbstractMatrix, algo::GreedyColoringAlgorithm)
 
-Compute a partial distance-2 coloring of the rows in the bipartite graph of the matrix `A`.
+Compute a partial distance-2 coloring of the rows in the bipartite graph of the matrix `A`, return a vector of integer colors.
 
 Function defined by ADTypes, re-exported by SparseMatrixColorings.
 
@@ -101,16 +101,52 @@ end
 """
     symmetric_coloring(A::AbstractMatrix, algo::GreedyColoringAlgorithm)
 
-Compute a star coloring of the columns in the adjacency graph of the symmetric matrix `A`.
+Compute a star coloring of the columns in the adjacency graph of the symmetric matrix `A`, return a vector of integer colors.
 
 Function defined by ADTypes, re-exported by SparseMatrixColorings.
 
 # Example
 
-!!! warning
-    Work in progress.
+```jldoctest
+using SparseMatrixColorings, SparseArrays
+
+algo = GreedyColoringAlgorithm(SparseMatrixColorings.LargestFirst())
+
+A = sparse([
+    1 1 1 1
+    1 1 0 0
+    1 0 1 0
+    1 0 0 1
+])
+
+symmetric_coloring(A, algo)
+
+# output
+
+4-element Vector{Int64}:
+ 1
+ 2
+ 2
+ 2
+```
 """
 function ADTypes.symmetric_coloring(A::AbstractMatrix, algo::GreedyColoringAlgorithm)
     ag = adjacency_graph(A)
     return star_coloring(ag, algo.order)
 end
+
+"""
+    symmetric_coloring_detailed(A::AbstractMatrix, algo::GreedyColoringAlgorithm)
+
+Do the same as [`symmetric_coloring`](@ref) but return a tuple `(color, star_set)`, where
+
+- `color` is the vector of integer colors
+- `star_set` is a [`StarSet`](@ref) encoding the set of 2-colored stars, which can be used to speed up decompression
+
+!!! warning
+    This function is not defined by ADTypes.
+"""
+function symmetric_coloring_detailed(A::AbstractMatrix, algo::GreedyColoringAlgorithm)
+    ag = adjacency_graph(A)
+    return star_coloring_detailed(ag, algo.order)
+end

src/coloring.jl

@@ -72,6 +72,51 @@ The vertices are colored in a greedy fashion, following the `order` supplied.
 > [_New Acyclic and Star Coloring Algorithms with Application to Computing Hessians_](https://epubs.siam.org/doi/abs/10.1137/050639879), Gebremedhin et al. (2007), Algorithm 4.1
 """
 function star_coloring(g::Graph, order::AbstractOrder)
+    color, _ = star_coloring_detailed(g, order)
+    return color
+end
+
+"""
+    StarSet
+
+Encode a set of 2-colored stars resulting from the star coloring algorithm.
+
+# Fields
+
+The fields are not part of the public API, even though the type is.
+
+- `star::Dict{Tuple{Int,Int},Int}`: a mapping from edges (pair of vertices) their to star index
+- `hub::Vector{Int}`: a mapping from star indices to their hub (the hub is `0` if the star only contains one edge)
+
+# References
+
+> [_New Acyclic and Star Coloring Algorithms with Application to Computing Hessians_](https://epubs.siam.org/doi/abs/10.1137/050639879), Gebremedhin et al. (2007), Algorithm 4.1
+"""
+struct StarSet
+    star::Dict{Tuple{Int,Int},Int}
+    hub::Vector{Int}
+end
+
+"""
+    star_coloring_detailed(g::Graph, order::AbstractOrder)
+
+Do the same as [`star_coloring`](@ref) but return a tuple `(color, star_set)`, where
+
+- `color` is the vector of integer colors
+- `star_set` is a [`StarSet`](@ref) encoding the set of 2-colored stars, which can be used to speed up decompression
+
+# See also
+
+- [`star_coloring`](@ref)
+- [`StarSet`](@ref)
+
+# References
+
+> [_New Acyclic and Star Coloring Algorithms with Application to Computing Hessians_](https://epubs.siam.org/doi/abs/10.1137/050639879), Gebremedhin et al. (2007), Algorithm 4.1
+
+> [_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 2
+"""
+function star_coloring_detailed(g::Graph, order::AbstractOrder)
     # Initialize data structures
     color = zeros(Int, length(g))
     forbidden_colors = zeros(Int, length(g))
@@ -112,7 +157,7 @@ function star_coloring(g::Graph, order::AbstractOrder)
         end
         _update_stars!(star, hub, g, v, color, first_neighbor)
     end
-    return color
+    return color, StarSet(star, hub)
 end
 
 _sort(u, v) = (min(u, v), max(u, v))

src/decompression.jl

@@ -8,7 +8,7 @@
         color::AbstractVector{<:Integer}
     ) where {R<:Real}
 
-Decompress the thin matrix `B` into the fat matrix `A` which must have the same sparsity pattern as `S`.
+Decompress the narrow matrix `B` into the wide matrix `A` which must have the same sparsity pattern as `S`.
 
 Here, `color` is a column coloring of `S`, while `B` is a compressed representation of matrix `A` obtained by summing the columns that share the same color.
 """
@@ -63,7 +63,7 @@ end
         color::AbstractVector{<:Integer}
     ) where {R<:Real}
 
-Decompress the thin matrix `B` into a new fat matrix `A` with the same sparsity pattern as `S`.
+Decompress the narrow matrix `B` into a new wide matrix `A` with the same sparsity pattern as `S`.
 
 Here, `color` is a column coloring of `S`, while `B` is a compressed representation of matrix `A` obtained by summing the columns that share the same color.
 """
@@ -158,16 +158,19 @@ end
         A::AbstractMatrix{R},
         S::AbstractMatrix{Bool},
         B::AbstractMatrix{R},
-        color::AbstractVector{<:Integer}
+        color::AbstractVector{<:Integer},
+        [star_set::StarSet],
     ) where {R<:Real}
 
-Decompress the thin matrix `B` into the symmetric matrix `A` which must have the same sparsity pattern as `S`.
+Decompress the narrow matrix `B` into the symmetric matrix `A` which must have the same sparsity pattern as `S`.
 
 Here, `color` is a symmetric coloring of `S`, while `B` is a compressed representation of matrix `A` obtained by summing the columns that share the same color.
 
+Decompression is faster when a [`StarSet`](@ref) is also provided.
+
 # 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), Figure 2
+> [_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
 """
 function decompress_symmetric! end
 
@@ -201,35 +204,96 @@ function decompress_symmetric!(
     return A
 end
 
+function decompress_symmetric!(
+    A::AbstractMatrix{R},
+    S::AbstractMatrix{Bool},
+    B::AbstractMatrix{R},
+    color::AbstractVector{<:Integer},
+    star_set::StarSet,
+) where {R<:Real}
+    @compat (; star, hub) = star_set
+    checksquare(A)
+    if !same_sparsity_pattern(A, S)
+        throw(DimensionMismatch("`A` and `S` must have the same sparsity pattern."))
+    end
+    A .= zero(R)
+    for i in axes(A, 1)
+        if !iszero(S[i, i])
+            A[i, i] = B[i, color[i]]
+        end
+    end
+    for ((i, j), star_id) in pairs(star)
+        i == j && continue
+        h = hub[star_id]
+        if h == 0
+            # pick arbitrary hub
+            h = i
+        end
+        if h == j
+            # i is the spoke
+            A[i, j] = A[j, i] = B[i, color[h]]
+        elseif h == i
+            # j is the spoke
+            A[i, j] = A[j, i] = B[j, color[h]]
+        end
+    end
+    return A
+end
+
 function decompress_symmetric!(
     A::Symmetric{R},
     S::AbstractMatrix{Bool},
     B::AbstractMatrix{R},
-    colors::AbstractVector{<:Integer},
+    color::AbstractVector{<:Integer},
+) where {R<:Real}
+    # requires parent decompression to handle both upper and lower triangles
+    decompress_symmetric!(parent(A), S, B, color)
+    return A
+end
+
+function decompress_symmetric!(
+    A::Symmetric{R},
+    S::AbstractMatrix{Bool},
+    B::AbstractMatrix{R},
+    color::AbstractVector{<:Integer},
+    star_set::StarSet,
 ) where {R<:Real}
     # requires parent decompression to handle both upper and lower triangles
-    decompress_symmetric!(parent(A), S, B, colors)
+    decompress_symmetric!(parent(A), S, B, color, star_set)
     return A
 end
 
 """
     decompress_symmetric(
         S::AbstractMatrix{Bool},
         B::AbstractMatrix{R},
-        colors::AbstractVector{<:Integer}
+        color::AbstractVector{<:Integer},
+        [star_set::StarSet],
     ) where {R<:Real}
 
-Decompress the thin matrix `B` into a new symmetric matrix `A` with the same sparsity pattern as `S`.
+Decompress the narrow matrix `B` into a new symmetric matrix `A` with the same sparsity pattern as `S`.
 
-Here, `colors` is a symmetric coloring of `S`, while `B` is a compressed representation of matrix `A` obtained by summing the columns that share the same color.
+Here, `color` is a symmetric coloring of `S`, while `B` is a compressed representation of matrix `A` obtained by summing the columns that share the same color.
+
+Decompression is faster when a [`StarSet`](@ref) is also provided.
 
 # 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)
+> [_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
 """
 function decompress_symmetric(
-    S::AbstractMatrix{Bool}, B::AbstractMatrix{R}, colors::AbstractVector{<:Integer}
+    S::AbstractMatrix{Bool}, B::AbstractMatrix{R}, color::AbstractVector{<:Integer}
+) where {R<:Real}
+    A = respectful_similar(S, R)
+    return decompress_symmetric!(A, S, B, color)
+end
+
+function decompress_symmetric(
+    S::AbstractMatrix{Bool},
+    B::AbstractMatrix{R},
+    color::AbstractVector{<:Integer},
+    star_set::StarSet,
 ) where {R<:Real}
     A = respectful_similar(S, R)
-    return decompress_symmetric!(A, S, B, colors)
+    return decompress_symmetric!(A, S, B, color, star_set)
 end

test/random.jl

@@ -7,6 +7,7 @@ using SparseMatrixColorings:
     GreedyColoringAlgorithm,
     structurally_orthogonal_columns,
     symmetrically_orthogonal_columns,
+    directly_recoverable_columns,
     matrix_versions
 using StableRNGs
 using Test
@@ -72,7 +73,9 @@ end;
         A0 = Symmetric(sprand(rng, Bool, n, n, p))
         S0 = map(!iszero, A0)
         @testset "A::$(typeof(A))" for A in matrix_versions(A0)
-            color = symmetric_coloring(A, algo)
+            # Naive decompression
+            color, star_set = symmetric_coloring_detailed(A, algo)
+            @test color == symmetric_coloring(A, algo)
             @test symmetrically_orthogonal_columns(A, color)
             @test directly_recoverable_columns(A, color)
             group = color_groups(color)
@@ -81,6 +84,7 @@ end;
             end
             @testset "S::$(typeof(S))" for S in matrix_versions(S0)
                 @test decompress_symmetric(S, B, color) == A
+                @test decompress_symmetric(S, B, color, star_set) == A
             end
         end
     end