Optimal coloring algorithm with JuMP formulation

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/optimal.jl

@@ -0,0 +1,34 @@
+"""
+    OptimalColoringAlgorithm
+
+Coloring algorithm that relies on mathematical programming with [JuMP](https://jump.dev/) to find an optimal coloring.
+
+!!! warning
+    This algorithm is only available when JuMP is loaded. If you encounter a method error, run `import JuMP` in your REPL and try again.
+    It only works for nonsymmetric, unidirectional colorings problems.
+
+!!! danger
+    The coloring problem is NP-hard, so it is unreasonable to expect an optimal solution in reasonable time for large instances.
+
+# Constructor
+
+    OptimalColoringAlgorithm(optimizer; silent::Bool=true, assert_solved::Bool=true)
+
+The `optimizer` argument can be any JuMP-compatible optimizer.
+However, the problem formulation is best suited to CP-SAT optimizers like [MiniZinc](https://github.com/jump-dev/MiniZinc.jl).
+You can use [`optimizer_with_attributes`](https://jump.dev/JuMP.jl/stable/api/JuMP/#optimizer_with_attributes) to set solver-specific parameters.
+
+# Keyword arguments
+
+- `silent`: whether to suppress solver output
+- `assert_solved`: whether to check that the solver found an optimal solution (as opposed to running out of time for example)
+"""
+struct OptimalColoringAlgorithm{O} <: ADTypes.AbstractColoringAlgorithm
+    optimizer::O
+    silent::Bool
+    assert_solved::Bool
+end
+
+function OptimalColoringAlgorithm(optimizer; silent::Bool=true, assert_solved::Bool=true)
+    return OptimalColoringAlgorithm(optimizer, silent, assert_solved)
+end

test/Project.toml

@@ -12,10 +12,13 @@ CliqueTrees = "60701a23-6482-424a-84db-faee86b9b1f8"
 Colors = "5ae59095-9a9b-59fe-a467-6f913c188581"
 DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
+HiGHS = "87dc4568-4c63-4d18-b0c0-bb2238e4078b"
 JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
+JuMP = "4076af6c-e467-56ae-b986-b466b2749572"
 JuliaFormatter = "98e50ef6-434e-11e9-1051-2b60c6c9e899"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MatrixDepot = "b51810bb-c9f3-55da-ae3c-350fc1fbce05"
+MiniZinc = "a7f392d2-6c35-496e-b8cc-0974fbfcbf91"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 SparseMatrixColorings = "0a514795-09f3-496d-8182-132a7b665d35"

test/optimal.jl

@@ -0,0 +1,46 @@
+using SparseArrays
+using SparseMatrixColorings
+using StableRNGs
+using Test
+using JuMP
+using MiniZinc
+using HiGHS
+
+rng = StableRNG(0)
+
+asymmetric_params = vcat(
+    [(10, 20, p) for p in (0.0:0.1:0.5)], [(20, 10, p) for p in (0.0:0.1:0.5)]
+)
+
+algo = GreedyColoringAlgorithm()
+optalgo = OptimalColoringAlgorithm(() -> MiniZinc.Optimizer{Float64}("highs"); silent=false)
+
+@testset "Column coloring" begin
+    problem = ColoringProblem(; structure=:nonsymmetric, partition=:column)
+    for (m, n, p) in asymmetric_params
+        A = sprand(rng, m, n, p)
+        result = coloring(A, problem, algo)
+        optresult = coloring(A, problem, optalgo)
+        @test ncolors(result) >= ncolors(optresult)
+    end
+end
+
+@testset "Row coloring" begin
+    problem = ColoringProblem(; structure=:nonsymmetric, partition=:row)
+    for (m, n, p) in asymmetric_params
+        A = sprand(rng, m, n, p)
+        result = coloring(A, problem, algo)
+        optresult = coloring(A, problem, optalgo)
+        @test ncolors(result) >= ncolors(optresult)
+    end
+end
+
+@testset "Too big" begin
+    A = sprand(rng, Bool, 100, 100, 0.1)
+    optalgo_timelimit = OptimalColoringAlgorithm(
+        optimizer_with_attributes(HiGHS.Optimizer, "time_limit" => 10.0); # 1 second
+        silent=false,
+        assert_solved=false,
+    )
+    @test_throws AssertionError coloring(A, ColoringProblem(), optalgo_timelimit)
+end

test/runtests.jl

@@ -59,6 +59,9 @@ include("utils.jl")
             @testset "Constant coloring" begin
                 include("constant.jl")
             end
+            @testset "Optimal coloring" begin
+                include("optimal.jl")
+            end
             @testset "ADTypes coloring algorithms" begin
                 include("adtypes.jl")
             end