Color by decreasing vertex degree (#262)

* Color vertices in descending degree order

* More docstrings

* Fix color grouping

* Uselesss tests

DifferentiationInterface/src/sparse/coloring.jl

@@ -2,10 +2,21 @@
 Everything in this file is taken from "What color is your Jacobian?"
 =#
 
-function get_groups(colors::AbstractVector{<:Integer})
-    return map(unique(colors)) do c
-        filter(j -> colors[j] == c, eachindex(colors))
+"""
+    color_groups(colors)
+
+Return `groups::Vector{Vector{Int}}` such that `i ∈ groups[c]` iff `colors[i] == c`.
+
+Assumes the colors are contiguously numbered from `1` to some `cmax`.
+"""
+function color_groups(colors::AbstractVector{<:Integer})
+    cmin, cmax = extrema(colors)
+    @assert cmin == 1
+    groups = [Int[] for c in 1:cmax]
+    for (k, c) in enumerate(colors)
+        push!(groups[c], k)
     end
+    return groups
 end
 
 abstract type AbstractMatrixGraph end
@@ -49,15 +60,15 @@ function distance2_column_coloring(g::BipartiteGraph)
     n = length(columns(g))
     colors = zeros(Int, n)
     forbidden_colors = zeros(Int, n)
-    for v in columns(g)  # default ordering
+    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 columns(g)
+        for c in eachindex(forbidden_colors)
             if forbidden_colors[c] != v
                 colors[v] = c
                 break
@@ -71,15 +82,15 @@ function distance2_row_coloring(g::BipartiteGraph)
     m = length(rows(g))
     colors = zeros(Int, m)
     forbidden_colors = zeros(Int, m)
-    for v in 1:m  # default ordering
+    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 rows(g)
+        for c in eachindex(forbidden_colors)
             if forbidden_colors[c] != v
                 colors[v] = c
                 break
@@ -129,7 +140,7 @@ function star_coloring(g::AdjacencyGraph)
     n = length(columns(g))
     colors = zeros(Int, n)
     forbidden_colors = zeros(Int, n)
-    for v in columns(g)  # default ordering
+    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
@@ -149,7 +160,7 @@ function star_coloring(g::AdjacencyGraph)
                 end
             end
         end
-        for c in columns(g)
+        for c in eachindex(forbidden_colors)
             if forbidden_colors[c] != v
                 colors[v] = c
                 break
@@ -164,33 +175,36 @@ end
 """
     GreedyColoringAlgorithm <: ADTypes.AbstractColoringAlgorithm
 
-Matrix coloring algorithm for sparse Jacobians and Hessians.
+Matrix coloring algorithm for sparse Jacobians and Hessians, in which vertices are colored sequentially by order of decreasing degree.
 
 Compatible with the [ADTypes.jl coloring framework](https://sciml.github.io/ADTypes.jl/stable/#Coloring-algorithm).
 
-# See also
+# Implements
+
+- [`ADTypes.column_coloring`](@extref ADTypes) with a partial distance-2 coloring of the bipartite graph
+- [`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
 
-- `ADTypes.column_coloring`
-- `ADTypes.row_coloring`
-- `ADTypes.symmetric_coloring`
+!!! warning
+    Symmetric coloring is not used by DifferentiationInterface.jl at the moment: Hessians are colored by columns just like Jacobians.
 
 # Reference
 
 > [What Color Is Your Jacobian? Graph Coloring for Computing Derivatives](https://epubs.siam.org/doi/abs/10.1137/S0036144504444711), Gebremedhin et al. (2005)
 """
 struct GreedyColoringAlgorithm <: ADTypes.AbstractColoringAlgorithm end
 
-function ADTypes.column_coloring(A, ::GreedyColoringAlgorithm)
+function ADTypes.column_coloring(A::AbstractMatrix, ::GreedyColoringAlgorithm)
     g = BipartiteGraph(A)
     return distance2_column_coloring(g)
 end
 
-function ADTypes.row_coloring(A, ::GreedyColoringAlgorithm)
+function ADTypes.row_coloring(A::AbstractMatrix, ::GreedyColoringAlgorithm)
     g = BipartiteGraph(A)
     return distance2_row_coloring(g)
 end
 
-function ADTypes.symmetric_coloring(A, ::GreedyColoringAlgorithm)
+function ADTypes.symmetric_coloring(A::AbstractMatrix, ::GreedyColoringAlgorithm)
     g = AdjacencyGraph(A)
     return star_coloring(g)
 end

DifferentiationInterface/src/sparse/detector.jl

@@ -6,12 +6,12 @@ Sparsity detection algorithm based on the [Symbolics.jl tracing system](https://
 Compatible with the [ADTypes.jl sparsity detection framework](https://sciml.github.io/ADTypes.jl/stable/#Sparsity-detector).
 
 !!! danger
-    This functionality is implemented in an extension, and requires [Symbolics.jl](https://github.com/JuliaSymbolics/Symbolics.jl) to be loaded.
+    This functionality is in a package extension, and requires [Symbolics.jl](https://github.com/JuliaSymbolics/Symbolics.jl) to be loaded.
 
-# See also
+# Implements
 
-- `ADTypes.jacobian_sparsity`
-- `ADTypes.hessian_sparsity`
+- [`ADTypes.jacobian_sparsity`](@extref ADTypes)
+- [`ADTypes.hessian_sparsity`](@extref ADTypes)
 
 # Reference
 

DifferentiationInterface/src/sparse/hessian.jl

@@ -12,8 +12,8 @@ end
 function prepare_hessian(f::F, backend::AutoSparse, x) where {F}
     initial_sparsity = hessian_sparsity(f, x, sparsity_detector(backend))
     sparsity = col_major(initial_sparsity)
-    colors = column_coloring(sparsity, coloring_algorithm(backend))
-    groups = get_groups(colors)
+    colors = column_coloring(sparsity, coloring_algorithm(backend))  # no star coloring
+    groups = color_groups(colors)
     seeds = map(groups) do group
         seed = zero(x)
         seed[group] .= one(eltype(x))

DifferentiationInterface/src/sparse/jacobian.jl

@@ -27,7 +27,7 @@ function prepare_jacobian(f::F, backend::AutoSparse, x) where {F}
     if Bool(pushforward_performance(backend))
         sparsity = col_major(initial_sparsity)
         colors = column_coloring(sparsity, coloring_algorithm(backend))
-        groups = get_groups(colors)
+        groups = color_groups(colors)
         seeds = map(groups) do group
             seed = zero(x)
             seed[group] .= one(eltype(x))
@@ -42,7 +42,7 @@ function prepare_jacobian(f::F, backend::AutoSparse, x) where {F}
     else
         sparsity = row_major(initial_sparsity)
         colors = row_coloring(sparsity, coloring_algorithm(backend))
-        groups = get_groups(colors)
+        groups = color_groups(colors)
         seeds = map(groups) do group
             seed = zero(y)
             seed[group] .= one(eltype(y))
@@ -110,7 +110,7 @@ function prepare_jacobian(f!::F, y, backend::AutoSparse, x) where {F}
     if Bool(pushforward_performance(backend))
         sparsity = col_major(initial_sparsity)
         colors = column_coloring(sparsity, coloring_algorithm(backend))
-        groups = get_groups(colors)
+        groups = color_groups(colors)
         seeds = map(groups) do group
             seed = zero(x)
             seed[group] .= one(eltype(x))
@@ -125,7 +125,7 @@ function prepare_jacobian(f!::F, y, backend::AutoSparse, x) where {F}
     else
         sparsity = row_major(initial_sparsity)
         colors = row_coloring(sparsity, coloring_algorithm(backend))
-        groups = get_groups(colors)
+        groups = color_groups(colors)
         seeds = map(groups) do group
             seed = zero(y)
             seed[group] .= one(eltype(y))

DifferentiationInterfaceTest/test/coloring.jl

@@ -3,20 +3,41 @@ import DifferentiationInterface as DI
 import DifferentiationInterfaceTest as DIT
 using LinearAlgebra: I, Symmetric
 using SparseArrays: sprand
+using Test
 
 alg = DI.GreedyColoringAlgorithm()
 
-A = sprand(Bool, 100, 200, 0.05)
+@testset "Grouping" begin
+    colors = [1, 3, 1, 3, 1, 2]
+    @test DI.color_groups(colors) == [[1, 3, 5], [6], [2, 4]]
+    colors = [2, 3, 2, 3, 2, 1]
+    @test DI.color_groups(colors) == [[6], [1, 3, 5], [2, 4]]
+    colors = [2, 3, 2, 3, 2]
+    @test_throws AssertionError DI.color_groups(colors)
+end
 
-column_colors = ADTypes.column_coloring(A, alg)
-@test DIT.check_structurally_orthogonal_columns(A, column_colors)
-@test maximum(column_colors) < size(A, 2) ÷ 2
+@testset "Column coloring" begin
+    for A in (sprand(Bool, 100, 200, 0.05), sprand(Bool, 200, 100, 0.05))
+        column_colors = ADTypes.column_coloring(A, alg)
+        @test DIT.check_structurally_orthogonal_columns(A, column_colors)
+        @test minimum(column_colors) == 1
+        @test maximum(column_colors) < size(A, 2) ÷ 2
+    end
+end
 
-row_colors = ADTypes.row_coloring(A, alg)
-@test DIT.check_structurally_orthogonal_rows(A, row_colors)
-@test maximum(row_colors) < size(A, 1) ÷ 2
+@testset "Row coloring" begin
+    for A in (sprand(Bool, 100, 200, 0.05), sprand(Bool, 200, 100, 0.05))
+        row_colors = ADTypes.row_coloring(A, alg)
+        @test DIT.check_structurally_orthogonal_rows(A, row_colors)
+        @test minimum(row_colors) == 1
+        @test maximum(row_colors) < size(A, 1) ÷ 2
+    end
+end
 
-S = Symmetric(sprand(Bool, 100, 100, 0.05)) + I
-symmetric_colors = ADTypes.symmetric_coloring(S, alg)
-@test DIT.check_symmetrically_structurally_orthogonal(S, symmetric_colors)
-@test maximum(symmetric_colors) < size(A, 2) ÷ 2
+@testset "Symmetric coloring" begin
+    S = Symmetric(sprand(Bool, 100, 100, 0.05)) + I
+    symmetric_colors = ADTypes.symmetric_coloring(S, alg)
+    @test DIT.check_symmetrically_structurally_orthogonal(S, symmetric_colors)
+    @test minimum(symmetric_colors) == 1
+    @test maximum(symmetric_colors) < size(S, 2) ÷ 2
+end

DifferentiationInterfaceTest/test/forwarddiff.jl

@@ -1,6 +1,14 @@
+using ADTypes
 using DifferentiationInterface
 using DifferentiationInterfaceTest
 using ForwardDiff: ForwardDiff
+using SparseConnectivityTracer
+
+function MyAutoSparse(backend::AbstractADType)
+    coloring_algorithm = GreedyColoringAlgorithm()
+    sparsity_detector = TracerSparsityDetector()
+    return AutoSparse(backend; sparsity_detector, coloring_algorithm)
+end
 
 test_differentiation(AutoForwardDiff(); logging=get(ENV, "CI", "false") == "false")
 

DifferentiationInterfaceTest/test/runtests.jl

@@ -10,12 +10,6 @@ if isdir(DI_PATH)
     Pkg.develop(; path=DI_PATH)
 end
 
-function MyAutoSparse(backend::AbstractADType)
-    coloring_algorithm = GreedyColoringAlgorithm()
-    sparsity_detector = TracerSparsityDetector()
-    return AutoSparse(backend; sparsity_detector, coloring_algorithm)
-end
-
 ## Main tests
 
 @testset verbose = true "DifferentiationInterfaceTest.jl" begin