[DO NOT MERGE] Stochastic plots, noisy_integrate, and other plot shits

bioptim/dynamics/configure_new_variable.py

@@ -518,6 +518,20 @@
                     self.nlp.variable_mappings[self.name],
                     node_index,
                 )
+                if not self.skip_plot:
+                    all_variables_in_one_subplot = True if self.name in ["m", "cov", "k"] else False
+                    self.nlp.plot[f"{self.name}_algebraic"] = CustomPlot(
+                        lambda t0, phases_dt, node_idx, x, u, p, a: (
+                            a[self.nlp.algebraic_states.key_index(self.name), :]
+                            if a.any()
+                            else np.ndarray((cx[0][0].shape[0], 1)) * np.nan
+                        ),
+                        plot_type=PlotType.STEP,
+                        axes_idx=self.axes_idx,
+                        legend=self.legend,
+                        combine_to=self.combine_name,
+                        all_variables_in_one_subplot=all_variables_in_one_subplot,
+                    )
 
 
 def _manage_fatigue_to_new_variable(

bioptim/dynamics/configure_problem.py

@@ -908,31 +908,85 @@
             A reference to the ocp
         nlp: NonLinearProgram
             A reference to the phase
-        dyn_func: Callable[time, states, controls, param, algebraic_states] | tuple[Callable[time, states, controls, param, algebraic_states], ...]
+        dyn_func: Callable[time, states, controls, param, algebraic_states]
             The function to get the derivative of the states
         """
 
         DynamicsFunctions.apply_parameters(nlp)
 
-        if not isinstance(dyn_func, (tuple, list)):
-            dyn_func = (dyn_func,)
+        dynamics_eval = dyn_func(
+            nlp.time_mx,
+            nlp.states.scaled.mx_reduced,
+            nlp.controls.scaled.mx_reduced,
+            nlp.parameters.scaled.mx_reduced,
+            nlp.algebraic_states.scaled.mx_reduced,
+            nlp,
+            **extra_params,
+        )
+        dynamics_dxdt = dynamics_eval.dxdt
+        if isinstance(dynamics_dxdt, (list, tuple)):
+            dynamics_dxdt = vertcat(*dynamics_dxdt)
 
-        for func in dyn_func:
-            dynamics_eval = func(
-                nlp.time_mx,
-                nlp.states.scaled.mx_reduced,
-                nlp.controls.scaled.mx_reduced,
-                nlp.parameters.scaled.mx_reduced,
-                nlp.algebraic_states.scaled.mx_reduced,
-                nlp,
-                **extra_params,
+        time_span_sym = vertcat(nlp.time_mx, nlp.dt_mx)
+        if nlp.dynamics_func is None:
+            nlp.dynamics_func = Function(
+                "ForwardDyn",
+                [
+                    time_span_sym,
+                    nlp.states.scaled.mx_reduced,
+                    nlp.controls.scaled.mx_reduced,
+                    nlp.parameters.scaled.mx_reduced,
+                    nlp.algebraic_states.scaled.mx_reduced,
+                ],
+                [dynamics_dxdt],
+                ["t_span", "x", "u", "p", "a"],
+                ["xdot"],
             )
-            dynamics_dxdt = dynamics_eval.dxdt
-            if isinstance(dynamics_dxdt, (list, tuple)):
-                dynamics_dxdt = vertcat(*dynamics_dxdt)
 
-            time_span_sym = vertcat(nlp.time_mx, nlp.dt_mx)
-            nlp.dynamics_func.append(
+            # TODO: allow expand for each dynamics independently
+            if nlp.dynamics_type.expand_dynamics:
+                try:
+                    nlp.dynamics_func = nlp.dynamics_func.expand()
+                except Exception as me:
+                    RuntimeError(
+                        f"An error occurred while executing the 'expand()' function for the dynamic function. "
+                        f"Please review the following casadi error message for more details.\n"
+                        "Several factors could be causing this issue. One of the most likely is the inability to "
+                        "use expand=True at all. In that case, try adding expand=False to the dynamics.\n"
+                        "Original casadi error message:\n"
+                        f"{me}"
+                    )
+
+            # Only possible for regular dynamics, not for extra_dynamics
+            if dynamics_eval.defects is not None:
+                nlp.implicit_dynamics_func = Function(
+                    "DynamicsDefects",
+                    [
+                        time_span_sym,
+                        nlp.states.scaled.mx_reduced,
+                        nlp.controls.scaled.mx_reduced,
+                        nlp.parameters.scaled.mx_reduced,
+                        nlp.algebraic_states.scaled.mx_reduced,
+                        nlp.states_dot.scaled.mx_reduced,
+                    ],
+                    [dynamics_eval.defects],
+                    ["t_span", "x", "u", "p", "a", "xdot"],
+                    ["defects"],
+                )
+                if nlp.dynamics_type.expand_dynamics:
+                    try:
+                        nlp.implicit_dynamics_func = nlp.implicit_dynamics_func.expand()
+                    except Exception as me:
+                        RuntimeError(
+                            f"An error occurred while executing the 'expand()' function for the dynamic function. "
+                            f"Please review the following casadi error message for more details.\n"
+                            "Several factors could be causing this issue. One of the most likely is the inability to "
+                            "use expand=True at all. In that case, try adding expand=False to the dynamics.\n"
+                            "Original casadi error message:\n"
+                            f"{me}"
+                        )
+        else:
+            nlp.extra_dynamics_func.append(
                 Function(
                     "ForwardDyn",
                     [
@@ -951,7 +1005,7 @@
             # TODO: allow expand for each dynamics independently
             if nlp.dynamics_type.expand_dynamics:
                 try:
-                    nlp.dynamics_func[-1] = nlp.dynamics_func[-1].expand()
+                    nlp.extra_dynamics_func[-1] = nlp.dynamics_funcextra_dynamics_func[-1].expand()
                 except Exception as me:
                     RuntimeError(
                         f"An error occurred while executing the 'expand()' function for the dynamic function. "
@@ -962,36 +1016,6 @@
                         f"{me}"
                     )
 
-            if dynamics_eval.defects is not None:
-                nlp.implicit_dynamics_func.append(
-                    Function(
-                        "DynamicsDefects",
-                        [
-                            time_span_sym,
-                            nlp.states.scaled.mx_reduced,
-                            nlp.controls.scaled.mx_reduced,
-                            nlp.parameters.scaled.mx_reduced,
-                            nlp.algebraic_states.scaled.mx_reduced,
-                            nlp.states_dot.scaled.mx_reduced,
-                        ],
-                        [dynamics_eval.defects],
-                        ["t_span", "x", "u", "p", "a", "xdot"],
-                        ["defects"],
-                    )
-                )
-                if nlp.dynamics_type.expand_dynamics:
-                    try:
-                        nlp.implicit_dynamics_func[-1] = nlp.implicit_dynamics_func[-1].expand()
-                    except Exception as me:
-                        RuntimeError(
-                            f"An error occurred while executing the 'expand()' function for the dynamic function. "
-                            f"Please review the following casadi error message for more details.\n"
-                            "Several factors could be causing this issue. One of the most likely is the inability to "
-                            "use expand=True at all. In that case, try adding expand=False to the dynamics.\n"
-                            "Original casadi error message:\n"
-                            f"{me}"
-                        )
-
     @staticmethod
     def configure_contact_function(ocp, nlp, dyn_func: Callable, **extra_params):
         """
@@ -1369,7 +1393,6 @@
             as_controls=False,
             as_states_dot=False,
             as_algebraic_states=True,
-            skip_plot=True,
         )
 
     @staticmethod
@@ -1493,7 +1516,6 @@
             as_controls=False,
             as_states_dot=False,
             as_algebraic_states=True,
-            skip_plot=True,
         )
 
     @staticmethod
@@ -1525,7 +1547,6 @@
             as_controls=False,
             as_states_dot=False,
             as_algebraic_states=True,
-            skip_plot=True,
         )
 
     @staticmethod
@@ -1554,7 +1575,6 @@
             as_controls=False,
             as_states_dot=False,
             as_algebraic_states=True,
-            skip_plot=True,
         )
 
     @staticmethod
@@ -1589,7 +1609,6 @@
             as_controls=False,
             as_states_dot=False,
             as_algebraic_states=True,
-            skip_plot=True,
         )
 
     @staticmethod

bioptim/dynamics/ode_solver.py

@@ -159,7 +159,9 @@
         """
         return nlp.parameters.scaled.cx_start
 
-    def initialize_integrator(self, ocp, nlp, dynamics_index: int, node_index: int, **extra_opt) -> Callable:
+    def initialize_integrator(
+        self, ocp, nlp, node_index: int, dynamics_index: int = 0, is_extra_dynamics: bool = False, **extra_opt
+    ) -> Callable:
         """
         Initialize the integrator
 
@@ -169,10 +171,12 @@
             The Optimal control program handler
         nlp
             The NonLinearProgram handler
-        dynamics_index
-            The current dynamics to resolve (that can be referred to nlp.dynamics_func[index])
         node_index
             The index of the node currently initialized
+        dynamics_index
+            The current extra dynamics to resolve (that can be referred to nlp.extra_dynamics_func[index])
+        is_extra_dynamics
+            If the dynamics is an extra dynamics
         extra_opt
             Any extra options to pass to the integrator
 
@@ -181,16 +185,20 @@
         The initialized integrator function
         """
 
+        if dynamics_index > 0 and not is_extra_dynamics:
+            raise RuntimeError("dynamics_index should be 0 if is_extra_dynamics is False")
+
         nlp.states.node_index = node_index
         nlp.states_dot.node_index = node_index
         nlp.controls.node_index = node_index
         nlp.algebraic_states.node_index = node_index
+        dynamics_func = nlp.dynamics_func if not is_extra_dynamics else nlp.extra_dynamics_func[dynamics_index]
         ode_opt = {
             "model": nlp.model,
             "cx": nlp.cx,
             "control_type": nlp.control_type,
             "defects_type": self.defects_type,
-            "ode_index": node_index if nlp.dynamics_func[dynamics_index].size2_out("xdot") > 1 else 0,
+            "ode_index": node_index if dynamics_func.size2_out("xdot") > 1 else 0,
             "duplicate_starting_point": self.duplicate_starting_point,
             **extra_opt,
         }
@@ -201,13 +209,8 @@
             "u": self.p_ode(nlp),
             "a": self.a_ode(nlp),
             "param": self.param_ode(nlp),
-            "ode": nlp.dynamics_func[dynamics_index],
-            # TODO this actually checks "not nlp.implicit_dynamics_func" (or that nlp.implicit_dynamics_func == [])
-            "implicit_ode": (
-                nlp.implicit_dynamics_func[dynamics_index]
-                if len(nlp.implicit_dynamics_func) > 0
-                else nlp.implicit_dynamics_func
-            ),
+            "ode": dynamics_func,
+            "implicit_ode": nlp.implicit_dynamics_func,
         }
 
         return nlp.ode_solver.integrator(ode, ode_opt)
@@ -235,14 +238,19 @@
 
         # Extra dynamics
         extra_dynamics = []
-        for i in range(1, len(nlp.dynamics_func)):
-            extra_dynamics += [nlp.ode_solver.initialize_integrator(ocp, nlp, dynamics_index=i, node_index=0)]
+        for i in range(len(nlp.extra_dynamics_func)):
+            extra_dynamics += [
+                nlp.ode_solver.initialize_integrator(ocp, nlp, dynamics_index=i, node_index=0, is_extra_dynamics=True)
+            ]
             if nlp.phase_dynamics == PhaseDynamics.SHARED_DURING_THE_PHASE:
                 extra_dynamics = extra_dynamics * nlp.ns
             else:
                 for node_index in range(1, nlp.ns):
-                    extra_dynamics += [nlp.ode_solver.initialize_integrator(ocp, nlp, dynamics_index=i, node_index=0)]
-            # TODO include this in nlp.dynamics so the index of nlp.dynamics_func and nlp.dynamics match
+                    extra_dynamics += [
+                        nlp.ode_solver.initialize_integrator(
+                            ocp, nlp, dynamics_index=i, node_index=node_index, is_extra_dynamics=True
+                        )
+                    ]
             nlp.extra_dynamics.append(extra_dynamics)
 
 
@@ -538,7 +546,9 @@
         def a_ode(self, nlp):
             return nlp.algebraic_states.scaled.cx
 
-        def initialize_integrator(self, ocp, nlp, dynamics_index: int, node_index: int, **extra_opt):
+        def initialize_integrator(
+            self, ocp, nlp, node_index: int, dynamics_index: int = 0, is_extra_dynamics: bool = False, **extra_opt
+        ):
             raise NotImplementedError("CVODES is not yet implemented")
 
             if extra_opt:
@@ -560,10 +570,11 @@
                 raise RuntimeError("CVODES cannot be used with algebraic_states variables")
 
             t = [self.t_ode(nlp)[0], self.t_ode(nlp)[1] - self.t_ode(nlp)[0]]
+            dynamics_func = nlp.dynamics_func if not is_extra_dynamics else nlp.extra_dynamics_func[dynamics_index]
             ode = {
                 "x": nlp.states.scaled.cx_start,
                 "u": nlp.controls.scaled.cx_start,  # todo: add p=parameters
-                "ode": nlp.dynamics_func[dynamics_index](
+                "ode": dynamics_func(
                     vertcat(*t), self.x_ode(nlp), self.p_ode(nlp), self.param_ode(nlp), self.a_ode(nlp)
                 ),
             }

bioptim/examples/getting_started/pendulum.py

@@ -136,7 +136,9 @@
     ocp = prepare_ocp(biorbd_model_path="models/pendulum.bioMod", final_time=1, n_shooting=400, n_threads=2)
 
     # Custom plots
-    ocp.add_plot_penalty(CostType.ALL)
+    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
 
     # --- If one is interested in checking the conditioning of the problem, they can uncomment the following line --- #
     # ocp.check_conditioning()
@@ -146,12 +148,60 @@
 
     # --- 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.graphs(show_bounds=True)
+    # sol.graphs(show_bounds=True, save_name="results.png")
     sol.animate(n_frames=100)
 
+    # # --- Save the solution --- #
+    # 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
+    # from datetime import date
+    #
+    # # Save the version of bioptim and the date of the optimization for future reference
+    # repo = git.Repo(search_parent_directories=True)
+    # commit_id = str(repo.commit())
+    # branch = str(repo.active_branch)
+    # tag = repo.git.describe("--tags")
+    # bioptim_version = repo.git.version_info
+    # git_date = repo.git.log("-1", "--format=%cd")
+    # version_dic = {
+    #     "commit_id": commit_id,
+    #     "git_date": git_date,
+    #     "branch": branch,
+    #     "tag": tag,
+    #     "bioptim_version": bioptim_version,
+    #     "date_of_the_optimization": date.today().strftime("%b-%d-%Y-%H-%M-%S"),
+    # }
+    #
+    # q = sol.decision_states(to_merge=SolutionMerge.NODES)["q"]
+    # qdot = sol.decision_states(to_merge=SolutionMerge.NODES)["qdot"]
+    # tau = sol.decision_controls(to_merge=SolutionMerge.NODES)["tau"]
+    #
+    # # Do everything you need with the solution here before we delete ocp
+    # integrated_sol = sol.integrate(to_merge=SolutionMerge.NODES)
+    # q_integrated = integrated_sol["q"]
+    # qdot_integrated = integrated_sol["qdot"]
+    #
+    # # Save the output of the optimization
+    # with open("pendulum_data.pkl", "wb") as file:
+    #     data = {"q": q,
+    #             "qdot": qdot,
+    #             "tau": tau,
+    #             "real_time_to_optimize": sol.real_time_to_optimize,
+    #             "version": version_dic,
+    #             "q_integrated": q_integrated,
+    #             "qdot_integrated": qdot_integrated}
+    #     pickle.dump(data, file)
+    #
+    # # Save the solution for future use, you will only need to do sol.ocp = prepare_ocp() to get the same solution object as above.
+    # with open("pendulum_sol.pkl", "wb") as file:
+    #     del sol.ocp
+    #     pickle.dump(sol, file)
+
 
 if __name__ == "__main__":
     main()