Working star coloring (#11)

* Working star coloring

* Fix verbosity

* More efficient recovery

Author: gdalle

README.md

@@ -18,11 +18,12 @@ pkg> add https://github.com/gdalle/SparseMatrixColorings.jl
 
 ## Background
 
-The algorithms implemented in this package are mainly taken from the following articles:
+The algorithms implemented in this package are taken from the following articles:
 
-> [_What Color Is Your Jacobian? Graph Coloring for Computing Derivatives_](https://epubs.siam.org/doi/10.1137/S0036144504444711), Gebremedhin et al. (2005)
-
-> [_ColPack: Software for graph coloring and related problems in scientific computing_](https://dl.acm.org/doi/10.1145/2513109.2513110), Gebremedhin et al. (2013)
+- [_What Color Is Your Jacobian? Graph Coloring for Computing Derivatives_](https://epubs.siam.org/doi/10.1137/S0036144504444711), Gebremedhin et al. (2005)
+- [_New Acyclic and Star Coloring Algorithms with Application to Computing Hessians_](https://epubs.siam.org/doi/abs/10.1137/050639879), Gebremedhin et al. (2007)
+- [_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)
+- [_ColPack: Software for graph coloring and related problems in scientific computing_](https://dl.acm.org/doi/10.1145/2513109.2513110), Gebremedhin et al. (2013)
 
 Some parts of the articles (like definitions) are thus copied verbatim in the documentation.
 

docs/src/api.md

@@ -10,6 +10,9 @@ CurrentModule = SparseMatrixColorings
 ```@docs
 SparseMatrixColorings
 GreedyColoringAlgorithm
+column_coloring
+row_coloring
+symmetric_coloring
 ```
 
 ## Public, not exported
@@ -60,13 +63,12 @@ vertices
 
 ```@docs
 partial_distance2_coloring
-star_coloring1
+star_coloring
 ```
 
 ### Testing
 
 ```@docs
 check_structurally_orthogonal_columns
-check_structurally_orthogonal_rows
-check_symmetrically_orthogonal
+check_symmetrically_orthogonal_columns
 ```

src/SparseMatrixColorings.jl

@@ -5,7 +5,8 @@ $README
 """
 module SparseMatrixColorings
 
-using ADTypes: ADTypes, AbstractColoringAlgorithm
+using ADTypes:
+    ADTypes, AbstractColoringAlgorithm, column_coloring, row_coloring, symmetric_coloring
 using Compat: @compat
 using DocStringExtensions: README
 using LinearAlgebra:
@@ -46,7 +47,9 @@ include("check.jl")
 @compat public decompress_columns, decompress_columns!
 @compat public decompress_rows, decompress_rows!
 @compat public decompress_symmetric, decompress_symmetric!
+@compat public column_coloring, row_coloring, symmetric_coloring
 
 export GreedyColoringAlgorithm
+export column_coloring, row_coloring, symmetric_coloring
 
 end

src/adtypes.jl

@@ -9,24 +9,42 @@ Compatible with the [ADTypes.jl coloring framework](https://sciml.github.io/ADTy
 
     GreedyColoringAlgorithm(order::AbstractOrder=NaturalOrder())
 
-# Implements
+# See also
+
+- [`AbstractOrder`](@ref)
+"""
+struct GreedyColoringAlgorithm{O<:AbstractOrder} <: ADTypes.AbstractColoringAlgorithm
+    order::O
+end
+
+GreedyColoringAlgorithm() = GreedyColoringAlgorithm(NaturalOrder())
+
+function Base.show(io::IO, algo::GreedyColoringAlgorithm)
+    return print(io, "GreedyColoringAlgorithm($(algo.order))")
+end
+
+"""
+    column_coloring(A::AbstractMatrix, algo::GreedyColoringAlgorithm)
+
+Compute a partial distance-2 coloring of the columns in the bipartite graph of the matrix `A`.
 
-- [`ADTypes.column_coloring`](@extref ADTypes) and [`ADTypes.row_coloring`](@extref ADTypes) with a partial distance-2 coloring of the bipartite graph
-- [`ADTypes.symmetric_coloring`](@extref ADTypes) with a star coloring of the adjacency graph
+Function defined by ADTypes, re-exported by SparseMatrixColorings.
 
-# Example use
+# Example
 
 ```jldoctest
-using ADTypes, SparseMatrixColorings, SparseArrays
+using SparseMatrixColorings, SparseArrays
 
 algo = GreedyColoringAlgorithm(SparseMatrixColorings.LargestFirst())
+
 A = sparse([
     0 0 1 1 0
     1 0 0 0 1
     0 1 1 0 0
     0 1 1 0 1
 ])
-ADTypes.column_coloring(A, algo)
+
+column_coloring(A, algo)
 
 # output
 
@@ -37,32 +55,62 @@ ADTypes.column_coloring(A, algo)
  2
  3
 ```
-
-# See also
-
-- [`AbstractOrder`](@ref)
 """
-struct GreedyColoringAlgorithm{O<:AbstractOrder} <: ADTypes.AbstractColoringAlgorithm
-    order::O
-end
-
-GreedyColoringAlgorithm() = GreedyColoringAlgorithm(NaturalOrder())
-
-function Base.show(io::IO, algo::GreedyColoringAlgorithm)
-    return print(io, "GreedyColoringAlgorithm($(algo.order))")
-end
-
 function ADTypes.column_coloring(A::AbstractMatrix, algo::GreedyColoringAlgorithm)
     bg = bipartite_graph(A)
     return partial_distance2_coloring(bg, Val(2), algo.order)
 end
 
+"""
+    row_coloring(A::AbstractMatrix, algo::GreedyColoringAlgorithm)
+
+Compute a partial distance-2 coloring of the rows in the bipartite graph of the matrix `A`.
+
+Function defined by ADTypes, re-exported by SparseMatrixColorings.
+
+# Example
+
+```jldoctest
+using SparseMatrixColorings, SparseArrays
+
+algo = GreedyColoringAlgorithm(SparseMatrixColorings.LargestFirst())
+
+A = sparse([
+    0 0 1 1 0
+    1 0 0 0 1
+    0 1 1 0 0
+    0 1 1 0 1
+])
+
+row_coloring(A, algo)
+
+# output
+
+4-element Vector{Int64}:
+ 2
+ 2
+ 3
+ 1
+```
+"""
 function ADTypes.row_coloring(A::AbstractMatrix, algo::GreedyColoringAlgorithm)
     bg = bipartite_graph(A)
     return partial_distance2_coloring(bg, Val(1), algo.order)
 end
 
+"""
+    symmetric_coloring(A::AbstractMatrix, algo::GreedyColoringAlgorithm)
+
+Compute a star coloring of the columns in the adjacency graph of the symmetric matrix `A`.
+
+Function defined by ADTypes, re-exported by SparseMatrixColorings.
+
+# Example
+
+!!! warning
+    Work in progress.
+"""
 function ADTypes.symmetric_coloring(A::AbstractMatrix, algo::GreedyColoringAlgorithm)
     ag = adjacency_graph(A)
-    return star_coloring1(ag, algo.order)
+    return star_coloring(ag, algo.order)
 end

src/check.jl

@@ -1,97 +1,95 @@
 """
     check_structurally_orthogonal_columns(
-        A::AbstractMatrix, colors::AbstractVector{<:Integer}
+        A::AbstractMatrix, color::AbstractVector{<:Integer}
         verbose=false
     )
 
-Return `true` if coloring the columns of the matrix `A` with the vector `colors` results in a partition that is structurally orthogonal, and `false` otherwise.
+Return `true` if coloring the columns of the matrix `A` with the vector `color` results in a partition that is structurally orthogonal, and `false` otherwise.
     
 A partition of the columns of a matrix `A` is _structurally orthogonal_ if, for every nonzero element `A[i, j]`, the group containing column `A[:, j]` has no other column with a nonzero in row `i`.
 
 !!! warning
     This function is not coded with efficiency in mind, it is designed for small-scale tests.
+
+# References
+
+> [_What Color Is Your Jacobian? Graph Coloring for Computing Derivatives_](https://epubs.siam.org/doi/10.1137/S0036144504444711), Gebremedhin et al. (2005)
 """
 function check_structurally_orthogonal_columns(
-    A::AbstractMatrix, colors::AbstractVector{<:Integer}; verbose::Bool=true
+    A::AbstractMatrix, color::AbstractVector{<:Integer}; verbose::Bool=false
 )
-    groups = color_groups(colors)
-    for (c, g) in enumerate(groups)
-        Ag = @view A[:, g]
+    if length(color) != size(A, 2)
+        if verbose
+            @warn "$(length(color)) colors provided for $(size(A, 2)) columns"
+        end
+        return false
+    end
+    group = color_groups(color)
+    for (c, g) in enumerate(group)
+        Ag = view(A, :, g)
         nonzeros_per_row = dropdims(count(!iszero, Ag; dims=2); dims=2)
         max_nonzeros_per_row, i = findmax(nonzeros_per_row)
         if max_nonzeros_per_row > 1
-            verbose && @warn "Columns $g (with color $c) share nonzeros in row $i"
-            return false
-        end
-    end
-    return true
-end
-
-"""
-    check_structurally_orthogonal_rows(
-        A::AbstractMatrix, colors::AbstractVector{<:Integer};
-        verbose=false
-    )
-
-Return `true` if coloring the rows of the matrix `A` with the vector `colors` results in a partition that is structurally orthogonal, and `false` otherwise.
-    
-A partition of the rows of a matrix `A` is _structurally orthogonal_ if, for every nonzero element `A[i, j]`, the group containing row `A[i, :]` has no other row with a nonzero in column `j`.
-
-!!! warning
-    This function is not coded with efficiency in mind, it is designed for small-scale tests.
-"""
-function check_structurally_orthogonal_rows(
-    A::AbstractMatrix, colors::AbstractVector{<:Integer}; verbose::Bool=true
-)
-    groups = color_groups(colors)
-    for (c, g) in enumerate(groups)
-        Ag = @view A[g, :]
-        nonzeros_per_col = dropdims(count(!iszero, Ag; dims=1); dims=1)
-        max_nonzeros_per_col, j = findmax(nonzeros_per_col)
-        if max_nonzeros_per_col > 1
-            verbose && @warn "Rows $g (with color $c) share nonzeros in column $j"
+            if verbose
+                incompatible_columns = g[findall(!iszero, view(Ag, i, :))]
+                @warn "In color $c, columns $incompatible_columns all have nonzeros in row $i"
+            end
             return false
         end
     end
     return true
 end
 
 """
-    check_symmetrically_orthogonal(
-        A::AbstractMatrix, colors::AbstractVector{<:Integer};
+    check_symmetrically_orthogonal_columns(
+        A::AbstractMatrix, color::AbstractVector{<:Integer};
         verbose=false
     )
 
-Return `true` if coloring the columns of the symmetric matrix `A` with the vector `colors` results in a partition that is symmetrically orthogonal, and `false` otherwise.
+Return `true` if coloring the columns of the symmetric matrix `A` with the vector `color` results in a partition that is symmetrically orthogonal, and `false` otherwise.
     
-A partition of the columns of a symmetrix matrix `A` is _symmetrically orthogonal_ if, for every nonzero element `A[i, j]`, either
+A partition of the columns of a symmetrix matrix `A` is _symmetrically orthogonal_ if, for every nonzero element `A[i, j]`, either of the following statements holds:
 
 1. the group containing the column `A[:, j]` has no other column with a nonzero in row `i`
 2. the group containing the column `A[:, i]` has no other column with a nonzero in row `j`
 
 !!! warning
     This function is not coded with efficiency in mind, it is designed for small-scale tests.
+
+# References
+
+> [_What Color Is Your Jacobian? Graph Coloring for Computing Derivatives_](https://epubs.siam.org/doi/10.1137/S0036144504444711), Gebremedhin et al. (2005)
 """
-function check_symmetrically_orthogonal(
-    A::AbstractMatrix, colors::AbstractVector{<:Integer}; verbose::Bool=true
+function check_symmetrically_orthogonal_columns(
+    A::AbstractMatrix, color::AbstractVector{<:Integer}; verbose::Bool=false
 )
     checksquare(A)
+    if length(color) != size(A, 2)
+        if verbose
+            @warn "$(length(color)) colors provided for $(size(A, 2)) columns"
+        end
+        return false
+    end
     issymmetric(A) || return false
-    groups = color_groups(colors)
+    group = color_groups(color)
     for i in axes(A, 2), j in axes(A, 2)
         iszero(A[i, j]) && continue
-        ki, kj = colors[i], colors[j]
-        gi, gj = groups[ki], groups[kj]
+        ci, cj = color[i], color[j]
+        gi, gj = group[ci], group[cj]
         A_gj_rowi = view(A, i, gj)
         A_gi_rowj = view(A, j, gi)
         nonzeros_gj_rowi = count(!iszero, A_gj_rowi)
         nonzeros_gi_rowj = count(!iszero, A_gi_rowj)
         if nonzeros_gj_rowi > 1 && nonzeros_gi_rowj > 1
-            verbose && @warn """
-            For coefficient $((i, j)):
-            - columns $gj (with color $kj) share nonzeros in row $i
-            - columns $gi (with color $ki) share nonzeros in row $j
-            """
+            if verbose
+                gj_incompatible_columns = gj[findall(!iszero, A_gj_rowi)]
+                gi_incompatible_columns = gi[findall(!iszero, A_gi_rowj)]
+                @warn """
+                For coefficient (i=$i, j=$j) with column colors (ci=$ci, cj=$cj):
+                - in color ci=$ci, columns $gi_incompatible_columns all have nonzeros in row j=$j
+                - in color cj=$cj, columns $gj_incompatible_columns all have nonzeros in row i=$i
+                """
+            end
             return false
         end
     end