Optimal coloring algorithm

.github/workflows/Test.yml

@@ -20,7 +20,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        julia-version: ['1.10', '1']
+        julia-version: ['1.10', '1.11']
 
     steps:
       - uses: actions/checkout@v5
@@ -40,4 +40,4 @@ jobs:
         with:
           files: lcov.info
           token: ${{ secrets.CODECOV_TOKEN }}
-          fail_ci_if_error: false
+          fail_ci_if_error: false

Project.toml

@@ -15,19 +15,24 @@ SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 CliqueTrees = "60701a23-6482-424a-84db-faee86b9b1f8"
 Colors = "5ae59095-9a9b-59fe-a467-6f913c188581"
+JuMP = "4076af6c-e467-56ae-b986-b466b2749572"
+MathOptInterface = "b8f27783-ece8-5eb3-8dc8-9495eed66fee"
 
 [extensions]
 SparseMatrixColoringsCUDAExt = "CUDA"
 SparseMatrixColoringsCliqueTreesExt = "CliqueTrees"
 SparseMatrixColoringsColorsExt = "Colors"
+SparseMatrixColoringsJuMPExt = ["JuMP", "MathOptInterface"]
 
 [compat]
 ADTypes = "1.2.1"
 CUDA = "5.8.2"
 CliqueTrees = "1"
 Colors = "0.12.11, 0.13"
 DocStringExtensions = "0.8,0.9"
+JuMP = "1.29.1"
 LinearAlgebra = "<0.0.1, 1"
+MathOptInterface = "1.45.0"
 PrecompileTools = "1.2.1"
 Random = "<0.0.1, 1"
 SparseArrays = "<0.0.1, 1"

docs/src/api.md

@@ -17,8 +17,14 @@ SparseMatrixColorings
 coloring
 fast_coloring
 ColoringProblem
+```
+
+## Coloring algorithms
+
+```@docs
 GreedyColoringAlgorithm
 ConstantColoringAlgorithm
+OptimalColoringAlgorithm
 ```
 
 ## Result analysis

ext/SparseMatrixColoringsJuMPExt.jl

@@ -0,0 +1,80 @@
+module SparseMatrixColoringsJuMPExt
+
+using ADTypes: ADTypes
+using JuMP:
+    Model,
+    is_solved_and_feasible,
+    optimize!,
+    primal_status,
+    set_silent,
+    set_start_value,
+    value,
+    @variable,
+    @constraint,
+    @objective
+using JuMP
+import MathOptInterface as MOI
+using SparseMatrixColorings:
+    BipartiteGraph, OptimalColoringAlgorithm, nb_vertices, neighbors, pattern, vertices
+
+function optimal_distance2_coloring(
+    bg::BipartiteGraph,
+    ::Val{side},
+    optimizer::O;
+    silent::Bool=true,
+    assert_solved::Bool=true,
+) where {side,O}
+    other_side = 3 - side
+    n = nb_vertices(bg, Val(side))
+    model = Model(optimizer)
+    silent && set_silent(model)
+    # one variable per vertex to color, removing some renumbering symmetries
+    @variable(model, 1 <= color[i=1:n] <= i, Int)
+    # one variable to count the number of distinct colors
+    @variable(model, ncolors, Int)
+    @constraint(model, [ncolors; color] in MOI.CountDistinct(n + 1))
+    # distance-2 coloring: neighbors of the same vertex must have distinct colors
+    for i in vertices(bg, Val(other_side))
+        neigh = neighbors(bg, Val(other_side), i)
+        @constraint(model, color[neigh] in MOI.AllDifferent(length(neigh)))
+    end
+    # minimize the number of distinct colors (can't use maximum because they are not necessarily numbered contiguously)
+    @objective(model, Min, ncolors)
+    # actual solving step where time is spent
+    optimize!(model)
+    if assert_solved
+        # assert feasibility and optimality
+        @assert is_solved_and_feasible(model)
+    else
+        # only assert feasibility
+        @assert primal_status(model) == MOI.FEASIBLE_POINT
+    end
+    # native solver solutions are floating point numbers
+    color_int = round.(Int, value.(color))
+    # remap to 1:cmax in case they are not contiguous
+    true_ncolors = 0
+    remap = fill(0, maximum(color_int))
+    for c in color_int
+        if remap[c] == 0
+            true_ncolors += 1
+            remap[c] = true_ncolors
+        end
+    end
+    return remap[color_int]
+end
+
+function ADTypes.column_coloring(A::AbstractMatrix, algo::OptimalColoringAlgorithm)
+    bg = BipartiteGraph(A)
+    return optimal_distance2_coloring(
+        bg, Val(2), algo.optimizer; algo.silent, algo.assert_solved
+    )
+end
+
+function ADTypes.row_coloring(A::AbstractMatrix, algo::OptimalColoringAlgorithm)
+    bg = BipartiteGraph(A)
+    return optimal_distance2_coloring(
+        bg, Val(1), algo.optimizer; algo.silent, algo.assert_solved
+    )
+end
+
+end

src/SparseMatrixColorings.jl

@@ -56,6 +56,7 @@ include("decompression.jl")
 include("check.jl")
 include("examples.jl")
 include("show_colors.jl")
+include("optimal.jl")
 
 include("precompile.jl")
 
@@ -64,6 +65,7 @@ export DynamicDegreeBasedOrder, SmallestLast, IncidenceDegree, DynamicLargestFir
 export PerfectEliminationOrder
 export ColoringProblem, GreedyColoringAlgorithm, AbstractColoringResult
 export ConstantColoringAlgorithm
+export OptimalColoringAlgorithm
 export coloring, fast_coloring
 export column_colors, row_colors, ncolors
 export column_groups, row_groups

src/adtypes.jl

@@ -1,21 +1,33 @@
-function coloring(
-    A::AbstractMatrix,
-    ::ColoringProblem{:nonsymmetric,:column},
-    algo::ADTypes.NoColoringAlgorithm;
-    kwargs...,
-)
-    bg = BipartiteGraph(A)
-    color = convert(Vector{eltype(bg)}, ADTypes.column_coloring(A, algo))
-    return ColumnColoringResult(A, bg, color)
-end
+## From ADTypes to SMC
 
 function coloring(
     A::AbstractMatrix,
-    ::ColoringProblem{:nonsymmetric,:row},
-    algo::ADTypes.NoColoringAlgorithm;
-    kwargs...,
-)
-    bg = BipartiteGraph(A)
-    color = convert(Vector{eltype(bg)}, ADTypes.row_coloring(A, algo))
-    return RowColoringResult(A, bg, color)
+    problem::ColoringProblem{structure,partition},
+    algo::ADTypes.AbstractColoringAlgorithm;
+    decompression_eltype::Type{R}=Float64,
+    symmetric_pattern::Bool=false,
+) where {structure,partition,R}
+    symmetric_pattern = symmetric_pattern || A isa Union{Symmetric,Hermitian}
+    if structure == :nonsymmetric
+        if partition == :column
+            forced_colors = ADTypes.column_coloring(A, algo)
+        elseif partition == :row
+            forced_colors = ADTypes.row_coloring(A, algo)
+        else
+            # TODO: improve once https://github.com/SciML/ADTypes.jl/issues/69 is done
+            A_and_Aᵀ, _ = bidirectional_pattern(A; symmetric_pattern)
+            forced_colors = ADTypes.symmetric_coloring(A_and_Aᵀ, algo)
+        end
+    else
+        forced_colors = ADTypes.symmetric_coloring(A, algo)
+    end
+    return _coloring(
+        WithResult(),
+        A,
+        problem,
+        GreedyColoringAlgorithm(),
+        R,
+        symmetric_pattern;
+        forced_colors,
+    )
 end

src/coloring.jl

@@ -1,12 +1,19 @@
+struct InvalidColoringError <: Exception end
+
 """
-    partial_distance2_coloring(bg::BipartiteGraph, ::Val{side}, vertices_in_order::AbstractVector)
+    partial_distance2_coloring(
+        bg::BipartiteGraph, ::Val{side}, vertices_in_order::AbstractVector;
+        forced_colors::Union{AbstractVector{<:Integer},Nothing}=nothing
+    )
 
 Compute a distance-2 coloring of the given `side` (`1` or `2`) in the bipartite graph `bg` 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 vertices are colored in a greedy fashion, following the order supplied.
 
+The optional `forced_colors` keyword argument is used to enforce predefined vertex colors (e.g. coming from another optimization algorithm) but still run the distance-2 coloring procedure to verify correctness.
+
 # See also
 
 - [`BipartiteGraph`](@ref)
@@ -17,11 +24,16 @@ The vertices are colored in a greedy fashion, following the order supplied.
 > [_What Color Is Your Jacobian? Graph Coloring for Computing Derivatives_](https://epubs.siam.org/doi/10.1137/S0036144504444711), Gebremedhin et al. (2005), Algorithm 3.2
 """
 function partial_distance2_coloring(
-    bg::BipartiteGraph{T}, ::Val{side}, vertices_in_order::AbstractVector{<:Integer}
+    bg::BipartiteGraph{T},
+    ::Val{side},
+    vertices_in_order::AbstractVector{<:Integer};
+    forced_colors::Union{AbstractVector{<:Integer},Nothing}=nothing,
 ) where {T,side}
     color = Vector{T}(undef, nb_vertices(bg, Val(side)))
     forbidden_colors = Vector{T}(undef, nb_vertices(bg, Val(side)))
-    partial_distance2_coloring!(color, forbidden_colors, bg, Val(side), vertices_in_order)
+    partial_distance2_coloring!(
+        color, forbidden_colors, bg, Val(side), vertices_in_order; forced_colors
+    )
     return color
 end
 
@@ -30,7 +42,8 @@ function partial_distance2_coloring!(
     forbidden_colors::AbstractVector{<:Integer},
     bg::BipartiteGraph,
     ::Val{side},
-    vertices_in_order::AbstractVector{<:Integer},
+    vertices_in_order::AbstractVector{<:Integer};
+    forced_colors::Union{AbstractVector{<:Integer},Nothing}=nothing,
 ) where {side}
     color .= 0
     forbidden_colors .= 0
@@ -44,17 +57,32 @@ function partial_distance2_coloring!(
                 end
             end
         end
-        for i in eachindex(forbidden_colors)
-            if forbidden_colors[i] != v
-                color[v] = i
-                break
+        if isnothing(forced_colors)
+            for i in eachindex(forbidden_colors)
+                if forbidden_colors[i] != v
+                    color[v] = i
+                    break
+                end
+            end
+        else
+            f = forced_colors[v]
+            if (
+                (f == 0 && length(neighbors(bg, Val(side), v)) > 0) ||
+                (f > 0 && forbidden_colors[f] == v)
+            )
+                throw(InvalidColoringError())
+            else
+                color[v] = f
             end
         end
     end
 end
 
 """
-    star_coloring(g::AdjacencyGraph, vertices_in_order::AbstractVector, postprocessing::Bool)
+    star_coloring(
+        g::AdjacencyGraph, vertices_in_order::AbstractVector, postprocessing::Bool;
+        forced_colors::Union{AbstractVector,Nothing}=nothing
+    )
 
 Compute a star coloring of all vertices in the adjacency graph `g` and return a tuple `(color, star_set)`, where
 
@@ -67,6 +95,8 @@ The vertices are colored in a greedy fashion, following the order supplied.
 
 If `postprocessing=true`, some colors might be replaced with `0` (the "neutral" color) as long as they are not needed during decompression.
 
+The optional `forced_colors` keyword argument is used to enforce predefined vertex colors (e.g. coming from another optimization algorithm) but still run the star coloring procedure to verify correctness and build auxiliary data structures, useful during decompression.
+
 # See also
 
 - [`AdjacencyGraph`](@ref)
@@ -77,7 +107,10 @@ If `postprocessing=true`, some colors might be replaced with `0` (the "neutral"
 > [_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::AdjacencyGraph{T}, vertices_in_order::AbstractVector{<:Integer}, postprocessing::Bool
+    g::AdjacencyGraph{T},
+    vertices_in_order::AbstractVector{<:Integer},
+    postprocessing::Bool;
+    forced_colors::Union{AbstractVector{<:Integer},Nothing}=nothing,
 ) where {T<:Integer}
     # Initialize data structures
     nv = nb_vertices(g)
@@ -115,10 +148,18 @@ function star_coloring(
                 end
             end
         end
-        for i in eachindex(forbidden_colors)
-            if forbidden_colors[i] != v
-                color[v] = i
-                break
+        if isnothing(forced_colors)
+            for i in eachindex(forbidden_colors)
+                if forbidden_colors[i] != v
+                    color[v] = i
+                    break
+                end
+            end
+        else
+            if forbidden_colors[forced_colors[v]] == v  # TODO: handle forced_colors[v] == 0
+                throw(InvalidColoringError())
+            else
+                color[v] = forced_colors[v]
             end
         end
         _update_stars!(star, hub, g, v, color, first_neighbor)
@@ -209,7 +250,10 @@ struct StarSet{T}
 end
 
 """
-    acyclic_coloring(g::AdjacencyGraph, vertices_in_order::AbstractVector, postprocessing::Bool)
+    acyclic_coloring(
+        g::AdjacencyGraph, vertices_in_order::AbstractVector, postprocessing::Bool;
+        forced_colors::Union{AbstractVector,Nothing}=nothing
+    )
 
 Compute an acyclic coloring of all vertices in the adjacency graph `g` and return a tuple `(color, tree_set)`, where
 
@@ -222,6 +266,8 @@ The vertices are colored in a greedy fashion, following the order supplied.
 
 If `postprocessing=true`, some colors might be replaced with `0` (the "neutral" color) as long as they are not needed during decompression.
 
+The optional `forced_colors` keyword argument is used to enforce predefined vertex colors (e.g. coming from another optimization algorithm) but still run the acyclic coloring procedure to verify correctness and build auxiliary data structures, useful during decompression.
+
 # See also
 
 - [`AdjacencyGraph`](@ref)
@@ -232,7 +278,10 @@ If `postprocessing=true`, some colors might be replaced with `0` (the "neutral"
 > [_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 3.1
 """
 function acyclic_coloring(
-    g::AdjacencyGraph{T}, vertices_in_order::AbstractVector{<:Integer}, postprocessing::Bool
+    g::AdjacencyGraph{T},
+    vertices_in_order::AbstractVector{<:Integer},
+    postprocessing::Bool;
+    forced_colors::Union{AbstractVector{<:Integer},Nothing}=nothing,
 ) where {T<:Integer}
     # Initialize data structures
     nv = nb_vertices(g)
@@ -271,10 +320,18 @@ function acyclic_coloring(
                 end
             end
         end
-        for i in eachindex(forbidden_colors)
-            if forbidden_colors[i] != v
-                color[v] = i
-                break
+        if isnothing(forced_colors)
+            for i in eachindex(forbidden_colors)
+                if forbidden_colors[i] != v
+                    color[v] = i
+                    break
+                end
+            end
+        else
+            if forbidden_colors[forced_colors[v]] == v  # TODO: handle forced_colors[v] == 0
+                throw(InvalidColoringError())
+            else
+                color[v] = forced_colors[v]
             end
         end
         for (w, index_vw) in neighbors_with_edge_indices(g, v)  # grow two-colored stars around the vertex v