Merge pull request #871 from EveCharbie/ipopt_output_save

Save IPOPT output

bioptim/examples/getting_started/pendulum.py

@@ -10,9 +10,6 @@
 """
 
 import platform
-import matplotlib.pyplot as plt
-import numpy as np
-
 
 from bioptim import (
     OptimalControlProgram,
@@ -29,8 +26,6 @@
     BiorbdModel,
     ControlType,
     PhaseDynamics,
-    SolutionMerge,
-    TimeAlignment,
 )
 
 
@@ -135,10 +130,18 @@ def main():
     # --- Prepare the ocp --- #
     ocp = prepare_ocp(biorbd_model_path="models/pendulum.bioMod", final_time=1, n_shooting=400, n_threads=2)
 
-    # Custom plots
+    # --- Live plots --- #
     ocp.add_plot_penalty(CostType.ALL)  # This will display the objectives and constraints at the current iteration
-    # ocp.add_plot_ipopt_outputs()  # This will display the solver's output at the current iteration
     # ocp.add_plot_check_conditioning()  # This will display the conditioning of the problem at the current iteration
+    # ocp.add_plot_ipopt_outputs()  # This will display the solver's output at the current iteration
+
+    # --- Saving the solver's output during the optimization --- #
+    # path_to_results = "temporary_results/"
+    # result_file_name = "pendulum"
+    # nb_iter_save = 10  # Save the solver's output every 10 iterations
+    # ocp.save_intermediary_ipopt_iterations(
+    #     path_to_results, result_file_name, nb_iter_save
+    # )  # This will save the solver's output at each iteration
 
     # --- If one is interested in checking the conditioning of the problem, they can uncomment the following line --- #
     # ocp.check_conditioning()
@@ -149,13 +152,12 @@ def main():
     # --- Solve the ocp. Please note that online graphics only works with the Linux operating system --- #
     sol = ocp.solve(Solver.IPOPT(show_online_optim=platform.system() == "Linux"))
 
-    sol.print_cost()
-
     # --- Show the results (graph or animation) --- #
+    sol.print_cost()
     # sol.graphs(show_bounds=True, save_name="results.png")
     sol.animate(n_frames=100)
 
-    # # --- Save the solution --- #
+    # # --- Saving the solver's output after the optimization --- #
     # Here is an example of how we recommend to save the solution. Please note that sol.ocp is not picklable and that sol will be loaded using the current bioptim version, not the version at the time of the generation of the results.
     # import pickle
     # import git

bioptim/gui/ipopt_output_plot.py

@@ -1,7 +1,9 @@
+import pickle
+
 import numpy as np
+from casadi import jacobian, gradient, sum1, Function
 from matplotlib import pyplot as plt
 from matplotlib.cm import get_cmap
-from casadi import jacobian, gradient, sum1, Function
 
 
 def create_ipopt_output_plot(ocp, interface):
@@ -131,3 +133,63 @@ def update_ipopt_output_plot(args, ocp):
         ocp.ipopt_plots["axs"][i].set_xlim(0, len(ocp.ipopt_plots["f"]))
 
     ocp.ipopt_plots["ipopt_fig"].canvas.draw()
+
+
+def save_ipopt_output(args, save_ipopt_iterations_info):
+    """
+    This function saves the ipopt outputs: x, f, g, lam_x, lam_g, lam_p every nb_iter_save iterations.
+    """
+    f = args["f"]
+
+    if len(save_ipopt_iterations_info.f_list) != 0 and save_ipopt_iterations_info.f_list[-1] == f:
+        return
+
+    save_ipopt_iterations_info.f_list += [f]
+    save_ipopt_iterations_info.current_iter += 1
+
+    if save_ipopt_iterations_info.current_iter % save_ipopt_iterations_info.nb_iter_save != 0:
+        return
+    else:
+        x = args["x"]
+        g = args["g"]
+        lam_x = args["lam_x"]
+        lam_g = args["lam_g"]
+        lam_p = args["lam_p"]
+
+        save_path = (
+            save_ipopt_iterations_info.path_to_results
+            + save_ipopt_iterations_info.result_file_name
+            + "_"
+            + str(save_ipopt_iterations_info.current_iter)
+            + ".pkl"
+        )
+        with open(save_path, "wb") as file:
+            pickle.dump({"x": x, "f": f, "g": g, "lam_x": lam_x, "lam_g": lam_g, "lam_p": lam_p}, file)
+
+
+class SaveIterationsInfo:
+    """
+    This class is used to store the ipopt outputs save info.
+    """
+
+    def __init__(self, path_to_results: str, result_file_name: str, nb_iter_save: int):
+
+        if not isinstance(path_to_results, str) or len(path_to_results) == 0:
+            raise ValueError("path_to_results should be a non-empty string")
+        if path_to_results[-1] != "/":
+            path_to_results += "/"
+
+        if not isinstance(result_file_name, str) or len(result_file_name) == 0:
+            raise ValueError("result_file_name should be a non-empty string")
+        if result_file_name[-4:] == ".pkl":
+            result_file_name = result_file_name[:-4]
+        result_file_name.replace(".", "-")
+
+        if not isinstance(nb_iter_save, int) or nb_iter_save <= 0:
+            raise ValueError("nb_iter_save should be a positive integer")
+
+        self.path_to_results = path_to_results
+        self.result_file_name = result_file_name
+        self.nb_iter_save = nb_iter_save
+        self.current_iter = 0
+        self.f_list = []

bioptim/gui/plot.py

@@ -756,6 +756,11 @@ def update_data(
 
             update_ipopt_output_plot(args, self.ocp)
 
+        if self.ocp.save_ipopt_iterations_info is not None:
+            from ..gui.ipopt_output_plot import save_ipopt_output
+
+            save_ipopt_output(args, self.ocp.save_ipopt_iterations_info)
+
         if self.ocp.plot_check_conditioning:
             from ..gui.check_conditioning import update_conditioning_plots
 

bioptim/optimization/optimal_control_program.py

@@ -14,6 +14,7 @@
 from ..gui.check_conditioning import check_conditioning
 from ..gui.graph import OcpToConsole, OcpToGraph
 from ..gui.plot import CustomPlot, PlotOcp
+from ..gui.ipopt_output_plot import SaveIterationsInfo
 from ..interfaces import Solver
 from ..interfaces.abstract_options import GenericSolver
 from ..limits.constraints import (
@@ -592,6 +593,7 @@ def _check_arguments_and_build_nlp(
         self.plot_ipopt_outputs = False
         # If we want the conditioning of the problem to be plotted live
         self.plot_check_conditioning = False
+        self.save_ipopt_iterations_info = None
 
         return (
             constraints,
@@ -1326,6 +1328,9 @@ def add_plot_ipopt_outputs(self):
     def add_plot_check_conditioning(self):
         self.plot_check_conditioning = True
 
+    def save_intermediary_ipopt_iterations(self, path_to_results, result_file_name, nb_iter_save):
+        self.save_ipopt_iterations_info = SaveIterationsInfo(path_to_results, result_file_name, nb_iter_save)
+
     def prepare_plots(
         self,
         automatically_organize: bool = True,

tests/shard1/test__global_plots.py

@@ -85,6 +85,24 @@ def test_plot_ipopt_output_live(phase_dynamics):
     ocp.add_plot_ipopt_outputs()
 
 
+def test_save_ipopt_output():
+    from bioptim.examples.getting_started import pendulum as ocp_module
+
+    bioptim_folder = os.path.dirname(ocp_module.__file__)
+
+    ocp = ocp_module.prepare_ocp(
+        biorbd_model_path=bioptim_folder + "/models/pendulum.bioMod",
+        final_time=1,
+        n_shooting=40,
+    )
+    path_to_results = bioptim_folder + "/temporary_results/"
+    if path_to_results not in os.listdir(bioptim_folder):
+        os.mkdir(path_to_results)
+    result_file_name = "pendulum"
+    nb_iter_save = 10
+    ocp.save_intermediary_ipopt_iterations(path_to_results, result_file_name, nb_iter_save)
+
+
 @pytest.mark.parametrize("phase_dynamics", [PhaseDynamics.SHARED_DURING_THE_PHASE, PhaseDynamics.ONE_PER_NODE])
 def test_plot_merged_graphs(phase_dynamics):
     # Load graphs_one_phase