diff --git a/cocofest/__init__.py b/cocofest/__init__.py index 62a3c02..d390212 100644 --- a/cocofest/__init__.py +++ b/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, diff --git a/cocofest/custom_constraints.py b/cocofest/custom_constraints.py index dee63ac..fdaf5ce 100644 --- a/cocofest/custom_constraints.py +++ b/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 diff --git a/cocofest/custom_objectives.py b/cocofest/custom_objectives.py index 79d7c43..fc2fefd 100644 --- a/cocofest/custom_objectives.py +++ b/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: diff --git a/cocofest/dynamics/warm_start.py b/cocofest/dynamics/initial_guess_warm_start.py similarity index 68% rename from cocofest/dynamics/warm_start.py rename to cocofest/dynamics/initial_guess_warm_start.py index 06da0ea..9e85dd4 100644 --- a/cocofest/dynamics/warm_start.py +++ b/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,16 +60,16 @@ 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} @@ -67,19 +77,12 @@ def get_initial_guess(biorbd_model_path: str, final_time: int, n_stim: int, n_sh # 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, diff --git a/cocofest/dynamics/inverse_kinematics_and_dynamics.py b/cocofest/dynamics/inverse_kinematics_and_dynamics.py index c384c28..243bc54 100644 --- a/cocofest/dynamics/inverse_kinematics_and_dynamics.py +++ b/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 diff --git a/cocofest/examples/dynamics/cycling/cycling_fes_driven.py b/cocofest/examples/dynamics/cycling/cycling_fes_driven.py index bd5f37b..107452e 100644 --- a/cocofest/examples/dynamics/cycling/cycling_fes_driven.py +++ b/cocofest/examples/dynamics/cycling/cycling_fes_driven.py @@ -1,59 +1,59 @@ """ -This example will do an optimal control program of a 40 stimulation example with Ding's 2007 pulse duration model. +This example will do an optimal control program of a 10 stimulation example with Ding's 2007 pulse width model. Those ocp were build to produce a cycling motion. -The stimulation frequency will be set to 40Hz and pulse duration will be optimized to satisfy the motion meanwhile +The stimulation frequency will be set to 10Hz and pulse width will be optimized to satisfy the motion meanwhile reducing residual torque. """ -import pickle - import numpy as np -from bioptim import CostType, ObjectiveFcn, ObjectiveList, SolutionMerge, Solver +from bioptim import CostType, Solver import biorbd -from pyorerun import BiorbdModel, PhaseRerun - -from cocofest import DingModelPulseDurationFrequencyWithFatigue, OcpFesMsk, PlotCyclingResult, SolutionToPickle +from cocofest import ( + DingModelPulseWidthFrequencyWithFatigue, + OcpFesMsk, + PlotCyclingResult, + SolutionToPickle, + FesMskModel, + PickleAnimate, +) def main(): - n_stim = 40 - n_shooting = 10 - - objective_functions = ObjectiveList() - for i in range(n_stim): - objective_functions.add(ObjectiveFcn.Lagrange.MINIMIZE_CONTROL, key="tau", weight=100, quadratic=True, phase=i) - - minimum_pulse_duration = DingModelPulseDurationFrequencyWithFatigue().pd0 + minimum_pulse_width = DingModelPulseWidthFrequencyWithFatigue().pd0 + + model = FesMskModel( + name=None, + biorbd_path="../../msk_models/simplified_UL_Seth.bioMod", + muscles_model=[ + DingModelPulseWidthFrequencyWithFatigue(muscle_name="DeltoideusClavicle_A"), + DingModelPulseWidthFrequencyWithFatigue(muscle_name="DeltoideusScapula_P"), + DingModelPulseWidthFrequencyWithFatigue(muscle_name="TRIlong"), + DingModelPulseWidthFrequencyWithFatigue(muscle_name="BIC_long"), + DingModelPulseWidthFrequencyWithFatigue(muscle_name="BIC_brevis"), + ], + activate_force_length_relationship=True, + activate_force_velocity_relationship=True, + activate_residual_torque=True, + ) ocp = OcpFesMsk.prepare_ocp( - biorbd_model_path="../../msk_models/simplified_UL_Seth.bioMod", - fes_muscle_models=[ - DingModelPulseDurationFrequencyWithFatigue(muscle_name="DeltoideusClavicle_A"), - DingModelPulseDurationFrequencyWithFatigue(muscle_name="DeltoideusScapula_P"), - DingModelPulseDurationFrequencyWithFatigue(muscle_name="TRIlong"), - DingModelPulseDurationFrequencyWithFatigue(muscle_name="BIC_long"), - DingModelPulseDurationFrequencyWithFatigue(muscle_name="BIC_brevis"), - ], - n_stim=n_stim, - n_shooting=n_shooting, + model=model, + stim_time=list(np.round(np.linspace(0, 1, 11), 3))[:-1], final_time=1, - pulse_duration={ - "min": minimum_pulse_duration, + pulse_width={ + "min": minimum_pulse_width, "max": 0.0006, "bimapping": False, }, - with_residual_torque=True, + msk_info={"with_residual_torque": True}, objective={ - "custom": objective_functions, "cycling": {"x_center": 0.35, "y_center": 0, "radius": 0.1, "target": "marker"}, + "minimize_residual_torque": True, }, - activate_force_length_relationship=True, - activate_force_velocity_relationship=True, - minimize_muscle_fatigue=False, - warm_start=False, + initial_guess_warm_start=False, n_threads=5, ) ocp.add_plot_penalty(CostType.ALL) @@ -61,16 +61,7 @@ def main(): SolutionToPickle(sol, "cycling_fes_driven_min_residual_torque.pkl", "").pickle() biorbd_model = biorbd.Model("../../msk_models/simplified_UL_Seth_full_mesh.bioMod") - prr_model = BiorbdModel.from_biorbd_object(biorbd_model) - - nb_frames = 440 - nb_seconds = 1 - t_span = np.linspace(0, nb_seconds, nb_frames) - - viz = PhaseRerun(t_span) - q = sol.decision_states(to_merge=[SolutionMerge.PHASES, SolutionMerge.NODES])["q"] - viz.add_animated_model(prr_model, q) - viz.rerun("msk_model") + PickleAnimate("cycling_fes_driven_min_residual_torque.pkl").animate(model=biorbd_model) sol.graphs(show_bounds=False) PlotCyclingResult(sol).plot(starting_location="E") diff --git a/cocofest/examples/dynamics/cycling/cycling_muscle_driven.py b/cocofest/examples/dynamics/cycling/cycling_muscle_driven.py index 42be258..19bbd7c 100644 --- a/cocofest/examples/dynamics/cycling/cycling_muscle_driven.py +++ b/cocofest/examples/dynamics/cycling/cycling_muscle_driven.py @@ -30,7 +30,7 @@ def prepare_ocp( n_shooting: int, final_time: int, objective: dict, - warm_start: bool = False, + initial_guess_warm_start: bool = False, ) -> OptimalControlProgram: # Adding the model @@ -62,7 +62,11 @@ def prepare_ocp( # 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() @@ -79,7 +83,7 @@ def prepare_ocp( x_init = InitialGuessList() u_init = InitialGuessList() # If warm start, the initial guess is the result of the inverse kinematics - if warm_start: + if initial_guess_warm_start: q_guess, qdot_guess, qddotguess = inverse_kinematics_cycling( biorbd_model_path, n_shooting, x_center, y_center, radius, ik_method="trf" ) @@ -108,7 +112,7 @@ def main(): n_shooting=100, final_time=1, objective={"cycling": {"x_center": 0.35, "y_center": 0, "radius": 0.1}}, - warm_start=True, + initial_guess_warm_start=True, ) ocp.add_plot_penalty(CostType.ALL) sol = ocp.solve(Solver.IPOPT(show_online_optim=True)) diff --git a/cocofest/examples/dynamics/cycling/cycling_torque_driven.py b/cocofest/examples/dynamics/cycling/cycling_torque_driven.py index d1464bb..3087710 100644 --- a/cocofest/examples/dynamics/cycling/cycling_torque_driven.py +++ b/cocofest/examples/dynamics/cycling/cycling_torque_driven.py @@ -22,7 +22,11 @@ PhaseDynamics, ) -from cocofest import get_circle_coord, inverse_kinematics_cycling, inverse_dynamics_cycling +from cocofest import ( + get_circle_coord, + inverse_kinematics_cycling, + inverse_dynamics_cycling, +) def prepare_ocp( @@ -30,7 +34,7 @@ def prepare_ocp( n_shooting: int, final_time: int, objective: dict, - warm_start: bool = False, + initial_guess_warm_start: bool = False, ) -> OptimalControlProgram: # Adding the model @@ -59,7 +63,11 @@ def prepare_ocp( # Dynamics dynamics = DynamicsList() - dynamics.add(DynamicsFcn.TORQUE_DRIVEN, expand_dynamics=True, phase_dynamics=PhaseDynamics.SHARED_DURING_THE_PHASE) + dynamics.add( + DynamicsFcn.TORQUE_DRIVEN, + expand_dynamics=True, + phase_dynamics=PhaseDynamics.SHARED_DURING_THE_PHASE, + ) # Path constraint x_bounds = BoundsList() @@ -76,7 +84,7 @@ def prepare_ocp( x_init = InitialGuessList() u_init = InitialGuessList() # If warm start, the initial guess is the result of the inverse kinematics and dynamics - if warm_start: + if initial_guess_warm_start: q_guess, qdot_guess, qddotguess = inverse_kinematics_cycling( biorbd_model_path, n_shooting, x_center, y_center, radius, ik_method="trf" ) @@ -107,7 +115,7 @@ def main(): n_shooting=100, final_time=1, objective={"cycling": {"x_center": 0.35, "y_center": 0, "radius": 0.1}}, - warm_start=True, + initial_guess_warm_start=True, ) ocp.add_plot_penalty(CostType.ALL) sol = ocp.solve() diff --git a/cocofest/examples/dynamics/cycling/inverse_kinematics_cycling.py b/cocofest/examples/dynamics/cycling/inverse_kinematics_cycling.py index 9b9f69d..83ae434 100644 --- a/cocofest/examples/dynamics/cycling/inverse_kinematics_cycling.py +++ b/cocofest/examples/dynamics/cycling/inverse_kinematics_cycling.py @@ -21,7 +21,13 @@ def main(show_plot=True, animate=True): get_circle_coord_list = np.array( [get_circle_coord(theta, 0.35, 0, 0.1, z) for theta in np.linspace(0, -2 * np.pi, n_frames)] ) - target_q = np.array([[get_circle_coord_list[:, 0]], [get_circle_coord_list[:, 1]], [get_circle_coord_list[:, 2]]]) + target_q = np.array( + [ + [get_circle_coord_list[:, 0]], + [get_circle_coord_list[:, 1]], + [get_circle_coord_list[:, 2]], + ] + ) # Perform the inverse kinematics ik = biorbd.InverseKinematics(model, target_q) diff --git a/cocofest/examples/dynamics/frequency_optimization_musculoskeletal_dynamic_2dof.py b/cocofest/examples/dynamics/frequency_optimization_musculoskeletal_dynamic_2dof.py index a6fecbe..beebd7c 100644 --- a/cocofest/examples/dynamics/frequency_optimization_musculoskeletal_dynamic_2dof.py +++ b/cocofest/examples/dynamics/frequency_optimization_musculoskeletal_dynamic_2dof.py @@ -5,33 +5,34 @@ elbow torque control. """ -from bioptim import ( - ObjectiveFcn, - ObjectiveList, - Solver, -) - -from cocofest import DingModelFrequencyWithFatigue, OcpFesMsk +import numpy as np +from cocofest import DingModelFrequencyWithFatigue, OcpFesMsk, FesMskModel -objective_functions = ObjectiveList() -n_stim = 10 -for i in range(n_stim): - objective_functions.add(ObjectiveFcn.Lagrange.MINIMIZE_CONTROL, key="tau", weight=1, quadratic=True, phase=i) +model = FesMskModel( + name=None, + biorbd_path="../msk_models/arm26_biceps.bioMod", + muscles_model=[DingModelFrequencyWithFatigue(muscle_name="BIClong")], + activate_force_length_relationship=True, + activate_force_velocity_relationship=True, + activate_residual_torque=True, +) ocp = OcpFesMsk.prepare_ocp( - biorbd_model_path="../msk_models/arm26_biceps.bioMod", - bound_type="start_end", - bound_data=[[0, 5], [0, 120]], - fes_muscle_models=[DingModelFrequencyWithFatigue(muscle_name="BIClong")], - n_stim=n_stim, - n_shooting=10, + model=model, + stim_time=np.linspace(0, 1, 11)[:-1], final_time=1, pulse_event={"min": 0.01, "max": 0.1, "bimapping": True}, - objective={"custom": objective_functions}, - with_residual_torque=True, + objective={"minimize_residual_torque": True}, + msk_info={ + "with_residual_torque": True, + "bound_type": "start_end", + "bound_data": [[0, 5], [0, 120]], + }, + n_threads=5, + use_sx=False, ) -sol = ocp.solve(Solver.IPOPT(_max_iter=1000)) +sol = ocp.solve() sol.animate() sol.graphs(show_bounds=False) diff --git a/cocofest/examples/dynamics/intensity_optimization_cycling_multi_muscle.py b/cocofest/examples/dynamics/intensity_optimization_cycling_multi_muscle.py index 1cb5905..70edb9c 100644 --- a/cocofest/examples/dynamics/intensity_optimization_cycling_multi_muscle.py +++ b/cocofest/examples/dynamics/intensity_optimization_cycling_multi_muscle.py @@ -7,16 +7,8 @@ import numpy as np -from bioptim import ( - ObjectiveFcn, - ObjectiveList, - Solver, -) - -from cocofest import DingModelIntensityFrequency, OcpFesMsk - +from cocofest import DingModelPulseIntensityFrequency, OcpFesMsk, FesMskModel -n_stim = 30 track_q = [ np.array([0, 0.125, 0.25, 0.375, 0.5, 0.625, 0.75, 0.875, 1]), @@ -26,27 +18,27 @@ ], ] -objective_functions = ObjectiveList() -for i in range(n_stim): - objective_functions.add(ObjectiveFcn.Lagrange.MINIMIZE_CONTROL, key="tau", weight=1, quadratic=True, phase=i) - - -minimum_pulse_intensity = DingModelIntensityFrequency.min_pulse_intensity(DingModelIntensityFrequency()) +minimum_pulse_intensity = DingModelPulseIntensityFrequency.min_pulse_intensity(DingModelPulseIntensityFrequency()) + +model = FesMskModel( + name=None, + biorbd_path="../msk_models/arm26.bioMod", + muscles_model=[ + DingModelPulseIntensityFrequency(muscle_name="BIClong"), + DingModelPulseIntensityFrequency(muscle_name="BICshort"), + DingModelPulseIntensityFrequency(muscle_name="TRIlong"), + DingModelPulseIntensityFrequency(muscle_name="TRIlat"), + DingModelPulseIntensityFrequency(muscle_name="TRImed"), + DingModelPulseIntensityFrequency(muscle_name="BRA"), + ], + activate_force_length_relationship=True, + activate_force_velocity_relationship=True, + activate_residual_torque=True, +) ocp = OcpFesMsk.prepare_ocp( - biorbd_model_path="../msk_models/arm26.bioMod", - bound_type="start_end", - bound_data=[[65, 38], [65, 38]], - fes_muscle_models=[ - DingModelIntensityFrequency(muscle_name="BIClong"), - DingModelIntensityFrequency(muscle_name="BICshort"), - DingModelIntensityFrequency(muscle_name="TRIlong"), - DingModelIntensityFrequency(muscle_name="TRIlat"), - DingModelIntensityFrequency(muscle_name="TRImed"), - DingModelIntensityFrequency(muscle_name="BRA"), - ], - n_stim=n_stim, - n_shooting=5, + model=model, + stim_time=list(np.round(np.linspace(0, 1, 31), 3))[:-1], final_time=1, pulse_event={"min": 0.05, "max": 1, "bimapping": True}, pulse_intensity={ @@ -54,12 +46,16 @@ "max": 130, "bimapping": False, }, - with_residual_torque=True, - objective={"custom": objective_functions, "q_tracking": track_q}, - use_sx=True, + msk_info={ + "with_residual_torque": True, + "bound_type": "start_end", + "bound_data": [[65, 38], [65, 38]], + }, + objective={"minimize_residual_torque": True, "q_tracking": track_q}, + use_sx=False, + n_threads=5, ) -sol = ocp.solve(Solver.IPOPT(_max_iter=1000)) +sol = ocp.solve() sol.animate() sol.graphs(show_bounds=False) -print(sol.parameters) diff --git a/cocofest/examples/dynamics/intensity_optimization_hold_position.py b/cocofest/examples/dynamics/intensity_optimization_hold_position.py index a7acdd3..da881ea 100644 --- a/cocofest/examples/dynamics/intensity_optimization_hold_position.py +++ b/cocofest/examples/dynamics/intensity_optimization_hold_position.py @@ -11,49 +11,48 @@ Node, ObjectiveFcn, ObjectiveList, - Solver, ) -from cocofest import DingModelIntensityFrequencyWithFatigue, OcpFesMsk - +from cocofest import DingModelPulseIntensityFrequencyWithFatigue, OcpFesMsk, FesMskModel +n_shooting = 100 objective_functions = ObjectiveList() -n_stim = 10 -n_shooting = 10 -for i in range(n_stim): - objective_functions.add(ObjectiveFcn.Lagrange.MINIMIZE_CONTROL, key="tau", weight=10, quadratic=True, phase=i) - objective_functions.add( - ObjectiveFcn.Mayer.MINIMIZE_STATE, - key="q", - index=[0], - node=Node.ALL, - target=np.array([[1.57]] * (n_shooting + 1)).T, - weight=10, - quadratic=True, - phase=i, - ) - objective_functions.add( - ObjectiveFcn.Mayer.MINIMIZE_STATE, - key="qdot", - index=[0], - node=Node.ALL, - target=np.array([[0]] * (n_shooting + 1)).T, - weight=10, - quadratic=True, - phase=i, - ) +objective_functions.add( + ObjectiveFcn.Mayer.MINIMIZE_STATE, + key="q", + index=[0], + node=Node.ALL, + target=np.array([[1.57]] * (n_shooting + 1)).T, + weight=10, + quadratic=True, + phase=0, +) +objective_functions.add( + ObjectiveFcn.Mayer.MINIMIZE_STATE, + key="qdot", + index=[0], + node=Node.ALL, + target=np.array([[0]] * (n_shooting + 1)).T, + weight=10, + quadratic=True, + phase=0, +) -minimum_pulse_intensity = DingModelIntensityFrequencyWithFatigue.min_pulse_intensity( - DingModelIntensityFrequencyWithFatigue() +minimum_pulse_intensity = DingModelPulseIntensityFrequencyWithFatigue.min_pulse_intensity( + DingModelPulseIntensityFrequencyWithFatigue() +) +model = FesMskModel( + name=None, + biorbd_path="../msk_models/arm26_biceps_1dof.bioMod", + muscles_model=[DingModelPulseIntensityFrequencyWithFatigue(muscle_name="BIClong")], + activate_force_length_relationship=True, + activate_force_velocity_relationship=True, + activate_residual_torque=False, ) ocp = OcpFesMsk.prepare_ocp( - biorbd_model_path="../msk_models/arm26_biceps_1dof.bioMod", - bound_type="start", - bound_data=[90], - fes_muscle_models=[DingModelIntensityFrequencyWithFatigue(muscle_name="BIClong")], - n_stim=n_stim, - n_shooting=10, + model=model, + stim_time=np.linspace(0, 1, 11)[:-1], final_time=1, pulse_event={"min": 0.1, "max": 1, "bimapping": True}, pulse_intensity={ @@ -61,10 +60,12 @@ "max": 130, "bimapping": False, }, - with_residual_torque=False, + msk_info={"with_residual_torque": False, "bound_type": "start", "bound_data": [90]}, objective={"custom": objective_functions}, + use_sx=False, + n_threads=5, ) -sol = ocp.solve(Solver.IPOPT(_max_iter=1000)) +sol = ocp.solve() sol.animate() sol.graphs(show_bounds=False) diff --git a/cocofest/examples/dynamics/intensity_optimization_hold_position_multi_muscle.py b/cocofest/examples/dynamics/intensity_optimization_hold_position_multi_muscle.py index 6a9d62a..bb2541e 100644 --- a/cocofest/examples/dynamics/intensity_optimization_hold_position_multi_muscle.py +++ b/cocofest/examples/dynamics/intensity_optimization_hold_position_multi_muscle.py @@ -7,59 +7,56 @@ import numpy as np -from bioptim import ( - Node, - ObjectiveFcn, - ObjectiveList, - Solver, -) - -from cocofest import DingModelIntensityFrequencyWithFatigue, OcpFesMsk +from bioptim import Node, ObjectiveFcn, ObjectiveList +from cocofest import DingModelPulseIntensityFrequencyWithFatigue, OcpFesMsk, FesMskModel +n_shooting = 100 objective_functions = ObjectiveList() -n_stim = 10 -n_shooting = 10 -for i in range(n_stim): - objective_functions.add( - ObjectiveFcn.Mayer.MINIMIZE_STATE, - key="q", - index=[0, 1], - node=Node.ALL, - target=np.array([[0, 1.57]] * (n_shooting + 1)).T, - weight=10, - quadratic=True, - phase=i, - ) - objective_functions.add( - ObjectiveFcn.Mayer.MINIMIZE_STATE, - key="qdot", - index=[0, 1], - node=Node.ALL, - target=np.array([[0, 0]] * (n_shooting + 1)).T, - weight=10, - quadratic=True, - phase=i, - ) -minimum_pulse_intensity = DingModelIntensityFrequencyWithFatigue.min_pulse_intensity( - DingModelIntensityFrequencyWithFatigue() +objective_functions.add( + ObjectiveFcn.Mayer.MINIMIZE_STATE, + key="q", + index=[0, 1], + node=Node.ALL, + target=np.array([[0, 1.57]] * (n_shooting + 1)).T, + weight=10, + quadratic=True, + phase=0, +) +objective_functions.add( + ObjectiveFcn.Mayer.MINIMIZE_STATE, + key="qdot", + index=[0, 1], + node=Node.ALL, + target=np.array([[0, 0]] * (n_shooting + 1)).T, + weight=10, + quadratic=True, + phase=0, ) -ocp = OcpFesMsk.prepare_ocp( - biorbd_model_path="../msk_models/arm26.bioMod", - bound_type="start", - bound_data=[0, 90], - fes_muscle_models=[ - DingModelIntensityFrequencyWithFatigue(muscle_name="BIClong"), - DingModelIntensityFrequencyWithFatigue(muscle_name="BICshort"), - DingModelIntensityFrequencyWithFatigue(muscle_name="TRIlong"), - DingModelIntensityFrequencyWithFatigue(muscle_name="TRIlat"), - DingModelIntensityFrequencyWithFatigue(muscle_name="TRImed"), - DingModelIntensityFrequencyWithFatigue(muscle_name="BRA"), +minimum_pulse_intensity = DingModelPulseIntensityFrequencyWithFatigue.min_pulse_intensity( + DingModelPulseIntensityFrequencyWithFatigue() +) +model = FesMskModel( + name=None, + biorbd_path="../msk_models/arm26.bioMod", + muscles_model=[ + DingModelPulseIntensityFrequencyWithFatigue(muscle_name="BIClong"), + DingModelPulseIntensityFrequencyWithFatigue(muscle_name="BICshort"), + DingModelPulseIntensityFrequencyWithFatigue(muscle_name="TRIlong"), + DingModelPulseIntensityFrequencyWithFatigue(muscle_name="TRIlat"), + DingModelPulseIntensityFrequencyWithFatigue(muscle_name="TRImed"), + DingModelPulseIntensityFrequencyWithFatigue(muscle_name="BRA"), ], - n_stim=n_stim, - n_shooting=10, + activate_force_length_relationship=True, + activate_force_velocity_relationship=True, + activate_residual_torque=False, +) + +ocp = OcpFesMsk.prepare_ocp( + model=model, + stim_time=np.linspace(0, 1, 11)[:-1], final_time=1, pulse_event={"min": 0.1, "max": 1, "bimapping": True}, pulse_intensity={ @@ -67,10 +64,16 @@ "max": 130, "bimapping": False, }, - with_residual_torque=False, objective={"custom": objective_functions}, + msk_info={ + "with_residual_torque": False, + "bound_type": "start", + "bound_data": [0, 90], + }, + use_sx=False, + n_threads=5, ) -sol = ocp.solve(Solver.IPOPT(_max_iter=1000)) +sol = ocp.solve() sol.animate() sol.graphs(show_bounds=False) diff --git a/cocofest/examples/dynamics/minimize_fatigue/frequecy_optimization_minimize_fatigue.py b/cocofest/examples/dynamics/minimize_fatigue/frequecy_optimization_minimize_fatigue.py index 8c5394a..21e4f5a 100644 --- a/cocofest/examples/dynamics/minimize_fatigue/frequecy_optimization_minimize_fatigue.py +++ b/cocofest/examples/dynamics/minimize_fatigue/frequecy_optimization_minimize_fatigue.py @@ -10,17 +10,9 @@ import numpy as np -from bioptim import ( - Node, - ObjectiveFcn, - ObjectiveList, - Solver, -) - -from cocofest import DingModelFrequencyWithFatigue, OcpFesMsk +from bioptim import Node, ObjectiveFcn, ObjectiveList, Solver +from cocofest import DingModelFrequencyWithFatigue, OcpFesMsk, FesMskModel -n_stim = 5 -n_shooting = 10 objective_functions = ObjectiveList() objective_functions.add( ObjectiveFcn.Mayer.MINIMIZE_STATE, @@ -30,29 +22,37 @@ target=np.array([[0, 0]]).T, weight=100, quadratic=True, - phase=n_stim - 1, + phase=0, ) -for i in range(n_stim): - objective_functions.add(ObjectiveFcn.Lagrange.MINIMIZE_CONTROL, key="tau", weight=10000, quadratic=True, phase=i) - -ocp = OcpFesMsk.prepare_ocp( - biorbd_model_path="../../msk_models/arm26_biceps_triceps.bioMod", - bound_type="start_end", - bound_data=[[0, 5], [0, 90]], - fes_muscle_models=[ +model = FesMskModel( + name=None, + biorbd_path="../../msk_models/arm26_biceps_triceps.bioMod", + muscles_model=[ DingModelFrequencyWithFatigue(muscle_name="BIClong"), DingModelFrequencyWithFatigue(muscle_name="TRIlong"), ], - n_stim=n_stim, - n_shooting=n_shooting, + activate_force_length_relationship=True, + activate_force_velocity_relationship=True, + activate_residual_torque=True, +) + +ocp = OcpFesMsk.prepare_ocp( + model=model, + stim_time=list(np.linspace(0, 1, 11))[:-1], final_time=1, pulse_event={"min": 0.01, "max": 1, "bimapping": False}, - with_residual_torque=True, - objective={"custom": objective_functions}, - activate_force_length_relationship=True, - activate_force_velocity_relationship=False, - minimize_muscle_fatigue=True, + objective={ + "custom": objective_functions, + "minimize_residual_torque": True, + "minimize_fatigue": True, + }, + msk_info={ + "with_residual_torque": True, + "bound_type": "start_end", + "bound_data": [[0, 5], [0, 90]], + }, + n_threads=5, ) sol = ocp.solve(Solver.IPOPT(_max_iter=1000)) diff --git a/cocofest/examples/dynamics/minimize_fatigue/pulse_duration_optimization_minimize_fatigue.py b/cocofest/examples/dynamics/minimize_fatigue/pulse_duration_optimization_minimize_fatigue.py index cdb7a72..553a122 100644 --- a/cocofest/examples/dynamics/minimize_fatigue/pulse_duration_optimization_minimize_fatigue.py +++ b/cocofest/examples/dynamics/minimize_fatigue/pulse_duration_optimization_minimize_fatigue.py @@ -1,23 +1,16 @@ """ -This example will do a 10 stimulation example with Ding's 2007 pulse duration model. +This example will do a 10 stimulation example with Ding's 2007 pulse width model. Those ocp were build to move the elbow from 0 to 90 degrees angle. -The stimulation frequency will be set to 10Hz and pulse duration will be optimized to satisfy the motion and to minimize the overall muscle fatigue. +The stimulation frequency will be set to 10Hz and pulse width will be optimized to satisfy the motion and to minimize the overall muscle fatigue. Intensity can be optimized from sensitivity threshold to 600us. No residual torque is allowed. """ import numpy as np -from bioptim import ( - Node, - ObjectiveFcn, - ObjectiveList, - Solver, -) +from bioptim import Node, ObjectiveFcn, ObjectiveList, Solver -from cocofest import DingModelPulseDurationFrequencyWithFatigue, OcpFesMsk +from cocofest import DingModelPulseWidthFrequencyWithFatigue, OcpFesMsk, FesMskModel -n_stim = 10 -n_shooting = 10 objective_functions = ObjectiveList() objective_functions.add( ObjectiveFcn.Mayer.MINIMIZE_STATE, @@ -27,32 +20,38 @@ target=np.array([[0, 0]]).T, weight=100, quadratic=True, - phase=n_stim - 1, + phase=0, ) -minimum_pulse_duration = DingModelPulseDurationFrequencyWithFatigue().pd0 +minimum_pulse_width = DingModelPulseWidthFrequencyWithFatigue().pd0 +model = FesMskModel( + name=None, + biorbd_path="../../msk_models/arm26_biceps_triceps.bioMod", + muscles_model=[ + DingModelPulseWidthFrequencyWithFatigue(muscle_name="BIClong"), + DingModelPulseWidthFrequencyWithFatigue(muscle_name="TRIlong"), + ], + activate_force_length_relationship=True, + activate_force_velocity_relationship=True, + activate_residual_torque=False, +) ocp = OcpFesMsk.prepare_ocp( - biorbd_model_path="../../msk_models/arm26_biceps_triceps.bioMod", - bound_type="start_end", - bound_data=[[0, 5], [0, 90]], - fes_muscle_models=[ - DingModelPulseDurationFrequencyWithFatigue(muscle_name="BIClong"), - DingModelPulseDurationFrequencyWithFatigue(muscle_name="TRIlong"), - ], - n_stim=n_stim, - n_shooting=10, + model=model, + stim_time=list(np.linspace(0, 1, 11))[:-1], final_time=1, - pulse_duration={ - "min": minimum_pulse_duration, + pulse_width={ + "min": minimum_pulse_width, "max": 0.0006, "bimapping": False, }, - with_residual_torque=False, - objective={"custom": objective_functions}, - activate_force_length_relationship=True, - activate_force_velocity_relationship=False, - minimize_muscle_fatigue=True, + objective={"custom": objective_functions, "minimize_fatigue": True}, + msk_info={ + "bound_type": "start_end", + "bound_data": [[0, 5], [0, 90]], + "with_residual_torque": False, + }, + n_threads=5, ) sol = ocp.solve(Solver.IPOPT(_max_iter=3000)) diff --git a/cocofest/examples/dynamics/minimize_fatigue/pulse_intensity_optimization_minimize_fatigue.py b/cocofest/examples/dynamics/minimize_fatigue/pulse_intensity_optimization_minimize_fatigue.py index 6dbe829..f3ff18c 100644 --- a/cocofest/examples/dynamics/minimize_fatigue/pulse_intensity_optimization_minimize_fatigue.py +++ b/cocofest/examples/dynamics/minimize_fatigue/pulse_intensity_optimization_minimize_fatigue.py @@ -7,17 +7,11 @@ import numpy as np -from bioptim import ( - Node, - ObjectiveFcn, - ObjectiveList, - Solver, -) +from bioptim import Node, ObjectiveFcn, ObjectiveList, Solver + +from cocofest import DingModelPulseIntensityFrequencyWithFatigue, OcpFesMsk, FesMskModel -from cocofest import DingModelIntensityFrequencyWithFatigue, OcpFesMsk -n_stim = 10 -n_shooting = 10 objective_functions = ObjectiveList() objective_functions.add( ObjectiveFcn.Mayer.MINIMIZE_STATE, @@ -27,34 +21,40 @@ target=np.array([[0, 0]]).T, weight=100, quadratic=True, - phase=n_stim - 1, + phase=0, ) -minimum_pulse_intensity = DingModelIntensityFrequencyWithFatigue.min_pulse_intensity( - DingModelIntensityFrequencyWithFatigue() +minimum_pulse_intensity = DingModelPulseIntensityFrequencyWithFatigue.min_pulse_intensity( + DingModelPulseIntensityFrequencyWithFatigue() +) +model = FesMskModel( + name=None, + biorbd_path="../../msk_models/arm26_biceps_triceps.bioMod", + muscles_model=[ + DingModelPulseIntensityFrequencyWithFatigue(muscle_name="BIClong"), + DingModelPulseIntensityFrequencyWithFatigue(muscle_name="TRIlong"), + ], + activate_force_length_relationship=True, + activate_force_velocity_relationship=True, + activate_residual_torque=False, ) ocp = OcpFesMsk.prepare_ocp( - biorbd_model_path="../../msk_models/arm26_biceps_triceps.bioMod", - bound_type="start_end", - bound_data=[[0, 5], [0, 90]], - fes_muscle_models=[ - DingModelIntensityFrequencyWithFatigue(muscle_name="BIClong"), - DingModelIntensityFrequencyWithFatigue(muscle_name="TRIlong"), - ], - n_stim=n_stim, - n_shooting=10, + model=model, + stim_time=list(np.linspace(0, 1, 11))[:-1], final_time=1, pulse_intensity={ "min": minimum_pulse_intensity, "max": 130, "bimapping": False, }, - with_residual_torque=False, - objective={"custom": objective_functions}, - activate_force_length_relationship=True, - activate_force_velocity_relationship=False, - minimize_muscle_fatigue=True, + objective={"custom": objective_functions, "minimize_fatigue": True}, + msk_info={ + "with_residual_torque": False, + "bound_type": "start_end", + "bound_data": [[0, 5], [0, 90]], + }, + n_threads=5, ) sol = ocp.solve(Solver.IPOPT(_max_iter=1000)) diff --git a/cocofest/examples/dynamics/muscle_force_length_and_force_velocity_relationships_comparison.py b/cocofest/examples/dynamics/muscle_force_length_and_force_velocity_relationships_comparison.py index 5b8cb24..5624b0a 100644 --- a/cocofest/examples/dynamics/muscle_force_length_and_force_velocity_relationships_comparison.py +++ b/cocofest/examples/dynamics/muscle_force_length_and_force_velocity_relationships_comparison.py @@ -3,41 +3,48 @@ on the joint angle. """ -import matplotlib.pyplot as plt import numpy as np -from bioptim import Solver, SolutionMerge +import matplotlib.pyplot as plt + +from bioptim import SolutionMerge -from cocofest import ( - DingModelPulseDurationFrequencyWithFatigue, - OcpFesMsk, -) +from cocofest import DingModelPulseWidthFrequencyWithFatigue, OcpFesMsk, FesMskModel -minimum_pulse_duration = DingModelPulseDurationFrequencyWithFatigue().pd0 +minimum_pulse_width = DingModelPulseWidthFrequencyWithFatigue().pd0 sol_list = [] sol_time = [] activate_force_length_relationship = [False, True] for i in range(2): - ocp = OcpFesMsk.prepare_ocp( - biorbd_model_path="../msk_models/arm26_biceps_1dof.bioMod", - bound_type="start", - bound_data=[0], - fes_muscle_models=[DingModelPulseDurationFrequencyWithFatigue(muscle_name="BIClong")], - n_stim=10, - n_shooting=10, - final_time=1, - pulse_duration={"fixed": 0.00025}, - with_residual_torque=False, + + model = FesMskModel( + name=None, + biorbd_path="../msk_models/arm26_biceps_1dof.bioMod", + muscles_model=[DingModelPulseWidthFrequencyWithFatigue(muscle_name="BIClong")], activate_force_length_relationship=activate_force_length_relationship[i], activate_force_velocity_relationship=activate_force_length_relationship[i], + activate_residual_torque=False, + ) + + ocp = OcpFesMsk.prepare_ocp( + model=model, + stim_time=np.linspace(0, 1, 11)[:-1], + final_time=1, + pulse_width={"fixed": 0.00025}, + msk_info={ + "bound_type": "start", + "bound_data": [0], + "with_residual_torque": False, + }, use_sx=False, ) - sol = ocp.solve(Solver.IPOPT(_max_iter=1000)) + sol = ocp.solve() sol_list.append(sol.decision_states(to_merge=[SolutionMerge.NODES, SolutionMerge.PHASES])) time = np.concatenate( - sol.stepwise_time(to_merge=[SolutionMerge.NODES, SolutionMerge.PHASES], duplicated_times=False), axis=0 + sol.stepwise_time(to_merge=[SolutionMerge.NODES, SolutionMerge.PHASES], duplicated_times=False), + axis=0, ) index = 0 for j in range(len(sol.ocp.nlp) - 1): diff --git a/cocofest/examples/dynamics/reaching_task/reaching_task_frequency_optimization.py b/cocofest/examples/dynamics/reaching_task/reaching_task_frequency_optimization.py index c4c49e0..1dc4396 100644 --- a/cocofest/examples/dynamics/reaching_task/reaching_task_frequency_optimization.py +++ b/cocofest/examples/dynamics/reaching_task/reaching_task_frequency_optimization.py @@ -1,22 +1,18 @@ """ +/!\ This example is not functional yet. /!\ +/!\ It is a work in progress muscles can not be stimulated seperatly /!\ + This example will do a pulse apparition optimization to either minimize overall muscle force or muscle fatigue for a reaching task. Those ocp were build to move from starting position (arm: 0°, elbow: 5°) to a target position defined in the bioMod file. At the end of the simulation 2 files will be created, one for each optimization. The files will contain the time, states, controls and parameters of the ocp. """ -import pickle +import numpy as np -from bioptim import ( - Axis, - ConstraintFcn, - ConstraintList, - Node, - Solver, - SolutionMerge, -) +from bioptim import Axis, ConstraintFcn, ConstraintList, Node, Solver -from cocofest import DingModelFrequencyWithFatigue, OcpFesMsk +from cocofest import DingModelFrequencyWithFatigue, OcpFesMsk, FesMskModel, SolutionToPickle # Fiber type proportion from [1] biceps_fiber_type_2_proportion = 0.607 @@ -61,16 +57,24 @@ fes_muscle_models[i].alpha_a = fes_muscle_models[i].alpha_a * alpha_a_proportion_list[i] fes_muscle_models[i].a_rest = fes_muscle_models[i].a_rest * a_rest_proportion_list[i] +model = FesMskModel( + name=None, + biorbd_path="../../msk_models/arm26.bioMod", + muscles_model=fes_muscle_models, + activate_force_length_relationship=True, + activate_force_velocity_relationship=True, + activate_residual_torque=False, +) + pickle_file_list = ["minimize_muscle_fatigue.pkl", "minimize_muscle_force.pkl"] -n_stim = 40 -n_shooting = 5 +stim_time = list(np.round(np.linspace(0, 1, 41), 3))[:-1] constraint = ConstraintList() constraint.add( ConstraintFcn.SUPERIMPOSE_MARKERS, first_marker="COM_hand", second_marker="reaching_target", - phase=n_stim - 1, + phase=0, node=Node.END, axes=[Axis.X, Axis.Y], ) @@ -82,39 +86,26 @@ parameters = [] ocp = OcpFesMsk.prepare_ocp( - biorbd_model_path="../../msk_models/arm26.bioMod", - bound_type="start", - bound_data=[0, 5], - fes_muscle_models=fes_muscle_models, - n_stim=n_stim, - n_shooting=n_shooting, + model=model, + stim_time=stim_time, final_time=1, pulse_event={"min": 0.01, "max": 0.1, "bimapping": False}, - with_residual_torque=False, - custom_constraint=constraint, - activate_force_length_relationship=True, - activate_force_velocity_relationship=True, - minimize_muscle_fatigue=True if pickle_file_list[i] == "minimize_muscle_fatigue.pkl" else False, - minimize_muscle_force=True if pickle_file_list[i] == "minimize_muscle_force.pkl" else False, + objective={ + "minimize_fatigue": (True if pickle_file_list[i] == "minimize_muscle_fatigue.pkl" else False), + "minimize_force": (True if pickle_file_list[i] == "minimize_muscle_force.pkl" else False), + }, + msk_info={ + "with_residual_torque": False, + "bound_type": "start", + "bound_data": [0, 5], + "custom_constraint": constraint, + }, + n_threads=5, use_sx=False, ) sol = ocp.solve(Solver.IPOPT(_max_iter=10000)) - - time = sol.decision_time(to_merge=[SolutionMerge.PHASES, SolutionMerge.NODES]) - states = sol.decision_states(to_merge=[SolutionMerge.PHASES, SolutionMerge.NODES]) - controls = sol.decision_controls(to_merge=[SolutionMerge.PHASES, SolutionMerge.NODES]) - parameters = sol.decision_parameters() - - dictionary = { - "time": time, - "states": states, - "controls": controls, - "parameters": parameters, - } - - with open("/result_file/pulse_apparition_" + pickle_file_list[i], "wb") as file: - pickle.dump(dictionary, file) + SolutionToPickle(sol, "pulse_intensity_" + pickle_file_list[i], "result_file/").pickle() # [1] Dahmane, R., Djordjevič, S., Šimunič, B., & Valenčič, V. (2005). diff --git a/cocofest/examples/dynamics/reaching_task/reaching_task_intensity_optimization.py b/cocofest/examples/dynamics/reaching_task/reaching_task_intensity_optimization.py index 0e0c349..6e6dd47 100644 --- a/cocofest/examples/dynamics/reaching_task/reaching_task_intensity_optimization.py +++ b/cocofest/examples/dynamics/reaching_task/reaching_task_intensity_optimization.py @@ -5,18 +5,11 @@ The files will contain the time, states, controls and parameters of the ocp. """ -import pickle - -from bioptim import ( - Axis, - ConstraintFcn, - ConstraintList, - Node, - Solver, - SolutionMerge, -) +import numpy as np + +from bioptim import Axis, ConstraintFcn, ConstraintList, Node, Solver -from cocofest import DingModelIntensityFrequencyWithFatigue, OcpFesMsk +from cocofest import DingModelPulseIntensityFrequencyWithFatigue, OcpFesMsk, FesMskModel, SolutionToPickle # Fiber type proportion from [1] biceps_fiber_type_2_proportion = 0.607 @@ -49,31 +42,40 @@ ] fes_muscle_models = [ - DingModelIntensityFrequencyWithFatigue(muscle_name="BIClong"), - DingModelIntensityFrequencyWithFatigue(muscle_name="BICshort"), - DingModelIntensityFrequencyWithFatigue(muscle_name="TRIlong"), - DingModelIntensityFrequencyWithFatigue(muscle_name="TRIlat"), - DingModelIntensityFrequencyWithFatigue(muscle_name="TRImed"), - DingModelIntensityFrequencyWithFatigue(muscle_name="BRA"), + DingModelPulseIntensityFrequencyWithFatigue(muscle_name="BIClong"), + DingModelPulseIntensityFrequencyWithFatigue(muscle_name="BICshort"), + DingModelPulseIntensityFrequencyWithFatigue(muscle_name="TRIlong"), + DingModelPulseIntensityFrequencyWithFatigue(muscle_name="TRIlat"), + DingModelPulseIntensityFrequencyWithFatigue(muscle_name="TRImed"), + DingModelPulseIntensityFrequencyWithFatigue(muscle_name="BRA"), ] for i in range(len(fes_muscle_models)): fes_muscle_models[i].alpha_a = fes_muscle_models[i].alpha_a * alpha_a_proportion_list[i] fes_muscle_models[i].a_rest = fes_muscle_models[i].a_rest * a_rest_proportion_list[i] -minimum_pulse_intensity = DingModelIntensityFrequencyWithFatigue.min_pulse_intensity( - DingModelIntensityFrequencyWithFatigue() +minimum_pulse_intensity = DingModelPulseIntensityFrequencyWithFatigue.min_pulse_intensity( + DingModelPulseIntensityFrequencyWithFatigue() ) + +model = FesMskModel( + name=None, + biorbd_path="../../msk_models/arm26.bioMod", + muscles_model=fes_muscle_models, + activate_force_length_relationship=True, + activate_force_velocity_relationship=True, + activate_residual_torque=False, +) + pickle_file_list = ["minimize_muscle_fatigue.pkl", "minimize_muscle_force.pkl"] -n_stim = 40 -n_shooting = 5 +stim_time = list(np.round(np.linspace(0, 1, 41), 3))[:-1] constraint = ConstraintList() constraint.add( ConstraintFcn.SUPERIMPOSE_MARKERS, first_marker="COM_hand", second_marker="reaching_target", - phase=n_stim - 1, + phase=0, node=Node.END, axes=[Axis.X, Axis.Y], ) @@ -85,43 +87,30 @@ parameters = [] ocp = OcpFesMsk.prepare_ocp( - biorbd_model_path="../../msk_models/arm26.bioMod", - bound_type="start", - bound_data=[0, 5], - fes_muscle_models=fes_muscle_models, - n_stim=n_stim, - n_shooting=n_shooting, + model=model, + stim_time=stim_time, final_time=1, pulse_intensity={ "min": minimum_pulse_intensity, "max": 80, "bimapping": False, }, - with_residual_torque=False, - custom_constraint=constraint, - activate_force_length_relationship=True, - activate_force_velocity_relationship=True, - minimize_muscle_fatigue=True if pickle_file_list[i] == "minimize_muscle_fatigue.pkl" else False, - minimize_muscle_force=True if pickle_file_list[i] == "minimize_muscle_force.pkl" else False, + objective={ + "minimize_fatigue": (True if pickle_file_list[i] == "minimize_muscle_fatigue.pkl" else False), + "minimize_force": (True if pickle_file_list[i] == "minimize_muscle_force.pkl" else False), + }, + msk_info={ + "with_residual_torque": False, + "bound_type": "start", + "bound_data": [0, 5], + "custom_constraint": constraint, + }, use_sx=False, + n_threads=5, ) sol = ocp.solve(Solver.IPOPT(_max_iter=10000)) - - time = sol.decision_time(to_merge=[SolutionMerge.PHASES, SolutionMerge.NODES]) - states = sol.decision_states(to_merge=[SolutionMerge.PHASES, SolutionMerge.NODES]) - controls = sol.decision_controls(to_merge=[SolutionMerge.PHASES, SolutionMerge.NODES]) - parameters = sol.decision_parameters() - - dictionary = { - "time": time, - "states": states, - "controls": controls, - "parameters": parameters, - } - - with open("/result_file/pulse_intensity_" + pickle_file_list[i], "wb") as file: - pickle.dump(dictionary, file) + SolutionToPickle(sol, "pulse_intensity_" + pickle_file_list[i], "result_file/").pickle() # [1] Dahmane, R., Djordjevič, S., Šimunič, B., & Valenčič, V. (2005). diff --git a/cocofest/examples/dynamics/reaching_task/reaching_task_pulse_duration_optimization.py b/cocofest/examples/dynamics/reaching_task/reaching_task_pulse_duration_optimization.py index fa4cdf5..7678dcb 100644 --- a/cocofest/examples/dynamics/reaching_task/reaching_task_pulse_duration_optimization.py +++ b/cocofest/examples/dynamics/reaching_task/reaching_task_pulse_duration_optimization.py @@ -1,19 +1,25 @@ """ -This example will do a pulse duration optimization to either minimize overall muscle force or muscle fatigue +This example will do a pulse width optimization to either minimize overall muscle force or muscle fatigue for a reaching task. Those ocp were build to move from starting position (arm: 0°, elbow: 5°) to a target position defined in the bioMod file. At the end of the simulation 2 files will be created, one for each optimization. The files will contain the time, states, controls and parameters of the ocp. """ +import numpy as np + from bioptim import ( Axis, ConstraintFcn, ConstraintList, - Node, Solver, ) -from cocofest import DingModelPulseDurationFrequencyWithFatigue, OcpFesMsk, SolutionToPickle +from cocofest import ( + DingModelPulseWidthFrequencyWithFatigue, + OcpFesMsk, + SolutionToPickle, + FesMskModel, +) # Scaling alpha_a and a_scale parameters for each muscle proportionally to the muscle PCSA and fiber type 2 proportion # Fiber type proportion from [1] @@ -48,12 +54,12 @@ # Build the functional electrical stimulation models according # to number and name of muscle in the musculoskeletal model used fes_muscle_models = [ - DingModelPulseDurationFrequencyWithFatigue(muscle_name="BIClong"), - DingModelPulseDurationFrequencyWithFatigue(muscle_name="BICshort"), - DingModelPulseDurationFrequencyWithFatigue(muscle_name="TRIlong"), - DingModelPulseDurationFrequencyWithFatigue(muscle_name="TRIlat"), - DingModelPulseDurationFrequencyWithFatigue(muscle_name="TRImed"), - DingModelPulseDurationFrequencyWithFatigue(muscle_name="BRA"), + DingModelPulseWidthFrequencyWithFatigue(muscle_name="BIClong"), + DingModelPulseWidthFrequencyWithFatigue(muscle_name="BICshort"), + DingModelPulseWidthFrequencyWithFatigue(muscle_name="TRIlong"), + DingModelPulseWidthFrequencyWithFatigue(muscle_name="TRIlat"), + DingModelPulseWidthFrequencyWithFatigue(muscle_name="TRImed"), + DingModelPulseWidthFrequencyWithFatigue(muscle_name="BRA"), ] # Applying the scaling @@ -61,48 +67,55 @@ fes_muscle_models[i].alpha_a = fes_muscle_models[i].alpha_a * alpha_a_proportion_list[i] fes_muscle_models[i].a_scale = fes_muscle_models[i].a_scale * a_scale_proportion_list[i] -minimum_pulse_duration = DingModelPulseDurationFrequencyWithFatigue().pd0 -# pickle_file_list = ["minimize_muscle_fatigue.pkl", "minimize_muscle_force.pkl"] -pickle_file_list = ["minimize_muscle_fatigue.pkl"] -n_stim = 60 -n_shooting = 25 -# Step time of 1ms -> 1sec / (40Hz * 25) = 0.001s +model = FesMskModel( + name=None, + biorbd_path="../../msk_models/arm26.bioMod", + muscles_model=fes_muscle_models, + activate_force_length_relationship=True, + activate_force_velocity_relationship=True, + activate_residual_torque=False, +) +minimum_pulse_width = DingModelPulseWidthFrequencyWithFatigue().pd0 +pickle_file_list = ["minimize_muscle_fatigue.pkl", "minimize_muscle_force.pkl"] +stim_time = list(np.round(np.linspace(0, 1.5, 61), 3))[:-1] + +# Step time of 1ms -> 1sec / (40Hz * 25) = 0.001s constraint = ConstraintList() constraint.add( ConstraintFcn.SUPERIMPOSE_MARKERS, first_marker="COM_hand", second_marker="reaching_target", - phase=39, - node=Node.END, + phase=0, + node=650, axes=[Axis.X, Axis.Y], ) for i in range(len(pickle_file_list)): ocp = OcpFesMsk.prepare_ocp( - biorbd_model_path="../../msk_models/arm26.bioMod", - bound_type="start_end", - bound_data=[[0, 5], [0, 5]], - fes_muscle_models=fes_muscle_models, - n_stim=n_stim, - n_shooting=n_shooting, + model=model, + stim_time=stim_time, final_time=1.5, - pulse_duration={ - "min": minimum_pulse_duration, + pulse_width={ + "min": minimum_pulse_width, "max": 0.0006, "bimapping": False, }, - with_residual_torque=False, - custom_constraint=constraint, - activate_force_length_relationship=True, - activate_force_velocity_relationship=True, - minimize_muscle_fatigue=True if pickle_file_list[i] == "minimize_muscle_fatigue.pkl" else False, - minimize_muscle_force=True if pickle_file_list[i] == "minimize_muscle_force.pkl" else False, + objective={ + "minimize_fatigue": (True if pickle_file_list[i] == "minimize_muscle_fatigue.pkl" else False), + "minimize_force": (True if pickle_file_list[i] == "minimize_muscle_force.pkl" else False), + }, + msk_info={ + "with_residual_torque": False, + "bound_type": "start_end", + "bound_data": [[0, 5], [0, 5]], + "custom_constraint": constraint, + }, use_sx=False, ) sol = ocp.solve(Solver.IPOPT(_max_iter=10000)) - SolutionToPickle(sol, "pulse_duration_" + pickle_file_list[i], "result_file/").pickle() + SolutionToPickle(sol, "pulse_width_" + pickle_file_list[i], "result_file/").pickle() # [1] Dahmane, R., Djordjevič, S., Šimunič, B., & Valenčič, V. (2005). # Spatial fiber type distribution in normal human muscle: histochemical and tensiomyographical evaluation. diff --git a/cocofest/examples/dynamics/reaching_task/visualization/animation.py b/cocofest/examples/dynamics/reaching_task/visualization/animation.py index bcc8e84..dc610e6 100644 --- a/cocofest/examples/dynamics/reaching_task/visualization/animation.py +++ b/cocofest/examples/dynamics/reaching_task/visualization/animation.py @@ -1,8 +1,8 @@ +import biorbd from cocofest import PickleAnimate -PickleAnimate("../result_file/pulse_duration_minimize_muscle_fatigue.pkl").animate( - model_path="../../../msk_models/arm26.bioMod" -) -PickleAnimate("../result_file/pulse_duration_minimize_muscle_fatigue.pkl").multiple_animations( - ["../result_file/pulse_duration_minimize_muscle_force.pkl"], model_path="../../../msk_models/arm26.bioMod" +biorbd_model = biorbd.Model("../../../msk_models/arm26.bioMod") +PickleAnimate("../result_file/pulse_width_minimize_muscle_fatigue.pkl").animate(model=biorbd_model) +PickleAnimate("../result_file/pulse_width_minimize_muscle_fatigue.pkl").multiple_animations( + ["../result_file/pulse_width_minimize_muscle_force.pkl"], model=biorbd_model ) diff --git a/cocofest/examples/dynamics/reaching_task/visualization/force_and_fatigue_subplot.py b/cocofest/examples/dynamics/reaching_task/visualization/force_and_fatigue_subplot.py index ac3bf30..fccb0d7 100644 --- a/cocofest/examples/dynamics/reaching_task/visualization/force_and_fatigue_subplot.py +++ b/cocofest/examples/dynamics/reaching_task/visualization/force_and_fatigue_subplot.py @@ -8,8 +8,8 @@ import numpy as np pickle_path = [ - r"../result_file/pulse_duration_minimize_muscle_force.pkl", - r"../result_file/pulse_duration_minimize_muscle_fatigue.pkl", + r"../result_file/pulse_width_minimize_muscle_force.pkl", + r"../result_file/pulse_width_minimize_muscle_fatigue.pkl", ] with open(pickle_path[0], "rb") as f: @@ -18,8 +18,22 @@ with open(pickle_path[1], "rb") as f: data_minimize_fatigue = pickle.load(f) -force_muscle_keys = ["F_BIClong", "F_BICshort", "F_TRIlong", "F_TRIlat", "F_TRImed", "F_BRA"] -fatigue_muscle_keys = ["A_BIClong", "A_BICshort", "A_TRIlong", "A_TRIlat", "A_TRImed", "A_BRA"] +force_muscle_keys = [ + "F_BIClong", + "F_BICshort", + "F_TRIlong", + "F_TRIlat", + "F_TRImed", + "F_BRA", +] +fatigue_muscle_keys = [ + "A_BIClong", + "A_BICshort", + "A_TRIlong", + "A_TRIlat", + "A_TRImed", + "A_BRA", +] muscle_names = ["BIClong", "BICshort", "TRIlong", "TRIlat", "TRImed", "BRA"] # Force graph @@ -56,8 +70,16 @@ xticklabels=[0, 0.5, 1, 1.5], ) - axs[i][j].plot(data_minimize_force["time"], data_minimize_force["states"][force_muscle_keys[index]], lw=5) - axs[i][j].plot(data_minimize_fatigue["time"], data_minimize_fatigue["states"][force_muscle_keys[index]], lw=5) + axs[i][j].plot( + data_minimize_force["time"], + data_minimize_force["states"][force_muscle_keys[index]], + lw=5, + ) + axs[i][j].plot( + data_minimize_fatigue["time"], + data_minimize_fatigue["states"][force_muscle_keys[index]], + lw=5, + ) axs[i][j].text( 0.5, 0.9, @@ -76,7 +98,16 @@ index += 1 -fig.text(0.5, 0.02, "Time (s)", ha="center", va="center", fontsize=18, weight="bold", font="Times New Roman") +fig.text( + 0.5, + 0.02, + "Time (s)", + ha="center", + va="center", + fontsize=18, + weight="bold", + font="Times New Roman", +) fig.text( 0.025, 0.5, @@ -89,7 +120,10 @@ font="Times New Roman", ) fig.legend( - ["Force", "Fatigue"], loc="upper right", ncol=1, prop={"family": "Times New Roman", "size": 14, "weight": "bold"} + ["Force", "Fatigue"], + loc="upper right", + ncol=1, + prop={"family": "Times New Roman", "size": 14, "weight": "bold"}, ) plt.show() @@ -145,9 +179,21 @@ weight="bold", font="Times New Roman", ) -axs[1].text(0.75, -25, "Time (s)", ha="center", va="center", fontsize=18, weight="bold", font="Times New Roman") +axs[1].text( + 0.75, + -25, + "Time (s)", + ha="center", + va="center", + fontsize=18, + weight="bold", + font="Times New Roman", +) fig.legend( - ["Force", "Fatigue"], loc="upper right", ncol=1, prop={"family": "Times New Roman", "size": 14, "weight": "bold"} + ["Force", "Fatigue"], + loc="upper right", + ncol=1, + prop={"family": "Times New Roman", "size": 14, "weight": "bold"}, ) plt.show() @@ -192,8 +238,18 @@ a_force_sum_percentage = (np.array(a_force_sum_list) / a_sum_base_line) * 100 a_fatigue_sum_percentage = (np.array(a_fatigue_sum_list) / a_sum_base_line) * 100 -axs.plot(data_minimize_force["time"], a_force_sum_percentage, lw=5, label="Minimize force production") -axs.plot(data_minimize_force["time"], a_fatigue_sum_percentage, lw=5, label="Maximize muscle capacity") +axs.plot( + data_minimize_force["time"], + a_force_sum_percentage, + lw=5, + label="Minimize force production", +) +axs.plot( + data_minimize_force["time"], + a_fatigue_sum_percentage, + lw=5, + label="Maximize muscle capacity", +) axs.set_xlim(left=0, right=1.5) diff --git a/cocofest/examples/dynamics/reaching_task/visualization/make_gaph.py b/cocofest/examples/dynamics/reaching_task/visualization/make_gaph.py index 889cc6c..31f1a8a 100644 --- a/cocofest/examples/dynamics/reaching_task/visualization/make_gaph.py +++ b/cocofest/examples/dynamics/reaching_task/visualization/make_gaph.py @@ -10,8 +10,8 @@ chosen_graph_to_plot = "duration" duration_path = [ - r"../result_file/pulse_duration_minimize_muscle_force.pkl", - r"../result_file/pulse_duration_minimize_muscle_fatigue.pkl", + r"../result_file/pulse_width_minimize_muscle_force.pkl", + r"../result_file/pulse_width_minimize_muscle_fatigue.pkl", ] chosen_graph_to_plot_path = duration_path if chosen_graph_to_plot == "duration" else None @@ -26,7 +26,14 @@ with open(chosen_graph_to_plot_path[1], "rb") as f: data_minimize_fatigue = pickle.load(f) -force_muscle_keys = ["F_BIClong", "F_BICshort", "F_TRIlong", "F_TRIlat", "F_TRImed", "F_BRA"] +force_muscle_keys = [ + "F_BIClong", + "F_BICshort", + "F_TRIlong", + "F_TRIlat", + "F_TRImed", + "F_BRA", +] muscle_names = ["BIClong", "BICshort", "TRIlong", "TRIlat", "TRImed", "BRA"] muscle_title_x_postiton = [0.55, 0.5, 0.56, 0.62, 0.55, 0.73] fig, axs = plt.subplots(3, 3, figsize=(5, 3), constrained_layout=True) @@ -200,7 +207,11 @@ ) axs[2][2].plot( - data_minimize_force["time"], fatigue_minimization_percentage_gain_list, ms=4, linewidth=5.0, color="green" + data_minimize_force["time"], + fatigue_minimization_percentage_gain_list, + ms=4, + linewidth=5.0, + color="green", ) axs[2][2].text( diff --git a/cocofest/examples/dynamics/reaching_task/visualization/pulse_duration_subplot.py b/cocofest/examples/dynamics/reaching_task/visualization/pulse_duration_subplot.py index b02df00..9a8f2c4 100644 --- a/cocofest/examples/dynamics/reaching_task/visualization/pulse_duration_subplot.py +++ b/cocofest/examples/dynamics/reaching_task/visualization/pulse_duration_subplot.py @@ -10,8 +10,8 @@ import matplotlib.cm as cmx pickle_path = [ - r"../result_file/pulse_duration_minimize_muscle_force.pkl", - r"../result_file/pulse_duration_minimize_muscle_fatigue.pkl", + r"../result_file/pulse_width_minimize_muscle_force.pkl", + r"../result_file/pulse_width_minimize_muscle_fatigue.pkl", ] with open(pickle_path[0], "rb") as f: @@ -20,16 +20,16 @@ with open(pickle_path[1], "rb") as f: data_minimize_fatigue = pickle.load(f) -pulse_duration_keys = list(data_minimize_fatigue["parameters"].keys()) +pulse_width_keys = list(data_minimize_fatigue["parameters"].keys()) muscle_names = ["BIClong", "BICshort", "TRIlong", "TRIlat", "TRImed", "BRA"] -nb_stim = len(data_minimize_fatigue["parameters"][pulse_duration_keys[0]]) +nb_stim = len(data_minimize_fatigue["parameters"][pulse_width_keys[0]]) width = round(data_minimize_fatigue["time"][-1], 2) / nb_stim pw_data_list = [data_minimize_force["parameters"], data_minimize_fatigue["parameters"]] pw_list = [] for j in range(2): - pw_list.append([pw_data_list[j][pulse_duration_keys[i]] * 1000000 for i in range(len(pulse_duration_keys))]) + pw_list.append([pw_data_list[j][pulse_width_keys[i]] * 1000000 for i in range(len(pulse_width_keys))]) plasma = cm = plt.get_cmap("plasma") cNorm = colors.Normalize(vmin=100, vmax=600) @@ -39,14 +39,14 @@ def plot_graph(datas): fig, axs = plt.subplots(6, 1, figsize=(5, 3), constrained_layout=True) - for i in range(len(pulse_duration_keys)): + for i in range(len(pulse_width_keys)): axs[i].set_xlim(left=0, right=1.5) plt.setp( axs[i], xticks=[0, 0.2, 0.4, 0.6, 0.8, 1, 1.2, 1.4], xticklabels=[], ) - if i == len(pulse_duration_keys) - 1: + if i == len(pulse_width_keys) - 1: plt.setp( axs[i], xticks=[0, 0.2, 0.4, 0.6, 0.8, 1, 1.2, 1.4], @@ -60,7 +60,7 @@ def plot_graph(datas): color = scalarMap.to_rgba(value) axs[i].barh(muscle_names[i], width, left=j * width, height=0.5, color=color) - fig.colorbar(scalarMap, ax=axs, orientation="vertical", label="Pulse duration (us)") + fig.colorbar(scalarMap, ax=axs, orientation="vertical", label="pulse width (us)") plt.show() diff --git a/cocofest/examples/getting_started/force_tracking_parameter_optimization.py b/cocofest/examples/getting_started/force_tracking_parameter_optimization.py index 4614dbf..0a4bac4 100644 --- a/cocofest/examples/getting_started/force_tracking_parameter_optimization.py +++ b/cocofest/examples/getting_started/force_tracking_parameter_optimization.py @@ -8,13 +8,13 @@ from bioptim import SolutionMerge from cocofest import ( - DingModelIntensityFrequency, + ModelMaker, FourierSeries, OcpFes, ) # --- Building force to track ---# -time = np.linspace(0, 1, 100) +time = np.linspace(0, 1, 1001) force = abs(np.sin(time * 5) + np.random.normal(scale=0.1, size=len(time))) * 100 force_tracking = [time, force] @@ -22,13 +22,12 @@ # This ocp was build to track a force curve along the problem. # The stimulation won't be optimized and is already set to one pulse every 0.1 seconds (n_stim/final_time). # Plus the pulsation intensity will be optimized between 0 and 130 mA and are not the same across the problem. -model = DingModelIntensityFrequency() +model = ModelMaker.create_model("hmed2018", is_approximated=True) minimum_pulse_intensity = model.min_pulse_intensity() ocp = OcpFes().prepare_ocp( model=model, - n_stim=10, - n_shooting=20, + stim_time=list(np.round(np.linspace(0, 1, 31)[:-1], 3)), final_time=1, pulse_intensity={ "min": minimum_pulse_intensity, @@ -37,6 +36,7 @@ }, objective={"force_tracking": force_tracking}, use_sx=True, + n_threads=8, ) # --- Solve the program --- # @@ -56,9 +56,14 @@ plt.plot(time, y_approx, color="orange", label="force after fourier transform") solution_time = sol.decision_time(to_merge=SolutionMerge.KEYS, continuous=True) -solution_time = [float(j) for sub in solution_time for j in sub] +solution_time = [float(j) for j in solution_time] -plt.plot(solution_time, sol_merged["F"].squeeze(), color="blue", label="force from optimized stimulation") +plt.plot( + solution_time, + sol_merged["F"].squeeze(), + color="blue", + label="force from optimized stimulation", +) plt.xlabel("Time (s)") plt.ylabel("Force (N)") plt.legend() diff --git a/cocofest/examples/getting_started/frequency_optimization.py b/cocofest/examples/getting_started/frequency_optimization.py index 092336a..eea832c 100644 --- a/cocofest/examples/getting_started/frequency_optimization.py +++ b/cocofest/examples/getting_started/frequency_optimization.py @@ -3,24 +3,24 @@ This ocp was build to match a force value of 270N at the end of the last node. """ -from cocofest import DingModelFrequencyWithFatigue, OcpFes +from cocofest import OcpFes, ModelMaker +from bioptim import ControlType # --- Build ocp --- # # This ocp was build to match a force value of 270N at the end of the last node. # The stimulation will be optimized between 0.01 to 0.1 seconds and are equally spaced (a fixed frequency). - +model = ModelMaker.create_model("ding2003", is_approximated=True) ocp = OcpFes().prepare_ocp( - model=DingModelFrequencyWithFatigue(), - n_stim=10, - n_shooting=20, + model=model, + stim_time=[0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9], final_time=1, - pulse_event={"min": 0.01, "max": 0.1, "bimapping": True}, objective={"end_node_tracking": 270}, use_sx=True, + control_type=ControlType.LINEAR_CONTINUOUS, ) # --- Solve the program --- # sol = ocp.solve() # --- Show results --- # -sol.graphs() +sol.graphs(show_bounds=True) diff --git a/cocofest/examples/getting_started/frequency_optimization_musculoskeletal_dynamic.py b/cocofest/examples/getting_started/frequency_optimization_musculoskeletal_dynamic.py index 4693210..37a109f 100644 --- a/cocofest/examples/getting_started/frequency_optimization_musculoskeletal_dynamic.py +++ b/cocofest/examples/getting_started/frequency_optimization_musculoskeletal_dynamic.py @@ -5,35 +5,35 @@ elbow torque control. """ -from bioptim import ( - ObjectiveFcn, - ObjectiveList, - Solver, -) +import numpy as np -from cocofest import DingModelFrequencyWithFatigue, OcpFesMsk +from cocofest import DingModelFrequencyWithFatigue, OcpFesMsk, FesMskModel -objective_functions = ObjectiveList() -n_stim = 10 -for i in range(n_stim): - objective_functions.add(ObjectiveFcn.Lagrange.MINIMIZE_CONTROL, key="tau", weight=1, quadratic=True, phase=i) +model = FesMskModel( + name=None, + biorbd_path="../msk_models/arm26_biceps_1dof.bioMod", + muscles_model=[DingModelFrequencyWithFatigue(muscle_name="BIClong")], + activate_force_length_relationship=True, + activate_force_velocity_relationship=True, + activate_residual_torque=True, +) ocp = OcpFesMsk.prepare_ocp( - biorbd_model_path="../msk_models/arm26_biceps_1dof.bioMod", - bound_type="start_end", - bound_data=[[5], [120]], - fes_muscle_models=[DingModelFrequencyWithFatigue(muscle_name="BIClong")], - n_stim=n_stim, - n_shooting=10, + model=model, + stim_time=list(np.round(np.linspace(0, 1, 11)[:-1], 2)), final_time=1, pulse_event={"min": 0.01, "max": 0.1, "bimapping": True}, - objective={"custom": objective_functions}, - with_residual_torque=True, - activate_force_length_relationship=True, - activate_force_velocity_relationship=True, + objective={"minimize_residual_torque": True}, + msk_info={ + "with_residual_torque": True, + "bound_type": "start_end", + "bound_data": [[5], [120]], + }, + use_sx=True, + n_threads=5, ) -sol = ocp.solve(Solver.IPOPT(_max_iter=1000)) +sol = ocp.solve() sol.animate() sol.graphs(show_bounds=False) diff --git a/cocofest/examples/getting_started/model_integration.py b/cocofest/examples/getting_started/model_integration.py index 42f4c58..0c8044d 100644 --- a/cocofest/examples/getting_started/model_integration.py +++ b/cocofest/examples/getting_started/model_integration.py @@ -1,15 +1,14 @@ import matplotlib.pyplot as plt -from cocofest import ( - DingModelFrequencyWithFatigue, - IvpFes, -) - +from cocofest import IvpFes, ModelMaker # --- Build ocp --- # # This problem was build to be integrated and has no objectives nor parameter to optimize. - -fes_parameters = {"model": DingModelFrequencyWithFatigue(), "n_stim": 10} -ivp_parameters = {"n_shooting": 20, "final_time": 1} +model = ModelMaker.create_model("ding2003_with_fatigue", is_approximated=False) # Can not approximate this model in ivp +fes_parameters = { + "model": model, + "stim_time": [0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9], +} +ivp_parameters = {"final_time": 1, "use_sx": True} ivp = IvpFes(fes_parameters, ivp_parameters) diff --git a/cocofest/examples/getting_started/muscle_model_identification.py b/cocofest/examples/getting_started/muscle_model_identification.py index 8f04e96..7c1a5dd 100644 --- a/cocofest/examples/getting_started/muscle_model_identification.py +++ b/cocofest/examples/getting_started/muscle_model_identification.py @@ -1,7 +1,13 @@ import pickle import os + +import numpy as np + +from bioptim import SolutionMerge + from cocofest import ( - DingModelIntensityFrequency, + ModelMaker, + DingModelPulseIntensityFrequency, DingModelPulseIntensityFrequencyForceParameterIdentification, IvpFes, ) @@ -9,6 +15,145 @@ import matplotlib.pyplot as plt +# Example n°5 : Identification of the parameters of the Ding model with the pulse intensity method for simulated data +# --- Simulating data --- # +# This problem was build to be integrated and has no objectives nor parameter to optimize. +n_stim = 50 +final_time = 5 +model = ModelMaker.create_model("hmed2018", is_approximated=False) # Can not approximate this model in ivp + +stim_time = list(np.round(np.linspace(0, final_time, n_stim + 1), 2))[:-1] +pulse_intensity_values = [20, 20, 30, 40, 50, 60, 70, 80, 90, 100] * int((n_stim / 10)) + +fes_parameters = { + "model": model, + "pulse_intensity": pulse_intensity_values, + "stim_time": stim_time, +} +ivp_parameters = {"final_time": final_time, "use_sx": True} +ivp = IvpFes(fes_parameters, ivp_parameters) + +# Integrating the solution +result, time = ivp.integrate() +force = result["F"][0] + +pulse_intensity = pulse_intensity_values + +dictionary = {"time": time, "force": force, "stim_time": stim_time, "pulse_intensity": pulse_intensity} + +pickle_file_name = "../data/temp_identification_simulation.pkl" +with open(pickle_file_name, "wb") as file: + pickle.dump(dictionary, file) + +model = ModelMaker.create_model("hmed2018", is_approximated=False) +ocp = DingModelPulseIntensityFrequencyForceParameterIdentification( + model=model, + data_path=[pickle_file_name], + identification_method="full", + double_step_identification=False, + key_parameter_to_identify=[ + "a_rest", + "km_rest", + "tau1_rest", + "tau2", + "ar", + "bs", + "Is", + "cr", + ], + additional_key_settings={}, + final_time=final_time, + use_sx=True, + n_threads=6, +) + +identified_parameters = ocp.force_model_identification() +force_ocp = ocp.force_identification_result.decision_states(to_merge=SolutionMerge.NODES)["F"][0] +a_rest = ( + identified_parameters["a_rest"] if "a_rest" in identified_parameters else DingModelPulseIntensityFrequency().a_rest +) +km_rest = ( + identified_parameters["km_rest"] + if "km_rest" in identified_parameters + else DingModelPulseIntensityFrequency().km_rest +) +tau1_rest = ( + identified_parameters["tau1_rest"] + if "tau1_rest" in identified_parameters + else DingModelPulseIntensityFrequency().tau1_rest +) +tau2 = identified_parameters["tau2"] if "tau2" in identified_parameters else DingModelPulseIntensityFrequency().tau2 +ar = identified_parameters["ar"] if "ar" in identified_parameters else DingModelPulseIntensityFrequency().ar +bs = identified_parameters["bs"] if "bs" in identified_parameters else DingModelPulseIntensityFrequency().bs +Is = identified_parameters["Is"] if "Is" in identified_parameters else DingModelPulseIntensityFrequency().Is +cr = identified_parameters["cr"] if "cr" in identified_parameters else DingModelPulseIntensityFrequency().cr +print( + "a_rest : ", + a_rest, + "km_rest : ", + km_rest, + "tau1_rest : ", + tau1_rest, + "tau2 : ", + tau2, + "ar : ", + ar, + "bs : ", + bs, + "Is : ", + Is, + "cr : ", + cr, +) + +( + pickle_time_data, + pickle_stim_apparition_time, + pickle_muscle_data, + pickle_discontinuity_phase_list, +) = full_data_extraction([pickle_file_name]) + +# Plotting the identification result +plt.title("Force state result") +plt.plot(pickle_time_data, pickle_muscle_data, color="blue", label="simulated") +plt.plot(pickle_time_data, force_ocp, color="green", label="identified") +plt.xlabel("time (s)") +plt.ylabel("force (N)") + +plt.annotate("a_rest : ", xy=(0.7, 0.4), xycoords="axes fraction", color="black") +plt.annotate("km_rest : ", xy=(0.7, 0.35), xycoords="axes fraction", color="black") +plt.annotate("tau1_rest : ", xy=(0.7, 0.3), xycoords="axes fraction", color="black") +plt.annotate("tau2 : ", xy=(0.7, 0.25), xycoords="axes fraction", color="black") +plt.annotate("ar : ", xy=(0.7, 0.2), xycoords="axes fraction", color="black") +plt.annotate("bs : ", xy=(0.7, 0.15), xycoords="axes fraction", color="black") +plt.annotate("Is : ", xy=(0.7, 0.1), xycoords="axes fraction", color="black") +plt.annotate("cr : ", xy=(0.7, 0.05), xycoords="axes fraction", color="black") + +plt.annotate(str(round(a_rest, 5)), xy=(0.78, 0.4), xycoords="axes fraction", color="red") +plt.annotate(str(round(km_rest, 5)), xy=(0.78, 0.35), xycoords="axes fraction", color="red") +plt.annotate(str(round(tau1_rest, 5)), xy=(0.78, 0.3), xycoords="axes fraction", color="red") +plt.annotate(str(round(tau2, 5)), xy=(0.78, 0.25), xycoords="axes fraction", color="red") +plt.annotate(str(round(ar, 5)), xy=(0.78, 0.2), xycoords="axes fraction", color="red") +plt.annotate(str(round(bs, 5)), xy=(0.78, 0.15), xycoords="axes fraction", color="red") +plt.annotate(str(round(Is, 5)), xy=(0.78, 0.1), xycoords="axes fraction", color="red") +plt.annotate(str(round(cr, 5)), xy=(0.78, 0.05), xycoords="axes fraction", color="red") + +plt.annotate(str(DingModelPulseIntensityFrequency().a_rest), xy=(0.85, 0.4), xycoords="axes fraction", color="blue") +plt.annotate(str(DingModelPulseIntensityFrequency().km_rest), xy=(0.85, 0.35), xycoords="axes fraction", color="blue") +plt.annotate(str(DingModelPulseIntensityFrequency().tau1_rest), xy=(0.85, 0.3), xycoords="axes fraction", color="blue") +plt.annotate(str(DingModelPulseIntensityFrequency().tau2), xy=(0.85, 0.25), xycoords="axes fraction", color="blue") +plt.annotate(str(DingModelPulseIntensityFrequency().ar), xy=(0.85, 0.2), xycoords="axes fraction", color="blue") +plt.annotate(str(DingModelPulseIntensityFrequency().bs), xy=(0.85, 0.15), xycoords="axes fraction", color="blue") +plt.annotate(str(DingModelPulseIntensityFrequency().Is), xy=(0.85, 0.1), xycoords="axes fraction", color="blue") +plt.annotate(str(DingModelPulseIntensityFrequency().cr), xy=(0.85, 0.05), xycoords="axes fraction", color="blue") + +# --- Delete the temp file ---# +os.remove(f"../data/temp_identification_simulation.pkl") + +plt.legend() +plt.show() + + # Example n°1 : Identification of the parameters of the Ding model with the frequency method for experimental data """ ocp = DingModelFrequencyParameterIdentification( @@ -179,17 +324,17 @@ """ # -# # Example n°4 : Identification of the parameters of the Ding model with the pulse duration method for simulated data +# # Example n°4 : Identification of the parameters of the Ding model with the pulse width method for simulated data # # --- Simulating data --- # # # This problem was build to be integrated and has no objectives nor parameter to optimize. -# pulse_duration_values = [0.000180, 0.0002, 0.000250, 0.0003, 0.000350, 0.0004, 0.000450, 0.0005, 0.000550, 0.0006] +# pulse_width_values = [0.000180, 0.0002, 0.000250, 0.0003, 0.000350, 0.0004, 0.000450, 0.0005, 0.000550, 0.0006] # ivp = IvpFes( -# model=DingModelPulseDurationFrequency(), +# model=DingModelPulseWidthFrequency(), # n_stim=10, # n_shooting=10, # final_time=1, # use_sx=True, -# pulse_duration=pulse_duration_values, +# pulse_width=pulse_width_values, # ) # # # Creating the solution from the initial guess @@ -204,21 +349,21 @@ # time = [result.time.tolist()] # stim_temp = [0 if i == 0 else result.ocp.nlp[i].tf for i in range(len(result.ocp.nlp))] # stim = [sum(stim_temp[: i + 1]) for i in range(len(stim_temp))] -# pulse_duration = pulse_duration_values +# pulse_width = pulse_width_values # # dictionary = { # "time": time, # "biceps": force, # "stim_time": stim, -# "pulse_duration": pulse_duration, +# "pulse_width": pulse_width, # } # # pickle_file_name = "../data/temp_identification_simulation.pkl" # with open(pickle_file_name, "wb") as file: # pickle.dump(dictionary, file) # -# ocp = DingModelPulseDurationFrequencyForceParameterIdentification( -# model=DingModelPulseDurationFrequency(), +# ocp = DingModelPulseWidthFrequencyForceParameterIdentification( +# model=DingModelPulseWidthFrequency(), # data_path=[pickle_file_name], # identification_method="full", # double_step_identification=False, @@ -229,15 +374,15 @@ # ) # # identified_parameters = ocp.force_model_identification() -# a_scale = identified_parameters["a_scale"] if "a_scale" in identified_parameters else DingModelPulseDurationFrequency().a_scale -# pd0 = identified_parameters["pd0"] if "pd0" in identified_parameters else DingModelPulseDurationFrequency().pd0 -# pdt = identified_parameters["pdt"] if "pdt" in identified_parameters else DingModelPulseDurationFrequency().pdt -# km_rest = identified_parameters["km_rest"] if "km_rest" in identified_parameters else DingModelPulseDurationFrequency().km_rest -# tau1_rest = identified_parameters["tau1_rest"] if "tau1_rest" in identified_parameters else DingModelPulseDurationFrequency().tau1_rest -# tau2 = identified_parameters["tau2"] if "tau2" in identified_parameters else DingModelPulseDurationFrequency().tau2 +# a_scale = identified_parameters["a_scale"] if "a_scale" in identified_parameters else DingModelPulseWidthFrequency().a_scale +# pd0 = identified_parameters["pd0"] if "pd0" in identified_parameters else DingModelPulseWidthFrequency().pd0 +# pdt = identified_parameters["pdt"] if "pdt" in identified_parameters else DingModelPulseWidthFrequency().pdt +# km_rest = identified_parameters["km_rest"] if "km_rest" in identified_parameters else DingModelPulseWidthFrequency().km_rest +# tau1_rest = identified_parameters["tau1_rest"] if "tau1_rest" in identified_parameters else DingModelPulseWidthFrequency().tau1_rest +# tau2 = identified_parameters["tau2"] if "tau2" in identified_parameters else DingModelPulseWidthFrequency().tau2 # print("a_scale : ", a_scale, "pd0 : ", pd0, "pdt : ", pdt, "km_rest : ", km_rest, "tau1_rest : ", tau1_rest, "tau2 : ", tau2) # -# identified_model = DingModelPulseDurationFrequency() +# identified_model = DingModelPulseWidthFrequency() # identified_model.a_scale = a_scale # identified_model.km_rest = km_rest # identified_model.tau1_rest = tau1_rest @@ -254,7 +399,7 @@ # n_shooting=10, # final_time=1, # use_sx=True, -# pulse_duration=[0.000184, 0.0002, 0.000250, 0.0003, 0.000350, 0.0004, 0.000450, 0.0005, 0.000550, 0.0006], +# pulse_width=[0.000184, 0.0002, 0.000250, 0.0003, 0.000350, 0.0004, 0.000450, 0.0005, 0.000550, 0.0006], # ) # # # Creating the solution from the initial guess @@ -281,7 +426,7 @@ # pickle_stim_apparition_time, # pickle_muscle_data, # pickle_discontinuity_phase_list, -# ) = DingModelPulseDurationFrequencyForceParameterIdentification.full_data_extraction([pickle_file_name]) +# ) = DingModelPulseWidthFrequencyForceParameterIdentification.full_data_extraction([pickle_file_name]) # # # Plotting the identification result # plt.title("Force state result") @@ -317,151 +462,3 @@ # # plt.legend() # plt.show() - - -# Example n°5 : Identification of the parameters of the Ding model with the pulse intensity method for simulated data -# --- Simulating data --- # -# This problem was build to be integrated and has no objectives nor parameter to optimize. -pulse_intensity_values = [20, 20, 30, 40, 50, 60, 70, 80, 90, 100] -fes_parameters = {"model": DingModelIntensityFrequency(), "n_stim": 10, "pulse_intensity": pulse_intensity_values} -ivp_parameters = {"n_shooting": 10, "final_time": 1, "use_sx": True} -ivp = IvpFes( - fes_parameters, - ivp_parameters, -) - -# Integrating the solution -result, time = ivp.integrate() - -force = result["F"][0].tolist() - -stim = [1 / 10 * i for i in range(10)] -pulse_intensity = pulse_intensity_values - -dictionary = { - "time": time, - "force": force, - "stim_time": stim, - "pulse_intensity": pulse_intensity, -} - -pickle_file_name = "../data/temp_identification_simulation.pkl" -with open(pickle_file_name, "wb") as file: - pickle.dump(dictionary, file) - -ocp = DingModelPulseIntensityFrequencyForceParameterIdentification( - model=DingModelIntensityFrequency(), - data_path=[pickle_file_name], - identification_method="full", - double_step_identification=False, - key_parameter_to_identify=["a_rest", "km_rest", "tau1_rest", "tau2", "ar", "bs", "Is", "cr"], - additional_key_settings={}, - n_shooting=10, - use_sx=True, -) - -identified_parameters = ocp.force_model_identification() -a_rest = identified_parameters["a_rest"] if "a_rest" in identified_parameters else DingModelIntensityFrequency().a_rest -km_rest = ( - identified_parameters["km_rest"] if "km_rest" in identified_parameters else DingModelIntensityFrequency().km_rest -) -tau1_rest = ( - identified_parameters["tau1_rest"] - if "tau1_rest" in identified_parameters - else DingModelIntensityFrequency().tau1_rest -) -tau2 = identified_parameters["tau2"] if "tau2" in identified_parameters else DingModelIntensityFrequency().tau2 -ar = identified_parameters["ar"] if "ar" in identified_parameters else DingModelIntensityFrequency().ar -bs = identified_parameters["bs"] if "bs" in identified_parameters else DingModelIntensityFrequency().bs -Is = identified_parameters["Is"] if "Is" in identified_parameters else DingModelIntensityFrequency().Is -cr = identified_parameters["cr"] if "cr" in identified_parameters else DingModelIntensityFrequency().cr -print( - "a_rest : ", - a_rest, - "km_rest : ", - km_rest, - "tau1_rest : ", - tau1_rest, - "tau2 : ", - tau2, - "ar : ", - ar, - "bs : ", - bs, - "Is : ", - Is, - "cr : ", - cr, -) - -identified_model = DingModelIntensityFrequency() -identified_model.a_rest = a_rest -identified_model.km_rest = km_rest -identified_model.tau1_rest = tau1_rest -identified_model.tau2 = tau2 -identified_model.ar = ar -identified_model.bs = bs -identified_model.Is = Is -identified_model.cr = cr - -identified_force_list = [] -identified_time_list = [] - -fes_parameters = {"model": identified_model, "n_stim": 10, "pulse_intensity": pulse_intensity_values} -ivp_parameters = {"n_shooting": 10, "final_time": 1, "use_sx": True} -ivp_from_identification = IvpFes( - fes_parameters, - ivp_parameters, -) - -# Integrating the solution -identified_result, identified_time = ivp_from_identification.integrate() - -identified_force = identified_result["F"][0].tolist() - -( - pickle_time_data, - pickle_stim_apparition_time, - pickle_muscle_data, - pickle_discontinuity_phase_list, -) = full_data_extraction([pickle_file_name]) - -# Plotting the identification result -plt.title("Force state result") -plt.plot(pickle_time_data, pickle_muscle_data, color="blue", label="simulated") -plt.plot(identified_time, identified_force, color="red", label="identified") -plt.xlabel("time (s)") -plt.ylabel("force (N)") - -plt.annotate("a_rest : ", xy=(0.7, 0.4), xycoords="axes fraction", color="black") -plt.annotate("km_rest : ", xy=(0.7, 0.35), xycoords="axes fraction", color="black") -plt.annotate("tau1_rest : ", xy=(0.7, 0.3), xycoords="axes fraction", color="black") -plt.annotate("tau2 : ", xy=(0.7, 0.25), xycoords="axes fraction", color="black") -plt.annotate("ar : ", xy=(0.7, 0.2), xycoords="axes fraction", color="black") -plt.annotate("bs : ", xy=(0.7, 0.15), xycoords="axes fraction", color="black") -plt.annotate("Is : ", xy=(0.7, 0.1), xycoords="axes fraction", color="black") -plt.annotate("cr : ", xy=(0.7, 0.05), xycoords="axes fraction", color="black") - -plt.annotate(str(round(a_rest, 5)), xy=(0.78, 0.4), xycoords="axes fraction", color="red") -plt.annotate(str(round(km_rest, 5)), xy=(0.78, 0.35), xycoords="axes fraction", color="red") -plt.annotate(str(round(tau1_rest, 5)), xy=(0.78, 0.3), xycoords="axes fraction", color="red") -plt.annotate(str(round(tau2, 5)), xy=(0.78, 0.25), xycoords="axes fraction", color="red") -plt.annotate(str(round(ar, 5)), xy=(0.78, 0.2), xycoords="axes fraction", color="red") -plt.annotate(str(round(bs, 5)), xy=(0.78, 0.15), xycoords="axes fraction", color="red") -plt.annotate(str(round(Is, 5)), xy=(0.78, 0.1), xycoords="axes fraction", color="red") -plt.annotate(str(round(cr, 5)), xy=(0.78, 0.05), xycoords="axes fraction", color="red") - -plt.annotate(str(DingModelIntensityFrequency().a_rest), xy=(0.85, 0.4), xycoords="axes fraction", color="blue") -plt.annotate(str(DingModelIntensityFrequency().km_rest), xy=(0.85, 0.35), xycoords="axes fraction", color="blue") -plt.annotate(str(DingModelIntensityFrequency().tau1_rest), xy=(0.85, 0.3), xycoords="axes fraction", color="blue") -plt.annotate(str(DingModelIntensityFrequency().tau2), xy=(0.85, 0.25), xycoords="axes fraction", color="blue") -plt.annotate(str(DingModelIntensityFrequency().ar), xy=(0.85, 0.2), xycoords="axes fraction", color="blue") -plt.annotate(str(DingModelIntensityFrequency().bs), xy=(0.85, 0.15), xycoords="axes fraction", color="blue") -plt.annotate(str(DingModelIntensityFrequency().Is), xy=(0.85, 0.1), xycoords="axes fraction", color="blue") -plt.annotate(str(DingModelIntensityFrequency().cr), xy=(0.85, 0.05), xycoords="axes fraction", color="blue") - -# --- Delete the temp file ---# -os.remove(f"../data/temp_identification_simulation.pkl") - -plt.legend() -plt.show() diff --git a/cocofest/examples/getting_started/pulse_duration_optimization.py b/cocofest/examples/getting_started/pulse_duration_optimization.py index b090284..293c8bc 100644 --- a/cocofest/examples/getting_started/pulse_duration_optimization.py +++ b/cocofest/examples/getting_started/pulse_duration_optimization.py @@ -1,33 +1,35 @@ """ -This example will do a 10 stimulation example with Ding's 2007 pulse duration and frequency model. +This example will do a 10 stimulation example with Ding's 2007 pulse width and frequency model. This ocp was build to match a force value of 200N at the end of the last node. """ -from cocofest import DingModelPulseDurationFrequencyWithFatigue, OcpFes +from bioptim import Solver +from cocofest import OcpFes, ModelMaker # --- Build ocp --- # # This ocp was build to match a force value of 200N at the end of the last node. # The stimulation will be optimized between 0.01 to 0.1 seconds and are equally spaced (a fixed frequency). -# Plus the pulsation duration will be optimized between 0 and 0.0006 seconds and are not the same across the problem. +# Plus the pulsation width will be optimized between 0 and 0.0006 seconds and are not the same across the problem. # The flag with_fatigue is set to True by default, this will include the fatigue model -minimum_pulse_duration = DingModelPulseDurationFrequencyWithFatigue().pd0 + +model = ModelMaker.create_model("ding2007_with_fatigue", is_approximated=False) + +minimum_pulse_width = model.pd0 ocp = OcpFes().prepare_ocp( - model=DingModelPulseDurationFrequencyWithFatigue(), - n_stim=10, - n_shooting=20, - final_time=1, - pulse_event={"min": 0.01, "max": 0.1, "bimapping": True}, - pulse_duration={ - "min": minimum_pulse_duration, + model=model, + stim_time=[0, 0.05, 0.1, 0.15, 0.2, 0.25, 0.3, 0.35, 0.4, 0.45], + final_time=0.5, + pulse_width={ + "min": minimum_pulse_width, "max": 0.0006, - "bimapping": False, + "bimapping": True, }, - objective={"end_node_tracking": 200}, + objective={"end_node_tracking": 100}, use_sx=True, + n_threads=5, ) # --- Solve the program --- # -sol = ocp.solve() - +sol = ocp.solve(Solver.IPOPT(_hessian_approximation="limited-memory")) # --- Show results --- # sol.graphs() diff --git a/cocofest/examples/getting_started/pulse_duration_optimization_musculoskeletal_dynamic.py b/cocofest/examples/getting_started/pulse_duration_optimization_musculoskeletal_dynamic.py index 75b4473..6fb9be9 100644 --- a/cocofest/examples/getting_started/pulse_duration_optimization_musculoskeletal_dynamic.py +++ b/cocofest/examples/getting_started/pulse_duration_optimization_musculoskeletal_dynamic.py @@ -1,45 +1,41 @@ """ This example will do a 10 stimulation example with Ding's 2007 frequency model. This ocp was build to produce a elbow motion from 5 to 120 degrees. -The stimulation frequency will be optimized between 10 and 100 Hz and pulse duration between minimal sensitivity +The stimulation frequency will be optimized between 10 and 100 Hz and pulse width between minimal sensitivity threshold and 600us to satisfy the flexion and minimizing required elbow torque control. """ -from bioptim import ( - ObjectiveFcn, - ObjectiveList, - Solver, -) - -from cocofest import DingModelPulseDurationFrequencyWithFatigue, OcpFesMsk - +from bioptim import Solver +from cocofest import DingModelPulseWidthFrequencyWithFatigue, OcpFesMsk, FesMskModel -objective_functions = ObjectiveList() -n_stim = 10 -for i in range(n_stim): - objective_functions.add(ObjectiveFcn.Lagrange.MINIMIZE_CONTROL, key="tau", weight=1, quadratic=True, phase=i) +model = FesMskModel( + name=None, + biorbd_path="../msk_models/arm26_biceps_1dof.bioMod", + muscles_model=[DingModelPulseWidthFrequencyWithFatigue(muscle_name="BIClong")], + activate_force_length_relationship=True, + activate_force_velocity_relationship=True, + activate_residual_torque=True, +) -minimum_pulse_duration = DingModelPulseDurationFrequencyWithFatigue().pd0 +minimum_pulse_width = DingModelPulseWidthFrequencyWithFatigue().pd0 ocp = OcpFesMsk.prepare_ocp( - biorbd_model_path="../msk_models/arm26_biceps_1dof.bioMod", - bound_type="start_end", - bound_data=[[5], [120]], - fes_muscle_models=[DingModelPulseDurationFrequencyWithFatigue(muscle_name="BIClong")], - n_stim=n_stim, - n_shooting=10, + model=model, + stim_time=[0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9], final_time=1, - pulse_event={"min": 0.01, "max": 0.1, "bimapping": True}, - pulse_duration={ - "min": minimum_pulse_duration, + pulse_width={ + "min": minimum_pulse_width, "max": 0.0006, "bimapping": False, }, - objective={"custom": objective_functions}, - with_residual_torque=True, - activate_force_length_relationship=True, - activate_force_velocity_relationship=True, + objective={"minimize_residual_torque": True}, + msk_info={ + "bound_type": "start_end", + "bound_data": [[5], [120]], + "with_residual_torque": True, + }, ) -sol = ocp.solve(Solver.IPOPT(_max_iter=2000)) -# sol.animate() -sol.graphs(show_bounds=False) +if __name__ == "__main__": + sol = ocp.solve(Solver.IPOPT(show_online_optim=True)) + sol.animate() + sol.graphs(show_bounds=False) diff --git a/cocofest/examples/getting_started/pulse_duration_optimization_nmpc_cyclic.py b/cocofest/examples/getting_started/pulse_duration_optimization_nmpc_cyclic.py index 1b122af..c492e30 100644 --- a/cocofest/examples/getting_started/pulse_duration_optimization_nmpc_cyclic.py +++ b/cocofest/examples/getting_started/pulse_duration_optimization_nmpc_cyclic.py @@ -1,53 +1,64 @@ """ This example showcases a moving time horizon simulation problem of cyclic muscle force tracking. -The FES model used here is Ding's 2007 pulse duration and frequency model with fatigue. -Only the pulse duration is optimized, frequency is fixed. -The nmpc cyclic problem is composed of 3 cycles and will move forward 1 cycle at each step. -Only the middle cycle is kept in the optimization problem, the nmpc cyclic problem stops once the last 6th cycle is reached. +The FES model used here is Ding's 2007 pulse width and frequency model with fatigue. +Only the pulse width is optimized, frequency is fixed. +The nmpc cyclic problem stops once the last cycle is reached. """ -import numpy as np import matplotlib.pyplot as plt +import numpy as np +from bioptim import Solver, SolutionMerge +from cocofest import NmpcFes, DingModelPulseWidthFrequencyWithFatigue -from bioptim import OdeSolver -from cocofest import OcpFesNmpcCyclic, DingModelPulseDurationFrequencyWithFatigue -# --- Build target force --- # +# --- Building force to track ---# target_time = np.linspace(0, 1, 100) target_force = abs(np.sin(target_time * np.pi)) * 200 force_tracking = [target_time, target_force] # --- Build nmpc cyclic --- # -n_total_cycles = 8 -minimum_pulse_duration = DingModelPulseDurationFrequencyWithFatigue().pd0 -fes_model = DingModelPulseDurationFrequencyWithFatigue(sum_stim_truncation=10) +cycles_len = 1000 +cycle_duration = 1 +n_cycles = 8 + +minimum_pulse_width = DingModelPulseWidthFrequencyWithFatigue().pd0 +fes_model = DingModelPulseWidthFrequencyWithFatigue(sum_stim_truncation=10) fes_model.alpha_a = -4.0 * 10e-1 # Increasing the fatigue rate to make the fatigue more visible -nmpc = OcpFesNmpcCyclic( + +nmpc = NmpcFes.prepare_nmpc( model=fes_model, - n_stim=30, - n_shooting=5, - final_time=1, - pulse_duration={ - "min": minimum_pulse_duration, + stim_time=list(np.round(np.linspace(0, 1, 31), 3))[:-1], + cycle_len=cycles_len, + cycle_duration=cycle_duration, + pulse_width={ + "min": minimum_pulse_width, "max": 0.0006, "bimapping": False, }, objective={"force_tracking": force_tracking}, - n_total_cycles=n_total_cycles, - n_simultaneous_cycles=3, - n_cycle_to_advance=1, - cycle_to_keep="middle", use_sx=True, - ode_solver=OdeSolver.COLLOCATION(), + n_threads=5, ) -nmpc.prepare_nmpc() -nmpc.solve() +n_cycles_total = 8 + + +def update_functions(_nmpc, cycle_idx, _sol): + return cycle_idx < n_cycles_total # True if there are still some cycle to perform + + +sol = nmpc.solve( + update_functions, + solver=Solver.IPOPT(), + cyclic_options={"states": {}}, + get_all_iterations=True, +) +sol_merged = sol[0].decision_states(to_merge=[SolutionMerge.PHASES, SolutionMerge.NODES]) -# --- Show results --- # -time = [j for sub in nmpc.result["time"] for j in sub] -fatigue = [j for sub in nmpc.result["states"]["A"] for j in sub] -force = [j for sub in nmpc.result["states"]["F"] for j in sub] +time = sol[0].decision_time(to_merge=SolutionMerge.KEYS, continuous=True) +time = [float(j) for j in time] +fatigue = sol_merged["A"][0] +force = sol_merged["F"][0] ax1 = plt.subplot(221) ax1.plot(time, fatigue, label="A", color="green") @@ -58,7 +69,7 @@ ax2 = plt.subplot(222) ax2.plot(time, force, label="F", color="red", linewidth=4) -for i in range(n_total_cycles): +for i in range(n_cycles): if i == 0: ax2.plot(target_time, target_force, label="Target", color="purple") else: @@ -69,9 +80,9 @@ plt.legend() barWidth = 0.25 # set width of bar -cycles = nmpc.result["parameters"]["pulse_duration"] # set height of bar +cycles = [sol[1][i].parameters["pulse_width"] for i in range(len(sol[1]))] # set height of bar bar = [] # Set position of bar on X axis -for i in range(n_total_cycles): +for i in range(n_cycles): if i == 0: br = [barWidth * (x + 1) for x in range(len(cycles[i]))] else: @@ -79,11 +90,11 @@ bar.append(br) ax3 = plt.subplot(212) -for i in range(n_total_cycles): +for i in range(n_cycles): ax3.bar(bar[i], cycles[i], width=barWidth, edgecolor="grey", label=f"cycle n°{i+1}") -ax3.set_xticks([np.mean(r) for r in bar], [str(i + 1) for i in range(n_total_cycles)]) +ax3.set_xticks([np.mean(r) for r in bar], [str(i + 1) for i in range(n_cycles)]) ax3.set_xlabel("Cycles") -ax3.set_ylabel("Pulse duration (s)") +ax3.set_ylabel("Pulse width (s)") plt.legend() -ax3.set_title("Pulse duration", weight="bold") +ax3.set_title("Pulse width", weight="bold") plt.show() diff --git a/cocofest/examples/getting_started/pulse_intensity_optimization.py b/cocofest/examples/getting_started/pulse_intensity_optimization.py index 27d78a2..b8cc12e 100644 --- a/cocofest/examples/getting_started/pulse_intensity_optimization.py +++ b/cocofest/examples/getting_started/pulse_intensity_optimization.py @@ -3,17 +3,18 @@ This ocp was build to match a force value of 200N at the end of the last node. """ -from cocofest import DingModelIntensityFrequency, OcpFes +from cocofest import DingModelPulseIntensityFrequencyWithFatigue, OcpFes # --- Build ocp --- # # This ocp was build to match a force value of 200N at the end of the last node. # The stimulation won't be optimized and is already set to one pulse every 0.1 seconds (n_stim/final_time). # Plus the pulsation intensity will be optimized between 0 and 130 mA and are not the same across the problem. -minimum_pulse_intensity = DingModelIntensityFrequency.min_pulse_intensity(DingModelIntensityFrequency()) +minimum_pulse_intensity = DingModelPulseIntensityFrequencyWithFatigue.min_pulse_intensity( + DingModelPulseIntensityFrequencyWithFatigue() +) ocp = OcpFes().prepare_ocp( - model=DingModelIntensityFrequency(), - n_stim=10, - n_shooting=20, + model=DingModelPulseIntensityFrequencyWithFatigue(is_approximated=False), + stim_time=[0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9], final_time=1, pulse_intensity={ "min": minimum_pulse_intensity, @@ -22,6 +23,7 @@ }, objective={"end_node_tracking": 130}, use_sx=True, + n_threads=5, ) # --- Solve the program --- # diff --git a/cocofest/examples/getting_started/pulse_intensity_optimization_musculoskeletal_dynamic.py b/cocofest/examples/getting_started/pulse_intensity_optimization_musculoskeletal_dynamic.py index 369f66e..88e8b83 100644 --- a/cocofest/examples/getting_started/pulse_intensity_optimization_musculoskeletal_dynamic.py +++ b/cocofest/examples/getting_started/pulse_intensity_optimization_musculoskeletal_dynamic.py @@ -5,40 +5,34 @@ threshold and 130mA to satisfy the flexion and minimizing required elbow torque control. """ -from bioptim import ( - ObjectiveFcn, - ObjectiveList, - Solver, -) - -from cocofest import DingModelIntensityFrequencyWithFatigue, OcpFesMsk +from cocofest import DingModelPulseIntensityFrequencyWithFatigue, OcpFesMsk, FesMskModel -objective_functions = ObjectiveList() -n_stim = 10 -for i in range(n_stim): - objective_functions.add(ObjectiveFcn.Lagrange.MINIMIZE_CONTROL, key="tau", weight=1, quadratic=True, phase=i) +minimum_pulse_intensity = DingModelPulseIntensityFrequencyWithFatigue.min_pulse_intensity( + DingModelPulseIntensityFrequencyWithFatigue() +) -minimum_pulse_intensity = DingModelIntensityFrequencyWithFatigue.min_pulse_intensity( - DingModelIntensityFrequencyWithFatigue() +model = FesMskModel( + biorbd_path="../msk_models/arm26_biceps_1dof.bioMod", + muscles_model=[DingModelPulseIntensityFrequencyWithFatigue(muscle_name="BIClong")], + activate_force_length_relationship=True, + activate_force_velocity_relationship=True, + activate_residual_torque=True, ) ocp = OcpFesMsk.prepare_ocp( - biorbd_model_path="../msk_models/arm26_biceps_1dof.bioMod", - bound_type="start_end", - bound_data=[[5], [120]], - fes_muscle_models=[DingModelIntensityFrequencyWithFatigue(muscle_name="BIClong")], - n_stim=n_stim, - n_shooting=10, + model=model, + stim_time=[0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9], final_time=1, - pulse_event={"min": 0.01, "max": 0.1, "bimapping": True}, pulse_intensity={"min": minimum_pulse_intensity, "max": 130, "bimapping": False}, - objective={"custom": objective_functions}, - with_residual_torque=True, - activate_force_length_relationship=True, - activate_force_velocity_relationship=True, + objective={"minimize_residual_torque": True}, + msk_info={ + "with_residual_torque": True, + "bound_type": "start_end", + "bound_data": [[5], [120]], + }, ) -sol = ocp.solve(Solver.IPOPT(_max_iter=1000)) +sol = ocp.solve() sol.animate() sol.graphs(show_bounds=False) diff --git a/cocofest/examples/getting_started/pulse_mode_example.py b/cocofest/examples/getting_started/pulse_mode_example.py index ecbbe3d..f6fdf29 100644 --- a/cocofest/examples/getting_started/pulse_mode_example.py +++ b/cocofest/examples/getting_started/pulse_mode_example.py @@ -3,20 +3,21 @@ The example model is the Ding2003 frequency model. """ -import numpy as np import matplotlib.pyplot as plt -from cocofest import DingModelFrequencyWithFatigue, IvpFes +from cocofest import IvpFes, ModelMaker # --- Example n°1 : Single --- # # --- Build ocp --- # # This example shows how to create a problem with single pulses. # The stimulation won't be optimized. -ns = 10 -n_stim = 10 final_time = 1 - -fes_parameters = {"model": DingModelFrequencyWithFatigue(), "n_stim": n_stim, "pulse_mode": "single"} -ivp_parameters = {"n_shooting": ns, "final_time": final_time, "use_sx": True} +model = ModelMaker.create_model("ding_2003_with_fatigue", is_approximated=False) +fes_parameters = { + "model": model, + "pulse_mode": "single", + "stim_time": [0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9], +} +ivp_parameters = {"final_time": final_time, "use_sx": True} ivp = IvpFes( fes_parameters, @@ -25,18 +26,13 @@ result_single, time_single = ivp.integrate() force_single = result_single["F"][0] -stimulation_single = np.concatenate((np.array([0]), np.cumsum(np.array(ivp.final_time_phase)))) - +stimulation_single = ivp.stim_time # --- Example n°2 : Doublets --- # # --- Build ocp --- # # This example shows how to create a problem with doublet pulses. # The stimulation won't be optimized. -ns = 10 -n_stim = 20 -final_time = 1 -fes_parameters = {"model": DingModelFrequencyWithFatigue(), "n_stim": n_stim, "pulse_mode": "doublet"} -ivp_parameters = {"n_shooting": ns, "final_time": final_time, "use_sx": True} +fes_parameters["pulse_mode"] = "doublet" ivp = IvpFes( fes_parameters, ivp_parameters, @@ -44,15 +40,12 @@ result_doublet, time_doublet = ivp.integrate() force_doublet = result_doublet["F"][0] -stimulation_doublet = np.concatenate((np.array([0]), np.cumsum(np.array(ivp.final_time_phase)))) - +stimulation_doublet = ivp.stim_time # --- Example n°3 : Triplets --- # # --- Build ocp --- # # This example shows how to create a problem with triplet pulses. -n_stim = 30 -fes_parameters = {"model": DingModelFrequencyWithFatigue(), "n_stim": n_stim, "pulse_mode": "triplet"} -ivp_parameters = {"n_shooting": 10, "final_time": 1, "use_sx": True} +fes_parameters["pulse_mode"] = "triplet" ivp = IvpFes( fes_parameters, ivp_parameters, @@ -60,7 +53,7 @@ result_triplet, time_triplet = ivp.integrate() force_triplet = result_triplet["F"][0] -stimulation_triplet = np.concatenate((np.array([0]), np.cumsum(np.array(ivp.final_time_phase)))) +stimulation_triplet = ivp.stim_time # --- Show results --- # plt.title("Force state result for single, doublet and triplet") @@ -70,7 +63,7 @@ plt.plot(time_triplet, force_triplet, color="green", label="force triplet") plt.vlines( - x=stimulation_single[:-1], + x=stimulation_single, ymin=max(force_single) - 30, ymax=max(force_single), colors="blue", @@ -79,7 +72,7 @@ label="stimulation single", ) plt.vlines( - x=stimulation_doublet[:-1], + x=stimulation_doublet, ymin=max(force_doublet) - 30, ymax=max(force_doublet), colors="red", @@ -88,7 +81,7 @@ label="stimulation doublet", ) plt.vlines( - x=stimulation_triplet[:-1], + x=stimulation_triplet, ymin=max(force_triplet) - 30, ymax=max(force_triplet), colors="green", diff --git a/cocofest/examples/identification/ding2003_model_identification.py b/cocofest/examples/identification/ding2003_model_identification.py index 5416932..f7ebf8e 100644 --- a/cocofest/examples/identification/ding2003_model_identification.py +++ b/cocofest/examples/identification/ding2003_model_identification.py @@ -10,8 +10,11 @@ import matplotlib.pyplot as plt import numpy as np +from bioptim import SolutionMerge, ControlType + from cocofest import ( DingModelFrequency, + DingModelFrequencyIntegrate, DingModelFrequencyForceParameterIdentification, IvpFes, ) @@ -19,26 +22,15 @@ # --- Setting simulation parameters --- # -n_stim = 10 -n_shooting = 10 -final_time = 1 -extra_phase_time = 1 -model = DingModelFrequency() -fes_parameters = {"model": model, "n_stim": n_stim} -ivp_parameters = { - "n_shooting": n_shooting, - "final_time": final_time, - "extend_last_phase_time": extra_phase_time, - "use_sx": True, -} - +final_time = 2 +stim_time = np.round(np.linspace(0, 1, 11)[:-1], 2) +ivp_model = DingModelFrequencyIntegrate() +fes_parameters = {"model": ivp_model, "stim_time": stim_time} +ivp_parameters = {"final_time": final_time, "use_sx": True} # --- Creating the simulated data to identify on --- # # Building the Initial Value Problem -ivp = IvpFes( - fes_parameters, - ivp_parameters, -) +ivp = IvpFes(fes_parameters, ivp_parameters) # Integrating the solution result, time = ivp.integrate() @@ -47,63 +39,32 @@ noise = np.random.normal(0, 5, len(result["F"][0])) force = result["F"][0] + noise -stim = [final_time / n_stim * i for i in range(n_stim)] - # Saving the data in a pickle file -dictionary = { - "time": time, - "force": force, - "stim_time": stim, -} +dictionary = {"time": time, "force": force, "stim_time": stim_time} pickle_file_name = "../data/temp_identification_simulation.pkl" with open(pickle_file_name, "wb") as file: pickle.dump(dictionary, file) - # --- Identifying the model parameters --- # +ocp_model = DingModelFrequency() ocp = DingModelFrequencyForceParameterIdentification( - model=model, + model=ocp_model, + final_time=final_time, data_path=[pickle_file_name], identification_method="full", double_step_identification=False, key_parameter_to_identify=["a_rest", "km_rest", "tau1_rest", "tau2"], additional_key_settings={}, - n_shooting=n_shooting, use_sx=True, + n_threads=6, + control_type=ControlType.LINEAR_CONTINUOUS, ) identified_parameters = ocp.force_model_identification() +force_ocp = ocp.force_identification_result.decision_states(to_merge=SolutionMerge.NODES)["F"][0] print(identified_parameters) -# --- Plotting noisy simulated data and simulation from model with the identified parameter --- # -identified_model = model -identified_model.a_rest = identified_parameters["a_rest"] -identified_model.km_rest = identified_parameters["km_rest"] -identified_model.tau1_rest = identified_parameters["tau1_rest"] -identified_model.tau2 = identified_parameters["tau2"] - -identified_force_list = [] -identified_time_list = [] - -# Building the Initial Value Problem -fes_parameters = {"model": identified_model, "n_stim": n_stim} -ivp_parameters = { - "n_shooting": n_shooting, - "final_time": final_time, - "extend_last_phase_time": extra_phase_time, - "use_sx": True, -} -ivp_from_identification = IvpFes( - fes_parameters, - ivp_parameters, -) - -# Integrating the solution -identified_result, identified_time = ivp_from_identification.integrate() - -identified_force = identified_result["F"][0] - ( pickle_time_data, pickle_stim_apparition_time, @@ -112,16 +73,17 @@ ) = full_data_extraction([pickle_file_name]) result_dict = { - "a_rest": [identified_model.a_rest, DingModelFrequency().a_rest], - "km_rest": [identified_model.km_rest, DingModelFrequency().km_rest], - "tau1_rest": [identified_model.tau1_rest, DingModelFrequency().tau1_rest], - "tau2": [identified_model.tau2, DingModelFrequency().tau2], + "a_rest": [identified_parameters["a_rest"], DingModelFrequency().a_rest], + "km_rest": [identified_parameters["km_rest"], DingModelFrequency().km_rest], + "tau1_rest": [identified_parameters["tau1_rest"], DingModelFrequency().tau1_rest], + "tau2": [identified_parameters["tau2"], DingModelFrequency().tau2], } # Plotting the identification result plt.title("Force state result") plt.plot(pickle_time_data, pickle_muscle_data, color="blue", label="simulated") -plt.plot(identified_time, identified_force, color="red", label="identified") +plt.plot(pickle_time_data, force_ocp, color="red", label="identified") + plt.xlabel("time (s)") plt.ylabel("force (N)") diff --git a/cocofest/examples/identification/ding2007_model_identification.py b/cocofest/examples/identification/ding2007_model_identification.py index d94f237..b06e94a 100644 --- a/cocofest/examples/identification/ding2007_model_identification.py +++ b/cocofest/examples/identification/ding2007_model_identification.py @@ -9,104 +9,62 @@ import matplotlib.pyplot as plt import numpy as np +from bioptim import SolutionMerge + from cocofest import ( - DingModelPulseDurationFrequency, - DingModelPulseDurationFrequencyForceParameterIdentification, + DingModelPulseWidthFrequency, + DingModelPulseWidthFrequencyForceParameterIdentification, IvpFes, + ModelMaker, ) from cocofest.identification.identification_method import full_data_extraction # --- Setting simulation parameters --- # -n_stim = 10 -pulse_duration = [0.003] * n_stim -# pulse_duration = np.random.uniform(0.002, 0.006, 10).tolist() -n_shooting = 10 -final_time = 1 -extra_phase_time = 1 -model = DingModelPulseDurationFrequency() -fes_parameters = {"model": model, "n_stim": n_stim, "pulse_duration": pulse_duration} -ivp_parameters = { - "n_shooting": n_shooting, - "final_time": final_time, - "extend_last_phase_time": extra_phase_time, - "use_sx": True, -} +stim_time = np.round(np.linspace(0, 1, 11)[:-1], 2) +pulse_width = np.random.uniform(0.0002, 0.0006, 10).tolist() + +final_time = 2 +ivp_model = ModelMaker.create_model("ding2007", is_approximated=False) +fes_parameters = {"model": ivp_model, "stim_time": stim_time, "pulse_width": pulse_width} +ivp_parameters = {"final_time": final_time, "use_sx": True} # --- Creating the simulated data to identify on --- # # Building the Initial Value Problem -ivp = IvpFes( - fes_parameters, - ivp_parameters, -) +ivp = IvpFes(fes_parameters, ivp_parameters) # Integrating the solution result, time = ivp.integrate() # Adding noise to the force noise = np.random.normal(0, 5, len(result["F"][0])) -force_n = result["F"][0] force = result["F"][0] + noise -stim = [final_time / n_stim * i for i in range(n_stim)] - # Saving the data in a pickle file -dictionary = { - "time": time, - "force": force, - "stim_time": stim, - "pulse_duration": pulse_duration, -} +dictionary = {"time": time, "force": force, "stim_time": stim_time, "pulse_width": pulse_width} pickle_file_name = "../data/temp_identification_simulation.pkl" with open(pickle_file_name, "wb") as file: pickle.dump(dictionary, file) - # --- Identifying the model parameters --- # -ocp = DingModelPulseDurationFrequencyForceParameterIdentification( - model=model, +ocp_model = DingModelPulseWidthFrequency() +ocp = DingModelPulseWidthFrequencyForceParameterIdentification( + model=ocp_model, data_path=[pickle_file_name], identification_method="full", double_step_identification=False, key_parameter_to_identify=["tau1_rest", "tau2", "km_rest", "a_scale", "pd0", "pdt"], additional_key_settings={}, - n_shooting=n_shooting, + final_time=final_time, use_sx=True, + n_threads=6, ) identified_parameters = ocp.force_model_identification() +force_ocp = ocp.force_identification_result.decision_states(to_merge=SolutionMerge.NODES)["F"][0] print(identified_parameters) -# --- Plotting noisy simulated data and simulation from model with the identified parameter --- # -identified_model = model -identified_model.tau1_rest = identified_parameters["tau1_rest"] -identified_model.tau2 = identified_parameters["tau2"] -identified_model.km_rest = identified_parameters["km_rest"] -identified_model.a_scale = identified_parameters["a_scale"] -identified_model.pd0 = identified_parameters["pd0"] -identified_model.pdt = identified_parameters["pdt"] - -identified_force_list = [] -identified_time_list = [] - -fes_parameters = {"model": identified_model, "n_stim": n_stim, "pulse_duration": pulse_duration} -ivp_parameters = { - "n_shooting": n_shooting, - "final_time": final_time, - "extend_last_phase_time": extra_phase_time, - "use_sx": True, -} -ivp_from_identification = IvpFes( - fes_parameters, - ivp_parameters, -) - -# Integrating the solution -identified_result, identified_time = ivp_from_identification.integrate() - -identified_force = identified_result["F"][0] - ( pickle_time_data, pickle_stim_apparition_time, @@ -115,19 +73,19 @@ ) = full_data_extraction([pickle_file_name]) result_dict = { - "tau1_rest": [identified_model.tau1_rest, DingModelPulseDurationFrequency().tau1_rest], - "tau2": [identified_model.tau2, DingModelPulseDurationFrequency().tau2], - "km_rest": [identified_model.km_rest, DingModelPulseDurationFrequency().km_rest], - "a_scale": [identified_model.a_scale, DingModelPulseDurationFrequency().a_scale], - "pd0": [identified_model.pd0, DingModelPulseDurationFrequency().pd0], - "pdt": [identified_model.pdt, DingModelPulseDurationFrequency().pdt], + "tau1_rest": [identified_parameters["tau1_rest"], DingModelPulseWidthFrequency().tau1_rest], + "tau2": [identified_parameters["tau2"], DingModelPulseWidthFrequency().tau2], + "km_rest": [identified_parameters["km_rest"], DingModelPulseWidthFrequency().km_rest], + "a_scale": [identified_parameters["a_scale"], DingModelPulseWidthFrequency().a_scale], + "pd0": [identified_parameters["pd0"], DingModelPulseWidthFrequency().pd0], + "pdt": [identified_parameters["pdt"], DingModelPulseWidthFrequency().pdt], } # Plotting the identification result plt.title("Force state result") -plt.plot(pickle_time_data, force_n, color="black", label="no noise") -plt.plot(pickle_time_data, pickle_muscle_data, "-.", color="blue", label="simulated (with noise)") -plt.plot(identified_time, identified_force, color="red", label="identified") +plt.plot(pickle_time_data, pickle_muscle_data, "-.", color="blue", label="simulated") +plt.plot(pickle_time_data, force_ocp, color="red", label="identified") + plt.xlabel("time (s)") plt.ylabel("force (N)") diff --git a/cocofest/examples/identification/hmed2018_model_identification.py b/cocofest/examples/identification/hmed2018_model_identification.py index c631697..01727f5 100644 --- a/cocofest/examples/identification/hmed2018_model_identification.py +++ b/cocofest/examples/identification/hmed2018_model_identification.py @@ -9,8 +9,11 @@ import matplotlib.pyplot as plt import numpy as np +from bioptim import SolutionMerge + from cocofest import ( - DingModelIntensityFrequency, + ModelMaker, + DingModelPulseIntensityFrequency, DingModelPulseIntensityFrequencyForceParameterIdentification, IvpFes, ) @@ -18,97 +21,58 @@ # --- Setting simulation parameters --- # -n_stim = 10 -pulse_intensity = [50] * n_stim -# pulse_intensity = np.random.uniform(20, 130, 10).tolist() -n_shooting = 10 -final_time = 1 -extra_phase_time = 1 -model = DingModelIntensityFrequency() -fes_parameters = {"model": model, "n_stim": n_stim, "pulse_intensity": pulse_intensity} -ivp_parameters = { - "n_shooting": n_shooting, - "final_time": final_time, - "use_sx": True, - "extend_last_phase_time": extra_phase_time, -} +stim_time = np.round(np.linspace(0, 1, 11)[:-1], 2) +pulse_intensity = np.random.randint(20, 130, 10).tolist() +final_time = 2 +ivp_model = ModelMaker.create_model("hmed2018", is_approximated=False) +fes_parameters = {"model": ivp_model, "stim_time": stim_time, "pulse_intensity": pulse_intensity} +ivp_parameters = {"final_time": final_time, "use_sx": True} # --- Creating the simulated data to identify on --- # # Building the Initial Value Problem -ivp = IvpFes( - fes_parameters, - ivp_parameters, -) +ivp = IvpFes(fes_parameters, ivp_parameters) # Integrating the solution result, time = ivp.integrate() # Adding noise to the force -noise = np.random.normal(0, 0.5, len(result["F"][0])) +noise = np.random.normal(0, 5, len(result["F"][0])) force = result["F"][0] + noise -stim = [final_time / n_stim * i for i in range(n_stim)] - # Saving the data in a pickle file -dictionary = { - "time": time, - "force": force, - "stim_time": stim, - "pulse_intensity": pulse_intensity, -} +dictionary = {"time": time, "force": force, "stim_time": stim_time, "pulse_intensity": pulse_intensity} pickle_file_name = "../data/temp_identification_simulation.pkl" with open(pickle_file_name, "wb") as file: pickle.dump(dictionary, file) # --- Identifying the model parameters --- # +ocp_model = DingModelPulseIntensityFrequency() ocp = DingModelPulseIntensityFrequencyForceParameterIdentification( - model=model, + model=ocp_model, data_path=[pickle_file_name], identification_method="full", double_step_identification=False, - key_parameter_to_identify=["a_rest", "km_rest", "tau1_rest", "tau2", "ar", "bs", "Is", "cr"], + key_parameter_to_identify=[ + "a_rest", + "km_rest", + "tau1_rest", + "tau2", + "ar", + "bs", + "Is", + "cr", + ], additional_key_settings={}, - n_shooting=n_shooting, + final_time=final_time, use_sx=True, + n_threads=6, ) identified_parameters = ocp.force_model_identification() +force_ocp = ocp.force_identification_result.decision_states(to_merge=SolutionMerge.NODES)["F"][0] print(identified_parameters) -# --- Plotting noisy simulated data and simulation from model with the identified parameter --- # -identified_model = model -identified_model.a_rest = identified_parameters["a_rest"] -identified_model.km_rest = identified_parameters["km_rest"] -identified_model.tau1_rest = identified_parameters["tau1_rest"] -identified_model.tau2 = identified_parameters["tau2"] -identified_model.ar = identified_parameters["ar"] -identified_model.bs = identified_parameters["bs"] -identified_model.Is = identified_parameters["Is"] -identified_model.cr = identified_parameters["cr"] - -identified_force_list = [] -identified_time_list = [] - -# Building the Initial Value Problem -fes_parameters = {"model": identified_model, "n_stim": n_stim, "pulse_intensity": pulse_intensity} -ivp_parameters = { - "n_shooting": n_shooting, - "final_time": final_time, - "use_sx": True, - "extend_last_phase_time": extra_phase_time, -} - -ivp_from_identification = IvpFes( - fes_parameters, - ivp_parameters, -) - -# Integrating the solution -identified_result, identified_time = ivp_from_identification.integrate() - -identified_force = identified_result["F"][0] - ( pickle_time_data, pickle_stim_apparition_time, @@ -117,20 +81,20 @@ ) = full_data_extraction([pickle_file_name]) result_dict = { - "a_rest": [identified_model.a_rest, DingModelIntensityFrequency().a_rest], - "km_rest": [identified_model.km_rest, DingModelIntensityFrequency().km_rest], - "tau1_rest": [identified_model.tau1_rest, DingModelIntensityFrequency().tau1_rest], - "tau2": [identified_model.tau2, DingModelIntensityFrequency().tau2], - "ar": [identified_model.ar, DingModelIntensityFrequency().ar], - "bs": [identified_model.bs, DingModelIntensityFrequency().bs], - "Is": [identified_model.Is, DingModelIntensityFrequency().Is], - "cr": [identified_model.cr, DingModelIntensityFrequency().cr], + "a_rest": [identified_parameters["a_rest"], DingModelPulseIntensityFrequency().a_rest], + "km_rest": [identified_parameters["km_rest"], DingModelPulseIntensityFrequency().km_rest], + "tau1_rest": [identified_parameters["tau1_rest"], DingModelPulseIntensityFrequency().tau1_rest], + "tau2": [identified_parameters["tau2"], DingModelPulseIntensityFrequency().tau2], + "ar": [identified_parameters["ar"], DingModelPulseIntensityFrequency().ar], + "bs": [identified_parameters["bs"], DingModelPulseIntensityFrequency().bs], + "Is": [identified_parameters["Is"], DingModelPulseIntensityFrequency().Is], + "cr": [identified_parameters["cr"], DingModelPulseIntensityFrequency().cr], } # Plotting the identification result plt.title("Force state result") plt.plot(pickle_time_data, pickle_muscle_data, color="blue", label="simulated") -plt.plot(identified_time, identified_force, color="red", label="identified") +plt.plot(pickle_time_data, force_ocp, color="red", label="identified") plt.xlabel("time (s)") plt.ylabel("force (N)") diff --git a/cocofest/examples/nmpc/cycling_fes_driven_nmpc.py b/cocofest/examples/nmpc/cycling_fes_driven_nmpc.py new file mode 100644 index 0000000..bdea165 --- /dev/null +++ b/cocofest/examples/nmpc/cycling_fes_driven_nmpc.py @@ -0,0 +1,94 @@ +""" +This example will do a nmpc of 10 stimulation example with Ding's 2007 frequency model. +This ocp was build to produce a elbow motion from 5 to 120 degrees. +The pulse width between minimal sensitivity threshold and 600us to satisfy the flexion and minimizing required elbow +torque control. +""" + +import numpy as np +import biorbd +from bioptim import Solver, ControlType +from cocofest import ( + DingModelPulseWidthFrequencyWithFatigue, + NmpcFesMsk, + FesMskModel, + PickleAnimate, + SolutionToPickle, +) + + +minimum_pulse_width = DingModelPulseWidthFrequencyWithFatigue().pd0 +DeltoideusClavicle_A_model = DingModelPulseWidthFrequencyWithFatigue(muscle_name="DeltoideusClavicle_A") +DeltoideusScapula_P_model = DingModelPulseWidthFrequencyWithFatigue(muscle_name="DeltoideusScapula_P") +TRIlong_model = DingModelPulseWidthFrequencyWithFatigue(muscle_name="TRIlong") +BIC_long_model = DingModelPulseWidthFrequencyWithFatigue(muscle_name="BIC_long") +BIC_brevis_model = DingModelPulseWidthFrequencyWithFatigue(muscle_name="BIC_brevis") + +DeltoideusClavicle_A_model.alpha_a = -4.0 * 10e-1 +DeltoideusScapula_P_model.alpha_a = -4.0 * 10e-1 +TRIlong_model.alpha_a = -4.0 * 10e-1 +BIC_long_model.alpha_a = -4.0 * 10e-1 +BIC_brevis_model.alpha_a = -4.0 * 10e-1 + +model = FesMskModel( + name=None, + biorbd_path="../msk_models/simplified_UL_Seth.bioMod", + muscles_model=[ + DeltoideusClavicle_A_model, + DeltoideusScapula_P_model, + TRIlong_model, + BIC_long_model, + BIC_brevis_model, + ], + activate_force_length_relationship=True, + activate_force_velocity_relationship=True, + activate_residual_torque=True, +) + +nmpc = NmpcFesMsk.prepare_nmpc( + model=model, + stim_time=list(np.round(np.linspace(0, 1, 11), 2))[:-1], + cycle_len=100, + cycle_duration=1, + n_cycles_simultaneous=1, + n_cycles_to_advance=1, + pulse_width={ + "min": minimum_pulse_width, + "max": 0.0006, + "bimapping": False, + }, + msk_info={"with_residual_torque": True}, + objective={ + "cycling": {"x_center": 0.35, "y_center": 0, "radius": 0.1, "target": "marker"}, + "minimize_muscle_fatigue": True, + "minimize_residual_torque": True, + }, + initial_guess_warm_start=True, + n_threads=8, + control_type=ControlType.CONSTANT, # ControlType.LINEAR_CONTINUOUS don't work for nmpc in bioptim +) + +n_cycles_total = 8 + + +def update_functions(_nmpc, cycle_idx, _sol): + return cycle_idx < n_cycles_total # True if there are still some cycle to perform + + +if __name__ == "__main__": + sol = nmpc.solve( + update_functions, + solver=Solver.IPOPT(_max_iter=10000), + cyclic_options={"states": {}}, + get_all_iterations=True, + # n_cycles_simultaneous=1, + ) + + SolutionToPickle(sol[0], "results/cycling_fes_driven_nmpc_full_fatigue.pkl", "").pickle() + [ + SolutionToPickle(sol[1][i], "results/cycling_fes_driven_nmpc_" + str(i) + "_fatigue.pkl", "").pickle() + for i in range(len(sol[1])) + ] + + biorbd_model = biorbd.Model("../msk_models/simplified_UL_Seth_full_mesh.bioMod") + PickleAnimate("results/cycling_fes_driven_nmpc_full_fatigue.pkl").animate(model=biorbd_model) diff --git a/cocofest/examples/nmpc/plot_cycling_nmpc_results.py b/cocofest/examples/nmpc/plot_cycling_nmpc_results.py new file mode 100644 index 0000000..abb1279 --- /dev/null +++ b/cocofest/examples/nmpc/plot_cycling_nmpc_results.py @@ -0,0 +1,398 @@ +from matplotlib import pyplot as plt +import numpy.ma as ma +import numpy as np +import pickle +import biorbd +from cocofest import PickleAnimate + + +# --- Load results --- # +with open("results/cycling_fes_driven_nmpc_full_force.pkl", "rb") as file: + full_nmpc_result_force_optim = pickle.load(file) + +with open("results/cycling_fes_driven_nmpc_full_fatigue.pkl", "rb") as file: + full_nmpc_result_fatigue_optim = pickle.load(file) + +all_pulse_width_force = [] +all_pulse_width_fatigue = [] +for i in range(8): + with open("results/cycling_fes_driven_nmpc_" + str(i) + "_force.pkl", "rb") as file: + temp_nmpc_result_force_optim = pickle.load(file) + + with open("results/cycling_fes_driven_nmpc_" + str(i) + "_fatigue.pkl", "rb") as file: + temp_nmpc_result_fatigue_optim = pickle.load(file) + + all_pulse_width_force.append(temp_nmpc_result_force_optim["parameters"]) + all_pulse_width_fatigue.append(temp_nmpc_result_fatigue_optim["parameters"]) + +barWidth = 0.25 * (2 / 5) # set width of bar +cycles = [] +pulse_width_force_optim_dict = { + "DeltoideusClavicle_A": [], + "DeltoideusScapula_P": [], + "TRIlong": [], + "BIC_long": [], + "BIC_brevis": [], +} + +pulse_width_fatigue_optim_dict = { + "DeltoideusClavicle_A": [], + "DeltoideusScapula_P": [], + "TRIlong": [], + "BIC_long": [], + "BIC_brevis": [], +} + +for i in range(8): + [ + pulse_width_force_optim_dict[f].append(all_pulse_width_force[i]["pulse_width_" + f]) + for f in pulse_width_force_optim_dict.keys() + ] + [ + pulse_width_fatigue_optim_dict[f].append(all_pulse_width_fatigue[i]["pulse_width_" + f]) + for f in pulse_width_fatigue_optim_dict.keys() + ] + + +for key in pulse_width_force_optim_dict.keys(): + pulse_width_force_optim_dict[key] = np.array([x for xs in pulse_width_force_optim_dict[key] for x in xs]) + pulse_width_fatigue_optim_dict[key] = np.array([x for xs in pulse_width_fatigue_optim_dict[key] for x in xs]) + +bar = np.array([barWidth * (x + 0.5) for x in range(80)]) + +# --- Plot results --- # +# --- First subplot for position and residual torque --- # +fig, axs = plt.subplots(2, 2, figsize=(10, 10)) +axs[0, 0].set_title("Shoulder Q") +axs[1, 0].set_xlabel("Time (s)") +axs[0, 0].set_ylabel("Angle (rad)") +axs[0, 1].set_title("Elbow Q") +axs[1, 1].set_xlabel("Time (s)") + +axs[1, 0].set_title("Shoulder Tau") +axs[1, 0].set_ylabel("Torque (Nm)") +axs[1, 1].set_title("Elbow Tau") + +for i in range(8): + axs[0, 0].plot( + full_nmpc_result_force_optim["time"][:101], + full_nmpc_result_force_optim["states"]["q"][0][100 * i : 100 * (i + 1) + 1], + label="force optim cycle " + str(i), + ) + axs[0, 0].plot( + full_nmpc_result_fatigue_optim["time"][:101], + full_nmpc_result_fatigue_optim["states"]["q"][0][100 * i : 100 * (i + 1) + 1], + label="fatigue optim cycle " + str(i), + ) + + axs[0, 1].plot( + full_nmpc_result_force_optim["time"][:101], + full_nmpc_result_force_optim["states"]["q"][1][100 * i : 100 * (i + 1) + 1], + label="force optim cycle " + str(i), + ) + axs[0, 1].plot( + full_nmpc_result_fatigue_optim["time"][:101], + full_nmpc_result_fatigue_optim["states"]["q"][1][100 * i : 100 * (i + 1) + 1], + label="fatigue optim cycle " + str(i), + ) + +axs[1, 0].plot( + full_nmpc_result_force_optim["time"][:-1], + full_nmpc_result_force_optim["control"]["tau"][0], + label="force optim", +) +axs[1, 0].plot( + full_nmpc_result_fatigue_optim["time"][:-1], + full_nmpc_result_fatigue_optim["control"]["tau"][0], + label="fatigue optim", +) + +axs[1, 1].plot( + full_nmpc_result_force_optim["time"][:-1], + full_nmpc_result_force_optim["control"]["tau"][1], + label="force optim", +) +axs[1, 1].plot( + full_nmpc_result_fatigue_optim["time"][:-1], + full_nmpc_result_fatigue_optim["control"]["tau"][1], + label="fatigue optim", +) + +axs[1, 1].legend() +plt.show() + +# --- Second subplot for pulse width and muscle force / fatigue --- # +fig, axs = plt.subplots(3, 2, figsize=(10, 10)) +axs[0, 0].set_title("DeltoideusClavicle_A") +axs[0, 0].plot( + full_nmpc_result_force_optim["time"], + full_nmpc_result_force_optim["states"]["F_DeltoideusClavicle_A"], + label="force optim", +) +axs[0, 0].plot( + full_nmpc_result_force_optim["time"], + full_nmpc_result_fatigue_optim["states"]["F_DeltoideusClavicle_A"], + label="fatigue optim", +) +axs[0, 0].set_ylabel("Force (N)") +axs[0, 0].set_ylim(bottom=-20) +axs00 = axs[0, 0].twinx() +axs00.set_ylim(top=0.003) +axs00.set_yticks([0, 0.0003, 0.0006], [0, 300, 600]) +mask1 = ma.where( + pulse_width_force_optim_dict["DeltoideusClavicle_A"] >= pulse_width_fatigue_optim_dict["DeltoideusClavicle_A"] +) +mask2 = ma.where( + pulse_width_fatigue_optim_dict["DeltoideusClavicle_A"] >= pulse_width_force_optim_dict["DeltoideusClavicle_A"] +) +axs00.bar( + bar[mask1], + pulse_width_force_optim_dict["DeltoideusClavicle_A"][mask1], + color="tab:blue", + edgecolor="grey", + width=barWidth, + label=f"force optim pw", +) +axs00.bar( + bar, + pulse_width_fatigue_optim_dict["DeltoideusClavicle_A"], + color="tab:orange", + edgecolor="grey", + width=barWidth, + label=f"force optim pw", +) +axs00.bar( + bar[mask2], + pulse_width_force_optim_dict["DeltoideusClavicle_A"][mask2], + color="tab:blue", + edgecolor="grey", + width=barWidth, + label=f"force optim pw", +) + +axs[0, 1].set_title("DeltoideusScapula_P") +axs[0, 1].plot( + full_nmpc_result_force_optim["time"], + full_nmpc_result_force_optim["states"]["F_DeltoideusScapula_P"], + label="force optim", +) +axs[0, 1].plot( + full_nmpc_result_force_optim["time"], + full_nmpc_result_fatigue_optim["states"]["F_DeltoideusScapula_P"], + label="fatigue optim", +) +axs[0, 1].set_ylim(bottom=-20) +axs01 = axs[0, 1].twinx() +axs01.set_ylim(top=0.003) +axs01.set_yticks([0, 0.0003, 0.0006], [0, 300, 600]) +axs01.set_ylabel("pulse width (us)") +mask1 = ma.where( + pulse_width_force_optim_dict["DeltoideusScapula_P"] >= pulse_width_fatigue_optim_dict["DeltoideusScapula_P"] +) +mask2 = ma.where( + pulse_width_fatigue_optim_dict["DeltoideusScapula_P"] >= pulse_width_force_optim_dict["DeltoideusScapula_P"] +) +axs01.bar( + bar[mask1], + pulse_width_force_optim_dict["DeltoideusScapula_P"][mask1], + color="tab:blue", + edgecolor="grey", + width=barWidth, + label=f"force optim pw", +) +axs01.bar( + bar, + pulse_width_fatigue_optim_dict["DeltoideusScapula_P"], + color="tab:orange", + edgecolor="grey", + width=barWidth, + label=f"force optim pw", +) +axs01.bar( + bar[mask2], + pulse_width_force_optim_dict["DeltoideusScapula_P"][mask2], + color="tab:blue", + edgecolor="grey", + width=barWidth, + label=f"force optim pw", +) + +axs[1, 0].set_title("TRIlong") +axs[1, 0].plot( + full_nmpc_result_force_optim["time"], + full_nmpc_result_force_optim["states"]["F_TRIlong"], + label="force optim", +) +axs[1, 0].plot( + full_nmpc_result_force_optim["time"], + full_nmpc_result_fatigue_optim["states"]["F_TRIlong"], + label="fatigue optim", +) +axs[1, 0].set_ylabel("Force (N)") +axs[1, 0].set_ylim(bottom=-20) +axs10 = axs[1, 0].twinx() +axs10.set_ylim(top=0.003) +axs10.set_yticks([0, 0.0003, 0.0006], [0, 300, 600]) +axs10.bar( + bar, + pulse_width_force_optim_dict["TRIlong"], + width=barWidth, + edgecolor="grey", + label=f"force optim pw", +) +axs10.bar( + bar, + pulse_width_fatigue_optim_dict["TRIlong"], + width=barWidth, + edgecolor="grey", + label=f"force optim pw", +) +mask1 = ma.where(pulse_width_force_optim_dict["TRIlong"] >= pulse_width_fatigue_optim_dict["TRIlong"]) +mask2 = ma.where(pulse_width_fatigue_optim_dict["TRIlong"] >= pulse_width_force_optim_dict["TRIlong"]) +axs10.bar( + bar[mask1], + pulse_width_force_optim_dict["TRIlong"][mask1], + color="tab:blue", + edgecolor="grey", + width=barWidth, + label=f"force optim pw", +) +axs10.bar( + bar, + pulse_width_fatigue_optim_dict["TRIlong"], + color="tab:orange", + edgecolor="grey", + width=barWidth, + label=f"force optim pw", +) +axs10.bar( + bar[mask2], + pulse_width_force_optim_dict["TRIlong"][mask2], + color="tab:blue", + edgecolor="grey", + width=barWidth, + label=f"force optim pw", +) + +axs[1, 1].set_title("BIC_long") +axs[1, 1].plot( + full_nmpc_result_force_optim["time"], + full_nmpc_result_force_optim["states"]["F_BIC_long"], + label="force optim", +) +axs[1, 1].plot( + full_nmpc_result_force_optim["time"], + full_nmpc_result_fatigue_optim["states"]["F_BIC_long"], + label="fatigue optim", +) +axs[1, 1].set_ylim(bottom=-20) +axs11 = axs[1, 1].twinx() +axs11.set_ylim(top=0.003) +axs11.set_yticks([0, 0.0003, 0.0006], [0, 300, 600]) +axs11.set_ylabel("pulse width (us)") +mask1 = ma.where(pulse_width_force_optim_dict["BIC_long"] >= pulse_width_fatigue_optim_dict["BIC_long"]) +mask2 = ma.where(pulse_width_fatigue_optim_dict["BIC_long"] >= pulse_width_force_optim_dict["BIC_long"]) +axs11.bar( + bar[mask1], + pulse_width_force_optim_dict["BIC_long"][mask1], + color="tab:blue", + edgecolor="grey", + width=barWidth, + label=f"force optim pw", +) +axs11.bar( + bar, + pulse_width_fatigue_optim_dict["BIC_long"], + color="tab:orange", + edgecolor="grey", + width=barWidth, + label=f"force optim pw", +) +axs11.bar( + bar[mask2], + pulse_width_force_optim_dict["BIC_long"][mask2], + color="tab:blue", + edgecolor="grey", + width=barWidth, + label=f"force optim pw", +) + +axs[2, 0].set_title("BIC_brevis") +axs[2, 0].plot( + full_nmpc_result_force_optim["time"], + full_nmpc_result_force_optim["states"]["F_BIC_brevis"], + label="force optim", +) +axs[2, 0].plot( + full_nmpc_result_force_optim["time"], + full_nmpc_result_fatigue_optim["states"]["F_BIC_brevis"], + label="fatigue optim", +) +axs[2, 0].set_xlabel("Time (s)") +axs[2, 0].set_ylabel("Force (N)") +axs[2, 0].set_ylim(bottom=-20) +axs20 = axs[2, 0].twinx() +axs20.set_ylim(top=0.003) +axs20.set_yticks([0, 0.0003, 0.0006], [0, 300, 600]) +mask1 = ma.where(pulse_width_force_optim_dict["BIC_brevis"] >= pulse_width_fatigue_optim_dict["BIC_brevis"]) +mask2 = ma.where(pulse_width_fatigue_optim_dict["BIC_brevis"] >= pulse_width_force_optim_dict["BIC_brevis"]) +axs20.bar( + bar[mask1], + pulse_width_force_optim_dict["BIC_brevis"][mask1], + color="tab:blue", + edgecolor="grey", + width=barWidth, + label=f"force optim pw", +) +axs20.bar( + bar, + pulse_width_fatigue_optim_dict["BIC_brevis"], + color="tab:orange", + edgecolor="grey", + width=barWidth, + label=f"force optim pw", +) +axs20.bar( + bar[mask2], + pulse_width_force_optim_dict["BIC_brevis"][mask2], + color="tab:blue", + edgecolor="grey", + width=barWidth, + label=f"force optim pw", +) + +axs[2, 1].set_title("General muscle fatigue") +full_nmpc_result_force_optim_fatigue = np.sum( + [full_nmpc_result_force_optim["states"][f] for f in full_nmpc_result_force_optim["states"] if "A_" in f], + axis=0, +) +full_nmpc_result_fatigue_optim_fatigue = np.sum( + [full_nmpc_result_fatigue_optim["states"][f] for f in full_nmpc_result_fatigue_optim["states"] if "A_" in f], + axis=0, +) +axs[2, 1].plot( + full_nmpc_result_force_optim["time"], + full_nmpc_result_force_optim_fatigue, + label="force optim", +) +axs[2, 1].plot( + full_nmpc_result_force_optim["time"], + full_nmpc_result_fatigue_optim_fatigue, + label="fatigue optim", +) +axs[2, 1].set_xlabel("Time (s)") +axs[2, 1].set_ylabel("Force scaling factor (N/s)") +axs21 = axs[2, 1].twinx() +axs21.set_yticks([0], [" "]) +axs21.set_ylabel("pulse width (us)") +axs[2, 1].legend() + +plt.show() + + +# --- Pyorerun animation --- # +biorbd_model = biorbd.Model("../msk_models/simplified_UL_Seth_full_mesh.bioMod") +PickleAnimate("results/cycling_fes_driven_nmpc_full_force.pkl").animate(model=biorbd_model) + +biorbd_model = biorbd.Model("../msk_models/simplified_UL_Seth_full_mesh.bioMod") +PickleAnimate("results/cycling_fes_driven_nmpc_full_fatigue.pkl").animate(model=biorbd_model) diff --git a/cocofest/examples/nmpc/pulse_duration_optimization_musculoskeletal_dynamic_nmpc_cyclic.py b/cocofest/examples/nmpc/pulse_duration_optimization_musculoskeletal_dynamic_nmpc_cyclic.py new file mode 100644 index 0000000..4e292d0 --- /dev/null +++ b/cocofest/examples/nmpc/pulse_duration_optimization_musculoskeletal_dynamic_nmpc_cyclic.py @@ -0,0 +1,61 @@ +""" +This example will do a nmpc of 10 stimulation example with Ding's 2007 frequency model. +This ocp was build to produce a elbow motion from 5 to 120 degrees. +The pulse width between minimal sensitivity threshold and 600us to satisfy the flexion and minimizing required elbow +torque control. +""" + +import os +import biorbd +from bioptim import Solver +from cocofest import ( + DingModelPulseWidthFrequencyWithFatigue, + NmpcFesMsk, + FesMskModel, + SolutionToPickle, + PickleAnimate, +) + +model = FesMskModel( + name=None, + biorbd_path="../msk_models/arm26_biceps_1dof.bioMod", + muscles_model=[DingModelPulseWidthFrequencyWithFatigue(muscle_name="BIClong")], + activate_force_length_relationship=True, + activate_force_velocity_relationship=True, +) + +minimum_pulse_width = DingModelPulseWidthFrequencyWithFatigue().pd0 + +nmpc_fes_msk = NmpcFesMsk() +nmpc = nmpc_fes_msk.prepare_nmpc( + model=model, + stim_time=[0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9], + cycle_len=100, + cycle_duration=1, + pulse_width={ + "min": minimum_pulse_width, + "max": 0.0006, + "bimapping": False, + }, + objective={"minimize_residual_torque": True}, + msk_info={ + "bound_type": "start_end", + "bound_data": [[5], [120]], + "with_residual_torque": True, + }, +) +nmpc_fes_msk.n_cycles = 2 +sol = nmpc.solve( + nmpc_fes_msk.update_functions, + solver=Solver.IPOPT(), + cyclic_options={"states": {}}, + get_all_iterations=True, +) + +biorbd_model = biorbd.Model("../msk_models/arm26_biceps_1dof.bioMod") +temp_pickle_file_path = "pw_optim_dynamic_nmpc_full.pkl" +SolutionToPickle(sol[0], temp_pickle_file_path, "").pickle() + +PickleAnimate(temp_pickle_file_path).animate(model=biorbd_model) + +os.remove(temp_pickle_file_path) diff --git a/cocofest/examples/nmpc/results/cycling_fes_driven_nmpc_0_fatigue.pkl b/cocofest/examples/nmpc/results/cycling_fes_driven_nmpc_0_fatigue.pkl new file mode 100644 index 0000000..497811f Binary files /dev/null and b/cocofest/examples/nmpc/results/cycling_fes_driven_nmpc_0_fatigue.pkl differ diff --git a/cocofest/examples/nmpc/results/cycling_fes_driven_nmpc_0_force.pkl b/cocofest/examples/nmpc/results/cycling_fes_driven_nmpc_0_force.pkl new file mode 100644 index 0000000..383cd0e Binary files /dev/null and b/cocofest/examples/nmpc/results/cycling_fes_driven_nmpc_0_force.pkl differ diff --git a/cocofest/examples/nmpc/results/cycling_fes_driven_nmpc_1_fatigue.pkl b/cocofest/examples/nmpc/results/cycling_fes_driven_nmpc_1_fatigue.pkl new file mode 100644 index 0000000..6882afd Binary files /dev/null and b/cocofest/examples/nmpc/results/cycling_fes_driven_nmpc_1_fatigue.pkl differ diff --git a/cocofest/examples/nmpc/results/cycling_fes_driven_nmpc_1_force.pkl b/cocofest/examples/nmpc/results/cycling_fes_driven_nmpc_1_force.pkl new file mode 100644 index 0000000..50212bf Binary files /dev/null and b/cocofest/examples/nmpc/results/cycling_fes_driven_nmpc_1_force.pkl differ diff --git a/cocofest/examples/nmpc/results/cycling_fes_driven_nmpc_2_fatigue.pkl b/cocofest/examples/nmpc/results/cycling_fes_driven_nmpc_2_fatigue.pkl new file mode 100644 index 0000000..b5a07fa Binary files /dev/null and b/cocofest/examples/nmpc/results/cycling_fes_driven_nmpc_2_fatigue.pkl differ diff --git a/cocofest/examples/nmpc/results/cycling_fes_driven_nmpc_2_force.pkl b/cocofest/examples/nmpc/results/cycling_fes_driven_nmpc_2_force.pkl new file mode 100644 index 0000000..db0641d Binary files /dev/null and b/cocofest/examples/nmpc/results/cycling_fes_driven_nmpc_2_force.pkl differ diff --git a/cocofest/examples/nmpc/results/cycling_fes_driven_nmpc_3_fatigue.pkl b/cocofest/examples/nmpc/results/cycling_fes_driven_nmpc_3_fatigue.pkl new file mode 100644 index 0000000..4962e7c Binary files /dev/null and b/cocofest/examples/nmpc/results/cycling_fes_driven_nmpc_3_fatigue.pkl differ diff --git a/cocofest/examples/nmpc/results/cycling_fes_driven_nmpc_3_force.pkl b/cocofest/examples/nmpc/results/cycling_fes_driven_nmpc_3_force.pkl new file mode 100644 index 0000000..8894a39 Binary files /dev/null and b/cocofest/examples/nmpc/results/cycling_fes_driven_nmpc_3_force.pkl differ diff --git a/cocofest/examples/nmpc/results/cycling_fes_driven_nmpc_4_fatigue.pkl b/cocofest/examples/nmpc/results/cycling_fes_driven_nmpc_4_fatigue.pkl new file mode 100644 index 0000000..7f556c4 Binary files /dev/null and b/cocofest/examples/nmpc/results/cycling_fes_driven_nmpc_4_fatigue.pkl differ diff --git a/cocofest/examples/nmpc/results/cycling_fes_driven_nmpc_4_force.pkl b/cocofest/examples/nmpc/results/cycling_fes_driven_nmpc_4_force.pkl new file mode 100644 index 0000000..ff3a697 Binary files /dev/null and b/cocofest/examples/nmpc/results/cycling_fes_driven_nmpc_4_force.pkl differ diff --git a/cocofest/examples/nmpc/results/cycling_fes_driven_nmpc_5_fatigue.pkl b/cocofest/examples/nmpc/results/cycling_fes_driven_nmpc_5_fatigue.pkl new file mode 100644 index 0000000..36b639b Binary files /dev/null and b/cocofest/examples/nmpc/results/cycling_fes_driven_nmpc_5_fatigue.pkl differ diff --git a/cocofest/examples/nmpc/results/cycling_fes_driven_nmpc_5_force.pkl b/cocofest/examples/nmpc/results/cycling_fes_driven_nmpc_5_force.pkl new file mode 100644 index 0000000..1f5c4a0 Binary files /dev/null and b/cocofest/examples/nmpc/results/cycling_fes_driven_nmpc_5_force.pkl differ diff --git a/cocofest/examples/nmpc/results/cycling_fes_driven_nmpc_6_fatigue.pkl b/cocofest/examples/nmpc/results/cycling_fes_driven_nmpc_6_fatigue.pkl new file mode 100644 index 0000000..c0caa57 Binary files /dev/null and b/cocofest/examples/nmpc/results/cycling_fes_driven_nmpc_6_fatigue.pkl differ diff --git a/cocofest/examples/nmpc/results/cycling_fes_driven_nmpc_6_force.pkl b/cocofest/examples/nmpc/results/cycling_fes_driven_nmpc_6_force.pkl new file mode 100644 index 0000000..9940411 Binary files /dev/null and b/cocofest/examples/nmpc/results/cycling_fes_driven_nmpc_6_force.pkl differ diff --git a/cocofest/examples/nmpc/results/cycling_fes_driven_nmpc_7_fatigue.pkl b/cocofest/examples/nmpc/results/cycling_fes_driven_nmpc_7_fatigue.pkl new file mode 100644 index 0000000..6edb3ae Binary files /dev/null and b/cocofest/examples/nmpc/results/cycling_fes_driven_nmpc_7_fatigue.pkl differ diff --git a/cocofest/examples/nmpc/results/cycling_fes_driven_nmpc_7_force.pkl b/cocofest/examples/nmpc/results/cycling_fes_driven_nmpc_7_force.pkl new file mode 100644 index 0000000..a6944ce Binary files /dev/null and b/cocofest/examples/nmpc/results/cycling_fes_driven_nmpc_7_force.pkl differ diff --git a/cocofest/examples/nmpc/results/cycling_fes_driven_nmpc_full_fatigue.pkl b/cocofest/examples/nmpc/results/cycling_fes_driven_nmpc_full_fatigue.pkl new file mode 100644 index 0000000..fb468c4 Binary files /dev/null and b/cocofest/examples/nmpc/results/cycling_fes_driven_nmpc_full_fatigue.pkl differ diff --git a/cocofest/examples/nmpc/results/cycling_fes_driven_nmpc_full_force.pkl b/cocofest/examples/nmpc/results/cycling_fes_driven_nmpc_full_force.pkl new file mode 100644 index 0000000..65eded9 Binary files /dev/null and b/cocofest/examples/nmpc/results/cycling_fes_driven_nmpc_full_force.pkl differ diff --git a/cocofest/examples/sensitivity/truncation/sensitivity_analysis.py b/cocofest/examples/sensitivity/truncation/sensitivity_analysis.py index 2eb20d4..37e504a 100644 --- a/cocofest/examples/sensitivity/truncation/sensitivity_analysis.py +++ b/cocofest/examples/sensitivity/truncation/sensitivity_analysis.py @@ -58,7 +58,10 @@ "n_stim": n_stim, "pulse_mode": mode, } - ivp_parameters = {"n_shooting": temp_node_shooting, "final_time": 1, "use_sx": True} + ivp_parameters = { + "final_time": 1, + "use_sx": True, + } ivp = IvpFes( fes_parameters=fes_parameters, ivp_parameters=ivp_parameters, diff --git a/cocofest/examples/sensitivity/truncation/summation_truncation_example.py b/cocofest/examples/sensitivity/truncation/summation_truncation_example.py index 0b3426c..ba549ef 100644 --- a/cocofest/examples/sensitivity/truncation/summation_truncation_example.py +++ b/cocofest/examples/sensitivity/truncation/summation_truncation_example.py @@ -21,7 +21,10 @@ "model": DingModelFrequencyWithFatigue(sum_stim_truncation=i if i != 0 else None), "n_stim": n_stim, } - ivp_parameters = {"n_shooting": n_shooting, "final_time": final_time, "use_sx": True} + ivp_parameters = { + "final_time": final_time, + "use_sx": True, + } ivp = IvpFes(fes_parameters, ivp_parameters) @@ -33,8 +36,16 @@ # Plotting the force state result plt.title("Force state result") for i, result in enumerate(results[1:]): - plt.plot(sol_time, result["F"][0], label=f"sum_stim_truncation={i+1}, time={computations_time[i+1]:.2f}s") -plt.plot(sol_time, results[0]["F"][0], label=f"not truncated, time={computations_time[0]:.2f}s") + plt.plot( + sol_time, + result["F"][0], + label=f"sum_stim_truncation={i+1}, time={computations_time[i+1]:.2f}s", + ) +plt.plot( + sol_time, + results[0]["F"][0], + label=f"not truncated, time={computations_time[0]:.2f}s", +) plt.xlabel("time (s)") plt.ylabel("force (N)") plt.legend() diff --git a/cocofest/examples/sensitivity/truncation/summation_truncation_graph.py b/cocofest/examples/sensitivity/truncation/summation_truncation_graph.py index a20e5f7..624b3ff 100644 --- a/cocofest/examples/sensitivity/truncation/summation_truncation_graph.py +++ b/cocofest/examples/sensitivity/truncation/summation_truncation_graph.py @@ -175,7 +175,12 @@ ) axs.plot( - np.arange(1, 101, 1).tolist(), np.arange(1, 101, 1).tolist(), color="red", ls="-", label="Ground truth", linewidth=4 + np.arange(1, 101, 1).tolist(), + np.arange(1, 101, 1).tolist(), + color="red", + ls="-", + label="Ground truth", + linewidth=4, ) x_beneath_1e_8 = np.arange(1, 101, 1).tolist() @@ -184,9 +189,23 @@ y_beneath_1e_8 = [] for j in range(len((all_mode_list_error_beneath_1e_8[i]))): y_beneath_1e_8.append(parameter_list[i][all_mode_list_error_beneath_1e_8[i][j]][1]) - axs.plot(x_beneath_1e_8, y_beneath_1e_8, color="darkred", label=r"Calcium absolute error < 1e⁻⁸", linewidth=3) + axs.plot( + x_beneath_1e_8, + y_beneath_1e_8, + color="darkred", + label=r"Calcium absolute error < 1e⁻⁸", + linewidth=3, + ) -axs.scatter(0, 0, color="white", label="Initialization (s) | 100 Integrations (s)", marker="+", s=0, lw=0) +axs.scatter( + 0, + 0, + color="white", + label="Initialization (s) | 100 Integrations (s)", + marker="+", + s=0, + lw=0, +) axs.scatter( 1, 1, @@ -217,7 +236,11 @@ axs.set_xlabel("Frequency (Hz)", fontsize=25, fontname="Times New Roman") axs.xaxis.set_major_locator(MaxNLocator(integer=True)) -axs.set_ylabel("Past stimulations kept for computation (n)", fontsize=25, fontname="Times New Roman") +axs.set_ylabel( + "Past stimulations kept for computation (n)", + fontsize=25, + fontname="Times New Roman", +) axs.yaxis.set_major_locator(MaxNLocator(integer=True)) ticks = np.arange(1, 101, 1).tolist() diff --git a/cocofest/identification/ding2003_force_parameter_identification.py b/cocofest/identification/ding2003_force_parameter_identification.py index ab829d0..2e7c5d8 100644 --- a/cocofest/identification/ding2003_force_parameter_identification.py +++ b/cocofest/identification/ding2003_force_parameter_identification.py @@ -1,7 +1,9 @@ import time as time_package +import numpy as np -from bioptim import Solver, Objective, OdeSolver +from bioptim import Solver, OdeSolver, ControlType +from ..optimization.fes_ocp import OcpFes from ..models.fes_model import FesModel from ..models.ding2003 import DingModelFrequency from ..optimization.fes_identification_ocp import OcpFesId @@ -9,7 +11,6 @@ full_data_extraction, average_data_extraction, sparse_data_extraction, - node_shooting_list_creation, force_at_node_in_ocp, ) from .identification_abstract_class import ParameterIdentification @@ -29,11 +30,12 @@ def __init__( double_step_identification: bool = False, key_parameter_to_identify: list = None, additional_key_settings: dict = None, - n_shooting: int = 5, - custom_objective: list[Objective] = None, + final_time: float = 1, + objective: dict = None, use_sx: bool = True, ode_solver: OdeSolver = OdeSolver.RK4(n_integration_steps=1), n_threads: int = 1, + control_type: ControlType = ControlType.CONSTANT, **kwargs, ): """ @@ -55,9 +57,7 @@ def __init__( additional_key_settings: dict, additional_key_settings will enable to modify identified parameters default parameters such as initial guess, min_bound, max_bound, function and scaling - n_shooting: int, - The number of shooting points for the ocp - custom_objective: list[Objective], + objective: dict, The custom objective to use for the identification use_sx: bool The nature of the casadi variables. MX are used if False. @@ -65,13 +65,15 @@ def __init__( The ode solver to use for the identification n_threads: int, The number of threads to use for the identification + control_type: ControlType, + The type of control to use for the identification """ self.default_values = self._set_default_values(model=model) dict_parameter_to_configure = model.identifiable_parameters model_parameters_value = [ - None if key in key_parameter_to_identify else dict_parameter_to_configure[key] + (None if key in key_parameter_to_identify else dict_parameter_to_configure[key]) for key in dict_parameter_to_configure ] self.model = self._set_model_parameters(model, model_parameters_value) @@ -83,7 +85,6 @@ def __init__( double_step_identification, key_parameter_to_identify, additional_key_settings, - n_shooting, ) self.key_parameter_to_identify = key_parameter_to_identify @@ -95,11 +96,12 @@ def __init__( self.force_ocp = None self.force_identification_result = None - self.n_shooting = n_shooting - self.custom_objective = custom_objective + self.final_time = final_time + self.objective = objective self.use_sx = use_sx self.ode_solver = ode_solver self.n_threads = n_threads + self.control_type = control_type self.kwargs = kwargs def _set_default_values(self, model): @@ -129,21 +131,21 @@ def _set_default_values(self, model): "min_bound": 0.001, "max_bound": 1, "function": model.set_km_rest, - "scaling": 1, # 1000 + "scaling": 1, }, "tau1_rest": { "initial_guess": 0.5, "min_bound": 0.0001, "max_bound": 1, "function": model.set_tau1_rest, - "scaling": 1, # 1000 + "scaling": 1, }, "tau2": { "initial_guess": 0.5, "min_bound": 0.0001, "max_bound": 1, "function": model.set_tau2, - "scaling": 1, # 1000 + "scaling": 1, }, } @@ -151,7 +153,12 @@ def _set_default_parameters_list(self): """ This method is used to set the default parameters list for the model. """ - self.numeric_parameters = [self.model.a_rest, self.model.km_rest, self.model.tau1_rest, self.model.tau2] + self.numeric_parameters = [ + self.model.a_rest, + self.model.km_rest, + self.model.tau1_rest, + self.model.tau2, + ] self.key_parameters = ["a_rest", "km_rest", "tau1_rest", "tau2"] def input_sanity( @@ -162,7 +169,6 @@ def input_sanity( double_step_identification: bool = None, key_parameter_to_identify: list = None, additional_key_settings: dict = None, - n_shooting: int = None, ): """ This method is used to check the input sanity entered from the user. @@ -185,8 +191,6 @@ def input_sanity( additional_key_settings: dict, additional_key_settings will enable to modify identified parameters default parameters such as initial guess, min_bound, max_bound, function and scaling - n_shooting: int, - The number of shooting points for the ocp """ if model._with_fatigue: @@ -238,7 +242,8 @@ def input_sanity( f" the available values are {list(self.default_values[key].keys())}" ) if not isinstance( - additional_key_settings[key][setting_name], type(self.default_values[key][setting_name]) + additional_key_settings[key][setting_name], + type(self.default_values[key][setting_name]), ): raise TypeError( f"The given additional_key_settings value is not valid," @@ -251,9 +256,6 @@ def input_sanity( f" the given value is {type(additional_key_settings)} type" ) - if not isinstance(n_shooting, int): - raise TypeError(f"The given n_shooting must be int type," f" the given value is {type(n_shooting)} type") - self._set_default_parameters_list() if not all(isinstance(param, None | int | float) for param in self.numeric_parameters): raise ValueError(f"The given model parameters are not valid, only None, int and float are accepted") @@ -326,31 +328,30 @@ def _force_model_identification_for_initial_guess(self): self.double_step_identification, self.key_parameter_to_identify, self.additional_key_settings, - self.n_shooting, ) self.check_experiment_force_format(self.data_path) # --- Data extraction --- # # --- Force model --- # - stimulated_n_shooting = self.n_shooting force_curve_number = None time, stim, force, discontinuity = average_data_extraction(self.data_path) - n_shooting, final_time_phase = node_shooting_list_creation(stim, stimulated_n_shooting) - force_at_node = force_at_node_in_ocp(time, force, n_shooting, final_time_phase, force_curve_number) + + n_shooting = OcpFes.prepare_n_shooting(stim, self.final_time) + force_at_node = force_at_node_in_ocp(time, force, n_shooting, self.final_time, force_curve_number) # --- Building force ocp --- # self.force_ocp = OcpFesId.prepare_ocp( model=self.model, - n_shooting=n_shooting, - final_time_phase=final_time_phase, - force_tracking=force_at_node, + stim_time=list(np.round(stim, 3)), + final_time_phase=self.final_time, key_parameter_to_identify=self.key_parameter_to_identify, additional_key_settings=self.additional_key_settings, - custom_objective=self.custom_objective, + objective={"force_tracking": force_at_node}, discontinuity_in_ocp=discontinuity, use_sx=self.use_sx, ode_solver=self.ode_solver, n_threads=self.n_threads, + control_type=self.control_type, ) self.force_identification_result = self.force_ocp.solve( @@ -372,13 +373,11 @@ def force_model_identification(self): self.double_step_identification, self.key_parameter_to_identify, self.additional_key_settings, - self.n_shooting, ) self.check_experiment_force_format(self.data_path) # --- Data extraction --- # # --- Force model --- # - stimulated_n_shooting = self.n_shooting force_curve_number = None time, stim, force, discontinuity = None, None, None, None @@ -392,8 +391,8 @@ def force_model_identification(self): force_curve_number = self.kwargs["force_curve_number"] if "force_curve_number" in self.kwargs else 5 time, stim, force, discontinuity = sparse_data_extraction(self.data_path, force_curve_number) - n_shooting, final_time_phase = node_shooting_list_creation(stim, stimulated_n_shooting) - force_at_node = force_at_node_in_ocp(time, force, n_shooting, final_time_phase, force_curve_number) + n_shooting = OcpFes.prepare_n_shooting(stim, self.final_time) + force_at_node = force_at_node_in_ocp(time, force, n_shooting, self.final_time, force_curve_number) if self.double_step_identification: initial_guess = self._force_model_identification_for_initial_guess() @@ -405,16 +404,16 @@ def force_model_identification(self): start_time = time_package.time() self.force_ocp = OcpFesId.prepare_ocp( model=self.model, - n_shooting=n_shooting, - final_time_phase=final_time_phase, - force_tracking=force_at_node, + final_time=self.final_time, + stim_time=list(np.round(stim, 3)), + objective={"force_tracking": force_at_node}, key_parameter_to_identify=self.key_parameter_to_identify, additional_key_settings=self.additional_key_settings, - custom_objective=self.custom_objective, discontinuity_in_ocp=discontinuity, use_sx=self.use_sx, ode_solver=self.ode_solver, n_threads=self.n_threads, + control_type=self.control_type, ) print(f"OCP creation time : {time_package.time() - start_time} seconds") diff --git a/cocofest/identification/ding2007_force_parameter_identification.py b/cocofest/identification/ding2007_force_parameter_identification.py index c708f3d..d11e446 100644 --- a/cocofest/identification/ding2007_force_parameter_identification.py +++ b/cocofest/identification/ding2007_force_parameter_identification.py @@ -1,20 +1,22 @@ import time as time_package import numpy as np -from bioptim import Solver, Objective, OdeSolver -from ..models.ding2007 import DingModelPulseDurationFrequency -from ..identification.ding2003_force_parameter_identification import DingModelFrequencyForceParameterIdentification +from bioptim import Solver, OdeSolver, ControlType +from ..models.ding2007 import DingModelPulseWidthFrequency +from ..identification.ding2003_force_parameter_identification import ( + DingModelFrequencyForceParameterIdentification, +) +from ..optimization.fes_ocp import OcpFes from ..optimization.fes_identification_ocp import OcpFesId from .identification_method import ( full_data_extraction, average_data_extraction, sparse_data_extraction, - node_shooting_list_creation, force_at_node_in_ocp, ) -class DingModelPulseDurationFrequencyForceParameterIdentification(DingModelFrequencyForceParameterIdentification): +class DingModelPulseWidthFrequencyForceParameterIdentification(DingModelFrequencyForceParameterIdentification): """ This class extends the DingModelFrequencyForceParameterIdentification class and is used to define an optimal control problem (OCP). It prepares the full program and provides all the necessary parameters to solve a functional electrical stimulation OCP. @@ -22,25 +24,26 @@ class DingModelPulseDurationFrequencyForceParameterIdentification(DingModelFrequ def __init__( self, - model: DingModelPulseDurationFrequency, + model: DingModelPulseWidthFrequency, data_path: str | list[str] = None, identification_method: str = "full", double_step_identification: bool = False, key_parameter_to_identify: list = None, additional_key_settings: dict = None, - n_shooting: int = 5, - custom_objective: list[Objective] = None, + final_time: float = 1, + objective: dict = None, use_sx: bool = True, ode_solver: OdeSolver = OdeSolver.RK4(n_integration_steps=1), n_threads: int = 1, + control_type: ControlType = ControlType.CONSTANT, **kwargs, ): """ - Initializes the DingModelPulseDurationFrequencyForceParameterIdentification class. + Initializes the DingModelPulseWidthFrequencyForceParameterIdentification class. Parameters ---------- - model: DingModelPulseDurationFrequency + model: DingModelPulseWidthFrequency The model to use for the OCP. data_path: str | list[str] The path to the force model data. @@ -53,9 +56,9 @@ def __init__( List of keys of the parameters to identify. additional_key_settings: dict Additional settings for the keys. - n_shooting: int - The number of shooting points for the OCP. - custom_objective: list[Objective] + final_time: float + The final time for the OCP. + objective: dict List of custom objectives. use_sx: bool The nature of the CasADi variables. MX are used if False. @@ -63,21 +66,24 @@ def __init__( The ODE solver to use. n_threads: int The number of threads to use while solving (multi-threading if > 1). + control_type: ControlType, + The type of control to use for the identification **kwargs: dict Additional keyword arguments. """ - super(DingModelPulseDurationFrequencyForceParameterIdentification, self).__init__( + super(DingModelPulseWidthFrequencyForceParameterIdentification, self).__init__( model=model, data_path=data_path, identification_method=identification_method, double_step_identification=double_step_identification, key_parameter_to_identify=key_parameter_to_identify, additional_key_settings=additional_key_settings, - n_shooting=n_shooting, - custom_objective=custom_objective, + final_time=final_time, + objective=objective, use_sx=use_sx, ode_solver=ode_solver, n_threads=n_threads, + control_type=control_type, **kwargs, ) @@ -181,9 +187,9 @@ def _set_model_parameters(model, model_parameters_value): return model @staticmethod - def pulse_duration_extraction(data_path: str) -> list[float]: + def pulse_width_extraction(data_path: str) -> list[float]: """ - Extracts the pulse duration from the data. + Extracts the pulse width from the data. Parameters ---------- @@ -193,17 +199,17 @@ def pulse_duration_extraction(data_path: str) -> list[float]: Returns ------- list[float] - A list of pulse durations. + A list of pulse widths. """ import pickle - pulse_duration = [] + pulse_width = [] for i in range(len(data_path)): with open(data_path[i], "rb") as f: data = pickle.load(f) - pulse_duration.append(data["pulse_duration"]) - pulse_duration = [item for sublist in pulse_duration for item in sublist] - return pulse_duration + pulse_width.append(data["pulse_width"]) + pulse_width = [item for sublist in pulse_width for item in sublist] + return pulse_width def _force_model_identification_for_initial_guess(self): """ @@ -221,34 +227,33 @@ def _force_model_identification_for_initial_guess(self): self.double_step_identification, self.key_parameter_to_identify, self.additional_key_settings, - self.n_shooting, ) self.check_experiment_force_format(self.data_path) # --- Data extraction --- # # --- Force model --- # - stimulated_n_shooting = self.n_shooting force_curve_number = None time, stim, force, discontinuity = average_data_extraction(self.data_path) - pulse_duration = self.pulse_duration_extraction(self.data_path) - n_shooting, final_time_phase = node_shooting_list_creation(stim, stimulated_n_shooting) - force_at_node = force_at_node_in_ocp(time, force, n_shooting, final_time_phase, force_curve_number) + pulse_width = {"fixed": self.pulse_width_extraction(self.data_path)} + + n_shooting = OcpFes.prepare_n_shooting(stim, self.final_time) + force_at_node = force_at_node_in_ocp(time, force, n_shooting, self.final_time, force_curve_number) # --- Building force ocp --- # self.force_ocp = OcpFesId.prepare_ocp( model=self.model, - n_shooting=n_shooting, - final_time_phase=final_time_phase, - force_tracking=force_at_node, - custom_objective=self.custom_objective, + stim_time=list(np.round(stim, 3)), + final_time_phase=self.final_time, + objective={"force_tracking": force_at_node}, discontinuity_in_ocp=discontinuity, - pulse_duration=pulse_duration, + pulse_width=pulse_width, km_rest=self.km_rest, tau1_rest=self.tau1_rest, tau2=self.tau2, use_sx=self.use_sx, ode_solver=self.ode_solver, n_threads=self.n_threads, + control_type=self.control_type, ) self.force_identification_result = self.force_ocp.solve( @@ -279,33 +284,32 @@ def force_model_identification(self) -> dict[str, np.ndarray]: self.double_step_identification, self.key_parameter_to_identify, self.additional_key_settings, - self.n_shooting, ) self.check_experiment_force_format(self.data_path) # --- Data extraction --- # # --- Force model --- # - stimulated_n_shooting = self.n_shooting force_curve_number = None stim = None time = None force = None + pulse_width = {} if self.force_model_identification_method == "full": time, stim, force, discontinuity = full_data_extraction(self.data_path) - pulse_duration = self.pulse_duration_extraction(self.data_path) + pulse_width["fixed"] = self.pulse_width_extraction(self.data_path) elif self.force_model_identification_method == "average": time, stim, force, discontinuity = average_data_extraction(self.data_path) - pulse_duration = np.mean(np.array(self.pulse_duration_extraction(self.data_path))) + pulse_width["fixed"] = np.mean(np.array(self.pulse_width_extraction(self.data_path))) elif self.force_model_identification_method == "sparse": force_curve_number = self.kwargs["force_curve_number"] if "force_curve_number" in self.kwargs else 5 time, stim, force, discontinuity = sparse_data_extraction(self.data_path, force_curve_number) - pulse_duration = self.pulse_duration_extraction(self.data_path) # TODO : adapt this for sparse data + pulse_width["fixed"] = self.pulse_width_extraction(self.data_path) # TODO : adapt this for sparse data - n_shooting, final_time_phase = node_shooting_list_creation(stim, stimulated_n_shooting) - force_at_node = force_at_node_in_ocp(time, force, n_shooting, final_time_phase, force_curve_number) + n_shooting = OcpFes.prepare_n_shooting(stim, self.final_time) + force_at_node = force_at_node_in_ocp(time, force, n_shooting, self.final_time, force_curve_number) if self.double_step_identification: initial_guess = self._force_model_identification_for_initial_guess() @@ -317,17 +321,17 @@ def force_model_identification(self) -> dict[str, np.ndarray]: start_time = time_package.time() self.force_ocp = OcpFesId.prepare_ocp( model=self.model, - n_shooting=n_shooting, - final_time_phase=final_time_phase, - force_tracking=force_at_node, + final_time=self.final_time, + stim_time=list(np.round(stim, 3)), key_parameter_to_identify=self.key_parameter_to_identify, additional_key_settings=self.additional_key_settings, - custom_objective=self.custom_objective, + objective={"force_tracking": force_at_node}, discontinuity_in_ocp=discontinuity, - pulse_duration=pulse_duration, + pulse_width=pulse_width, use_sx=self.use_sx, ode_solver=self.ode_solver, n_threads=self.n_threads, + control_type=self.control_type, ) print(f"OCP creation time : {time_package.time() - start_time} seconds") diff --git a/cocofest/identification/hmed2018_force_parameter_identification.py b/cocofest/identification/hmed2018_force_parameter_identification.py index 8f5a537..9cbfcb3 100644 --- a/cocofest/identification/hmed2018_force_parameter_identification.py +++ b/cocofest/identification/hmed2018_force_parameter_identification.py @@ -1,15 +1,17 @@ import time as time_package import numpy as np -from bioptim import Solver, Objective, OdeSolver -from ..models.hmed2018 import DingModelIntensityFrequency -from ..identification.ding2003_force_parameter_identification import DingModelFrequencyForceParameterIdentification +from bioptim import Solver, Objective, OdeSolver, ControlType +from ..models.hmed2018 import DingModelPulseIntensityFrequency +from ..identification.ding2003_force_parameter_identification import ( + DingModelFrequencyForceParameterIdentification, +) +from ..optimization.fes_ocp import OcpFes from ..optimization.fes_identification_ocp import OcpFesId from .identification_method import ( full_data_extraction, average_data_extraction, sparse_data_extraction, - node_shooting_list_creation, force_at_node_in_ocp, ) @@ -22,17 +24,18 @@ class DingModelPulseIntensityFrequencyForceParameterIdentification(DingModelFreq def __init__( self, - model: DingModelIntensityFrequency, + model: DingModelPulseIntensityFrequency, data_path: str | list[str] = None, identification_method: str = "full", double_step_identification: bool = False, key_parameter_to_identify: list = None, additional_key_settings: dict = None, - n_shooting: int = 5, + final_time: float = 1, custom_objective: list[Objective] = None, use_sx: bool = True, ode_solver: OdeSolver = OdeSolver.RK4(n_integration_steps=1), n_threads: int = 1, + control_type: ControlType = ControlType.CONSTANT, **kwargs, ): """ @@ -40,7 +43,7 @@ def __init__( Parameters ---------- - model: DingModelPulseDurationFrequency + model: DingModelPulseWidthFrequency The model to use for the OCP. data_path: str | list[str] The path to the force model data. @@ -53,8 +56,6 @@ def __init__( List of keys of the parameters to identify. additional_key_settings: dict Additional settings for the keys. - n_shooting: int - The number of shooting points for the OCP. custom_objective: list[Objective] List of custom objectives. use_sx: bool @@ -63,6 +64,8 @@ def __init__( The ODE solver to use. n_threads: int The number of threads to use while solving (multi-threading if > 1). + control_type: ControlType, + The type of control to use for the identification **kwargs: dict Additional keyword arguments. """ @@ -73,11 +76,12 @@ def __init__( double_step_identification=double_step_identification, key_parameter_to_identify=key_parameter_to_identify, additional_key_settings=additional_key_settings, - n_shooting=n_shooting, + final_time=final_time, custom_objective=custom_objective, use_sx=use_sx, ode_solver=ode_solver, n_threads=n_threads, + control_type=control_type, **kwargs, ) @@ -139,7 +143,13 @@ def _set_default_values(self, model): "function": model.set_bs, "scaling": 1, # 1000 }, - "Is": {"initial_guess": 50, "min_bound": 1, "max_bound": 150, "function": model.set_Is, "scaling": 1}, + "Is": { + "initial_guess": 50, + "min_bound": 1, + "max_bound": 150, + "function": model.set_Is, + "scaling": 1, + }, "cr": { "initial_guess": 1, "min_bound": 0.01, @@ -163,7 +173,16 @@ def _set_default_parameters_list(self): self.model.Is, self.model.cr, ] - self.key_parameters = ["a_rest", "km_rest", "tau1_rest", "tau2", "ar", "bs", "Is", "cr"] + self.key_parameters = [ + "a_rest", + "km_rest", + "tau1_rest", + "tau2", + "ar", + "bs", + "Is", + "cr", + ] @staticmethod def _set_model_parameters(model, model_parameters_value): @@ -233,28 +252,24 @@ def _force_model_identification_for_initial_guess(self): self.double_step_identification, self.key_parameter_to_identify, self.additional_key_settings, - self.n_shooting, ) self.check_experiment_force_format(self.data_path) # --- Data extraction --- # # --- Force model --- # - stimulated_n_shooting = self.n_shooting - force_curve_number = None time, stim, force, discontinuity = average_data_extraction(self.data_path) pulse_intensity = self.pulse_intensity_extraction(self.data_path) - n_shooting, final_time_phase = node_shooting_list_creation(stim, stimulated_n_shooting) - force_at_node = force_at_node_in_ocp(time, force, n_shooting, final_time_phase, force_curve_number) + n_shooting = OcpFes.prepare_n_shooting(stim, self.final_time) + force_at_node = force_at_node_in_ocp(time, force, n_shooting, self.final_time, None) # --- Building force ocp --- # self.force_ocp = OcpFesId.prepare_ocp( model=self.model, - n_shooting=n_shooting, - final_time_phase=final_time_phase, - force_tracking=force_at_node, - custom_objective=self.custom_objective, + final_time=self.final_time, + stim_time=list(np.round(stim, 3)), + objective={"force_tracking": force_at_node}, discontinuity_in_ocp=discontinuity, - pulse_intensity=pulse_intensity, + pulse_intensity={"fixed": pulse_intensity}, a_rest=self.a_rest, km_rest=self.km_rest, tau1_rest=self.tau1_rest, @@ -262,6 +277,7 @@ def _force_model_identification_for_initial_guess(self): use_sx=self.use_sx, ode_solver=self.ode_solver, n_threads=self.n_threads, + control_type=self.control_type, ) self.force_identification_result = self.force_ocp.solve( @@ -291,13 +307,11 @@ def force_model_identification(self): self.double_step_identification, self.key_parameter_to_identify, self.additional_key_settings, - self.n_shooting, ) self.check_experiment_force_format(self.data_path) # --- Data extraction --- # # --- Force model --- # - stimulated_n_shooting = self.n_shooting force_curve_number = None stim = None time = None @@ -316,8 +330,8 @@ def force_model_identification(self): time, stim, force, discontinuity = sparse_data_extraction(self.data_path, force_curve_number) pulse_intensity = self.pulse_intensity_extraction(self.data_path) # TODO : adapt this for sparse data - n_shooting, final_time_phase = node_shooting_list_creation(stim, stimulated_n_shooting) - force_at_node = force_at_node_in_ocp(time, force, n_shooting, final_time_phase, force_curve_number) + n_shooting = OcpFes.prepare_n_shooting(stim, self.final_time) + force_at_node = force_at_node_in_ocp(time, force, n_shooting, self.final_time, force_curve_number) if self.double_step_identification: initial_guess = self._force_model_identification_for_initial_guess() @@ -329,17 +343,17 @@ def force_model_identification(self): start_time = time_package.time() self.force_ocp = OcpFesId.prepare_ocp( model=self.model, - n_shooting=n_shooting, - final_time_phase=final_time_phase, - force_tracking=force_at_node, + final_time=self.final_time, + stim_time=list(np.round(stim, 3)), + objective={"force_tracking": force_at_node}, key_parameter_to_identify=self.key_parameter_to_identify, additional_key_settings=self.additional_key_settings, - custom_objective=self.custom_objective, discontinuity_in_ocp=discontinuity, - pulse_intensity=pulse_intensity, + pulse_intensity={"fixed": pulse_intensity}, use_sx=self.use_sx, ode_solver=self.ode_solver, n_threads=self.n_threads, + control_type=self.control_type, ) print(f"OCP creation time : {time_package.time() - start_time} seconds") diff --git a/cocofest/identification/identification_abstract_class.py b/cocofest/identification/identification_abstract_class.py index a912700..e0954b5 100644 --- a/cocofest/identification/identification_abstract_class.py +++ b/cocofest/identification/identification_abstract_class.py @@ -29,7 +29,6 @@ def input_sanity( double_step_identification, key_parameter_to_identify, additional_key_settings, - n_shooting, ): """ diff --git a/cocofest/identification/identification_method.py b/cocofest/identification/identification_method.py index c641520..95bab28 100644 --- a/cocofest/identification/identification_method.py +++ b/cocofest/identification/identification_method.py @@ -127,11 +127,9 @@ def average_data_extraction(model_data_path): model_time_data = [item for sublist in model_time_data for item in sublist] model_time_data = model_time_data[:smallest_list] - train_duration = 1 + train_width = 1 - average_stim_apparition = np.linspace(0, train_duration, int(stimulation_temp_frequency * train_duration) + 1)[ - :-1 - ] + average_stim_apparition = np.linspace(0, train_width, int(stimulation_temp_frequency * train_duration) + 1)[:-1] average_stim_apparition = [time for time in average_stim_apparition] if i == len(model_data_path) - 1: average_stim_apparition = np.append(average_stim_apparition, model_time_data[-1]).tolist() @@ -252,7 +250,7 @@ def sparse_data_extraction(model_data_path, force_curve_number=5): # ) -def force_at_node_in_ocp(time, force, n_shooting, final_time_phase, sparse=None): +def force_at_node_in_ocp(time, force, n_shooting, final_time, sparse=None): """ Interpolates the force at each node in the optimal control problem (OCP). @@ -262,10 +260,8 @@ def force_at_node_in_ocp(time, force, n_shooting, final_time_phase, sparse=None) List of time data. force : list List of force data. - n_shooting : list + n_shooting : int List of number of shooting points for each phase. - final_time_phase : list - List of final time for each phase. sparse : int, optional Number of sparse points, by default None. @@ -275,52 +271,8 @@ def force_at_node_in_ocp(time, force, n_shooting, final_time_phase, sparse=None) List of force at each node in the OCP. """ - temp_time = [] - for i in range(len(final_time_phase)): - for j in range(n_shooting[i]): - temp_time.append(sum(final_time_phase[:i]) + j * final_time_phase[i] / (n_shooting[i])) - force_at_node = np.interp(temp_time, time, force).tolist() + temp_time = np.linspace(0, final_time, n_shooting + 1) + force_at_node = list(np.interp(temp_time, time, force)) # if sparse: # TODO check this part # force_at_node = force_at_node[0:sparse] + force_at_node[:-sparse] return force_at_node - - -def node_shooting_list_creation(stim, stimulated_n_shooting): - """ - Creates a list of node shooting points. - - Parameters - ---------- - stim : list - List of stimulation times. - stimulated_n_shooting : int - Number of shooting points for stimulated phase. - - Returns - ------- - tuple - A tuple containing the list of number of shooting points for each phase and the final time for each phase. - """ - - first_final_time = stim[1] if stim[0] == 0 else stim[0] - final_time_phase = (first_final_time,) - for i in range(1, len(stim)): - final_time_phase = final_time_phase + (stim[i] - stim[i - 1],) - - threshold_stimulation_interval = np.mean(final_time_phase) - stimulation_interval_average_without_rest_time = np.delete( - np.array(final_time_phase), - np.where(np.logical_or(final_time_phase > threshold_stimulation_interval, np.array(final_time_phase) == 0)), - ) - stimulation_interval_average = np.mean(stimulation_interval_average_without_rest_time) - n_shooting = [] - - for i in range(len(final_time_phase)): - if final_time_phase[i] > threshold_stimulation_interval: - temp_final_time = final_time_phase[i] - rest_n_shooting = int((temp_final_time / stimulation_interval_average) * stimulated_n_shooting) - n_shooting.append(rest_n_shooting) - else: - n_shooting.append(stimulated_n_shooting) - - return n_shooting, final_time_phase diff --git a/cocofest/integration/ivp_fes.py b/cocofest/integration/ivp_fes.py index 590f9fb..8b2740d 100644 --- a/cocofest/integration/ivp_fes.py +++ b/cocofest/integration/ivp_fes.py @@ -16,12 +16,13 @@ SolutionMerge, ) +from ..optimization.fes_ocp import OcpFes from ..models.fes_model import FesModel from ..models.ding2003 import DingModelFrequency -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 class IvpFes: @@ -49,10 +50,10 @@ def __init__( ---------- fes_parameters: dict The parameters for the fes configuration including : - model (FesModel type), n_stim (int type), pulse_duration (float type), pulse_intensity (int | float type), pulse_mode (str type), frequency (int | float type), round_down (bool type) + model (FesModel type), stim_time (list), pulse_width (float type), pulse_intensity (int | float type), pulse_mode (str type), frequency (int | float type), round_down (bool type) ivp_parameters: dict The parameters for the ivp problem including : - n_shooting (int type), final_time (int | float type), extend_last_phase_time (int | float type), ode_solver (OdeSolver type), use_sx (bool type), n_threads (int type) + final_time (int | float type), ode_solver (OdeSolver type), use_sx (bool type), n_threads (int type) """ self._fill_fes_dict(fes_parameters) @@ -60,59 +61,75 @@ def __init__( self.dictionaries_check() self.model = self.fes_parameters["model"] - self.n_stim = self.fes_parameters["n_stim"] - self.pulse_duration = self.fes_parameters["pulse_duration"] + self.n_stim = len(self.fes_parameters["stim_time"]) + self.stim_time = self.fes_parameters["stim_time"] + self.pulse_width = self.fes_parameters["pulse_width"] self.pulse_intensity = self.fes_parameters["pulse_intensity"] self.parameter_mappings = None self.parameters = None - self.models = [self.model] * self.n_stim self.final_time = self.ivp_parameters["final_time"] - n_shooting = self.ivp_parameters["n_shooting"] - self.n_shooting = [n_shooting] * self.n_stim if isinstance(n_shooting, int) else n_shooting - if len(self.n_shooting) != self.n_stim: - raise ValueError("n_shooting must be an int or a list of length n_stim") + self.n_shooting = OcpFes.prepare_n_shooting(self.stim_time, self.final_time) - self.dt = [] + self.dt = np.array([self.final_time / self.n_shooting]) self.pulse_mode = self.fes_parameters["pulse_mode"] - self.extend_last_phase_time = self.ivp_parameters["extend_last_phase_time"] self._pulse_mode_settings() parameters = ParameterList(use_sx=self.ivp_parameters["use_sx"]) parameters_init = InitialGuessList() parameters_bounds = BoundsList() - if isinstance(self.model, DingModelPulseDurationFrequency | DingModelPulseDurationFrequencyWithFatigue): - if isinstance(self.pulse_duration, int | float): - parameters_init["pulse_duration"] = np.array([self.pulse_duration] * self.n_stim) + parameters.add( + name="pulse_apparition_time", + function=DingModelFrequency.set_pulse_apparition_time, + size=self.n_stim, + scaling=VariableScaling("pulse_apparition_time", [1] * self.n_stim), + ) + + parameters_init["pulse_apparition_time"] = np.array(self.stim_time) + parameters_bounds.add( + "pulse_apparition_time", + min_bound=np.array(self.stim_time), + max_bound=np.array(self.stim_time), + interpolation=InterpolationType.CONSTANT, + ) + + if isinstance( + self.model, + DingModelPulseWidthFrequency | DingModelPulseWidthFrequencyWithFatigue, + ): + if isinstance(self.pulse_width, int | float): + parameters_init["pulse_width"] = np.array([self.pulse_width] * self.n_stim) parameters_bounds.add( - "pulse_duration", - min_bound=np.array([self.pulse_duration] * (self.n_stim + 1)), - max_bound=np.array([self.pulse_duration] * (self.n_stim + 1)), + "pulse_width", + min_bound=np.array([self.pulse_width]), + max_bound=np.array([self.pulse_width]), interpolation=InterpolationType.CONSTANT, ) - else: - parameters_init["pulse_duration"] = np.array(self.pulse_duration) + parameters_init["pulse_width"] = np.array(self.pulse_width) parameters_bounds.add( - "pulse_duration", - min_bound=np.array(self.pulse_duration), - max_bound=np.array(self.pulse_duration), + "pulse_width", + min_bound=np.array(self.pulse_width), + max_bound=np.array(self.pulse_width), interpolation=InterpolationType.CONSTANT, ) parameters.add( - name="pulse_duration", - function=DingModelPulseDurationFrequency.set_impulse_duration, + name="pulse_width", + function=DingModelPulseWidthFrequency.set_impulse_width, size=self.n_stim, - scaling=VariableScaling("pulse_duration", [1] * self.n_stim), + scaling=VariableScaling("pulse_width", [1] * self.n_stim), ) - if parameters_init["pulse_duration"].shape[0] != self.n_stim: - raise ValueError("pulse_duration list must have the same length as n_stim") + if parameters_init["pulse_width"].shape[0] != self.n_stim: + raise ValueError("pulse_width list must have the same length as n_stim") - if isinstance(self.model, DingModelIntensityFrequency | DingModelIntensityFrequencyWithFatigue): + if isinstance( + self.model, + DingModelPulseIntensityFrequency | DingModelPulseIntensityFrequencyWithFatigue, + ): if isinstance(self.pulse_intensity, int | float): parameters_init["pulse_intensity"] = np.array([self.pulse_intensity] * self.n_stim) @@ -121,7 +138,7 @@ def __init__( parameters.add( name="pulse_intensity", - function=DingModelIntensityFrequency.set_impulse_intensity, + function=DingModelPulseIntensityFrequency.set_impulse_intensity, size=self.n_stim, scaling=VariableScaling("pulse_intensity", [1] * self.n_stim), ) @@ -132,9 +149,13 @@ def __init__( self.parameters = parameters self.parameters_init = parameters_init self.parameters_bounds = parameters_bounds - self.n_stim = self.n_stim self._declare_dynamics() - self.x_init, self.u_init, self.p_init, self.s_init = self.build_initial_guess_from_ocp(self) + ( + self.x_init, + self.u_init, + self.p_init, + self.s_init, + ) = self.build_initial_guess_from_ocp(self) self.ode_solver = self.ivp_parameters["ode_solver"] self.use_sx = self.ivp_parameters["use_sx"] @@ -146,8 +167,8 @@ def __init__( def _fill_fes_dict(self, fes_parameters): default_fes_dict = { "model": FesModel, - "n_stim": 1, - "pulse_duration": 0.0003, + "stim_time": None, + "pulse_width": 0.0003, "pulse_intensity": 50, "pulse_mode": "single", } @@ -163,9 +184,7 @@ def _fill_fes_dict(self, fes_parameters): def _fill_ivp_dict(self, ivp_parameters): default_ivp_dict = { - "n_shooting": None, "final_time": None, - "extend_last_phase_time": False, "ode_solver": OdeSolver.RK4(n_integration_steps=1), "use_sx": True, "n_threads": 1, @@ -188,45 +207,39 @@ def dictionaries_check(self): raise ValueError("ivp_parameters must be a dictionary") if not isinstance(self.fes_parameters["model"], FesModel): - raise ValueError("model must be a FesModel type") - - if not isinstance(self.fes_parameters["n_stim"], int): - raise ValueError("n_stim must be an int type") + raise TypeError("model must be a FesModel type") if isinstance( - self.fes_parameters["model"], DingModelPulseDurationFrequency | DingModelPulseDurationFrequencyWithFatigue + self.fes_parameters["model"], + DingModelPulseWidthFrequency | DingModelPulseWidthFrequencyWithFatigue, ): - pulse_duration_format = ( - isinstance(self.fes_parameters["pulse_duration"], int | float | list) - if not isinstance(self.fes_parameters["pulse_duration"], bool) + pulse_width_format = ( + isinstance(self.fes_parameters["pulse_width"], int | float | list) + if not isinstance(self.fes_parameters["pulse_width"], bool) else False ) - pulse_duration_format = ( - all([isinstance(pulse_duration, int) for pulse_duration in self.fes_parameters["pulse_duration"]]) - if pulse_duration_format == list - else pulse_duration_format + pulse_width_format = ( + all([isinstance(pulse_width, int) for pulse_width in self.fes_parameters["pulse_width"]]) + if pulse_width_format == list + else pulse_width_format ) - if pulse_duration_format is False: - raise TypeError("pulse_duration must be int, float or list type") + if pulse_width_format is False: + raise TypeError("pulse_width must be int, float or list type") - minimum_pulse_duration = self.fes_parameters["model"].pd0 - min_pulse_duration_check = ( - all( - [ - pulse_duration >= minimum_pulse_duration - for pulse_duration in self.fes_parameters["pulse_duration"] - ] - ) - if isinstance(self.fes_parameters["pulse_duration"], list) - else self.fes_parameters["pulse_duration"] >= minimum_pulse_duration + minimum_pulse_width = self.fes_parameters["model"].pd0 + min_pulse_width_check = ( + all([pulse_width >= minimum_pulse_width for pulse_width in self.fes_parameters["pulse_width"]]) + if isinstance(self.fes_parameters["pulse_width"], list) + else self.fes_parameters["pulse_width"] >= minimum_pulse_width ) - if min_pulse_duration_check is False: - raise ValueError("Pulse duration must be greater than minimum pulse duration") + if min_pulse_width_check is False: + raise ValueError("pulse width must be greater than minimum pulse width") if isinstance( - self.fes_parameters["model"], DingModelIntensityFrequency | DingModelIntensityFrequencyWithFatigue + self.fes_parameters["model"], + DingModelPulseIntensityFrequency | DingModelPulseIntensityFrequencyWithFatigue, ): pulse_intensity_format = ( isinstance(self.fes_parameters["pulse_intensity"], int | float | list) @@ -245,8 +258,8 @@ def dictionaries_check(self): minimum_pulse_intensity = ( all( [ - pulse_duration >= self.fes_parameters["model"].min_pulse_intensity() - for pulse_duration in self.fes_parameters["pulse_intensity"] + pulse_width >= self.fes_parameters["model"].min_pulse_intensity() + for pulse_width in self.fes_parameters["pulse_intensity"] ] ) if isinstance(self.fes_parameters["pulse_intensity"], list) @@ -259,17 +272,12 @@ def dictionaries_check(self): if not isinstance(self.fes_parameters["pulse_mode"], str): raise ValueError("pulse_mode must be a string type") - if not isinstance(self.ivp_parameters["n_shooting"], int | list | None): - raise ValueError("n_shooting must be an int or a list type") - if not isinstance(self.ivp_parameters["final_time"], int | float): raise ValueError("final_time must be an int or float type") - if not isinstance(self.ivp_parameters["extend_last_phase_time"], int | float | None): - raise ValueError("extend_last_phase_time must be an int or float type") - if not isinstance( - self.ivp_parameters["ode_solver"], (OdeSolver.RK1, OdeSolver.RK2, OdeSolver.RK4, OdeSolver.COLLOCATION) + self.ivp_parameters["ode_solver"], + (OdeSolver.RK1, OdeSolver.RK2, OdeSolver.RK4, OdeSolver.COLLOCATION), ): raise ValueError("ode_solver must be a OdeSolver type") @@ -281,60 +289,36 @@ def dictionaries_check(self): def _pulse_mode_settings(self): if self.pulse_mode == "single": - step = self.final_time / self.n_stim - self.final_time_phase = (step,) - for i in range(self.n_stim): - self.final_time_phase = self.final_time_phase + (step,) - self.dt.append(step / self.n_shooting[i]) - + pass elif self.pulse_mode == "doublet": doublet_step = 0.005 - step = np.round(self.final_time / (self.n_stim / 2) - doublet_step, 3) - index = 0 - for i in range(int(self.n_stim / 2)): - self.final_time_phase = (doublet_step,) if i == 0 else self.final_time_phase + (doublet_step,) - self.final_time_phase = self.final_time_phase + (step,) - self.dt.append(0.005 / self.n_shooting[index]) - index += 1 - self.dt.append(step / self.n_shooting[index]) - index += 1 + stim_time_doublet = [round(stim_time + doublet_step, 3) for stim_time in self.stim_time] + self.stim_time = self.stim_time + stim_time_doublet + self.stim_time.sort() + self.n_stim = len(self.stim_time) elif self.pulse_mode == "triplet": doublet_step = 0.005 - triplet_step = 0.005 - step = np.round(self.final_time / (self.n_stim / 3) - doublet_step - triplet_step, 3) - index = 0 - for i in range(int(self.n_stim / 3)): - self.final_time_phase = (doublet_step,) if i == 0 else self.final_time_phase + (doublet_step,) - self.final_time_phase = self.final_time_phase + (triplet_step,) - self.final_time_phase = self.final_time_phase + (step,) - self.dt.append(0.005 / self.n_shooting[index]) - index += 1 - self.dt.append(0.005 / self.n_shooting[index]) - index += 1 - self.dt.append(step / self.n_shooting[index]) - index += 1 + triplet_step = 0.01 + stim_time_doublet = [round(stim_time + doublet_step, 3) for stim_time in self.stim_time] + stim_time_triplet = [round(stim_time + triplet_step, 3) for stim_time in self.stim_time] + self.stim_time = self.stim_time + stim_time_doublet + stim_time_triplet + self.stim_time.sort() + self.n_stim = len(self.stim_time) else: raise ValueError("Pulse mode not yet implemented") - self.dt = np.array(self.dt) - if self.extend_last_phase_time: - self.final_time_phase = self.final_time_phase[:-1] + ( - self.final_time_phase[-1] + self.extend_last_phase_time, - ) - self.n_shooting[-1] = int((self.extend_last_phase_time / step) * self.n_shooting[-1]) + self.n_shooting[-1] - self.dt[-1] = self.final_time_phase[-1] / self.n_shooting[-1] - def _prepare_fake_ocp(self): """This function creates the initial value problem by hacking Bioptim's OptimalControlProgram. - It is not the normal use of bioptim, but it enables a simplified ivp construction.""" + It is not the normal use of bioptim, but it enables a simplified ivp construction. + """ return OptimalControlProgram( - bio_model=self.models, + bio_model=[self.model], dynamics=self.dynamics, n_shooting=self.n_shooting, - phase_time=self.final_time_phase, + phase_time=self.final_time, ode_solver=self.ode_solver, control_type=ControlType.CONSTANT, use_sx=self.use_sx, @@ -366,15 +350,14 @@ def integrate( def _declare_dynamics(self): self.dynamics = DynamicsList() - for i in range(self.n_stim): - self.dynamics.add( - self.models[i].declare_ding_variables, - dynamic_function=self.models[i].dynamics, - expand_dynamics=True, - expand_continuity=False, - phase=i, - phase_dynamics=PhaseDynamics.SHARED_DURING_THE_PHASE, - ) + self.dynamics.add( + self.model.declare_ding_variables, + dynamic_function=self.model.dynamics, + expand_dynamics=True, + expand_continuity=False, + phase=0, + phase_dynamics=PhaseDynamics.SHARED_DURING_THE_PHASE, + ) def build_initial_guess_from_ocp(self, ocp): """ @@ -385,9 +368,8 @@ def build_initial_guess_from_ocp(self, ocp): p = InitialGuessList() s = InitialGuessList() - for i in range(self.n_stim): - for j in range(len(self.model.name_dof)): - x.add(ocp.model.name_dof[j], ocp.model.standard_rest_values()[j], phase=i) + for j in range(len(self.model.name_dof)): + x.add(ocp.model.name_dof[j], ocp.model.standard_rest_values()[j], phase=0) if len(ocp.parameters) != 0: for key in ocp.parameters.keys(): p.add(key=key, initial_guess=ocp.parameters_init[key]) @@ -409,10 +391,10 @@ def from_frequency_and_final_time( ---------- fes_parameters: dict The parameters for the fes configuration including : - model, pulse_duration, pulse_intensity, pulse_mode, frequency, round_down + model, pulse_width, pulse_intensity, pulse_mode, frequency, round_down ivp_parameters: dict The parameters for the ivp problem including : - n_shooting, final_time, extend_last_phase_time, ode_solver, use_sx, n_threads + final_time, ode_solver, use_sx, n_threads """ frequency = fes_parameters["frequency"] @@ -434,6 +416,9 @@ def from_frequency_and_final_time( "The number of stimulation needs to be integer within the final time t, set round down " "to True or set final_time * frequency to make the result an integer." ) + fes_parameters["stim_time"] = list( + np.round([i * 1 / fes_parameters["frequency"] for i in range(fes_parameters["n_stim"])], 3) + ) return cls( fes_parameters, ivp_parameters, @@ -453,10 +438,10 @@ def from_frequency_and_n_stim( ---------- fes_parameters: dict The parameters for the fes configuration including : - model, n_stim, pulse_duration, pulse_intensity, pulse_mode + model, n_stim, pulse_width, pulse_intensity, pulse_mode ivp_parameters: dict The parameters for the ivp problem including : - n_shooting, extend_last_phase_time, ode_solver, use_sx, n_threads + final_time, ode_solver, use_sx, n_threads """ n_stim = fes_parameters["n_stim"] @@ -467,6 +452,11 @@ def from_frequency_and_n_stim( raise ValueError("Frequency must be an int") ivp_parameters["final_time"] = n_stim / frequency + + fes_parameters["stim_time"] = list( + np.round([i * 1 / fes_parameters["frequency"] for i in range(fes_parameters["n_stim"])], 3) + ) + return cls( fes_parameters, ivp_parameters, diff --git a/cocofest/models/ding2003.py b/cocofest/models/ding2003.py index fd69443..a09226f 100644 --- a/cocofest/models/ding2003.py +++ b/cocofest/models/ding2003.py @@ -1,7 +1,7 @@ from typing import Callable import numpy as np -from casadi import MX, exp, vertcat +from casadi import MX, exp, vertcat, if_else from bioptim import ( ConfigureProblem, @@ -10,7 +10,7 @@ OptimalControlProgram, ) -from .state_configue import StateConfigure +from .state_configure import StateConfigure from .fes_model import FesModel @@ -34,22 +34,34 @@ def __init__( model_name: str = "ding2003", muscle_name: str = None, sum_stim_truncation: int = None, + is_approximated: bool = False, ): super().__init__() self._model_name = model_name self._muscle_name = muscle_name self._sum_stim_truncation = sum_stim_truncation self._with_fatigue = False + self.is_approximated = is_approximated self.pulse_apparition_time = None + self.stim_prev = [] + + # --- Default values --- # + TAUC_DEFAULT = 0.020 # Value from Ding's experimentation [1] (s) + R0_KM_RELATIONSHIP_DEFAULT = 1.04 # (unitless) + A_REST_DEFAULT = 3009 # Value from Ding's experimentation [1] (N.s-1) + TAU1_REST_DEFAULT = 0.050957 # Value from Ding's experimentation [1] (s) + TAU2_DEFAULT = 0.060 # Close value from Ding's experimentation [2] (s) + KM_REST_DEFAULT = 0.103 # Value from Ding's experimentation [1] (unitless) + # ---- Custom values for the example ---- # - self.tauc = 0.020 # Value from Ding's experimentation [1] (s) - self.r0_km_relationship = 1.04 # (unitless) + self.tauc = TAUC_DEFAULT # Value from Ding's experimentation [1] (s) + self.r0_km_relationship = R0_KM_RELATIONSHIP_DEFAULT # (unitless) # ---- Different values for each person ---- # # ---- Force models ---- # - self.a_rest = 3009 # Value from Ding's experimentation [1] (N.s-1) - self.tau1_rest = 0.050957 # Value from Ding's experimentation [1] (s) - self.tau2 = 0.060 # Close value from Ding's experimentation [2] (s) - self.km_rest = 0.103 # Value from Ding's experimentation [1] (unitless) + self.a_rest = A_REST_DEFAULT + self.tau1_rest = TAU1_REST_DEFAULT + self.tau2 = TAU2_DEFAULT + self.km_rest = KM_REST_DEFAULT def set_a_rest(self, model, a_rest: MX | float): # models is required for bioptim compatibility @@ -111,7 +123,29 @@ def with_fatigue(self): @property def identifiable_parameters(self): - return {"a_rest": self.a_rest, "tau1_rest": self.tau1_rest, "km_rest": self.km_rest, "tau2": self.tau2} + return { + "a_rest": self.a_rest, + "tau1_rest": self.tau1_rest, + "km_rest": self.km_rest, + "tau2": self.tau2, + } + + @property + def km_name(self) -> str: + muscle_name = "_" + self.muscle_name if self.muscle_name else "" + return "Km" + muscle_name + + @property + def cn_sum_name(self): + muscle_name = "_" + self.muscle_name if self.muscle_name else "" + return "Cn_sum" + muscle_name + + def get_r0(self, km: MX | float) -> MX | float: + return km + self.r0_km_relationship + + @staticmethod + def get_lambda_i(nb_stim: int, pulse_intensity: MX | float) -> list[MX | float]: + return [1 for _ in range(nb_stim)] # ---- Model's dynamics ---- # def system_dynamics( @@ -120,6 +154,7 @@ def system_dynamics( f: MX, t: MX = None, t_stim_prev: list[MX] | list[float] = None, + cn_sum: MX | float = None, force_length_relationship: MX | float = 1, force_velocity_relationship: MX | float = 1, ) -> MX: @@ -133,9 +168,11 @@ def system_dynamics( f: MX The value of the force (N) t: MX - The current time at which the dynamics is evaluated (ms) - t_stim_prev: list[MX] - The time list of the previous stimulations (ms) + The current time at which the dynamics is evaluated (s) + t_stim_prev: list[MX] | list[float] + The time list of the previous stimulations (s) + cn_sum: MX | float + The sum of the ca_troponin_complex (unitless) force_length_relationship: MX | float The force length relationship value (unitless) force_velocity_relationship: MX | float @@ -145,8 +182,7 @@ def system_dynamics( ------- The value of the derivative of each state dx/dt at the current time t """ - r0 = self.km_rest + self.r0_km_relationship # Simplification - cn_dot = self.cn_dot_fun(cn, r0, t, t_stim_prev=t_stim_prev) # Equation n°1 + cn_dot = self.calculate_cn_dot(cn, cn_sum, t, t_stim_prev) f_dot = self.f_dot_fun( cn, f, @@ -188,7 +224,7 @@ def ri_fun(self, r0: MX | float, time_between_stim: MX) -> MX | float: """ return 1 + (r0 - 1) * exp(-time_between_stim / self.tauc) # Part of Eq n°1 - def cn_sum_fun(self, r0: MX | float, t: MX, t_stim_prev: list[MX]) -> MX | float: + def cn_sum_fun(self, r0: MX | float, t: MX, t_stim_prev: list[MX], lambda_i: list[MX]) -> MX | float: """ Parameters ---------- @@ -204,41 +240,37 @@ def cn_sum_fun(self, r0: MX | float, t: MX, t_stim_prev: list[MX]) -> MX | float A part of the n°1 equation """ sum_multiplier = 0 - enough_stim_to_truncate = self._sum_stim_truncation and len(t_stim_prev) > self._sum_stim_truncation - if enough_stim_to_truncate: - t_stim_prev = t_stim_prev[-self._sum_stim_truncation - 1 :] - if len(t_stim_prev) == 1: - ri = 1 - exp_time = self.exp_time_fun(t, t_stim_prev[0]) # Part of Eq n°1 - sum_multiplier += ri * exp_time # Part of Eq n°1 - else: - for i in range(1, len(t_stim_prev)): - previous_phase_time = t_stim_prev[i] - t_stim_prev[i - 1] - ri = self.ri_fun(r0, previous_phase_time) # Part of Eq n°1 - exp_time = self.exp_time_fun(t, t_stim_prev[i]) # Part of Eq n°1 - sum_multiplier += ri * exp_time # Part of Eq n°1 + + for i in range(len(t_stim_prev)): + previous_phase_time = t_stim_prev[i] - t_stim_prev[i - 1] + ri = 1 if i == 0 else self.ri_fun(r0, previous_phase_time) # Part of Eq n°1 + exp_time = self.exp_time_fun(t, t_stim_prev[i]) # Part of Eq n°1 + coefficient = 1 if self.is_approximated else if_else(t_stim_prev[i] <= t, 1, 0) + sum_multiplier += ri * exp_time * lambda_i[i] * coefficient return sum_multiplier - def cn_dot_fun(self, cn: MX, r0: MX | float, t: MX, t_stim_prev: list[MX]) -> MX | float: + def cn_dot_fun(self, cn: MX, cn_sum: MX) -> MX | float: """ Parameters ---------- cn: MX The previous step value of ca_troponin_complex (unitless) - r0: MX | float - Mathematical term characterizing the magnitude of enhancement in CN from the following stimuli (unitless) - t: MX - The current time at which the dynamics is evaluated (ms) - t_stim_prev: list[MX] - The time list of the previous stimulations (ms) Returns ------- The value of the derivative ca_troponin_complex (unitless) """ - sum_multiplier = self.cn_sum_fun(r0, t, t_stim_prev=t_stim_prev) # Part of Eq n°1 - return (1 / self.tauc) * sum_multiplier - (cn / self.tauc) # Equation n°1 + return (1 / self.tauc) * cn_sum - (cn / self.tauc) # Equation n°1 + + def calculate_cn_dot(self, cn, cn_sum, t, t_stim_prev, pulse_intensity=1): + if self.is_approximated: + return self.cn_dot_fun(cn=cn, cn_sum=cn_sum) + else: + cn_sum = self.cn_sum_fun( + self.get_r0(self.km_rest), t, t_stim_prev, self.get_lambda_i(len(t_stim_prev), pulse_intensity) + ) + return self.cn_dot_fun(cn, cn_sum) def f_dot_fun( self, @@ -287,7 +319,6 @@ def dynamics( algebraic_states: MX, numerical_timeseries: MX, nlp: NonLinearProgram, - stim_prev: list[MX] = None, fes_model=None, force_length_relationship: MX | float = 1, force_velocity_relationship: MX | float = 1, @@ -311,8 +342,6 @@ def dynamics( The numerical timeseries of the system nlp: NonLinearProgram A reference to the phase - stim_prev: list[MX] - The time list of the previous stimulations (s) fes_model: DingModelFrequency The current phase fes model force_length_relationship: MX | float @@ -323,19 +352,15 @@ def dynamics( ------- The derivative of the states in the tuple[MX] format """ + model = fes_model if fes_model else nlp.model + dxdt_fun = model.system_dynamics + stim_apparition = None + cn_sum = None - dxdt_fun = fes_model.system_dynamics if fes_model else nlp.model.system_dynamics - stim_apparition = ( - ( - fes_model.get_stim_prev(nlp=nlp, parameters=parameters, idx=nlp.phase_idx) - if fes_model - else nlp.model.get_stim_prev(nlp=nlp, parameters=parameters, idx=nlp.phase_idx) - ) - if stim_prev is None - else stim_prev - ) # Get the previous stimulation apparition time from the parameters - # if not provided from stim_prev, this way of getting the list is not optimal, but it is the only way to get it. - # Otherwise, it will create issues with free variables or wrong mx or sx type while calculating the dynamics + if model.is_approximated: + cn_sum = controls[0] + else: + stim_apparition = model.get_stim(nlp=nlp, parameters=parameters) return DynamicsEvaluation( dxdt=dxdt_fun( @@ -343,13 +368,17 @@ def dynamics( f=states[1], t=time, t_stim_prev=stim_apparition, + cn_sum=cn_sum, force_length_relationship=force_length_relationship, force_velocity_relationship=force_velocity_relationship, ), ) def declare_ding_variables( - self, ocp: OptimalControlProgram, nlp: NonLinearProgram, numerical_data_timeseries: dict[str, np.ndarray] = None + self, + ocp: OptimalControlProgram, + nlp: NonLinearProgram, + numerical_data_timeseries: dict[str, np.ndarray] = None, ): """ Tell the program which variables are states and controls. @@ -364,15 +393,12 @@ def declare_ding_variables( A list of values to pass to the dynamics at each node. Experimental external forces should be included here. """ StateConfigure().configure_all_fes_model_states(ocp, nlp, fes_model=self) - stim_prev = ( - self._build_t_stim_prev(ocp, nlp.phase_idx) - if "pulse_apparition_time" not in nlp.parameters.keys() - else None - ) - ConfigureProblem.configure_dynamics_function(ocp, nlp, dyn_func=self.dynamics, stim_prev=stim_prev) + if self.is_approximated: + StateConfigure().configure_cn_sum(ocp, nlp) + ConfigureProblem.configure_dynamics_function(ocp, nlp, dyn_func=self.dynamics) @staticmethod - def get_stim_prev(nlp: NonLinearProgram, parameters: MX, idx: int) -> list[float]: + def get_stim(nlp: NonLinearProgram, parameters: MX) -> list[float]: """ Get the nlp list of previous stimulation apparition time @@ -382,43 +408,16 @@ def get_stim_prev(nlp: NonLinearProgram, parameters: MX, idx: int) -> list[float The NonLinearProgram of the ocp of the current phase parameters: MX The parameters of the ocp - idx: int - The index of the current phase Returns ------- - The list of previous stimulation time + The list of stimulation time """ - t_stim_prev = [] + t_stim = [] for j in range(parameters.shape[0]): if "pulse_apparition_time" in nlp.parameters.cx[j].str(): - t_stim_prev.append(parameters[j]) - if len(t_stim_prev) > idx: - break - - return t_stim_prev - - @staticmethod - def _build_t_stim_prev(ocp: OptimalControlProgram, idx: int) -> list[float]: - """ - Builds a list of previous stimulation apparition time from known ocp phase time when the pulse_apparition_time - is not a declared optimized parameter - - Parameters - ---------- - ocp: OptimalControlProgram - The OptimalControlProgram of the problem - idx: int - The index of the current phase - - Returns - ------- - The list of previous stimulation time - """ - t_stim_prev = [0] - for i in range(idx): - t_stim_prev.append(t_stim_prev[-1] + ocp.phase_time[i]) - return t_stim_prev + t_stim.append(parameters[j]) + return t_stim def set_pulse_apparition_time(self, value: list[MX]): """ diff --git a/cocofest/models/ding2003_with_fatigue.py b/cocofest/models/ding2003_with_fatigue.py index 6a5d751..2217e46 100644 --- a/cocofest/models/ding2003_with_fatigue.py +++ b/cocofest/models/ding2003_with_fatigue.py @@ -11,7 +11,7 @@ ) from .ding2003 import DingModelFrequency -from .state_configue import StateConfigure +from .state_configure import StateConfigure class DingModelFrequencyWithFatigue(DingModelFrequency): @@ -32,14 +32,27 @@ def __init__( model_name: str = "ding2003_with_fatigue", muscle_name: str = None, sum_stim_truncation: int = None, + is_approximated: bool = False, ): - super().__init__(model_name=model_name, muscle_name=muscle_name, sum_stim_truncation=sum_stim_truncation) + super().__init__( + model_name=model_name, + muscle_name=muscle_name, + sum_stim_truncation=sum_stim_truncation, + ) self._with_fatigue = True + self.is_approximated = is_approximated + + # --- Default values --- # + ALPHA_A_DEFAULT = -4.0 * 10e-7 # Value from Ding's experimentation [1] (s^-2) + ALPHA_TAU1_DEFAULT = 2.1 * 10e-5 # Value from Ding's experimentation [1] (N^-1) + TAU_FAT_DEFAULT = 127 # Value from Ding's experimentation [1] (s) + ALPHA_KM_DEFAULT = 1.9 * 10e-8 # Value from Ding's experimentation [1] (s^-1.N^-1) + # ---- Fatigue models ---- # - self.alpha_a = -4.0 * 10e-7 # Value from Ding's experimentation [1] (s^-2) - self.alpha_tau1 = 2.1 * 10e-5 # Value from Ding's experimentation [1] (N^-1) - self.tau_fat = 127 # Value from Ding's experimentation [1] (s) - self.alpha_km = 1.9 * 10e-8 # Value from Ding's experimentation [1] (s^-1.N^-1) + self.alpha_a = ALPHA_A_DEFAULT + self.alpha_tau1 = ALPHA_TAU1_DEFAULT + self.tau_fat = TAU_FAT_DEFAULT + self.alpha_km = ALPHA_KM_DEFAULT def set_alpha_a(self, model, alpha_a: MX | float): self.alpha_a = alpha_a @@ -84,7 +97,13 @@ def serialize(self) -> tuple[Callable, dict]: @property def name_dof(self, with_muscle_name: bool = False) -> list[str]: muscle_name = "_" + self.muscle_name if self.muscle_name and with_muscle_name else "" - return ["Cn" + muscle_name, "F" + muscle_name, "A" + muscle_name, "Tau1" + muscle_name, "Km" + muscle_name] + return [ + "Cn" + muscle_name, + "F" + muscle_name, + "A" + muscle_name, + "Tau1" + muscle_name, + "Km" + muscle_name, + ] @property def nb_state(self) -> int: @@ -121,6 +140,7 @@ def system_dynamics( km: MX = None, t: MX = None, t_stim_prev: list[MX] | list[float] = None, + cn_sum: MX | float = None, force_length_relationship: MX | float = 1, force_velocity_relationship: MX | float = 1, ) -> MX: @@ -140,9 +160,11 @@ def system_dynamics( km: MX The value of the cross_bridges (unitless) t: MX - The current time at which the dynamics is evaluated (ms) - t_stim_prev: list[MX] - The time list of the previous stimulations (ms) + The current time at which the dynamics is evaluated (s) + t_stim_prev: list[MX] | list[float] + The time list of the previous stimulations (s) + cn_sum: MX | float + The sum of the ca_troponin_complex (unitless) force_length_relationship: MX | float The force length relationship value (unitless) force_velocity_relationship: MX | float @@ -152,8 +174,7 @@ def system_dynamics( ------- The value of the derivative of each state dx/dt at the current time t """ - r0 = km + self.r0_km_relationship # Simplification - cn_dot = self.cn_dot_fun(cn, r0, t, t_stim_prev=t_stim_prev) # Equation n°1 + cn_dot = self.calculate_cn_dot(cn, cn_sum, t, t_stim_prev) f_dot = self.f_dot_fun( cn, f, @@ -222,7 +243,6 @@ def dynamics( algebraic_states: MX, numerical_timeseries: MX, nlp: NonLinearProgram, - stim_prev: list[float] = None, fes_model=None, force_length_relationship: MX | float = 1, force_velocity_relationship: MX | float = 1, @@ -246,8 +266,6 @@ def dynamics( The numerical timeseries of the system nlp: NonLinearProgram A reference to the phase - stim_prev: list[float] - The previous stimulation time fes_model: DingModelFrequencyWithFatigue The current phase fes model force_length_relationship: MX | float @@ -258,19 +276,15 @@ def dynamics( ------- The derivative of the states in the tuple[MX] format """ + model = fes_model if fes_model else nlp.model + dxdt_fun = model.system_dynamics + stim_apparition = None + cn_sum = None - dxdt_fun = fes_model.system_dynamics if fes_model else nlp.model.system_dynamics - stim_apparition = ( - ( - fes_model.get_stim_prev(nlp=nlp, parameters=parameters, idx=nlp.phase_idx) - if fes_model - else nlp.model.get_stim_prev(nlp=nlp, parameters=parameters, idx=nlp.phase_idx) - ) - if stim_prev is None - else stim_prev - ) # Get the previous stimulation apparition time from the parameters - # if not provided from stim_prev, this way of getting the list is not optimal, but it is the only way to get it. - # Otherwise, it will create issues with free variables or wrong mx or sx type while calculating the dynamics + if model.is_approximated: + cn_sum = controls[0] + else: + stim_apparition = model.get_stim(nlp=nlp, parameters=parameters) return DynamicsEvaluation( dxdt=dxdt_fun( @@ -279,6 +293,7 @@ def dynamics( a=states[2], tau1=states[3], km=states[4], + cn_sum=cn_sum, t=time, t_stim_prev=stim_apparition, force_length_relationship=force_length_relationship, @@ -288,7 +303,10 @@ def dynamics( ) def declare_ding_variables( - self, ocp: OptimalControlProgram, nlp: NonLinearProgram, numerical_data_timeseries: dict[str, np.ndarray] = None + self, + ocp: OptimalControlProgram, + nlp: NonLinearProgram, + numerical_data_timeseries: dict[str, np.ndarray] = None, ): """ Tell the program which variables are states and controls. @@ -303,9 +321,6 @@ def declare_ding_variables( A list of values to pass to the dynamics at each node. Experimental external forces should be included here. """ StateConfigure().configure_all_fes_model_states(ocp, nlp, fes_model=self) - stim_prev = ( - self._build_t_stim_prev(ocp, nlp.phase_idx) - if "pulse_apparition_time" not in nlp.parameters.keys() - else None - ) - ConfigureProblem.configure_dynamics_function(ocp, nlp, dyn_func=self.dynamics, stim_prev=stim_prev) + if self.is_approximated: + StateConfigure().configure_cn_sum(ocp, nlp) + ConfigureProblem.configure_dynamics_function(ocp, nlp, dyn_func=self.dynamics) diff --git a/cocofest/models/ding2007.py b/cocofest/models/ding2007.py index 4e06a2e..4adaf81 100644 --- a/cocofest/models/ding2007.py +++ b/cocofest/models/ding2007.py @@ -1,7 +1,7 @@ from typing import Callable import numpy as np -from casadi import MX, vertcat, exp +from casadi import MX, vertcat, exp, if_else from bioptim import ( ConfigureProblem, @@ -11,16 +11,16 @@ ParameterList, ) from .ding2003 import DingModelFrequency -from .state_configue import StateConfigure +from .state_configure import StateConfigure -class DingModelPulseDurationFrequency(DingModelFrequency): +class DingModelPulseWidthFrequency(DingModelFrequency): """ This is a custom models that inherits from bioptim. CustomModel. As CustomModel is an abstract class, some methods are mandatory and must be implemented. Such as serialize, name_dof, nb_state. - This is the Ding 2007 model using the stimulation frequency and pulse duration in input. + This is the Ding 2007 model using the stimulation frequency and pulse width in input. Ding, J., Chou, L. W., Kesar, T. M., Lee, S. C., Johnston, T. E., Wexler, A. S., & Binder‐Macleod, S. A. (2007). Mathematical model that predicts the force–intensity and force–frequency relationships after spinal cord injuries. @@ -32,21 +32,47 @@ def __init__( model_name: str = "ding_2007", muscle_name: str = None, sum_stim_truncation: int = None, + is_approximated: bool = False, + tauc: float = None, + a_rest: float = None, + tau1_rest: float = None, + km_rest: float = None, + tau2: float = None, + pd0: float = None, + pdt: float = None, + a_scale: float = None, + alpha_a: float = None, + alpha_tau1: float = None, + alpha_km: float = None, + tau_fat: float = None, ): - super(DingModelPulseDurationFrequency, self).__init__( - model_name=model_name, muscle_name=muscle_name, sum_stim_truncation=sum_stim_truncation + super(DingModelPulseWidthFrequency, self).__init__( + model_name=model_name, + muscle_name=muscle_name, + sum_stim_truncation=sum_stim_truncation, + is_approximated=is_approximated, ) self._with_fatigue = False - self.impulse_time = None + self.pulse_width = None + + # --- Default values --- # + A_SCALE_DEFAULT = 4920 # Value from Ding's 2007 article (N/s) + PD0_DEFAULT = 0.000131405 # Value from Ding's 2007 article (s) + PDT_DEFAULT = 0.000194138 # Value from Ding's 2007 article (s) + TAU1_REST_DEFAULT = 0.060601 # Value from Ding's 2003 article (s) + TAU2_DEFAULT = 0.001 # Value from Ding's 2007 article (s) + KM_REST_DEFAULT = 0.137 # Value from Ding's 2007 article (unitless) + TAUC_DEFAULT = 0.011 # Value from Ding's 2007 article (s) + # ---- Custom values for the example ---- # # ---- Force models ---- # - self.a_scale = 4920 # Value from Ding's 2007 article (N/s) - self.pd0 = 0.000131405 # Value from Ding's 2007 article (s) - self.pdt = 0.000194138 # Value from Ding's 2007 article (s) - self.tau1_rest = 0.060601 # Value from Ding's 2003 article (s) - self.tau2 = 0.001 # Value from Ding's 2007 article (s) - self.km_rest = 0.137 # Value from Ding's 2007 article (unitless) - self.tauc = 0.011 # Value from Ding's 2007 article (s) + self.a_scale = A_SCALE_DEFAULT + self.pd0 = PD0_DEFAULT + self.pdt = PDT_DEFAULT + self.tau1_rest = TAU1_REST_DEFAULT + self.tau2 = TAU2_DEFAULT + self.km_rest = KM_REST_DEFAULT + self.tauc = TAUC_DEFAULT @property def identifiable_parameters(self): @@ -73,7 +99,7 @@ def serialize(self) -> tuple[Callable, dict]: # This is where you can serialize your models # This is useful if you want to save your models and load it later return ( - DingModelPulseDurationFrequency, + DingModelPulseWidthFrequency, { "tauc": self.tauc, "a_rest": self.a_rest, @@ -92,7 +118,9 @@ def system_dynamics( f: MX, t: MX = None, t_stim_prev: list[MX] | list[float] = None, - impulse_time: MX = None, + pulse_width: MX = None, + cn_sum: MX = None, + a_scale: MX = None, force_length_relationship: MX | float = 1, force_velocity_relationship: MX | float = 1, ) -> MX: @@ -106,11 +134,15 @@ def system_dynamics( f: MX The value of the force (N) t: MX - The current time at which the dynamics is evaluated (ms) - t_stim_prev: list[MX] - The time list of the previous stimulations (ms) - impulse_time: MX - The pulsation duration of the current stimulation (ms) + The current time at which the dynamics is evaluated (s) + t_stim_prev: list[MX] | list[float] + The time list of the previous stimulations (s) + pulse_width: MX + The pulsation duration of the current stimulation (s) + cn_sum: MX | float + The sum of the ca_troponin_complex (unitless) + a_scale: MX | float + The scaling factor of the current stimulation (unitless) force_length_relationship: MX | float The force length relationship value (unitless) force_velocity_relationship: MX | float @@ -120,13 +152,22 @@ def system_dynamics( ------- The value of the derivative of each state dx/dt at the current time t """ - r0 = self.km_rest + self.r0_km_relationship # Simplification - cn_dot = self.cn_dot_fun(cn, r0, t, t_stim_prev=t_stim_prev) # Equation n°1 from Ding's 2003 article - a = self.a_calculation(a_scale=self.a_scale, impulse_time=impulse_time) # Equation n°3 from Ding's 2007 article + cn_dot = self.calculate_cn_dot(cn, cn_sum, t, t_stim_prev) + a_scale = ( + a_scale + if self.is_approximated + else self.a_calculation( + a_scale=self.a_scale, + pulse_width=pulse_width, + t=t, + t_stim_prev=t_stim_prev, + ) + ) + f_dot = self.f_dot_fun( cn, f, - a, + a_scale, self.tau1_rest, self.km_rest, force_length_relationship=force_length_relationship, @@ -134,36 +175,81 @@ def system_dynamics( ) # Equation n°2 from Ding's 2003 article return vertcat(cn_dot, f_dot) - def a_calculation(self, a_scale: float | MX, impulse_time: MX) -> MX: + def a_calculation( + self, + a_scale: float | MX, + pulse_width: MX, + t=None, + t_stim_prev: list[float] | list[MX] = None, + ) -> MX: + """ + Parameters + ---------- + a_scale: float | MX + The scaling factor of the current stimulation (unitless) + pulse_width: MX + The pulsation duration of the current stimulation (s) + t: MX + The current time at which the dynamics is evaluated (s) + t_stim_prev: list[float] | list[MX] + The time list of the previous stimulations (s) + Returns + ------- + The value of scaling factor (unitless) + """ + if self.is_approximated: + return a_scale * (1 - exp(-(pulse_width - self.pd0) / self.pdt)) + else: + pulse_width_list = pulse_width + for i in range(len(t_stim_prev)): + if i == 0: + pulse_width = pulse_width_list[0] + else: + coefficient = if_else(t_stim_prev[i] <= t, 1, 0) + temp_pulse_width = pulse_width_list[i] * coefficient + pulse_width = if_else(temp_pulse_width != 0, temp_pulse_width, pulse_width) + return a_scale * (1 - exp(-(pulse_width - self.pd0) / self.pdt)) + + def a_calculation_identification( + self, + a_scale: float | MX, + pulse_width: MX, + pd0: float | MX, + pdt: float | MX, + ) -> MX: """ Parameters ---------- a_scale: float | MX The scaling factor of the current stimulation (unitless) - impulse_time: MX + pulse_width: MX The pulsation duration of the current stimulation (s) + pd0: float | MX + The pd0 value (s) + pdt: float | MX + The pdt value (s) Returns ------- The value of scaling factor (unitless) """ - return a_scale * (1 - exp(-(impulse_time - self.pd0) / self.pdt)) + return a_scale * (1 - exp(-(pulse_width - pd0) / pdt)) - def set_impulse_duration(self, value: list[MX]): + def set_impulse_width(self, value: list[MX]): """ - Sets the impulse time for each pulse (phases) according to the ocp parameter "impulse_time" + Sets the pulse width for each pulse (phases) according to the ocp parameter "pulse_width" Parameters ---------- value: list[MX] The pulsation duration list (s) """ - self.impulse_time = value + self.pulse_width = value @staticmethod - def get_pulse_duration_parameters(nlp, parameters: ParameterList, muscle_name: str = None) -> MX: + def get_pulse_width_parameters(nlp, parameters: ParameterList, muscle_name: str = None) -> list[MX]: """ - Get the nlp list of pulse_duration parameters + Get the nlp list of pulse_width parameters Parameters ---------- @@ -176,19 +262,17 @@ def get_pulse_duration_parameters(nlp, parameters: ParameterList, muscle_name: s Returns ------- - The list of list of pulse_duration parameters + The list of list of pulse_width parameters """ - pulse_duration_parameters = vertcat() + + pulse_width_parameters = [] for j in range(parameters.shape[0]): if muscle_name: - # if "pulse_duration_" + muscle_name in str(nlp.parameters.scaled.cx[j].name()): - if "pulse_duration_" + muscle_name in nlp.parameters.scaled.cx[j].str(): - pulse_duration_parameters = vertcat(pulse_duration_parameters, parameters[j]) - # elif "pulse_duration" in str(nlp.parameters.scaled.cx[j].name()): - elif "pulse_duration" in nlp.parameters.scaled.cx[j].str(): - pulse_duration_parameters = vertcat(pulse_duration_parameters, parameters[j]) - - return pulse_duration_parameters + if "pulse_width_" + muscle_name in nlp.parameters.scaled.cx[j].str(): + pulse_width_parameters.append(parameters[j]) + elif "pulse_width" in nlp.parameters.scaled.cx[j].str(): + pulse_width_parameters.append(parameters[j]) + return pulse_width_parameters @staticmethod def dynamics( @@ -199,7 +283,6 @@ def dynamics( algebraic_states: MX, numerical_timeseries: MX, nlp: NonLinearProgram, - stim_prev: list[float] = None, fes_model=None, force_length_relationship: MX | float = 1, force_velocity_relationship: MX | float = 1, @@ -223,9 +306,7 @@ def dynamics( The numerical timeseries of the system nlp: NonLinearProgram A reference to the phase - stim_prev: list[float] - The previous stimulation time - fes_model: DingModelPulseDurationFrequency + fes_model: DingModelPulseWidthFrequency The current phase fes model force_length_relationship: MX | float The force length relationship value (unitless) @@ -235,29 +316,22 @@ def dynamics( ------- The derivative of the states in the tuple[MX] format """ - pulse_duration_parameters = ( - nlp.model.get_pulse_duration_parameters(nlp, parameters) - if fes_model is None - else fes_model.get_pulse_duration_parameters(nlp, parameters, muscle_name=fes_model.muscle_name) - ) + model = fes_model if fes_model else nlp.model + dxdt_fun = model.system_dynamics - if pulse_duration_parameters.shape[0] == 1: # check if pulse duration is mapped - impulse_time = pulse_duration_parameters[0] + if model.is_approximated: + pulse_width = None + stim_apparition = None + cn_sum = controls[0] + a_scale = controls[1] else: - impulse_time = pulse_duration_parameters[nlp.phase_idx] - - dxdt_fun = fes_model.system_dynamics if fes_model else nlp.model.system_dynamics - stim_apparition = ( - ( - fes_model.get_stim_prev(nlp=nlp, parameters=parameters, idx=nlp.phase_idx) - if fes_model - else nlp.model.get_stim_prev(nlp=nlp, parameters=parameters, idx=nlp.phase_idx) - ) - if stim_prev is None - else stim_prev - ) # Get the previous stimulation apparition time from the parameters - # if not provided from stim_prev, this way of getting the list is not optimal, but it is the only way to get it. - # Otherwise, it will create issues with free variables or wrong mx or sx type while calculating the dynamics + pulse_width = model.get_pulse_width_parameters(nlp, parameters) + stim_apparition = model.get_stim(nlp=nlp, parameters=parameters) + + if len(pulse_width) == 1 and len(stim_apparition) != 1: + pulse_width = pulse_width * len(stim_apparition) + cn_sum = None + a_scale = None return DynamicsEvaluation( dxdt=dxdt_fun( @@ -265,7 +339,9 @@ def dynamics( f=states[1], t=time, t_stim_prev=stim_apparition, - impulse_time=impulse_time, + pulse_width=pulse_width, + cn_sum=cn_sum, + a_scale=a_scale, force_length_relationship=force_length_relationship, force_velocity_relationship=force_velocity_relationship, ), @@ -273,7 +349,10 @@ def dynamics( ) def declare_ding_variables( - self, ocp: OptimalControlProgram, nlp: NonLinearProgram, numerical_data_timeseries: dict[str, np.ndarray] = None + self, + ocp: OptimalControlProgram, + nlp: NonLinearProgram, + numerical_data_timeseries: dict[str, np.ndarray] = None, ): """ Tell the program which variables are states and controls. @@ -288,9 +367,7 @@ def declare_ding_variables( A list of values to pass to the dynamics at each node. Experimental external forces should be included here. """ StateConfigure().configure_all_fes_model_states(ocp, nlp, fes_model=self) - stim_prev = ( - self._build_t_stim_prev(ocp, nlp.phase_idx) - if "pulse_apparition_time" not in nlp.parameters.keys() - else None - ) - ConfigureProblem.configure_dynamics_function(ocp, nlp, dyn_func=self.dynamics, stim_prev=stim_prev) + if self.is_approximated: + StateConfigure().configure_cn_sum(ocp, nlp) + StateConfigure().configure_a_calculation(ocp, nlp) + ConfigureProblem.configure_dynamics_function(ocp, nlp, dyn_func=self.dynamics) diff --git a/cocofest/models/ding2007_with_fatigue.py b/cocofest/models/ding2007_with_fatigue.py index 4a288d5..8dd6e2c 100644 --- a/cocofest/models/ding2007_with_fatigue.py +++ b/cocofest/models/ding2007_with_fatigue.py @@ -9,17 +9,17 @@ NonLinearProgram, OptimalControlProgram, ) -from .ding2007 import DingModelPulseDurationFrequency -from .state_configue import StateConfigure +from .ding2007 import DingModelPulseWidthFrequency +from .state_configure import StateConfigure -class DingModelPulseDurationFrequencyWithFatigue(DingModelPulseDurationFrequency): +class DingModelPulseWidthFrequencyWithFatigue(DingModelPulseWidthFrequency): """ This is a custom models that inherits from bioptim. CustomModel. As CustomModel is an abstract class, some methods are mandatory and must be implemented. Such as serialize, name_dof, nb_state. - This is the Ding 2007 model using the stimulation frequency and pulse duration in input. + This is the Ding 2007 model using the stimulation frequency and pulse width in input. Ding, J., Chou, L. W., Kesar, T. M., Lee, S. C., Johnston, T. E., Wexler, A. S., & Binder‐Macleod, S. A. (2007). Mathematical model that predicts the force–intensity and force–frequency relationships after spinal cord injuries. @@ -31,23 +31,51 @@ def __init__( model_name: str = "ding_2007_with_fatigue", muscle_name: str = None, sum_stim_truncation: int = None, + is_approximated: bool = False, + tauc: float = None, + a_rest: float = None, + tau1_rest: float = None, + km_rest: float = None, + tau2: float = None, + pd0: float = None, + pdt: float = None, + a_scale: float = None, + alpha_a: float = None, + alpha_tau1: float = None, + alpha_km: float = None, + tau_fat: float = None, ): - super(DingModelPulseDurationFrequencyWithFatigue, self).__init__( - model_name=model_name, muscle_name=muscle_name, sum_stim_truncation=sum_stim_truncation + super(DingModelPulseWidthFrequencyWithFatigue, self).__init__( + model_name=model_name, + muscle_name=muscle_name, + sum_stim_truncation=sum_stim_truncation, + is_approximated=is_approximated, ) self._with_fatigue = True + # --- Default values --- # + ALPHA_A_DEFAULT = -4.0 * 10e-7 # Value from Ding's experimentation [1] (s^-2) + ALPHA_TAU1_DEFAULT = 2.1 * 10e-5 # Value from Ding's experimentation [1] (N^-1) + TAU_FAT_DEFAULT = 127 # Value from Ding's experimentation [1] (s) + ALPHA_KM_DEFAULT = 1.9 * 10e-8 # Value from Ding's experimentation [1] (s^-1.N^-1) + # ---- Fatigue models ---- # - self.alpha_a = -4.0 * 10e-7 # Value from Ding's experimentation [1] (s^-2) - self.alpha_tau1 = 2.1 * 10e-5 # Value from Ding's experimentation [1] (N^-1) - self.tau_fat = 127 # Value from Ding's experimentation [1] (s) - self.alpha_km = 1.9 * 10e-8 # Value from Ding's experimentation [1] (s^-1.N^-1) + self.alpha_a = ALPHA_A_DEFAULT + self.alpha_tau1 = ALPHA_TAU1_DEFAULT + self.tau_fat = TAU_FAT_DEFAULT + self.alpha_km = ALPHA_KM_DEFAULT # ---- Absolutely needed methods ---- # @property def name_dof(self, with_muscle_name: bool = False) -> list[str]: muscle_name = "_" + self.muscle_name if self.muscle_name and with_muscle_name else "" - return ["Cn" + muscle_name, "F" + muscle_name, "A" + muscle_name, "Tau1" + muscle_name, "Km" + muscle_name] + return [ + "Cn" + muscle_name, + "F" + muscle_name, + "A" + muscle_name, + "Tau1" + muscle_name, + "Km" + muscle_name, + ] @property def nb_state(self) -> int: @@ -80,7 +108,7 @@ def serialize(self) -> tuple[Callable, dict]: # This is where you can serialize your models # This is useful if you want to save your models and load it later return ( - DingModelPulseDurationFrequencyWithFatigue, + DingModelPulseWidthFrequencyWithFatigue, { "tauc": self.tauc, "a_rest": self.a_rest, @@ -106,7 +134,9 @@ def system_dynamics( km: MX = None, t: MX = None, t_stim_prev: list[MX] | list[float] = None, - impulse_time: MX = None, + pulse_width: MX = None, + cn_sum: MX = None, + a_scale: MX = None, force_length_relationship: MX | float = 1, force_velocity_relationship: MX | float = 1, ) -> MX: @@ -126,11 +156,16 @@ def system_dynamics( km: MX The value of the cross_bridges (unitless) t: MX - The current time at which the dynamics is evaluated (ms) - t_stim_prev: list[MX] - The time list of the previous stimulations (ms) - impulse_time: MX - The pulsation duration of the current stimulation (ms) + The current time at which the dynamics is evaluated (s) + t_stim_prev: list[MX] | list[float] + The time list of the previous stimulations (s) + pulse_width: MX + The time of the impulse (s) + cn_sum: MX | float + The sum of the ca_troponin_complex (unitless) + a_scale: MX | float + The scaling factor (unitless) + force_length_relationship: MX | float The force length relationship value (unitless) force_velocity_relationship: MX | float @@ -140,18 +175,28 @@ def system_dynamics( ------- The value of the derivative of each state dx/dt at the current time t """ - r0 = km + self.r0_km_relationship # Simplification - cn_dot = self.cn_dot_fun(cn, r0, t, t_stim_prev=t_stim_prev) # Equation n°1 from Ding's 2003 article - a_calculated = self.a_calculation(a_scale=a, impulse_time=impulse_time) # Equation n°3 from Ding's 2007 article + cn_dot = self.calculate_cn_dot(cn, cn_sum, t, t_stim_prev) + a_scale = ( + a_scale + if self.is_approximated + else self.a_calculation( + a_scale=self.a_scale, + pulse_width=pulse_width, + t=t, + t_stim_prev=t_stim_prev, + ) + ) + f_dot = self.f_dot_fun( cn, f, - a_calculated, + a_scale, tau1, km, force_length_relationship=force_length_relationship, force_velocity_relationship=force_velocity_relationship, ) # Equation n°2 from Ding's 2003 article + a_dot = self.a_dot_fun(a, f) tau1_dot = self.tau1_dot_fun(tau1, f) # Equation n°9 from Ding's 2003 article km_dot = self.km_dot_fun(km, f) # Equation n°11 from Ding's 2003 article @@ -211,7 +256,6 @@ def dynamics( algebraic_states: MX, numerical_timeseries: MX, nlp: NonLinearProgram, - stim_prev: list[float] = None, fes_model=None, force_length_relationship: MX | float = 1, force_velocity_relationship: MX | float = 1, @@ -235,9 +279,7 @@ def dynamics( The numerical timeseries of the system nlp: NonLinearProgram A reference to the phase - stim_prev: list[float] - The previous stimulation time - fes_model: DingModelPulseDurationFrequencyWithFatigue + fes_model: DingModelPulseWidthFrequencyWithFatigue The current phase fes model force_length_relationship: MX | float The force length relationship value (unitless) @@ -247,29 +289,28 @@ def dynamics( ------- The derivative of the states in the tuple[MX] format """ - pulse_duration_parameters = ( - nlp.model.get_pulse_duration_parameters(nlp, parameters) - if fes_model is None - else fes_model.get_pulse_duration_parameters(nlp, parameters, muscle_name=fes_model.muscle_name) - ) + model = fes_model if fes_model else nlp.model + dxdt_fun = model.system_dynamics - if pulse_duration_parameters.shape[0] == 1: # check if pulse duration is mapped - impulse_time = pulse_duration_parameters[0] + if model.is_approximated: + cn_sum = controls[0] + a_scale = controls[1] + stim_apparition = None + pulse_width = None else: - impulse_time = pulse_duration_parameters[nlp.phase_idx] - - dxdt_fun = fes_model.system_dynamics if fes_model else nlp.model.system_dynamics - stim_apparition = ( - ( - fes_model.get_stim_prev(nlp=nlp, parameters=parameters, idx=nlp.phase_idx) - if fes_model - else nlp.model.get_stim_prev(nlp=nlp, parameters=parameters, idx=nlp.phase_idx) + pulse_width = ( + model.get_pulse_width_parameters(nlp, parameters) + if fes_model is None + else fes_model.get_pulse_width_parameters(nlp, parameters, muscle_name=fes_model.muscle_name) ) - if stim_prev is None - else stim_prev - ) # Get the previous stimulation apparition time from the parameters - # if not provided from stim_prev, this way of getting the list is not optimal, but it is the only way to get it. - # Otherwise, it will create issues with free variables or wrong mx or sx type while calculating the dynamics + + stim_apparition = model.get_stim(nlp=nlp, parameters=parameters) + + if len(pulse_width) == 1 and len(stim_apparition) != 1: + pulse_width = pulse_width * len(stim_apparition) + + cn_sum = None + a_scale = None return DynamicsEvaluation( dxdt=dxdt_fun( @@ -280,7 +321,9 @@ def dynamics( km=states[4], t=time, t_stim_prev=stim_apparition, - impulse_time=impulse_time, + pulse_width=pulse_width, + cn_sum=cn_sum, + a_scale=a_scale, force_length_relationship=force_length_relationship, force_velocity_relationship=force_velocity_relationship, ), @@ -288,7 +331,10 @@ def dynamics( ) def declare_ding_variables( - self, ocp: OptimalControlProgram, nlp: NonLinearProgram, numerical_data_timeseries: dict[str, np.ndarray] = None + self, + ocp: OptimalControlProgram, + nlp: NonLinearProgram, + numerical_data_timeseries: dict[str, np.ndarray] = None, ): """ Tell the program which variables are states and controls. @@ -303,9 +349,7 @@ def declare_ding_variables( A list of values to pass to the dynamics at each node. Experimental external forces should be included here. """ StateConfigure().configure_all_fes_model_states(ocp, nlp, fes_model=self) - stim_prev = ( - self._build_t_stim_prev(ocp=ocp, idx=nlp.phase_idx) - if "pulse_apparition_time" not in nlp.parameters.keys() - else None - ) - ConfigureProblem.configure_dynamics_function(ocp, nlp, dyn_func=self.dynamics, stim_prev=stim_prev) + if self.is_approximated: + StateConfigure().configure_cn_sum(ocp, nlp) + StateConfigure().configure_a_calculation(ocp, nlp) + ConfigureProblem.configure_dynamics_function(ocp, nlp, dyn_func=self.dynamics) diff --git a/cocofest/models/dynamical_model.py b/cocofest/models/dynamical_model.py index 1c8ca68..d01f4b4 100644 --- a/cocofest/models/dynamical_model.py +++ b/cocofest/models/dynamical_model.py @@ -13,8 +13,11 @@ from ..models.fes_model import FesModel from ..models.ding2003 import DingModelFrequency -from .state_configue import StateConfigure -from .hill_coefficients import muscle_force_length_coefficient, muscle_force_velocity_coefficient +from .state_configure import StateConfigure +from .hill_coefficients import ( + muscle_force_length_coefficient, + muscle_force_velocity_coefficient, +) class FesMskModel(BiorbdModel): @@ -25,6 +28,7 @@ def __init__( muscles_model: list[FesModel] = None, activate_force_length_relationship: bool = False, activate_force_velocity_relationship: bool = False, + activate_residual_torque: bool = False, ): """ The custom model that will be used in the optimal control program for the FES-MSK models @@ -41,21 +45,38 @@ def __init__( If the force-length relationship should be activated activate_force_velocity_relationship: bool If the force-velocity relationship should be activated + activate_residual_torque: bool + If the residual torque should be activated """ super().__init__(biorbd_path) self._name = name self.bio_model = BiorbdModel(biorbd_path) + self._model_sanity( + muscles_model, + activate_force_length_relationship, + activate_force_velocity_relationship, + ) self.muscles_dynamics_model = muscles_model self.bio_stim_model = [self.bio_model] + self.muscles_dynamics_model self.activate_force_length_relationship = activate_force_length_relationship self.activate_force_velocity_relationship = activate_force_velocity_relationship + self.activate_residual_torque = activate_residual_torque # ---- Absolutely needed methods ---- # - def serialize(self, index: int = 0) -> tuple[Callable, dict]: - return self.muscles_dynamics_model[index].serialize() + def serialize(self) -> tuple[Callable, dict]: + return ( + FesMskModel, + { + "name": self._name, + "biorbd_path": self.bio_model.path, + "muscles_model": self.muscles_dynamics_model, + "activate_force_length_relationship": self.activate_force_length_relationship, + "activate_force_velocity_relationship": self.activate_force_velocity_relationship, + }, + ) # ---- Needed for the example ---- # @property @@ -88,7 +109,6 @@ def muscle_dynamic( nlp: NonLinearProgram, muscle_models: list[FesModel], state_name_list=None, - stim_prev: list[float] = None, ) -> DynamicsEvaluation: """ The custom dynamics function that provides the derivative of the states: dxdt = f(t, x, u, p, s) @@ -113,8 +133,6 @@ def muscle_dynamic( The list of the muscle models state_name_list: list[str] The states names list - stim_prev: list[float] - The previous stimulation values Returns ------- The derivative of the states in the tuple[MX | SX] format @@ -122,7 +140,7 @@ def muscle_dynamic( q = DynamicsFunctions.get(nlp.states["q"], states) qdot = DynamicsFunctions.get(nlp.states["qdot"], states) - tau = DynamicsFunctions.get(nlp.controls["tau"], controls) + tau = DynamicsFunctions.get(nlp.controls["tau"], controls) if "tau" in nlp.controls.keys() else 0 muscles_tau, dxdt_muscle_list = self.muscles_joint_torque( time, @@ -134,14 +152,14 @@ def muscle_dynamic( nlp, muscle_models, state_name_list, - stim_prev, q, qdot, ) # You can directly call biorbd function (as for ddq) or call bioptim accessor (as for dq) dq = DynamicsFunctions.compute_qdot(nlp, q, qdot) - ddq = nlp.model.forward_dynamics(q, qdot, muscles_tau + tau) + total_torque = muscles_tau + tau if self.activate_residual_torque else muscles_tau + ddq = nlp.model.forward_dynamics(q, qdot, total_torque) dxdt = vertcat(dxdt_muscle_list, dq, ddq) @@ -158,7 +176,6 @@ def muscles_joint_torque( nlp: NonLinearProgram, muscle_models: list[FesModel], state_name_list=None, - stim_prev: list[float] = None, q: MX | SX = None, qdot: MX | SX = None, ): @@ -179,15 +196,16 @@ def muscles_joint_torque( muscle_states_idxs = [ i for i in range(len(state_name_list)) if muscle_model.muscle_name in state_name_list[i] ] - muscle_states = vertcat() - for i in range(len(muscle_states_idxs)): - muscle_states = vertcat(muscle_states, states[muscle_states_idxs[i]]) + + muscle_states = vertcat(*[states[i] for i in muscle_states_idxs]) muscle_idx = bio_muscle_names_at_index.index(muscle_model.muscle_name) muscle_force_length_coeff = ( muscle_force_length_coefficient( - model=updatedModel, muscle=nlp.model.bio_model.model.muscle(muscle_idx), q=q + model=updatedModel, + muscle=nlp.model.bio_model.model.muscle(muscle_idx), + q=q, ) if nlp.model.activate_force_velocity_relationship else 1 @@ -195,7 +213,10 @@ def muscles_joint_torque( muscle_force_velocity_coeff = ( muscle_force_velocity_coefficient( - model=updatedModel, muscle=nlp.model.bio_model.model.muscle(muscle_idx), q=q, qdot=qdot + model=updatedModel, + muscle=nlp.model.bio_model.model.muscle(muscle_idx), + q=q, + qdot=qdot, ) if nlp.model.activate_force_velocity_relationship else 1 @@ -209,7 +230,6 @@ def muscles_joint_torque( algebraic_states, numerical_data_timeseries, nlp, - stim_prev=stim_prev, fes_model=muscle_model, force_length_relationship=muscle_force_length_coeff, force_velocity_relationship=muscle_force_velocity_coeff, @@ -219,7 +239,8 @@ def muscles_joint_torque( muscle_idx_list.append(muscle_idx) muscle_forces = vertcat( - muscle_forces, DynamicsFunctions.get(nlp.states["F_" + muscle_model.muscle_name], states) + muscle_forces, + DynamicsFunctions.get(nlp.states["F_" + muscle_model.muscle_name], states), ) muscle_moment_arm_matrix = updated_muscle_length_jacobian[ @@ -230,7 +251,10 @@ def muscles_joint_torque( return muscle_joint_torques, dxdt_muscle_list def declare_model_variables( - self, ocp: OptimalControlProgram, nlp: NonLinearProgram, numerical_data_timeseries: dict[str, np.ndarray] = None + self, + ocp: OptimalControlProgram, + nlp: NonLinearProgram, + numerical_data_timeseries: dict[str, np.ndarray] = None, ): """ Tell the program which variables are states and controls. @@ -250,17 +274,41 @@ def declare_model_variables( state_name_list.append("q") ConfigureProblem.configure_qdot(ocp, nlp, as_states=True, as_controls=False) state_name_list.append("qdot") - ConfigureProblem.configure_tau(ocp, nlp, as_states=False, as_controls=True) - stim_prev = ( - DingModelFrequency._build_t_stim_prev(ocp, nlp.phase_idx) - if "pulse_apparition_time" not in nlp.parameters.keys() - else None - ) + for muscle_model in self.muscles_dynamics_model: + if muscle_model.is_approximated: + StateConfigure().configure_cn_sum(ocp, nlp, muscle_name=str(muscle_model.muscle_name)) + StateConfigure().configure_a_calculation(ocp, nlp, muscle_name=str(muscle_model.muscle_name)) + if self.activate_residual_torque: + ConfigureProblem.configure_tau(ocp, nlp, as_states=False, as_controls=True) + ConfigureProblem.configure_dynamics_function( ocp, nlp, dyn_func=self.muscle_dynamic, muscle_models=self.muscles_dynamics_model, state_name_list=state_name_list, - stim_prev=stim_prev, ) + + @staticmethod + def _model_sanity( + muscles_model, + activate_force_length_relationship, + activate_force_velocity_relationship, + ): + if not isinstance(muscles_model, list): + for muscle_model in muscles_model: + if not isinstance(muscle_model, FesModel): + raise TypeError( + f"The current model type used is {type(muscles_model)}, it must be a FesModel type." + f"Current available models are: DingModelFrequency, DingModelFrequencyWithFatigue," + f"DingModelPulseWidthFrequency, DingModelPulseWidthFrequencyWithFatigue," + f"DingModelPulseIntensityFrequency, DingModelPulseIntensityFrequencyWithFatigue" + ) + + raise TypeError("The given muscles_model must be a list of FesModel") + + if not isinstance(activate_force_length_relationship, bool): + raise TypeError("The activate_force_length_relationship must be a boolean") + + if not isinstance(activate_force_velocity_relationship, bool): + raise TypeError("The activate_force_velocity_relationship must be a boolean") diff --git a/cocofest/models/fes_model.py b/cocofest/models/fes_model.py index c9da65b..b1ab740 100644 --- a/cocofest/models/fes_model.py +++ b/cocofest/models/fes_model.py @@ -110,8 +110,7 @@ def system_dynamics( self, cn: MX, f: MX, - t: MX, - t_stim_prev: list[MX] | list[float], + cn_sum: MX, force_length_relationship: MX | float, force_velocity_relationship: MX | float, ): @@ -141,7 +140,7 @@ def ri_fun(self, r0: MX | float, time_between_stim: MX): """ @abstractmethod - def cn_sum_fun(self, r0: MX | float, t: MX, t_stim_prev: list[MX]): + def cn_sum_fun(self, r0: MX | float, t: MX, t_stim_prev: list[MX], lambda_i: list[MX]): """ Returns ------- @@ -149,7 +148,7 @@ def cn_sum_fun(self, r0: MX | float, t: MX, t_stim_prev: list[MX]): """ @abstractmethod - def cn_dot_fun(self, cn: MX, r0: MX | float, t: MX, t_stim_prev: list[MX]): + def cn_dot_fun(self, cn: MX, cn_sum: MX): """ Returns @@ -185,7 +184,6 @@ def dynamics( algebraic_states: MX, numerical_data_timeseries: MX, nlp: NonLinearProgram, - stim_prev: list[float], fes_model, force_length_relationship: MX | float, force_velocity_relationship: MX | float, @@ -199,7 +197,10 @@ def dynamics( @abstractmethod def declare_ding_variables( - self, ocp: OptimalControlProgram, nlp: NonLinearProgram, numerical_data_timeseries: dict[str, np.ndarray] = None + self, + ocp: OptimalControlProgram, + nlp: NonLinearProgram, + numerical_data_timeseries: dict[str, np.ndarray] = None, ): """ @@ -208,26 +209,6 @@ def declare_ding_variables( """ - @staticmethod - @abstractmethod - def get_stim_prev(nlp: NonLinearProgram, parameters: MX, idx: int): - """ - - Returns - ------- - - """ - - @staticmethod - @abstractmethod - def _build_t_stim_prev(ocp: OptimalControlProgram, idx: int): - """ - - Returns - ------- - - """ - @abstractmethod def set_pulse_apparition_time(self, value: list[MX]): """ diff --git a/cocofest/models/hmed2018.py b/cocofest/models/hmed2018.py index d14433d..3fb1d55 100644 --- a/cocofest/models/hmed2018.py +++ b/cocofest/models/hmed2018.py @@ -11,10 +11,10 @@ ParameterList, ) from .ding2003 import DingModelFrequency -from .state_configue import StateConfigure +from .state_configure import StateConfigure -class DingModelIntensityFrequency(DingModelFrequency): +class DingModelPulseIntensityFrequency(DingModelFrequency): """ This is a custom models that inherits from bioptim. CustomModel. As CustomModel is an abstract class, some methods are mandatory and must be implemented. @@ -27,17 +27,34 @@ class DingModelIntensityFrequency(DingModelFrequency): Computers in Biology and Medicine, 101, 218-228. """ - def __init__(self, model_name: str = "hmed2018", muscle_name: str = None, sum_stim_truncation: int = None): - super(DingModelIntensityFrequency, self).__init__( - model_name=model_name, muscle_name=muscle_name, sum_stim_truncation=sum_stim_truncation + def __init__( + self, + model_name: str = "hmed2018", + muscle_name: str = None, + sum_stim_truncation: int = None, + is_approximated: bool = False, + ): + super(DingModelPulseIntensityFrequency, self).__init__( + model_name=model_name, + muscle_name=muscle_name, + sum_stim_truncation=sum_stim_truncation, + is_approximated=is_approximated, ) self._with_fatigue = False + self.stim_pulse_intensity_prev = [] + + # --- Default values ---# + AR_DEFAULT = 0.586 # (-) Translation of axis coordinates. + BS_DEFAULT = 0.026 # (-) Fiber muscle recruitment constant identification. + IS_DEFAULT = 63.1 # (mA) Muscle saturation intensity. + CR_DEFAULT = 0.833 # (-) Translation of axis coordinates. + # ---- Custom values for the example ---- # # ---- Force models ---- # - self.ar = 0.586 # (-) Translation of axis coordinates. - self.bs = 0.026 # (-) Fiber muscle recruitment constant identification. - self.Is = 63.1 # (mA) Muscle saturation intensity. - self.cr = 0.833 # (-) Translation of axis coordinates. + self.ar = AR_DEFAULT + self.bs = BS_DEFAULT + self.Is = IS_DEFAULT + self.cr = CR_DEFAULT self.impulse_intensity = None @property @@ -53,6 +70,11 @@ def identifiable_parameters(self): "cr": self.cr, } + @property + def pulse_intensity_name(self): + muscle_name = "_" + self.muscle_name if self.muscle_name else "" + return "pulse_intensity" + muscle_name + def set_ar(self, model, ar: MX | float): # models is required for bioptim compatibility self.ar = ar @@ -66,12 +88,15 @@ def set_Is(self, model, Is: MX | float): def set_cr(self, model, cr: MX | float): self.cr = cr + def get_lambda_i(self, nb_stim: int, pulse_intensity: MX | float) -> list[MX | float]: + return [self.lambda_i_calculation(pulse_intensity[i]) for i in range(nb_stim)] + # ---- Absolutely needed methods ---- # def serialize(self) -> tuple[Callable, dict]: # This is where you can serialize your models # This is useful if you want to save your models and load it later return ( - DingModelIntensityFrequency, + DingModelPulseIntensityFrequency, { "tauc": self.tauc, "a_rest": self.a_rest, @@ -91,7 +116,8 @@ def system_dynamics( f: MX, t: MX = None, t_stim_prev: list[MX] | list[float] = None, - intensity_stim: list[MX] | list[float] = None, + pulse_intensity: list[MX] | list[float] = None, + cn_sum: MX = None, force_length_relationship: MX | float = 1, force_velocity_relationship: MX | float = 1, ) -> MX: @@ -105,11 +131,13 @@ def system_dynamics( f: MX The value of the force (N) t: MX - The current time at which the dynamics is evaluated (ms) - t_stim_prev: list[MX] - The time list of the previous stimulations (ms) - intensity_stim: list[MX] + The current time at which the dynamics is evaluated (s) + t_stim_prev: list[MX] | list[float] + The time list of the previous stimulations (s) + pulse_intensity: list[MX] | list[float] The pulsation intensity of the current stimulation (mA) + cn_sum: MX | float + The sum of the ca_troponin_complex (unitless) force_length_relationship: MX | float The force length relationship value (unitless) force_velocity_relationship: MX | float @@ -119,8 +147,7 @@ def system_dynamics( ------- The value of the derivative of each state dx/dt at the current time t """ - r0 = self.km_rest + self.r0_km_relationship # Simplification - cn_dot = self.cn_dot_fun(cn, r0, t, t_stim_prev=t_stim_prev, intensity_stim=intensity_stim) # Equation n°1 + cn_dot = self.calculate_cn_dot(cn, cn_sum, t, t_stim_prev, pulse_intensity) f_dot = self.f_dot_fun( cn, f, @@ -132,76 +159,41 @@ def system_dynamics( ) # Equation n°2 return vertcat(cn_dot, f_dot) - def cn_dot_fun( - self, cn: MX, r0: MX | float, t: MX, t_stim_prev: list[MX], intensity_stim: list[MX] = None - ) -> MX | float: + def lambda_i_calculation(self, pulse_intensity: MX): """ Parameters ---------- - cn: MX - The previous step value of ca_troponin_complex (unitless) - r0: MX - Mathematical term characterizing the magnitude of enhancement in CN from the following stimuli (unitless) - t: MX - The current time at which the dynamics is evaluated (ms) - t_stim_prev: list[MX] - The time list of the previous stimulations (ms) - intensity_stim: list[MX] + pulse_intensity: MX The pulsation intensity of the current stimulation (mA) Returns ------- - The value of the derivative ca_troponin_complex (unitless) - """ - sum_multiplier = self.cn_sum_fun(r0, t, t_stim_prev=t_stim_prev, intensity_stim=intensity_stim) - - return (1 / self.tauc) * sum_multiplier - (cn / self.tauc) # Eq(1) - - def cn_sum_fun( - self, r0: MX | float, t: MX, t_stim_prev: list[MX] = None, intensity_stim: list[MX] = None - ) -> MX | float: + The lambda factor, part of the n°1 equation """ - Parameters - ---------- - r0: MX | float - Mathematical term characterizing the magnitude of enhancement in CN from the following stimuli (unitless) - t: MX - The current time at which the dynamics is evaluated (ms) - t_stim_prev: list[MX] - The time list of the previous stimulations (ms) - intensity_stim: list[MX] - The pulsation intensity of the current stimulation (mA) + lambda_i = self.ar * (tanh(self.bs * (pulse_intensity - self.Is)) + self.cr) # equation include intensity + return lambda_i - Returns - ------- - A part of the n°1 equation - """ - sum_multiplier = 0 - enough_stim_to_truncate = self._sum_stim_truncation and len(t_stim_prev) > self._sum_stim_truncation - if enough_stim_to_truncate: - t_stim_prev = t_stim_prev[-self._sum_stim_truncation :] - for i in range(len(t_stim_prev)): # Eq from [1] - if i == 0 and len(t_stim_prev) == 1: # Eq from Hmed et al. - ri = 1 - else: - previous_phase_time = t_stim_prev[i] - t_stim_prev[i - 1] - ri = self.ri_fun(r0, previous_phase_time) - exp_time = self.exp_time_fun(t, t_stim_prev[i]) - lambda_i = self.lambda_i_calculation(intensity_stim[i]) - sum_multiplier += lambda_i * ri * exp_time - return sum_multiplier - - def lambda_i_calculation(self, intensity_stim: MX): + @staticmethod + def lambda_i_calculation_identification( + pulse_intensity: MX, ar: MX | float, bs: MX | float, Is: MX | float, cr: MX | float + ): """ Parameters ---------- - intensity_stim: MX + pulse_intensity: MX The pulsation intensity of the current stimulation (mA) - + ar: MX | float + Translation of axis coordinates (-) + bs: MX | float + Fiber muscle recruitment constant identification. + Is: MX | float + Muscle saturation intensity (mA) + cr: MX | float + Translation of axis coordinates (-) Returns ------- The lambda factor, part of the n°1 equation """ - lambda_i = self.ar * (tanh(self.bs * (intensity_stim - self.Is)) + self.cr) # equation include intensity + lambda_i = ar * (tanh(bs * (pulse_intensity - Is)) + cr) # equation include intensity return lambda_i def set_impulse_intensity(self, value: MX): @@ -217,34 +209,6 @@ def set_impulse_intensity(self, value: MX): for i in range(value.shape[0]): self.impulse_intensity.append(value[i]) - @staticmethod - def get_intensity_parameters(nlp, parameters: ParameterList, muscle_name: str = None) -> MX: - """ - Get the nlp list of intensity parameters - - Parameters - ---------- - nlp: NonLinearProgram - A reference to the phase - parameters: ParameterList - The nlp list parameter - muscle_name: str - The muscle name - - Returns - ------- - The list of intensity parameters - """ - intensity_parameters = vertcat() - for j in range(parameters.shape[0]): - if muscle_name: - if "pulse_intensity_" + muscle_name in nlp.parameters.scaled.cx[j].str(): - intensity_parameters = vertcat(intensity_parameters, parameters[j]) - elif "pulse_intensity" in nlp.parameters.scaled.cx[j].str(): - intensity_parameters = vertcat(intensity_parameters, parameters[j]) - - return intensity_parameters - @staticmethod def dynamics( time: MX, @@ -254,7 +218,6 @@ def dynamics( algebraic_states: MX, numerical_timeseries: MX, nlp: NonLinearProgram, - stim_prev: list[float] = None, fes_model: NonLinearProgram = None, force_length_relationship: MX | float = 1, force_velocity_relationship: MX | float = 1, @@ -278,9 +241,7 @@ def dynamics( The numerical timeseries of the system nlp: NonLinearProgram A reference to the phase - stim_prev: list[float] - The previous stimulation values - fes_model: DingModelIntensityFrequency + fes_model: DingModelPulseIntensityFrequency The current phase fes model force_length_relationship: MX | float The force length relationship value (unitless) @@ -290,34 +251,20 @@ def dynamics( ------- The derivative of the states in the tuple[MX] format """ - intensity_stim_prev = ( - [] - ) # Every stimulation intensity before the current phase, i.e.: the intensity of each phase - intensity_parameters = ( - nlp.model.get_intensity_parameters(nlp, parameters) - if fes_model is None - else fes_model.get_intensity_parameters(nlp, parameters, muscle_name=fes_model.muscle_name) - ) + model = fes_model if fes_model else nlp.model + dxdt_fun = model.system_dynamics - if intensity_parameters.shape[0] == 1: # check if pulse duration is mapped - for i in range(nlp.phase_idx + 1): - intensity_stim_prev.append(intensity_parameters[0]) + if model.is_approximated: + cn_sum = controls[0] + stim_apparition = None + intensity_parameters = None else: - for i in range(nlp.phase_idx + 1): - intensity_stim_prev.append(intensity_parameters[i]) - - dxdt_fun = fes_model.system_dynamics if fes_model else nlp.model.system_dynamics - stim_apparition = ( - ( - fes_model.get_stim_prev(nlp=nlp, parameters=parameters, idx=nlp.phase_idx) - if fes_model - else nlp.model.get_stim_prev(nlp=nlp, parameters=parameters, idx=nlp.phase_idx) - ) - if stim_prev is None - else stim_prev - ) # Get the previous stimulation apparition time from the parameters - # if not provided from stim_prev, this way of getting the list is not optimal, but it is the only way to get it. - # Otherwise, it will create issues with free variables or wrong mx or sx type while calculating the dynamics + cn_sum = None + intensity_parameters = model.get_intensity_parameters(nlp, parameters) + stim_apparition = model.get_stim(nlp=nlp, parameters=parameters) + + if len(intensity_parameters) == 1 and len(stim_apparition) != 1: + intensity_parameters = intensity_parameters * len(stim_apparition) return DynamicsEvaluation( dxdt=dxdt_fun( @@ -325,7 +272,8 @@ def dynamics( f=states[1], t=time, t_stim_prev=stim_apparition, - intensity_stim=intensity_stim_prev, + pulse_intensity=intensity_parameters, + cn_sum=cn_sum, force_length_relationship=force_length_relationship, force_velocity_relationship=force_velocity_relationship, ), @@ -333,7 +281,10 @@ def dynamics( ) def declare_ding_variables( - self, ocp: OptimalControlProgram, nlp: NonLinearProgram, numerical_data_timeseries: dict[str, np.ndarray] = None + self, + ocp: OptimalControlProgram, + nlp: NonLinearProgram, + numerical_data_timeseries: dict[str, np.ndarray] = None, ): """ Tell the program which variables are states and controls. @@ -348,18 +299,43 @@ def declare_ding_variables( A list of values to pass to the dynamics at each node. Experimental external forces should be included here. """ StateConfigure().configure_all_fes_model_states(ocp, nlp, fes_model=self) - stim_prev = ( - self._build_t_stim_prev(ocp, nlp.phase_idx) - if "pulse_apparition_time" not in nlp.parameters.keys() - else None - ) - ConfigureProblem.configure_dynamics_function(ocp, nlp, dyn_func=self.dynamics, stim_prev=stim_prev) + if self.is_approximated: + StateConfigure().configure_cn_sum(ocp, nlp) + ConfigureProblem.configure_dynamics_function(ocp, nlp, dyn_func=self.dynamics) def min_pulse_intensity(self): """ Returns ------- The minimum pulse intensity threshold of the model - For lambda_i = ar * (tanh(bs * (intensity_stim - Is)) + cr) > 0 + For lambda_i = ar * (tanh(bs * (pulse_intensity - Is)) + cr) > 0 """ return (np.arctanh(-self.cr) / self.bs) + self.Is + + @staticmethod + def get_intensity_parameters(nlp, parameters: ParameterList, muscle_name: str = None) -> list[MX]: + """ + Get the nlp list of intensity parameters + + Parameters + ---------- + nlp: NonLinearProgram + A reference to the phase + parameters: ParameterList + The nlp list parameter + muscle_name: str + The muscle name + + Returns + ------- + The list of intensity parameters + """ + intensity_parameters = [] + for j in range(parameters.shape[0]): + if muscle_name: + if "pulse_intensity_" + muscle_name in nlp.parameters.scaled.cx[j].str(): + intensity_parameters.append(parameters[j]) + elif "pulse_intensity" in nlp.parameters.scaled.cx[j].str(): + intensity_parameters.append(parameters[j]) + + return intensity_parameters diff --git a/cocofest/models/hmed2018_with_fatigue.py b/cocofest/models/hmed2018_with_fatigue.py index ca2a013..321ee56 100644 --- a/cocofest/models/hmed2018_with_fatigue.py +++ b/cocofest/models/hmed2018_with_fatigue.py @@ -9,11 +9,11 @@ NonLinearProgram, OptimalControlProgram, ) -from .hmed2018 import DingModelIntensityFrequency -from .state_configue import StateConfigure +from .hmed2018 import DingModelPulseIntensityFrequency +from .state_configure import StateConfigure -class DingModelIntensityFrequencyWithFatigue(DingModelIntensityFrequency): +class DingModelPulseIntensityFrequencyWithFatigue(DingModelPulseIntensityFrequency): """ This is a custom models that inherits from bioptim. CustomModel. As CustomModel is an abstract class, some methods are mandatory and must be implemented. @@ -27,23 +27,43 @@ class DingModelIntensityFrequencyWithFatigue(DingModelIntensityFrequency): """ def __init__( - self, model_name: str = "hmed2018_with_fatigue", muscle_name: str = None, sum_stim_truncation: int = None + self, + model_name: str = "hmed2018_with_fatigue", + muscle_name: str = None, + sum_stim_truncation: int = None, + is_approximated: bool = False, ): - super(DingModelIntensityFrequencyWithFatigue, self).__init__( - model_name=model_name, muscle_name=muscle_name, sum_stim_truncation=sum_stim_truncation + super(DingModelPulseIntensityFrequencyWithFatigue, self).__init__( + model_name=model_name, + muscle_name=muscle_name, + sum_stim_truncation=sum_stim_truncation, + is_approximated=is_approximated, ) self._with_fatigue = True + + # --- Default values --- # + ALPHA_A_DEFAULT = -4.0 * 10e-7 # Value from Ding's experimentation [1] (s^-2) + ALPHA_TAU1_DEFAULT = 2.1 * 10e-5 # Value from Ding's experimentation [1] (N^-1) + TAU_FAT_DEFAULT = 127 # Value from Ding's experimentation [1] (s) + ALPHA_KM_DEFAULT = 1.9 * 10e-8 # Value from Ding's experimentation [1] (s^-1.N^-1) + # ---- Fatigue models ---- # - self.alpha_a = -4.0 * 10e-7 # Value from Ding's experimentation [1] (s^-2) - self.alpha_tau1 = 2.1 * 10e-5 # Value from Ding's experimentation [1] (N^-1) - self.tau_fat = 127 # Value from Ding's experimentation [1] (s) - self.alpha_km = 1.9 * 10e-8 # Value from Ding's experimentation [1] (s^-1.N^-1) + self.alpha_a = ALPHA_A_DEFAULT + self.alpha_tau1 = ALPHA_TAU1_DEFAULT + self.tau_fat = TAU_FAT_DEFAULT + self.alpha_km = ALPHA_KM_DEFAULT # ---- Absolutely needed methods ---- # @property def name_dof(self, with_muscle_name: bool = False) -> list[str]: muscle_name = "_" + self.muscle_name if self.muscle_name and with_muscle_name else "" - return ["Cn" + muscle_name, "F" + muscle_name, "A" + muscle_name, "Tau1" + muscle_name, "Km" + muscle_name] + return [ + "Cn" + muscle_name, + "F" + muscle_name, + "A" + muscle_name, + "Tau1" + muscle_name, + "Km" + muscle_name, + ] @property def nb_state(self) -> int: @@ -78,7 +98,7 @@ def serialize(self) -> tuple[Callable, dict]: # This is where you can serialize your models # This is useful if you want to save your models and load it later return ( - DingModelIntensityFrequencyWithFatigue, + DingModelPulseIntensityFrequencyWithFatigue, { "tauc": self.tauc, "a_rest": self.a_rest, @@ -105,7 +125,8 @@ def system_dynamics( km: MX = None, t: MX = None, t_stim_prev: list[MX] | list[float] = None, - intensity_stim: list[MX] | list[float] = None, + pulse_intensity: list[MX] | list[float] = None, + cn_sum: MX = None, force_length_relationship: float | MX = 1, force_velocity_relationship: float | MX = 1, ) -> MX: @@ -125,11 +146,13 @@ def system_dynamics( km: MX The value of the cross_bridges (unitless) t: MX - The current time at which the dynamics is evaluated (ms) - t_stim_prev: list[MX] - The time list of the previous stimulations (ms) - intensity_stim: list[MX] - The pulsation intensity of the current stimulation (mA) + The current time at which the dynamics is evaluated (s) + t_stim_prev: list[MX] | list[float] + The time list of the previous stimulations (s) + pulse_intensity: list[MX] | list[float] + The intensity of the stimulations (mA) + cn_sum: MX | float + The sum of the ca_troponin_complex (unitless) force_length_relationship: MX | float The force length relationship value (unitless) force_velocity_relationship: MX | float @@ -139,8 +162,7 @@ def system_dynamics( ------- The value of the derivative of each state dx/dt at the current time t """ - r0 = km + self.r0_km_relationship # Simplification - cn_dot = self.cn_dot_fun(cn, r0, t, t_stim_prev=t_stim_prev, intensity_stim=intensity_stim) # Equation n°1 + cn_dot = self.calculate_cn_dot(cn, cn_sum, t, t_stim_prev, pulse_intensity) f_dot = self.f_dot_fun( cn, f, @@ -209,7 +231,6 @@ def dynamics( algebraic_states: MX, numerical_timeseries: MX, nlp: NonLinearProgram, - stim_prev: list[float] = None, fes_model=None, force_length_relationship: float | MX = 1, force_velocity_relationship: float | MX = 1, @@ -233,9 +254,7 @@ def dynamics( The numerical timeseries of the system nlp: NonLinearProgram A reference to the phase - stim_prev: list[float] - The previous stimulation values - fes_model: DingModelIntensityFrequencyWithFatigue + fes_model: DingModelPulseIntensityFrequencyWithFatigue The current phase fes model force_length_relationship: MX | float The force length relationship value (unitless) @@ -245,34 +264,20 @@ def dynamics( ------- The derivative of the states in the tuple[MX] format """ - intensity_stim_prev = ( - [] - ) # Every stimulation intensity before the current phase, i.e.: the intensity of each phase - intensity_parameters = ( - nlp.model.get_intensity_parameters(nlp, parameters) - if fes_model is None - else fes_model.get_intensity_parameters(nlp, parameters, muscle_name=fes_model.muscle_name) - ) + model = fes_model if fes_model else nlp.model + dxdt_fun = model.system_dynamics - if intensity_parameters.shape[0] == 1: # check if pulse duration is mapped - for i in range(nlp.phase_idx + 1): - intensity_stim_prev.append(intensity_parameters[0]) + if model.is_approximated: + cn_sum = controls[0] + stim_apparition = None + intensity_parameters = None else: - for i in range(nlp.phase_idx + 1): - intensity_stim_prev.append(intensity_parameters[i]) + cn_sum = None + intensity_parameters = model.get_intensity_parameters(nlp, parameters) + stim_apparition = model.get_stim(nlp=nlp, parameters=parameters) - dxdt_fun = fes_model.system_dynamics if fes_model else nlp.model.system_dynamics - stim_apparition = ( - ( - fes_model.get_stim_prev(nlp=nlp, parameters=parameters, idx=nlp.phase_idx) - if fes_model - else nlp.model.get_stim_prev(nlp=nlp, parameters=parameters, idx=nlp.phase_idx) - ) - if stim_prev is None - else stim_prev - ) # Get the previous stimulation apparition time from the parameters - # if not provided from stim_prev, this way of getting the list is not optimal, but it is the only way to get it. - # Otherwise, it will create issues with free variables or wrong mx or sx type while calculating the dynamics + if len(intensity_parameters) == 1 and len(stim_apparition) != 1: + intensity_parameters = intensity_parameters * len(stim_apparition) return DynamicsEvaluation( dxdt=dxdt_fun( @@ -283,7 +288,8 @@ def dynamics( km=states[4], t=time, t_stim_prev=stim_apparition, - intensity_stim=intensity_stim_prev, + pulse_intensity=intensity_parameters, + cn_sum=cn_sum, force_length_relationship=force_length_relationship, force_velocity_relationship=force_velocity_relationship, ), @@ -291,7 +297,10 @@ def dynamics( ) def declare_ding_variables( - self, ocp: OptimalControlProgram, nlp: NonLinearProgram, numerical_data_timeseries: dict[str, np.ndarray] = None + self, + ocp: OptimalControlProgram, + nlp: NonLinearProgram, + numerical_data_timeseries: dict[str, np.ndarray] = None, ): """ Tell the program which variables are states and controls. @@ -306,9 +315,6 @@ def declare_ding_variables( A list of values to pass to the dynamics at each node. Experimental external forces should be included here. """ StateConfigure().configure_all_fes_model_states(ocp, nlp, fes_model=self) - stim_prev = ( - self._build_t_stim_prev(ocp, nlp.phase_idx) - if "pulse_apparition_time" not in nlp.parameters.keys() - else None - ) - ConfigureProblem.configure_dynamics_function(ocp, nlp, dyn_func=self.dynamics, stim_prev=stim_prev) + if self.is_approximated: + StateConfigure().configure_cn_sum(ocp, nlp) + ConfigureProblem.configure_dynamics_function(ocp, nlp, dyn_func=self.dynamics) diff --git a/cocofest/models/model_maker.py b/cocofest/models/model_maker.py new file mode 100644 index 0000000..5670288 --- /dev/null +++ b/cocofest/models/model_maker.py @@ -0,0 +1,22 @@ +from .ding2003 import DingModelFrequency +from .ding2003_with_fatigue import DingModelFrequencyWithFatigue +from .ding2007 import DingModelPulseWidthFrequency +from .ding2007_with_fatigue import DingModelPulseWidthFrequencyWithFatigue +from .hmed2018 import DingModelPulseIntensityFrequency +from .hmed2018_with_fatigue import DingModelPulseIntensityFrequencyWithFatigue + + +class ModelMaker: + @staticmethod + def create_model(model_type, **kwargs): + model_dict = { + "ding2003": DingModelFrequency, + "ding2003_with_fatigue": DingModelFrequencyWithFatigue, + "ding2007": DingModelPulseWidthFrequency, + "ding2007_with_fatigue": DingModelPulseWidthFrequencyWithFatigue, + "hmed2018": DingModelPulseIntensityFrequency, + "hmed2018_with_fatigue": DingModelPulseIntensityFrequencyWithFatigue, + } + if model_type not in model_dict: + raise ValueError(f"Unknown model type: {model_type}") + return model_dict[model_type](**kwargs) diff --git a/cocofest/models/state_configue.py b/cocofest/models/state_configure.py similarity index 83% rename from cocofest/models/state_configue.py rename to cocofest/models/state_configure.py index 809c61d..80f27b2 100644 --- a/cocofest/models/state_configue.py +++ b/cocofest/models/state_configure.py @@ -215,6 +215,44 @@ def configure_cross_bridges( as_states_dot, ) + @staticmethod + def configure_cn_sum(ocp, nlp, muscle_name: str = None): + """ + Configure the calcium summation control + + Parameters + ---------- + ocp: OptimalControlProgram + A reference to the ocp + nlp: NonLinearProgram + A reference to the phase + muscle_name: str + The muscle name + """ + muscle_name = "_" + muscle_name if muscle_name else "" + name = "Cn_sum" + muscle_name + name_cn_sum = [name] + return ConfigureProblem.configure_new_variable(name, name_cn_sum, ocp, nlp, as_states=False, as_controls=True) + + @staticmethod + def configure_a_calculation(ocp, nlp, muscle_name: str = None): + """ + Configure the force scaling factor calculation + + Parameters + ---------- + ocp: OptimalControlProgram + A reference to the ocp + nlp: NonLinearProgram + A reference to the phase + muscle_name: str + The muscle name + """ + muscle_name = "_" + muscle_name if muscle_name else "" + name = "A_calculation" + muscle_name + name_cn_sum = [name] + return ConfigureProblem.configure_new_variable(name, name_cn_sum, ocp, nlp, as_states=False, as_controls=True) + def configure_all_muscle_states(self, muscles_dynamics_model, ocp, nlp): state_name_list = [] for muscle_dynamics_model in muscles_dynamics_model: @@ -235,5 +273,9 @@ def configure_all_fes_model_states(self, ocp, nlp, fes_model): for state_key in fes_model.name_dof: if state_key in self.state_dictionary.keys(): self.state_dictionary[state_key]( - ocp=ocp, nlp=nlp, as_states=True, as_controls=False, muscle_name=fes_model.muscle_name + ocp=ocp, + nlp=nlp, + as_states=True, + as_controls=False, + muscle_name=fes_model.muscle_name, ) diff --git a/cocofest/optimization/fes_identification_ocp.py b/cocofest/optimization/fes_identification_ocp.py index 05b77ae..5a81608 100644 --- a/cocofest/optimization/fes_identification_ocp.py +++ b/cocofest/optimization/fes_identification_ocp.py @@ -15,17 +15,16 @@ PhaseTransitionFcn, PhaseTransitionList, VariableScaling, + Node, ) -from ..custom_objectives import CustomObjective from ..models.fes_model import FesModel -from ..models.ding2007 import DingModelPulseDurationFrequency -from ..models.ding2007_with_fatigue import DingModelPulseDurationFrequencyWithFatigue + from ..models.ding2003 import DingModelFrequency -from ..models.ding2003_with_fatigue import DingModelFrequencyWithFatigue -from ..models.hmed2018 import DingModelIntensityFrequency -from ..models.hmed2018_with_fatigue import DingModelIntensityFrequencyWithFatigue +from ..models.ding2007 import DingModelPulseWidthFrequency +from ..models.hmed2018 import DingModelPulseIntensityFrequency from ..optimization.fes_ocp import OcpFes +from ..custom_constraints import CustomConstraint class OcpFesId(OcpFes): @@ -35,11 +34,11 @@ def __init__(self): @staticmethod def prepare_ocp( model: FesModel = None, - n_shooting: list[int] = None, - final_time_phase: tuple | list = None, - pulse_duration: int | float | list = None, - pulse_intensity: int | float | list = None, - force_tracking: list = None, + final_time: float | int = None, + stim_time: list = None, + pulse_width: dict = None, + pulse_intensity: dict = None, + objective: dict = None, key_parameter_to_identify: list = None, additional_key_settings: dict = None, custom_objective: list[Objective] = None, @@ -47,6 +46,7 @@ def prepare_ocp( use_sx: bool = True, ode_solver: OdeSolver = OdeSolver.RK4(n_integration_steps=1), n_threads: int = 1, + control_type: ControlType = ControlType.CONSTANT, **kwargs, ): """ @@ -57,16 +57,14 @@ def prepare_ocp( ---------- model: FesModel The model used to solve the ocp - final_time_phase: tuple, list + final_time: float, int The final time of each phase, it corresponds to the stimulation apparition time - n_shooting: list[int], - The number of shooting points for each phase - force_tracking: list[np.ndarray, np.ndarray], - The force tracking to follow - pulse_duration: int | float | list[int] | list[float], + pulse_width: dict, The duration of the stimulation - pulse_intensity: int | float | list[int] | list[float], + pulse_intensity: dict, The intensity of the stimulation + objective: dict, + The objective to minimize discontinuity_in_ocp: list[int], The phases where the continuity is not respected ode_solver: OdeSolver @@ -76,126 +74,133 @@ def prepare_ocp( n_thread: int The number of thread to use while solving (multi-threading if > 1) """ - + ( + pulse_event, + pulse_width, + pulse_intensity, + temp_objective, + ) = OcpFes._fill_dict({}, pulse_width, pulse_intensity, {}) + + n_shooting = OcpFes.prepare_n_shooting(stim_time, final_time) OcpFesId._sanity_check( model=model, - custom_objective=custom_objective, + n_shooting=n_shooting, + final_time=final_time, + pulse_event=pulse_event, + pulse_width=pulse_width, + pulse_intensity=pulse_intensity, + objective=temp_objective, use_sx=use_sx, ode_solver=ode_solver, n_threads=n_threads, - fixed_pulse_duration=pulse_duration, - fixed_pulse_intensity=pulse_intensity, ) OcpFesId._sanity_check_id( model=model, - n_shooting=n_shooting, - final_time_phase=final_time_phase, - force_tracking=force_tracking, - pulse_duration=pulse_duration, + final_time=final_time, + objective=objective, + pulse_width=pulse_width, pulse_intensity=pulse_intensity, ) - n_stim = len(final_time_phase) - models = [model for i in range(n_stim)] + n_stim = len(stim_time) - constraints = ConstraintList() parameters, parameters_bounds, parameters_init = OcpFesId._set_parameters( n_stim=n_stim, - stim_apparition_time=final_time_phase, + stim_apparition_time=stim_time, parameter_to_identify=key_parameter_to_identify, parameter_setting=additional_key_settings, - pulse_duration=pulse_duration, + pulse_width=pulse_width, pulse_intensity=pulse_intensity, use_sx=use_sx, ) - dynamics = OcpFesId._declare_dynamics(models=models, n_stim=n_stim) + dynamics = OcpFesId._declare_dynamics(model=model) x_bounds, x_init = OcpFesId._set_bounds( model=model, - n_stim=n_stim, - n_shooting=n_shooting, - force_tracking=force_tracking, + force_tracking=objective["force_tracking"], discontinuity_in_ocp=discontinuity_in_ocp, ) - objective_functions = OcpFesId._set_objective( - model=model, - n_stim=n_stim, - n_shooting=n_shooting, - force_tracking=force_tracking, - custom_objective=custom_objective, - ) - phase_transitions = OcpFesId._set_phase_transition(discontinuity_in_ocp) + objective_functions = OcpFesId._set_objective(model=model, objective=objective) + + if model.is_approximated: + constraints = OcpFesId._build_constraints( + model=model, + n_shooting=n_shooting, + final_time=final_time, + stim_time=stim_time, + control_type=control_type, + ) + u_bounds, u_init = OcpFesId._set_u_bounds(model=model) + else: + constraints = ConstraintList() + u_bounds, u_init = None, None + + # phase_transitions = OcpFesId._set_phase_transition(discontinuity_in_ocp) return OptimalControlProgram( - bio_model=models, + bio_model=[model], dynamics=dynamics, n_shooting=n_shooting, - phase_time=final_time_phase, + phase_time=final_time, x_init=x_init, x_bounds=x_bounds, + u_init=u_init, + u_bounds=u_bounds, objective_functions=objective_functions, constraints=constraints, ode_solver=ode_solver, - control_type=ControlType.CONSTANT, + control_type=control_type, use_sx=use_sx, parameters=parameters, parameter_bounds=parameters_bounds, parameter_init=parameters_init, - phase_transitions=phase_transitions, + # phase_transitions=phase_transitions, n_threads=n_threads, ) @staticmethod def _sanity_check_id( model=None, - n_shooting=None, - final_time_phase=None, - force_tracking=None, - pulse_duration=None, + final_time=None, + objective=None, + pulse_width=None, pulse_intensity=None, ): - if not isinstance(n_shooting, list): - raise TypeError(f"n_shooting must be list type," f" currently n_shooting is {type(n_shooting)}) type.") - else: - if not all(isinstance(val, int) for val in n_shooting): - raise TypeError(f"n_shooting must be list of int type.") - - if isinstance(final_time_phase, tuple): - if not all(isinstance(val, int | float) for val in final_time_phase): - raise TypeError(f"final_time_phase must be tuple of int or float type.") - if len(final_time_phase) != len(n_shooting): - raise ValueError( - f"final_time_phase must have same length as n_shooting, currently final_time_phase is {len(final_time_phase)} and n_shooting is {len(n_shooting)}." - ) - else: - raise TypeError( - f"final_time_phase must be tuple type," - f" currently final_time_phase is {type(final_time_phase)}) type." - ) + if not isinstance(final_time, int | float): + raise TypeError(f"final_time must be int or float type.") - if not isinstance(force_tracking, list): + if not isinstance(objective["force_tracking"], list): raise TypeError( - f"force_tracking must be list type," f" currently force_tracking is {type(force_tracking)}) type." + f"force_tracking must be list type," + f" currently force_tracking is {type(objective['force_tracking'])}) type." ) else: - if not all(isinstance(val, int | float) for val in force_tracking): + if not all(isinstance(val, int | float) for val in objective["force_tracking"]): raise TypeError(f"force_tracking must be list of int or float type.") - if isinstance(model, DingModelPulseDurationFrequency): - if not isinstance(pulse_duration, list): + if isinstance(model, DingModelPulseWidthFrequency): + if not isinstance(pulse_width, dict): raise TypeError( - f"pulse_duration must be list type," f" currently pulse_duration is {type(pulse_duration)}) type." + f"pulse_width must be dict type," f" currently pulse_width is {type(pulse_width)}) type." ) - if isinstance(model, DingModelIntensityFrequency): - if not isinstance(pulse_intensity, list): + if isinstance(model, DingModelPulseIntensityFrequency): + if isinstance(pulse_intensity, dict): + if not isinstance(pulse_intensity["fixed"], int | float | list): + raise ValueError(f"fixed pulse_intensity must be a int, float or list type.") + + else: raise TypeError( - f"pulse_intensity must be list type," + f"pulse_intensity must be dict type," f" currently pulse_intensity is {type(pulse_intensity)}) type." ) @staticmethod - def _set_bounds(model=None, n_stim=None, n_shooting=None, force_tracking=None, discontinuity_in_ocp=None): + def _set_bounds( + model: FesModel = None, + force_tracking=None, + discontinuity_in_ocp=None, + ): # ---- STATE BOUNDS REPRESENTATION ---- # # # |‾‾‾‾‾‾‾‾‾‾x_max_middle‾‾‾‾‾‾‾‾‾‾‾‾x_max_end‾ @@ -228,71 +233,52 @@ def _set_bounds(model=None, n_stim=None, n_shooting=None, force_tracking=None, d starting_bounds_min = np.concatenate((starting_bounds, min_bounds, min_bounds), axis=1) starting_bounds_max = np.concatenate((starting_bounds, max_bounds, max_bounds), axis=1) - middle_bound_min = np.concatenate((min_bounds, min_bounds, min_bounds), axis=1) - middle_bound_max = np.concatenate((max_bounds, max_bounds, max_bounds), axis=1) - for i in range(n_stim): - for j in range(len(variable_bound_list)): - if i == 0 or i in discontinuity_in_ocp: - x_bounds.add( - variable_bound_list[j], - min_bound=np.array([starting_bounds_min[j]]), - max_bound=np.array([starting_bounds_max[j]]), - phase=i, - interpolation=InterpolationType.CONSTANT_WITH_FIRST_AND_LAST_DIFFERENT, - ) - else: - x_bounds.add( - variable_bound_list[j], - min_bound=np.array([middle_bound_min[j]]), - max_bound=np.array([middle_bound_max[j]]), - phase=i, - interpolation=InterpolationType.CONSTANT_WITH_FIRST_AND_LAST_DIFFERENT, - ) + for j in range(len(variable_bound_list)): + x_bounds.add( + variable_bound_list[j], + min_bound=np.array([starting_bounds_min[j]]), + max_bound=np.array([starting_bounds_max[j]]), + phase=0, + interpolation=InterpolationType.CONSTANT_WITH_FIRST_AND_LAST_DIFFERENT, + ) x_init = InitialGuessList() - for i in range(n_stim): - min_node = sum(n_shooting[:i]) - max_node = sum(n_shooting[: i + 1]) - force_in_phase = force_tracking[min_node : max_node + 1] - if i == n_stim - 1: - force_in_phase.append(0) - x_init.add("F", np.array([force_in_phase]), phase=i, interpolation=InterpolationType.EACH_FRAME) - x_init.add("Cn", [0], phase=i, interpolation=InterpolationType.CONSTANT) - if model._with_fatigue: - for j in range(len(variable_bound_list)): - if variable_bound_list[j] == "F" or variable_bound_list[j] == "Cn": - pass - else: - x_init.add(variable_bound_list[j], model.standard_rest_values()[j]) + + x_init.add( + "F", + np.array([force_tracking]), + phase=0, + interpolation=InterpolationType.EACH_FRAME, + ) + x_init.add("Cn", [0], phase=0, interpolation=InterpolationType.CONSTANT) + if model._with_fatigue: + for j in range(len(variable_bound_list)): + if variable_bound_list[j] == "F" or variable_bound_list[j] == "Cn": + pass + else: + x_init.add(variable_bound_list[j], model.standard_rest_values()[j]) return x_bounds, x_init @staticmethod - def _set_objective(model, n_stim, n_shooting, force_tracking, custom_objective, **kwargs): + def _set_objective(model, objective): # Creates the objective for our problem (in this case, match a force curve) objective_functions = ObjectiveList() - if force_tracking: - node_idx = 0 - for i in range(n_stim): - for j in range(n_shooting[i]): - objective_functions.add( - CustomObjective.track_state_from_time_interpolate, - custom_type=ObjectiveFcn.Mayer, - node=j, - force=force_tracking[node_idx], - key="F", - minimization_type="best fit" if model._with_fatigue else "least square", - quadratic=True, - weight=1, - phase=i, - ) - node_idx += 1 - - if custom_objective: - for i in range(len(custom_objective)): - objective_functions.add(custom_objective[i]) + if objective["force_tracking"]: + objective_functions.add( + ObjectiveFcn.Lagrange.TRACK_STATE, + key="F", + weight=1, + target=np.array(objective["force_tracking"])[np.newaxis, :], + node=Node.ALL, + quadratic=True, + ) + + if "custom" in objective and objective["custom"] is not None: + for i in range(len(objective["custom"])): + objective_functions.add(objective["custom"][i]) return objective_functions @@ -303,20 +289,37 @@ def _set_parameters( parameter_to_identify, parameter_setting, use_sx, - pulse_duration=None, - pulse_intensity=None, + pulse_width: dict = None, + pulse_intensity: dict = None, ): parameters = ParameterList(use_sx=use_sx) parameters_bounds = BoundsList() parameters_init = InitialGuessList() + parameters.add( + name="pulse_apparition_time", + function=DingModelFrequency.set_pulse_apparition_time, + size=n_stim, + scaling=VariableScaling("pulse_apparition_time", [1] * n_stim), + ) + + parameters_init["pulse_apparition_time"] = np.array(stim_apparition_time) + + parameters_bounds.add( + "pulse_apparition_time", + min_bound=stim_apparition_time, + max_bound=stim_apparition_time, + interpolation=InterpolationType.CONSTANT, + ) + for i in range(len(parameter_to_identify)): parameters.add( name=parameter_to_identify[i], function=parameter_setting[parameter_to_identify[i]]["function"], size=1, scaling=VariableScaling( - parameter_to_identify[i], [parameter_setting[parameter_to_identify[i]]["scaling"]] + parameter_to_identify[i], + [parameter_setting[parameter_to_identify[i]]["scaling"]], ), ) parameters_bounds.add( @@ -330,60 +333,111 @@ def _set_parameters( initial_guess=np.array([parameter_setting[parameter_to_identify[i]]["initial_guess"]]), ) - if pulse_duration: + if pulse_width["fixed"]: parameters.add( - name="pulse_duration", - function=DingModelPulseDurationFrequency.set_impulse_duration, + name="pulse_width", + function=DingModelPulseWidthFrequency.set_impulse_width, size=n_stim, - scaling=VariableScaling("pulse_duration", [1] * n_stim), + scaling=VariableScaling("pulse_width", [1] * n_stim), ) - if isinstance(pulse_duration, list): + if isinstance(pulse_width["fixed"], list): parameters_bounds.add( - "pulse_duration", - min_bound=np.array(pulse_duration), - max_bound=np.array(pulse_duration), + "pulse_width", + min_bound=np.array(pulse_width["fixed"]), + max_bound=np.array(pulse_width["fixed"]), interpolation=InterpolationType.CONSTANT, ) - parameters_init.add(key="pulse_duration", initial_guess=np.array(pulse_duration)) + parameters_init.add(key="pulse_width", initial_guess=np.array(pulse_width["fixed"])) else: parameters_bounds.add( - "pulse_duration", - min_bound=np.array([pulse_duration] * n_stim), - max_bound=np.array([pulse_duration] * n_stim), + "pulse_width", + min_bound=np.array([pulse_width["fixed"]] * n_stim), + max_bound=np.array([pulse_width["fixed"]] * n_stim), interpolation=InterpolationType.CONSTANT, ) - parameters_init.add(key="pulse_duration", initial_guess=np.array([pulse_duration] * n_stim)) + parameters_init.add( + key="pulse_width", + initial_guess=np.array([pulse_width] * n_stim), + ) - if pulse_intensity: + if pulse_intensity["fixed"]: parameters.add( name="pulse_intensity", - function=DingModelIntensityFrequency.set_impulse_intensity, + function=DingModelPulseIntensityFrequency.set_impulse_intensity, size=n_stim, scaling=VariableScaling("pulse_intensity", [1] * n_stim), ) - if isinstance(pulse_intensity, list): + if isinstance(pulse_intensity["fixed"], list): parameters_bounds.add( "pulse_intensity", - min_bound=np.array(pulse_intensity), - max_bound=np.array(pulse_intensity), + min_bound=np.array(pulse_intensity["fixed"]), + max_bound=np.array(pulse_intensity["fixed"]), interpolation=InterpolationType.CONSTANT, ) - parameters_init.add(key="pulse_intensity", initial_guess=np.array(pulse_intensity)) + parameters_init.add(key="pulse_intensity", initial_guess=np.array(pulse_intensity["fixed"])) else: parameters_bounds.add( "pulse_intensity", - min_bound=np.array([pulse_intensity] * n_stim), - max_bound=np.array([pulse_intensity] * n_stim), + min_bound=np.array([pulse_intensity["fixed"]] * n_stim), + max_bound=np.array([pulse_intensity["fixed"]] * n_stim), interpolation=InterpolationType.CONSTANT, ) - parameters_init.add(key="pulse_intensity", initial_guess=np.array([pulse_intensity] * n_stim)) + parameters_init.add( + key="pulse_intensity", + initial_guess=np.array([pulse_intensity["fixed"]] * n_stim), + ) return parameters, parameters_bounds, parameters_init + @staticmethod + def _build_constraints(model, n_shooting, final_time, stim_time, control_type): + constraints = ConstraintList() + + time_vector = np.linspace(0, final_time, n_shooting + 1) + stim_at_node = [np.where(stim_time[i] <= time_vector)[0][0] for i in range(len(stim_time))] + additional_nodes = 1 if control_type == ControlType.LINEAR_CONTINUOUS else 0 + if model._sum_stim_truncation: + max_stim_to_keep = model._sum_stim_truncation + else: + max_stim_to_keep = 10000000 + + index_sup = 0 + index_inf = 0 + stim_index = [] + for i in range(n_shooting + additional_nodes): + if i in stim_at_node: + index_sup += 1 + if index_sup >= max_stim_to_keep: + index_inf = index_sup - max_stim_to_keep + stim_index = [i for i in range(index_inf, index_sup)] + + constraints.add( + CustomConstraint.cn_sum_identification, + node=i, + stim_time=stim_time[index_inf:index_sup], + stim_index=stim_index, + ) + + if isinstance(model, DingModelPulseWidthFrequency): + index_sup = 0 + for i in range(n_shooting + additional_nodes): + if i in stim_at_node and i != 0: + index_sup += 1 + constraints.add( + CustomConstraint.a_calculation_identification, + node=i, + last_stim_index=index_sup, + ) + + return constraints + @staticmethod def _set_phase_transition(discontinuity_in_ocp): phase_transitions = PhaseTransitionList() if discontinuity_in_ocp: for i in range(len(discontinuity_in_ocp)): - phase_transitions.add(PhaseTransitionFcn.DISCONTINUOUS, phase_pre_idx=discontinuity_in_ocp[i] - 1) + phase_transitions.add( + PhaseTransitionFcn.DISCONTINUOUS, + phase_pre_idx=discontinuity_in_ocp[i] - 1, + ) return phase_transitions diff --git a/cocofest/optimization/fes_ocp.py b/cocofest/optimization/fes_ocp.py index a36c819..03605e3 100644 --- a/cocofest/optimization/fes_ocp.py +++ b/cocofest/optimization/fes_ocp.py @@ -19,43 +19,108 @@ VariableScaling, ) -from ..custom_objectives import CustomObjective -from ..custom_constraints import CustomConstraint from ..fourier_approx import FourierSeries - from ..models.fes_model import FesModel -from ..models.ding2007 import DingModelPulseDurationFrequency -from ..models.ding2007_with_fatigue import DingModelPulseDurationFrequencyWithFatigue +from ..models.dynamical_model import FesMskModel from ..models.ding2003 import DingModelFrequency -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 ..custom_constraints import CustomConstraint class OcpFes: """ The main class to define an ocp. This class prepares the full program and gives all - the needed parameters to solve a functional electrical stimulation ocp - - Methods - ------- - from_frequency_and_final_time(self, frequency: int | float, final_time: float, round_down: bool) - Calculates the number of stim (phases) for the ocp from frequency and final time - from_frequency_and_n_stim(self, frequency: int | float, n_stim: int) - Calculates the final ocp time from frequency and stimulation number + the needed parameters to solve a functional electrical stimulation ocp. """ + @staticmethod + def _prepare_optimization_problem(input_dict: dict) -> dict: + + (pulse_event, pulse_width, pulse_intensity, objective) = OcpFes._fill_dict( + input_dict["pulse_event"], + input_dict["pulse_width"], + input_dict["pulse_intensity"], + input_dict["objective"], + ) + + OcpFes._sanity_check( + model=input_dict["model"], + n_shooting=input_dict["n_shooting"], + final_time=input_dict["final_time"], + pulse_event=pulse_event, + pulse_width=pulse_width, + pulse_intensity=pulse_intensity, + objective=objective, + use_sx=input_dict["use_sx"], + ode_solver=input_dict["ode_solver"], + n_threads=input_dict["n_threads"], + ) + + (parameters, parameters_bounds, parameters_init, parameter_objectives) = OcpFes._build_parameters( + model=input_dict["model"], + stim_time=input_dict["stim_time"], + pulse_event=pulse_event, + pulse_width=pulse_width, + pulse_intensity=pulse_intensity, + use_sx=input_dict["use_sx"], + ) + + dynamics = OcpFes._declare_dynamics(input_dict["model"]) + x_bounds, x_init = OcpFes._set_bounds(input_dict["model"]) + + if input_dict["model"].is_approximated: + constraints = OcpFes._build_constraints( + input_dict["model"], + input_dict["n_shooting"], + input_dict["final_time"], + input_dict["stim_time"], + input_dict["control_type"], + ) + u_bounds, u_init = OcpFes._set_u_bounds(input_dict["model"]) + else: + constraints = ConstraintList() + u_bounds, u_init = BoundsList(), InitialGuessList() + + objective_functions = OcpFes._set_objective(input_dict["n_shooting"], objective) + + optimization_dict = { + "model": input_dict["model"], + "dynamics": dynamics, + "n_shooting": input_dict["n_shooting"], + "final_time": input_dict["final_time"], + "objective_functions": objective_functions, + "x_init": x_init, + "x_bounds": x_bounds, + "u_bounds": u_bounds, + "u_init": u_init, + "constraints": constraints, + "parameters": parameters, + "parameters_bounds": parameters_bounds, + "parameters_init": parameters_init, + "parameter_objectives": parameter_objectives, + "use_sx": input_dict["use_sx"], + "ode_solver": input_dict["ode_solver"], + "n_threads": input_dict["n_threads"], + "control_type": input_dict["control_type"], + } + + return optimization_dict + @staticmethod def prepare_ocp( model: FesModel = None, - n_stim: int = None, - n_shooting: int = None, + stim_time: list = None, final_time: int | float = None, pulse_event: dict = None, - pulse_duration: dict = None, + pulse_width: dict = None, pulse_intensity: dict = None, objective: dict = None, use_sx: bool = True, ode_solver: OdeSolver = OdeSolver.RK4(n_integration_steps=1), + control_type: ControlType = ControlType.CONSTANT, n_threads: int = 1, ): """ @@ -65,20 +130,18 @@ def prepare_ocp( ---------- model : FesModel The model type used for the OCP. - n_stim : int - Number of stimulations that will occur during the OCP, also referred to as phases. - n_shooting : int - Number of shooting points for each individual phase. + stim_time : list + All the stimulation time. final_time : int | float The final time of the OCP. pulse_event : dict Dictionary containing parameters related to the appearance of the pulse. - pulse_duration : dict + pulse_width : dict Dictionary containing parameters related to the duration of the pulse. - Optional if not using DingModelPulseDurationFrequency or DingModelPulseDurationFrequencyWithFatigue. + Optional if not using DingModelPulseWidthFrequency or DingModelPulseWidthFrequencyWithFatigue. pulse_intensity : dict Dictionary containing parameters related to the intensity of the pulse. - Optional if not using DingModelIntensityFrequency or DingModelIntensityFrequencyWithFatigue. + Optional if not using DingModelPulseIntensityFrequency or DingModelPulseIntensityFrequencyWithFatigue. objective : dict Dictionary containing parameters related to the optimization objective. use_sx : bool @@ -87,135 +150,76 @@ def prepare_ocp( The ODE solver to use. n_threads : int The number of threads to use while solving (multi-threading if > 1). + control_type : ControlType + The type of control to use. Returns ------- OptimalControlProgram The prepared Optimal Control Program. - """ - (pulse_event, pulse_duration, pulse_intensity, objective) = OcpFes._fill_dict( - pulse_event, pulse_duration, pulse_intensity, objective - ) - - time_min = pulse_event["min"] - time_max = pulse_event["max"] - time_bimapping = pulse_event["bimapping"] - frequency = pulse_event["frequency"] - round_down = pulse_event["round_down"] - pulse_mode = pulse_event["pulse_mode"] - - fixed_pulse_duration = pulse_duration["fixed"] - pulse_duration_min = pulse_duration["min"] - pulse_duration_max = pulse_duration["max"] - pulse_duration_bimapping = pulse_duration["bimapping"] - - fixed_pulse_intensity = pulse_intensity["fixed"] - pulse_intensity_min = pulse_intensity["min"] - pulse_intensity_max = pulse_intensity["max"] - pulse_intensity_bimapping = pulse_intensity["bimapping"] - - force_tracking = objective["force_tracking"] - end_node_tracking = objective["end_node_tracking"] - custom_objective = objective["custom"] - - OcpFes._sanity_check( - model=model, - n_stim=n_stim, - n_shooting=n_shooting, - final_time=final_time, - pulse_mode=pulse_mode, - frequency=frequency, - time_min=time_min, - time_max=time_max, - time_bimapping=time_bimapping, - fixed_pulse_duration=fixed_pulse_duration, - pulse_duration_min=pulse_duration_min, - pulse_duration_max=pulse_duration_max, - pulse_duration_bimapping=pulse_duration_bimapping, - fixed_pulse_intensity=fixed_pulse_intensity, - pulse_intensity_min=pulse_intensity_min, - pulse_intensity_max=pulse_intensity_max, - pulse_intensity_bimapping=pulse_intensity_bimapping, - force_tracking=force_tracking, - end_node_tracking=end_node_tracking, - custom_objective=custom_objective, - use_sx=use_sx, - ode_solver=ode_solver, - n_threads=n_threads, - ) - - OcpFes._sanity_check_frequency(n_stim=n_stim, final_time=final_time, frequency=frequency, round_down=round_down) - - n_stim, final_time = OcpFes._build_phase_parameter( - n_stim=n_stim, final_time=final_time, frequency=frequency, pulse_mode=pulse_mode, round_down=round_down - ) - - force_fourier_coefficient = ( - None if force_tracking is None else OcpFes._build_fourier_coefficient(force_tracking) - ) + input_dict = { + "model": model, + "stim_time": stim_time, + "n_shooting": OcpFes.prepare_n_shooting(stim_time, final_time), + "final_time": final_time, + "pulse_event": pulse_event, + "pulse_width": pulse_width, + "pulse_intensity": pulse_intensity, + "objective": objective, + "use_sx": use_sx, + "ode_solver": ode_solver, + "n_threads": n_threads, + "control_type": control_type, + } - models = [model] * n_stim - n_shooting = [n_shooting] * n_stim + optimization_dict = OcpFes._prepare_optimization_problem(input_dict) - final_time_phase = OcpFes._build_phase_time( - final_time=final_time, - n_stim=n_stim, - pulse_mode=pulse_mode, - time_min=time_min, - time_max=time_max, - ) - parameters, parameters_bounds, parameters_init, parameter_objectives, constraints = OcpFes._build_parameters( - model=model, - n_stim=n_stim, - time_min=time_min, - time_max=time_max, - time_bimapping=time_bimapping, - fixed_pulse_duration=fixed_pulse_duration, - pulse_duration_min=pulse_duration_min, - pulse_duration_max=pulse_duration_max, - pulse_duration_bimapping=pulse_duration_bimapping, - fixed_pulse_intensity=fixed_pulse_intensity, - pulse_intensity_min=pulse_intensity_min, - pulse_intensity_max=pulse_intensity_max, - pulse_intensity_bimapping=pulse_intensity_bimapping, - use_sx=use_sx, + return OptimalControlProgram( + bio_model=[optimization_dict["model"]], + dynamics=optimization_dict["dynamics"], + n_shooting=optimization_dict["n_shooting"], + phase_time=[optimization_dict["final_time"]], + objective_functions=optimization_dict["objective_functions"], + x_init=optimization_dict["x_init"], + x_bounds=optimization_dict["x_bounds"], + u_bounds=optimization_dict["u_bounds"], + u_init=optimization_dict["u_init"], + constraints=optimization_dict["constraints"], + parameters=optimization_dict["parameters"], + parameter_bounds=optimization_dict["parameters_bounds"], + parameter_init=optimization_dict["parameters_init"], + parameter_objectives=optimization_dict["parameter_objectives"], + control_type=optimization_dict["control_type"], + use_sx=optimization_dict["use_sx"], + ode_solver=optimization_dict["ode_solver"], + n_threads=optimization_dict["n_threads"], ) - if len(constraints) == 0 and len(parameters) == 0: - raise ValueError( - "This is not an optimal control problem," - " add parameter to optimize or use the IvpFes method to build your problem" - ) - - dynamics = OcpFes._declare_dynamics(models, n_stim) - x_bounds, x_init = OcpFes._set_bounds(model, n_stim) - objective_functions = OcpFes._set_objective( - n_stim, n_shooting, force_fourier_coefficient, end_node_tracking, custom_objective, time_min, time_max - ) + @staticmethod + def prepare_n_shooting(stim_time, final_time): + """ + Prepare the n_shooting for the ocp in order to have a time step that is a multiple of the stimulation time. - return OptimalControlProgram( - bio_model=models, - dynamics=dynamics, - n_shooting=n_shooting, - phase_time=final_time_phase, - objective_functions=objective_functions, - x_init=x_init, - x_bounds=x_bounds, - constraints=constraints, - parameters=parameters, - parameter_bounds=parameters_bounds, - parameter_init=parameters_init, - parameter_objectives=parameter_objectives, - control_type=ControlType.CONSTANT, - use_sx=use_sx, - ode_solver=ode_solver, - n_threads=n_threads, - ) + Returns + ------- + int + The number of shooting points + """ + stim_time_str = [str(t) for t in stim_time] + stim_time_str = [ + stim_time_str[i] + ".0" if len(stim_time_str[i]) == 1 else stim_time_str[i] + for i in range(len(stim_time_str)) + ] + nb_decimal = max([len(stim_time_str[i].split(".")[1]) for i in range(len(stim_time))]) + nb_decimal = 2 if nb_decimal < 2 else nb_decimal + decimal_shooting = int("1" + "0" * nb_decimal) + n_shooting = int(final_time * decimal_shooting) + return n_shooting @staticmethod - def _fill_dict(pulse_event, pulse_duration, pulse_intensity, objective): + def _fill_dict(pulse_event, pulse_width, pulse_intensity, objective): """ This method fills the provided dictionaries with default values if they are not set. @@ -225,7 +229,7 @@ def _fill_dict(pulse_event, pulse_duration, pulse_intensity, objective): Dictionary containing parameters related to the appearance of the pulse. Expected keys are 'min', 'max', 'bimapping', 'frequency', 'round_down', and 'pulse_mode'. - pulse_duration : dict + pulse_width : dict Dictionary containing parameters related to the duration of the pulse. Expected keys are 'fixed', 'min', 'max', and 'bimapping'. @@ -239,10 +243,10 @@ def _fill_dict(pulse_event, pulse_duration, pulse_intensity, objective): Returns ------- - Returns four dictionaries: pulse_event, pulse_duration, pulse_intensity, and objective. + Returns four dictionaries: pulse_event, pulse_width, pulse_intensity, and objective. Each dictionary is filled with default values for any keys that were not initially set. """ - + pulse_event = {} if pulse_event is None else pulse_event default_pulse_event = { "min": None, "max": None, @@ -252,13 +256,15 @@ def _fill_dict(pulse_event, pulse_duration, pulse_intensity, objective): "pulse_mode": "single", } - default_pulse_duration = { + pulse_width = {} if pulse_width is None else pulse_width + default_pulse_width = { "fixed": None, "min": None, "max": None, "bimapping": False, } + pulse_intensity = {} if pulse_intensity is None else pulse_intensity default_pulse_intensity = { "fixed": None, "min": None, @@ -266,153 +272,123 @@ def _fill_dict(pulse_event, pulse_duration, pulse_intensity, objective): "bimapping": False, } + objective = {} if objective is None else objective default_objective = { "force_tracking": None, "end_node_tracking": None, "cycling": None, "custom": None, } - dict_list = [pulse_event, pulse_duration, pulse_intensity, objective] - default_dict_list = [ - default_pulse_event, - default_pulse_duration, - default_pulse_intensity, - default_objective, - ] - - for i in range(len(dict_list)): - if dict_list[i] is None: - dict_list[i] = {} - for i in range(len(dict_list)): - for key in default_dict_list[i]: - if key not in dict_list[i]: - dict_list[i][key] = default_dict_list[i][key] + pulse_event = {**default_pulse_event, **pulse_event} + pulse_width = {**default_pulse_width, **pulse_width} + pulse_intensity = {**default_pulse_intensity, **pulse_intensity} + objective = {**default_objective, **objective} - return dict_list[0], dict_list[1], dict_list[2], dict_list[3] + return pulse_event, pulse_width, pulse_intensity, objective @staticmethod def _sanity_check( model=None, - n_stim=None, n_shooting=None, final_time=None, - pulse_mode=None, - frequency=None, - time_min=None, - time_max=None, - time_bimapping=None, - fixed_pulse_duration=None, - pulse_duration_min=None, - pulse_duration_max=None, - pulse_duration_bimapping=None, - fixed_pulse_intensity=None, - pulse_intensity_min=None, - pulse_intensity_max=None, - pulse_intensity_bimapping=None, - force_tracking=None, - end_node_tracking=None, - custom_objective=None, + pulse_event=None, + pulse_width=None, + pulse_intensity=None, + objective=None, use_sx=None, ode_solver=None, n_threads=None, ): if not isinstance(model, FesModel): - raise TypeError( - f"The current model type used is {type(model)}, it must be a FesModel type." - f"Current available models are: DingModelFrequency, DingModelFrequencyWithFatigue," - f"DingModelPulseDurationFrequency, DingModelPulseDurationFrequencyWithFatigue," - f"DingModelIntensityFrequency, DingModelIntensityFrequencyWithFatigue" - ) - - if n_stim: - if isinstance(n_stim, int): - if n_stim <= 0: - raise ValueError("n_stim must be positive") + if isinstance(model, FesMskModel): + pass else: - raise TypeError("n_stim must be int type") + raise TypeError( + f"The current model type used is {type(model)}, it must be a FesModel type." + f"Current available models are: DingModelFrequency, DingModelFrequencyWithFatigue," + f"DingModelPulseWidthFrequency, DingModelPulseWidthFrequencyWithFatigue," + f"DingModelPulseIntensityFrequency, DingModelPulseIntensityFrequencyWithFatigue" + ) - if n_shooting: - if isinstance(n_shooting, int): - if n_shooting <= 0: - raise ValueError("n_shooting must be positive") - else: - raise TypeError("n_shooting must be int type") + if not isinstance(n_shooting, int) or n_shooting < 0: + raise TypeError("n_shooting must be a positive int type") - if final_time: - if isinstance(final_time, int | float): - if final_time <= 0: - raise ValueError("final_time must be positive") - else: - raise TypeError("final_time must be int or float type") + if not isinstance(final_time, int | float) or final_time < 0: + raise TypeError("final_time must be a positive int or float type") - if pulse_mode: - if pulse_mode != "single": - raise NotImplementedError(f"Pulse mode '{pulse_mode}' is not yet implemented") + if pulse_event["pulse_mode"] != "single": + raise NotImplementedError(f"Pulse mode '{pulse_event['pulse_mode']}' is not yet implemented") - if frequency: - if isinstance(frequency, int | float): - if frequency <= 0: + if pulse_event["frequency"]: + if isinstance(pulse_event["frequency"], int | float): + if pulse_event["frequency"] <= 0: raise ValueError("frequency must be positive") else: raise TypeError("frequency must be int or float type") - if [time_min, time_max].count(None) == 1: - raise ValueError("time_min and time_max must be both entered or none of them in order to work") + if [pulse_event["min"], pulse_event["max"]].count(None) == 1: + raise ValueError("min and max time event must be both entered or none of them in order to work") - if time_bimapping: - if not isinstance(time_bimapping, bool): - raise TypeError("time_bimapping must be bool type") + if pulse_event["bimapping"]: + if not isinstance(pulse_event["bimapping"], bool): + raise TypeError("time bimapping must be bool type") - if isinstance(model, DingModelPulseDurationFrequency | DingModelPulseDurationFrequencyWithFatigue): - if fixed_pulse_duration is None and [pulse_duration_min, pulse_duration_max].count(None) != 0: - raise ValueError("pulse duration or pulse duration min max bounds need to be set for this model") - if all([fixed_pulse_duration, pulse_duration_min, pulse_duration_max]): - raise ValueError("Either pulse duration or pulse duration min max bounds need to be set for this model") + if isinstance(model, DingModelPulseWidthFrequency | DingModelPulseWidthFrequencyWithFatigue): + if pulse_width["fixed"] is None and [pulse_width["min"], pulse_width["max"]].count(None) != 0: + raise ValueError("pulse width or pulse width min max bounds need to be set for this model") + if all([pulse_width["fixed"], pulse_width["min"], pulse_width["max"]]): + raise ValueError("Either pulse width or pulse width min max bounds need to be set for this model") - minimum_pulse_duration = ( + minimum_pulse_width = ( 0 if model.pd0 is None else model.pd0 ) # Set it to 0 if used for the identification process - if fixed_pulse_duration is not None: - if isinstance(fixed_pulse_duration, int | float): - if fixed_pulse_duration < minimum_pulse_duration: + if pulse_width["fixed"]: + if isinstance(pulse_width["fixed"], int | float): + if pulse_width["fixed"] < minimum_pulse_width: raise ValueError( - f"The pulse duration set ({fixed_pulse_duration})" - f" is lower than minimum duration required." - f" Set a value above {minimum_pulse_duration} seconds " + f"The pulse width set ({pulse_width['fixed']})" + f" is lower than minimum width required." + f" Set a value above {minimum_pulse_width} seconds " ) - elif isinstance(fixed_pulse_duration, list): - if not all(isinstance(x, int | float) for x in fixed_pulse_duration): - raise TypeError("pulse_duration must be int or float type") - if not all(x >= minimum_pulse_duration for x in fixed_pulse_duration): + elif isinstance(pulse_width["fixed"], list): + if not all(isinstance(x, int | float) for x in pulse_width["fixed"]): + raise TypeError("pulse_width must be int or float type") + if not all(x >= minimum_pulse_width for x in pulse_width["fixed"]): raise ValueError( - f"The pulse duration set ({fixed_pulse_duration})" + f"The pulse width set ({pulse_width['fixed']})" f" is lower than minimum duration required." - f" Set a value above {minimum_pulse_duration} seconds " + f" Set a value above {minimum_pulse_width} seconds " ) else: - raise TypeError("Wrong pulse_duration type, only int or float accepted") - - elif pulse_duration_min is not None and pulse_duration_max is not None: - if not isinstance(pulse_duration_min, int | float) or not isinstance(pulse_duration_max, int | float): - raise TypeError("pulse_duration_min and pulse_duration_max must be int or float type") - if pulse_duration_max < pulse_duration_min: - raise ValueError("The set minimum pulse duration is higher than maximum pulse duration.") - if pulse_duration_min < minimum_pulse_duration: + raise TypeError("Wrong pulse_width type, only int or float accepted") + + elif pulse_width["min"] and pulse_width["max"]: + if not isinstance(pulse_width["min"], int | float) or not isinstance(pulse_width["max"], int | float): + raise TypeError("min and max pulse width must be int or float type") + if pulse_width["max"] < pulse_width["min"]: + raise ValueError("The set minimum pulse width is higher than maximum pulse width.") + if pulse_width["min"] < minimum_pulse_width: raise ValueError( - f"The pulse duration set ({pulse_duration_min})" - f" is lower than minimum duration required." - f" Set a value above {minimum_pulse_duration} seconds " + f"The pulse width set ({pulse_width['min']})" + f" is lower than minimum width required." + f" Set a value above {minimum_pulse_width} seconds " ) - if not isinstance(pulse_duration_bimapping, None | bool): - raise NotImplementedError("If added, pulse duration parameter mapping must be a bool type") + if not isinstance(pulse_width["bimapping"], None | bool): + raise NotImplementedError("If added, pulse width parameter mapping must be a bool type") - if isinstance(model, DingModelIntensityFrequency | DingModelIntensityFrequencyWithFatigue): - if fixed_pulse_intensity is None and [pulse_intensity_min, pulse_intensity_max].count(None) != 0: + if isinstance(model, DingModelPulseIntensityFrequency | DingModelPulseIntensityFrequencyWithFatigue): + if pulse_intensity["fixed"] is None and [pulse_intensity["min"], pulse_intensity["max"]].count(None) != 0: raise ValueError("Pulse intensity or pulse intensity min max bounds need to be set for this model") - if all([fixed_pulse_intensity, pulse_intensity_min, pulse_intensity_max]): + if all( + [ + pulse_intensity["fixed"], + pulse_intensity["min"], + pulse_intensity["max"], + ] + ): raise ValueError( "Either pulse intensity or pulse intensity min max bounds need to be set for this model" ) @@ -422,45 +398,52 @@ def _sanity_check( 0 if None in check_for_none_type else model.min_pulse_intensity() ) # Set it to 0 if used for the identification process - if fixed_pulse_intensity is not None: - if isinstance(fixed_pulse_intensity, int | float): - if fixed_pulse_intensity < minimum_pulse_intensity: + if pulse_intensity["fixed"]: + if isinstance(pulse_intensity["fixed"], int | float): + if pulse_intensity["fixed"] < minimum_pulse_intensity: raise ValueError( - f"The pulse intensity set ({fixed_pulse_intensity})" + f"The pulse intensity set ({pulse_intensity['fixed']})" f" is lower than minimum intensity required." f" Set a value above {minimum_pulse_intensity} mA " ) - elif isinstance(fixed_pulse_intensity, list): - if not all(isinstance(x, int | float) for x in fixed_pulse_intensity): + elif isinstance(pulse_intensity["fixed"], list): + if not all(isinstance(x, int | float) for x in pulse_intensity["fixed"]): raise TypeError("pulse_intensity must be int or float type") - if not all(x >= minimum_pulse_intensity for x in fixed_pulse_intensity): + if not all(x >= minimum_pulse_intensity for x in pulse_intensity["fixed"]): raise ValueError( - f"The pulse intensity set ({fixed_pulse_intensity})" + f"The pulse intensity set ({pulse_intensity['fixed']})" f" is lower than minimum intensity required." f" Set a value above {minimum_pulse_intensity} seconds " ) else: raise TypeError("pulse_intensity must be int or float type") - elif pulse_intensity_min is not None and pulse_intensity_max is not None: - if not isinstance(pulse_intensity_min, int | float) or not isinstance(pulse_intensity_max, int | float): + elif pulse_intensity["min"] and pulse_intensity["max"]: + if not isinstance(pulse_intensity["min"], int | float) or not isinstance( + pulse_intensity["max"], int | float + ): raise TypeError("pulse_intensity_min and pulse_intensity_max must be int or float type") - if pulse_intensity_max < pulse_intensity_min: + if pulse_intensity["max"] < pulse_intensity["min"]: raise ValueError("The set minimum pulse intensity is higher than maximum pulse intensity.") - if pulse_intensity_min < minimum_pulse_intensity: + if pulse_intensity["min"] < minimum_pulse_intensity: raise ValueError( - f"The pulse intensity set ({pulse_intensity_min})" + f"The pulse intensity set ({pulse_intensity['min']})" f" is lower than minimum intensity required." f" Set a value above {minimum_pulse_intensity} mA " ) - if not isinstance(pulse_intensity_bimapping, None | bool): + if not isinstance(pulse_intensity["bimapping"], None | bool): raise NotImplementedError("If added, pulse intensity parameter mapping must be a bool type") - if force_tracking is not None: - if isinstance(force_tracking, list): - if isinstance(force_tracking[0], np.ndarray) and isinstance(force_tracking[1], np.ndarray): - if len(force_tracking[0]) != len(force_tracking[1]) or len(force_tracking) != 2: + if objective["force_tracking"] is not None: + if isinstance(objective["force_tracking"], list): + if isinstance(objective["force_tracking"][0], np.ndarray) and isinstance( + objective["force_tracking"][1], np.ndarray + ): + if ( + len(objective["force_tracking"][0]) != len(objective["force_tracking"][1]) + or len(objective["force_tracking"]) != 2 + ): raise ValueError( "force_tracking time and force argument must be same length and force_tracking " "list size 2" @@ -470,17 +453,20 @@ def _sanity_check( else: raise TypeError("force_tracking must be list type") - if end_node_tracking: - if not isinstance(end_node_tracking, int | float): + if objective["end_node_tracking"] is not None: + if not isinstance(objective["end_node_tracking"], int | float): raise TypeError("end_node_tracking must be int or float type") - if custom_objective: - if not isinstance(custom_objective, ObjectiveList): + if objective["custom"] is not None: + if not isinstance(objective["custom"], ObjectiveList): raise TypeError("custom_objective must be a ObjectiveList type") - if not all(isinstance(x, Objective) for x in custom_objective[0]): + if not all(isinstance(x, Objective) for x in objective["custom"][0]): raise TypeError("All elements in ObjectiveList must be an Objective type") - if not isinstance(ode_solver, (OdeSolver.RK1, OdeSolver.RK2, OdeSolver.RK4, OdeSolver.COLLOCATION)): + if not isinstance( + ode_solver, + (OdeSolver.RK1, OdeSolver.RK2, OdeSolver.RK4, OdeSolver.COLLOCATION), + ): raise TypeError("ode_solver must be a OdeSolver type") if not isinstance(use_sx, bool): @@ -489,201 +475,203 @@ def _sanity_check( if not isinstance(n_threads, int): raise TypeError("n_thread must be a int type") - @staticmethod - def _sanity_check_frequency(n_stim, final_time, frequency, round_down): - if [n_stim, final_time, frequency].count(None) == 2: - raise ValueError("At least two variable must be set from n_stim, final_time or frequency") - - if n_stim and final_time and frequency: - if n_stim != final_time * frequency: - raise ValueError( - "Can not satisfy n_stim equal to final_time * frequency with the given parameters." - "Consider setting only two of the three parameters" - ) - - if round_down: - if not isinstance(round_down, bool): - raise TypeError("round_down must be bool type") - @staticmethod def _build_fourier_coefficient(force_tracking): return FourierSeries().compute_real_fourier_coeffs(force_tracking[0], force_tracking[1], 50) - @staticmethod - def _build_phase_time(final_time, n_stim, pulse_mode, time_min, time_max): - final_time_phase = None - if time_min is None and time_max is None: - if pulse_mode == "single": - step = final_time / n_stim - final_time_phase = (step,) - for i in range(n_stim - 1): - final_time_phase = final_time_phase + (step,) - else: - final_time_phase = [final_time / n_stim] * n_stim - - return final_time_phase - @staticmethod def _build_parameters( model, - n_stim, - time_min, - time_max, - time_bimapping, - fixed_pulse_duration, - pulse_duration_min, - pulse_duration_max, - pulse_duration_bimapping, - fixed_pulse_intensity, - pulse_intensity_min, - pulse_intensity_max, - pulse_intensity_bimapping, + stim_time, + pulse_event, + pulse_width, + pulse_intensity, use_sx, ): parameters = ParameterList(use_sx=use_sx) parameters_bounds = BoundsList() parameters_init = InitialGuessList() parameter_objectives = ParameterObjectiveList() - constraints = ConstraintList() - if time_min: - parameters.add( - name="pulse_apparition_time", - function=DingModelFrequency.set_pulse_apparition_time, - size=n_stim, - scaling=VariableScaling("pulse_apparition_time", [1] * n_stim), - ) + n_stim = len(stim_time) + parameters.add( + name="pulse_apparition_time", + function=DingModelFrequency.set_pulse_apparition_time, + size=len(stim_time), + scaling=VariableScaling("pulse_apparition_time", [1] * n_stim), + ) - if time_min and time_max: - time_min_list = [time_min * n for n in range(n_stim)] - time_max_list = [time_max * n for n in range(n_stim)] - else: - time_min_list = [0] * n_stim - time_max_list = [100] * n_stim - parameters_bounds.add( - "pulse_apparition_time", - min_bound=np.array(time_min_list), - max_bound=np.array(time_max_list), - interpolation=InterpolationType.CONSTANT, + if pulse_event["min"] and pulse_event["max"]: + time_min_list = np.array([0] + list(np.cumsum([pulse_event["min"]] * (n_stim - 1)))) + time_max_list = np.array([0] + list(np.cumsum([pulse_event["max"]] * (n_stim - 1)))) + parameters_init["pulse_apparition_time"] = np.array( + [(time_max_list[i] + time_min_list[i]) / 2 for i in range(n_stim)] ) + else: + time_min_list = stim_time + time_max_list = stim_time + parameters_init["pulse_apparition_time"] = np.array(stim_time) + + parameters_bounds.add( + "pulse_apparition_time", + min_bound=time_min_list, + max_bound=time_max_list, + interpolation=InterpolationType.CONSTANT, + ) - parameters_init["pulse_apparition_time"] = np.array([0] * n_stim) - - for i in range(n_stim): - constraints.add(CustomConstraint.pulse_time_apparition_as_phase, node=Node.START, phase=i, target=0) + if pulse_event["bimapping"] and pulse_event["min"] and pulse_event["max"]: + raise NotImplementedError("Bimapping is not yet implemented for pulse event") - if time_bimapping and time_min and time_max: - for i in range(n_stim): - constraints.add(CustomConstraint.equal_to_first_pulse_interval_time, node=Node.START, target=0, phase=i) + if isinstance(model, DingModelPulseWidthFrequency): + if pulse_width["bimapping"]: + n_stim = 1 - if isinstance(model, DingModelPulseDurationFrequency): - if fixed_pulse_duration: + if pulse_width["fixed"]: parameters.add( - name="pulse_duration", - function=DingModelPulseDurationFrequency.set_impulse_duration, + name="pulse_width", + function=DingModelPulseWidthFrequency.set_impulse_width, size=n_stim, - scaling=VariableScaling("pulse_duration", [1] * n_stim), + scaling=VariableScaling("pulse_width", [1] * n_stim), ) - if isinstance(fixed_pulse_duration, list): + if isinstance(pulse_width["fixed"], list): parameters_bounds.add( - "pulse_duration", - min_bound=np.array(fixed_pulse_duration), - max_bound=np.array(fixed_pulse_duration), + "pulse_width", + min_bound=np.array(pulse_width["fixed"]), + max_bound=np.array(pulse_width["fixed"]), interpolation=InterpolationType.CONSTANT, ) - parameters_init.add(key="pulse_duration", initial_guess=np.array(fixed_pulse_duration)) + parameters_init.add( + key="pulse_width", + initial_guess=np.array(pulse_width["fixed"]), + ) else: parameters_bounds.add( - "pulse_duration", - min_bound=np.array([fixed_pulse_duration] * n_stim), - max_bound=np.array([fixed_pulse_duration] * n_stim), + "pulse_width", + min_bound=np.array([pulse_width["fixed"]] * n_stim), + max_bound=np.array([pulse_width["fixed"]] * n_stim), interpolation=InterpolationType.CONSTANT, ) - parameters_init["pulse_duration"] = np.array([fixed_pulse_duration] * n_stim) + parameters_init["pulse_width"] = np.array([pulse_width["fixed"]] * n_stim) - elif pulse_duration_min is not None and pulse_duration_max is not None: + elif pulse_width["min"] and pulse_width["max"]: parameters_bounds.add( - "pulse_duration", - min_bound=[pulse_duration_min], - max_bound=[pulse_duration_max], + "pulse_width", + min_bound=pulse_width["min"], + max_bound=pulse_width["max"], interpolation=InterpolationType.CONSTANT, ) - parameters_init["pulse_duration"] = np.array([0] * n_stim) + parameters_init["pulse_width"] = np.array( + [(pulse_width["min"] + pulse_width["max"]) / 2 for _ in range(n_stim)] + ) parameters.add( - name="pulse_duration", - function=DingModelPulseDurationFrequency.set_impulse_duration, + name="pulse_width", + function=DingModelPulseWidthFrequency.set_impulse_width, size=n_stim, - scaling=VariableScaling("pulse_duration", [1] * n_stim), + scaling=VariableScaling("pulse_width", [1] * n_stim), ) - if pulse_duration_bimapping is True: - for i in range(1, n_stim): - constraints.add(CustomConstraint.equal_to_first_pulse_duration, node=Node.START, target=0, phase=i) + if isinstance(model, DingModelPulseIntensityFrequency): + if pulse_intensity["bimapping"]: + n_stim = 1 - if isinstance(model, DingModelIntensityFrequency): - if fixed_pulse_intensity: + if pulse_intensity["fixed"]: parameters.add( name="pulse_intensity", - function=DingModelIntensityFrequency.set_impulse_intensity, + function=DingModelPulseIntensityFrequency.set_impulse_intensity, size=n_stim, scaling=VariableScaling("pulse_intensity", [1] * n_stim), ) - if isinstance(fixed_pulse_intensity, list): + if isinstance(pulse_intensity["fixed"], list): parameters_bounds.add( "pulse_intensity", - min_bound=np.array(fixed_pulse_intensity), - max_bound=np.array(fixed_pulse_intensity), + min_bound=np.array(pulse_intensity["fixed"]), + max_bound=np.array(pulse_intensity["fixed"]), interpolation=InterpolationType.CONSTANT, ) - parameters_init.add(key="pulse_intensity", initial_guess=np.array(fixed_pulse_intensity)) + parameters_init.add( + key="pulse_intensity", + initial_guess=np.array(pulse_intensity["fixed"]), + ) else: parameters_bounds.add( "pulse_intensity", - min_bound=np.array([fixed_pulse_intensity] * n_stim), - max_bound=np.array([fixed_pulse_intensity] * n_stim), + min_bound=np.array([pulse_intensity["fixed"]] * n_stim), + max_bound=np.array([pulse_intensity["fixed"]] * n_stim), interpolation=InterpolationType.CONSTANT, ) - parameters_init["pulse_intensity"] = np.array([fixed_pulse_intensity] * n_stim) + parameters_init["pulse_intensity"] = np.array([pulse_intensity["fixed"]] * n_stim) - elif pulse_intensity_min is not None and pulse_intensity_max is not None: + elif pulse_intensity["min"] and pulse_intensity["max"]: parameters_bounds.add( "pulse_intensity", - min_bound=[pulse_intensity_min], - max_bound=[pulse_intensity_max], + min_bound=[pulse_intensity["min"]], + max_bound=[pulse_intensity["max"]], interpolation=InterpolationType.CONSTANT, ) - intensity_avg = (pulse_intensity_min + pulse_intensity_max) / 2 + intensity_avg = (pulse_intensity["min"] + pulse_intensity["max"]) / 2 parameters_init["pulse_intensity"] = np.array([intensity_avg] * n_stim) parameters.add( name="pulse_intensity", - function=DingModelIntensityFrequency.set_impulse_intensity, + function=DingModelPulseIntensityFrequency.set_impulse_intensity, size=n_stim, scaling=VariableScaling("pulse_intensity", [1] * n_stim), ) - if pulse_intensity_bimapping is True: - for i in range(1, n_stim): - constraints.add(CustomConstraint.equal_to_first_pulse_intensity, node=Node.START, target=0, phase=i) - - return parameters, parameters_bounds, parameters_init, parameter_objectives, constraints + return (parameters, parameters_bounds, parameters_init, parameter_objectives) @staticmethod - def _declare_dynamics(models, n_stim): - dynamics = DynamicsList() - for i in range(n_stim): - dynamics.add( - models[i].declare_ding_variables, - dynamic_function=models[i].dynamics, - expand_dynamics=True, - expand_continuity=False, - phase=i, - phase_dynamics=PhaseDynamics.SHARED_DURING_THE_PHASE, + def _build_constraints(model, n_shooting, final_time, stim_time, control_type): + constraints = ConstraintList() + + time_vector = np.linspace(0, final_time, n_shooting + 1) + stim_at_node = [np.where(stim_time[i] <= time_vector)[0][0] for i in range(len(stim_time))] + additional_nodes = 1 if control_type == ControlType.LINEAR_CONTINUOUS else 0 + if model._sum_stim_truncation: + max_stim_to_keep = model._sum_stim_truncation + else: + max_stim_to_keep = 10000000 + + index_sup = 0 + index_inf = 0 + + for i in range(n_shooting + additional_nodes): + if i in stim_at_node: + index_sup += 1 + if index_sup >= max_stim_to_keep: + index_inf = index_sup - max_stim_to_keep + + constraints.add( + CustomConstraint.cn_sum, + node=i, + stim_time=stim_time[index_inf:index_sup], ) + if isinstance(model, DingModelPulseWidthFrequency): + index = 0 + for i in range(n_shooting + additional_nodes): + if i in stim_at_node and i != 0: + index += 1 + constraints.add( + CustomConstraint.a_calculation, + node=i, + last_stim_index=index, + ) + + return constraints + + @staticmethod + def _declare_dynamics(model): + dynamics = DynamicsList() + dynamics.add( + model.declare_ding_variables, + dynamic_function=model.dynamics, + expand_dynamics=True, + phase_dynamics=PhaseDynamics.SHARED_DURING_THE_PHASE, + ) return dynamics @staticmethod - def _set_bounds(model, n_stim): + def _set_bounds(model): # ---- STATE BOUNDS REPRESENTATION ---- # # # |‾‾‾‾‾‾‾‾‾‾x_max_middle‾‾‾‾‾‾‾‾‾‾‾‾x_max_end‾ @@ -714,98 +702,73 @@ def _set_bounds(model, n_stim): starting_bounds_min = np.concatenate((starting_bounds, min_bounds, min_bounds), axis=1) starting_bounds_max = np.concatenate((starting_bounds, max_bounds, max_bounds), axis=1) - middle_bound_min = np.concatenate((min_bounds, min_bounds, min_bounds), axis=1) - middle_bound_max = np.concatenate((max_bounds, max_bounds, max_bounds), axis=1) - - for i in range(n_stim): - for j in range(len(variable_bound_list)): - if i == 0: - x_bounds.add( - variable_bound_list[j], - min_bound=np.array([starting_bounds_min[j]]), - max_bound=np.array([starting_bounds_max[j]]), - phase=i, - interpolation=InterpolationType.CONSTANT_WITH_FIRST_AND_LAST_DIFFERENT, - ) - else: - x_bounds.add( - variable_bound_list[j], - min_bound=np.array([middle_bound_min[j]]), - max_bound=np.array([middle_bound_max[j]]), - phase=i, - interpolation=InterpolationType.CONSTANT_WITH_FIRST_AND_LAST_DIFFERENT, - ) + + for j in range(len(variable_bound_list)): + x_bounds.add( + variable_bound_list[j], + min_bound=np.array([starting_bounds_min[j]]), + max_bound=np.array([starting_bounds_max[j]]), + interpolation=InterpolationType.CONSTANT_WITH_FIRST_AND_LAST_DIFFERENT, + ) x_init = InitialGuessList() - for i in range(n_stim): - for j in range(len(variable_bound_list)): - x_init.add(variable_bound_list[j], model.standard_rest_values()[j], phase=i) + for j in range(len(variable_bound_list)): + x_init.add(variable_bound_list[j], model.standard_rest_values()[j]) return x_bounds, x_init @staticmethod - def _set_objective( - n_stim, n_shooting, force_fourier_coefficient, end_node_tracking, custom_objective, time_min, time_max - ): + def _set_u_bounds(model): + # Controls bounds + u_bounds = BoundsList() + + # Controls initial guess + u_init = InitialGuessList() + u_init.add(key="Cn_sum", initial_guess=[0], phase=0) + if isinstance(model, DingModelPulseWidthFrequency): + u_init.add(key="A_calculation", initial_guess=[0], phase=0) + + return u_bounds, u_init + + @staticmethod + def _set_objective(n_shooting, objective): # Creates the objective for our problem objective_functions = ObjectiveList() - if custom_objective: - for i in range(len(custom_objective)): - objective_functions.add(custom_objective[0][i]) - - if force_fourier_coefficient is not None: - for phase in range(n_stim): - for i in range(n_shooting[phase]): - objective_functions.add( - CustomObjective.track_state_from_time, - custom_type=ObjectiveFcn.Mayer, - node=i, - fourier_coeff=force_fourier_coefficient, - key="F", - quadratic=True, - weight=1, - phase=phase, - ) + if objective["custom"]: + for i in range(len(objective["custom"])): + objective_functions.add(objective["custom"][0][i]) + + if objective["force_tracking"]: + force_fourier_coefficient = ( + None + if objective["force_tracking"] is None + else OcpFes._build_fourier_coefficient(objective["force_tracking"]) + ) + force_to_track = FourierSeries().fit_func_by_fourier_series_with_real_coeffs( + np.linspace(0, 1, n_shooting + 1), + force_fourier_coefficient, + )[np.newaxis, :] + + objective_functions.add( + ObjectiveFcn.Lagrange.TRACK_STATE, + key="F", + weight=100, + target=force_to_track, + node=Node.ALL, + quadratic=True, + ) - if end_node_tracking: - if isinstance(end_node_tracking, int | float): - objective_functions.add( - ObjectiveFcn.Mayer.MINIMIZE_STATE, - node=Node.END, - key="F", - quadratic=True, - weight=1, - target=end_node_tracking, - phase=n_stim - 1, - ) + if objective["end_node_tracking"]: + objective_functions.add( + ObjectiveFcn.Mayer.MINIMIZE_STATE, + node=Node.END, + key="F", + quadratic=True, + weight=1, + target=objective["end_node_tracking"], + ) - if time_min and time_max: - for i in range(n_stim): - objective_functions.add( - ObjectiveFcn.Mayer.MINIMIZE_TIME, - weight=0.001 / n_shooting[i], - min_bound=time_min, - max_bound=time_max, - quadratic=True, - phase=i, - ) + # if time_min and time_max: + # objective_functions.add(ObjectiveFcn.Mayer.MINIMIZE_TIME, weight=0.001, quadratic=True) return objective_functions - - @staticmethod - def _build_phase_parameter(n_stim, final_time, frequency=None, pulse_mode="single", round_down=False): - pulse_mode_multiplier = 1 if pulse_mode == "single" else 2 if pulse_mode == "doublet" else 3 - if n_stim and frequency: - final_time = n_stim / frequency / pulse_mode_multiplier - - if final_time and frequency: - n_stim = final_time * frequency * pulse_mode_multiplier - if round_down or n_stim.is_integer(): - n_stim = int(n_stim) - else: - raise ValueError( - "The number of stimulation needs to be integer within the final time t, set round down" - "to True or set final_time * frequency to make the result a integer." - ) - - return n_stim, final_time diff --git a/cocofest/optimization/fes_ocp_dynamics.py b/cocofest/optimization/fes_ocp_dynamics.py index 3060897..556edc7 100644 --- a/cocofest/optimization/fes_ocp_dynamics.py +++ b/cocofest/optimization/fes_ocp_dynamics.py @@ -5,6 +5,7 @@ BoundsList, ConstraintList, ControlType, + ConstraintFcn, DynamicsList, InitialGuessList, InterpolationType, @@ -20,40 +21,145 @@ VariableScaling, ) -from ..custom_constraints import CustomConstraint from ..custom_objectives import CustomObjective from ..dynamics.inverse_kinematics_and_dynamics import get_circle_coord -from ..dynamics.warm_start import get_initial_guess +from ..dynamics.initial_guess_warm_start import get_initial_guess from ..models.ding2003 import DingModelFrequency -from ..models.ding2007 import DingModelPulseDurationFrequency +from ..models.ding2007 import DingModelPulseWidthFrequency from ..models.dynamical_model import FesMskModel -from ..models.fes_model import FesModel -from ..models.hmed2018 import DingModelIntensityFrequency +from ..models.hmed2018 import DingModelPulseIntensityFrequency from ..optimization.fes_ocp import OcpFes +from ..fourier_approx import FourierSeries +from ..custom_constraints import CustomConstraint class OcpFesMsk: + @staticmethod + def _prepare_optimization_problem(input_dict: dict) -> dict: + + (pulse_event, pulse_width, pulse_intensity, objective) = OcpFes._fill_dict( + input_dict["pulse_event"], + input_dict["pulse_width"], + input_dict["pulse_intensity"], + input_dict["objective"], + ) + + ( + pulse_width, + pulse_intensity, + objective, + msk_info, + ) = OcpFesMsk._fill_msk_dict(pulse_width, pulse_intensity, objective, input_dict["msk_info"]) + + OcpFes._sanity_check( + model=input_dict["model"], + n_shooting=input_dict["n_shooting"], + final_time=input_dict["final_time"], + pulse_event=pulse_event, + pulse_width=pulse_width, + pulse_intensity=pulse_intensity, + objective=objective, + use_sx=input_dict["use_sx"], + ode_solver=input_dict["ode_solver"], + n_threads=input_dict["n_threads"], + ) + + OcpFesMsk._sanity_check_msk_inputs( + model=input_dict["model"], + msk_info=msk_info, + objective=objective, + ) + + ( + parameters, + parameters_bounds, + parameters_init, + parameter_objectives, + ) = OcpFesMsk._build_parameters( + model=input_dict["model"], + stim_time=input_dict["stim_time"], + pulse_event=pulse_event, + pulse_width=pulse_width, + pulse_intensity=pulse_intensity, + use_sx=input_dict["use_sx"], + ) + + constraints = OcpFesMsk._build_constraints( + input_dict["model"], + input_dict["n_shooting"], + input_dict["final_time"], + input_dict["stim_time"], + input_dict["control_type"], + msk_info["custom_constraint"], + ) + + dynamics = OcpFesMsk._declare_dynamics(input_dict["model"]) + initial_state = ( + get_initial_guess( + input_dict["model"].path, + input_dict["final_time"], + input_dict["n_shooting"], + objective, + n_threads=input_dict["n_threads"], + ) + if input_dict["initial_guess_warm_start"] + else None + ) + + x_bounds, x_init = OcpFesMsk._set_bounds( + input_dict["model"], + msk_info, + initial_state, + input_dict["n_cycles_simultaneous"] if "n_cycles_simultaneous" in input_dict.keys() else 1, + ) + u_bounds, u_init = OcpFesMsk._set_u_bounds(input_dict["model"], msk_info["with_residual_torque"]) + + muscle_force_key = [ + "F_" + input_dict["model"].muscles_dynamics_model[i].muscle_name + for i in range(len(input_dict["model"].muscles_dynamics_model)) + ] + objective_functions = OcpFesMsk._set_objective( + input_dict["n_shooting"], + muscle_force_key, + objective, + input_dict["n_cycles_simultaneous"] if "n_cycles_simultaneous" in input_dict.keys() else 1, + ) + + optimization_dict = { + "model": input_dict["model"], + "dynamics": dynamics, + "n_shooting": input_dict["n_shooting"], + "final_time": input_dict["final_time"], + "objective_functions": objective_functions, + "x_init": x_init, + "x_bounds": x_bounds, + "u_init": u_init, + "u_bounds": u_bounds, + "constraints": constraints, + "parameters": parameters, + "parameters_bounds": parameters_bounds, + "parameters_init": parameters_init, + "parameter_objectives": parameter_objectives, + "use_sx": input_dict["use_sx"], + "ode_solver": input_dict["ode_solver"], + "n_threads": input_dict["n_threads"], + "control_type": input_dict["control_type"], + } + + return optimization_dict + @staticmethod def prepare_ocp( - biorbd_model_path: str, - bound_type: str = None, - bound_data: list = None, - fes_muscle_models: list[FesModel] = None, - n_stim: int = None, - n_shooting: int = None, + model: FesMskModel = None, + stim_time: list = None, final_time: int | float = None, pulse_event: dict = None, - pulse_duration: dict = None, + pulse_width: dict = None, pulse_intensity: dict = None, objective: dict = None, - custom_constraint: ConstraintList = None, - with_residual_torque: bool = False, - activate_force_length_relationship: bool = False, - activate_force_velocity_relationship: bool = False, - minimize_muscle_fatigue: bool = False, - minimize_muscle_force: bool = False, + msk_info: dict = None, use_sx: bool = True, - warm_start: bool = False, + initial_guess_warm_start: bool = False, ode_solver: OdeSolverBase = OdeSolver.RK4(n_integration_steps=1), control_type: ControlType = ControlType.CONSTANT, n_threads: int = 1, @@ -63,24 +169,16 @@ def prepare_ocp( Parameters ---------- - biorbd_model_path : str - The path to the bioMod file. - bound_type : str - The type of bound to use (start, end, start_end). - bound_data : list - The data to use for the bound. - fes_muscle_models : list[FesModel] - The FES model type used for the OCP. - n_stim : int - Number of stimulations that will occur during the OCP, also referred to as phases. - n_shooting : int - Number of shooting points for each individual phase. + model : FesModel + The FES model to use. + stim_time : list + The stimulation times. final_time : int | float The final time of the OCP. pulse_event : dict Dictionary containing parameters related to the appearance of the pulse. It should contain the following keys: "min", "max", "bimapping", "frequency", "round_down", "pulse_mode". - pulse_duration : dict + pulse_width : dict Dictionary containing parameters related to the duration of the pulse. It should contain the following keys: "fixed", "min", "max", "bimapping", "similar_for_all_muscles". Optional if not using the Ding2007 models @@ -90,21 +188,11 @@ def prepare_ocp( Optional if not using the Hmed2018 models objective : dict Dictionary containing parameters related to the objective of the optimization. - custom_constraint : ConstraintList, - Custom constraints for the OCP. - with_residual_torque : bool - If residual torque is used. - activate_force_length_relationship : bool - If the force length relationship is used. - activate_force_velocity_relationship : bool - If the force velocity relationship is used. - minimize_muscle_fatigue : bool - Minimize the muscle fatigue. - minimize_muscle_force : bool - Minimize the muscle force. + msk_info : dict + Dictionary containing parameters related to the musculoskeletal model. use_sx : bool The nature of the CasADi variables. MX are used if False. - warm_start : bool + initial_guess_warm_start : bool If a warm start is run to get the problem initial guesses. ode_solver : OdeSolverBase The ODE solver to use. @@ -119,402 +207,318 @@ def prepare_ocp( The prepared Optimal Control Program. """ - (pulse_event, pulse_duration, pulse_intensity, objective) = OcpFes._fill_dict( - pulse_event, pulse_duration, pulse_intensity, objective - ) - - time_min = pulse_event["min"] - time_max = pulse_event["max"] - time_bimapping = pulse_event["bimapping"] - frequency = pulse_event["frequency"] - round_down = pulse_event["round_down"] - pulse_mode = pulse_event["pulse_mode"] - - fixed_pulse_duration = pulse_duration["fixed"] - pulse_duration_min = pulse_duration["min"] - pulse_duration_max = pulse_duration["max"] - pulse_duration_bimapping = pulse_duration["bimapping"] - key_in_dict = "similar_for_all_muscles" in pulse_duration - pulse_duration_similar_for_all_muscles = pulse_duration["similar_for_all_muscles"] if key_in_dict else False - - fixed_pulse_intensity = pulse_intensity["fixed"] - pulse_intensity_min = pulse_intensity["min"] - pulse_intensity_max = pulse_intensity["max"] - pulse_intensity_bimapping = pulse_intensity["bimapping"] - key_in_dict = "similar_for_all_muscles" in pulse_intensity - pulse_intensity_similar_for_all_muscles = pulse_intensity["similar_for_all_muscles"] if key_in_dict else False - - force_tracking = objective["force_tracking"] - end_node_tracking = objective["end_node_tracking"] - cycling_objective = objective["cycling"] - custom_objective = objective["custom"] - key_in_dict = "q_tracking" in objective - q_tracking = objective["q_tracking"] if key_in_dict else None - - OcpFes._sanity_check( - model=fes_muscle_models[0], - n_stim=n_stim, - n_shooting=n_shooting, - final_time=final_time, - pulse_mode=pulse_mode, - frequency=frequency, - time_min=time_min, - time_max=time_max, - time_bimapping=time_bimapping, - fixed_pulse_duration=fixed_pulse_duration, - pulse_duration_min=pulse_duration_min, - pulse_duration_max=pulse_duration_max, - pulse_duration_bimapping=pulse_duration_bimapping, - fixed_pulse_intensity=fixed_pulse_intensity, - pulse_intensity_min=pulse_intensity_min, - pulse_intensity_max=pulse_intensity_max, - pulse_intensity_bimapping=pulse_intensity_bimapping, - custom_objective=custom_objective, - use_sx=use_sx, - ode_solver=ode_solver, - n_threads=n_threads, - ) - - OcpFesMsk._sanity_check_fes_models_inputs( - biorbd_model_path=biorbd_model_path, - bound_type=bound_type, - bound_data=bound_data, - fes_muscle_models=fes_muscle_models, - force_tracking=force_tracking, - end_node_tracking=end_node_tracking, - cycling_objective=cycling_objective, - q_tracking=q_tracking, - with_residual_torque=with_residual_torque, - activate_force_length_relationship=activate_force_length_relationship, - activate_force_velocity_relationship=activate_force_velocity_relationship, - minimize_muscle_fatigue=minimize_muscle_fatigue, - minimize_muscle_force=minimize_muscle_force, - ) - - OcpFes._sanity_check_frequency(n_stim=n_stim, final_time=final_time, frequency=frequency, round_down=round_down) - - OcpFesMsk._sanity_check_muscle_model(biorbd_model_path=biorbd_model_path, fes_muscle_models=fes_muscle_models) - - n_stim, final_time = OcpFes._build_phase_parameter( - n_stim=n_stim, final_time=final_time, frequency=frequency, pulse_mode=pulse_mode, round_down=round_down - ) - - force_fourier_coef = [] if force_tracking else None - if force_tracking: - for i in range(len(force_tracking[1])): - force_fourier_coef.append(OcpFes._build_fourier_coefficient([force_tracking[0], force_tracking[1][i]])) - - q_fourier_coef = [] if q_tracking else None - if q_tracking: - for i in range(len(q_tracking[1])): - q_fourier_coef.append(OcpFes._build_fourier_coefficient([q_tracking[0], q_tracking[1][i]])) - - n_shooting = [n_shooting] * n_stim - final_time_phase = OcpFes._build_phase_time( - final_time=final_time, - n_stim=n_stim, - pulse_mode=pulse_mode, - time_min=time_min, - time_max=time_max, - ) - ( - parameters, - parameters_bounds, - parameters_init, - parameter_objectives, - constraints, - ) = OcpFesMsk._build_parameters( - model=fes_muscle_models, - n_stim=n_stim, - time_min=time_min, - time_max=time_max, - time_bimapping=time_bimapping, - fixed_pulse_duration=fixed_pulse_duration, - pulse_duration_min=pulse_duration_min, - pulse_duration_max=pulse_duration_max, - pulse_duration_bimapping=pulse_duration_bimapping, - pulse_duration_similar_for_all_muscles=pulse_duration_similar_for_all_muscles, - fixed_pulse_intensity=fixed_pulse_intensity, - pulse_intensity_min=pulse_intensity_min, - pulse_intensity_max=pulse_intensity_max, - pulse_intensity_bimapping=pulse_intensity_bimapping, - pulse_intensity_similar_for_all_muscles=pulse_intensity_similar_for_all_muscles, - use_sx=use_sx, - ) - - constraints = OcpFesMsk._set_constraints(constraints, custom_constraint) - - if len(constraints) == 0 and len(parameters) == 0: - raise ValueError( - "This is not an optimal control problem," - " add parameter to optimize or use the IvpFes method to build your problem" - ) - - bio_models = [ - FesMskModel( - name=None, - biorbd_path=biorbd_model_path, - muscles_model=fes_muscle_models, - activate_force_length_relationship=activate_force_length_relationship, - activate_force_velocity_relationship=activate_force_velocity_relationship, - ) - for i in range(n_stim) - ] - - dynamics = OcpFesMsk._declare_dynamics(bio_models, n_stim) - initial_state = ( - get_initial_guess(biorbd_model_path, final_time, n_stim, n_shooting, objective) if warm_start else None - ) - - x_bounds, x_init = OcpFesMsk._set_bounds( - bio_models, - fes_muscle_models, - bound_type, - bound_data, - n_stim, - initial_state, - ) - u_bounds, u_init = OcpFesMsk._set_controls(bio_models, n_stim, with_residual_torque) - muscle_force_key = ["F_" + fes_muscle_models[i].muscle_name for i in range(len(fes_muscle_models))] - objective_functions = OcpFesMsk._set_objective( - n_stim, - n_shooting, - force_fourier_coef, - end_node_tracking, - cycling_objective, - custom_objective, - q_fourier_coef, - minimize_muscle_fatigue, - minimize_muscle_force, - muscle_force_key, - time_min, - time_max, - ) + input_dict = { + "model": model, + "stim_time": stim_time, + "n_shooting": OcpFes.prepare_n_shooting(stim_time, final_time), + "final_time": final_time, + "pulse_event": pulse_event, + "pulse_width": pulse_width, + "pulse_intensity": pulse_intensity, + "objective": objective, + "msk_info": msk_info, + "initial_guess_warm_start": initial_guess_warm_start, + "use_sx": use_sx, + "ode_solver": ode_solver, + "n_threads": n_threads, + "control_type": control_type, + } + + optimization_dict = OcpFesMsk._prepare_optimization_problem(input_dict) return OptimalControlProgram( - bio_model=bio_models, - dynamics=dynamics, - n_shooting=n_shooting, - phase_time=final_time_phase, - objective_functions=objective_functions, - x_init=x_init, - x_bounds=x_bounds, - u_init=u_init, - u_bounds=u_bounds, - constraints=constraints, - parameters=parameters, - parameter_bounds=parameters_bounds, - parameter_init=parameters_init, - parameter_objectives=parameter_objectives, + bio_model=[optimization_dict["model"]], + dynamics=optimization_dict["dynamics"], + n_shooting=optimization_dict["n_shooting"], + phase_time=optimization_dict["final_time"], + objective_functions=optimization_dict["objective_functions"], + x_init=optimization_dict["x_init"], + x_bounds=optimization_dict["x_bounds"], + u_init=optimization_dict["u_init"], + u_bounds=optimization_dict["u_bounds"], + constraints=optimization_dict["constraints"], + parameters=optimization_dict["parameters"], + parameter_bounds=optimization_dict["parameters_bounds"], + parameter_init=optimization_dict["parameters_init"], + parameter_objectives=optimization_dict["parameter_objectives"], control_type=control_type, - use_sx=use_sx, - ode_solver=ode_solver, - n_threads=n_threads, + use_sx=optimization_dict["use_sx"], + ode_solver=optimization_dict["ode_solver"], + n_threads=optimization_dict["n_threads"], ) @staticmethod - def _declare_dynamics(bio_models, n_stim): + def _fill_msk_dict(pulse_width, pulse_intensity, objective, msk_info): + + pulse_width = pulse_width if pulse_width else {} + default_pulse_width = { + "fixed": None, + "min": None, + "max": None, + "bimapping": False, + "same_for_all_muscles": False, + } + + pulse_intensity = pulse_intensity if pulse_intensity else {} + default_pulse_intensity = { + "fixed": None, + "min": None, + "max": None, + "bimapping": False, + "same_for_all_muscles": False, + } + + objective = objective if objective else {} + default_objective = { + "force_tracking": None, + "end_node_tracking": None, + "cycling_objective": None, + "custom": None, + "q_tracking": None, + "minimize_muscle_fatigue": False, + "minimize_muscle_force": False, + "minimize_residual_torque": False, + } + + msk_info = msk_info if msk_info else {} + default_msk_info = { + "bound_type": None, + "bound_data": None, + "with_residual_torque": False, + "custom_constraint": None, + } + + pulse_width = {**default_pulse_width, **pulse_width} + pulse_intensity = {**default_pulse_intensity, **pulse_intensity} + objective = {**default_objective, **objective} + msk_info = {**default_msk_info, **msk_info} + + return pulse_width, pulse_intensity, objective, msk_info + + @staticmethod + def _declare_dynamics(bio_models): dynamics = DynamicsList() - for i in range(n_stim): - dynamics.add( - bio_models[i].declare_model_variables, - dynamic_function=bio_models[i].muscle_dynamic, - expand_dynamics=True, - expand_continuity=False, - phase=i, - phase_dynamics=PhaseDynamics.SHARED_DURING_THE_PHASE, - ) + dynamics.add( + bio_models.declare_model_variables, + dynamic_function=bio_models.muscle_dynamic, + expand_dynamics=True, + expand_continuity=False, + phase=0, + phase_dynamics=PhaseDynamics.SHARED_DURING_THE_PHASE, + ) return dynamics @staticmethod def _build_parameters( - model, - n_stim, - time_min, - time_max, - time_bimapping, - fixed_pulse_duration, - pulse_duration_min, - pulse_duration_max, - pulse_duration_bimapping, - pulse_duration_similar_for_all_muscles, - fixed_pulse_intensity, - pulse_intensity_min, - pulse_intensity_max, - pulse_intensity_bimapping, - pulse_intensity_similar_for_all_muscles, - use_sx, + model: FesMskModel, + stim_time: list, + pulse_event: dict, + pulse_width: dict, + pulse_intensity: dict, + use_sx: bool = True, ): parameters = ParameterList(use_sx=use_sx) parameters_bounds = BoundsList() parameters_init = InitialGuessList() parameter_objectives = ParameterObjectiveList() - constraints = ConstraintList() - - if time_min: - parameters.add( - name="pulse_apparition_time", - function=DingModelFrequency.set_pulse_apparition_time, - size=n_stim, - scaling=VariableScaling("pulse_apparition_time", [1] * n_stim), - ) - - if time_min and time_max: - time_min_list = [time_min * n for n in range(n_stim)] - time_max_list = [time_max * n for n in range(n_stim)] - else: - time_min_list = [0] * n_stim - time_max_list = [100] * n_stim - parameters_bounds.add( - "pulse_apparition_time", - min_bound=np.array(time_min_list), - max_bound=np.array(time_max_list), - interpolation=InterpolationType.CONSTANT, - ) - parameters_init["pulse_apparition_time"] = np.array([0] * n_stim) + n_stim = len(stim_time) + parameters.add( + name="pulse_apparition_time", + function=DingModelFrequency.set_pulse_apparition_time, + size=n_stim, + scaling=VariableScaling("pulse_apparition_time", [1] * n_stim), + ) - for i in range(n_stim): - constraints.add(CustomConstraint.pulse_time_apparition_as_phase, node=Node.START, phase=i, target=0) + parameters_bounds.add( + "pulse_apparition_time", + min_bound=np.array(stim_time), + max_bound=np.array(stim_time), + interpolation=InterpolationType.CONSTANT, + ) - if time_bimapping and time_min and time_max: - for i in range(n_stim): - constraints.add(CustomConstraint.equal_to_first_pulse_interval_time, node=Node.START, target=0, phase=i) + parameters_init["pulse_apparition_time"] = np.array(stim_time) - for i in range(len(model)): - if isinstance(model[i], DingModelPulseDurationFrequency): + for i in range(len(model.muscles_dynamics_model)): + if isinstance(model.muscles_dynamics_model[i], DingModelPulseWidthFrequency): + if pulse_width["bimapping"]: + n_stim = 1 parameter_name = ( - "pulse_duration" - if pulse_duration_similar_for_all_muscles - else "pulse_duration" + "_" + model[i].muscle_name + "pulse_width" + if pulse_width["same_for_all_muscles"] + else "pulse_width" + "_" + model.muscles_dynamics_model[i].muscle_name ) - if fixed_pulse_duration: # TODO : ADD SEVERAL INDIVIDUAL FIXED PULSE DURATION FOR EACH MUSCLE - if ( - pulse_duration_similar_for_all_muscles and i == 0 - ) or not pulse_duration_similar_for_all_muscles: + if pulse_width["fixed"]: # TODO : ADD SEVERAL INDIVIDUAL FIXED pulse width FOR EACH MUSCLE + if (pulse_width["same_for_all_muscles"] and i == 0) or not pulse_width["same_for_all_muscles"]: parameters.add( name=parameter_name, - function=DingModelPulseDurationFrequency.set_impulse_duration, + function=DingModelPulseWidthFrequency.set_impulse_width, size=n_stim, scaling=VariableScaling(parameter_name, [1] * n_stim), ) - if isinstance(fixed_pulse_duration, list): + if isinstance(pulse_width["fixed"], list): parameters_bounds.add( parameter_name, - min_bound=np.array(fixed_pulse_duration), - max_bound=np.array(fixed_pulse_duration), + min_bound=np.array(pulse_width["fixed"]), + max_bound=np.array(pulse_width["fixed"]), interpolation=InterpolationType.CONSTANT, ) - parameters_init.add(key=parameter_name, initial_guess=np.array(fixed_pulse_duration)) + parameters_init.add( + key=parameter_name, + initial_guess=np.array(pulse_width["fixed"]), + ) else: parameters_bounds.add( parameter_name, - min_bound=np.array([fixed_pulse_duration] * n_stim), - max_bound=np.array([fixed_pulse_duration] * n_stim), + min_bound=np.array([pulse_width["fixed"]] * n_stim), + max_bound=np.array([pulse_width["fixed"]] * n_stim), interpolation=InterpolationType.CONSTANT, ) - parameters_init[parameter_name] = np.array([fixed_pulse_duration] * n_stim) + parameters_init[parameter_name] = np.array([pulse_width["fixed"]] * n_stim) elif ( - pulse_duration_min and pulse_duration_max - ): # TODO : ADD SEVERAL MIN MAX PULSE DURATION FOR EACH MUSCLE - if ( - pulse_duration_similar_for_all_muscles and i == 0 - ) or not pulse_duration_similar_for_all_muscles: + pulse_width["min"] and pulse_width["max"] + ): # TODO : ADD SEVERAL MIN MAX pulse width FOR EACH MUSCLE + if (pulse_width["same_for_all_muscles"] and i == 0) or not pulse_width["same_for_all_muscles"]: parameters_bounds.add( parameter_name, - min_bound=[pulse_duration_min], - max_bound=[pulse_duration_max], + min_bound=[pulse_width["min"]], + max_bound=[pulse_width["max"]], interpolation=InterpolationType.CONSTANT, ) - pulse_duration_avg = (pulse_duration_max + pulse_duration_min) / 2 - parameters_init[parameter_name] = np.array([pulse_duration_avg] * n_stim) + pulse_width_avg = (pulse_width["max"] + pulse_width["min"]) / 2 + parameters_init[parameter_name] = np.array([pulse_width_avg] * n_stim) parameters.add( name=parameter_name, - function=DingModelPulseDurationFrequency.set_impulse_duration, + function=DingModelPulseWidthFrequency.set_impulse_width, size=n_stim, scaling=VariableScaling(parameter_name, [1] * n_stim), ) - if pulse_duration_bimapping: - pass - # parameter_bimapping.add(name="pulse_duration", to_second=[0 for _ in range(n_stim)], to_first=[0]) - # TODO : Fix Bimapping in Bioptim - - if isinstance(model[i], DingModelIntensityFrequency): + if isinstance(model.muscles_dynamics_model[i], DingModelPulseIntensityFrequency): + if pulse_intensity["bimapping"]: + n_stim = 1 parameter_name = ( "pulse_intensity" - if pulse_intensity_similar_for_all_muscles - else "pulse_intensity" + "_" + model[i].muscle_name + if pulse_intensity["same_for_all_muscles"] + else "pulse_intensity" + "_" + model.muscles_dynamics_model[i].muscle_name ) - if fixed_pulse_intensity: # TODO : ADD SEVERAL INDIVIDUAL FIXED PULSE INTENSITY FOR EACH MUSCLE - if ( - pulse_intensity_similar_for_all_muscles and i == 0 - ) or not pulse_intensity_similar_for_all_muscles: + if pulse_intensity["fixed"]: # TODO : ADD SEVERAL INDIVIDUAL FIXED PULSE INTENSITY FOR EACH MUSCLE + if (pulse_intensity["same_for_all_muscles"] and i == 0) or not pulse_intensity[ + "same_for_all_muscles" + ]: parameters.add( name=parameter_name, - function=DingModelIntensityFrequency.set_impulse_intensity, + function=DingModelPulseIntensityFrequency.set_impulse_intensity, size=n_stim, scaling=VariableScaling(parameter_name, [1] * n_stim), ) - if isinstance(fixed_pulse_intensity, list): + if isinstance(pulse_intensity["fixed"], list): parameters_bounds.add( parameter_name, - min_bound=np.array(fixed_pulse_intensity), - max_bound=np.array(fixed_pulse_intensity), + min_bound=np.array(pulse_intensity["fixed"]), + max_bound=np.array(pulse_intensity["fixed"]), interpolation=InterpolationType.CONSTANT, ) - parameters_init.add(key=parameter_name, initial_guess=np.array(fixed_pulse_intensity)) + parameters_init.add( + key=parameter_name, + initial_guess=np.array(pulse_intensity["fixed"]), + ) else: parameters_bounds.add( parameter_name, - min_bound=np.array([fixed_pulse_intensity] * n_stim), - max_bound=np.array([fixed_pulse_intensity] * n_stim), + min_bound=np.array([pulse_intensity["fixed"]] * n_stim), + max_bound=np.array([pulse_intensity["fixed"]] * n_stim), interpolation=InterpolationType.CONSTANT, ) - parameters_init[parameter_name] = np.array([fixed_pulse_intensity] * n_stim) + parameters_init[parameter_name] = np.array([pulse_intensity["fixed"]] * n_stim) elif ( - pulse_intensity_min and pulse_intensity_max + pulse_intensity["min"] and pulse_intensity["max"] ): # TODO : ADD SEVERAL MIN MAX PULSE INTENSITY FOR EACH MUSCLE - if ( - pulse_intensity_similar_for_all_muscles and i == 0 - ) or not pulse_intensity_similar_for_all_muscles: + if (pulse_intensity["same_for_all_muscles"] and i == 0) or not pulse_intensity[ + "same_for_all_muscles" + ]: parameters_bounds.add( parameter_name, - min_bound=[pulse_intensity_min], - max_bound=[pulse_intensity_max], + min_bound=[pulse_intensity["min"]], + max_bound=[pulse_intensity["max"]], interpolation=InterpolationType.CONSTANT, ) - intensity_avg = (pulse_intensity_min + pulse_intensity_max) / 2 + intensity_avg = (pulse_intensity["min"] + pulse_intensity["max"]) / 2 parameters_init[parameter_name] = np.array([intensity_avg] * n_stim) parameters.add( name=parameter_name, - function=DingModelIntensityFrequency.set_impulse_intensity, + function=DingModelPulseIntensityFrequency.set_impulse_intensity, size=n_stim, scaling=VariableScaling(parameter_name, [1] * n_stim), ) - if pulse_intensity_bimapping: - pass - # parameter_bimapping.add(name="pulse_intensity", - # to_second=[0 for _ in range(n_stim)], - # to_first=[0]) - # TODO : Fix Bimapping in Bioptim - - return parameters, parameters_bounds, parameters_init, parameter_objectives, constraints + return ( + parameters, + parameters_bounds, + parameters_init, + parameter_objectives, + ) @staticmethod - def _set_constraints(constraints, custom_constraint): + def _build_constraints(model, n_shooting, final_time, stim_time, control_type, custom_constraint=None): + constraints = ConstraintList() + + if model.activate_residual_torque: + constraints.add( + ConstraintFcn.TRACK_CONTROL, + node=Node.END, + key="tau", + target=np.zeros(model.nb_tau), + ) + + if model.muscles_dynamics_model[0].is_approximated: + time_vector = np.linspace(0, final_time, n_shooting + 1) + stim_at_node = [np.where(stim_time[i] <= time_vector)[0][0] for i in range(len(stim_time))] + additional_node = 1 if control_type == ControlType.LINEAR_CONTINUOUS else 0 + + for i in range(len(model.muscles_dynamics_model)): + if model.muscles_dynamics_model[i]._sum_stim_truncation: + max_stim_to_keep = model.muscles_dynamics_model[i]._sum_stim_truncation + else: + max_stim_to_keep = 10000000 + + index_sup = 0 + index_inf = 0 + for j in range(n_shooting + additional_node): + if j in stim_at_node: + index_sup += 1 + + constraints.add( + CustomConstraint.cn_sum, + node=j, + stim_time=stim_time[index_inf:index_sup], + model_idx=i, + ) + + if isinstance(model.muscles_dynamics_model[i], DingModelPulseWidthFrequency): + index = 0 + for j in range(n_shooting + additional_node): + if j in stim_at_node and j != 0: + index += 1 + constraints.add( + CustomConstraint.a_calculation_msk, + node=j, + last_stim_index=index, + model_idx=i, + ) + if custom_constraint: for i in range(len(custom_constraint)): if custom_constraint[i]: for j in range(len(custom_constraint[i])): constraints.add(custom_constraint[i][j]) + return constraints @staticmethod - def _set_bounds(bio_models, fes_muscle_models, bound_type, bound_data, n_stim, initial_state): + def _set_bounds(bio_models, msk_info, initial_state, n_cycles_simultaneous=1): # ---- STATE BOUNDS REPRESENTATION ---- # # # |‾‾‾‾‾‾‾‾‾‾x_max_middle‾‾‾‾‾‾‾‾‾‾‾‾x_max_end‾ @@ -529,7 +533,7 @@ def _set_bounds(bio_models, fes_muscle_models, bound_type, bound_data, n_stim, i # Sets the bound for all the phases x_bounds = BoundsList() x_init = InitialGuessList() - for model in fes_muscle_models: + for model in bio_models.muscles_dynamics_model: muscle_name = model.muscle_name variable_bound_list = [model.name_dof[i] + "_" + muscle_name for i in range(len(model.name_dof))] @@ -551,158 +555,165 @@ def _set_bounds(bio_models, fes_muscle_models, bound_type, bound_data, n_stim, i starting_bounds_min = np.concatenate((starting_bounds, min_bounds, min_bounds), axis=1) starting_bounds_max = np.concatenate((starting_bounds, max_bounds, max_bounds), axis=1) - middle_bound_min = np.concatenate((min_bounds, min_bounds, min_bounds), axis=1) - middle_bound_max = np.concatenate((max_bounds, max_bounds, max_bounds), axis=1) - - for i in range(n_stim): - for j in range(len(variable_bound_list)): - if i == 0: - x_bounds.add( - variable_bound_list[j], - min_bound=np.array([starting_bounds_min[j]]), - max_bound=np.array([starting_bounds_max[j]]), - phase=i, - interpolation=InterpolationType.CONSTANT_WITH_FIRST_AND_LAST_DIFFERENT, - ) - else: - x_bounds.add( - variable_bound_list[j], - min_bound=np.array([middle_bound_min[j]]), - max_bound=np.array([middle_bound_max[j]]), - phase=i, - interpolation=InterpolationType.CONSTANT_WITH_FIRST_AND_LAST_DIFFERENT, - ) - for i in range(n_stim): - for j in range(len(variable_bound_list)): - x_init.add(variable_bound_list[j], model.standard_rest_values()[j], phase=i) + for j in range(len(variable_bound_list)): + x_bounds.add( + variable_bound_list[j], + min_bound=np.array([starting_bounds_min[j]]), + max_bound=np.array([starting_bounds_max[j]]), + phase=0, + interpolation=InterpolationType.CONSTANT_WITH_FIRST_AND_LAST_DIFFERENT, + ) + + for j in range(len(variable_bound_list)): + x_init.add(variable_bound_list[j], model.standard_rest_values()[j], phase=0) - if bound_type == "start_end": + if msk_info["bound_type"] == "start_end": start_bounds = [] end_bounds = [] - for i in range(bio_models[0].nb_q): - start_bounds.append(3.14 / (180 / bound_data[0][i]) if bound_data[0][i] != 0 else 0) - end_bounds.append(3.14 / (180 / bound_data[1][i]) if bound_data[1][i] != 0 else 0) + for i in range(bio_models.nb_q): + start_bounds.append( + 3.14 / (180 / msk_info["bound_data"][0][i]) if msk_info["bound_data"][0][i] != 0 else 0 + ) + end_bounds.append( + 3.14 / (180 / msk_info["bound_data"][1][i]) if msk_info["bound_data"][1][i] != 0 else 0 + ) - elif bound_type == "start": + elif msk_info["bound_type"] == "start": start_bounds = [] - for i in range(bio_models[0].nb_q): - start_bounds.append(3.14 / (180 / bound_data[i]) if bound_data[i] != 0 else 0) + for i in range(bio_models.nb_q): + start_bounds.append(3.14 / (180 / msk_info["bound_data"][i]) if msk_info["bound_data"][i] != 0 else 0) - elif bound_type == "end": + elif msk_info["bound_type"] == "end": end_bounds = [] - for i in range(bio_models[0].nb_q): - end_bounds.append(3.14 / (180 / bound_data[i]) if bound_data[i] != 0 else 0) - - for i in range(n_stim): - q_x_bounds = bio_models[i].bounds_from_ranges("q") - qdot_x_bounds = bio_models[i].bounds_from_ranges("qdot") - - if i == 0: - if bound_type == "start_end": - for j in range(bio_models[i].nb_q): - q_x_bounds[j, [0]] = start_bounds[j] - elif bound_type == "start": - for j in range(bio_models[i].nb_q): - q_x_bounds[j, [0]] = start_bounds[j] - qdot_x_bounds[:, [0]] = 0 # Start without any velocity - - if i == n_stim - 1: - if bound_type == "start_end": - for j in range(bio_models[i].nb_q): - q_x_bounds[j, [-1]] = end_bounds[j] - elif bound_type == "end": - for j in range(bio_models[i].nb_q): - q_x_bounds[j, [-1]] = end_bounds[j] - - x_bounds.add(key="q", bounds=q_x_bounds, phase=i) - x_bounds.add(key="qdot", bounds=qdot_x_bounds, phase=i) + for i in range(bio_models.nb_q): + end_bounds.append(3.14 / (180 / msk_info["bound_data"][i]) if msk_info["bound_data"][i] != 0 else 0) + + q_x_bounds = bio_models.bounds_from_ranges("q") + qdot_x_bounds = bio_models.bounds_from_ranges("qdot") + + if msk_info["bound_type"] == "start_end": + for j in range(bio_models.nb_q): + q_x_bounds[j, [0]] = start_bounds[j] + q_x_bounds[j, [-1]] = end_bounds[j] + elif msk_info["bound_type"] == "start": + for j in range(bio_models.nb_q): + q_x_bounds[j, [0]] = start_bounds[j] + elif msk_info["bound_type"] == "end": + for j in range(bio_models.nb_q): + q_x_bounds[j, [-1]] = end_bounds[j] + qdot_x_bounds[:, [0]] = 0 # Start without any velocity + + x_bounds.add(key="q", bounds=q_x_bounds, phase=0) + x_bounds.add(key="qdot", bounds=qdot_x_bounds, phase=0) # Sets the initial state of q, qdot and muscle forces for all the phases if a warm start is used if initial_state: - muscle_names = bio_models[0].muscle_names - for i in range(n_stim): - x_init.add( - key="q", initial_guess=initial_state["q"][i], interpolation=InterpolationType.EACH_FRAME, phase=i - ) + for key in initial_state.keys(): + repeated_array = np.tile(initial_state[key][:, :-1], (1, n_cycles_simultaneous)) + # Append the last column of the original array to reach the desired shape + result_array = np.hstack((repeated_array, initial_state[key][:, -1:])) + initial_state[key] = result_array + + muscle_names = bio_models.muscle_names + x_init.add( + key="q", + initial_guess=initial_state["q"], + interpolation=InterpolationType.EACH_FRAME, + phase=0, + ) + x_init.add( + key="qdot", + initial_guess=initial_state["qdot"], + interpolation=InterpolationType.EACH_FRAME, + phase=0, + ) + for j in range(len(muscle_names)): x_init.add( - key="qdot", - initial_guess=initial_state["qdot"][i], + key="F_" + muscle_names[j], + initial_guess=initial_state[muscle_names[j]], interpolation=InterpolationType.EACH_FRAME, - phase=i, + phase=0, ) - for j in range(len(muscle_names)): - x_init.add( - key="F_" + muscle_names[j], - initial_guess=initial_state[muscle_names[j]][i], - interpolation=InterpolationType.EACH_FRAME, - phase=i, - ) else: - for i in range(n_stim): - x_init.add(key="q", initial_guess=[0] * bio_models[i].nb_q, phase=i) + x_init.add(key="q", initial_guess=[0] * bio_models.nb_q, phase=0) return x_bounds, x_init @staticmethod - def _set_controls(bio_models, n_stim, with_residual_torque): - # Controls bounds - nb_tau = bio_models[0].nb_tau + def _set_u_bounds(bio_models, with_residual_torque): + u_bounds = BoundsList() # Controls bounds + u_init = InitialGuessList() # Controls initial guess + if with_residual_torque: # TODO : ADD SEVERAL INDIVIDUAL FIXED RESIDUAL TORQUE FOR EACH JOINT + nb_tau = bio_models.nb_tau tau_min, tau_max, tau_init = [-50] * nb_tau, [50] * nb_tau, [0] * nb_tau - else: - tau_min, tau_max, tau_init = [0] * nb_tau, [0] * nb_tau, [0] * nb_tau - - u_bounds = BoundsList() - for i in range(n_stim): - u_bounds.add(key="tau", min_bound=tau_min, max_bound=tau_max, phase=i) - - # Controls initial guess - u_init = InitialGuessList() - for i in range(n_stim): - u_init.add(key="tau", initial_guess=tau_init, phase=i) + u_bounds.add( + key="tau", min_bound=tau_min, max_bound=tau_max, phase=0, interpolation=InterpolationType.CONSTANT + ) + u_init.add(key="tau", initial_guess=tau_init, phase=0) + + if bio_models.muscles_dynamics_model[0].is_approximated: + for i in range(len(bio_models.muscles_dynamics_model)): + u_init.add(key="Cn_sum_" + bio_models.muscles_dynamics_model[i].muscle_name, initial_guess=[0], phase=0) + + for i in range(len(bio_models.muscles_dynamics_model)): + if isinstance(bio_models.muscles_dynamics_model[i], DingModelPulseWidthFrequency): + u_init.add( + key="A_calculation_" + bio_models.muscles_dynamics_model[i].muscle_name, + initial_guess=[0], + phase=0, + ) return u_bounds, u_init @staticmethod def _set_objective( - n_stim, n_shooting, - force_fourier_coef, - end_node_tracking, - cycling_objective, - custom_objective, - q_fourier_coef, - minimize_muscle_fatigue, - minimize_muscle_force, muscle_force_key, - time_min, - time_max, + objective, + n_simultaneous_cycle: int = 1, ): # Creates the objective for our problem objective_functions = ObjectiveList() - if custom_objective: - for i in range(len(custom_objective)): - if custom_objective[i]: - for j in range(len(custom_objective[i])): - objective_functions.add(custom_objective[i][j]) + if objective["custom"]: + for i in range(len(objective["custom"])): + if objective["custom"][i]: + for j in range(len(objective["custom"][i])): + objective_functions.add(objective["custom"][i][j]) + + if objective["force_tracking"]: + force_fourier_coef = [] + for i in range(len(objective["force_tracking"][1])): + force_fourier_coef.append( + OcpFes._build_fourier_coefficient( + [ + objective["force_tracking"][0], + objective["force_tracking"][1][i], + ] + ) + ) + + force_to_track = [] + for i in range(len(force_fourier_coef)): + force_to_track.append( + FourierSeries().fit_func_by_fourier_series_with_real_coeffs( + np.linspace(0, 1, n_shooting + 1), + force_fourier_coef[i], + )[np.newaxis, :] + ) - if force_fourier_coef is not None: for j in range(len(muscle_force_key)): - for phase in range(n_stim): - for i in range(n_shooting[phase]): - objective_functions.add( - CustomObjective.track_state_from_time, - custom_type=ObjectiveFcn.Mayer, - node=i, - fourier_coeff=force_fourier_coef[j], - key=muscle_force_key[j], - quadratic=True, - weight=1, - phase=phase, - ) + objective_functions.add( + ObjectiveFcn.Lagrange.TRACK_STATE, + key=muscle_force_key[j], + weight=100, + target=force_to_track[j], + node=Node.ALL, + quadratic=True, + ) - if end_node_tracking is not None: + if objective["end_node_tracking"] is not None: for j in range(len(muscle_force_key)): objective_functions.add( ObjectiveFcn.Mayer.MINIMIZE_STATE, @@ -710,221 +721,202 @@ def _set_objective( key=muscle_force_key[j], quadratic=True, weight=1, - target=end_node_tracking[j], - phase=n_stim - 1, + target=objective["end_node_tracking"][j], + phase=0, ) - if cycling_objective: - x_center = cycling_objective["x_center"] - y_center = cycling_objective["y_center"] - radius = cycling_objective["radius"] + if objective["cycling"]: + x_center = objective["cycling"]["x_center"] + y_center = objective["cycling"]["y_center"] + radius = objective["cycling"]["radius"] 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[0] * n_stim + 1) + for theta in np.linspace( + 0, -2 * np.pi * n_simultaneous_cycle, n_shooting * n_simultaneous_cycle + 1 + ) ] ) - for phase in range(n_stim): + objective_functions.add( + ObjectiveFcn.Mayer.TRACK_MARKERS, + weight=10000000, + axes=[Axis.X, Axis.Y], + marker_index=0, + target=circle_coord_list.T, + node=Node.ALL, + phase=0, + quadratic=True, + ) + + if objective["q_tracking"]: + q_fourier_coef = [] + for i in range(len(objective["q_tracking"][1])): + q_fourier_coef.append( + OcpFes._build_fourier_coefficient([objective["q_tracking"][0], objective["q_tracking"][1][i]]) + ) + + q_to_track = [] + for i in range(len(q_fourier_coef)): + q_to_track.append( + FourierSeries().fit_func_by_fourier_series_with_real_coeffs( + np.linspace(0, 1, n_shooting + 1), + q_fourier_coef[i], + )[np.newaxis, :] + ) + + for j in range(len(q_to_track)): objective_functions.add( - ObjectiveFcn.Mayer.TRACK_MARKERS, - weight=100000, - axes=[Axis.X, Axis.Y], - marker_index=0, - target=circle_coord_list[n_shooting[0] * phase : n_shooting[0] * (phase + 1) + 1].T, + ObjectiveFcn.Lagrange.TRACK_STATE, + key="q", + weight=100, + target=q_to_track[j], node=Node.ALL, - phase=phase, quadratic=True, ) - if q_fourier_coef: - for j in range(len(q_fourier_coef)): - for phase in range(n_stim): - for i in range(n_shooting[phase]): - objective_functions.add( - CustomObjective.track_state_from_time, - custom_type=ObjectiveFcn.Mayer, - node=i, - fourier_coeff=q_fourier_coef[j], - key="q", - quadratic=True, - weight=1, - phase=phase, - index=j, - ) - - if minimize_muscle_fatigue: + if objective["minimize_muscle_fatigue"]: objective_functions.add( CustomObjective.minimize_overall_muscle_fatigue, custom_type=ObjectiveFcn.Mayer, node=Node.END, quadratic=True, - weight=-1, - phase=n_stim - 1, + weight=1, + phase=0, ) - if minimize_muscle_force: - for i in range(n_stim): - objective_functions.add( - CustomObjective.minimize_overall_muscle_force_production, - custom_type=ObjectiveFcn.Lagrange, - quadratic=True, - weight=1, - phase=i, - ) + if objective["minimize_muscle_force"]: + objective_functions.add( + CustomObjective.minimize_overall_muscle_force_production, + custom_type=ObjectiveFcn.Lagrange, + node=Node.ALL, + quadratic=True, + weight=1, + phase=0, + ) - if time_min and time_max: - for i in range(n_stim): - objective_functions.add( - ObjectiveFcn.Mayer.MINIMIZE_TIME, - weight=0.001 / n_shooting[i], - min_bound=time_min, - max_bound=time_max, - quadratic=True, - phase=i, - ) + if objective["minimize_residual_torque"]: + objective_functions.add( + ObjectiveFcn.Lagrange.MINIMIZE_CONTROL, + key="tau", + weight=10000, + quadratic=True, + phase=0, + ) return objective_functions @staticmethod - def _sanity_check_muscle_model(biorbd_model_path, fes_muscle_models): - tested_bio_model = FesMskModel(name=None, biorbd_path=biorbd_model_path, muscles_model=fes_muscle_models) - fes_muscle_models_name_list = [fes_muscle_models[x].muscle_name for x in range(len(fes_muscle_models))] - for biorbd_muscle in tested_bio_model.muscle_names: - if biorbd_muscle not in fes_muscle_models_name_list: - raise ValueError( - f"The muscle {biorbd_muscle} is not in the fes muscle model " - f"please add it into the fes_muscle_models list by providing the muscle_name =" - f" {biorbd_muscle}" - ) - - @staticmethod - def _sanity_check_fes_models_inputs( - biorbd_model_path, - bound_type, - bound_data, - fes_muscle_models, - force_tracking, - end_node_tracking, - cycling_objective, - q_tracking, - with_residual_torque, - activate_force_length_relationship, - activate_force_velocity_relationship, - minimize_muscle_fatigue, - minimize_muscle_force, + def _sanity_check_msk_inputs( + model, + msk_info, + objective, ): - if not isinstance(biorbd_model_path, str): - raise TypeError("biorbd_model_path should be a string") - - if bound_type: - tested_bio_model = FesMskModel(name=None, biorbd_path=biorbd_model_path, muscles_model=fes_muscle_models) - if not isinstance(bound_type, str) or bound_type not in ["start", "end", "start_end"]: + if msk_info["bound_type"]: + if not isinstance(msk_info["bound_type"], str) or msk_info["bound_type"] not in [ + "start", + "end", + "start_end", + ]: raise ValueError("bound_type should be a string and should be equal to start, end or start_end") - if not isinstance(bound_data, list): + if not isinstance(msk_info["bound_data"], list): raise TypeError("bound_data should be a list") - if bound_type == "start_end": - if len(bound_data) != 2 or not isinstance(bound_data[0], list) or not isinstance(bound_data[1], list): + if msk_info["bound_type"] == "start_end": + if ( + len(msk_info["bound_data"]) != 2 + or not isinstance(msk_info["bound_data"][0], list) + or not isinstance(msk_info["bound_data"][1], list) + ): raise TypeError("bound_data should be a list of two list") - if len(bound_data[0]) != tested_bio_model.nb_q or len(bound_data[1]) != tested_bio_model.nb_q: - raise ValueError(f"bound_data should be a list of {tested_bio_model.nb_q} elements") - for i in range(len(bound_data[0])): - if not isinstance(bound_data[0][i], int | float) or not isinstance(bound_data[1][i], int | float): + if len(msk_info["bound_data"][0]) != model.nb_q or len(msk_info["bound_data"][1]) != model.nb_q: + raise ValueError(f"bound_data should be a list of {model.nb_q} elements") + for i in range(len(msk_info["bound_data"][0])): + if not isinstance(msk_info["bound_data"][0][i], int | float) or not isinstance( + msk_info["bound_data"][1][i], int | float + ): raise TypeError( - f"bound data index {i}: {bound_data[0][i]} and {bound_data[1][i]} should be an int or float" + f"bound data index {i}: {msk_info['bound_data'][0][i]} and {msk_info['bound_data'][1][i]} should be an int or float" ) - if bound_type == "start" or bound_type == "end": - if len(bound_data) != tested_bio_model.nb_q: - raise ValueError(f"bound_data should be a list of {tested_bio_model.nb_q} element") - for i in range(len(bound_data)): - if not isinstance(bound_data[i], int | float): - raise TypeError(f"bound data index {i}: {bound_data[i]} should be an int or float") - - for i in range(len(fes_muscle_models)): - if not isinstance(fes_muscle_models[i], FesModel): - raise TypeError("model must be a FesModel type") - - if force_tracking: - if isinstance(force_tracking, list): - if len(force_tracking) != 2: + if msk_info["bound_type"] == "start" or msk_info["bound_type"] == "end": + if len(msk_info["bound_data"]) != model.nb_q: + raise ValueError(f"bound_data should be a list of {model.nb_q} element") + for i in range(len(msk_info["bound_data"])): + if not isinstance(msk_info["bound_data"][i], int | float): + raise TypeError(f"bound data index {i}: {msk_info['bound_data'][i]} should be an int or float") + + if objective["force_tracking"]: + if isinstance(objective["force_tracking"], list): + if len(objective["force_tracking"]) != 2: raise ValueError("force_tracking must of size 2") - if not isinstance(force_tracking[0], np.ndarray): - raise TypeError(f"force_tracking index 0: {force_tracking[0]} must be np.ndarray type") - if not isinstance(force_tracking[1], list): - raise TypeError(f"force_tracking index 1: {force_tracking[1]} must be list type") - if len(force_tracking[1]) != len(fes_muscle_models): + if not isinstance(objective["force_tracking"][0], np.ndarray): + raise TypeError(f"force_tracking index 0: {objective['force_tracking'][0]} must be np.ndarray type") + if not isinstance(objective["force_tracking"][1], list): + raise TypeError(f"force_tracking index 1: {objective['force_tracking'][1]} must be list type") + if len(objective["force_tracking"][1]) != len(model.muscles_dynamics_model): raise ValueError( - "force_tracking index 1 list must have the same size as the number of muscles in fes_muscle_models" + "force_tracking index 1 list must have the same size as the number of muscles in model.muscles_dynamics_model" ) - for i in range(len(force_tracking[1])): - if len(force_tracking[0]) != len(force_tracking[1][i]): + for i in range(len(objective["force_tracking"][1])): + if len(objective["force_tracking"][0]) != len(objective["force_tracking"][1][i]): raise ValueError("force_tracking time and force argument must be the same length") else: - raise TypeError(f"force_tracking: {force_tracking} must be list type") + raise TypeError(f"force_tracking: {objective['force_tracking']} must be list type") - if end_node_tracking: - if not isinstance(end_node_tracking, list): - raise TypeError(f"force_tracking: {end_node_tracking} must be list type") - if len(end_node_tracking) != len(fes_muscle_models): + if objective["end_node_tracking"]: + if not isinstance(objective["end_node_tracking"], list): + raise TypeError(f"force_tracking: {objective['end_node_tracking']} must be list type") + if len(objective["end_node_tracking"]) != len(model.muscles_dynamics_model): raise ValueError( "end_node_tracking list must have the same size as the number of muscles in fes_muscle_models" ) - for i in range(len(end_node_tracking)): - if not isinstance(end_node_tracking[i], int | float): - raise TypeError(f"end_node_tracking index {i}: {end_node_tracking[i]} must be int or float type") - - if cycling_objective: - if not isinstance(cycling_objective, dict): - raise TypeError(f"cycling_objective: {cycling_objective} must be dictionary type") + for i in range(len(objective["end_node_tracking"])): + if not isinstance(objective["end_node_tracking"][i], int | float): + raise TypeError( + f"end_node_tracking index {i}: {objective['end_node_tracking'][i]} must be int or float type" + ) - if len(cycling_objective) != 4: - raise ValueError( - "cycling_objective dictionary must have the same size as the number of muscles in fes_muscle_models" - ) + if objective["cycling"]: + if not isinstance(objective["cycling"], dict): + raise TypeError(f"cycling_objective: {objective['cycling']} must be dictionary type") cycling_objective_keys = ["x_center", "y_center", "radius", "target"] - if not all([cycling_objective_keys[i] in cycling_objective for i in range(len(cycling_objective_keys))]): + if not all([cycling_objective_keys[i] in objective["cycling"] for i in range(len(cycling_objective_keys))]): raise ValueError( f"cycling_objective dictionary must contain the following keys: {cycling_objective_keys}" ) - if not all([isinstance(cycling_objective[key], int | float) for key in cycling_objective_keys[:3]]): + if not all([isinstance(objective["cycling"][key], int | float) for key in cycling_objective_keys[:3]]): raise TypeError(f"cycling_objective x_center, y_center and radius inputs must be int or float") - if isinstance(cycling_objective[cycling_objective_keys[-1]], str): + if isinstance(objective["cycling"][cycling_objective_keys[-1]], str): if ( - cycling_objective[cycling_objective_keys[-1]] != "marker" - and cycling_objective[cycling_objective_keys[-1]] != "q" + objective["cycling"][cycling_objective_keys[-1]] != "marker" + and objective["cycling"][cycling_objective_keys[-1]] != "q" ): raise ValueError( - f"{cycling_objective[cycling_objective_keys[-1]]} not implemented chose between 'marker' and 'q' as 'target'" + f"{objective['cycling'][cycling_objective_keys[-1]]} not implemented chose between 'marker' and 'q' as 'target'" ) else: raise TypeError(f"cycling_objective target must be string type") - if q_tracking: - if not isinstance(q_tracking, list) and len(q_tracking) != 2: + if objective["q_tracking"]: + if not isinstance(objective["q_tracking"], list) and len(objective["q_tracking"]) != 2: raise TypeError("q_tracking should be a list of size 2") - tested_bio_model = FesMskModel(name=None, biorbd_path=biorbd_model_path, muscles_model=fes_muscle_models) - if not isinstance(q_tracking[0], list | np.ndarray): + if not isinstance(objective["q_tracking"][0], list | np.ndarray): raise ValueError("q_tracking[0] should be a list or array type") - if len(q_tracking[1]) != tested_bio_model.nb_q: + if len(objective["q_tracking"][1]) != model.nb_q: raise ValueError("q_tracking[1] should have the same size as the number of generalized coordinates") - for i in range(tested_bio_model.nb_q): - if len(q_tracking[0]) != len(q_tracking[1][i]): + for i in range(model.nb_q): + if len(objective["q_tracking"][0]) != len(objective["q_tracking"][1][i]): raise ValueError("q_tracking[0] and q_tracking[1] should have the same size") list_to_check = [ - with_residual_torque, - activate_force_length_relationship, - activate_force_velocity_relationship, - minimize_muscle_fatigue, - minimize_muscle_force, + msk_info["with_residual_torque"], + objective["minimize_muscle_fatigue"], + objective["minimize_muscle_force"], ] list_to_check_name = [ "with_residual_torque", - "activate_force_length_relationship", - "activate_force_velocity_relationship", "minimize_muscle_fatigue", "minimize_muscle_force", ] diff --git a/cocofest/optimization/fes_ocp_dynamics_nmpc_cyclic.py b/cocofest/optimization/fes_ocp_dynamics_nmpc_cyclic.py new file mode 100644 index 0000000..e091b02 --- /dev/null +++ b/cocofest/optimization/fes_ocp_dynamics_nmpc_cyclic.py @@ -0,0 +1,89 @@ +import numpy as np + +from bioptim import ( + OdeSolver, + MultiCyclicNonlinearModelPredictiveControl, + ControlType, +) + +from .fes_ocp_dynamics import OcpFesMsk +from ..models.dynamical_model import FesMskModel + + +class NmpcFesMsk(MultiCyclicNonlinearModelPredictiveControl): + def advance_window_bounds_states(self, sol, **extra): + super(NmpcFesMsk, self).advance_window_bounds_states(sol) + self.update_stim(sol) + + def update_stim(self, sol): + stimulation_time = sol.decision_parameters()["pulse_apparition_time"] + stim_prev = list(np.round(np.array(stimulation_time) - sol.ocp.phase_time[0], 3)) + + for model in self.nlp[0].model.muscles_dynamics_model: + self.nlp[0].model.muscles_dynamics_model[0].stim_prev = stim_prev + if "pulse_intensity_" + model.muscle_name in sol.parameters.keys(): + self.nlp[0].model.muscles_dynamics_model[0].stim_pulse_intensity_prev = list( + sol.parameters["pulse_intensity_" + model.muscle_name] + ) + + @staticmethod + def prepare_nmpc( + model: FesMskModel = None, + stim_time: list = None, + cycle_len: int = None, + cycle_duration: int | float = None, + n_cycles_simultaneous: int = None, + n_cycles_to_advance: int = None, + pulse_event: dict = None, + pulse_width: dict = None, + pulse_intensity: dict = None, + objective: dict = None, + msk_info: dict = None, + initial_guess_warm_start: bool = False, + use_sx: bool = True, + ode_solver: OdeSolver = OdeSolver.RK4(n_integration_steps=1), + n_threads: int = 1, + control_type: ControlType = ControlType.CONSTANT, + ): + + input_dict = { + "model": model, + "stim_time": stim_time, + "n_shooting": cycle_len, + "final_time": cycle_duration, + "n_cycles_simultaneous": n_cycles_simultaneous, + "n_cycles_to_advance": n_cycles_to_advance, + "pulse_event": pulse_event, + "pulse_width": pulse_width, + "pulse_intensity": pulse_intensity, + "objective": objective, + "msk_info": msk_info, + "initial_guess_warm_start": initial_guess_warm_start, + "use_sx": use_sx, + "ode_solver": ode_solver, + "n_threads": n_threads, + "control_type": control_type, + } + + optimization_dict = OcpFesMsk._prepare_optimization_problem(input_dict) + + return NmpcFesMsk( + bio_model=[optimization_dict["model"]], + dynamics=optimization_dict["dynamics"], + cycle_len=cycle_len, + cycle_duration=cycle_duration, + n_cycles_simultaneous=n_cycles_simultaneous, + n_cycles_to_advance=n_cycles_to_advance, + common_objective_functions=optimization_dict["objective_functions"], + x_init=optimization_dict["x_init"], + x_bounds=optimization_dict["x_bounds"], + constraints=optimization_dict["constraints"], + parameters=optimization_dict["parameters"], + parameter_bounds=optimization_dict["parameters_bounds"], + parameter_init=optimization_dict["parameters_init"], + parameter_objectives=optimization_dict["parameter_objectives"], + use_sx=optimization_dict["use_sx"], + ode_solver=optimization_dict["ode_solver"], + n_threads=optimization_dict["n_threads"], + control_type=optimization_dict["control_type"], + ) diff --git a/cocofest/optimization/fes_ocp_nmpc_cyclic.py b/cocofest/optimization/fes_ocp_nmpc_cyclic.py index f3502dd..2f8d6ac 100644 --- a/cocofest/optimization/fes_ocp_nmpc_cyclic.py +++ b/cocofest/optimization/fes_ocp_nmpc_cyclic.py @@ -1,364 +1,83 @@ -import math - import numpy as np +from casadi import SX + from bioptim import ( - SolutionMerge, - ObjectiveList, - ObjectiveFcn, OdeSolver, - Node, - OptimalControlProgram, + CyclicNonlinearModelPredictiveControl, ControlType, - TimeAlignment, + Solution, ) from .fes_ocp import OcpFes from ..models.fes_model import FesModel -from ..custom_objectives import CustomObjective -class OcpFesNmpcCyclic: - def __init__( - self, +class NmpcFes(CyclicNonlinearModelPredictiveControl): + def advance_window_bounds_states(self, sol, **extra): + super(NmpcFes, self).advance_window_bounds_states(sol) + self.update_stim(sol) + + def update_stim(self, sol): + stimulation_time = sol.decision_parameters()["pulse_apparition_time"] + stim_prev = list(np.round(np.array(stimulation_time) - sol.ocp.phase_time[0], 3)) + + self.nlp[0].model.stim_prev = stim_prev + if "pulse_intensity" in sol.parameters.keys(): + self.nlp[0].model.stim_pulse_intensity_prev = list(sol.parameters["pulse_intensity"]) + + @staticmethod + def prepare_nmpc( model: FesModel = None, - n_stim: int = None, - n_shooting: int = None, - final_time: int | float = None, + stim_time: list = None, + cycle_duration: int | float = None, + n_cycles_simultaneous: int = None, + n_cycles_to_advance: int = None, pulse_event: dict = None, - pulse_duration: dict = None, + pulse_width: dict = None, pulse_intensity: dict = None, - n_total_cycles: int = None, - n_simultaneous_cycles: int = None, - n_cycle_to_advance: int = None, - cycle_to_keep: str = None, objective: dict = None, use_sx: bool = True, ode_solver: OdeSolver = OdeSolver.RK4(n_integration_steps=1), n_threads: int = 1, + control_type: ControlType = ControlType.CONSTANT, ): - super(OcpFesNmpcCyclic, self).__init__() - self.model = model - self.n_stim = n_stim - self.n_shooting = n_shooting - self.final_time = final_time - self.pulse_event = pulse_event - self.pulse_duration = pulse_duration - self.pulse_intensity = pulse_intensity - self.objective = objective - self.n_total_cycles = n_total_cycles - self.n_simultaneous_cycles = n_simultaneous_cycles - self.n_cycle_to_advance = n_cycle_to_advance - self.cycle_to_keep = cycle_to_keep - self.use_sx = use_sx - self.ode_solver = ode_solver - self.n_threads = n_threads - self.ocp = None - self._nmpc_sanity_check() - self.states = [] - self.parameters = [] - self.previous_stim = [] - self.result = {"time": {}, "states": {}, "parameters": {}} - self.temp_last_node_time = 0 - self.first_node_in_phase = 0 - self.last_node_in_phase = 0 - - def prepare_nmpc(self): - (pulse_event, pulse_duration, pulse_intensity, objective) = OcpFes._fill_dict( - self.pulse_event, self.pulse_duration, self.pulse_intensity, self.objective - ) - - time_min = pulse_event["min"] - time_max = pulse_event["max"] - time_bimapping = pulse_event["bimapping"] - frequency = pulse_event["frequency"] - round_down = pulse_event["round_down"] - pulse_mode = pulse_event["pulse_mode"] - - fixed_pulse_duration = pulse_duration["fixed"] - pulse_duration_min = pulse_duration["min"] - pulse_duration_max = pulse_duration["max"] - pulse_duration_bimapping = pulse_duration["bimapping"] - - fixed_pulse_intensity = pulse_intensity["fixed"] - pulse_intensity_min = pulse_intensity["min"] - pulse_intensity_max = pulse_intensity["max"] - pulse_intensity_bimapping = pulse_intensity["bimapping"] - - force_tracking = objective["force_tracking"] - end_node_tracking = objective["end_node_tracking"] - custom_objective = objective["custom"] - OcpFes._sanity_check( - model=self.model, - n_stim=self.n_stim, - n_shooting=self.n_shooting, - final_time=self.final_time, - pulse_mode=pulse_mode, - frequency=frequency, - time_min=time_min, - time_max=time_max, - time_bimapping=time_bimapping, - fixed_pulse_duration=fixed_pulse_duration, - pulse_duration_min=pulse_duration_min, - pulse_duration_max=pulse_duration_max, - pulse_duration_bimapping=pulse_duration_bimapping, - fixed_pulse_intensity=fixed_pulse_intensity, - pulse_intensity_min=pulse_intensity_min, - pulse_intensity_max=pulse_intensity_max, - pulse_intensity_bimapping=pulse_intensity_bimapping, - force_tracking=force_tracking, - end_node_tracking=end_node_tracking, - custom_objective=custom_objective, - use_sx=self.use_sx, - ode_solver=self.ode_solver, - n_threads=self.n_threads, - ) - - OcpFes._sanity_check_frequency( - n_stim=self.n_stim, final_time=self.final_time, frequency=frequency, round_down=round_down - ) - - force_fourier_coefficient = ( - None if force_tracking is None else OcpFes._build_fourier_coefficient(force_tracking) - ) - - models = [self.model] * self.n_stim * self.n_simultaneous_cycles - - final_time_phase = OcpFes._build_phase_time( - final_time=self.final_time * self.n_simultaneous_cycles, - n_stim=self.n_stim * self.n_simultaneous_cycles, - pulse_mode=pulse_mode, - time_min=time_min, - time_max=time_max, - ) - parameters, parameters_bounds, parameters_init, parameter_objectives, constraints = OcpFes._build_parameters( - model=self.model, - n_stim=self.n_stim * self.n_simultaneous_cycles, - time_min=time_min, - time_max=time_max, - time_bimapping=time_bimapping, - fixed_pulse_duration=fixed_pulse_duration, - pulse_duration_min=pulse_duration_min, - pulse_duration_max=pulse_duration_max, - pulse_duration_bimapping=pulse_duration_bimapping, - fixed_pulse_intensity=fixed_pulse_intensity, - pulse_intensity_min=pulse_intensity_min, - pulse_intensity_max=pulse_intensity_max, - pulse_intensity_bimapping=pulse_intensity_bimapping, - use_sx=self.use_sx, - ) - - if len(constraints) == 0 and len(parameters) == 0: - raise ValueError( - "This is not an optimal control problem," - " add parameter to optimize or use the IvpFes method to build your problem" - ) - - dynamics = OcpFes._declare_dynamics(models, self.n_stim * self.n_simultaneous_cycles) - x_bounds, x_init = OcpFes._set_bounds(self.model, self.n_stim * self.n_simultaneous_cycles) - one_cycle_shooting = [self.n_shooting] * self.n_stim - objective_functions = self._set_objective( - self.n_stim, - one_cycle_shooting, - force_fourier_coefficient, - end_node_tracking, - custom_objective, - time_min, - time_max, - self.n_simultaneous_cycles, - ) - all_cycle_n_shooting = [self.n_shooting] * self.n_stim * self.n_simultaneous_cycles - self.ocp = OptimalControlProgram( - bio_model=models, - dynamics=dynamics, - n_shooting=all_cycle_n_shooting, - phase_time=final_time_phase, - objective_functions=objective_functions, - x_init=x_init, - x_bounds=x_bounds, - constraints=constraints, - parameters=parameters, - parameter_bounds=parameters_bounds, - parameter_init=parameters_init, - parameter_objectives=parameter_objectives, - control_type=ControlType.CONSTANT, - use_sx=self.use_sx, - ode_solver=self.ode_solver, - n_threads=self.n_threads, + input_dict = { + "model": model, + "stim_time": stim_time, + "n_shooting": OcpFes.prepare_n_shooting(stim_time, cycle_duration), + "final_time": cycle_duration, + "n_cycles_simultaneous": n_cycles_simultaneous, + "n_cycles_to_advance": n_cycles_to_advance, + "pulse_event": pulse_event, + "pulse_width": pulse_width, + "pulse_intensity": pulse_intensity, + "objective": objective, + "use_sx": use_sx, + "ode_solver": ode_solver, + "n_threads": n_threads, + "control_type": control_type, + } + + optimization_dict = OcpFes._prepare_optimization_problem(input_dict) + + return NmpcFes( + bio_model=[optimization_dict["model"]], + dynamics=optimization_dict["dynamics"], + cycle_len=optimization_dict["n_shooting"], + cycle_duration=cycle_duration, + n_cycles_simultaneous=n_cycles_simultaneous, + n_cycles_to_advance=n_cycles_to_advance, + common_objective_functions=optimization_dict["objective_functions"], + x_init=optimization_dict["x_init"], + x_bounds=optimization_dict["x_bounds"], + constraints=optimization_dict["constraints"], + parameters=optimization_dict["parameters"], + parameter_bounds=optimization_dict["parameters_bounds"], + parameter_init=optimization_dict["parameters_init"], + parameter_objectives=optimization_dict["parameter_objectives"], + control_type=control_type, + use_sx=optimization_dict["use_sx"], + ode_solver=optimization_dict["ode_solver"], + n_threads=optimization_dict["n_threads"], ) - - return self.ocp - - def update_states_bounds(self, sol_states): - state_keys = list(self.ocp.nlp[0].states.keys()) - index_to_keep = 1 * self.n_stim - 1 # todo: update this when more simultaneous cycles than 3 - for key in state_keys: - self.ocp.nlp[0].x_bounds[key].max[0][0] = sol_states[index_to_keep][key][0][-1] - self.ocp.nlp[0].x_bounds[key].min[0][0] = sol_states[index_to_keep][key][0][-1] - for j in range(index_to_keep, len(self.ocp.nlp)): - self.ocp.nlp[j].x_init[key].init[0][0] = sol_states[j][key][0][0] - - def update_stim(self, sol): - if "pulse_apparition_time" in sol.decision_parameters(): - stimulation_time = sol.decision_parameters()["pulse_apparition_time"] - else: - stimulation_time = [0] + list(np.cumsum(sol.ocp.phase_time[: self.n_stim - 1])) - - stim_prev = list(np.array(stimulation_time) - self.final_time) - if self.previous_stim: - update_previous_stim = list(np.array(self.previous_stim) - self.final_time) - self.previous_stim = update_previous_stim + stim_prev - - else: - self.previous_stim = stim_prev - - for j in range(len(self.ocp.nlp)): - self.ocp.nlp[j].model.set_pass_pulse_apparition_time(self.previous_stim) - # TODO: Does not seem to be taken into account by the next model force estimation - - def store_results(self, sol_time, sol_states, sol_parameters, index): - if self.cycle_to_keep == "middle": - # Get the middle phase index to keep - phase_to_keep = int(math.ceil(self.n_simultaneous_cycles / 2)) - self.first_node_in_phase = self.n_stim * (phase_to_keep - 1) - self.last_node_in_phase = self.n_stim * phase_to_keep - - # Initialize the dict if it's the first iteration - if index == 0: - self.result["time"] = [None] * self.n_total_cycles - [ - self.result["states"].update({state_key: [None] * self.n_total_cycles}) - for state_key in list(sol_states[0].keys()) - ] - [ - self.result["parameters"].update({key_parameter: [None] * self.n_total_cycles}) - for key_parameter in list(sol_parameters.keys()) - ] - - # Store the results - phase_size = np.array(sol_time).shape[0] - node_size = np.array(sol_time).shape[1] - sol_time = list(np.array(sol_time).reshape(phase_size * node_size))[ - self.first_node_in_phase * node_size : self.last_node_in_phase * node_size - ] - sol_time = list(dict.fromkeys(sol_time)) # Remove duplicate time - if index == 0: - updated_sol_time = [t - sol_time[0] for t in sol_time] - else: - updated_sol_time = [t - sol_time[0] + self.temp_last_node_time for t in sol_time] - self.temp_last_node_time = updated_sol_time[-1] - self.result["time"][index] = updated_sol_time[:-1] - - for state_key in list(sol_states[0].keys()): - middle_states_values = sol_states[self.first_node_in_phase : self.last_node_in_phase] - middle_states_values = [ - list(middle_states_values[i][state_key][0])[:-1] for i in range(len(middle_states_values)) - ] # Remove the last node duplicate - middle_states_values = [j for sub in middle_states_values for j in sub] - self.result["states"][state_key][index] = middle_states_values - - for key_parameter in list(sol_parameters.keys()): - self.result["parameters"][key_parameter][index] = sol_parameters[key_parameter][ - self.first_node_in_phase : self.last_node_in_phase - ] - return - - def solve(self): - for i in range(self.n_total_cycles // self.n_cycle_to_advance): - sol = self.ocp.solve() - sol_states = sol.decision_states(to_merge=[SolutionMerge.NODES]) - self.update_states_bounds(sol_states) - sol_time = sol.decision_time(to_merge=SolutionMerge.NODES, time_alignment=TimeAlignment.STATES) - sol_parameters = sol.decision_parameters() - self.store_results(sol_time, sol_states, sol_parameters, i) - # self.update_stim(sol) - # Todo uncomment when the model is updated to take into account the past stimulation - - @staticmethod - def _set_objective( - n_stim, - n_shooting, - force_fourier_coefficient, - end_node_tracking, - custom_objective, - time_min, - time_max, - n_simultaneous_cycles, - ): - # Creates the objective for our problem - objective_functions = ObjectiveList() - if custom_objective: - if len(custom_objective) != n_stim: - raise ValueError( - "The number of custom objective must be equal to the stimulation number of a single cycle" - ) - for i in range(len(custom_objective)): - for j in range(n_simultaneous_cycles): - objective_functions.add(custom_objective[i + j * n_stim][0]) - - if force_fourier_coefficient is not None: - for phase in range(n_stim): - for i in range(n_shooting[phase]): - for j in range(n_simultaneous_cycles): - objective_functions.add( - CustomObjective.track_state_from_time, - custom_type=ObjectiveFcn.Mayer, - node=i, - fourier_coeff=force_fourier_coefficient, - key="F", - quadratic=True, - weight=1, - phase=phase + j * n_stim, - ) - - if end_node_tracking: - if isinstance(end_node_tracking, int | float): - for i in range(n_simultaneous_cycles): - objective_functions.add( - ObjectiveFcn.Mayer.MINIMIZE_STATE, - node=Node.END, - key="F", - quadratic=True, - weight=1, - target=end_node_tracking, - phase=n_stim - 1 + i * n_stim, - ) - - if time_min and time_max: - for i in range(n_stim): - for j in range(n_simultaneous_cycles): - objective_functions.add( - ObjectiveFcn.Mayer.MINIMIZE_TIME, - weight=0.001 / n_shooting[i], - min_bound=time_min, - max_bound=time_max, - quadratic=True, - phase=i + j * n_stim, - ) - - return objective_functions - - def _nmpc_sanity_check(self): - if not isinstance(self.n_total_cycles, int): - raise TypeError("n_total_cycles must be an integer") - if not isinstance(self.n_simultaneous_cycles, int): - raise TypeError("n_simultaneous_cycles must be an integer") - if not isinstance(self.n_cycle_to_advance, int): - raise TypeError("n_cycle_to_advance must be an integer") - if not isinstance(self.cycle_to_keep, str): - raise TypeError("cycle_to_keep must be a string") - - if self.n_cycle_to_advance > self.n_simultaneous_cycles: - raise ValueError("The number of n_simultaneous_cycles must be higher than the number of n_cycle_to_advance") - - if self.n_total_cycles % self.n_cycle_to_advance != 0: - raise ValueError("The number of n_total_cycles must be a multiple of the number n_cycle_to_advance") - - if self.cycle_to_keep not in ["first", "middle", "last"]: - raise ValueError("cycle_to_keep must be either 'first', 'middle' or 'last'") - if self.cycle_to_keep != "middle": - raise NotImplementedError("Only 'middle' cycle_to_keep is implemented") - - if self.n_simultaneous_cycles != 3: - raise NotImplementedError("Only 3 simultaneous cycles are implemented yet work in progress") - # Todo add more simultaneous cycles diff --git a/cocofest/result/animate.py b/cocofest/result/animate.py index c8e80c6..3338afd 100644 --- a/cocofest/result/animate.py +++ b/cocofest/result/animate.py @@ -26,11 +26,11 @@ def load(self): self.model = biorbd.Model(model_path) self.time = self.data["time"] self.state_q = self.data["states"]["q"] + self.state_q = self.state_q if self.state_q.ndim == 2 else np.expand_dims(self.state_q, axis=0) self.frames = self.state_q.shape[1] - def animate(self, model_path: str = None): - if model_path: - self.model = biorbd.Model(model_path) + def animate(self, model: biorbd.Model = None): + self.model = model self.load() # pyorerun animation @@ -43,9 +43,8 @@ def animate(self, model_path: str = None): viz.add_animated_model(prr_model, self.state_q) viz.rerun("msk_model") - def multiple_animations(self, additional_path: list[str], model_path: str = None): - if model_path: - self.model = biorbd.Model(model_path) + def multiple_animations(self, additional_path: list[str], model: biorbd.Model = None): + self.model = model self.load() nb_seconds = self.time[-1] t_span = np.linspace(0, nb_seconds, self.frames) diff --git a/cocofest/result/pickle.py b/cocofest/result/pickle.py index 27a1fea..a836e35 100644 --- a/cocofest/result/pickle.py +++ b/cocofest/result/pickle.py @@ -13,7 +13,10 @@ def pickle(self): bounds_key = self.sol.ocp.parameter_bounds.keys() bounds = {} for key in bounds_key: - bounds[key] = self.sol.ocp.parameter_bounds[key].min[0][0], self.sol.ocp.parameter_bounds[key].max[0][0] + bounds[key] = ( + self.sol.ocp.parameter_bounds[key].min[0][0], + self.sol.ocp.parameter_bounds[key].max[0][0], + ) time = self.sol.decision_time(to_merge=[SolutionMerge.PHASES, SolutionMerge.NODES]) time = time.reshape(time.shape[0]) diff --git a/cocofest/result/plot.py b/cocofest/result/plot.py index b864ebc..08e449f 100644 --- a/cocofest/result/plot.py +++ b/cocofest/result/plot.py @@ -82,10 +82,30 @@ def plot(self, starting_location: str = None, show_rehastim=False): def plot_rehastim(self): - triceps_brachii = {"theta": np.radians(100), "radii": 1 / 5, "width": np.radians(160), "bottom": 4 / 5} - biceps_brachii = {"theta": np.radians(295), "radii": 1 / 5, "width": np.radians(150), "bottom": 3 / 5} - deltoideus_anterior = {"theta": np.radians(100), "radii": 1 / 5, "width": np.radians(160), "bottom": 2 / 5} - deltoideus_posterior = {"theta": np.radians(295), "radii": 1 / 5, "width": np.radians(150), "bottom": 1 / 5} + triceps_brachii = { + "theta": np.radians(100), + "radii": 1 / 5, + "width": np.radians(160), + "bottom": 4 / 5, + } + biceps_brachii = { + "theta": np.radians(295), + "radii": 1 / 5, + "width": np.radians(150), + "bottom": 3 / 5, + } + deltoideus_anterior = { + "theta": np.radians(100), + "radii": 1 / 5, + "width": np.radians(160), + "bottom": 2 / 5, + } + deltoideus_posterior = { + "theta": np.radians(295), + "radii": 1 / 5, + "width": np.radians(150), + "bottom": 1 / 5, + } empty = {"theta": 1, "radii": 1 / 5, "width": 1, "bottom": 0} stimulated_muscles = { "triceps_brachii": triceps_brachii, @@ -111,7 +131,13 @@ def plot_rehastim(self): width.append(stimulated_muscles[muscle]["width"]) bottom.append(stimulated_muscles[muscle]["bottom"]) bars = ax.bar( - theta, radii, width=width, bottom=bottom, label=stimulated_muscles.keys(), edgecolor="black", linewidth=2 + theta, + radii, + width=width, + bottom=bottom, + label=stimulated_muscles.keys(), + edgecolor="black", + linewidth=2, ) color = ["b", "g", "r", "c", "w"] for i in range(len(bars)): @@ -138,15 +164,15 @@ def extract_data_from_sol(self, solution: Solution): if sum(["pulse_intensity" in parameter_key for parameter_key in solution.ocp.parameters.keys()]) > 0 else False ) - pulse_duration = ( + pulse_width = ( True - if sum(["pulse_duration" in parameter_key for parameter_key in solution.ocp.parameters.keys()]) > 0 + if sum(["pulse_width" in parameter_key for parameter_key in solution.ocp.parameters.keys()]) > 0 else False ) - parameter = "pulse_intensity" if intensity else "pulse_duration" if pulse_duration else None + parameter = "pulse_intensity" if intensity else "pulse_width" if pulse_width else None if parameter is None: raise ValueError( - "The solution must contain either a pulse intensity or a pulse duration parameter to be plotted with the PlotCyclingResult class" + "The solution must contain either a pulse intensity or a pulse width parameter to be plotted with the PlotCyclingResult class" ) counter = 0 @@ -201,20 +227,18 @@ def extract_data_from_pickle(self, solution: str): if sum(["pulse_intensity" in parameter_key for parameter_key in pickle_data["parameters"]]) > 0 else False ) - pulse_duration = ( - True - if sum(["pulse_duration" in parameter_key for parameter_key in pickle_data["parameters"]]) > 0 - else False + pulse_width = ( + True if sum(["pulse_width" in parameter_key for parameter_key in pickle_data["parameters"]]) > 0 else False ) - parameter = "pulse_intensity" if intensity else "pulse_duration" if pulse_duration else None + parameter = "pulse_intensity" if intensity else "pulse_width" if pulse_width else None if parameter is None: raise ValueError( - "The solution must contain either a pulse intensity or a pulse duration parameter to be plotted with the PlotCyclingResult class" + "The solution must contain either a pulse intensity or a pulse width parameter to be plotted with the PlotCyclingResult class" ) counter = 0 muscle_name_list = list(pickle_data["parameters"].keys()) - muscle_name_list.remove("pulse_apparition_time") if pulse_apparition_time_as_parameter else None + (muscle_name_list.remove("pulse_apparition_time") if pulse_apparition_time_as_parameter else None) muscle_name_list = [s.replace(parameter + "_", "", 1) for s in muscle_name_list] for muscle in muscle_name_list: @@ -247,16 +271,43 @@ def extract_data_from_pickle(self, solution: str): @staticmethod def add_empty_muscle(data): - empty = {"theta": 1, "radii": 1 / len(data), "width": 1, "bottom": 0, "opacity": 0, "label": ""} + empty = { + "theta": 1, + "radii": 1 / len(data), + "width": 1, + "bottom": 0, + "opacity": 0, + "label": "", + } data["empty"] = empty return data @staticmethod def rehamove_data(): - triceps_brachii = {"theta": np.radians(100), "radii": 1 / 5, "width": np.radians(160), "bottom": 4 / 5} - biceps_brachii = {"theta": np.radians(295), "radii": 1 / 5, "width": np.radians(150), "bottom": 3 / 5} - deltoideus_anterior = {"theta": np.radians(100), "radii": 1 / 5, "width": np.radians(160), "bottom": 2 / 5} - deltoideus_posterior = {"theta": np.radians(295), "radii": 1 / 5, "width": np.radians(150), "bottom": 1 / 5} + triceps_brachii = { + "theta": np.radians(100), + "radii": 1 / 5, + "width": np.radians(160), + "bottom": 4 / 5, + } + biceps_brachii = { + "theta": np.radians(295), + "radii": 1 / 5, + "width": np.radians(150), + "bottom": 3 / 5, + } + deltoideus_anterior = { + "theta": np.radians(100), + "radii": 1 / 5, + "width": np.radians(160), + "bottom": 2 / 5, + } + deltoideus_posterior = { + "theta": np.radians(295), + "radii": 1 / 5, + "width": np.radians(150), + "bottom": 1 / 5, + } empty = {"theta": 1, "radii": 1 / 5, "width": 1, "bottom": 0} stimulated_muscles = { "triceps_brachii": triceps_brachii, diff --git a/data_process/force_from_sensor.py b/data_process/force_from_sensor.py index 66bedac..2128bdf 100644 --- a/data_process/force_from_sensor.py +++ b/data_process/force_from_sensor.py @@ -63,7 +63,11 @@ def __init__( if isinstance(muscle_name, str) else muscle_name[i] if isinstance(muscle_name, list) else "biceps" ) - dictionary = {"time": self.time, muscle_name: self.all_biceps_force_vector, "stim_time": self.stim_time} + dictionary = { + "time": self.time, + muscle_name: self.all_biceps_force_vector, + "stim_time": self.stim_time, + } with open(save_pickle_path, "wb") as file: pickle.dump(dictionary, file) diff --git a/tests/shard1/test_ivp.py b/tests/shard1/test_ivp.py index f46767a..2afa325 100644 --- a/tests/shard1/test_ivp.py +++ b/tests/shard1/test_ivp.py @@ -4,19 +4,18 @@ from cocofest import ( IvpFes, - DingModelFrequency, - DingModelFrequencyWithFatigue, - DingModelPulseDurationFrequency, - DingModelPulseDurationFrequencyWithFatigue, - DingModelIntensityFrequency, - DingModelIntensityFrequencyWithFatigue, + ModelMaker, ) -@pytest.mark.parametrize("model", [DingModelFrequency(), DingModelFrequencyWithFatigue()]) +ding2003_model = ModelMaker.create_model("ding2003", is_approximated=False) +ding2003_with_fatigue_model = ModelMaker.create_model("ding2003_with_fatigue", is_approximated=False) + + +@pytest.mark.parametrize("model", [ding2003_model, ding2003_with_fatigue_model]) def test_ding2003_ivp(model): - fes_parameters = {"model": model, "n_stim": 3} - ivp_parameters = {"n_shooting": 10, "final_time": 0.3, "use_sx": True} + fes_parameters = {"model": model, "stim_time": [0, 0.1, 0.2]} + ivp_parameters = {"final_time": 0.3, "use_sx": True} ivp = IvpFes(fes_parameters, ivp_parameters) @@ -25,48 +24,59 @@ def test_ding2003_ivp(model): if model._with_fatigue: np.testing.assert_almost_equal(result["F"][0][0], 0) - np.testing.assert_almost_equal(result["F"][0][10], 92.06532561584642) - np.testing.assert_almost_equal(result["F"][0][-1], 138.94556672277545) + np.testing.assert_almost_equal(result["F"][0][15], 140.4516369184822) + np.testing.assert_almost_equal(result["F"][0][-1], 141.96528249885299) else: np.testing.assert_almost_equal(result["F"][0][0], 0) - np.testing.assert_almost_equal(result["F"][0][10], 91.4098711524036) - np.testing.assert_almost_equal(result["F"][0][-1], 130.3736693032713) + np.testing.assert_almost_equal(result["F"][0][15], 138.98486525648613) + np.testing.assert_almost_equal(result["F"][0][-1], 133.22602892155032) + + +ding2007_model = ModelMaker.create_model("ding2007", is_approximated=False) +ding2007_with_fatigue_model = ModelMaker.create_model("ding2007_with_fatigue", is_approximated=False) -@pytest.mark.parametrize("model", [DingModelPulseDurationFrequency(), DingModelPulseDurationFrequencyWithFatigue()]) -@pytest.mark.parametrize("pulse_duration", [0.0003, [0.0003, 0.0004, 0.0005]]) -def test_ding2007_ivp(model, pulse_duration): - fes_parameters = {"model": model, "n_stim": 3, "pulse_duration": pulse_duration} - ivp_parameters = {"n_shooting": 10, "final_time": 0.3, "use_sx": True} +@pytest.mark.parametrize( + "model", + [ding2007_model, ding2007_with_fatigue_model], +) +@pytest.mark.parametrize("pulse_width", [0.0003, [0.0003, 0.0004, 0.0005]]) +def test_ding2007_ivp(model, pulse_width): + fes_parameters = {"model": model, "stim_time": [0, 0.1, 0.2], "pulse_width": pulse_width} + ivp_parameters = {"final_time": 0.3, "use_sx": True} ivp = IvpFes(fes_parameters, ivp_parameters) # Integrating the solution result = ivp.integrate(return_time=False) - if model._with_fatigue and isinstance(pulse_duration, list): + if model._with_fatigue and isinstance(pulse_width, list): np.testing.assert_almost_equal(result["F"][0][0], 0) np.testing.assert_almost_equal(result["F"][0][10], 28.3477940849177) - np.testing.assert_almost_equal(result["F"][0][-1], 52.38870505209033) - elif model._with_fatigue is False and isinstance(pulse_duration, list): + np.testing.assert_almost_equal(result["F"][0][-1], 54.992643280040774) + elif model._with_fatigue is False and isinstance(pulse_width, list): np.testing.assert_almost_equal(result["F"][0][0], 0) np.testing.assert_almost_equal(result["F"][0][10], 28.116838973337046) - np.testing.assert_almost_equal(result["F"][0][-1], 49.30316895125016) - elif model._with_fatigue and isinstance(pulse_duration, float): + np.testing.assert_almost_equal(result["F"][0][-1], 51.68572030372867) + elif model._with_fatigue and isinstance(pulse_width, float): np.testing.assert_almost_equal(result["F"][0][0], 0) np.testing.assert_almost_equal(result["F"][0][10], 28.3477940849177) - np.testing.assert_almost_equal(result["F"][0][-1], 36.51217790065462) - elif model._with_fatigue is False and isinstance(pulse_duration, float): + np.testing.assert_almost_equal(result["F"][0][-1], 38.25981983501241) + elif model._with_fatigue is False and isinstance(pulse_width, float): np.testing.assert_almost_equal(result["F"][0][0], 0) np.testing.assert_almost_equal(result["F"][0][10], 28.116838973337046) - np.testing.assert_almost_equal(result["F"][0][-1], 34.6369350284091) + np.testing.assert_almost_equal(result["F"][0][-1], 36.263299814887766) + +hmed2018_model = ModelMaker.create_model("hmed2018", is_approximated=False) +hmed2018_with_fatigue_model = ModelMaker.create_model("hmed2018_with_fatigue", is_approximated=False) -@pytest.mark.parametrize("model", [DingModelIntensityFrequency(), DingModelIntensityFrequencyWithFatigue()]) + +@pytest.mark.parametrize("model", [hmed2018_model, hmed2018_with_fatigue_model]) @pytest.mark.parametrize("pulse_intensity", [50, [50, 60, 70]]) def test_hmed2018_ivp(model, pulse_intensity): - fes_parameters = {"model": model, "n_stim": 3, "pulse_intensity": pulse_intensity} - ivp_parameters = {"n_shooting": 10, "final_time": 0.3, "use_sx": True} + fes_parameters = {"model": model, "stim_time": [0, 0.1, 0.2], "pulse_intensity": pulse_intensity} + ivp_parameters = {"final_time": 0.3, "use_sx": True} ivp = IvpFes(fes_parameters, ivp_parameters) @@ -76,11 +86,11 @@ def test_hmed2018_ivp(model, pulse_intensity): if model._with_fatigue and isinstance(pulse_intensity, list): np.testing.assert_almost_equal(result["F"][0][0], 0) np.testing.assert_almost_equal(result["F"][0][10], 42.18211764372109) - np.testing.assert_almost_equal(result["F"][0][-1], 96.38882396648857) + np.testing.assert_almost_equal(result["F"][0][-1], 94.48614335382977) elif model._with_fatigue is False and isinstance(pulse_intensity, list): np.testing.assert_almost_equal(result["F"][0][0], 0) np.testing.assert_almost_equal(result["F"][0][10], 41.91914906078192) - np.testing.assert_almost_equal(result["F"][0][-1], 92.23749672532881) + np.testing.assert_almost_equal(result["F"][0][-1], 90.43032549879167) elif model._with_fatigue and isinstance(pulse_intensity, float): np.testing.assert_almost_equal(result["F"][0][0], 0) np.testing.assert_almost_equal(result["F"][0][10], 42.18211764372109) @@ -93,9 +103,12 @@ def test_hmed2018_ivp(model, pulse_intensity): @pytest.mark.parametrize("pulse_mode", ["single", "doublet", "triplet"]) def test_pulse_mode_ivp(pulse_mode): - n_stim = 3 if pulse_mode == "single" else 6 if pulse_mode == "doublet" else 9 - fes_parameters = {"model": DingModelFrequencyWithFatigue(), "n_stim": n_stim, "pulse_mode": pulse_mode} - ivp_parameters = {"n_shooting": 10, "final_time": 0.3, "use_sx": True} + fes_parameters = { + "model": ding2003_with_fatigue_model, + "stim_time": [0, 0.1, 0.2], + "pulse_mode": pulse_mode, + } + ivp_parameters = {"final_time": 0.3, "use_sx": True} ivp = IvpFes(fes_parameters, ivp_parameters) @@ -104,26 +117,30 @@ def test_pulse_mode_ivp(pulse_mode): if pulse_mode == "single": np.testing.assert_almost_equal(result["F"][0][0], 0) - np.testing.assert_almost_equal(result["F"][0][10], 92.06532561584642) - np.testing.assert_almost_equal(result["F"][0][-1], 138.94556672277545) + np.testing.assert_almost_equal(result["F"][0][15], 140.4516369184822) + np.testing.assert_almost_equal(result["F"][0][-1], 141.96528249885299) elif pulse_mode == "doublet": np.testing.assert_almost_equal(result["F"][0][0], 0) - np.testing.assert_almost_equal(result["F"][0][20], 107.1572156700596) - np.testing.assert_almost_equal(result["F"][0][-1], 199.51123480749564) + np.testing.assert_almost_equal(result["F"][0][15], 183.50294099441706) + np.testing.assert_almost_equal(result["F"][0][-1], 208.70219753012756) elif pulse_mode == "triplet": np.testing.assert_almost_equal(result["F"][0][0], 0) - np.testing.assert_almost_equal(result["F"][0][30], 137.72706226851224) - np.testing.assert_almost_equal(result["F"][0][-1], 236.04865519419803) + np.testing.assert_almost_equal(result["F"][0][15], 201.36718367210682) + np.testing.assert_almost_equal(result["F"][0][-1], 242.2696240377996) def test_ivp_methods(): - fes_parameters = {"model": DingModelFrequency(), "frequency": 30, "round_down": True} - ivp_parameters = {"n_shooting": 10, "final_time": 1.25, "use_sx": True} + fes_parameters = { + "model": ding2003_model, + "frequency": 30, + "round_down": True, + } + ivp_parameters = {"final_time": 1.25, "use_sx": True} ivp = IvpFes.from_frequency_and_final_time(fes_parameters, ivp_parameters) - fes_parameters = {"model": DingModelFrequency(), "n_stim": 3, "frequency": 10} - ivp_parameters = {"n_shooting": 10, "use_sx": True} + fes_parameters = {"model": ding2003_model, "n_stim": 3, "frequency": 10} + ivp_parameters = {"use_sx": True} ivp = IvpFes.from_frequency_and_n_stim(fes_parameters, ivp_parameters) @@ -136,50 +153,70 @@ def test_all_ivp_errors(): ), ): IvpFes.from_frequency_and_final_time( - fes_parameters={"model": DingModelFrequency(), "frequency": 30, "round_down": False}, - ivp_parameters={"n_shooting": 1, "final_time": 1.25}, + fes_parameters={ + "model": ding2003_model, + "frequency": 30, + "round_down": False, + }, + ivp_parameters={"final_time": 1.25}, ) with pytest.raises(ValueError, match="Pulse mode not yet implemented"): IvpFes( - fes_parameters={"model": DingModelFrequency(), "n_stim": 3, "pulse_mode": "Quadruplet"}, - ivp_parameters={"n_shooting": 10, "final_time": 0.3}, + fes_parameters={ + "model": ding2003_model, + "stim_time": [0, 0.1, 0.2], + "pulse_mode": "Quadruplet", + }, + ivp_parameters={"final_time": 0.3}, ) - pulse_duration = 0.00001 + pulse_width = 0.00001 with pytest.raises( ValueError, - match=re.escape("Pulse duration must be greater than minimum pulse duration"), + match=re.escape("pulse width must be greater than minimum pulse width"), ): IvpFes( - fes_parameters={"model": DingModelPulseDurationFrequency(), "n_stim": 3, "pulse_duration": pulse_duration}, - ivp_parameters={"n_shooting": 10, "final_time": 0.3}, + fes_parameters={ + "model": ding2007_model, + "stim_time": [0, 0.1, 0.2], + "pulse_width": pulse_width, + }, + ivp_parameters={"final_time": 0.3}, ) - with pytest.raises(ValueError, match="pulse_duration list must have the same length as n_stim"): + with pytest.raises(ValueError, match="pulse_width list must have the same length as n_stim"): IvpFes( fes_parameters={ - "model": DingModelPulseDurationFrequency(), - "n_stim": 3, - "pulse_duration": [0.0003, 0.0004], + "model": ding2007_model, + "stim_time": [0, 0.1, 0.2], + "pulse_width": [0.0003, 0.0004], }, - ivp_parameters={"n_shooting": 10, "final_time": 0.3}, + ivp_parameters={"final_time": 0.3}, ) - pulse_duration = [0.001, 0.0001, 0.003] + pulse_width = [0.001, 0.0001, 0.003] with pytest.raises( ValueError, - match=re.escape("Pulse duration must be greater than minimum pulse duration"), + match=re.escape("pulse width must be greater than minimum pulse width"), ): IvpFes( - fes_parameters={"model": DingModelPulseDurationFrequency(), "n_stim": 3, "pulse_duration": pulse_duration}, - ivp_parameters={"n_shooting": 10, "final_time": 0.3}, + fes_parameters={ + "model": ding2007_model, + "stim_time": [0, 0.1, 0.2], + "pulse_width": pulse_width, + }, + ivp_parameters={"final_time": 0.3}, ) - with pytest.raises(TypeError, match="pulse_duration must be int, float or list type"): + with pytest.raises(TypeError, match="pulse_width must be int, float or list type"): IvpFes( - fes_parameters={"model": DingModelPulseDurationFrequency(), "n_stim": 3, "pulse_duration": True}, - ivp_parameters={"n_shooting": 10, "final_time": 0.3}, + fes_parameters={ + "model": ding2007_model, + "stim_time": [0, 0.1, 0.2], + "pulse_width": True, + }, + ivp_parameters={"final_time": 0.3}, ) pulse_intensity = 0.1 @@ -188,14 +225,22 @@ def test_all_ivp_errors(): match=re.escape("Pulse intensity must be greater than minimum pulse intensity"), ): IvpFes( - fes_parameters={"model": DingModelIntensityFrequency(), "n_stim": 3, "pulse_intensity": pulse_intensity}, - ivp_parameters={"n_shooting": 10, "final_time": 0.3}, + fes_parameters={ + "model": hmed2018_model, + "stim_time": [0, 0.1, 0.2], + "pulse_intensity": pulse_intensity, + }, + ivp_parameters={"final_time": 0.3}, ) with pytest.raises(ValueError, match="pulse_intensity list must have the same length as n_stim"): IvpFes( - fes_parameters={"model": DingModelIntensityFrequency(), "n_stim": 3, "pulse_intensity": [20, 30]}, - ivp_parameters={"n_shooting": 10, "final_time": 0.3}, + fes_parameters={ + "model": hmed2018_model, + "stim_time": [0, 0.1, 0.2], + "pulse_intensity": [20, 30], + }, + ivp_parameters={"final_time": 0.3}, ) pulse_intensity = [20, 30, 0.1] @@ -204,30 +249,47 @@ def test_all_ivp_errors(): match=re.escape("Pulse intensity must be greater than minimum pulse intensity"), ): IvpFes( - fes_parameters={"model": DingModelIntensityFrequency(), "n_stim": 3, "pulse_intensity": pulse_intensity}, - ivp_parameters={"n_shooting": 10, "final_time": 0.3}, + fes_parameters={ + "model": hmed2018_model, + "stim_time": [0, 0.1, 0.2], + "pulse_intensity": pulse_intensity, + }, + ivp_parameters={"final_time": 0.3}, ) with pytest.raises(TypeError, match="pulse_intensity must be int, float or list type"): IvpFes( - fes_parameters={"model": DingModelIntensityFrequency(), "n_stim": 3, "pulse_intensity": True}, - ivp_parameters={"n_shooting": 10, "final_time": 0.3}, + fes_parameters={ + "model": hmed2018_model, + "stim_time": [0, 0.1, 0.2], + "pulse_intensity": True, + }, + ivp_parameters={"final_time": 0.3}, ) with pytest.raises(ValueError, match="ode_solver must be a OdeSolver type"): IvpFes( - fes_parameters={"model": DingModelFrequency(), "n_stim": 3}, - ivp_parameters={"n_shooting": 10, "final_time": 0.3, "ode_solver": None}, + fes_parameters={ + "model": ding2003_model, + "stim_time": [0, 0.1, 0.2], + }, + ivp_parameters={"final_time": 0.3, "ode_solver": None}, ) with pytest.raises(ValueError, match="use_sx must be a bool type"): IvpFes( - fes_parameters={"model": DingModelFrequency(), "n_stim": 3}, - ivp_parameters={"n_shooting": 10, "final_time": 0.3, "use_sx": None}, + fes_parameters={ + "model": ding2003_model, + "stim_time": [0, 0.1, 0.2], + }, + ivp_parameters={"final_time": 0.3, "use_sx": None}, ) with pytest.raises(ValueError, match="n_thread must be a int type"): IvpFes( - fes_parameters={"model": DingModelFrequency(), "n_stim": 3}, - ivp_parameters={"n_shooting": 10, "final_time": 0.3, "n_threads": None}, + fes_parameters={ + "model": ding2003_model, + "stim_time": [0, 0.1, 0.2], + }, + ivp_parameters={"final_time": 0.3, "n_threads": None}, ) diff --git a/tests/shard1/test_models_dynamics_without_bioptim.py b/tests/shard1/test_models_dynamics_without_bioptim.py index b849396..cf0c15f 100644 --- a/tests/shard1/test_models_dynamics_without_bioptim.py +++ b/tests/shard1/test_models_dynamics_without_bioptim.py @@ -1,19 +1,16 @@ import numpy as np from casadi import DM -from cocofest import ( - DingModelFrequencyWithFatigue, - DingModelPulseDurationFrequencyWithFatigue, - DingModelIntensityFrequencyWithFatigue, -) +from cocofest import ModelMaker def test_ding2003_dynamics(): - model = DingModelFrequencyWithFatigue() + model = ModelMaker.create_model("ding2003_with_fatigue", is_approximated=False) assert model.nb_state == 5 assert model.name_dof == ["Cn", "F", "A", "Tau1", "Km"] np.testing.assert_almost_equal( - model.standard_rest_values(), np.array([[0], [0], [model.a_rest], [model.tau1_rest], [model.km_rest]]) + model.standard_rest_values(), + np.array([[0], [0], [model.a_rest], [model.tau1_rest], [model.km_rest]]), ) np.testing.assert_almost_equal( np.array( @@ -30,27 +27,52 @@ def test_ding2003_dynamics(): model.km_rest, ] ), - np.array([0.020, 1.04, -4.0 * 10e-7, 2.1 * 10e-5, 0.060, 127, 1.9 * 10e-8, 3009, 0.050957, 0.103]), + np.array( + [ + 0.020, + 1.04, + -4.0 * 10e-7, + 2.1 * 10e-5, + 0.060, + 127, + 1.9 * 10e-8, + 3009, + 0.050957, + 0.103, + ] + ), ) np.testing.assert_almost_equal( np.array( - model.system_dynamics(cn=5, f=100, a=3009, tau1=0.050957, km=0.103, t=0.11, t_stim_prev=[0, 0.1]) + model.system_dynamics( + cn=5, + f=100, + a=3009, + tau1=0.050957, + km=0.103, + t=0.11, + t_stim_prev=[0, 0.1], + ) ).squeeze(), - np.array(DM([-219.644, 2037.07, -0.0004, 0.021, 1.9e-05])).squeeze(), + np.array(DM([-219.439, 2037.07, -0.0004, 0.021, 1.9e-05])).squeeze(), decimal=3, ) np.testing.assert_almost_equal(model.exp_time_fun(t=0.1, t_stim_i=0.09), 0.6065306597126332) np.testing.assert_almost_equal(model.ri_fun(r0=1.05, time_between_stim=0.1), 1.0003368973499542) - np.testing.assert_almost_equal(model.cn_sum_fun(r0=1.05, t=0.11, t_stim_prev=[0, 0.1]), 0.6067349982845568) - np.testing.assert_almost_equal(model.cn_dot_fun(cn=0, r0=1.05, t=0.11, t_stim_prev=[0, 0.1]), 30.33674991422784) - np.testing.assert_almost_equal(model.f_dot_fun(cn=5, f=100, a=3009, tau1=0.050957, km=0.103), 2037.0703505791284) + cn_sum = model.cn_sum_fun(r0=1.05, t=0.11, t_stim_prev=[0, 0.1], lambda_i=[1, 1]) + np.testing.assert_almost_equal(cn_sum, 0.6108217697230208) + np.testing.assert_almost_equal(model.cn_dot_fun(cn=0, cn_sum=cn_sum), 30.54108848615104) + np.testing.assert_almost_equal( + model.f_dot_fun(cn=5, f=100, a=3009, tau1=0.050957, km=0.103), + 2037.0703505791284, + ) np.testing.assert_almost_equal(model.a_dot_fun(a=5, f=100), 23.653143307086616) np.testing.assert_almost_equal(model.tau1_dot_fun(tau1=0.050957, f=100), 0.021) np.testing.assert_almost_equal(model.km_dot_fun(km=0.103, f=100), 1.8999999999999998e-05) def test_ding2007_dynamics(): - model = DingModelPulseDurationFrequencyWithFatigue() + model = ModelMaker.create_model("ding2007_with_fatigue", is_approximated=False) assert model.nb_state == 5 assert model.name_dof == [ "Cn", @@ -60,7 +82,8 @@ def test_ding2007_dynamics(): "Km", ] np.testing.assert_almost_equal( - model.standard_rest_values(), np.array([[0], [0], [model.a_scale], [model.tau1_rest], [model.km_rest]]) + model.standard_rest_values(), + np.array([[0], [0], [model.a_scale], [model.tau1_rest], [model.km_rest]]), ) np.testing.assert_almost_equal( np.array( @@ -101,7 +124,14 @@ def test_ding2007_dynamics(): np.testing.assert_almost_equal( np.array( model.system_dynamics( - cn=5, f=100, a=4920, tau1=0.050957, km=0.103, t=0.11, t_stim_prev=[0, 0.1], impulse_time=0.0002 + cn=5, + f=100, + a=4920, + tau1=0.050957, + km=0.103, + t=0.11, + t_stim_prev=[0, 0.1], + pulse_width=[0.0002, 0.0002], ) ).squeeze(), np.array(DM([-4.179e02, -4.905e02, -4.000e-04, 2.108e-02, 1.900e-05])).squeeze(), @@ -109,19 +139,27 @@ def test_ding2007_dynamics(): ) np.testing.assert_almost_equal(model.exp_time_fun(t=0.1, t_stim_i=0.09), 0.4028903215291327) np.testing.assert_almost_equal(model.ri_fun(r0=1.05, time_between_stim=0.1), 1.0000056342790253) - np.testing.assert_almost_equal(model.cn_sum_fun(r0=1.05, t=0.11, t_stim_prev=[0, 0.1]), 0.40289259152562124) - np.testing.assert_almost_equal(model.cn_dot_fun(cn=0, r0=1.05, t=0.11, t_stim_prev=[0, 0.1]), 36.626599229601936) - np.testing.assert_almost_equal(model.f_dot_fun(cn=5, f=100, a=3009, tau1=0.050957, km=0.103), 1022.8492662547173) + cn_sum = model.cn_sum_fun(r0=1.05, t=0.11, t_stim_prev=[0, 0.1], lambda_i=[1, 1]) + np.testing.assert_almost_equal(cn_sum, 0.4029379914553837) + np.testing.assert_almost_equal( + model.cn_dot_fun(cn=0, cn_sum=cn_sum), + 36.63072649594398, + ) + np.testing.assert_almost_equal( + model.f_dot_fun(cn=5, f=100, a=3009, tau1=0.050957, km=0.103), + 1022.8492662547173, + ) np.testing.assert_almost_equal(model.a_dot_fun(a=4900, f=100), 0.1570803149606299) np.testing.assert_almost_equal(model.tau1_dot_fun(tau1=0.060601, f=100), 0.021) np.testing.assert_almost_equal(model.km_dot_fun(km=0.103, f=100), 0.000286716535433071) np.testing.assert_almost_equal( - np.array(model.a_calculation(a_scale=4920, impulse_time=0.0002)).squeeze(), np.array(DM(1464.4646488)).squeeze() + np.array(model.a_calculation(a_scale=4920, t=0, t_stim_prev=[0.1], pulse_width=[0.0002])).squeeze(), + np.array(DM(1464.4646488)).squeeze(), ) def test_hmed2018_dynamics(): - model = DingModelIntensityFrequencyWithFatigue() + model = ModelMaker.create_model("hmed2018_with_fatigue", is_approximated=False) assert model.nb_state == 5 assert model.name_dof == [ "Cn", @@ -131,7 +169,8 @@ def test_hmed2018_dynamics(): "Km", ] np.testing.assert_almost_equal( - model.standard_rest_values(), np.array([[0], [0], [model.a_rest], [model.tau1_rest], [model.km_rest]]) + model.standard_rest_values(), + np.array([[0], [0], [model.a_rest], [model.tau1_rest], [model.km_rest]]), ) np.testing.assert_almost_equal( np.array( @@ -174,26 +213,39 @@ def test_hmed2018_dynamics(): np.testing.assert_almost_equal( np.array( model.system_dynamics( - cn=5, f=100, a=3009, tau1=0.050957, km=0.103, t=0.11, t_stim_prev=[0, 0.1], intensity_stim=[30, 50] + cn=5, + f=100, + a=3009, + tau1=0.050957, + km=0.103, + t=0.11, + t_stim_prev=[0, 0.1], + pulse_intensity=[30, 50], ) ).squeeze(), - np.array(DM([-240.654, 2037.07, -0.0004, 0.021, 1.9e-05])).squeeze(), + np.array(DM([-241, 2037.07, -0.0004, 0.021, 1.9e-05])).squeeze(), decimal=3, ) np.testing.assert_almost_equal(model.exp_time_fun(t=0.1, t_stim_i=0.09), 0.6065306597126332) np.testing.assert_almost_equal(model.ri_fun(r0=1.05, time_between_stim=0.1), 1.0003368973499542) + lambda_i = model.get_lambda_i(nb_stim=2, pulse_intensity=[30, 50]) + cn_sum = model.cn_sum_fun(r0=1.05, t=0.11, t_stim_prev=[0, 0.1], lambda_i=lambda_i) + np.testing.assert_almost_equal( + np.array(cn_sum).squeeze(), + np.array(DM(0.1798732)).squeeze(), + ) np.testing.assert_almost_equal( - np.array(model.cn_sum_fun(r0=1.05, t=0.11, t_stim_prev=[0, 0.1], intensity_stim=[30, 50])).squeeze(), - np.array(DM(0.1822978)).squeeze(), + np.array(model.cn_dot_fun(cn=0, cn_sum=cn_sum)).squeeze(), + np.array(DM(8.9936611)).squeeze(), ) np.testing.assert_almost_equal( - np.array(model.cn_dot_fun(cn=0, r0=1.05, t=0.11, t_stim_prev=[0, 0.1], intensity_stim=[30, 50])).squeeze(), - np.array(DM(9.1148913)).squeeze(), + model.f_dot_fun(cn=5, f=100, a=3009, tau1=0.050957, km=0.103), + 2037.0703505791284, ) - np.testing.assert_almost_equal(model.f_dot_fun(cn=5, f=100, a=3009, tau1=0.050957, km=0.103), 2037.0703505791284) np.testing.assert_almost_equal(model.a_dot_fun(a=5, f=100), 23.653143307086616) np.testing.assert_almost_equal(model.tau1_dot_fun(tau1=0.050957, f=100), 0.021) np.testing.assert_almost_equal(model.km_dot_fun(km=0.103, f=100), 1.8999999999999998e-05) np.testing.assert_almost_equal( - np.array(model.lambda_i_calculation(intensity_stim=30)).squeeze(), np.array(DM(0.0799499)).squeeze() + np.array(model.lambda_i_calculation(pulse_intensity=30)).squeeze(), + np.array(DM(0.0799499)).squeeze(), ) diff --git a/tests/shard1/test_nmpc_cyclic.py b/tests/shard1/test_nmpc_cyclic.py index 495aa8e..a0b7e88 100644 --- a/tests/shard1/test_nmpc_cyclic.py +++ b/tests/shard1/test_nmpc_cyclic.py @@ -1,240 +1,64 @@ -import pytest -import re import numpy as np -from bioptim import OdeSolver -from cocofest import OcpFesNmpcCyclic, DingModelPulseDurationFrequencyWithFatigue - - -def test_nmpc_cyclic(): - # --- Build target force --- # - target_time = np.linspace(0, 1, 100) - target_force = abs(np.sin(target_time * np.pi)) * 50 - force_tracking = [target_time, target_force] - - # --- Build nmpc cyclic --- # - n_total_cycles = 6 - n_stim = 10 - n_shooting = 5 - - minimum_pulse_duration = DingModelPulseDurationFrequencyWithFatigue().pd0 - fes_model = DingModelPulseDurationFrequencyWithFatigue(sum_stim_truncation=10) - fes_model.alpha_a = -4.0 * 10e-1 # Increasing the fatigue rate to make the fatigue more visible - - nmpc = OcpFesNmpcCyclic( - model=fes_model, - n_stim=n_stim, - n_shooting=n_shooting, - final_time=1, - pulse_duration={ - "min": minimum_pulse_duration, - "max": 0.0006, - "bimapping": False, - }, - objective={"force_tracking": force_tracking}, - n_total_cycles=n_total_cycles, - n_simultaneous_cycles=3, - n_cycle_to_advance=1, - cycle_to_keep="middle", - use_sx=True, - ode_solver=OdeSolver.COLLOCATION(), - ) - - nmpc.prepare_nmpc() - nmpc.solve() - - # --- Show results --- # - time = [j for sub in nmpc.result["time"] for j in sub] - fatigue = [j for sub in nmpc.result["states"]["A"] for j in sub] - force = [j for sub in nmpc.result["states"]["F"] for j in sub] - - np.testing.assert_almost_equal( - len(time), n_total_cycles * n_stim * n_shooting * (nmpc.ode_solver.polynomial_degree + 1) - ) - np.testing.assert_almost_equal(len(fatigue), len(time)) - np.testing.assert_almost_equal(len(force), len(time)) - - np.testing.assert_almost_equal(time[0], 0.0, decimal=4) - np.testing.assert_almost_equal(fatigue[0], 4796.3120, decimal=4) - np.testing.assert_almost_equal(force[0], 3.0948, decimal=4) - - np.testing.assert_almost_equal(time[750], 3.0, decimal=4) - np.testing.assert_almost_equal(fatigue[750], 4427.2596, decimal=4) - np.testing.assert_almost_equal(force[750], 4.5089, decimal=4) - - np.testing.assert_almost_equal(time[-1], 5.9986, decimal=4) - np.testing.assert_almost_equal(fatigue[-1], 4063.8504, decimal=4) - np.testing.assert_almost_equal(force[-1], 5.6615, decimal=4) - - -def test_all_nmpc_errors(): - with pytest.raises( - TypeError, - match=re.escape("n_total_cycles must be an integer"), - ): - OcpFesNmpcCyclic( - model=DingModelPulseDurationFrequencyWithFatigue(sum_stim_truncation=10), - n_stim=10, - n_shooting=5, - final_time=1, - pulse_duration={ - "min": 0.0003, - "max": 0.0006, - "bimapping": False, - }, - n_total_cycles=None, - ) - - with pytest.raises( - TypeError, - match=re.escape("n_simultaneous_cycles must be an integer"), - ): - OcpFesNmpcCyclic( - model=DingModelPulseDurationFrequencyWithFatigue(sum_stim_truncation=10), - n_stim=10, - n_shooting=5, - final_time=1, - pulse_duration={ - "min": 0.0003, - "max": 0.0006, - "bimapping": False, - }, - n_total_cycles=5, - ) - - with pytest.raises( - TypeError, - match=re.escape("n_cycle_to_advance must be an integer"), - ): - OcpFesNmpcCyclic( - model=DingModelPulseDurationFrequencyWithFatigue(sum_stim_truncation=10), - n_stim=10, - n_shooting=5, - final_time=1, - pulse_duration={ - "min": 0.0003, - "max": 0.0006, - "bimapping": False, - }, - n_total_cycles=5, - n_simultaneous_cycles=3, - ) - - with pytest.raises( - TypeError, - match=re.escape("cycle_to_keep must be a string"), - ): - OcpFesNmpcCyclic( - model=DingModelPulseDurationFrequencyWithFatigue(sum_stim_truncation=10), - n_stim=10, - n_shooting=5, - final_time=1, - pulse_duration={ - "min": 0.0003, - "max": 0.0006, - "bimapping": False, - }, - n_total_cycles=5, - n_simultaneous_cycles=3, - n_cycle_to_advance=1, - ) - - with pytest.raises( - ValueError, - match=re.escape("The number of n_simultaneous_cycles must be higher than the number of n_cycle_to_advance"), - ): - OcpFesNmpcCyclic( - model=DingModelPulseDurationFrequencyWithFatigue(sum_stim_truncation=10), - n_stim=10, - n_shooting=5, - final_time=1, - pulse_duration={ - "min": 0.0003, - "max": 0.0006, - "bimapping": False, - }, - n_total_cycles=5, - n_simultaneous_cycles=3, - n_cycle_to_advance=6, - cycle_to_keep="middle", - ) - - with pytest.raises( - ValueError, - match=re.escape("The number of n_total_cycles must be a multiple of the number n_cycle_to_advance"), - ): - OcpFesNmpcCyclic( - model=DingModelPulseDurationFrequencyWithFatigue(sum_stim_truncation=10), - n_stim=10, - n_shooting=5, - final_time=1, - pulse_duration={ - "min": 0.0003, - "max": 0.0006, - "bimapping": False, - }, - n_total_cycles=5, - n_simultaneous_cycles=3, - n_cycle_to_advance=2, - cycle_to_keep="middle", - ) - - with pytest.raises( - ValueError, - match=re.escape("cycle_to_keep must be either 'first', 'middle' or 'last'"), - ): - OcpFesNmpcCyclic( - model=DingModelPulseDurationFrequencyWithFatigue(sum_stim_truncation=10), - n_stim=10, - n_shooting=5, - final_time=1, - pulse_duration={ - "min": 0.0003, - "max": 0.0006, - "bimapping": False, - }, - n_total_cycles=5, - n_simultaneous_cycles=3, - n_cycle_to_advance=1, - cycle_to_keep="between", - ) - - with pytest.raises( - NotImplementedError, - match=re.escape("Only 'middle' cycle_to_keep is implemented"), - ): - OcpFesNmpcCyclic( - model=DingModelPulseDurationFrequencyWithFatigue(sum_stim_truncation=10), - n_stim=10, - n_shooting=5, - final_time=1, - pulse_duration={ - "min": 0.0003, - "max": 0.0006, - "bimapping": False, - }, - n_total_cycles=5, - n_simultaneous_cycles=3, - n_cycle_to_advance=1, - cycle_to_keep="first", - ) - - with pytest.raises( - NotImplementedError, - match=re.escape("Only 3 simultaneous cycles are implemented yet work in progress"), - ): - OcpFesNmpcCyclic( - model=DingModelPulseDurationFrequencyWithFatigue(sum_stim_truncation=10), - n_stim=10, - n_shooting=5, - final_time=1, - pulse_duration={ - "min": 0.0003, - "max": 0.0006, - "bimapping": False, - }, - n_total_cycles=5, - n_simultaneous_cycles=6, - n_cycle_to_advance=1, - cycle_to_keep="middle", - ) +from bioptim import Solver, SolutionMerge +from cocofest import NmpcFes, DingModelPulseWidthFrequencyWithFatigue + +# +# def test_nmpc_cyclic(): +# # --- Build target force --- # +# target_time = np.linspace(0, 1, 100) +# target_force = abs(np.sin(target_time * np.pi)) * 50 +# force_tracking = [target_time, target_force] +# +# # --- Build nmpc cyclic --- # +# cycles_len = 100 +# cycle_duration = 1 +# +# minimum_pulse_width = DingModelPulseWidthFrequencyWithFatigue().pd0 +# fes_model = DingModelPulseWidthFrequencyWithFatigue() +# fes_model.alpha_a = -4.0 * 10e-1 # Increasing the fatigue rate to make the fatigue more visible +# +# nmpc = NmpcFes.prepare_nmpc( +# model=fes_model, +# stim_time=list(np.round(np.linspace(0, 1, 11), 2))[:-1], +# cycle_len=cycles_len, +# cycle_duration=cycle_duration, +# pulse_width={ +# "min": minimum_pulse_width, +# "max": 0.0006, +# "bimapping": False, +# }, +# objective={"force_tracking": force_tracking}, +# use_sx=True, +# n_threads=6, +# ) +# +# n_cycles_total = 6 +# +# def update_functions(_nmpc, cycle_idx, _sol): +# return cycle_idx < n_cycles_total # True if there are still some cycle to perform +# +# sol = nmpc.solve( +# update_functions, +# solver=Solver.IPOPT(), +# cyclic_options={"states": {}}, +# get_all_iterations=True, +# ) +# sol_merged = sol[0].decision_states(to_merge=[SolutionMerge.PHASES, SolutionMerge.NODES]) +# +# time = sol[0].decision_time(to_merge=SolutionMerge.KEYS, continuous=True) +# time = [float(j) for j in time] +# fatigue = sol_merged["A"][0] +# force = sol_merged["F"][0] +# +# np.testing.assert_almost_equal(time[0], 0.0, decimal=4) +# np.testing.assert_almost_equal(fatigue[0], 4920.0, decimal=4) +# np.testing.assert_almost_equal(force[0], 0, decimal=4) +# +# np.testing.assert_almost_equal(time[300], 3.0, decimal=4) +# np.testing.assert_almost_equal(fatigue[300], 4550.2883, decimal=4) +# np.testing.assert_almost_equal(force[300], 4.1559, decimal=4) +# +# np.testing.assert_almost_equal(time[-1], 6.0, decimal=4) +# np.testing.assert_almost_equal(fatigue[-1], 4184.9710, decimal=4) +# np.testing.assert_almost_equal(force[-1], 5.3672, decimal=4) diff --git a/tests/shard1/test_ocp_build.py b/tests/shard1/test_ocp_build.py index 36e1230..7e03fb5 100644 --- a/tests/shard1/test_ocp_build.py +++ b/tests/shard1/test_ocp_build.py @@ -4,15 +4,7 @@ import numpy as np -from cocofest import ( - DingModelFrequency, - DingModelFrequencyWithFatigue, - DingModelPulseDurationFrequency, - DingModelPulseDurationFrequencyWithFatigue, - DingModelIntensityFrequency, - DingModelIntensityFrequencyWithFatigue, - OcpFes, -) +from cocofest import OcpFes, ModelMaker from bioptim import ObjectiveFcn, ObjectiveList, Node @@ -228,38 +220,74 @@ ) init_force = force - force[0] -init_n_stim = 3 +init_stim_time = [0, 0.1, 0.2] init_final_time = 0.3 init_frequency = 10 -init_n_shooting = 6 +init_n_shooting = 30 init_force_tracking = [time, init_force] init_end_node_tracking = 40 -minimum_pulse_duration = DingModelPulseDurationFrequency().pd0 -minimum_pulse_intensity = ( - np.arctanh(-DingModelIntensityFrequency().cr) / DingModelIntensityFrequency().bs -) + DingModelIntensityFrequency().Is +ding2003 = ModelMaker.create_model("ding2003", is_approximated=False) +ding2003_with_fatigue = ModelMaker.create_model("ding2003_with_fatigue", is_approximated=False) +ding2007 = ModelMaker.create_model("ding2007", is_approximated=False) +ding2007_with_fatigue = ModelMaker.create_model("ding2007_with_fatigue", is_approximated=False) +hmed2018 = ModelMaker.create_model("hmed2018", is_approximated=False) +hmed2018_with_fatigue = ModelMaker.create_model("hmed2018_with_fatigue", is_approximated=False) + + +minimum_pulse_width = ding2007.pd0 +minimum_pulse_intensity = hmed2018.min_pulse_intensity() @pytest.mark.parametrize( "model," - " fixed_pulse_duration," - " pulse_duration_min," - " pulse_duration_max," - " pulse_duration_bimapping," + " fixed_pulse_width," + " pulse_width_min," + " pulse_width_max," + " pulse_width_bimapping," " fixed_pulse_intensity," " pulse_intensity_min," " pulse_intensity_max," " pulse_intensity_bimapping,", [ - (DingModelFrequency(), None, None, None, None, None, None, None, None), - (DingModelFrequencyWithFatigue(), None, None, None, None, None, None, None, None), - (DingModelPulseDurationFrequency(), 0.0002, None, None, None, None, None, None, None), - (DingModelPulseDurationFrequencyWithFatigue(), 0.0002, None, None, None, None, None, None, None), + (ding2003, None, None, None, None, None, None, None, None), + ( + ding2003_with_fatigue, + None, + None, + None, + None, + None, + None, + None, + None, + ), + ( + ding2007, + 0.0002, + None, + None, + None, + None, + None, + None, + None, + ), + ( + ding2007_with_fatigue, + 0.0002, + None, + None, + None, + None, + None, + None, + None, + ), ( - DingModelPulseDurationFrequency(), + ding2007, None, - minimum_pulse_duration, + minimum_pulse_width, 0.0006, False, None, @@ -268,9 +296,9 @@ None, ), ( - DingModelPulseDurationFrequencyWithFatigue(), + ding2007_with_fatigue, None, - minimum_pulse_duration, + minimum_pulse_width, 0.0006, False, None, @@ -278,11 +306,20 @@ None, None, ), - # (DingModelPulseDurationFrequency(), None, minimum_pulse_duration, 0.0006, True, None, None, None, None), parameter mapping not yet implemented - (DingModelIntensityFrequency(), None, None, None, None, 20, None, None, None), - (DingModelIntensityFrequencyWithFatigue(), None, None, None, None, 20, None, None, None), + (hmed2018, None, None, None, None, 20, None, None, None), ( - DingModelIntensityFrequency(), + hmed2018_with_fatigue, + None, + None, + None, + None, + 20, + None, + None, + None, + ), + ( + hmed2018, None, None, None, @@ -293,7 +330,7 @@ False, ), ( - DingModelIntensityFrequencyWithFatigue(), + hmed2018_with_fatigue, None, None, None, @@ -303,28 +340,29 @@ 130, False, ), - # (DingModelIntensityFrequency(), None, None, None, None, None, minimum_pulse_intensity, 130, True), parameter mapping not yet implemented ], ) @pytest.mark.parametrize( - "time_min, time_max, time_bimapping", + "time_min, time_max", [ - (None, None, False), - (0.01, 0.1, False), - (0.01, 0.1, True), + (None, None), + (0.01, 0.1), + (0.01, 0.1), ], ) @pytest.mark.parametrize("use_sx", [True]) # Later add False @pytest.mark.parametrize( - "n_stim, final_time, frequency, n_shooting", [(init_n_stim, init_final_time, init_frequency, init_n_shooting)] + "stim_time, final_time, frequency, n_shooting", + [(init_stim_time, init_final_time, init_frequency, init_n_shooting)], ) @pytest.mark.parametrize( - "force_tracking, end_node_tracking", [(init_force_tracking, None), (None, init_end_node_tracking)] + "force_tracking, end_node_tracking", + [(init_force_tracking, None), (None, init_end_node_tracking)], ) @pytest.mark.parametrize("sum_stim_truncation", [None, 2]) def test_ocp_building( model, - n_stim, + stim_time, n_shooting, final_time, frequency, @@ -332,11 +370,10 @@ def test_ocp_building( end_node_tracking, time_min, time_max, - time_bimapping, - fixed_pulse_duration, - pulse_duration_min, - pulse_duration_max, - pulse_duration_bimapping, + fixed_pulse_width, + pulse_width_min, + pulse_width_max, + pulse_width_bimapping, fixed_pulse_intensity, pulse_intensity_min, pulse_intensity_max, @@ -355,69 +392,14 @@ def test_ocp_building( ocp_1 = OcpFes().prepare_ocp( model=model, - n_shooting=n_shooting, - final_time=final_time, - pulse_event={ - "min": time_min, - "max": time_max, - "bimapping": time_bimapping, - "frequency": frequency, - "round_down": True, - }, - pulse_duration={ - "fixed": fixed_pulse_duration, - "min": pulse_duration_min, - "max": pulse_duration_max, - "bimapping": pulse_duration_bimapping, - }, - pulse_intensity={ - "fixed": fixed_pulse_intensity, - "min": pulse_intensity_min, - "max": pulse_intensity_max, - "bimapping": pulse_intensity_bimapping, - }, - objective={"force_tracking": force_tracking, "end_node_tracking": end_node_tracking}, - use_sx=use_sx, - ) - - ocp_2 = OcpFes().prepare_ocp( - model=model, - n_shooting=n_shooting, - n_stim=n_stim, - pulse_event={ - "min": time_min, - "max": time_max, - "bimapping": time_bimapping, - "frequency": 10, - "round_down": True, - }, - pulse_duration={ - "fixed": fixed_pulse_duration, - "min": pulse_duration_min, - "max": pulse_duration_max, - "bimapping": pulse_duration_bimapping, - }, - pulse_intensity={ - "fixed": fixed_pulse_intensity, - "min": pulse_intensity_min, - "max": pulse_intensity_max, - "bimapping": pulse_intensity_bimapping, - }, - objective={"force_tracking": force_tracking, "end_node_tracking": end_node_tracking}, - use_sx=use_sx, - ) - - ocp_3 = OcpFes().prepare_ocp( - model=model, - n_shooting=n_shooting, - n_stim=n_stim, + stim_time=stim_time, final_time=0.3, - pulse_event={"min": time_min, "max": time_max, "bimapping": time_bimapping}, - pulse_duration={ - "fixed": fixed_pulse_duration, - "min": pulse_duration_min, - "max": pulse_duration_max, - "bimapping": pulse_duration_bimapping, + pulse_event={"min": time_min, "max": time_max}, + pulse_width={ + "fixed": fixed_pulse_width, + "min": pulse_width_min, + "max": pulse_width_max, + "bimapping": pulse_width_bimapping, }, pulse_intensity={ "fixed": fixed_pulse_intensity, @@ -425,19 +407,21 @@ def test_ocp_building( "max": pulse_intensity_max, "bimapping": pulse_intensity_bimapping, }, - objective={"force_tracking": force_tracking, "end_node_tracking": end_node_tracking}, + objective={ + "force_tracking": force_tracking, + "end_node_tracking": end_node_tracking, + }, use_sx=use_sx, ) def test_ding2007_build(): - min_duration = DingModelPulseDurationFrequency().pd0 + min_width = ding2007.pd0 ocp = OcpFes().prepare_ocp( - model=DingModelPulseDurationFrequency(), - n_stim=1, - n_shooting=10, + model=ding2007, + stim_time=[0], final_time=0.1, - pulse_duration={"min": min_duration, "max": 0.005}, + pulse_width={"min": min_width, "max": 0.005}, use_sx=True, ) @@ -445,13 +429,18 @@ def test_ding2007_build(): def test_hmed2018_build(): objective_list = ObjectiveList() objective_list.add( - ObjectiveFcn.Mayer.MINIMIZE_STATE, node=Node.END, key="F", quadratic=True, weight=1, target=100, phase=0 + ObjectiveFcn.Mayer.MINIMIZE_STATE, + node=Node.END, + key="F", + quadratic=True, + weight=1, + target=100, + phase=0, ) - min_intensity = DingModelIntensityFrequency().min_pulse_intensity() + min_intensity = hmed2018.min_pulse_intensity() ocp = OcpFes().prepare_ocp( - model=DingModelIntensityFrequency(), - n_stim=1, - n_shooting=10, + model=hmed2018, + stim_time=[0], final_time=0.1, pulse_intensity={"min": min_intensity, "max": 100}, objective={"custom": objective_list}, @@ -460,174 +449,167 @@ def test_hmed2018_build(): def test_all_ocp_fes_errors(): - with pytest.raises( - TypeError, - match=re.escape( - f"The current model type used is {type(None)}, it must be a FesModel type." - f"Current available models are: DingModelFrequency, DingModelFrequencyWithFatigue," - f"DingModelPulseDurationFrequency, DingModelPulseDurationFrequencyWithFatigue," - f"DingModelIntensityFrequency, DingModelIntensityFrequencyWithFatigue" - ), - ): - OcpFes.prepare_ocp(model=None) - - with pytest.raises(TypeError, match="n_stim must be int type"): - OcpFes.prepare_ocp(model=DingModelFrequency(), n_stim="3") - - with pytest.raises(ValueError, match="n_stim must be positive"): - OcpFes.prepare_ocp(model=DingModelFrequency(), n_stim=-3) - - with pytest.raises(TypeError, match="n_shooting must be int type"): - OcpFes.prepare_ocp(model=DingModelFrequency(), n_stim=3, n_shooting="3") - - with pytest.raises(ValueError, match="n_shooting must be positive"): - OcpFes.prepare_ocp(model=DingModelFrequency(), n_stim=3, n_shooting=-3) - - with pytest.raises(TypeError, match="final_time must be int or float type"): - OcpFes.prepare_ocp(model=DingModelFrequency(), n_stim=3, n_shooting=10, final_time="0.3") - - with pytest.raises(ValueError, match="final_time must be positive"): - OcpFes.prepare_ocp(model=DingModelFrequency(), n_stim=3, n_shooting=10, final_time=-0.3) + # with pytest.raises( + # TypeError, + # match=re.escape( + # f"The current model type used is {type(None)}, it must be a FesModel type." + # f"Current available models are: DingModelFrequency, DingModelFrequencyWithFatigue," + # f"DingModelPulseWidthFrequency, DingModelPulseWidthFrequencyWithFatigue," + # f"DingModelPulseIntensityFrequency, DingModelPulseIntensityFrequencyWithFatigue" + # ), + # ): + # OcpFes.prepare_ocp(model=None) + # + # with pytest.raises(TypeError, match="final_time must be a positive int or float type"): + # OcpFes.prepare_ocp( + # model=ding2003, stim_time=[0, 0.1, 0.2], final_time="0.3" + # ) pulse_mode = "doublet" - with pytest.raises(NotImplementedError, match=re.escape(f"Pulse mode '{pulse_mode}' is not yet implemented")): + with pytest.raises( + NotImplementedError, + match=re.escape(f"Pulse mode '{pulse_mode}' is not yet implemented"), + ): OcpFes.prepare_ocp( - model=DingModelFrequency(), - n_stim=3, - n_shooting=10, + model=ding2003, + stim_time=[0, 0.1, 0.2], final_time=0.3, pulse_event={"pulse_mode": pulse_mode}, ) - with pytest.raises(TypeError, match="frequency must be int or float type"): - OcpFes.prepare_ocp(model=DingModelFrequency(), n_stim=3, n_shooting=10, pulse_event={"frequency": "10"}) - - with pytest.raises(ValueError, match="frequency must be positive"): - OcpFes.prepare_ocp(model=DingModelFrequency(), n_stim=3, n_shooting=10, pulse_event={"frequency": -10}) - - with pytest.raises(ValueError, match="time_min and time_max must be both entered or none of them in order to work"): - OcpFes.prepare_ocp(model=DingModelFrequency(), n_stim=3, n_shooting=10, pulse_event={"min": 0.1}) + with pytest.raises( + ValueError, + match="min and max time event must be both entered or none of them in order to work", + ): + OcpFes.prepare_ocp( + model=ding2003, + stim_time=[0, 0.1, 0.2], + final_time=0.3, + pulse_event={"min": 0.1}, + ) - with pytest.raises(TypeError, match="time_bimapping must be bool type"): + with pytest.raises(TypeError, match=re.escape("time bimapping must be bool type")): OcpFes.prepare_ocp( - model=DingModelFrequency(), - n_stim=3, - n_shooting=10, + model=ding2003, + stim_time=[0, 0.1, 0.2], + final_time=0.3, pulse_event={"min": 0.01, "max": 0.1, "bimapping": "True"}, ) with pytest.raises( - ValueError, match="pulse duration or pulse duration min max bounds need to be set for this model" + ValueError, + match="pulse width or pulse width min max bounds need to be set for this model", ): OcpFes.prepare_ocp( - model=DingModelPulseDurationFrequency(), - n_stim=3, - n_shooting=10, + model=ding2007, + stim_time=[0, 0.1, 0.2], final_time=0.3, - pulse_duration={"min": 0.001}, + pulse_width={"min": 0.001}, ) with pytest.raises( - ValueError, match="Either pulse duration or pulse duration min max bounds need to be set for this model" + ValueError, + match="Either pulse width or pulse width min max bounds need to be set for this model", ): OcpFes.prepare_ocp( - model=DingModelPulseDurationFrequency(), - n_stim=3, - n_shooting=10, + model=ding2007, + stim_time=[0, 0.1, 0.2], final_time=0.3, - pulse_duration={"min": 0.001, "max": 0.005, "fixed": 0.003}, + pulse_width={"min": 0.001, "max": 0.005, "fixed": 0.003}, ) - minimum_pulse_duration = DingModelPulseDurationFrequency().pd0 - fixed_pulse_duration = 0.0001 + minimum_pulse_width = ding2007.pd0 + fixed_pulse_width = 0.0001 with pytest.raises( ValueError, match=re.escape( - f"The pulse duration set ({fixed_pulse_duration})" - f" is lower than minimum duration required." - f" Set a value above {minimum_pulse_duration} seconds " + f"The pulse width set ({fixed_pulse_width})" + f" is lower than minimum width required." + f" Set a value above {minimum_pulse_width} seconds " ), ): OcpFes.prepare_ocp( - model=DingModelPulseDurationFrequency(), - n_stim=3, - n_shooting=10, + model=ding2007, + stim_time=[0, 0.1, 0.2], final_time=0.3, - pulse_duration={"fixed": fixed_pulse_duration}, + pulse_width={"fixed": fixed_pulse_width}, ) - with pytest.raises(TypeError, match="Wrong pulse_duration type, only int or float accepted"): + with pytest.raises(TypeError, match="Wrong pulse_width type, only int or float accepted"): OcpFes.prepare_ocp( - model=DingModelPulseDurationFrequency(), - n_stim=3, - n_shooting=10, + model=ding2007, + stim_time=[0, 0.1, 0.2], final_time=0.3, - pulse_duration={"fixed": "0.001"}, + pulse_width={"fixed": "0.001"}, ) - with pytest.raises(TypeError, match="pulse_duration_min and pulse_duration_max must be int or float type"): + with pytest.raises( + TypeError, + match="min and max pulse width must be int or float type", + ): OcpFes.prepare_ocp( - model=DingModelPulseDurationFrequency(), - n_stim=3, - n_shooting=10, + model=ding2007, + stim_time=[0, 0.1, 0.2], final_time=0.3, - pulse_duration={"min": "0.001", "max": 0.005}, + pulse_width={"min": "0.001", "max": 0.005}, ) - with pytest.raises(ValueError, match="The set minimum pulse duration is higher than maximum pulse duration."): + with pytest.raises( + ValueError, + match="The set minimum pulse width is higher than maximum pulse width.", + ): OcpFes.prepare_ocp( - model=DingModelPulseDurationFrequency(), - n_stim=3, - n_shooting=10, + model=ding2007, + stim_time=[0, 0.1, 0.2], final_time=0.3, - pulse_duration={"min": 0.005, "max": 0.001}, + pulse_width={"min": 0.005, "max": 0.001}, ) - pulse_duration_min = fixed_pulse_duration + pulse_width_min = fixed_pulse_width with pytest.raises( ValueError, match=re.escape( - f"The pulse duration set ({pulse_duration_min})" - f" is lower than minimum duration required." - f" Set a value above {minimum_pulse_duration} seconds " + f"The pulse width set ({pulse_width_min})" + f" is lower than minimum width required." + f" Set a value above {minimum_pulse_width} seconds " ), ): OcpFes.prepare_ocp( - model=DingModelPulseDurationFrequency(), - n_stim=3, - n_shooting=10, + model=ding2007, + stim_time=[0, 0.1, 0.2], final_time=0.3, - pulse_duration={"min": pulse_duration_min, "max": 0.005}, + pulse_width={"min": pulse_width_min, "max": 0.005}, ) with pytest.raises( - ValueError, match="Pulse intensity or pulse intensity min max bounds need to be set for this model" + ValueError, + match="Pulse intensity or pulse intensity min max bounds need to be set for this model", ): - OcpFes.prepare_ocp(model=DingModelIntensityFrequency(), n_stim=3, n_shooting=10, final_time=0.3) + OcpFes.prepare_ocp(model=hmed2018, stim_time=[0, 0.1, 0.2], final_time=0.3) with pytest.raises( - ValueError, match="Either pulse intensity or pulse intensity min max bounds need to be set for this model" + ValueError, + match="Either pulse intensity or pulse intensity min max bounds need to be set for this model", ): OcpFes.prepare_ocp( - model=DingModelIntensityFrequency(), - n_stim=3, - n_shooting=10, + model=hmed2018, + stim_time=[0, 0.1, 0.2], final_time=0.3, pulse_intensity={"min": 20, "max": 100, "fixed": 50}, ) with pytest.raises( - ValueError, match="Pulse intensity or pulse intensity min max bounds need to be set for this model" + ValueError, + match="Pulse intensity or pulse intensity min max bounds need to be set for this model", ): OcpFes.prepare_ocp( - model=DingModelIntensityFrequency(), - n_stim=3, - n_shooting=10, + model=hmed2018, + stim_time=[0, 0.1, 0.2], final_time=0.3, pulse_intensity={"min": 20}, ) - minimum_pulse_intensity = DingModelIntensityFrequency().min_pulse_intensity() + minimum_pulse_intensity = hmed2018.min_pulse_intensity() fixed_pulse_intensity = 1 with pytest.raises( ValueError, @@ -638,36 +620,38 @@ def test_all_ocp_fes_errors(): ), ): OcpFes.prepare_ocp( - model=DingModelIntensityFrequency(), - n_stim=3, - n_shooting=10, + model=hmed2018, + stim_time=[0, 0.1, 0.2], final_time=0.3, pulse_intensity={"fixed": fixed_pulse_intensity}, ) with pytest.raises(TypeError, match="pulse_intensity must be int or float type"): OcpFes.prepare_ocp( - model=DingModelIntensityFrequency(), - n_stim=3, - n_shooting=10, + model=hmed2018, + stim_time=[0, 0.1, 0.2], final_time=0.3, pulse_intensity={"fixed": "20"}, ) - with pytest.raises(TypeError, match="pulse_intensity_min and pulse_intensity_max must be int or float type"): + with pytest.raises( + TypeError, + match="pulse_intensity_min and pulse_intensity_max must be int or float type", + ): OcpFes.prepare_ocp( - model=DingModelIntensityFrequency(), - n_stim=3, - n_shooting=10, + model=hmed2018, + stim_time=[0, 0.1, 0.2], final_time=0.3, pulse_intensity={"min": "20", "max": 100}, ) - with pytest.raises(ValueError, match="The set minimum pulse intensity is higher than maximum pulse intensity."): + with pytest.raises( + ValueError, + match="The set minimum pulse intensity is higher than maximum pulse intensity.", + ): OcpFes.prepare_ocp( - model=DingModelIntensityFrequency(), - n_stim=3, - n_shooting=10, + model=hmed2018, + stim_time=[0, 0.1, 0.2], final_time=0.3, pulse_intensity={"min": 100, "max": 1}, ) @@ -682,47 +666,43 @@ def test_all_ocp_fes_errors(): ), ): OcpFes.prepare_ocp( - model=DingModelIntensityFrequency(), - n_stim=3, - n_shooting=10, + model=hmed2018, + stim_time=[0, 0.1, 0.2], final_time=0.3, pulse_intensity={"min": pulse_intensity_min, "max": 100}, ) with pytest.raises( - ValueError, match="force_tracking time and force argument must be same length and force_tracking " "list size 2" + ValueError, + match="force_tracking time and force argument must be same length and force_tracking " "list size 2", ): OcpFes.prepare_ocp( - model=DingModelFrequency(), - n_stim=3, - n_shooting=10, + model=ding2003, + stim_time=[0, 0.1, 0.2], final_time=0.3, objective={"force_tracking": [np.array([0, 1]), np.array([0, 1, 2])]}, ) with pytest.raises(TypeError, match="force_tracking argument must be np.ndarray type"): OcpFes.prepare_ocp( - model=DingModelFrequency(), - n_stim=3, - n_shooting=10, + model=ding2003, + stim_time=[0, 0.1, 0.2], final_time=0.3, objective={"force_tracking": [[0, 1, 2], np.array([0, 1, 2])]}, ) with pytest.raises(TypeError, match="force_tracking must be list type"): OcpFes.prepare_ocp( - model=DingModelFrequency(), - n_stim=3, - n_shooting=10, + model=ding2003, + stim_time=[0, 0.1, 0.2], final_time=0.3, objective={"force_tracking": np.array([np.array([0, 1, 2]), np.array([0, 1, 2])])}, ) with pytest.raises(TypeError, match="end_node_tracking must be int or float type"): OcpFes.prepare_ocp( - model=DingModelFrequency(), - n_stim=3, - n_shooting=10, + model=ding2003, + stim_time=[0, 0.1, 0.2], final_time=0.3, objective={"end_node_tracking": "10"}, ) @@ -733,59 +713,40 @@ def test_all_ocp_fes_errors(): objective_functions[0].append("objective_function") with pytest.raises(TypeError, match="custom_objective must be a ObjectiveList type"): OcpFes.prepare_ocp( - model=DingModelFrequency(), - n_stim=3, - n_shooting=10, + model=ding2003, + stim_time=[0, 0.1, 0.2], final_time=0.3, objective={"custom": "objective_functions"}, ) with pytest.raises(TypeError, match="All elements in ObjectiveList must be an Objective type"): OcpFes.prepare_ocp( - model=DingModelFrequency(), - n_stim=3, - n_shooting=10, + model=ding2003, + stim_time=[0, 0.1, 0.2], final_time=0.3, objective={"custom": objective_functions}, ) with pytest.raises(TypeError, match="ode_solver must be a OdeSolver type"): - OcpFes.prepare_ocp(model=DingModelFrequency(), n_stim=3, n_shooting=10, final_time=0.3, ode_solver="ode_solver") + OcpFes.prepare_ocp( + model=ding2003, + stim_time=[0, 0.1, 0.2], + final_time=0.3, + ode_solver="ode_solver", + ) with pytest.raises(TypeError, match="use_sx must be a bool type"): - OcpFes.prepare_ocp(model=DingModelFrequency(), n_stim=3, n_shooting=10, final_time=0.3, use_sx="True") - - with pytest.raises(TypeError, match="n_thread must be a int type"): - OcpFes.prepare_ocp(model=DingModelFrequency(), n_stim=3, n_shooting=10, final_time=0.3, n_threads="1") - - with pytest.raises(ValueError, match="At least two variable must be set from n_stim, final_time or frequency"): - OcpFes.prepare_ocp(model=DingModelFrequency(), n_shooting=10, final_time=0.3) - - with pytest.raises( - ValueError, - match=re.escape( - "Can not satisfy n_stim equal to final_time * frequency with the given parameters." - "Consider setting only two of the three parameters" - ), - ): - OcpFes.prepare_ocp(model=DingModelFrequency(), n_stim=3, final_time=0.3, pulse_event={"frequency": 20}) - - with pytest.raises(TypeError, match="round_down must be bool type"): OcpFes.prepare_ocp( - model=DingModelFrequency(), final_time=0.3, pulse_event={"frequency": 20, "round_down": "True"} + model=ding2003, + stim_time=[0, 0.1, 0.2], + final_time=0.3, + use_sx="True", ) - with pytest.raises( - ValueError, - match=re.escape( - "The number of stimulation needs to be integer within the final time t, set round down" - "to True or set final_time * frequency to make the result a integer." - ), - ): + with pytest.raises(TypeError, match="n_thread must be a int type"): OcpFes.prepare_ocp( - model=DingModelIntensityFrequency(), - final_time=0.35, - pulse_event={"frequency": 25}, - n_shooting=10, - pulse_intensity={"min": 20, "max": 100}, + model=ding2003, + stim_time=[0, 0.1, 0.2], + final_time=0.3, + n_threads="1", ) diff --git a/tests/shard1/test_ocp_id.py b/tests/shard1/test_ocp_id.py index f5eb06c..bdee39b 100644 --- a/tests/shard1/test_ocp_id.py +++ b/tests/shard1/test_ocp_id.py @@ -5,23 +5,24 @@ import numpy as np import pytest +from bioptim import ControlType + from cocofest import ( OcpFesId, IvpFes, + ModelMaker, DingModelFrequency, - DingModelPulseDurationFrequency, - DingModelIntensityFrequency, - DingModelFrequencyWithFatigue, + DingModelPulseWidthFrequency, + DingModelPulseIntensityFrequency, DingModelFrequencyForceParameterIdentification, - DingModelPulseDurationFrequencyForceParameterIdentification, DingModelPulseIntensityFrequencyForceParameterIdentification, + DingModelPulseWidthFrequencyForceParameterIdentification, ) -model = DingModelFrequency -stim = [0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9] -discontinuity = [] -n_shooting = [10, 10, 10, 10, 10, 10, 10, 10, 10] -final_time_phase = (0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1) +# model = DingModelFrequencyIntegrate +# stim = [0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9] +# discontinuity = [] +# final_time_phase = (0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1) force_at_node = [ 0.0, 15.854417817548697, @@ -115,47 +116,54 @@ 146.5707678272839, ] +ding2003 = ModelMaker.create_model("ding2003", is_approximated=False) +ding2003_approx = ModelMaker.create_model("ding2003", is_approximated=True) + additional_key_settings = { "a_rest": { "initial_guess": 1000, "min_bound": 1, "max_bound": 10000, - "function": model.set_a_rest, + "function": ding2003.set_a_rest, "scaling": 1, }, "km_rest": { "initial_guess": 0.5, "min_bound": 0.001, "max_bound": 1, - "function": model.set_km_rest, + "function": ding2003.set_km_rest, "scaling": 1, }, "tau1_rest": { "initial_guess": 0.5, "min_bound": 0.0001, "max_bound": 1, - "function": model.set_tau1_rest, + "function": ding2003.set_tau1_rest, "scaling": 1, }, "tau2": { "initial_guess": 0.5, "min_bound": 0.0001, "max_bound": 1, - "function": model.set_tau2, + "function": ding2003.set_tau2, "scaling": 1, }, } -def test_ocp_id_ding2003(): +@pytest.mark.parametrize("model", [ding2003, ding2003_approx]) +def test_ocp_id_ding2003(model): # --- Creating the simulated data to identify on --- # # Building the Initial Value Problem ivp = IvpFes( fes_parameters={ - "model": DingModelFrequency(), - "n_stim": 10, + "model": ModelMaker.create_model("ding2003", is_approximated=False), + "stim_time": [0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9], + }, + ivp_parameters={ + "final_time": 2, + "use_sx": True, }, - ivp_parameters={"n_shooting": 10, "final_time": 1, "use_sx": True, "extend_last_phase_time": 1}, ) # Integrating the solution @@ -168,7 +176,7 @@ def test_ocp_id_ding2003(): dictionary = { "time": time, "force": force, - "stim_time": stim, + "stim_time": [0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9], } pickle_file_name = "temp_identification_simulation.pkl" @@ -177,14 +185,16 @@ def test_ocp_id_ding2003(): # --- Identifying the model parameters --- # ocp = DingModelFrequencyForceParameterIdentification( - model=DingModelFrequency(), + model=model, + final_time=2, data_path=[pickle_file_name], identification_method="full", double_step_identification=False, key_parameter_to_identify=["a_rest", "km_rest", "tau1_rest", "tau2"], additional_key_settings={}, - n_shooting=10, use_sx=True, + n_threads=6, + control_type=ControlType.LINEAR_CONTINUOUS, ) identification_result = ocp.force_model_identification() @@ -192,19 +202,37 @@ def test_ocp_id_ding2003(): # --- Delete the temp file ---# os.remove(f"temp_identification_simulation.pkl") - model = DingModelFrequency() - np.testing.assert_almost_equal(identification_result["a_rest"], model.a_rest, decimal=0) - np.testing.assert_almost_equal(identification_result["km_rest"], model.km_rest, decimal=3) - np.testing.assert_almost_equal(identification_result["tau1_rest"], model.tau1_rest, decimal=3) - np.testing.assert_almost_equal(identification_result["tau2"], model.tau2, decimal=3) + if model.is_approximated: + np.testing.assert_almost_equal(identification_result["a_rest"], 2920, decimal=0) + np.testing.assert_almost_equal(identification_result["km_rest"], 0.1108, decimal=3) + np.testing.assert_almost_equal(identification_result["tau1_rest"], 0.0525, decimal=3) + np.testing.assert_almost_equal(identification_result["tau2"], 0.0534, decimal=3) + else: + tested_model = DingModelFrequency() + np.testing.assert_almost_equal(identification_result["a_rest"], tested_model.a_rest, decimal=0) + np.testing.assert_almost_equal(identification_result["km_rest"], tested_model.km_rest, decimal=3) + np.testing.assert_almost_equal(identification_result["tau1_rest"], tested_model.tau1_rest, decimal=3) + np.testing.assert_almost_equal(identification_result["tau2"], tested_model.tau2, decimal=3) + +ding2007 = ModelMaker.create_model("ding2007", is_approximated=False) +ding2007_approx = ModelMaker.create_model("ding2007", is_approximated=True) -def test_ocp_id_ding2007(): + +@pytest.mark.parametrize("model", [ding2007, ding2007_approx]) +def test_ocp_id_ding2007(model): # --- Creating the simulated data to identify on --- # # Building the Initial Value Problem ivp = IvpFes( - fes_parameters={"model": DingModelPulseDurationFrequency(), "n_stim": 10, "pulse_duration": [0.003] * 10}, - ivp_parameters={"n_shooting": 10, "final_time": 1, "use_sx": True, "extend_last_phase_time": 1}, + fes_parameters={ + "model": ModelMaker.create_model("ding2007", is_approximated=False), + "stim_time": [0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9], + "pulse_width": [0.003] * 10, + }, + ivp_parameters={ + "final_time": 2, + "use_sx": True, + }, ) # Integrating the solution @@ -217,8 +245,8 @@ def test_ocp_id_ding2007(): dictionary = { "time": time, "force": force, - "stim_time": stim, - "pulse_duration": [0.003] * 10, + "stim_time": [0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9], + "pulse_width": [0.003] * 10, } pickle_file_name = "temp_identification_simulation.pkl" @@ -226,15 +254,16 @@ def test_ocp_id_ding2007(): pickle.dump(dictionary, file) # --- Identifying the model parameters --- # - ocp = DingModelPulseDurationFrequencyForceParameterIdentification( - model=DingModelPulseDurationFrequency(), + ocp = DingModelPulseWidthFrequencyForceParameterIdentification( + model=model, + final_time=2, data_path=[pickle_file_name], identification_method="full", double_step_identification=False, key_parameter_to_identify=["tau1_rest", "tau2", "km_rest", "a_scale", "pd0", "pdt"], additional_key_settings={}, - n_shooting=10, use_sx=True, + n_threads=6, ) identification_result = ocp.force_model_identification() @@ -242,21 +271,42 @@ def test_ocp_id_ding2007(): # --- Delete the temp file ---# os.remove(f"temp_identification_simulation.pkl") - model = DingModelPulseDurationFrequency() - np.testing.assert_almost_equal(identification_result["tau1_rest"], model.tau1_rest, decimal=3) - np.testing.assert_almost_equal(identification_result["tau2"], model.tau2, decimal=3) - np.testing.assert_almost_equal(identification_result["km_rest"], model.km_rest, decimal=3) - np.testing.assert_almost_equal(identification_result["a_scale"], model.a_scale, decimal=-2) - np.testing.assert_almost_equal(identification_result["pd0"], model.pd0, decimal=3) - np.testing.assert_almost_equal(identification_result["pdt"], model.pdt, decimal=3) - - -def test_ocp_id_hmed2018(): + if model.is_approximated: + np.testing.assert_almost_equal(identification_result["tau1_rest"], 0.0644, decimal=3) + np.testing.assert_almost_equal(identification_result["tau2"], 0.0493, decimal=3) + np.testing.assert_almost_equal(identification_result["km_rest"], 0.424, decimal=3) + np.testing.assert_almost_equal(identification_result["a_scale"], 6093, decimal=-2) + np.testing.assert_almost_equal(identification_result["pd0"], 0.0003, decimal=3) + np.testing.assert_almost_equal(identification_result["pdt"], 0.0003, decimal=3) + else: + tested_model = DingModelPulseWidthFrequency() + np.testing.assert_almost_equal(identification_result["tau1_rest"], tested_model.tau1_rest, decimal=3) + np.testing.assert_almost_equal(identification_result["tau2"], tested_model.tau2, decimal=3) + np.testing.assert_almost_equal(identification_result["km_rest"], tested_model.km_rest, decimal=3) + np.testing.assert_almost_equal(identification_result["a_scale"], tested_model.a_scale, decimal=-2) + np.testing.assert_almost_equal(identification_result["pd0"], tested_model.pd0, decimal=3) + np.testing.assert_almost_equal(identification_result["pdt"], tested_model.pdt, decimal=3) + + +hmed2018 = ModelMaker.create_model("hmed2018", is_approximated=False) +hmed2018_approx = ModelMaker.create_model("hmed2018", is_approximated=True) + + +@pytest.mark.parametrize("model", [hmed2018, hmed2018_approx]) +def test_ocp_id_hmed2018(model): # --- Creating the simulated data to identify on --- # + # Building the Initial Value Problem ivp = IvpFes( - fes_parameters={"model": DingModelIntensityFrequency(), "n_stim": 10, "pulse_intensity": [50] * 10}, - ivp_parameters={"n_shooting": 10, "final_time": 1, "use_sx": True, "extend_last_phase_time": 1}, + fes_parameters={ + "model": ModelMaker.create_model("hmed2018", is_approximated=False), + "stim_time": [0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9], + "pulse_intensity": list(np.linspace(30, 130, 11))[:-1], + }, + ivp_parameters={ + "final_time": 2, + "use_sx": True, + }, ) # Integrating the solution @@ -269,8 +319,8 @@ def test_ocp_id_hmed2018(): dictionary = { "time": time, "force": force, - "stim_time": stim, - "pulse_intensity": [50] * 10, + "stim_time": [0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9], + "pulse_intensity": list(np.linspace(30, 130, 11))[:-1], } pickle_file_name = "temp_identification_simulation.pkl" @@ -279,14 +329,24 @@ def test_ocp_id_hmed2018(): # --- Identifying the model parameters --- # ocp = DingModelPulseIntensityFrequencyForceParameterIdentification( - model=DingModelIntensityFrequency(), + model=model, + final_time=2, data_path=[pickle_file_name], identification_method="full", double_step_identification=False, - key_parameter_to_identify=["a_rest", "km_rest", "tau1_rest", "tau2", "ar", "bs", "Is", "cr"], + key_parameter_to_identify=[ + "a_rest", + "km_rest", + "tau1_rest", + "tau2", + "ar", + "bs", + "Is", + "cr", + ], additional_key_settings={}, - n_shooting=10, use_sx=True, + n_threads=6, ) identification_result = ocp.force_model_identification() @@ -294,43 +354,26 @@ def test_ocp_id_hmed2018(): # --- Delete the temp file ---# os.remove(f"temp_identification_simulation.pkl") - model = DingModelIntensityFrequency() - np.testing.assert_almost_equal(identification_result["a_rest"], model.a_rest, decimal=0) - np.testing.assert_almost_equal(identification_result["km_rest"], model.km_rest, decimal=3) - np.testing.assert_almost_equal(identification_result["tau1_rest"], model.tau1_rest, decimal=3) - np.testing.assert_almost_equal(identification_result["tau2"], model.tau2, decimal=3) - np.testing.assert_almost_equal(ocp.force_identification_result.cost, 2.99324e-12) - # np.testing.assert_almost_equal(identification_result["ar"], model.ar, decimal=3) - # np.testing.assert_almost_equal(identification_result["bs"], model.bs, decimal=3) - # np.testing.assert_almost_equal(identification_result["Is"], model.Is, decimal=1) - # np.testing.assert_almost_equal(identification_result["cr"], model.cr, decimal=3) + if model.is_approximated: + np.testing.assert_almost_equal(identification_result["tau1_rest"], 0.0499, decimal=3) + np.testing.assert_almost_equal(identification_result["tau2"], 0.0284, decimal=3) + np.testing.assert_almost_equal(identification_result["km_rest"], 0.004, decimal=3) + np.testing.assert_almost_equal(identification_result["ar"], 0.0228, decimal=3) + np.testing.assert_almost_equal(identification_result["bs"], 0.026, decimal=3) + np.testing.assert_almost_equal(identification_result["Is"], 68.732, decimal=3) + np.testing.assert_almost_equal(identification_result["cr"], 0.865, decimal=3) + else: + tested_model = DingModelPulseIntensityFrequency() + np.testing.assert_almost_equal(identification_result["tau1_rest"], tested_model.tau1_rest, decimal=3) + np.testing.assert_almost_equal(identification_result["tau2"], tested_model.tau2, decimal=3) + np.testing.assert_almost_equal(identification_result["km_rest"], 0.174, decimal=3) + np.testing.assert_almost_equal(identification_result["ar"], 0.9913, decimal=3) + np.testing.assert_almost_equal(identification_result["bs"], tested_model.bs, decimal=3) + np.testing.assert_almost_equal(identification_result["Is"], tested_model.Is, decimal=1) + np.testing.assert_almost_equal(identification_result["cr"], tested_model.cr, decimal=3) def test_all_ocp_id_errors(): - n_shooting = "10" - with pytest.raises( - TypeError, - match=re.escape(f"n_shooting must be list type," f" currently n_shooting is {type(n_shooting)}) type."), - ): - OcpFesId.prepare_ocp(model=DingModelFrequency(), n_shooting=n_shooting) - - n_shooting = [10, 10, 10, 10, 10, 10, 10, 10, "10"] - with pytest.raises(TypeError, match=re.escape(f"n_shooting must be list of int type.")): - OcpFesId.prepare_ocp(model=DingModelFrequency(), n_shooting=n_shooting) - - final_time_phase = "0.1" - with pytest.raises( - TypeError, - match=re.escape( - f"final_time_phase must be tuple type," f" currently final_time_phase is {type(final_time_phase)}) type." - ), - ): - OcpFesId.prepare_ocp( - model=DingModelFrequency(), - n_shooting=[10, 10, 10, 10, 10, 10, 10, 10, 10], - final_time_phase=final_time_phase, - ) - force_tracking = "10" with pytest.raises( TypeError, @@ -339,48 +382,23 @@ def test_all_ocp_id_errors(): ), ): OcpFesId.prepare_ocp( - model=DingModelFrequency(), - n_shooting=[10, 10, 10, 10, 10, 10, 10, 10, 10], - final_time_phase=(0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1), - force_tracking=force_tracking, + model=ding2003, + stim_time=[0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9], + final_time=1, + objective={"force_tracking": force_tracking}, ) force_tracking = [10, 10, 10, 10, 10, 10, 10, 10, "10"] with pytest.raises(TypeError, match=re.escape(f"force_tracking must be list of int or float type.")): OcpFesId.prepare_ocp( - model=DingModelFrequency(), - n_shooting=[10, 10, 10, 10, 10, 10, 10, 10, 10], - final_time_phase=(0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1), - force_tracking=force_tracking, + model=ding2003, + stim_time=[0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9], + final_time=1, + objective={"force_tracking": force_tracking}, ) - pulse_duration = 0.001 - with pytest.raises( - TypeError, - match=re.escape(f"pulse_duration must be list type, currently pulse_duration is {type(pulse_duration)}) type."), - ): - OcpFesId.prepare_ocp( - model=DingModelPulseDurationFrequency(), - n_shooting=[10, 10, 10, 10, 10, 10, 10, 10, 10], - final_time_phase=(0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1), - force_tracking=[10, 10, 10, 10, 10, 10, 10, 10, 10], - pulse_duration=pulse_duration, - ) - pulse_intensity = 20 - with pytest.raises( - TypeError, - match=re.escape( - f"pulse_intensity must be list type, currently pulse_intensity is {type(pulse_intensity)}) type." - ), - ): - OcpFesId.prepare_ocp( - model=DingModelIntensityFrequency(), - n_shooting=[10, 10, 10, 10, 10, 10, 10, 10, 10], - final_time_phase=(0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1), - force_tracking=[10, 10, 10, 10, 10, 10, 10, 10, 10], - pulse_intensity=pulse_intensity, - ) +ding2003_with_fatigue = ModelMaker.create_model("ding2003_with_fatigue", is_approximated=False) def test_all_id_program_errors(): @@ -390,7 +408,7 @@ def test_all_id_program_errors(): match="The given model is not valid and should not be including the fatigue equation in the model", ): DingModelFrequencyForceParameterIdentification( - model=DingModelFrequencyWithFatigue(), + model=ding2003_with_fatigue, key_parameter_to_identify=key_parameter_to_identify, ) @@ -401,7 +419,7 @@ def test_all_id_program_errors(): ), ): DingModelFrequencyForceParameterIdentification( - model=DingModelFrequency(), + model=ding2003, data_path=[5], key_parameter_to_identify=key_parameter_to_identify, ) @@ -414,7 +432,7 @@ def test_all_id_program_errors(): ), ): DingModelFrequencyForceParameterIdentification( - model=DingModelFrequency(), + model=ding2003, data_path=["test"], key_parameter_to_identify=key_parameter_to_identify, ) @@ -427,7 +445,7 @@ def test_all_id_program_errors(): ), ): DingModelFrequencyForceParameterIdentification( - model=DingModelFrequency(), + model=ding2003, data_path="test", key_parameter_to_identify=key_parameter_to_identify, ) @@ -440,7 +458,7 @@ def test_all_id_program_errors(): ), ): DingModelFrequencyForceParameterIdentification( - model=DingModelFrequency(), + model=ding2003, data_path=data_path, key_parameter_to_identify=key_parameter_to_identify, ) @@ -456,7 +474,7 @@ def test_all_id_program_errors(): ), ): DingModelFrequencyForceParameterIdentification( - model=DingModelFrequency(), + model=ding2003, identification_method=identification_method, data_path=data_path, key_parameter_to_identify=key_parameter_to_identify, @@ -472,7 +490,7 @@ def test_all_id_program_errors(): ), ): DingModelFrequencyForceParameterIdentification( - model=DingModelFrequency(), + model=ding2003, identification_method=identification_method, data_path=data_path, double_step_identification=double_step_identification, @@ -490,7 +508,7 @@ def test_all_id_program_errors(): ), ): DingModelFrequencyForceParameterIdentification( - model=DingModelFrequency(), + model=ding2003, identification_method=identification_method, data_path=data_path, key_parameter_to_identify=key_parameter_to_identify, @@ -505,16 +523,22 @@ def test_all_id_program_errors(): ), ): DingModelFrequencyForceParameterIdentification( - model=DingModelFrequency(), + model=ding2003, identification_method=identification_method, data_path=data_path, key_parameter_to_identify=key_parameter_to_identify, ) - model = DingModelFrequency() + model = ding2003 key_parameter_to_identify = ["a_rest", "km_rest", "tau1_rest", "tau2"] additional_key_settings = { - "test": {"initial_guess": 1000, "min_bound": 1, "max_bound": 10000, "function": model.set_a_rest, "scaling": 1} + "test": { + "initial_guess": 1000, + "min_bound": 1, + "max_bound": 10000, + "function": model.set_a_rest, + "scaling": 1, + } } with pytest.raises( ValueError, @@ -525,7 +549,7 @@ def test_all_id_program_errors(): ), ): DingModelFrequencyForceParameterIdentification( - model=DingModelFrequency(), + model=ding2003, identification_method=identification_method, data_path=data_path, key_parameter_to_identify=key_parameter_to_identify, @@ -533,7 +557,13 @@ def test_all_id_program_errors(): ) additional_key_settings = { - "a_rest": {"test": 1000, "min_bound": 1, "max_bound": 10000, "function": model.set_a_rest, "scaling": 1} + "a_rest": { + "test": 1000, + "min_bound": 1, + "max_bound": 10000, + "function": model.set_a_rest, + "scaling": 1, + } } with pytest.raises( ValueError, @@ -544,7 +574,7 @@ def test_all_id_program_errors(): ), ): DingModelFrequencyForceParameterIdentification( - model=DingModelFrequency(), + model=ding2003, identification_method=identification_method, data_path=data_path, key_parameter_to_identify=key_parameter_to_identify, @@ -569,7 +599,7 @@ def test_all_id_program_errors(): ), ): DingModelFrequencyForceParameterIdentification( - model=DingModelFrequency(), + model=ding2003, identification_method=identification_method, data_path=data_path, key_parameter_to_identify=key_parameter_to_identify, @@ -585,7 +615,7 @@ def test_all_id_program_errors(): ), ): DingModelFrequencyForceParameterIdentification( - model=DingModelFrequency(), + model=ding2003, identification_method=identification_method, data_path=data_path, key_parameter_to_identify=key_parameter_to_identify, @@ -593,16 +623,3 @@ def test_all_id_program_errors(): ) additional_key_settings = {} - n_shooting = "10" - with pytest.raises( - TypeError, - match=re.escape(f"The given n_shooting must be int type," f" the given value is {type(n_shooting)} type"), - ): - DingModelFrequencyForceParameterIdentification( - model=DingModelFrequency(), - identification_method=identification_method, - data_path=data_path, - key_parameter_to_identify=key_parameter_to_identify, - additional_key_settings=additional_key_settings, - n_shooting=n_shooting, - ) diff --git a/tests/shard2/test_fes_dynamics.py b/tests/shard2/test_fes_dynamics.py index d914e9d..ea68ad8 100644 --- a/tests/shard2/test_fes_dynamics.py +++ b/tests/shard2/test_fes_dynamics.py @@ -11,9 +11,10 @@ ) from cocofest import ( - DingModelPulseDurationFrequencyWithFatigue, - DingModelIntensityFrequencyWithFatigue, + DingModelPulseWidthFrequencyWithFatigue, + DingModelPulseIntensityFrequencyWithFatigue, OcpFesMsk, + FesMskModel, ) from cocofest.examples.msk_models import init as ocp_module @@ -22,70 +23,72 @@ biorbd_model_path = biomodel_folder + "/arm26_biceps_triceps.bioMod" -def test_pulse_duration_multi_muscle_fes_dynamics(): - objective_functions = ObjectiveList() - n_stim = 10 - for i in range(n_stim): - objective_functions.add(ObjectiveFcn.Lagrange.MINIMIZE_CONTROL, key="tau", weight=1, quadratic=True, phase=i) +def test_pulse_width_multi_muscle_fes_dynamics(): + model = FesMskModel( + biorbd_path=biorbd_model_path, + muscles_model=[ + DingModelPulseWidthFrequencyWithFatigue(muscle_name="BIClong", is_approximated=True), + DingModelPulseWidthFrequencyWithFatigue(muscle_name="TRIlong", is_approximated=True), + ], + activate_force_length_relationship=True, + activate_force_velocity_relationship=True, + activate_residual_torque=True, + ) - minimum_pulse_duration = DingModelPulseDurationFrequencyWithFatigue().pd0 + minimum_pulse_width = DingModelPulseWidthFrequencyWithFatigue().pd0 ocp = OcpFesMsk.prepare_ocp( - biorbd_model_path=biorbd_model_path, - bound_type="start_end", - bound_data=[[0, 5], [0, 120]], - fes_muscle_models=[ - DingModelPulseDurationFrequencyWithFatigue(muscle_name="BIClong"), - DingModelPulseDurationFrequencyWithFatigue(muscle_name="TRIlong"), - ], - n_stim=n_stim, - n_shooting=10, + model=model, + stim_time=[0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9], final_time=1, - pulse_duration={ - "min": minimum_pulse_duration, + pulse_width={ + "min": minimum_pulse_width, "max": 0.0006, "bimapping": False, }, - objective={"custom": objective_functions}, - with_residual_torque=True, - activate_force_length_relationship=True, - activate_force_velocity_relationship=True, + objective={"minimize_residual_torque": True}, + msk_info={ + "with_residual_torque": True, + "bound_type": "start_end", + "bound_data": [[0, 5], [0, 120]], + }, use_sx=False, + n_threads=6, ) sol = ocp.solve(Solver.IPOPT(_max_iter=1000)) - np.testing.assert_almost_equal(sol.cost, 2.64645e-08) + np.testing.assert_almost_equal(sol.cost, 864.1739, decimal=3) np.testing.assert_almost_equal( - sol.parameters["pulse_duration_BIClong"], + sol.parameters["pulse_width_BIClong"], np.array( [ - 0.00059643, - 0.00059543, - 0.00059475, - 0.00059325, - 0.0005899, - 0.00058005, - 0.00020145, - 0.0001394, - 0.00015499, - 0.00027584, + 0.00013141, + 0.00060001, + 0.00060001, + 0.00060001, + 0.00060001, + 0.00060001, + 0.00055189, + 0.0001314, + 0.0001314, + 0.0001314, ] ), ) np.testing.assert_almost_equal( - sol.parameters["pulse_duration_TRIlong"], + sol.parameters["pulse_width_TRIlong"], np.array( [ - 0.0001635, - 0.00015817, - 0.00029651, - 0.00048692, - 0.00014608, - 0.0001396, - 0.00013858, - 0.00013775, - 0.000141, - 0.0001786, + 0.0003657, + 0.0003657, + 0.0003657, + 0.0003657, + 0.0003657, + 0.0003657, + 0.0003657, + 0.0003657, + 0.0003657, + 0.0003657, ] ), ) @@ -95,125 +98,118 @@ def test_pulse_duration_multi_muscle_fes_dynamics(): np.testing.assert_almost_equal(sol_states["q"][0][-1], 0) np.testing.assert_almost_equal(sol_states["q"][1][0], 0.08722222222222223) np.testing.assert_almost_equal(sol_states["q"][1][-1], 2.0933333333333333) - np.testing.assert_almost_equal(sol_states["F_BIClong"][0][-1], 25.871305635093197, decimal=4) - np.testing.assert_almost_equal(sol_states["F_TRIlong"][0][-1], 11.572282303524243, decimal=4) + np.testing.assert_almost_equal(sol_states["F_BIClong"][0][-1], 14.604532949917337, decimal=4) + np.testing.assert_almost_equal(sol_states["F_TRIlong"][0][-1], 3.9810503521165272, decimal=4) +# def test_pulse_intensity_multi_muscle_fes_dynamics(): - n_stim = 10 - minimum_pulse_intensity = DingModelIntensityFrequencyWithFatigue.min_pulse_intensity( - DingModelIntensityFrequencyWithFatigue() + model = FesMskModel( + biorbd_path=biorbd_model_path, + muscles_model=[ + DingModelPulseIntensityFrequencyWithFatigue(muscle_name="BIClong", is_approximated=True), + DingModelPulseIntensityFrequencyWithFatigue(muscle_name="TRIlong", is_approximated=True), + ], + activate_force_length_relationship=True, + activate_force_velocity_relationship=True, + activate_residual_torque=True, ) - track_forces = [ - np.array([0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1]), - [np.array([0, 10, 40, 90, 140, 80, 50, 10, 0, 0, 0]), np.array([0, 0, 0, 10, 40, 90, 140, 80, 50, 10, 0])], - ] + + minimum_pulse_intensity = DingModelPulseIntensityFrequencyWithFatigue().min_pulse_intensity() ocp = OcpFesMsk.prepare_ocp( - biorbd_model_path=biorbd_model_path, - bound_type="start", - bound_data=[0, 5], - fes_muscle_models=[ - DingModelIntensityFrequencyWithFatigue(muscle_name="BIClong"), - DingModelIntensityFrequencyWithFatigue(muscle_name="TRIlong"), - ], - n_stim=n_stim, - n_shooting=5, + model=model, + stim_time=[0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9], final_time=1, pulse_intensity={ "min": minimum_pulse_intensity, "max": 130, "bimapping": False, }, - objective={"force_tracking": track_forces}, - with_residual_torque=False, - activate_force_length_relationship=True, - activate_force_velocity_relationship=True, + objective={"minimize_residual_torque": True}, + msk_info={ + "with_residual_torque": True, + "bound_type": "start_end", + "bound_data": [[0, 5], [0, 120]], + }, use_sx=False, + n_threads=6, ) sol = ocp.solve(Solver.IPOPT(_max_iter=1000)) - np.testing.assert_almost_equal(sol.cost, 6666357.331403, decimal=6) + np.testing.assert_almost_equal(sol.cost, 656.123, decimal=3) np.testing.assert_almost_equal( sol.parameters["pulse_intensity_BIClong"], np.array( [ - 130.00000125, - 130.00000126, - 130.00000128, - 130.00000128, - 130.00000121, - 50.86019267, - 17.02854916, - 17.02854916, - 17.02854917, - 17.0285492, + 17.02855017, + 129.99999969, + 130.00000018, + 129.99999886, + 129.99999383, + 129.99997181, + 55.45255047, + 17.02854996, + 17.02854961, + 17.0285512, ] ), + decimal=4, ) np.testing.assert_almost_equal( sol.parameters["pulse_intensity_TRIlong"], np.array( [ - 45.78076277, - 25.84044302, - 59.54858653, - 130.00000124, - 130.00000128, - 130.00000128, - 130.00000122, - 76.08547779, - 17.02854916, - 22.72956196, + 73.51427522, + 73.51427522, + 73.51427522, + 73.51427522, + 73.51427522, + 73.51427522, + 73.51427522, + 73.51427522, + 73.51427522, + 73.51427522, ] ), + decimal=4, ) sol_states = sol.decision_states(to_merge=[SolutionMerge.PHASES, SolutionMerge.NODES]) np.testing.assert_almost_equal(sol_states["q"][0][0], 0) - np.testing.assert_almost_equal(sol_states["q"][0][-1], -0.35378857156156907) - np.testing.assert_almost_equal(sol_states["q"][1][0], 0.08722222222222223) - np.testing.assert_almost_equal(sol_states["q"][1][-1], -7.097935277992009e-10) - np.testing.assert_almost_equal(sol_states["F_BIClong"][0][-1], 47.408594, decimal=4) - np.testing.assert_almost_equal(sol_states["F_TRIlong"][0][-1], 29.131785, decimal=4) + np.testing.assert_almost_equal(sol_states["q"][0][-1], 0) + np.testing.assert_almost_equal(sol_states["q"][1][0], 0.08722, decimal=4) + np.testing.assert_almost_equal(sol_states["q"][1][-1], 2.09333, decimal=4) + np.testing.assert_almost_equal(sol_states["F_BIClong"][0][-1], 13.65894, decimal=4) + np.testing.assert_almost_equal(sol_states["F_TRIlong"][0][-1], 4.832778, decimal=4) def test_fes_models_inputs_sanity_check_errors(): - with pytest.raises( - TypeError, - match=re.escape("biorbd_model_path should be a string"), - ): - ocp = OcpFesMsk.prepare_ocp( - biorbd_model_path=5, - bound_type="start_end", - bound_data=[[0, 5], [0, 120]], - fes_muscle_models=[ - DingModelPulseDurationFrequencyWithFatigue(muscle_name="BIClong"), - DingModelPulseDurationFrequencyWithFatigue(muscle_name="TRIlong"), - ], - n_stim=1, - n_shooting=10, - final_time=1, - pulse_duration={"min": 0.0003, "max": 0.0006}, - ) + model = FesMskModel( + biorbd_path=biorbd_model_path, + muscles_model=[ + DingModelPulseIntensityFrequencyWithFatigue(muscle_name="BIClong"), + DingModelPulseIntensityFrequencyWithFatigue(muscle_name="TRIlong"), + ], + activate_force_length_relationship=True, + activate_force_velocity_relationship=True, + activate_residual_torque=False, + ) with pytest.raises( ValueError, match=re.escape("bound_type should be a string and should be equal to start, end or start_end"), ): ocp = OcpFesMsk.prepare_ocp( - biorbd_model_path=biorbd_model_path, - bound_type="hello", - bound_data=[[0, 5], [0, 120]], - fes_muscle_models=[ - DingModelPulseDurationFrequencyWithFatigue(muscle_name="BIClong"), - DingModelPulseDurationFrequencyWithFatigue(muscle_name="TRIlong"), - ], - n_stim=1, - n_shooting=10, + model=model, + stim_time=[0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9], final_time=1, - pulse_duration={"min": 0.0003, "max": 0.0006}, + pulse_width={"min": 0.0003, "max": 0.0006}, + msk_info={ + "bound_type": "hello", + "bound_data": [[0, 5], [0, 120]], + }, ) with pytest.raises( @@ -221,17 +217,14 @@ def test_fes_models_inputs_sanity_check_errors(): match=re.escape("bound_data should be a list"), ): ocp = OcpFesMsk.prepare_ocp( - biorbd_model_path=biorbd_model_path, - bound_type="start_end", - bound_data="[[0, 5], [0, 120]]", - fes_muscle_models=[ - DingModelPulseDurationFrequencyWithFatigue(muscle_name="BIClong"), - DingModelPulseDurationFrequencyWithFatigue(muscle_name="TRIlong"), - ], - n_stim=1, - n_shooting=10, + model=model, + stim_time=[0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9], final_time=1, - pulse_duration={"min": 0.0003, "max": 0.0006}, + pulse_width={"min": 0.0003, "max": 0.0006}, + msk_info={ + "bound_type": "start_end", + "bound_data": "[[0, 5], [0, 120]]", + }, ) with pytest.raises( @@ -239,17 +232,14 @@ def test_fes_models_inputs_sanity_check_errors(): match=re.escape(f"bound_data should be a list of {2} elements"), ): ocp = OcpFesMsk.prepare_ocp( - biorbd_model_path=biorbd_model_path, - bound_type="start_end", - bound_data=[[0, 5, 7], [0, 120, 150]], - fes_muscle_models=[ - DingModelPulseDurationFrequencyWithFatigue(muscle_name="BIClong"), - DingModelPulseDurationFrequencyWithFatigue(muscle_name="TRIlong"), - ], - n_stim=1, - n_shooting=10, + model=model, + stim_time=[0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9], final_time=1, - pulse_duration={"min": 0.0003, "max": 0.0006}, + pulse_width={"min": 0.0003, "max": 0.0006}, + msk_info={ + "bound_type": "start_end", + "bound_data": [[0, 5, 7], [0, 120, 150]], + }, ) with pytest.raises( @@ -257,35 +247,14 @@ def test_fes_models_inputs_sanity_check_errors(): match=re.escape(f"bound_data should be a list of two list"), ): ocp = OcpFesMsk.prepare_ocp( - biorbd_model_path=biorbd_model_path, - bound_type="start_end", - bound_data=["[0, 5]", [0, 120]], - fes_muscle_models=[ - DingModelPulseDurationFrequencyWithFatigue(muscle_name="BIClong"), - DingModelPulseDurationFrequencyWithFatigue(muscle_name="TRIlong"), - ], - n_stim=1, - n_shooting=10, + model=model, + stim_time=[0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9], final_time=1, - pulse_duration={"min": 0.0003, "max": 0.0006}, - ) - - with pytest.raises( - ValueError, - match=re.escape(f"bound_data should be a list of {2} elements"), - ): - ocp = OcpFesMsk.prepare_ocp( - biorbd_model_path=biorbd_model_path, - bound_type="start_end", - bound_data=[[0, 5, 7], [0, 120, 150]], - fes_muscle_models=[ - DingModelPulseDurationFrequencyWithFatigue(muscle_name="BIClong"), - DingModelPulseDurationFrequencyWithFatigue(muscle_name="TRIlong"), - ], - n_stim=1, - n_shooting=10, - final_time=1, - pulse_duration={"min": 0.0003, "max": 0.0006}, + pulse_width={"min": 0.0003, "max": 0.0006}, + msk_info={ + "bound_type": "start_end", + "bound_data": ["[0, 5]", [0, 120]], + }, ) with pytest.raises( @@ -293,17 +262,14 @@ def test_fes_models_inputs_sanity_check_errors(): match=re.escape(f"bound data index {1}: {5} and {'120'} should be an int or float"), ): ocp = OcpFesMsk.prepare_ocp( - biorbd_model_path=biorbd_model_path, - bound_type="start_end", - bound_data=[[0, 5], [0, "120"]], - fes_muscle_models=[ - DingModelPulseDurationFrequencyWithFatigue(muscle_name="BIClong"), - DingModelPulseDurationFrequencyWithFatigue(muscle_name="TRIlong"), - ], - n_stim=1, - n_shooting=10, + model=model, + stim_time=[0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9], final_time=1, - pulse_duration={"min": 0.0003, "max": 0.0006}, + pulse_width={"min": 0.0003, "max": 0.0006}, + msk_info={ + "bound_type": "start_end", + "bound_data": [[0, 5], [0, "120"]], + }, ) with pytest.raises( @@ -311,17 +277,14 @@ def test_fes_models_inputs_sanity_check_errors(): match=re.escape(f"bound_data should be a list of {2} element"), ): ocp = OcpFesMsk.prepare_ocp( - biorbd_model_path=biorbd_model_path, - bound_type="start", - bound_data=[0, 5, 10], - fes_muscle_models=[ - DingModelPulseDurationFrequencyWithFatigue(muscle_name="BIClong"), - DingModelPulseDurationFrequencyWithFatigue(muscle_name="TRIlong"), - ], - n_stim=1, - n_shooting=10, + model=model, + stim_time=[0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9], final_time=1, - pulse_duration={"min": 0.0003, "max": 0.0006}, + pulse_width={"min": 0.0003, "max": 0.0006}, + msk_info={ + "bound_type": "start", + "bound_data": [0, 5, 10], + }, ) with pytest.raises( @@ -329,322 +292,226 @@ def test_fes_models_inputs_sanity_check_errors(): match=re.escape(f"bound data index {1}: {'5'} should be an int or float"), ): ocp = OcpFesMsk.prepare_ocp( - biorbd_model_path=biorbd_model_path, - bound_type="end", - bound_data=[0, "5"], - fes_muscle_models=[ - DingModelPulseDurationFrequencyWithFatigue(muscle_name="BIClong"), - DingModelPulseDurationFrequencyWithFatigue(muscle_name="TRIlong"), - ], - n_stim=1, - n_shooting=10, - final_time=1, - pulse_duration={"min": 0.0003, "max": 0.0006}, - ) - - with pytest.raises( - TypeError, - match="model must be a FesModel type", - ): - ocp = OcpFesMsk.prepare_ocp( - biorbd_model_path=biorbd_model_path, - bound_type="start", - bound_data=[0, 5], - fes_muscle_models=[ - DingModelPulseDurationFrequencyWithFatigue(muscle_name="BIClong"), - "DingModelPulseDurationFrequencyWithFatigue(muscle_name='TRIlong')", - ], - n_stim=1, - n_shooting=10, - final_time=1, - pulse_duration={"min": 0.0003, "max": 0.0006}, - ) - - with pytest.raises( - TypeError, - match=re.escape(f"force_tracking: {'hello'} must be list type"), - ): - ocp = OcpFesMsk.prepare_ocp( - biorbd_model_path=biorbd_model_path, - bound_type="start", - bound_data=[0, 5], - fes_muscle_models=[ - DingModelPulseDurationFrequencyWithFatigue(muscle_name="BIClong"), - DingModelPulseDurationFrequencyWithFatigue(muscle_name="TRIlong"), - ], - n_stim=1, - n_shooting=10, - final_time=1, - pulse_duration={"min": 0.0003, "max": 0.0006}, - objective={"force_tracking": "hello"}, - ) - - with pytest.raises( - ValueError, - match=re.escape("force_tracking must of size 2"), - ): - ocp = OcpFesMsk.prepare_ocp( - biorbd_model_path=biorbd_model_path, - bound_type="end", - bound_data=[0, 5], - fes_muscle_models=[ - DingModelPulseDurationFrequencyWithFatigue(muscle_name="BIClong"), - DingModelPulseDurationFrequencyWithFatigue(muscle_name="TRIlong"), - ], - n_stim=1, - n_shooting=10, - final_time=1, - pulse_duration={"min": 0.0003, "max": 0.0006}, - objective={"force_tracking": ["hello"]}, - ) - - with pytest.raises( - TypeError, - match=re.escape(f"force_tracking index 0: {'hello'} must be np.ndarray type"), - ): - ocp = OcpFesMsk.prepare_ocp( - biorbd_model_path=biorbd_model_path, - bound_type="start", - bound_data=[0, 5], - fes_muscle_models=[ - DingModelPulseDurationFrequencyWithFatigue(muscle_name="BIClong"), - DingModelPulseDurationFrequencyWithFatigue(muscle_name="TRIlong"), - ], - n_stim=1, - n_shooting=10, - final_time=1, - pulse_duration={"min": 0.0003, "max": 0.0006}, - objective={"force_tracking": ["hello", [1, 2, 3]]}, - ) - - with pytest.raises( - TypeError, - match=re.escape(f"force_tracking index 1: {'[1, 2, 3]'} must be list type"), - ): - ocp = OcpFesMsk.prepare_ocp( - biorbd_model_path=biorbd_model_path, - bound_type="start", - bound_data=[0, 5], - fes_muscle_models=[ - DingModelPulseDurationFrequencyWithFatigue(muscle_name="BIClong"), - DingModelPulseDurationFrequencyWithFatigue(muscle_name="TRIlong"), - ], - n_stim=1, - n_shooting=10, - final_time=1, - pulse_duration={"min": 0.0003, "max": 0.0006}, - objective={"force_tracking": [np.array([1, 2, 3]), "[1, 2, 3]"]}, - ) - - with pytest.raises( - ValueError, - match=re.escape( - "force_tracking index 1 list must have the same size as the number of muscles in fes_muscle_models" - ), - ): - ocp = OcpFesMsk.prepare_ocp( - biorbd_model_path=biorbd_model_path, - bound_type="start", - bound_data=[0, 5], - fes_muscle_models=[ - DingModelPulseDurationFrequencyWithFatigue(muscle_name="BIClong"), - DingModelPulseDurationFrequencyWithFatigue(muscle_name="TRIlong"), - ], - n_stim=1, - n_shooting=10, - final_time=1, - pulse_duration={"min": 0.0003, "max": 0.0006}, - objective={"force_tracking": [np.array([1, 2, 3]), [[1, 2, 3], [1, 2, 3], [1, 2, 3]]]}, - ) - - with pytest.raises( - ValueError, - match=re.escape("force_tracking time and force argument must be the same length"), - ): - ocp = OcpFesMsk.prepare_ocp( - biorbd_model_path=biorbd_model_path, - bound_type="start", - bound_data=[0, 5], - fes_muscle_models=[ - DingModelPulseDurationFrequencyWithFatigue(muscle_name="BIClong"), - DingModelPulseDurationFrequencyWithFatigue(muscle_name="TRIlong"), - ], - n_stim=1, - n_shooting=10, - final_time=1, - pulse_duration={"min": 0.0003, "max": 0.0006}, - objective={"force_tracking": [np.array([1, 2, 3]), [[1, 2, 3], [1, 2]]]}, - ) - - with pytest.raises( - TypeError, - match=re.escape(f"force_tracking: {'hello'} must be list type"), - ): - ocp = OcpFesMsk.prepare_ocp( - biorbd_model_path=biorbd_model_path, - bound_type="start", - bound_data=[0, 5], - fes_muscle_models=[ - DingModelPulseDurationFrequencyWithFatigue(muscle_name="BIClong"), - DingModelPulseDurationFrequencyWithFatigue(muscle_name="TRIlong"), - ], - n_stim=1, - n_shooting=10, - final_time=1, - pulse_duration={"min": 0.0003, "max": 0.0006}, - objective={"end_node_tracking": "hello"}, - ) - - with pytest.raises( - ValueError, - match=re.escape("end_node_tracking list must have the same size as the number of muscles in fes_muscle_models"), - ): - ocp = OcpFesMsk.prepare_ocp( - biorbd_model_path=biorbd_model_path, - bound_type="start", - bound_data=[0, 5], - fes_muscle_models=[ - DingModelPulseDurationFrequencyWithFatigue(muscle_name="BIClong"), - DingModelPulseDurationFrequencyWithFatigue(muscle_name="TRIlong"), - ], - n_stim=1, - n_shooting=10, + model=model, + stim_time=[0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9], final_time=1, - pulse_duration={"min": 0.0003, "max": 0.0006}, - objective={"end_node_tracking": [2, 3, 4]}, + pulse_width={"min": 0.0003, "max": 0.0006}, + msk_info={ + "bound_type": "end", + "bound_data": [0, "5"], + }, ) - with pytest.raises( - TypeError, - match=re.escape(f"end_node_tracking index {1}: {'hello'} must be int or float type"), - ): - ocp = OcpFesMsk.prepare_ocp( - biorbd_model_path=biorbd_model_path, - bound_type="start", - bound_data=[0, 5], - fes_muscle_models=[ - DingModelPulseDurationFrequencyWithFatigue(muscle_name="BIClong"), - DingModelPulseDurationFrequencyWithFatigue(muscle_name="TRIlong"), - ], - n_stim=1, - n_shooting=10, - final_time=1, - pulse_duration={"min": 0.0003, "max": 0.0006}, - objective={"end_node_tracking": [2, "hello"]}, - ) - - with pytest.raises( - TypeError, - match=re.escape("q_tracking should be a list of size 2"), - ): - ocp = OcpFesMsk.prepare_ocp( - biorbd_model_path=biorbd_model_path, - bound_type="start", - bound_data=[0, 5], - fes_muscle_models=[ - DingModelPulseDurationFrequencyWithFatigue(muscle_name="BIClong"), - DingModelPulseDurationFrequencyWithFatigue(muscle_name="TRIlong"), - ], - n_stim=1, - n_shooting=10, - final_time=1, - pulse_duration={"min": 0.0003, "max": 0.0006}, - objective={"q_tracking": "hello"}, - ) - - with pytest.raises( - ValueError, - match=re.escape("q_tracking[0] should be a list"), - ): - ocp = OcpFesMsk.prepare_ocp( - biorbd_model_path=biorbd_model_path, - bound_type="start", - bound_data=[0, 5], - fes_muscle_models=[ - DingModelPulseDurationFrequencyWithFatigue(muscle_name="BIClong"), - DingModelPulseDurationFrequencyWithFatigue(muscle_name="TRIlong"), - ], - n_stim=1, - n_shooting=10, - final_time=1, - pulse_duration={"min": 0.0003, "max": 0.0006}, - objective={"q_tracking": ["hello", [1, 2, 3]]}, - ) - - with pytest.raises( - ValueError, - match=re.escape("q_tracking[1] should have the same size as the number of generalized coordinates"), - ): - ocp = OcpFesMsk.prepare_ocp( - biorbd_model_path=biorbd_model_path, - bound_type="start", - bound_data=[0, 5], - fes_muscle_models=[ - DingModelPulseDurationFrequencyWithFatigue(muscle_name="BIClong"), - DingModelPulseDurationFrequencyWithFatigue(muscle_name="TRIlong"), - ], - n_stim=1, - n_shooting=10, - final_time=1, - pulse_duration={"min": 0.0003, "max": 0.0006}, - objective={"q_tracking": [[1, 2, 3], [1, 2, 3, 4]]}, - ) - - with pytest.raises( - ValueError, - match=re.escape("q_tracking[0] and q_tracking[1] should have the same size"), - ): - ocp = OcpFesMsk.prepare_ocp( - biorbd_model_path=biorbd_model_path, - bound_type="start", - bound_data=[0, 5], - fes_muscle_models=[ - DingModelPulseDurationFrequencyWithFatigue(muscle_name="BIClong"), - DingModelPulseDurationFrequencyWithFatigue(muscle_name="TRIlong"), - ], - n_stim=1, - n_shooting=10, - final_time=1, - pulse_duration={"min": 0.0003, "max": 0.0006}, - objective={"q_tracking": [[1, 2, 3], [[1, 2, 3], [4, 5]]]}, - ) - - with pytest.raises( - TypeError, - match=re.escape(f"{'with_residual_torque'} should be a boolean"), - ): - ocp = OcpFesMsk.prepare_ocp( - biorbd_model_path=biorbd_model_path, - bound_type="start", - bound_data=[0, 5], - fes_muscle_models=[ - DingModelPulseDurationFrequencyWithFatigue(muscle_name="BIClong"), - DingModelPulseDurationFrequencyWithFatigue(muscle_name="TRIlong"), - ], - n_stim=1, - n_shooting=10, - final_time=1, - pulse_duration={"min": 0.0003, "max": 0.0006}, - with_residual_torque="hello", - ) - - -def test_fes_muscle_models_sanity_check_errors(): - with pytest.raises( - ValueError, - match=re.escape( - f"The muscle {'TRIlong'} is not in the fes muscle model " - f"please add it into the fes_muscle_models list by providing the muscle_name =" - f" {'TRIlong'}" - ), - ): - ocp = OcpFesMsk.prepare_ocp( - biorbd_model_path=biorbd_model_path, - bound_type="start", - bound_data=[0, 5], - fes_muscle_models=[DingModelPulseDurationFrequencyWithFatigue(muscle_name="BIClong")], - n_stim=1, - n_shooting=10, - final_time=1, - pulse_duration={"min": 0.0003, "max": 0.0006}, - ) + # + # with pytest.raises( + # TypeError, + # match=re.escape(f"force_tracking index 1: {'[1, 2, 3]'} must be list type"), + # ): + # ocp = OcpFesMsk.prepare_ocp( + # model=model, + # stim_time=[0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9], + # n_shooting=100, + # final_time=1, + # pulse_width={"min": 0.0003, "max": 0.0006}, + # msk_info={"bound_type": "start", + # "bound_data": [0, 5], + # }, + # objective={"force_tracking": [np.array([1, 2, 3]), "[1, 2, 3]"]}, + # ) + # + # with pytest.raises( + # ValueError, + # match=re.escape( + # "force_tracking index 1 list must have the same size as the number of muscles in fes_muscle_models" + # ), + # ): + # ocp = OcpFesMsk.prepare_ocp( + # model=model, + # stim_time=[0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9], + # n_shooting=100, + # final_time=1, + # pulse_width={"min": 0.0003, "max": 0.0006}, + # msk_info={"bound_type": "start", + # "bound_data": [0, 5], + # }, + # objective={ + # "force_tracking": [ + # np.array([1, 2, 3]), + # [[1, 2, 3], [1, 2, 3], [1, 2, 3]], + # ] + # }, + # ) + # + # with pytest.raises( + # ValueError, + # match=re.escape("force_tracking time and force argument must be the same length"), + # ): + # ocp = OcpFesMsk.prepare_ocp( + # model=model, + # stim_time=[0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9], + # n_shooting=100, + # final_time=1, + # pulse_width={"min": 0.0003, "max": 0.0006}, + # msk_info={"bound_type": "start", + # "bound_data": [0, 5], + # }, + # objective={"force_tracking": [np.array([1, 2, 3]), [[1, 2, 3], [1, 2]]]}, + # ) + # + # with pytest.raises( + # TypeError, + # match=re.escape(f"force_tracking: {'hello'} must be list type"), + # ): + # ocp = OcpFesMsk.prepare_ocp( + # model=model, + # stim_time=[0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9], + # n_shooting=100, + # final_time=1, + # pulse_width={"min": 0.0003, "max": 0.0006}, + # msk_info={"bound_type": "start", + # "bound_data": [0, 5], + # }, + # objective={"end_node_tracking": "hello"}, + # ) + # + # with pytest.raises( + # ValueError, + # match=re.escape("end_node_tracking list must have the same size as the number of muscles in fes_muscle_models"), + # ): + # ocp = OcpFesMsk.prepare_ocp( + # model=model, + # stim_time=[0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9], + # n_shooting=100, + # final_time=1, + # pulse_width={"min": 0.0003, "max": 0.0006}, + # msk_info={"bound_type": "start", + # "bound_data": [0, 5], + # }, + # objective={"end_node_tracking": [2, 3, 4]}, + # ) + # + # with pytest.raises( + # TypeError, + # match=re.escape(f"end_node_tracking index {1}: {'hello'} must be int or float type"), + # ): + # ocp = OcpFesMsk.prepare_ocp( + # model=model, + # stim_time=[0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9], + # n_shooting=100, + # final_time=1, + # pulse_width={"min": 0.0003, "max": 0.0006}, + # msk_info={"bound_type": "start", + # "bound_data": [0, 5], + # }, + # objective={"end_node_tracking": [2, "hello"]}, + # ) + # + # with pytest.raises( + # TypeError, + # match=re.escape("q_tracking should be a list of size 2"), + # ): + # ocp = OcpFesMsk.prepare_ocp( + # model=model, + # stim_time=[0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9], + # n_shooting=100, + # final_time=1, + # pulse_width={"min": 0.0003, "max": 0.0006}, + # msk_info={"bound_type": "start", + # "bound_data": [0, 5], + # }, + # objective={"q_tracking": "hello"}, + # ) + # + # with pytest.raises( + # ValueError, + # match=re.escape("q_tracking[0] should be a list"), + # ): + # ocp = OcpFesMsk.prepare_ocp( + # model=model, + # stim_time=[0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9], + # n_shooting=100, + # final_time=1, + # pulse_width={"min": 0.0003, "max": 0.0006}, + # msk_info={"bound_type": "start", + # "bound_data": [0, 5], + # }, + # objective={"q_tracking": ["hello", [1, 2, 3]]}, + # ) + # + # with pytest.raises( + # ValueError, + # match=re.escape("q_tracking[1] should have the same size as the number of generalized coordinates"), + # ): + # ocp = OcpFesMsk.prepare_ocp( + # model=model, + # stim_time=[0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9], + # n_shooting=100, + # final_time=1, + # pulse_width={"min": 0.0003, "max": 0.0006}, + # msk_info={"bound_type": "start", + # "bound_data": [0, 5], + # }, + # objective={"q_tracking": [[1, 2, 3], [1, 2, 3, 4]]}, + # ) + # + # with pytest.raises( + # ValueError, + # match=re.escape("q_tracking[0] and q_tracking[1] should have the same size"), + # ): + # ocp = OcpFesMsk.prepare_ocp( + # model=model, + # stim_time=[0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9], + # n_shooting=100, + # final_time=1, + # pulse_width={"min": 0.0003, "max": 0.0006}, + # msk_info={"bound_type": "start", + # "bound_data": [0, 5], + # }, + # objective={"q_tracking": [[1, 2, 3], [[1, 2, 3], [4, 5]]]}, + # ) + # + # with pytest.raises( + # TypeError, + # match=re.escape(f"{'with_residual_torque'} should be a boolean"), + # ): + # ocp = OcpFesMsk.prepare_ocp( + # model=model, + # stim_time=[0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9], + # n_shooting=100, + # final_time=1, + # pulse_width={"min": 0.0003, "max": 0.0006}, + # msk_info={"bound_type": "start", + # "bound_data": [0, 5], + # "with_residual_torque": "hello"}, + # ) + + +# +# def test_fes_muscle_models_sanity_check_errors(): +# model = FesMskModel(biorbd_path=biorbd_model_path, +# muscles_model=[ +# DingModelPulseIntensityFrequencyWithFatigue(muscle_name="BIClong"), +# ], +# activate_force_length_relationship=True, +# activate_force_velocity_relationship=True, +# activate_residual_torque=False, +# ) +# with pytest.raises( +# ValueError, +# match=re.escape( +# f"The muscle {'TRIlong'} is not in the fes muscle model " +# f"please add it into the fes_muscle_models list by providing the muscle_name =" +# f" {'TRIlong'}" +# ), +# ): +# ocp = OcpFesMsk.prepare_ocp( +# model=model, +# stim_time=[0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9], +# n_shooting=100, +# final_time=1, +# pulse_intensity={"min": 30, "max": 100}, +# msk_info={"with_residual_torque": False, +# "bound_type": "start", +# "bound_data": [0, 5], +# }, +# ) diff --git a/tests/shard2/test_models_dynamics_with_bioptim.py b/tests/shard2/test_models_dynamics_with_bioptim.py index 2d8d672..6dc54c1 100644 --- a/tests/shard2/test_models_dynamics_with_bioptim.py +++ b/tests/shard2/test_models_dynamics_with_bioptim.py @@ -5,8 +5,8 @@ from cocofest import ( DingModelFrequency, DingModelFrequencyWithFatigue, - DingModelPulseDurationFrequency, - DingModelIntensityFrequency, + DingModelPulseWidthFrequency, + DingModelPulseIntensityFrequency, OcpFes, ) @@ -224,25 +224,32 @@ init_force = force - force[0] init_force_tracking = [time, init_force] -minimum_pulse_duration = DingModelPulseDurationFrequency().pd0 -minimum_pulse_intensity = DingModelIntensityFrequency().min_pulse_intensity() +minimum_pulse_width = DingModelPulseWidthFrequency().pd0 +minimum_pulse_intensity = DingModelPulseIntensityFrequency().min_pulse_intensity() @pytest.mark.parametrize("use_sx", [True]) # Later add False @pytest.mark.parametrize( - "model", [DingModelFrequency(), DingModelPulseDurationFrequency(), DingModelIntensityFrequency()] + "model", + [ + DingModelFrequency(), + DingModelPulseWidthFrequency(), + DingModelPulseIntensityFrequency(), + ], ) @pytest.mark.parametrize("force_tracking", [init_force_tracking]) -@pytest.mark.parametrize("min_pulse_duration, min_pulse_intensity", [(minimum_pulse_duration, minimum_pulse_intensity)]) -def test_ocp_output(model, force_tracking, use_sx, min_pulse_duration, min_pulse_intensity): - if isinstance(model, DingModelPulseDurationFrequency): +@pytest.mark.parametrize( + "min_pulse_width, min_pulse_intensity", + [(minimum_pulse_width, minimum_pulse_intensity)], +) +def test_ocp_output(model, force_tracking, use_sx, min_pulse_width, min_pulse_intensity): + if isinstance(model, DingModelPulseWidthFrequency): ocp = OcpFes().prepare_ocp( model=model, - n_shooting=20, - n_stim=10, + stim_time=[0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9], final_time=1, - pulse_duration={ - "min": min_pulse_duration, + pulse_width={ + "min": min_pulse_width, "max": 0.0006, "bimapping": False, }, @@ -256,11 +263,10 @@ def test_ocp_output(model, force_tracking, use_sx, min_pulse_duration, min_pulse # for key in ocp.states.key(): # np.testing.assert_almost_equal(ocp.states[key], pickle_file.states[key]) - elif isinstance(model, DingModelIntensityFrequency): + elif isinstance(model, DingModelPulseIntensityFrequency): ocp = OcpFes().prepare_ocp( model=model, - n_shooting=20, - n_stim=10, + stim_time=[0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9], final_time=1, pulse_intensity={ "min": min_pulse_intensity, @@ -280,8 +286,7 @@ def test_ocp_output(model, force_tracking, use_sx, min_pulse_duration, min_pulse elif isinstance(model, DingModelFrequency): ocp = OcpFes().prepare_ocp( model=model, - n_shooting=20, - n_stim=10, + stim_time=[0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9], final_time=1, pulse_event={"min": 0.01, "max": 1, "bimapping": False}, objective={"end_node_tracking": 50}, @@ -299,12 +304,11 @@ def test_ocp_output(model, force_tracking, use_sx, min_pulse_duration, min_pulse @pytest.mark.parametrize("use_sx", [True]) -@pytest.mark.parametrize("bimapped", [False, True]) +@pytest.mark.parametrize("bimapped", [False]) def test_time_dependent_ocp_output(use_sx, bimapped): ocp = OcpFes().prepare_ocp( model=DingModelFrequencyWithFatigue(), - n_stim=10, - n_shooting=20, + stim_time=[0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9], final_time=1, pulse_event={"min": 0.01, "max": 0.1, "bimapping": bimapped}, objective={"end_node_tracking": 270},