First algorithms

Author: gdalle

Project.toml

@@ -4,10 +4,14 @@ authors = ["Guillaume Dalle <22795598+gdalle@users.noreply.github.com>"]
 version = "0.1.0"
 
 [deps]
+ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
+LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 
 [compat]
+ADTypes = "1.2.1"
+LinearAlgebra = "1"
 Random = "1"
 SparseArrays = "1"
 julia = "1.6"

README.md

@@ -21,7 +21,9 @@ The algorithms implemented in this package are all taken from the following surv
 
 > [_What Color Is Your Jacobian? Graph Coloring for Computing Derivatives_](https://epubs.siam.org/doi/10.1137/S0036144504444711), Gebremedhin et al. (2005)
 
+Some parts of the survey (like definitions and theorems) are also copied verbatim or referred to by their number in the documentation.
+
 ## Alternatives
 
-- [ColPack.jl](https://github.com/michel2323/ColPack.jl)
-- [SparseDiffTools.jl](https://github.com/JuliaDiff/SparseDiffTools.jl)
+- [ColPack.jl](https://github.com/michel2323/ColPack.jl): a Julia interface to the C++ library [ColPack](https://github.com/CSCsw/ColPack)
+- [SparseDiffTools.jl](https://github.com/JuliaDiff/SparseDiffTools.jl): contains Julia implementations of some coloring algorithms

docs/src/api.md

@@ -15,5 +15,5 @@ Private = false
 
 ```@autodocs
 Modules = [SparseMatrixColorings]
-Private = false
+Public = false
 ```

src/SparseMatrixColorings.jl

@@ -1,6 +1,32 @@
+"""
+    SparseMatrixColorings
+
+Coloring algorithms for sparse Jacobian and Hessian matrices.
+
+The algorithms implemented in this package are all taken from the following survey:
+
+> [_What Color Is Your Jacobian? Graph Coloring for Computing Derivatives_](https://epubs.siam.org/doi/10.1137/S0036144504444711), Gebremedhin et al. (2005)
+
+Some parts of the survey (like definitions and theorems) are also copied verbatim or referred to by their number in the documentation.
+"""
 module SparseMatrixColorings
 
+using ADTypes: ADTypes, AbstractColoringAlgorithm
+using LinearAlgebra: Transpose, parent, transpose
 using Random: AbstractRNG
-using SparseArrays: SparseMatrixCSC
+using SparseArrays: SparseMatrixCSC, nzrange, rowvals
+
+include("utils.jl")
+
+include("bipartite_graph.jl")
+include("adjacency_graph.jl")
+
+include("distance2_coloring.jl")
+include("star_coloring.jl")
+
+include("adtypes.jl")
+include("check.jl")
+
+export GreedyColoringAlgorithm
 
 end

src/adjacency_graph.jl

@@ -0,0 +1,26 @@
+"""
+    AdjacencyGraph
+
+Represent a graph between the columns of a symmetric `n × n` matrix `A` with nonzero diagonal elements.
+
+This graph is defined as `G = (C, E)` where `C = 1:n` is the set of columns and `(i, j) ∈ E` whenever `A[i, j]` is nonzero for some `j ∈ 1:m`, `j ≠ i`.
+
+# Fields
+
+- `A_colmajor::AbstractMatrix`: output of [`col_major`](@ref) applied to `A`
+"""
+struct AdjacencyGraph{M<:AbstractMatrix}
+    A_colmajor::M
+
+    function AdjacencyGraph(A::AbstractMatrix)
+        A_colmajor = col_major(A)
+        return new{typeof(A_colmajor)}(A_colmajor)
+    end
+end
+
+rows(g::AdjacencyGraph) = axes(g.A_colmajor, 1)
+columns(g::AdjacencyGraph) = axes(g.A_colmajor, 2)
+
+function neighbors(g::AdjacencyGraph, j::Integer)
+    return filter(!isequal(j), nz_in_col(g.A_colmajor, j))
+end

src/adtypes.jl

@@ -0,0 +1,29 @@
+"""
+    GreedyColoringAlgorithm <: ADTypes.AbstractColoringAlgorithm
+
+Matrix coloring algorithm for sparse Jacobians and Hessians.
+
+Compatible with the [ADTypes.jl coloring framework](https://sciml.github.io/ADTypes.jl/stable/#Coloring-algorithm).
+
+# Implements
+
+- `ADTypes.column_coloring` with a partial distance-2 coloring of the bipartite graph
+- `ADTypes.row_coloring` with a partial distance-2 coloring of the bipartite graph
+- `ADTypes.symmetric_coloring` with a star coloring of the adjacency graph
+"""
+struct GreedyColoringAlgorithm <: ADTypes.AbstractColoringAlgorithm end
+
+function ADTypes.column_coloring(A::AbstractMatrix, ::GreedyColoringAlgorithm)
+    g = BipartiteGraph(A)
+    return distance2_column_coloring(g)
+end
+
+function ADTypes.row_coloring(A::AbstractMatrix, ::GreedyColoringAlgorithm)
+    g = BipartiteGraph(A)
+    return distance2_row_coloring(g)
+end
+
+function ADTypes.symmetric_coloring(A::AbstractMatrix, ::GreedyColoringAlgorithm)
+    g = AdjacencyGraph(A)
+    return star_coloring(g)
+end

src/bipartite_graph.jl

@@ -0,0 +1,28 @@
+"""
+    BipartiteGraph
+
+Represent a bipartite graph between the rows and the columns of a non-symmetric `m × n` matrix `A`.
+
+This graph is defined as `G = (R, C, E)` where `R = 1:m` is the set of row indices, `C = 1:n` is the set of column indices and `(i, j) ∈ E` whenever `A[i, j]` is nonzero.
+
+# Fields
+
+- `A_colmajor::AbstractMatrix`: output of [`col_major`](@ref) applied to `A` (useful to get neighbors of a column)
+- `A_rowmajor::AbstractMatrix`: output of [`row_major`](@ref) applied to `A` (useful to get neighbors of a row)
+"""
+struct BipartiteGraph{M1<:AbstractMatrix,M2<:AbstractMatrix}
+    A_colmajor::M1
+    A_rowmajor::M2
+
+    function BipartiteGraph(A::AbstractMatrix)
+        A_colmajor = col_major(A)
+        A_rowmajor = row_major(A)
+        return new{typeof(A_colmajor),typeof(A_rowmajor)}(A_colmajor, A_rowmajor)
+    end
+end
+
+rows(g::BipartiteGraph) = axes(g.A_colmajor, 1)
+columns(g::BipartiteGraph) = axes(g.A_colmajor, 2)
+
+neighbors_of_column(g::BipartiteGraph, j::Integer) = nz_in_col(g.A_colmajor, j)
+neighbors_of_row(g::BipartiteGraph, i::Integer) = nz_in_row(g.A_rowmajor, i)

src/check.jl

@@ -0,0 +1,81 @@
+"""
+    check_structurally_orthogonal_columns(A, colors; 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.
+    
+Def 3.2: 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`.
+
+Thm 3.5: The function [`distance2_column_coloring`](@ref) applied to the [`BipartiteGraph`](@ref) of `A` should return a suitable coloring.
+"""
+function check_structurally_orthogonal_columns(
+    A::AbstractMatrix, colors::AbstractVector{<:Integer}; verbose=false
+)
+    for c in unique(colors)
+        js = filter(j -> colors[j] == c, axes(A, 2))
+        Ajs = @view A[:, js]
+        nonzeros_per_row = count(!iszero, Ajs; dims=2)
+        if maximum(nonzeros_per_row) > 1
+            verbose && @warn "Color $c has columns $js sharing nonzeros"
+            return false
+        end
+    end
+    return true
+end
+
+"""
+    check_structurally_orthogonal_rows(A, colors; 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.
+    
+Def 3.2: 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`.
+
+Thm 3.5: The function [`distance2_row_coloring`](@ref) applied to the [`BipartiteGraph`](@ref) of `A` should return a suitable coloring.
+"""
+function check_structurally_orthogonal_rows(
+    A::AbstractMatrix, colors::AbstractVector{<:Integer}; verbose=false
+)
+    for c in unique(colors)
+        is = filter(i -> colors[i] == c, axes(A, 1))
+        Ais = @view A[is, :]
+        nonzeros_per_column = count(!iszero, Ais; dims=1)
+        if maximum(nonzeros_per_column) > 1
+            verbose && @warn "Color $c has rows $is sharing nonzeros"
+            return false
+        end
+    end
+    return true
+end
+
+"""
+    check_symmetrically_orthogonal(A, colors; 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.
+    
+Def 4.2: A partition of the columns of a symmetrix matrix `A` is _symmetrically orthogonal_ if, for every nonzero element `A[i, j]`, either
+
+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`
+"""
+function check_symmetrically_orthogonal(
+    A::AbstractMatrix, colors::AbstractVector{<:Integer}; verbose=false
+)
+    for i in axes(A, 2), j in axes(A, 2)
+        if !iszero(A[i, j])
+            group_i = filter(i2 -> (i2 != i) && (colors[i2] == colors[i]), axes(A, 2))
+            group_j = filter(j2 -> (j2 != j) && (colors[j2] == colors[j]), axes(A, 2))
+            A_group_i_column_j = @view A[group_i, j]
+            A_group_j_column_i = @view A[group_j, i]
+            nonzeros_group_i_column_j = count(!iszero, A_group_i_column_j)
+            nonzeros_group_j_column_i = count(!iszero, A_group_j_column_i)
+            if nonzeros_group_i_column_j > 0 && nonzeros_group_j_column_i > 0
+                verbose && @warn """
+                For coefficient $((i, j)), both of the following have confounding zeros:
+                - color $(colors[j]) with group $group_j
+                - color $(colors[i]) with group $group_i
+                """
+                return false
+            end
+        end
+    end
+    return true
+end

src/distance2_coloring.jl

@@ -0,0 +1,61 @@
+"""
+    distance2_column_coloring(g::BipartiteGraph)
+
+Compute a distance-2 coloring of the column vertices in the [`BipartiteGraph`](@ref) `g` and return a vector of integer colors.
+
+A _distance-2 coloring_ is such that two vertices have different colors if they are at distance at most 2.
+
+The algorithm used is the greedy Algorithm 3.2.
+"""
+function distance2_column_coloring(g::BipartiteGraph)
+    n = length(columns(g))
+    colors = zeros(Int, n)
+    forbidden_colors = zeros(Int, n)
+    for v in sort(columns(g); by=j -> length(neighbors_of_column(g, j)), rev=true)
+        for w in neighbors_of_column(g, v)
+            for x in neighbors_of_row(g, w)
+                if !iszero(colors[x])
+                    forbidden_colors[colors[x]] = v
+                end
+            end
+        end
+        for c in eachindex(forbidden_colors)
+            if forbidden_colors[c] != v
+                colors[v] = c
+                break
+            end
+        end
+    end
+    return colors
+end
+
+"""
+    distance2_row_coloring(g::BipartiteGraph)
+
+Compute a distance-2 coloring of the row vertices in the [`BipartiteGraph`](@ref) `g` and return a vector of integer colors.
+
+A _distance-2 coloring_ is such that two vertices have different colors if they are at distance at most 2.
+
+The algorithm used is the greedy Algorithm 3.2.
+"""
+function distance2_row_coloring(g::BipartiteGraph)
+    m = length(rows(g))
+    colors = zeros(Int, m)
+    forbidden_colors = zeros(Int, m)
+    for v in sort(rows(g); by=i -> length(neighbors_of_row(g, i)), rev=true)
+        for w in neighbors_of_row(g, v)
+            for x in neighbors_of_column(g, w)
+                if !iszero(colors[x])
+                    forbidden_colors[colors[x]] = v
+                end
+            end
+        end
+        for c in eachindex(forbidden_colors)
+            if forbidden_colors[c] != v
+                colors[v] = c
+                break
+            end
+        end
+    end
+    return colors
+end

src/star_coloring.jl

@@ -0,0 +1,42 @@
+"""
+    star_coloring(g::BipartiteGraph)
+
+Compute a star coloring of the column vertices in the [`AdjacencyGraph`](@ref) `g` and return a vector of integer colors.
+
+Def 4.5: A _star coloring_ is a distance-1 coloring such that every path on four vertices uses at least three colors.
+
+The algorithm used is the greedy Algorithm 4.1.
+"""
+function star_coloring(g::AdjacencyGraph)
+    n = length(columns(g))
+    colors = zeros(Int, n)
+    forbidden_colors = zeros(Int, n)
+    for v in sort(columns(g); by=j -> length(neighbors(g, j)), rev=true)
+        for w in neighbors(g, v)
+            if !iszero(colors[w])  # w is colored
+                forbidden_colors[colors[w]] = v
+            end
+            for x in neighbors(g, w)
+                if !iszero(colors[x]) && iszero(colors[w])  # w is not colored
+                    forbidden_colors[colors[x]] = v
+                else
+                    for y in neighbors(g, x)
+                        if !iszero(colors[y]) && y != w
+                            if colors[y] == colors[w]
+                                forbidden_colors[colors[x]] = v
+                                break
+                            end
+                        end
+                    end
+                end
+            end
+        end
+        for c in eachindex(forbidden_colors)
+            if forbidden_colors[c] != v
+                colors[v] = c
+                break
+            end
+        end
+    end
+    return colors
+end

src/utils.jl

@@ -0,0 +1,42 @@
+## Conversion between row- and col-major
+
+"""
+    col_major(A::AbstractMatrix)
+
+Construct a column-major representation of the matrix `A`.
+"""
+col_major(A::M) where {M<:AbstractMatrix} = A
+col_major(A::Transpose{<:Any,M}) where {M<:AbstractMatrix} = M(A)
+
+"""
+    row_major(A::AbstractMatrix)
+
+Construct a row-major representation of the matrix `A`.
+"""
+row_major(A::M) where {M<:AbstractMatrix} = transpose(M(transpose(A)))
+row_major(A::Transpose{<:Any,M}) where {M<:AbstractMatrix} = A
+
+## Generic nz
+
+function nz_in_col(A_colmajor::AbstractMatrix, j::Integer)
+    return filter(i -> !iszero(A_colmajor[i, j]), axes(A_colmajor, 1))
+end
+
+function nz_in_row(A_rowmajor::AbstractMatrix, i::Integer)
+    return filter(j -> !iszero(A_rowmajor[i, j]), axes(A_rowmajor, 2))
+end
+
+## Sparse nz
+
+function nz_in_col(A_colmajor::SparseMatrixCSC{T}, j::Integer) where {T}
+    rv = rowvals(A_colmajor)
+    ind = nzrange(A_colmajor, j)
+    return view(rv, ind)
+end
+
+function nz_in_row(A_rowmajor::Transpose{T,<:SparseMatrixCSC{T}}, i::Integer) where {T}
+    A_transpose_colmajor = parent(A_rowmajor)
+    rv = rowvals(A_transpose_colmajor)
+    ind = nzrange(A_transpose_colmajor, i)
+    return view(rv, ind)
+end

test/Project.toml

@@ -1,6 +1,9 @@
 [deps]
+ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
 JuliaFormatter = "98e50ef6-434e-11e9-1051-2b60c6c9e899"
+LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"