Merge branch 'test' of https://github.com/EveCharbie/BiorbdOptim into…

… test

bioptim/dynamics/configure_new_variable.py

@@ -418,7 +418,7 @@ def _declare_cx_and_plot(self):
                 )
                 if not self.skip_plot:
                     self.nlp.plot[f"{self.name}_states"] = CustomPlot(
-                        lambda t0, phases_dt, node_idx, x, u, p, a: (
+                        lambda t0, phases_dt, node_idx, x, u, p, a, d: (
                             x[self.nlp.states.key_index(self.name), :]
                             if x.any()
                             else np.ndarray((cx[0][0].shape[0], 1)) * np.nan
@@ -455,7 +455,7 @@ def _declare_cx_and_plot(self):
                 plot_type = PlotType.PLOT if self.nlp.control_type == ControlType.LINEAR_CONTINUOUS else PlotType.STEP
                 if not self.skip_plot:
                     self.nlp.plot[f"{self.name}_controls"] = CustomPlot(
-                        lambda t0, phases_dt, node_idx, x, u, p, a: (
+                        lambda t0, phases_dt, node_idx, x, u, p, a, d: (
                             u[self.nlp.controls.key_index(self.name), :]
                             if u.any()
                             else np.ndarray((cx[0][0].shape[0], 1)) * np.nan
@@ -521,7 +521,7 @@ def _declare_cx_and_plot(self):
                 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: (
+                        lambda t0, phases_dt, node_idx, x, u, p, a, d: (
                             a[self.nlp.algebraic_states.key_index(self.name), :]
                             if a.any()
                             else np.ndarray((cx[0][0].shape[0], 1)) * np.nan
@@ -591,7 +591,7 @@ def _manage_fatigue_to_new_variable(
     legend = [f"{name}_{i}" for i in name_elements]
     fatigue_plot_name = f"fatigue_{name}"
     nlp.plot[fatigue_plot_name] = CustomPlot(
-        lambda t0, phases_dt, node_idx, x, u, p, a: (
+        lambda t0, phases_dt, node_idx, x, u, p, a, d: (
             x[:n_elements, :] if x.any() else np.ndarray((len(name_elements), 1))
         )
         * np.nan,
@@ -601,7 +601,7 @@ def _manage_fatigue_to_new_variable(
     )
     control_plot_name = f"{name}_controls" if not multi_interface and split_controls else f"{name}"
     nlp.plot[control_plot_name] = CustomPlot(
-        lambda t0, phases_dt, node_idx, x, u, p, a: (
+        lambda t0, phases_dt, node_idx, x, u, p, a, d: (
             u[:n_elements, :] if u.any() else np.ndarray((len(name_elements), 1))
         )
         * np.nan,
@@ -621,7 +621,7 @@ def _manage_fatigue_to_new_variable(
                 var_names_with_suffix[-1], name_elements, ocp, nlp, as_states, as_controls, skip_plot=True
             )
             nlp.plot[f"{var_names_with_suffix[-1]}_controls"] = CustomPlot(
-                lambda t0, phases_dt, node_idx, x, u, p, a, key: (
+                lambda t0, phases_dt, node_idx, x, u, p, a, d, key: (
                     u[nlp.controls.key_index(key), :] if u.any() else np.ndarray((len(name_elements), 1)) * np.nan
                 ),
                 plot_type=PlotType.STEP,
@@ -632,7 +632,7 @@ def _manage_fatigue_to_new_variable(
         elif i == 0:
             NewVariableConfiguration(f"{name}", name_elements, ocp, nlp, as_states, as_controls, skip_plot=True)
             nlp.plot[f"{name}_controls"] = CustomPlot(
-                lambda t0, phases_dt, node_idx, x, u, p, a, key: (
+                lambda t0, phases_dt, node_idx, x, u, p, a, d, key: (
                     u[nlp.controls.key_index(key), :] if u.any() else np.ndarray((len(name_elements), 1)) * np.nan
                 ),
                 plot_type=PlotType.STEP,
@@ -645,7 +645,7 @@ def _manage_fatigue_to_new_variable(
             name_tp = f"{var_names_with_suffix[-1]}_{params}"
             NewVariableConfiguration(name_tp, name_elements, ocp, nlp, True, False, skip_plot=True)
             nlp.plot[name_tp] = CustomPlot(
-                lambda t0, phases_dt, node_idx, x, u, p, a, key, mod: (
+                lambda t0, phases_dt, node_idx, x, u, p, a, d, key, mod: (
                     mod * x[nlp.states.key_index(key), :] if x.any() else np.ndarray((len(name_elements), 1)) * np.nan
                 ),
                 plot_type=PlotType.INTEGRATED,

bioptim/dynamics/configure_problem.py

@@ -134,10 +134,7 @@ def initialize(ocp, nlp):
         nlp.dynamics_type.type(ocp, nlp, **nlp.dynamics_type.extra_parameters)
 
     @staticmethod
-    def custom(ocp,
-               nlp,
-               dynamics_constants_used_at_each_nodes: dict[list] = {},
-               **extra_params):
+    def custom(ocp, nlp, dynamics_constants_used_at_each_nodes: dict[list] = {}, **extra_params):
         """
         Call the user-defined dynamics configuration function
 
@@ -149,7 +146,7 @@ def custom(ocp,
             A reference to the phase
         """
 
-        nlp.dynamics_type.configure(ocp, nlp, dynamics_constants_used_at_each_nodes, **extra_params)
+        nlp.dynamics_type.configure(ocp, nlp, **extra_params)
 
     @staticmethod
     def torque_driven(
@@ -317,6 +314,7 @@ def torque_driven_free_floating_base(
         with_passive_torque: bool = False,
         with_ligament: bool = False,
         with_friction: bool = False,
+        dynamics_constants_used_at_each_nodes: dict[list] = {},
     ):
         """
         Configure the dynamics for a torque driven program with a free floating base.
@@ -336,6 +334,8 @@ def torque_driven_free_floating_base(
             If the dynamic with ligament should be used
         with_friction: bool
             If the dynamic with joint friction should be used (friction = coefficients * qdot)
+        dynamics_constants_used_at_each_nodes: dict[np.ndarray]
+            A list of values to pass to the dynamics at each node.
         """
 
         nb_q = nlp.model.nb_q
@@ -440,6 +440,7 @@ def stochastic_torque_driven(
         with_friction: bool = False,
         with_cholesky: bool = False,
         initial_matrix: DM = None,
+        dynamics_constants_used_at_each_nodes: dict[list] = {},
     ):
         """
         Configure the dynamics for a torque driven stochastic program (states are q and qdot, controls are tau)
@@ -454,6 +455,12 @@ def stochastic_torque_driven(
             If the dynamic with contact should be used
         with_friction: bool
             If the dynamic with joint friction should be used (friction = coefficient * qdot)
+        with_cholesky: bool
+            If the Cholesky decomposition should be used for the covariance matrix.
+        initial_matrix: DM
+            The initial value for the covariance matrix
+        dynamics_constants_used_at_each_nodes: dict[np.ndarray]
+            A list of values to pass to the dynamics at each node. Experimental external forces should be included here.
         """
 
         if "tau" in nlp.model.motor_noise_mapping:
@@ -513,10 +520,10 @@ def stochastic_torque_driven_free_floating_base(
         ocp,
         nlp,
         problem_type,
-        with_contact: bool = False,
         with_friction: bool = False,
         with_cholesky: bool = False,
         initial_matrix: DM = None,
+        dynamics_constants_used_at_each_nodes: dict[list] = {},
     ):
         """
         Configure the dynamics for a stochastic torque driven program with a free floating base.
@@ -528,10 +535,14 @@ def stochastic_torque_driven_free_floating_base(
             A reference to the ocp
         nlp: NonLinearProgram
             A reference to the phase
-        with_contact: bool
-            If the dynamic with contact should be used
         with_friction: bool
             If the dynamic with joint friction should be used (friction = coefficient * qdot)
+        with_cholesky: bool
+            If the Cholesky decomposition should be used for the covariance matrix.
+        initial_matrix: DM
+            The initial value for the covariance matrix
+        dynamics_constants_used_at_each_nodes: dict[np.ndarray]
+            A list of values to pass to the dynamics at each node.
         """
         n_noised_tau = nlp.model.n_noised_controls
         n_noise = nlp.model.motor_noise_magnitude.shape[0] + nlp.model.sensory_noise_magnitude.shape[0]
@@ -754,7 +765,11 @@ def torque_activations_driven(
         ConfigureProblem.configure_soft_contact_function(ocp, nlp)
 
     @staticmethod
-    def joints_acceleration_driven(ocp, nlp, rigidbody_dynamics: RigidBodyDynamics = RigidBodyDynamics.ODE):
+    def joints_acceleration_driven(ocp,
+                                   nlp,
+                                   rigidbody_dynamics: RigidBodyDynamics = RigidBodyDynamics.ODE,
+                                   dynamics_constants_used_at_each_nodes: dict[list] = {}
+                                   ):
         """
         Configure the dynamics for a joints acceleration driven program
         (states are q and qdot, controls are qddot_joints)
@@ -767,7 +782,8 @@ def joints_acceleration_driven(ocp, nlp, rigidbody_dynamics: RigidBodyDynamics =
             A reference to the phase
         rigidbody_dynamics: RigidBodyDynamics
             which rigidbody dynamics should be used
-
+        dynamics_constants_used_at_each_nodes: dict[np.ndarray]
+            A list of values to pass to the dynamics at each node. Experimental external forces should be included here.
         """
         if rigidbody_dynamics != RigidBodyDynamics.ODE:
             raise NotImplementedError("Implicit dynamics not implemented yet.")
@@ -905,7 +921,10 @@ def muscle_driven(
         ConfigureProblem.configure_soft_contact_function(ocp, nlp)
 
     @staticmethod
-    def holonomic_torque_driven(ocp, nlp):
+    def holonomic_torque_driven(ocp,
+                                nlp,
+                                dynamics_constants_used_at_each_nodes: dict[list] = {}
+                                ):
         """
         Tell the program which variables are states and controls.
 
@@ -1077,6 +1096,7 @@ def configure_contact_function(ocp, nlp, dyn_func: Callable, **extra_params):
                 nlp.controls.scaled.mx_reduced,
                 nlp.parameters.scaled.mx_reduced,
                 nlp.algebraic_states.scaled.mx_reduced,
+                nlp.dynamics_constants.mx,
             ],
             [
                 dyn_func(
@@ -1085,11 +1105,12 @@ def configure_contact_function(ocp, nlp, dyn_func: Callable, **extra_params):
                     nlp.controls.scaled.mx_reduced,
                     nlp.parameters.scaled.mx_reduced,
                     nlp.algebraic_states.scaled.mx_reduced,
+                    nlp.dynamics_constants.mx,
                     nlp,
                     **extra_params,
                 )
             ],
-            ["t_span", "x", "u", "p", "a"],
+            ["t_span", "x", "u", "p", "a", "dynamics_constants"],
             ["contact_forces"],
         ).expand()
 
@@ -1111,8 +1132,8 @@ def configure_contact_function(ocp, nlp, dyn_func: Callable, **extra_params):
             )
 
         nlp.plot["contact_forces"] = CustomPlot(
-            lambda t0, phases_dt, node_idx, x, u, p, a: nlp.contact_forces_func(
-                [t0, t0 + phases_dt[nlp.phase_idx]], x, u, p, a
+            lambda t0, phases_dt, node_idx, x, u, p, a, d: nlp.contact_forces_func(
+                [t0, t0 + phases_dt[nlp.phase_idx]], x, u, p, a, d
             ),
             plot_type=PlotType.INTEGRATED,
             axes_idx=axes_idx,
@@ -1184,8 +1205,8 @@ def configure_soft_contact_function(ocp, nlp):
                     to_second=[i for i, c in enumerate(all_soft_contact_names) if c in soft_contact_names_in_phase],
                 )
             nlp.plot[f"soft_contact_forces_{nlp.model.soft_contact_names[i_sc]}"] = CustomPlot(
-                lambda t0, phases_dt, node_idx, x, u, p, a: nlp.soft_contact_forces_func(
-                    [t0, t0 + phases_dt[nlp.phase_idx]], x, u, p, a
+                lambda t0, phases_dt, node_idx, x, u, p, a, d: nlp.soft_contact_forces_func(
+                    [t0, t0 + phases_dt[nlp.phase_idx]], x, u, p, a, d
                 )[(i_sc * 6) : ((i_sc + 1) * 6), :],
                 plot_type=PlotType.INTEGRATED,
                 axes_idx=phase_mappings,