Adding some callback retrieving time statistics for solver that can produce intermediate results

discrete_optimization/generic_tools/callbacks/stats_retrievers.py

@@ -0,0 +1,88 @@
+#  Copyright (c) 2024 AIRBUS and its affiliates.
+#  This source code is licensed under the MIT license found in the
+#  LICENSE file in the root directory of this source tree.
+from time import perf_counter
+from typing import Optional
+
+from discrete_optimization.generic_tools.callbacks.callback import Callback
+from discrete_optimization.generic_tools.do_solver import SolverDO
+from discrete_optimization.generic_tools.ortools_cpsat_tools import OrtoolsCpSatSolver
+from discrete_optimization.generic_tools.result_storage.result_storage import (
+    ResultStorage,
+)
+
+
+class BasicStatsCallback(Callback):
+    """
+    This callback is storing the computation time at different step of the solving process,
+    this can help to display the evolution of the best solution through time, and compare easily different solvers.
+    """
+
+    def __init__(self):
+        self.starting_time: int = None
+        self.end_time: int = None
+        self.stats: list[dict] = []
+
+    def on_step_end(
+        self, step: int, res: ResultStorage, solver: SolverDO
+    ) -> Optional[bool]:
+        t = perf_counter()
+        best_sol, fit = res.get_best_solution_fit()
+        self.stats.append({"sol": best_sol, "fit": fit, "time": t - self.starting_time})
+
+    def on_solve_start(self, solver: SolverDO):
+        self.starting_time = perf_counter()
+
+    def on_solve_end(self, res: ResultStorage, solver: SolverDO):
+        """Called at the end of solve.
+        Args:
+        res: current result storage
+        solver: solvers using the callback
+        """
+        self.on_step_end(None, res, solver)
+
+
+class StatsCpsatCallback(BasicStatsCallback):
+    """
+    This callback is specific to cpsat solver.
+    """
+
+    def __init__(self):
+        super().__init__()
+        self.final_status: str = None
+
+    def on_step_end(
+        self, step: int, res: ResultStorage, solver: OrtoolsCpSatSolver
+    ) -> Optional[bool]:
+        super().on_step_end(step=step, res=res, solver=solver)
+        self.stats[-1].update(
+            {
+                "obj": solver.clb.ObjectiveValue(),
+                "bound": solver.clb.BestObjectiveBound(),
+                "time-cpsat": {
+                    "user-time": solver.clb.UserTime(),
+                    "wall-time": solver.clb.WallTime(),
+                },
+            }
+        )
+        if solver.clb.ObjectiveValue() == solver.clb.BestObjectiveBound():
+            return False
+
+    def on_solve_start(self, solver: OrtoolsCpSatSolver):
+        self.starting_time = perf_counter()
+
+    def on_solve_end(self, res: ResultStorage, solver: OrtoolsCpSatSolver):
+        # super().on_solve_end(res=res, solver=solver)
+        status_name = solver.solver.status_name()
+        if len(self.stats) > 0:
+            self.stats[-1].update(
+                {
+                    "obj": solver.solver.ObjectiveValue(),
+                    "bound": solver.solver.BestObjectiveBound(),
+                    "time-cpsat": {
+                        "user-time": solver.clb.UserTime(),
+                        "wall-time": solver.clb.WallTime(),
+                    },
+                }
+            )
+        self.final_status = status_name

tests/generic_tools/callbacks/test_stats_callbacks.py

@@ -0,0 +1,58 @@
+#  Copyright (c) 2024 AIRBUS and its affiliates.
+#  This source code is licensed under the MIT license found in the
+#  LICENSE file in the root directory of this source tree.
+from matplotlib import pyplot as plt
+
+from discrete_optimization.generic_rcpsp_tools.solvers.ls import LsGenericRcpspSolver
+from discrete_optimization.generic_tools.callbacks.stats_retrievers import (
+    BasicStatsCallback,
+    StatsCpsatCallback,
+)
+from discrete_optimization.rcpsp.parser import get_data_available, parse_file
+from discrete_optimization.rcpsp.solvers.cpsat import CpSatRcpspSolver
+from discrete_optimization.rcpsp.solvers.dp import DpRcpspSolver, dp
+
+
+def test_basic_stats_callback():
+    file = [f for f in get_data_available() if "j301_1.sm" in f][0]
+    problem = parse_file(file)
+    callback = BasicStatsCallback()
+    solver = LsGenericRcpspSolver(problem=problem)
+    res = solver.solve(
+        callbacks=[callback],
+        nb_iteration_max=10000,
+        retrieve_intermediate_solutions=True,
+    )
+    fig, ax = plt.subplots(1)
+    ax.plot([x["time"] for x in callback.stats], [x["fit"] for x in callback.stats])
+    ax.set_xlabel("Time (s)")
+    ax.set_ylabel("Fitness")
+    plt.show()
+    assert len(callback.stats) > 0
+
+
+def test_cpsat_callback():
+    file = [f for f in get_data_available() if "j1201_5.sm" in f][0]
+    problem = parse_file(file)
+    callback = StatsCpsatCallback()
+    solver = CpSatRcpspSolver(problem=problem)
+    res = solver.solve(callbacks=[callback], time_limit=20)
+    fig, ax = plt.subplots(1)
+    ax.plot(
+        [x["time"] for x in callback.stats],
+        [x["obj"] for x in callback.stats],
+        label="obj function",
+    )
+    ax.plot(
+        [x["time"] for x in callback.stats],
+        [x["bound"] for x in callback.stats],
+        label="bound",
+    )
+    ax.legend()
+    ax.set_xlabel("Time (s)")
+    ax.set_ylabel("Objective")
+    plt.show()
+    assert len(callback.stats) > 0
+    assert all("bound" in x for x in callback.stats)
+    assert all("obj" in x for x in callback.stats)
+    assert all("time-cpsat" in x for x in callback.stats)