Building examples and features for the cycling motion

cocofest/__init__.py

@@ -19,3 +19,4 @@
 from .identification.hmed2018_force_parameter_identification import (
     DingModelPulseIntensityFrequencyForceParameterIdentification,
 )
+from .dynamics.inverse_kinematics_and_dynamics import get_circle_coord, inverse_kinematics_cycling, inverse_dynamics_cycling

cocofest/dynamics/inverse_kinematics_and_dynamics.py

@@ -0,0 +1,133 @@
+import math
+import numpy as np
+
+import biorbd
+
+
+# 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)-> list:
+    """
+    Get the x, y, z coordinates of a circle based on the angle theta and the center of the circle
+    
+    Parameters
+    ----------
+    theta: int | float
+        The angle of the circle in radians
+    x_center: int | float
+        The x coordinate of the center of the circle
+    y_center: int | float
+        The y coordinate of the center of the circle
+    radius: int | float
+        The radius of the circle
+    z: int | float
+        The z coordinate of the center of the circle. If None, the z coordinate of the circle will be 0
+        
+    Returns
+    ----------
+    x, y, z : list
+        The x, y, z coordinates of the circle
+    
+    """
+    x = radius * math.cos(theta) + x_center
+    y = radius * math.sin(theta) + y_center
+    if z is None:
+        return [x, y, 0]
+    else:
+        return [x, y, z]
+
+
+# This function gives the inverse kinematics q of a cycling movement for a given model
+def inverse_kinematics_cycling(
+    model_path: str,
+    n_shooting: int,
+    x_center: int|float,
+    y_center: int|float,
+    radius: int|float,
+    ik_method: str = "trf",
+) -> tuple:
+    """
+    Perform the inverse kinematics of a cycling movement
+
+    Parameters
+    ----------
+    model_path: str
+        The path to the model
+    n_shooting: int
+        The number of shooting points
+    x_center: int | float
+        The x coordinate of the center of the circle
+    y_center: int | float
+        The y coordinate of the center of the circle
+    radius: int | float
+        The radius of the circle
+    ik_method: str
+        The method to solve the inverse kinematics problem
+        If ik_method = 'lm', the 'trf' method will be used for the first frame, in order to respect the bounds of the model.
+        Then, the 'lm' method will be used for the following frames.
+        If ik_method = 'trf', the 'trf' method will be used for all the frames.
+        If ik_method = 'only_lm', the 'lm' method will be used for all the frames.
+
+        In least_square:
+            -‘trf’ : Trust Region Reflective algorithm, particularly suitable for large sparse problems
+                    with bounds.
+                    Generally robust method.
+            -‘lm’ : Levenberg-Marquardt algorithm as implemented in MINPACK.
+                    Does not handle bounds and sparse Jacobians.
+                    Usually the most efficient method for small unconstrained problems.
+
+    Returns
+    ----------
+    q : np.array
+        joints angles
+    """
+
+    model = biorbd.Model(model_path)
+    if 0 <= model.nbMarkers() > 1:
+        raise ValueError("The model must have only one markers to perform the inverse kinematics")
+
+    z = model.markers(np.array([0, 0]))[0].to_array()[2]
+    if z != model.markers(np.array([np.pi/2, np.pi/2]))[0].to_array()[2]:
+        print("The model not strictly 2d. Warm start not optimal.")
+
+    x_y_z_coord = np.array([get_circle_coord(theta, x_center, y_center, radius, z=z) for theta in np.linspace(0, -2 * np.pi + (-2 * np.pi / n_shooting), n_shooting+1)]).T
+    target_q = x_y_z_coord.reshape((3, 1, n_shooting+1))
+    ik = biorbd.InverseKinematics(model, target_q)
+    ik_q = ik.solve(method=ik_method)
+    ik_qdot = np.array([np.gradient(ik_q[i], (1/n_shooting)) for i in range(ik_q.shape[0])])
+    ik_qddot = np.array([np.gradient(ik_qdot[i], (1 / n_shooting)) for i in range(ik_qdot.shape[0])])
+    return ik_q, ik_qdot, ik_qddot
+
+
+# This function gives the inverse dynamics Tau of a cycling motion for a given model and q, qdot, qddot
+def inverse_dynamics_cycling(
+    model_path: str,
+    q: np.array,
+    qdot: np.array,
+    qddot: np.array,
+) -> np.array:
+    """
+    Perform the inverse dynamics of a cycling movement
+
+    Parameters
+    ----------
+    model_path: str
+        The path to the used model
+    q: np.array
+        joints angles
+    qdot: np.array
+        joints velocities
+    qddot: np.array
+        joints accelerations
+
+    Returns
+    ----------
+    Tau : np.array
+        joints torques
+    """
+    model = biorbd.Model(model_path)
+    tau_shape = (model.nbQ(), q.shape[1]-1)
+    tau = np.zeros(tau_shape)
+    for i in range(tau.shape[1]):
+        tau_i = model.InverseDynamics(q[:, i], qdot[:, i], qddot[:, i])
+        tau[:, i] = tau_i.to_array()
+    return tau

cocofest/dynamics/warm_start.py

@@ -0,0 +1,168 @@
+import numpy as np
+
+import biorbd
+
+from bioptim import (
+    Axis,
+    BiorbdModel,
+    BoundsList,
+    DynamicsFcn,
+    DynamicsList,
+    InitialGuessList,
+    InterpolationType,
+    ObjectiveFcn,
+    ObjectiveList,
+    OdeSolver,
+    OptimalControlProgram,
+    PhaseDynamics,
+    SolutionMerge,
+    Solver
+)
+
+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:
+    """
+    Get the initial guess for the ocp
+
+    Parameters
+    ----------
+    biorbd_model_path: str
+        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
+
+    Returns
+    -------
+    dict
+        The initial guess for the ocp
+
+    """
+    # Checking if the objective is a cycling objective
+    if objective["cycling"] is None:
+        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)
+
+    # Solving the ocp to get muscle controls
+    sol = ocp.solve(Solver.IPOPT(_tol=1e-4))
+    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)]
+
+    # 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
+
+    return states_init
+
+
+def prepare_muscle_driven_ocp(
+        biorbd_model_path: str,
+        n_shooting: int,
+        final_time: int,
+        objective: dict,
+) -> tuple:
+    """
+    Prepare the muscle driven ocp with a cycling objective
+
+    Parameters
+    ----------
+    biorbd_model_path: str
+        The path to the model
+    n_shooting: int
+        The number of shooting points
+    final_time: int
+        The ocp final time
+    objective: dict
+        The ocp objective
+
+    Returns
+    -------
+    OptimalControlProgram
+        The muscle driven ocp
+    np.array
+        The joints angles
+    np.array
+        The joints velocities
+    """
+
+    # Adding the models to the same phase
+    bio_model = BiorbdModel(biorbd_model_path, )
+
+    # Add objective functions
+    x_center = objective["cycling"]["x_center"]
+    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)])
+    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,
+        )
+
+    # Dynamics
+    dynamics = DynamicsList()
+    dynamics.add(DynamicsFcn.MUSCLE_DRIVEN, expand_dynamics=True,
+                 phase_dynamics=PhaseDynamics.SHARED_DURING_THE_PHASE)
+
+    # Path constraint
+    x_bounds = BoundsList()
+    q_x_bounds = bio_model.bounds_from_ranges("q")
+    qdot_x_bounds = bio_model.bounds_from_ranges("qdot")
+    x_bounds.add(key="q", bounds=q_x_bounds, phase=0)
+    x_bounds.add(key="qdot", bounds=qdot_x_bounds, phase=0)
+
+    # Define control path constraint
+    u_bounds = BoundsList()
+    u_bounds["muscles"] = [0] * bio_model.nb_muscles, [1] * bio_model.nb_muscles
+
+    # Initial q guess
+    x_init = InitialGuessList()
+    u_init = InitialGuessList()
+
+    q_guess, qdot_guess, qddotguess = inverse_kinematics_cycling(biorbd_model_path, n_shooting, x_center,
+                                                                 y_center, radius, ik_method="trf")
+    x_init.add("q", q_guess, interpolation=InterpolationType.EACH_FRAME)
+    x_init.add("qdot", qdot_guess, interpolation=InterpolationType.EACH_FRAME)
+
+    return OptimalControlProgram(
+        bio_model,
+        dynamics,
+        n_shooting,
+        final_time,
+        x_bounds=x_bounds,
+        u_bounds=u_bounds,
+        x_init=x_init,
+        u_init=u_init,
+        objective_functions=objective_functions,
+        ode_solver=OdeSolver.RK4(),
+    ), q_guess, qdot_guess