Better comments

Author: gdalle

ext/SparseMatrixColoringsJuMPExt.jl

@@ -3,36 +3,55 @@ module SparseMatrixColoringsJuMPExt
 using ADTypes: ADTypes
 using JuMP:
     Model,
-    assert_is_solved_and_feasible,
+    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
+    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)
-    @variable(model, 1 <= color[i = 1:n] <= i, Int)
+    # 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)
-    assert_is_solved_and_feasible(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
+    # remap to 1:cmax in case they are not contiguous
     true_ncolors = 0
     remap = fill(0, maximum(color_int))
     for c in color_int
@@ -46,12 +65,16 @@ end
 
 function ADTypes.column_coloring(A::AbstractMatrix, algo::OptimalColoringAlgorithm)
     bg = BipartiteGraph(A)
-    return optimal_distance2_coloring(bg, Val(2), algo.optimizer; algo.silent)
+    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)
+    return optimal_distance2_coloring(
+        bg, Val(1), algo.optimizer; algo.silent, algo.assert_solved
+    )
 end
 
 end

src/optimal.jl

@@ -11,17 +11,23 @@ Coloring algorithm that relies on mathematical programming with [JuMP](https://j
 
 # Constructor
 
-    OptimalColoringAlgorithm(optimizer; silent::Bool=true)
+    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)
-    return OptimalColoringAlgorithm(optimizer, silent)
+function OptimalColoringAlgorithm(optimizer; silent::Bool=true, assert_solved::Bool=true)
+    return OptimalColoringAlgorithm(optimizer, silent, assert_solved)
 end

test/optimal.jl

@@ -34,3 +34,13 @@ end
         @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