[DO NOT MERGE] multi-threading for external forces

README.md

@@ -1872,17 +1872,8 @@ The solver must minimize the force to lift the box while reaching the marker in
 It is designed to show how to use external forces. An example of external forces that depends on the state (for
 example, a spring) can be found at 'examples/torque_driven_ocp/spring_load.py'
 
-`Bioptim` expects `external_forces` to be a list (for each shooting node) of np.ndarray [6 x n],
-where the six components are [Mx, My, Mz, Fx, Fy, Fz], expressed at the origin of the global reference frame for each node n.
-Let us take a look at the definition of the external forces in
-this example :
-
-```python
-external_forces = [[["Seg1", (0, 0, 0, 0, 0, -2)], ["Test", (0, 0, 0, 0, 0, 5)]] for _ in range(n_shooting)]
-```
-
-`external_forces` is 30-element long, and each sub list array are composed of a string, the name of the segment, and a tuple, the external force.
-The tuple is [Mx, My, Mz, Fx, Fy, Fz] for each node (in this example, we take 30 shooting nodes).
+`Bioptim` expects `external_forces` to be a np.ndarray [6 x n x n_shooting], where the six components are 
+[Mx, My, Mz, Fx, Fy, Fz], expressed at the origin of the global reference frame for each node.
 
 ### The [example_implicit_dynamics.py](./bioptim/examples/getting_started/example_implicit_dynamics.py) file
 *#TODO*

bioptim/dynamics/configure_problem.py

@@ -1,6 +1,7 @@
 from typing import Callable, Any
 
-from casadi import vertcat, Function, DM
+from casadi import vertcat, Function, DM, horzcat
+import numpy as np
 
 from .configure_new_variable import NewVariableConfiguration
 from .dynamics_functions import DynamicsFunctions
@@ -133,7 +134,10 @@ def initialize(ocp, nlp):
         nlp.dynamics_type.type(ocp, nlp, **nlp.dynamics_type.extra_parameters)
 
     @staticmethod
-    def custom(ocp, nlp, **extra_params):
+    def custom(ocp,
+               nlp,
+               dynamics_constants_used_at_each_nodes: dict[list] = {},
+               **extra_params):
         """
         Call the user-defined dynamics configuration function
 
@@ -145,7 +149,7 @@ def custom(ocp, nlp, **extra_params):
             A reference to the phase
         """
 
-        nlp.dynamics_type.configure(ocp, nlp, **extra_params)
+        nlp.dynamics_type.configure(ocp, nlp, dynamics_constants_used_at_each_nodes, **extra_params)
 
     @staticmethod
     def torque_driven(
@@ -158,7 +162,7 @@ def torque_driven(
         rigidbody_dynamics: RigidBodyDynamics = RigidBodyDynamics.ODE,
         soft_contacts_dynamics: SoftContactDynamics = SoftContactDynamics.ODE,
         fatigue: FatigueList = None,
-        external_forces: list = None,
+        dynamics_constants_used_at_each_nodes: dict[list] = {},
     ):
         """
         Configure the dynamics for a torque driven program (states are q and qdot, controls are tau)
@@ -183,16 +187,24 @@ def torque_driven(
             which soft contact dynamic should be used
         fatigue: FatigueList
             A list of fatigue elements
-        external_forces: list[Any]
-            A list of external forces
+        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.
         """
 
         _check_contacts_in_biorbd_model(with_contact, nlp.model.nb_contacts, nlp.phase_idx)
         _check_soft_contacts_dynamics(
             rigidbody_dynamics, soft_contacts_dynamics, nlp.model.nb_soft_contacts, nlp.phase_idx
         )
-        _check_external_forces_format(external_forces, nlp.ns, nlp.phase_idx)
-        _check_external_forces_and_phase_dynamics(external_forces, nlp.phase_dynamics, nlp.phase_idx)
+        external_forces = None
+        for key in dynamics_constants_used_at_each_nodes.keys():
+            if key != "external_forces":
+                raise RuntimeError(
+                    "The only dynamics_constants_used_at_each_nodes allowed for torque_driven dynamics is external_forces."
+                )
+            _check_dynamics_constants_format(dynamics_constants_used_at_each_nodes[key], nlp.ns, nlp.phase_idx)
+            external_forces = nlp.dynamics_constants[0].mx
+            for i in range(1, dynamics_constants_used_at_each_nodes[key].shape[1]):
+                external_forces = horzcat(external_forces, nlp.dynamics_constants[i].mx)
 
         # Declared rigidbody states and controls
         ConfigureProblem.configure_q(ocp, nlp, as_states=True, as_controls=False)
@@ -577,7 +589,7 @@ def torque_derivative_driven(
         with_ligament: bool = False,
         rigidbody_dynamics: RigidBodyDynamics = RigidBodyDynamics.ODE,
         soft_contacts_dynamics: SoftContactDynamics = SoftContactDynamics.ODE,
-        external_forces: list = None,
+        dynamics_constants_used_at_each_nodes: dict[list] = {},
     ):
         """
         Configure the dynamics for a torque driven program (states are q and qdot, controls are tau)
@@ -598,8 +610,8 @@ def torque_derivative_driven(
             which rigidbody dynamics should be used
         soft_contacts_dynamics: SoftContactDynamics
             which soft contact dynamic should be used
-        external_forces: list[Any]
-            A list of external forces
+        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.
 
         """
         _check_contacts_in_biorbd_model(with_contact, nlp.model.nb_contacts, nlp.phase_idx)
@@ -610,8 +622,16 @@ def torque_derivative_driven(
         _check_soft_contacts_dynamics(
             rigidbody_dynamics, soft_contacts_dynamics, nlp.model.nb_soft_contacts, nlp.phase_idx
         )
-        _check_external_forces_format(external_forces, nlp.ns, nlp.phase_idx)
-        _check_external_forces_and_phase_dynamics(external_forces, nlp.phase_dynamics, nlp.phase_idx)
+        external_forces = None
+        for key in dynamics_constants_used_at_each_nodes.keys():
+            if key != "external_forces":
+                raise RuntimeError(
+                    "The only dynamics_constants_used_at_each_nodes allowed for torque_driven dynamics is external_forces."
+                )
+            _check_dynamics_constants_format(dynamics_constants_used_at_each_nodes[key], nlp.ns, nlp.phase_idx)
+            external_forces = nlp.dynamics_constants[0].mx
+            for i in range(1, dynamics_constants_used_at_each_nodes[key].shape[1]):
+                external_forces = horzcat(external_forces, nlp.dynamics_constants[i].mx)
 
         ConfigureProblem.configure_q(ocp, nlp, True, False)
         ConfigureProblem.configure_qdot(ocp, nlp, True, False)
@@ -669,7 +689,7 @@ def torque_activations_driven(
         with_passive_torque: bool = False,
         with_residual_torque: bool = False,
         with_ligament: bool = False,
-        external_forces: list[Any] = None,
+        dynamics_constants_used_at_each_nodes: dict[list] = {},
     ):
         """
         Configure the dynamics for a torque driven program (states are q and qdot, controls are tau activations).
@@ -690,14 +710,21 @@ def torque_activations_driven(
             If the dynamic with a residual torque should be used
         with_ligament: bool
             If the dynamic with ligament should be used
-        external_forces: list[Any]
-            A list of external forces
-
+        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.
         """
 
         _check_contacts_in_biorbd_model(with_contact, nlp.model.nb_contacts, nlp.phase_idx)
-        _check_external_forces_format(external_forces, nlp.ns, nlp.phase_idx)
-        _check_external_forces_and_phase_dynamics(external_forces, nlp.phase_dynamics, nlp.phase_idx)
+        external_forces = None
+        for key in dynamics_constants_used_at_each_nodes.keys():
+            if key != "external_forces":
+                raise RuntimeError(
+                    "The only dynamics_constants_used_at_each_nodes allowed for torque_driven dynamics is external_forces."
+                )
+            _check_dynamics_constants_format(dynamics_constants_used_at_each_nodes[key], nlp.ns, nlp.phase_idx)
+            external_forces = nlp.dynamics_constants[0].mx
+            for i in range(1, dynamics_constants_used_at_each_nodes[key].shape[1]):
+                external_forces = horzcat(external_forces, nlp.dynamics_constants[i].mx)
 
         ConfigureProblem.configure_q(ocp, nlp, True, False)
         ConfigureProblem.configure_qdot(ocp, nlp, True, False)
@@ -787,7 +814,7 @@ def muscle_driven(
         with_passive_torque: bool = False,
         with_ligament: bool = False,
         rigidbody_dynamics: RigidBodyDynamics = RigidBodyDynamics.ODE,
-        external_forces: list[Any] = None,
+        dynamics_constants_used_at_each_nodes: dict[list] = {},
     ):
         """
         Configure the dynamics for a muscle driven program.
@@ -816,13 +843,20 @@ def muscle_driven(
             If the dynamic with ligament should be used
         rigidbody_dynamics: RigidBodyDynamics
             which rigidbody dynamics should be used
-        external_forces: list[Any]
-            A list of external forces
+        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.
         """
         _check_contacts_in_biorbd_model(with_contact, nlp.model.nb_contacts, nlp.phase_idx)
-        _check_external_forces_format(external_forces, nlp.ns, nlp.phase_idx)
-        _check_external_forces_and_phase_dynamics(external_forces, nlp.phase_dynamics, nlp.phase_idx)
-
+        external_forces = None
+        for key in dynamics_constants_used_at_each_nodes.keys():
+            if key != "external_forces":
+                raise RuntimeError(
+                    "The only dynamics_constants_used_at_each_nodes allowed for torque_driven dynamics is external_forces."
+                )
+            _check_dynamics_constants_format(dynamics_constants_used_at_each_nodes[key], nlp.ns, nlp.phase_idx)
+            external_forces = nlp.dynamics_constants[0].mx
+            for i in range(1, dynamics_constants_used_at_each_nodes[key].shape[1]):
+                external_forces = horzcat(external_forces, nlp.dynamics_constants[i].mx)
         if fatigue is not None and "tau" in fatigue and not with_residual_torque:
             raise RuntimeError("Residual torques need to be used to apply fatigue on torques")
 
@@ -924,6 +958,7 @@ def configure_dynamics_function(ocp, nlp, dyn_func, **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,
             )
@@ -941,9 +976,10 @@ def configure_dynamics_function(ocp, nlp, dyn_func, **extra_params):
                         nlp.controls.scaled.mx_reduced,
                         nlp.parameters.scaled.mx_reduced,
                         nlp.algebraic_states.scaled.mx_reduced,
+                        nlp.dynamics_constants.mx,
                     ],
                     [dynamics_dxdt],
-                    ["t_span", "x", "u", "p", "a"],
+                    ["t_span", "x", "u", "p", "a", "dynamics_constants"],
                     ["xdot"],
                 ),
             )
@@ -1903,29 +1939,19 @@ def print(self):
         raise NotImplementedError("Printing of DynamicsList is not ready yet")
 
 
-def _check_external_forces_format(external_forces: list[Any], n_shooting: int, phase_idx: int):
-    """Check if the external_forces is of the right format"""
-    if external_forces is not None and len(external_forces) != n_shooting:
+def _check_dynamics_constants_format(dynamics_constant: np.ndarray, n_shooting: int, phase_idx: int):
+    """Check if the dynamics_constant_at_each_node is of the right format"""
+    if type(dynamics_constant) is not np.ndarray:
         raise RuntimeError(
-            f"Phase {phase_idx} has {n_shooting} shooting points but the external_forces "
-            f"has {len(external_forces)} shooting points."
-            f"The external_forces should be of format list[Any] "
-            f"where the list is the number of shooting points of the phase and the dict is the Any "
-            f"is the specific way to add external_force for the specific implementation of the biomodel."
+            f"Phase {phase_idx} has dynamics_constant_at_each_node of type {type(dynamics_constant)} "
+            f"but it should be of type np.ndarray"
         )
-
-
-def _check_external_forces_and_phase_dynamics(
-    external_forces: list[Any], phase_dynamics: PhaseDynamics, phase_idx: int
-):
-    """If external_forces is not None, phase_dynamics should be PhaseDynamics.ONE_PER_NODE"""
-    # Note to the developer: External_forces doesn't necessitate ONE_PER_NODE, we made it anyway
-    # because at this stage it makes more sens to send full time series of external forces
-    # but one day someone could be interested in sending a unique value that could be applied to all nodes
-    if external_forces is not None and phase_dynamics != PhaseDynamics.ONE_PER_NODE:
+    if dynamics_constant is not None and dynamics_constant.shape[2] != n_shooting + 1:
         raise RuntimeError(
-            f"Phase {phase_idx} has external_forces but the phase_dynamics is {phase_dynamics}."
-            f"Please set phase_dynamics=PhaseDynamics.ONE_PER_NODE"
+            f"Phase {phase_idx} has {n_shooting}+1 shooting points but the dynamics_constant_at_each_node "
+            f"has {dynamics_constant.shape[2]} shooting points."
+            f"The dynamics_constant_at_each_node should be of format dict[np.ndarray] "
+            f"where the list is the number of shooting points of the phase "
         )
 
 

bioptim/dynamics/dynamics_functions.py

@@ -1,3 +1,4 @@
+import numpy as np
 from casadi import horzcat, vertcat, MX, SX
 
 from ..misc.enums import RigidBodyDynamics, DefectType
@@ -87,7 +88,7 @@ def torque_driven(
         with_friction: bool,
         rigidbody_dynamics: RigidBodyDynamics,
         fatigue: FatigueList,
-        external_forces: list = None,
+        external_forces: np.ndarray = None,
     ) -> DynamicsEvaluation:
         """
         Forward dynamics driven by joint torques, optional external forces can be declared.
@@ -118,7 +119,7 @@ def torque_driven(
             which rigidbody dynamics should be used
         fatigue : FatigueList
             A list of fatigue elements
-        external_forces: list[Any]
+        external_forces: np.ndarray
             The external forces
 
         Returns
@@ -461,7 +462,7 @@ def torque_activations_driven(
         with_passive_torque: bool,
         with_residual_torque: bool,
         with_ligament: bool,
-        external_forces: list = None,
+        external_forces: np.ndarray = None,
     ):
         """
         Forward dynamics driven by joint torques activations.
@@ -488,7 +489,7 @@ def torque_activations_driven(
             If the dynamic should be added with residual torques
         with_ligament: bool
             If the dynamic with ligament should be used
-        external_forces: list[Any]
+        external_forces: np.ndarray
             The external forces
 
         Returns
@@ -528,7 +529,7 @@ def torque_derivative_driven(
         with_contact: bool,
         with_passive_torque: bool,
         with_ligament: bool,
-        external_forces: list = None,
+        external_forces: np.ndarray = None,
     ) -> DynamicsEvaluation:
         """
         Forward dynamics driven by joint torques, optional external forces can be declared.
@@ -555,7 +556,7 @@ def torque_derivative_driven(
             If the dynamic with passive torque should be used
         with_ligament: bool
             If the dynamic with ligament should be used
-        external_forces: list[Any]
+        external_forces: np.ndarray
             The external forces
 
         Returns
@@ -605,7 +606,7 @@ def forces_from_torque_driven(
         nlp,
         with_passive_torque: bool = False,
         with_ligament: bool = False,
-        external_forces: list = None,
+        external_forces: np.ndarray = None,
     ) -> MX:
         """
         Contact forces of a forward dynamics driven by joint torques with contact constraints.
@@ -628,7 +629,7 @@ def forces_from_torque_driven(
             If the dynamic with passive torque should be used
         with_ligament: bool
             If the dynamic with ligament should be used
-        external_forces: list[Any]
+        external_forces: np.ndarray
             The external forces
 
         Returns
@@ -655,7 +656,7 @@ def forces_from_torque_activation_driven(
         nlp,
         with_passive_torque: bool = False,
         with_ligament: bool = False,
-        external_forces: list = None,
+        external_forces: np.ndarray = None,
     ) -> MX:
         """
         Contact forces of a forward dynamics driven by joint torques with contact constraints.
@@ -678,7 +679,7 @@ def forces_from_torque_activation_driven(
             If the dynamic with passive torque should be used
         with_ligament: bool
             If the dynamic with ligament should be used
-        external_forces: list[Any]
+        external_forces: np.ndarray
             The external forces
 
         Returns
@@ -709,7 +710,7 @@ def muscles_driven(
         rigidbody_dynamics: RigidBodyDynamics = RigidBodyDynamics.ODE,
         with_residual_torque: bool = False,
         fatigue=None,
-        external_forces: list = None,
+        external_forces: np.ndarray = None,
     ) -> DynamicsEvaluation:
         """
         Forward dynamics driven by muscle.
@@ -740,7 +741,7 @@ def muscles_driven(
             To define fatigue elements
         with_residual_torque: bool
             If the dynamic should be added with residual torques
-        external_forces: list[Any]
+        external_forces: np.ndarray
             The external forces
 
         Returns
@@ -1092,16 +1093,11 @@ def forward_dynamics(
 
             return qdot_var_mapping.map(qddot)
         else:
-            dxdt = MX(len(qdot_var_mapping), nlp.ns)
-            # Todo: Should be added to pass f_ext in controls (as a symoblic value)
-            #  this would avoid to create multiple equations of motions per node
-            for i, f_ext in enumerate(external_forces):
-                if with_contact:
-                    qddot = nlp.model.constrained_forward_dynamics(q, qdot, tau, f_ext)
-                else:
-                    qddot = nlp.model.forward_dynamics(q, qdot, tau, f_ext)
-                dxdt[:, i] = qdot_var_mapping.map(qddot)
-            return dxdt
+            if with_contact:
+                qddot = nlp.model.constrained_forward_dynamics(q, qdot, tau, external_forces)
+            else:
+                qddot = nlp.model.forward_dynamics(q, qdot, tau, external_forces)
+            return qdot_var_mapping.map(qddot)
 
     @staticmethod
     def inverse_dynamics(