Extract solutions at the end of mathopt solve for OrtoolsMathOptMilpSolver + HiGHS

discrete_optimization/generic_tools/lp_tools.py

@@ -371,6 +371,7 @@ class OrtoolsMathOptMilpSolver(MilpSolver, WarmstartMixin):
 
     """
     random_seed: Optional[int] = None
+    mathopt_res: Optional[mathopt.SolveResult] = None
 
     def remove_constraints(self, constraints: Iterable[Any]) -> None:
         for cstr in constraints:
@@ -442,6 +443,8 @@ def solve(
         mathopt_enable_output: bool = False,
         mathopt_model_parameters: Optional[mathopt.ModelSolveParameters] = None,
         mathopt_additional_solve_parameters: Optional[mathopt.SolveParameters] = None,
+        store_mathopt_res: bool = False,
+        extract_solutions_from_mathopt_res: Optional[bool] = None,
         **kwargs: Any,
     ) -> ResultStorage:
         """Solve with OR-Tools MathOpt API
@@ -458,11 +461,21 @@ def solve(
             mathopt_additional_solve_parameters: passed to `mathopt.solve()` as `params`,
                 except that parameters defined by above `time_limit` and `parameters_milp`
                 will be overriden by them.
+            store_mathopt_res: whether to store the `mathopt.SolveResult` generated by `mathopt.solve()`
+            extract_solutions_from_mathopt_res: whether to extract solutions from the `mathopt.SolveResult` generated by `mathopt.solve()`.
+                If False, the solutions are rather extracted on the fly inside a callback at each mip solution improvement.
+                By default, this is False except for HiGHS solver type, as its OR-Tools wrapper does not yet integrate callbacks.
             **kwargs: passed to init_model() if model not yet existing
 
         Returns:
 
         """
+        # Default parameters
+        if extract_solutions_from_mathopt_res is None:
+            extract_solutions_from_mathopt_res = (
+                mathopt_solver_type == mathopt.SolverType.HIGHS
+            )
+
         self.early_stopping_exception = None
         callbacks_list = CallbackList(callbacks=callbacks)
 
@@ -473,7 +486,7 @@ def solve(
         mathopt_cb = MathOptCallback(do_solver=self, callback=callbacks_list)
 
         # optimize
-        self.optimize_model(
+        mathopt_res = self.optimize_model(
             parameters_milp=parameters_milp,
             time_limit=time_limit,
             mathopt_solver_type=mathopt_solver_type,
@@ -483,9 +496,14 @@ def solve(
             mathopt_additional_solve_parameters=mathopt_additional_solve_parameters,
             **kwargs,
         )
+        if store_mathopt_res:
+            self.mathopt_res = mathopt_res
 
         # get result storage
-        res = mathopt_cb.res
+        if extract_solutions_from_mathopt_res:
+            res = self._extract_result_storage(mathopt_res)
+        else:
+            res = mathopt_cb.res
 
         # callback: solve end
         callbacks_list.on_solve_end(res=res, solver=self)
@@ -630,6 +648,25 @@ def construct_linear_sum(expr: Iterable) -> Any:
         """Generate a linear sum (with variables) ready for the internal model."""
         return mathopt.LinearSum(expr)
 
+    def _extract_result_storage(
+        self, mathopt_res: mathopt.SolveResult
+    ) -> ResultStorage:
+        list_solution_fits = []
+        for internal_sol in mathopt_res.solutions:
+            get_var_value_for_current_solution = (
+                lambda var: internal_sol.primal_solution.variable_values[var]
+            )
+            get_obj_value_for_current_solution = (
+                lambda: internal_sol.primal_solution.objective_value
+            )
+            sol = self.retrieve_current_solution(
+                get_var_value_for_current_solution=get_var_value_for_current_solution,
+                get_obj_value_for_current_solution=get_obj_value_for_current_solution,
+            )
+            fit = self.aggreg_from_sol(sol)
+            list_solution_fits.append((sol, fit))
+        return self.create_result_storage(list(reversed(list_solution_fits)))
+
 
 map_mathopt_status_to_do_status: dict[mathopt.TerminationReason, StatusSolver] = {
     mathopt.TerminationReason.OPTIMAL: StatusSolver.OPTIMAL,

tests/coloring/solvers/test_solvers.py

@@ -8,6 +8,7 @@
 import numpy as np
 import pytest
 from minizinc.solver import Solver
+from ortools.math_opt.python import mathopt
 
 from discrete_optimization.coloring.parser import get_data_available, parse_file
 from discrete_optimization.coloring.problem import (
@@ -563,7 +564,10 @@ def test_color_lp_gurobi_cb_exception():
 
 @pytest.mark.parametrize("use_cliques", [False, True])
 @pytest.mark.parametrize("greedy_start", [True, False])
-def test_color_lp_ortools_mathopt(use_cliques, greedy_start):
+@pytest.mark.parametrize(
+    "solver_type", [mathopt.SolverType.CP_SAT, mathopt.SolverType.HIGHS]
+)
+def test_color_lp_ortools_mathopt(use_cliques, greedy_start, solver_type):
     file = [f for f in get_data_available() if "gc_70_1" in f][0]
     color_problem = parse_file(file)
     solver = MathOptColoringSolver(
@@ -574,13 +578,14 @@ def test_color_lp_ortools_mathopt(use_cliques, greedy_start):
         use_cliques=use_cliques,
         greedy_start=greedy_start,
         parameters_milp=ParametersMilp.default(),
+        mathopt_solver_type=solver_type,
     )
     result_store = solver.solve(**kwargs)
     solution = result_store.get_best_solution_fit()[0]
     assert color_problem.satisfy(solution)
 
-    # Test warm-start only once
-    if greedy_start and not use_cliques:
+    # Test warm-start only once (and not for HiGHS as only 1 solution found)
+    if greedy_start and not use_cliques and solver_type != mathopt.SolverType.HIGHS:
         # first solution is not start_solution
         assert result_store[0][0].colors != solver.start_solution.colors