Single phase problem and calcium summation in control

cocofest/__init__.py

@@ -3,20 +3,24 @@
 from .models.fes_model import FesModel
 from .models.ding2003 import DingModelFrequency
 from .models.ding2003_with_fatigue import DingModelFrequencyWithFatigue
-from .models.ding2007 import DingModelPulseDurationFrequency
-from .models.ding2007_with_fatigue import DingModelPulseDurationFrequencyWithFatigue
-from .models.hmed2018 import DingModelIntensityFrequency
-from .models.hmed2018_with_fatigue import DingModelIntensityFrequencyWithFatigue
+from .models.ding2007 import DingModelPulseWidthFrequency
+from .models.ding2007_with_fatigue import DingModelPulseWidthFrequencyWithFatigue
+from .models.hmed2018 import DingModelPulseIntensityFrequency
+from .models.hmed2018_with_fatigue import DingModelPulseIntensityFrequencyWithFatigue
 from .models.dynamical_model import FesMskModel
+from .models.model_maker import ModelMaker
 from .optimization.fes_ocp import OcpFes
 from .optimization.fes_identification_ocp import OcpFesId
 from .optimization.fes_ocp_dynamics import OcpFesMsk
-from .optimization.fes_ocp_nmpc_cyclic import OcpFesNmpcCyclic
+from .optimization.fes_ocp_nmpc_cyclic import NmpcFes
+from .optimization.fes_ocp_dynamics_nmpc_cyclic import NmpcFesMsk
 from .integration.ivp_fes import IvpFes
 from .fourier_approx import FourierSeries
-from .identification.ding2003_force_parameter_identification import DingModelFrequencyForceParameterIdentification
+from .identification.ding2003_force_parameter_identification import (
+    DingModelFrequencyForceParameterIdentification,
+)
 from .identification.ding2007_force_parameter_identification import (
-    DingModelPulseDurationFrequencyForceParameterIdentification,
+    DingModelPulseWidthFrequencyForceParameterIdentification,
 )
 from .identification.hmed2018_force_parameter_identification import (
     DingModelPulseIntensityFrequencyForceParameterIdentification,

cocofest/custom_constraints.py

@@ -2,42 +2,97 @@
 This class regroups all the custom constraints that are used in the optimization problem.
 """
 
-from casadi import MX, SX
+from casadi import MX, SX, vertcat
 
 from bioptim import PenaltyController
 
+from .models.hmed2018 import DingModelPulseIntensityFrequency
+
 
 class CustomConstraint:
     @staticmethod
-    def pulse_time_apparition_as_phase(controller: PenaltyController) -> MX | SX:
-        return controller.time.cx - controller.parameters["pulse_apparition_time"].cx[controller.phase_idx]
+    def cn_sum(controller: PenaltyController, stim_time: list, model_idx: int = None) -> MX | SX:
+        model = controller.model.muscles_dynamics_model[model_idx] if isinstance(model_idx, int) else controller.model
+        cn_sum_key = model.cn_sum_name
+        km_key = model.km_name
+        intensity_in_model = True if isinstance(model, DingModelPulseIntensityFrequency) else False
+        pulse_intensity_key = model.pulse_intensity_name if intensity_in_model else None
+        pulse_intensity = controller.parameters[pulse_intensity_key].cx if intensity_in_model else None
+        lambda_i = model.get_lambda_i(nb_stim=len(stim_time), pulse_intensity=pulse_intensity)
+        km = controller.states[km_key].cx if model._with_fatigue else model.km_rest
+        r0 = model.get_r0(km=km)
 
-    @staticmethod
-    def equal_to_first_pulse_interval_time(controller: PenaltyController) -> MX | SX:
-        if controller.ocp.n_phases <= 1:
-            RuntimeError("There is only one phase, the bimapping constraint is not possible")
+        return controller.controls[cn_sum_key].cx - model.cn_sum_fun(
+            r0=r0, t=controller.time.cx, t_stim_prev=stim_time, lambda_i=lambda_i
+        )
 
-        first_phase_tf = controller.ocp.node_time(0, controller.ocp.nlp[controller.phase_idx].ns)
-        current_phase_tf = controller.ocp.nlp[controller.phase_idx].node_time(
-            controller.ocp.nlp[controller.phase_idx].ns
+    @staticmethod
+    def cn_sum_identification(controller: PenaltyController, stim_time: list, stim_index: list) -> MX | SX:
+        intensity_in_model = True if isinstance(controller.model, DingModelPulseIntensityFrequency) else False
+        ar, bs, Is, cr = None, None, None, None
+        if intensity_in_model:
+            ar = controller.parameters["ar"].cx if "ar" in controller.parameters.keys() else controller.model.ar
+            bs = controller.parameters["bs"].cx if "bs" in controller.parameters.keys() else controller.model.bs
+            Is = controller.parameters["Is"].cx if "Is" in controller.parameters.keys() else controller.model.Is
+            cr = controller.parameters["cr"].cx if "cr" in controller.parameters.keys() else controller.model.cr
+        lambda_i = (
+            [
+                controller.model.lambda_i_calculation_identification(
+                    controller.parameters["pulse_intensity"].cx[i], ar, bs, Is, cr
+                )
+                for i in stim_index
+            ]
+            if intensity_in_model
+            else [1 for _ in range(len(stim_time))]
+        )
+        km = (
+            controller.parameters["km_rest"].cx
+            if "km_rest" in controller.parameters.keys()
+            else controller.model.km_rest
+        )
+        r0 = km + controller.model.r0_km_relationship
+        return controller.controls["Cn_sum"].cx - controller.model.cn_sum_fun(
+            r0=r0, t=controller.time.cx, t_stim_prev=stim_time, lambda_i=lambda_i
         )
 
-        return first_phase_tf - current_phase_tf
+    @staticmethod
+    def a_calculation(controller: PenaltyController, last_stim_index: int) -> MX | SX:
+        a = controller.states["A"].cx if controller.model.with_fatigue else controller.model.a_scale
+        last_stim_index = 0 if controller.parameters["pulse_width"].cx.shape == (1, 1) else last_stim_index
+        a_calculation = controller.model.a_calculation(
+            a_scale=a,
+            pulse_width=controller.parameters["pulse_width"].cx[last_stim_index],
+        )
+        return controller.controls["A_calculation"].cx - a_calculation
 
     @staticmethod
-    def equal_to_first_pulse_duration(controller: PenaltyController) -> MX | SX:
-        if controller.ocp.n_phases <= 1:
-            RuntimeError("There is only one phase, the bimapping constraint is not possible")
-        return (
-            controller.parameters["pulse_duration"].cx[0]
-            - controller.parameters["pulse_duration"].cx[controller.phase_idx]
+    def a_calculation_msk(controller: PenaltyController, last_stim_index: int, model_idx: int) -> MX | SX:
+        model = controller.model.muscles_dynamics_model[model_idx]
+        muscle_name = model.muscle_name
+        a = controller.states["A_" + muscle_name].cx if model.with_fatigue else model.a_scale
+        last_stim_index = (
+            0 if controller.parameters["pulse_width_" + muscle_name].cx.shape == (1, 1) else last_stim_index
         )
+        a_calculation = model.a_calculation(
+            a_scale=a,
+            pulse_width=controller.parameters["pulse_width_" + muscle_name].cx[last_stim_index],
+        )
+        return controller.controls["A_calculation_" + muscle_name].cx - a_calculation
 
     @staticmethod
-    def equal_to_first_pulse_intensity(controller: PenaltyController) -> MX | SX:
-        if controller.ocp.n_phases <= 1:
-            RuntimeError("There is only one phase, the bimapping constraint is not possible")
-        return (
-            controller.parameters["pulse_intensity"].cx[0]
-            - controller.parameters["pulse_intensity"].cx[controller.phase_idx]
+    def a_calculation_identification(controller: PenaltyController, last_stim_index: int) -> MX | SX:
+        a = (
+            controller.parameters["a_scale"].cx
+            if "a_scale" in controller.parameters.keys()
+            else controller.model.a_scale
+        )
+        pd0 = controller.parameters["pd0"].cx if "pd0" in controller.parameters.keys() else controller.model.pd0
+        pdt = controller.parameters["pdt"].cx if "pdt" in controller.parameters.keys() else controller.model.pdt
+        last_stim_index = 0 if controller.parameters["pulse_width"].cx.shape == (1, 1) else last_stim_index
+        a_calculation = controller.model.a_calculation_identification(
+            a_scale=a,
+            pulse_width=controller.parameters["pulse_width"].cx[last_stim_index],
+            pd0=pd0,
+            pdt=pdt,
         )
+        return controller.controls["A_calculation"].cx - a_calculation

cocofest/custom_objectives.py

@@ -41,7 +41,10 @@ def track_state_from_time(controller: PenaltyController, fourier_coeff: np.ndarr
 
     @staticmethod
     def track_state_from_time_interpolate(
-        controller: PenaltyController, force: np.ndarray, key: str, minimization_type: str = "least square"
+        controller: PenaltyController,
+        force: np.ndarray,
+        key: str,
+        minimization_type: str = "least square",
     ) -> MX:
         """
         Minimize the states variables.
@@ -85,10 +88,13 @@ def minimize_overall_muscle_fatigue(controller: PenaltyController) -> MX:
         The sum of each force scaling factor
         """
         muscle_name_list = controller.model.bio_model.muscle_names
+        muscle_fatigue_rest = horzcat(
+            *[controller.model.bio_stim_model[x].a_rest for x in range(1, len(muscle_name_list) + 1)]
+        )
         muscle_fatigue = horzcat(
             *[controller.states["A_" + muscle_name_list[x]].cx for x in range(len(muscle_name_list))]
         )
-        return sum1(muscle_fatigue)
+        return sum1(muscle_fatigue_rest) / sum1(muscle_fatigue)
 
     @staticmethod
     def minimize_overall_muscle_force_production(controller: PenaltyController) -> MX:

cocofest/dynamics/warm_start.py → cocofest/dynamics/initial_guess_warm_start.py

@@ -10,6 +10,7 @@
     DynamicsList,
     InitialGuessList,
     InterpolationType,
+    Node,
     ObjectiveFcn,
     ObjectiveList,
     OdeSolver,
@@ -19,10 +20,19 @@
     Solver,
 )
 
-from ..dynamics.inverse_kinematics_and_dynamics import get_circle_coord, inverse_kinematics_cycling
+from ..dynamics.inverse_kinematics_and_dynamics import (
+    get_circle_coord,
+    inverse_kinematics_cycling,
+)
 
 
-def get_initial_guess(biorbd_model_path: str, final_time: int, n_stim: int, n_shooting: int, objective: dict) -> dict:
+def get_initial_guess(
+    biorbd_model_path: str,
+    final_time: int,
+    n_shooting: int,
+    objective: dict,
+    n_threads: int,
+) -> dict:
     """
     Get the initial guess for the ocp
 
@@ -32,12 +42,12 @@ def get_initial_guess(biorbd_model_path: str, final_time: int, n_stim: int, n_sh
         The path to the model
     final_time: int
         The ocp final time
-    n_stim: int
-        The number of stimulation
     n_shooting: list
         The shooting points number
     objective: dict
         The ocp objective
+    n_threads: int
+        The number of threads
 
     Returns
     -------
@@ -50,36 +60,29 @@ def get_initial_guess(biorbd_model_path: str, final_time: int, n_stim: int, n_sh
         raise ValueError("Only a cycling objective is implemented for the warm start")
 
     # Getting q and qdot from the inverse kinematics
-    ocp, q, qdot = prepare_muscle_driven_ocp(biorbd_model_path, n_shooting[0] * n_stim, final_time, objective)
+    ocp, q, qdot = prepare_muscle_driven_ocp(biorbd_model_path, n_shooting, final_time, objective, n_threads)
 
     # Solving the ocp to get muscle controls
-    sol = ocp.solve(Solver.IPOPT(_tol=1e-4))
+    sol = ocp.solve()
     muscles_control = sol.decision_controls(to_merge=[SolutionMerge.PHASES, SolutionMerge.NODES])
     model = biorbd.Model(biorbd_model_path)
 
     # Reorganizing the q and qdot shape for the warm start
-    q_init = [q[:, n_shooting[0] * i : n_shooting[0] * (i + 1) + 1] for i in range(n_stim)]
-    qdot_init = [qdot[:, n_shooting[0] * i : n_shooting[0] * (i + 1) + 1] for i in range(n_stim)]
+    q_init = q
+    qdot_init = qdot
 
     # Building the initial guess dictionary
     states_init = {"q": q_init, "qdot": qdot_init}
 
     # Creating initial muscle forces guess from the muscle controls and the muscles max iso force characteristics
     for i in range(muscles_control["muscles"].shape[0]):
         fmax = model.muscle(i).characteristics().forceIsoMax()  # Get the max iso force of the muscle
-        states_init[model.muscle(i).name().to_string()] = [
-            np.array([muscles_control["muscles"][i][n_shooting[0] * j : n_shooting[0] * (j + 1) + 1]]) * fmax
-            for j in range(n_stim)
-        ]  # Multiply the muscle control by the max iso force to get the muscle force for each shooting point
-        states_init[model.muscle(i).name().to_string()][-1] = np.array(
-            [
-                np.append(
-                    states_init[model.muscle(i).name().to_string()][-1],
-                    states_init[model.muscle(i).name().to_string()][-1][0][-1],
-                )
-            ]
-        )  # Adding a last value to the end for each interpolation frames
-
+        states_init[model.muscle(i).name().to_string()] = np.array(muscles_control["muscles"][i]) * fmax
+        states_init[model.muscle(i).name().to_string()] = np.append(
+            states_init[model.muscle(i).name().to_string()],
+            states_init[model.muscle(i).name().to_string()][-1],
+        )
+        states_init[model.muscle(i).name().to_string()] = np.array([states_init[model.muscle(i).name().to_string()]])
     return states_init
 
 
@@ -88,6 +91,7 @@ def prepare_muscle_driven_ocp(
     n_shooting: int,
     final_time: int,
     objective: dict,
+    n_threads: int,
 ) -> tuple:
     """
     Prepare the muscle driven ocp with a cycling objective
@@ -102,6 +106,8 @@ def prepare_muscle_driven_ocp(
         The ocp final time
     objective: dict
         The ocp objective
+    n_threads: int
+        The number of threads
 
     Returns
     -------
@@ -123,24 +129,31 @@ def prepare_muscle_driven_ocp(
     y_center = objective["cycling"]["y_center"]
     radius = objective["cycling"]["radius"]
     get_circle_coord_list = np.array(
-        [get_circle_coord(theta, x_center, y_center, radius)[:-1] for theta in np.linspace(0, -2 * np.pi, n_shooting)]
+        [
+            get_circle_coord(theta, x_center, y_center, radius)[:-1]
+            for theta in np.linspace(0, -2 * np.pi, n_shooting + 1)
+        ]
     )
+
     objective_functions = ObjectiveList()
-    for i in range(n_shooting):
-        objective_functions.add(
-            ObjectiveFcn.Mayer.TRACK_MARKERS,
-            weight=100,
-            axes=[Axis.X, Axis.Y],
-            marker_index=0,
-            target=np.array(get_circle_coord_list[i]),
-            node=i,
-            phase=0,
-            quadratic=True,
-        )
+    objective_functions.add(
+        ObjectiveFcn.Mayer.TRACK_MARKERS,
+        weight=100,
+        axes=[Axis.X, Axis.Y],
+        marker_index=0,
+        target=get_circle_coord_list.T,
+        node=Node.ALL,
+        phase=0,
+        quadratic=True,
+    )
 
     # Dynamics
     dynamics = DynamicsList()
-    dynamics.add(DynamicsFcn.MUSCLE_DRIVEN, expand_dynamics=True, phase_dynamics=PhaseDynamics.SHARED_DURING_THE_PHASE)
+    dynamics.add(
+        DynamicsFcn.MUSCLE_DRIVEN,
+        expand_dynamics=True,
+        phase_dynamics=PhaseDynamics.SHARED_DURING_THE_PHASE,
+    )
 
     # Path constraint
     x_bounds = BoundsList()
@@ -175,6 +188,7 @@ def prepare_muscle_driven_ocp(
             u_init=u_init,
             objective_functions=objective_functions,
             ode_solver=OdeSolver.RK4(),
+            n_threads=n_threads,
         ),
         q_guess,
         qdot_guess,

cocofest/dynamics/inverse_kinematics_and_dynamics.py

@@ -6,7 +6,11 @@
 
 # This function gets the x, y, z circle coordinates based on the angle theta
 def get_circle_coord(
-    theta: int | float, x_center: int | float, y_center: int | float, radius: int | float, z: int | float = None
+    theta: int | float,
+    x_center: int | float,
+    y_center: int | float,
+    radius: int | float,
+    z: int | float = None,
 ) -> list:
     """
     Get the x, y, z coordinates of a circle based on the angle theta and the center of the circle