Start working on optimal coloring

Author: gdalle

Project.toml

@@ -15,11 +15,14 @@ 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"

ext/SparseMatrixColoringsJuMPExt.jl

@@ -0,0 +1,41 @@
+module SparseMatrixColoringsJuMPExt
+
+using ADTypes: ADTypes
+using JuMP
+using LinearAlgebra
+import MathOptInterface as MOI
+using SparseArrays
+using SparseMatrixColorings:
+    BipartiteGraph, OptimalColoringAlgorithm, nb_vertices, neighbors, pattern, vertices
+
+function optimal_distance2_coloring(
+    bg::BipartiteGraph, ::Val{side}, optimizer::Type{O}
+) where {side,O<:MOI.AbstractOptimizer}
+    other_side = 3 - side
+    n = nb_vertices(bg, Val(side))
+    model = Model(optimizer)
+    set_silent(model)
+    @variable(model, 1 <= color[i=1:n] <= i, Int)
+    @variable(model, ncolors, Int)
+    @constraint(model, [ncolors; color] in MOI.CountDistinct(n + 1))
+    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
+    @objective(model, Min, ncolors)
+    optimize!(model)
+    assert_is_solved_and_feasible(model)
+    return round.(Int, value.(color))
+end
+
+function ADTypes.column_coloring(A::AbstractMatrix, algo::OptimalColoringAlgorithm)
+    bg = BipartiteGraph(A)
+    return optimal_distance2_coloring(bg, Val(2), algo.optimizer)
+end
+
+function ADTypes.row_coloring(A::AbstractMatrix, algo::OptimalColoringAlgorithm)
+    bg = BipartiteGraph(A)
+    return optimal_distance2_coloring(bg, Val(1), algo.optimizer)
+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,18 @@
+"""
+    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.
+
+# Constructor
+
+    OptimalColoringAlgorithm(optimizer)
+
+The `optimizer` argument can be any JuMP-compatible optimizer, like `HiGHS.Optimizer`.
+You can use [`optimizer_with_attributes`](https://jump.dev/JuMP.jl/stable/api/JuMP/#optimizer_with_attributes) to set solver-specific parameters.
+"""
+struct OptimalColoringAlgorithm{S} <: ADTypes.AbstractColoringAlgorithm
+    optimizer::S
+end

test/Project.toml

@@ -12,7 +12,9 @@ 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"

test/optimal.jl

@@ -0,0 +1,37 @@
+using SparseArrays
+using SparseMatrixColorings
+using StableRNGs
+using Test
+
+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)],
+    [(100, 200, p) for p in (0.01:0.01:0.05)],
+    [(200, 100, p) for p in (0.01:0.01:0.05)],
+)
+
+@testset "Column coloring" begin
+    problem = ColoringProblem(; structure=:nonsymmetric, partition=:column)
+    algo = GreedyColoringAlgorithm()
+    optalgo = OptimalColoringAlgorithm(HiGHS.Optimizer)
+    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)
+    algo = GreedyColoringAlgorithm()
+    optalgo = OptimalColoringAlgorithm(HiGHS.Optimizer)
+    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

test/runtests.jl

@@ -58,6 +58,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