From 48f316bc281afa2d4311ac0d84e954dda29fadba Mon Sep 17 00:00:00 2001 From: p-shg Date: Thu, 30 Jan 2025 11:58:36 +0100 Subject: [PATCH 1/2] Added with_friction support for muscle driven dynamics. Ported funtionalities from torque driven dynamics --- bioptim/dynamics/configure_problem.py | 606 +++++++++++++++++++------ bioptim/dynamics/dynamics_functions.py | 305 ++++++++++--- 2 files changed, 715 insertions(+), 196 deletions(-) diff --git a/bioptim/dynamics/configure_problem.py b/bioptim/dynamics/configure_problem.py index 10231287e..e246f5151 100644 --- a/bioptim/dynamics/configure_problem.py +++ b/bioptim/dynamics/configure_problem.py @@ -95,23 +95,42 @@ def _get_kinematics_based_names(nlp, var_type: str) -> list[str]: The list of str to display on figures """ - idx = nlp.phase_mapping.to_first.map_idx if nlp.phase_mapping else range(nlp.model.nb_q) + idx = ( + nlp.phase_mapping.to_first.map_idx + if nlp.phase_mapping + else range(nlp.model.nb_q) + ) if nlp.model.nb_quaternions == 0: new_names = [nlp.model.name_dof[i] for i in idx] else: new_names = [] for i in nlp.phase_mapping.to_first.map_idx: - if nlp.model.name_dof[i][-4:-1] == "Rot" or nlp.model.name_dof[i][-6:-1] == "Trans": + if ( + nlp.model.name_dof[i][-4:-1] == "Rot" + or nlp.model.name_dof[i][-6:-1] == "Trans" + ): new_names += [nlp.model.name_dof[i]] else: if nlp.model.name_dof[i][-5:] != "QuatW": if var_type == "qdot": - new_names += [nlp.model.name_dof[i][:-5] + "omega" + nlp.model.name_dof[i][-1]] + new_names += [ + nlp.model.name_dof[i][:-5] + + "omega" + + nlp.model.name_dof[i][-1] + ] elif var_type == "qddot": - new_names += [nlp.model.name_dof[i][:-5] + "omegadot" + nlp.model.name_dof[i][-1]] + new_names += [ + nlp.model.name_dof[i][:-5] + + "omegadot" + + nlp.model.name_dof[i][-1] + ] elif var_type == "qdddot": - new_names += [nlp.model.name_dof[i][:-5] + "omegaddot" + nlp.model.name_dof[i][-1]] + new_names += [ + nlp.model.name_dof[i][:-5] + + "omegaddot" + + nlp.model.name_dof[i][-1] + ] elif var_type == "tau" or var_type == "taudot": new_names += [nlp.model.name_dof[i]] @@ -189,13 +208,21 @@ def torque_driven( A list of values to pass to the dynamics at each node. Experimental external forces should be included here. """ - _check_contacts_in_biorbd_model(with_contact, nlp.model.nb_contacts, nlp.phase_idx) - _check_soft_contacts_dynamics(soft_contacts_dynamics, nlp.model.nb_soft_contacts, nlp.phase_idx) + _check_contacts_in_biorbd_model( + with_contact, nlp.model.nb_contacts, nlp.phase_idx + ) + _check_soft_contacts_dynamics( + soft_contacts_dynamics, nlp.model.nb_soft_contacts, nlp.phase_idx + ) # Declared rigidbody states and controls ConfigureProblem.configure_q(ocp, nlp, as_states=True, as_controls=False) - ConfigureProblem.configure_qdot(ocp, nlp, as_states=True, as_controls=False, as_states_dot=True) - ConfigureProblem.configure_tau(ocp, nlp, as_states=False, as_controls=True, fatigue=fatigue) + ConfigureProblem.configure_qdot( + ocp, nlp, as_states=True, as_controls=False, as_states_dot=True + ) + ConfigureProblem.configure_tau( + ocp, nlp, as_states=False, as_controls=True, fatigue=fatigue + ) ConfigureProblem.configure_qddot(ocp, nlp, False, False, True) # Declared soft contacts controls @@ -204,7 +231,9 @@ def torque_driven( # Configure the actual ODE of the dynamics if nlp.dynamics_type.dynamic_function: - ConfigureProblem.configure_dynamics_function(ocp, nlp, DynamicsFunctions.custom) + ConfigureProblem.configure_dynamics_function( + ocp, nlp, DynamicsFunctions.custom + ) else: ConfigureProblem.configure_dynamics_function( ocp, @@ -219,7 +248,9 @@ def torque_driven( # Configure the contact forces if with_contact: - ConfigureProblem.configure_contact_function(ocp, nlp, DynamicsFunctions.forces_from_torque_driven) + ConfigureProblem.configure_contact_function( + ocp, nlp, DynamicsFunctions.forces_from_torque_driven + ) # Configure the soft contact forces ConfigureProblem.configure_soft_contact_function(ocp, nlp) @@ -346,7 +377,9 @@ def torque_driven_free_floating_base( # Configure the actual ODE of the dynamics if nlp.dynamics_type.dynamic_function: - ConfigureProblem.configure_dynamics_function(ocp, nlp, DynamicsFunctions.custom) + ConfigureProblem.configure_dynamics_function( + ocp, nlp, DynamicsFunctions.custom + ) else: ConfigureProblem.configure_dynamics_function( ocp, @@ -393,19 +426,30 @@ def stochastic_torque_driven( n_noised_tau = len(nlp.model.motor_noise_mapping["tau"].to_first.map_idx) else: n_noised_tau = nlp.model.nb_tau - n_noise = nlp.model.motor_noise_magnitude.shape[0] + nlp.model.sensory_noise_magnitude.shape[0] + n_noise = ( + nlp.model.motor_noise_magnitude.shape[0] + + nlp.model.sensory_noise_magnitude.shape[0] + ) n_noised_states = 2 * n_noised_tau # Algebraic states variables ConfigureProblem.configure_stochastic_k( - ocp, nlp, n_noised_controls=n_noised_tau, n_references=nlp.model.n_references + ocp, + nlp, + n_noised_controls=n_noised_tau, + n_references=nlp.model.n_references, + ) + ConfigureProblem.configure_stochastic_ref( + ocp, nlp, n_references=nlp.model.n_references ) - ConfigureProblem.configure_stochastic_ref(ocp, nlp, n_references=nlp.model.n_references) n_collocation_points = 1 if isinstance(problem_type, SocpType.COLLOCATION): n_collocation_points += problem_type.polynomial_degree ConfigureProblem.configure_stochastic_m( - ocp, nlp, n_noised_states=n_noised_states, n_collocation_points=n_collocation_points + ocp, + nlp, + n_noised_states=n_noised_states, + n_collocation_points=n_collocation_points, ) if isinstance(problem_type, SocpType.TRAPEZOIDAL_EXPLICIT): @@ -418,13 +462,21 @@ def stochastic_torque_driven( ) else: if with_cholesky: - ConfigureProblem.configure_stochastic_cholesky_cov(ocp, nlp, n_noised_states=n_noised_states) + ConfigureProblem.configure_stochastic_cholesky_cov( + ocp, nlp, n_noised_states=n_noised_states + ) else: - ConfigureProblem.configure_stochastic_cov_implicit(ocp, nlp, n_noised_states=n_noised_states) + ConfigureProblem.configure_stochastic_cov_implicit( + ocp, nlp, n_noised_states=n_noised_states + ) if isinstance(problem_type, SocpType.TRAPEZOIDAL_IMPLICIT): - ConfigureProblem.configure_stochastic_a(ocp, nlp, n_noised_states=n_noised_states) - ConfigureProblem.configure_stochastic_c(ocp, nlp, n_noised_states=n_noised_states, n_noise=n_noise) + ConfigureProblem.configure_stochastic_a( + ocp, nlp, n_noised_states=n_noised_states + ) + ConfigureProblem.configure_stochastic_c( + ocp, nlp, n_noised_states=n_noised_states, n_noise=n_noise + ) ConfigureProblem.torque_driven( ocp=ocp, @@ -471,19 +523,30 @@ def stochastic_torque_driven_free_floating_base( A list of values to pass to the dynamics at each node. """ n_noised_tau = nlp.model.n_noised_controls - n_noise = nlp.model.motor_noise_magnitude.shape[0] + nlp.model.sensory_noise_magnitude.shape[0] + n_noise = ( + nlp.model.motor_noise_magnitude.shape[0] + + nlp.model.sensory_noise_magnitude.shape[0] + ) n_noised_states = nlp.model.n_noised_states # Stochastic variables ConfigureProblem.configure_stochastic_k( - ocp, nlp, n_noised_controls=n_noised_tau, n_references=nlp.model.n_references + ocp, + nlp, + n_noised_controls=n_noised_tau, + n_references=nlp.model.n_references, + ) + ConfigureProblem.configure_stochastic_ref( + ocp, nlp, n_references=nlp.model.n_references ) - ConfigureProblem.configure_stochastic_ref(ocp, nlp, n_references=nlp.model.n_references) n_collocation_points = 1 if isinstance(problem_type, SocpType.COLLOCATION): n_collocation_points += problem_type.polynomial_degree ConfigureProblem.configure_stochastic_m( - ocp, nlp, n_noised_states=n_noised_states, n_collocation_points=n_collocation_points + ocp, + nlp, + n_noised_states=n_noised_states, + n_collocation_points=n_collocation_points, ) if isinstance(problem_type, SocpType.TRAPEZOIDAL_EXPLICIT): @@ -496,13 +559,21 @@ def stochastic_torque_driven_free_floating_base( ) else: if with_cholesky: - ConfigureProblem.configure_stochastic_cholesky_cov(ocp, nlp, n_noised_states=n_noised_states) + ConfigureProblem.configure_stochastic_cholesky_cov( + ocp, nlp, n_noised_states=n_noised_states + ) else: - ConfigureProblem.configure_stochastic_cov_implicit(ocp, nlp, n_noised_states=n_noised_states) + ConfigureProblem.configure_stochastic_cov_implicit( + ocp, nlp, n_noised_states=n_noised_states + ) if isinstance(problem_type, SocpType.TRAPEZOIDAL_IMPLICIT): - ConfigureProblem.configure_stochastic_a(ocp, nlp, n_noised_states=n_noised_states) - ConfigureProblem.configure_stochastic_c(ocp, nlp, n_noised_states=n_noised_states, n_noise=n_noise) + ConfigureProblem.configure_stochastic_a( + ocp, nlp, n_noised_states=n_noised_states + ) + ConfigureProblem.configure_stochastic_c( + ocp, nlp, n_noised_states=n_noised_states, n_noise=n_noise + ) ConfigureProblem.torque_driven_free_floating_base( ocp=ocp, @@ -548,9 +619,13 @@ def torque_derivative_driven( A list of values to pass to the dynamics at each node. Experimental external forces should be included here. """ - _check_contacts_in_biorbd_model(with_contact, nlp.model.nb_contacts, nlp.phase_idx) + _check_contacts_in_biorbd_model( + with_contact, nlp.model.nb_contacts, nlp.phase_idx + ) - _check_soft_contacts_dynamics(soft_contacts_dynamics, nlp.model.nb_soft_contacts, nlp.phase_idx) + _check_soft_contacts_dynamics( + soft_contacts_dynamics, nlp.model.nb_soft_contacts, nlp.phase_idx + ) ConfigureProblem.configure_q(ocp, nlp, True, False) ConfigureProblem.configure_qdot(ocp, nlp, True, False) @@ -561,7 +636,9 @@ def torque_derivative_driven( ConfigureProblem.configure_soft_contact_forces(ocp, nlp, False, True) if nlp.dynamics_type.dynamic_function: - ConfigureProblem.configure_dynamics_function(ocp, nlp, DynamicsFunctions.custom) + ConfigureProblem.configure_dynamics_function( + ocp, nlp, DynamicsFunctions.custom + ) else: ConfigureProblem.configure_dynamics_function( ocp, @@ -621,7 +698,9 @@ def torque_activations_driven( A list of values to pass to the dynamics at each node. Experimental external forces should be included here. """ - _check_contacts_in_biorbd_model(with_contact, nlp.model.nb_contacts, nlp.phase_idx) + _check_contacts_in_biorbd_model( + with_contact, nlp.model.nb_contacts, nlp.phase_idx + ) ConfigureProblem.configure_q(ocp, nlp, True, False) ConfigureProblem.configure_qdot(ocp, nlp, True, False) @@ -631,7 +710,9 @@ def torque_activations_driven( ConfigureProblem.configure_residual_tau(ocp, nlp, False, True) if nlp.dynamics_type.dynamic_function: - ConfigureProblem.configure_dynamics_function(ocp, nlp, DynamicsFunctions.custom) + ConfigureProblem.configure_dynamics_function( + ocp, nlp, DynamicsFunctions.custom + ) else: ConfigureProblem.configure_dynamics_function( ocp, @@ -669,7 +750,9 @@ def joints_acceleration_driven( A list of values to pass to the dynamics at each node. Experimental external forces should be included here. """ ConfigureProblem.configure_q(ocp, nlp, as_states=True, as_controls=False) - ConfigureProblem.configure_qdot(ocp, nlp, as_states=True, as_controls=False, as_states_dot=True) + ConfigureProblem.configure_qdot( + ocp, nlp, as_states=True, as_controls=False, as_states_dot=True + ) # Configure qddot joints nb_root = nlp.model.nb_root if not nb_root > 0: @@ -686,7 +769,9 @@ def joints_acceleration_driven( as_states_dot=True, ) - name_qddot_joints = [str(i + nb_root) for i in range(nlp.model.nb_qddot - nb_root)] + name_qddot_joints = [ + str(i + nb_root) for i in range(nlp.model.nb_qddot - nb_root) + ] ConfigureProblem.configure_new_variable( "qddot_joints", name_qddot_joints, @@ -697,7 +782,9 @@ def joints_acceleration_driven( as_states_dot=True, ) - ConfigureProblem.configure_dynamics_function(ocp, nlp, DynamicsFunctions.joints_acceleration_driven) + ConfigureProblem.configure_dynamics_function( + ocp, nlp, DynamicsFunctions.joints_acceleration_driven + ) @staticmethod def muscle_driven( @@ -709,6 +796,7 @@ def muscle_driven( with_contact: bool = False, with_passive_torque: bool = False, with_ligament: bool = False, + with_friction: bool = False, numerical_data_timeseries: dict[str, np.ndarray] = None, ): """ @@ -736,13 +824,19 @@ def muscle_driven( If the dynamic with passive torque should be used with_ligament: bool If the dynamic with ligament should be used + with_friction: bool + If the dynamic with joint friction should be used (friction = coefficients * qdot) numerical_data_timeseries: dict[str, np.ndarray] A list of values to pass to the dynamics at each node. Experimental external forces should be included here. """ - _check_contacts_in_biorbd_model(with_contact, nlp.model.nb_contacts, nlp.phase_idx) + _check_contacts_in_biorbd_model( + with_contact, nlp.model.nb_contacts, nlp.phase_idx + ) if fatigue is not None and "tau" in fatigue and not with_residual_torque: - raise RuntimeError("Residual torques need to be used to apply fatigue on torques") + raise RuntimeError( + "Residual torques need to be used to apply fatigue on torques" + ) ConfigureProblem.configure_q(ocp, nlp, True, False) ConfigureProblem.configure_qdot(ocp, nlp, True, False, True) @@ -750,10 +844,14 @@ def muscle_driven( if with_residual_torque: ConfigureProblem.configure_tau(ocp, nlp, False, True, fatigue=fatigue) - ConfigureProblem.configure_muscles(ocp, nlp, with_excitations, True, fatigue=fatigue) + ConfigureProblem.configure_muscles( + ocp, nlp, with_excitations, True, fatigue=fatigue + ) if nlp.dynamics_type.dynamic_function: - ConfigureProblem.configure_dynamics_function(ocp, nlp, DynamicsFunctions.custom) + ConfigureProblem.configure_dynamics_function( + ocp, nlp, DynamicsFunctions.custom + ) else: ConfigureProblem.configure_dynamics_function( ocp, @@ -764,6 +862,7 @@ def muscle_driven( with_residual_torque=with_residual_torque, with_passive_torque=with_passive_torque, with_ligament=with_ligament, + with_friction=with_friction, ) if with_contact: @@ -775,7 +874,9 @@ def muscle_driven( ConfigureProblem.configure_soft_contact_function(ocp, nlp) @staticmethod - def holonomic_torque_driven(ocp, nlp, numerical_data_timeseries: dict[str, np.ndarray] = None): + def holonomic_torque_driven( + ocp, nlp, numerical_data_timeseries: dict[str, np.ndarray] = None + ): """ Tell the program which variables are states and controls. @@ -823,12 +924,18 @@ def holonomic_torque_driven(ocp, nlp, numerical_data_timeseries: dict[str, np.nd # extra plots ConfigureProblem.configure_qv(ocp, nlp, nlp.model.compute_q_v) ConfigureProblem.configure_qdotv(ocp, nlp, nlp.model._compute_qdot_v) - ConfigureProblem.configure_lagrange_multipliers_function(ocp, nlp, nlp.model.compute_the_lagrangian_multipliers) + ConfigureProblem.configure_lagrange_multipliers_function( + ocp, nlp, nlp.model.compute_the_lagrangian_multipliers + ) - ConfigureProblem.configure_dynamics_function(ocp, nlp, DynamicsFunctions.holonomic_torque_driven) + ConfigureProblem.configure_dynamics_function( + ocp, nlp, DynamicsFunctions.holonomic_torque_driven + ) @staticmethod - def configure_lagrange_multipliers_function(ocp, nlp, dyn_func: Callable, **extra_params): + def configure_lagrange_multipliers_function( + ocp, nlp, dyn_func: Callable, **extra_params + ): """ Configure the contact points @@ -855,8 +962,12 @@ def configure_lagrange_multipliers_function(ocp, nlp, dyn_func: Callable, **extr ], [ dyn_func()( - nlp.get_var_from_states_or_controls("q_u", nlp.states.scaled.cx, nlp.controls.scaled.cx), - nlp.get_var_from_states_or_controls("qdot_u", nlp.states.scaled.cx, nlp.controls.scaled.cx), + nlp.get_var_from_states_or_controls( + "q_u", nlp.states.scaled.cx, nlp.controls.scaled.cx + ), + nlp.get_var_from_states_or_controls( + "qdot_u", nlp.states.scaled.cx, nlp.controls.scaled.cx + ), DM.zeros(nlp.model.nb_dependent_joints, 1), DynamicsFunctions.get(nlp.controls["tau"], nlp.controls.scaled.cx), ) @@ -867,20 +978,39 @@ def configure_lagrange_multipliers_function(ocp, nlp, dyn_func: Callable, **extr all_multipliers_names = [] for nlp_i in ocp.nlp: - if hasattr(nlp_i.model, "has_holonomic_constraints"): # making sure we have a HolonomicBiorbdModel - nlp_i_multipliers_names = [nlp_i.model.name_dof[i] for i in nlp_i.model.dependent_joint_index] + if hasattr( + nlp_i.model, "has_holonomic_constraints" + ): # making sure we have a HolonomicBiorbdModel + nlp_i_multipliers_names = [ + nlp_i.model.name_dof[i] for i in nlp_i.model.dependent_joint_index + ] all_multipliers_names.extend( - [name for name in nlp_i_multipliers_names if name not in all_multipliers_names] + [ + name + for name in nlp_i_multipliers_names + if name not in all_multipliers_names + ] ) - all_multipliers_names = [f"lagrange_multiplier_{name}" for name in all_multipliers_names] + all_multipliers_names = [ + f"lagrange_multiplier_{name}" for name in all_multipliers_names + ] all_multipliers_names_in_phase = [ - f"lagrange_multiplier_{nlp.model.name_dof[i]}" for i in nlp.model.dependent_joint_index + f"lagrange_multiplier_{nlp.model.name_dof[i]}" + for i in nlp.model.dependent_joint_index ] axes_idx = BiMapping( - to_first=[i for i, c in enumerate(all_multipliers_names) if c in all_multipliers_names_in_phase], - to_second=[i for i, c in enumerate(all_multipliers_names) if c in all_multipliers_names_in_phase], + to_first=[ + i + for i, c in enumerate(all_multipliers_names) + if c in all_multipliers_names_in_phase + ], + to_second=[ + i + for i, c in enumerate(all_multipliers_names) + if c in all_multipliers_names_in_phase + ], ) nlp.plot["lagrange_multipliers"] = CustomPlot( @@ -920,7 +1050,9 @@ def configure_qv(ocp, nlp, dyn_func: Callable, **extra_params): ], [ dyn_func()( - nlp.get_var_from_states_or_controls("q_u", nlp.states.cx, nlp.controls.cx), + nlp.get_var_from_states_or_controls( + "q_u", nlp.states.cx, nlp.controls.cx + ), DM.zeros(nlp.model.nb_dependent_joints, 1), ) ], @@ -930,16 +1062,34 @@ def configure_qv(ocp, nlp, dyn_func: Callable, **extra_params): all_multipliers_names = [] for nlp_i in ocp.nlp: - if hasattr(nlp_i.model, "has_holonomic_constraints"): # making sure we have a HolonomicBiorbdModel - nlp_i_multipliers_names = [nlp_i.model.name_dof[i] for i in nlp_i.model.dependent_joint_index] + if hasattr( + nlp_i.model, "has_holonomic_constraints" + ): # making sure we have a HolonomicBiorbdModel + nlp_i_multipliers_names = [ + nlp_i.model.name_dof[i] for i in nlp_i.model.dependent_joint_index + ] all_multipliers_names.extend( - [name for name in nlp_i_multipliers_names if name not in all_multipliers_names] + [ + name + for name in nlp_i_multipliers_names + if name not in all_multipliers_names + ] ) - all_multipliers_names_in_phase = [nlp.model.name_dof[i] for i in nlp.model.dependent_joint_index] + all_multipliers_names_in_phase = [ + nlp.model.name_dof[i] for i in nlp.model.dependent_joint_index + ] axes_idx = BiMapping( - to_first=[i for i, c in enumerate(all_multipliers_names) if c in all_multipliers_names_in_phase], - to_second=[i for i, c in enumerate(all_multipliers_names) if c in all_multipliers_names_in_phase], + to_first=[ + i + for i, c in enumerate(all_multipliers_names) + if c in all_multipliers_names_in_phase + ], + to_second=[ + i + for i, c in enumerate(all_multipliers_names) + if c in all_multipliers_names_in_phase + ], ) nlp.plot["q_v"] = CustomPlot( @@ -979,8 +1129,12 @@ def configure_qdotv(ocp, nlp, dyn_func: Callable, **extra_params): ], [ dyn_func()( - nlp.get_var_from_states_or_controls("q_u", nlp.states.scaled.cx, nlp.controls.scaled.cx), - nlp.get_var_from_states_or_controls("qdot_u", nlp.states.scaled.cx, nlp.controls.scaled.cx), + nlp.get_var_from_states_or_controls( + "q_u", nlp.states.scaled.cx, nlp.controls.scaled.cx + ), + nlp.get_var_from_states_or_controls( + "qdot_u", nlp.states.scaled.cx, nlp.controls.scaled.cx + ), DM.zeros(nlp.model.nb_dependent_joints, 1), ) ], @@ -990,16 +1144,34 @@ def configure_qdotv(ocp, nlp, dyn_func: Callable, **extra_params): all_multipliers_names = [] for nlp_i in ocp.nlp: - if hasattr(nlp_i.model, "has_holonomic_constraints"): # making sure we have a HolonomicBiorbdModel - nlp_i_multipliers_names = [nlp_i.model.name_dof[i] for i in nlp_i.model.dependent_joint_index] + if hasattr( + nlp_i.model, "has_holonomic_constraints" + ): # making sure we have a HolonomicBiorbdModel + nlp_i_multipliers_names = [ + nlp_i.model.name_dof[i] for i in nlp_i.model.dependent_joint_index + ] all_multipliers_names.extend( - [name for name in nlp_i_multipliers_names if name not in all_multipliers_names] + [ + name + for name in nlp_i_multipliers_names + if name not in all_multipliers_names + ] ) - all_multipliers_names_in_phase = [nlp.model.name_dof[i] for i in nlp.model.dependent_joint_index] + all_multipliers_names_in_phase = [ + nlp.model.name_dof[i] for i in nlp.model.dependent_joint_index + ] axes_idx = BiMapping( - to_first=[i for i, c in enumerate(all_multipliers_names) if c in all_multipliers_names_in_phase], - to_second=[i for i, c in enumerate(all_multipliers_names) if c in all_multipliers_names_in_phase], + to_first=[ + i + for i, c in enumerate(all_multipliers_names) + if c in all_multipliers_names_in_phase + ], + to_second=[ + i + for i, c in enumerate(all_multipliers_names) + if c in all_multipliers_names_in_phase + ], ) nlp.plot["qdot_v"] = CustomPlot( @@ -1173,19 +1345,37 @@ def configure_contact_function(ocp, nlp, dyn_func: Callable, **extra_params): all_contact_names = [] for elt in ocp.nlp: - all_contact_names.extend([name for name in elt.model.contact_names if name not in all_contact_names]) + all_contact_names.extend( + [ + name + for name in elt.model.contact_names + if name not in all_contact_names + ] + ) if "contact_forces" in nlp.plot_mapping: contact_names_in_phase = [name for name in nlp.model.contact_names] axes_idx = BiMapping( to_first=nlp.plot_mapping["contact_forces"].map_idx, - to_second=[i for i, c in enumerate(all_contact_names) if c in contact_names_in_phase], + to_second=[ + i + for i, c in enumerate(all_contact_names) + if c in contact_names_in_phase + ], ) else: contact_names_in_phase = [name for name in nlp.model.contact_names] axes_idx = BiMapping( - to_first=[i for i, c in enumerate(all_contact_names) if c in contact_names_in_phase], - to_second=[i for i, c in enumerate(all_contact_names) if c in contact_names_in_phase], + to_first=[ + i + for i, c in enumerate(all_contact_names) + if c in contact_names_in_phase + ], + to_second=[ + i + for i, c in enumerate(all_contact_names) + if c in contact_names_in_phase + ], ) nlp.plot["contact_forces"] = CustomPlot( @@ -1249,7 +1439,11 @@ def configure_soft_contact_function(ocp, nlp): ] phase_mappings = BiMapping( to_first=nlp.plot_mapping["soft_contact_forces"].map_idx, - to_second=[i for i, c in enumerate(all_soft_contact_names) if c in soft_contact_names_in_phase], + to_second=[ + i + for i, c in enumerate(all_soft_contact_names) + if c in soft_contact_names_in_phase + ], ) else: soft_contact_names_in_phase = [ @@ -1258,16 +1452,33 @@ def configure_soft_contact_function(ocp, nlp): if nlp.model.soft_contact_names[i_sc] not in all_soft_contact_names ] phase_mappings = BiMapping( - to_first=[i for i, c in enumerate(all_soft_contact_names) if c in soft_contact_names_in_phase], - to_second=[i for i, c in enumerate(all_soft_contact_names) if c in soft_contact_names_in_phase], + to_first=[ + i + for i, c in enumerate(all_soft_contact_names) + if c in soft_contact_names_in_phase + ], + to_second=[ + i + for i, c in enumerate(all_soft_contact_names) + if c in soft_contact_names_in_phase + ], + ) + nlp.plot[f"soft_contact_forces_{nlp.model.soft_contact_names[i_sc]}"] = ( + CustomPlot( + lambda t0, phases_dt, node_idx, x, u, p, a, d: nlp.soft_contact_forces_func( + np.concatenate([t0, t0 + phases_dt[nlp.phase_idx]]), + x, + u, + p, + a, + d, + )[ + (i_sc * 6) : ((i_sc + 1) * 6), : + ], + plot_type=PlotType.INTEGRATED, + axes_idx=phase_mappings, + legend=all_soft_contact_names, ) - nlp.plot[f"soft_contact_forces_{nlp.model.soft_contact_names[i_sc]}"] = CustomPlot( - lambda t0, phases_dt, node_idx, x, u, p, a, d: nlp.soft_contact_forces_func( - np.concatenate([t0, t0 + phases_dt[nlp.phase_idx]]), x, u, p, a, d - )[(i_sc * 6) : ((i_sc + 1) * 6), :], - plot_type=PlotType.INTEGRATED, - axes_idx=phase_mappings, - legend=all_soft_contact_names, ) @staticmethod @@ -1363,7 +1574,11 @@ def configure_integrated_value( # TODO: compute values at collocation points # but for now only cx_start can be used - n_cx = nlp.ode_solver.n_cx - 1 if isinstance(nlp.ode_solver, OdeSolver.COLLOCATION) else 3 + n_cx = ( + nlp.ode_solver.n_cx - 1 + if isinstance(nlp.ode_solver, OdeSolver.COLLOCATION) + else 3 + ) if n_cx < 3: n_cx = 3 @@ -1377,8 +1592,13 @@ def configure_integrated_value( mapping=dummy_mapping, node_index=0, ) - for node_index in range(1, nlp.ns + 1): # cannot use phase_dynamics == PhaseDynamics.SHARED_DURING_THE_PHASE - cx_scaled_next = [nlp.integrated_value_functions[name](nlp, node_index) for _ in range(n_cx)] + for node_index in range( + 1, nlp.ns + 1 + ): # cannot use phase_dynamics == PhaseDynamics.SHARED_DURING_THE_PHASE + cx_scaled_next = [ + nlp.integrated_value_functions[name](nlp, node_index) + for _ in range(n_cx) + ] nlp.integrated_values.append( name, cx_scaled_next_formatted, @@ -1388,7 +1608,9 @@ def configure_integrated_value( ) @staticmethod - def configure_q(ocp, nlp, as_states: bool, as_controls: bool, as_states_dot: bool = False): + def configure_q( + ocp, nlp, as_states: bool, as_controls: bool, as_states_dot: bool = False + ): """ Configure the generalized coordinates @@ -1407,11 +1629,20 @@ def configure_q(ocp, nlp, as_states: bool, as_controls: bool, as_states_dot: boo name_q = nlp.model.name_dof axes_idx = ConfigureProblem._apply_phase_mapping(ocp, nlp, name) ConfigureProblem.configure_new_variable( - name, name_q, ocp, nlp, as_states, as_controls, as_states_dot, axes_idx=axes_idx + name, + name_q, + ocp, + nlp, + as_states, + as_controls, + as_states_dot, + axes_idx=axes_idx, ) @staticmethod - def configure_qdot(ocp, nlp, as_states: bool, as_controls: bool, as_states_dot: bool = False): + def configure_qdot( + ocp, nlp, as_states: bool, as_controls: bool, as_states_dot: bool = False + ): """ Configure the generalized velocities @@ -1431,11 +1662,20 @@ def configure_qdot(ocp, nlp, as_states: bool, as_controls: bool, as_states_dot: name_qdot = ConfigureProblem._get_kinematics_based_names(nlp, name) axes_idx = ConfigureProblem._apply_phase_mapping(ocp, nlp, name) ConfigureProblem.configure_new_variable( - name, name_qdot, ocp, nlp, as_states, as_controls, as_states_dot, axes_idx=axes_idx + name, + name_qdot, + ocp, + nlp, + as_states, + as_controls, + as_states_dot, + axes_idx=axes_idx, ) @staticmethod - def configure_qddot(ocp, nlp, as_states: bool, as_controls: bool, as_states_dot: bool = False): + def configure_qddot( + ocp, nlp, as_states: bool, as_controls: bool, as_states_dot: bool = False + ): """ Configure the generalized accelerations @@ -1455,7 +1695,14 @@ def configure_qddot(ocp, nlp, as_states: bool, as_controls: bool, as_states_dot: name_qddot = ConfigureProblem._get_kinematics_based_names(nlp, name) axes_idx = ConfigureProblem._apply_phase_mapping(ocp, nlp, name) ConfigureProblem.configure_new_variable( - name, name_qddot, ocp, nlp, as_states, as_controls, as_states_dot, axes_idx=axes_idx + name, + name_qddot, + ocp, + nlp, + as_states, + as_controls, + as_states_dot, + axes_idx=axes_idx, ) @staticmethod @@ -1476,7 +1723,9 @@ def configure_qdddot(ocp, nlp, as_states: bool, as_controls: bool): name = "qdddot" name_qdddot = ConfigureProblem._get_kinematics_based_names(nlp, name) axes_idx = ConfigureProblem._apply_phase_mapping(ocp, nlp, name) - ConfigureProblem.configure_new_variable(name, name_qdddot, ocp, nlp, as_states, as_controls, axes_idx=axes_idx) + ConfigureProblem.configure_new_variable( + name, name_qdddot, ocp, nlp, as_states, as_controls, axes_idx=axes_idx + ) @staticmethod def configure_stochastic_k(ocp, nlp, n_noised_controls: int, n_references: int): @@ -1490,7 +1739,9 @@ def configure_stochastic_k(ocp, nlp, n_noised_controls: int, n_references: int): name = "k" if name in nlp.variable_mappings: - raise NotImplementedError(f"Algebraic states and mapping cannot be use together for now.") + raise NotImplementedError( + f"Algebraic states and mapping cannot be use together for now." + ) name_k = [] control_names = [f"control_{i}" for i in range(n_noised_controls)] @@ -1499,7 +1750,8 @@ def configure_stochastic_k(ocp, nlp, n_noised_controls: int, n_references: int): for name_2 in ref_names: name_k += [name_1 + "_&_" + name_2] nlp.variable_mappings[name] = BiMapping( - list(range(len(control_names) * len(ref_names))), list(range(len(control_names) * len(ref_names))) + list(range(len(control_names) * len(ref_names))), + list(range(len(control_names) * len(ref_names))), ) ConfigureProblem.configure_new_variable( name, @@ -1524,14 +1776,17 @@ def configure_stochastic_c(ocp, nlp, n_noised_states: int, n_noise: int): name = "c" if name in nlp.variable_mappings: - raise NotImplementedError(f"Algebraic states variables and mapping cannot be use together for now.") + raise NotImplementedError( + f"Algebraic states variables and mapping cannot be use together for now." + ) name_c = [] for name_1 in [f"X_{i}" for i in range(n_noised_states)]: for name_2 in [f"X_{i}" for i in range(n_noise)]: name_c += [name_1 + "_&_" + name_2] nlp.variable_mappings[name] = BiMapping( - list(range(n_noised_states * n_noise)), list(range(n_noised_states * n_noise)) + list(range(n_noised_states * n_noise)), + list(range(n_noised_states * n_noise)), ) ConfigureProblem.configure_new_variable( @@ -1558,13 +1813,17 @@ def configure_stochastic_a(ocp, nlp, n_noised_states: int): name = "a" if name in nlp.variable_mappings: - raise NotImplementedError(f"Algebraic states and mapping cannot be use together for now.") + raise NotImplementedError( + f"Algebraic states and mapping cannot be use together for now." + ) name_a = [] for name_1 in [f"X_{i}" for i in range(n_noised_states)]: for name_2 in [f"X_{i}" for i in range(n_noised_states)]: name_a += [name_1 + "_&_" + name_2] - nlp.variable_mappings[name] = BiMapping(list(range(n_noised_states**2)), list(range(n_noised_states**2))) + nlp.variable_mappings[name] = BiMapping( + list(range(n_noised_states**2)), list(range(n_noised_states**2)) + ) ConfigureProblem.configure_new_variable( name, @@ -1579,7 +1838,9 @@ def configure_stochastic_a(ocp, nlp, n_noised_states: int): ) @staticmethod - def configure_stochastic_cov_explicit(ocp, nlp, n_noised_states: int, initial_matrix: DM): + def configure_stochastic_cov_explicit( + ocp, nlp, n_noised_states: int, initial_matrix: DM + ): """ Configure the covariance matrix P representing the motor noise. Parameters @@ -1590,13 +1851,17 @@ def configure_stochastic_cov_explicit(ocp, nlp, n_noised_states: int, initial_ma name = "cov" if name in nlp.variable_mappings: - raise NotImplementedError(f"Algebraic states and mapping cannot be use together for now.") + raise NotImplementedError( + f"Algebraic states and mapping cannot be use together for now." + ) name_cov = [] for name_1 in [f"X_{i}" for i in range(n_noised_states)]: for name_2 in [f"X_{i}" for i in range(n_noised_states)]: name_cov += [name_1 + "_&_" + name_2] - nlp.variable_mappings[name] = BiMapping(list(range(n_noised_states**2)), list(range(n_noised_states**2))) + nlp.variable_mappings[name] = BiMapping( + list(range(n_noised_states**2)), list(range(n_noised_states**2)) + ) ConfigureProblem.configure_integrated_value( name, name_cov, @@ -1617,13 +1882,17 @@ def configure_stochastic_cov_implicit(ocp, nlp, n_noised_states: int): name = "cov" if name in nlp.variable_mappings: - raise NotImplementedError(f"Algebraic states and mapping cannot be use together for now.") + raise NotImplementedError( + f"Algebraic states and mapping cannot be use together for now." + ) name_cov = [] for name_1 in [f"X_{i}" for i in range(n_noised_states)]: for name_2 in [f"X_{i}" for i in range(n_noised_states)]: name_cov += [name_1 + "_&_" + name_2] - nlp.variable_mappings[name] = BiMapping(list(range(n_noised_states**2)), list(range(n_noised_states**2))) + nlp.variable_mappings[name] = BiMapping( + list(range(n_noised_states**2)), list(range(n_noised_states**2)) + ) ConfigureProblem.configure_new_variable( name, name_cov, @@ -1648,13 +1917,17 @@ def configure_stochastic_cholesky_cov(ocp, nlp, n_noised_states: int): name = "cholesky_cov" if name in nlp.variable_mappings: - raise NotImplementedError(f"Algebraic states and mapping cannot be use together for now.") + raise NotImplementedError( + f"Algebraic states and mapping cannot be use together for now." + ) name_cov = [] for nb_1, name_1 in enumerate([f"X_{i}" for i in range(n_noised_states)]): for name_2 in [f"X_{i}" for i in range(nb_1 + 1)]: name_cov += [name_1 + "_&_" + name_2] - nlp.variable_mappings[name] = BiMapping(list(range(len(name_cov))), list(range(len(name_cov)))) + nlp.variable_mappings[name] = BiMapping( + list(range(len(name_cov))), list(range(len(name_cov))) + ) ConfigureProblem.configure_new_variable( name, name_cov, @@ -1679,10 +1952,14 @@ def configure_stochastic_ref(ocp, nlp, n_references: int): name = "ref" if name in nlp.variable_mappings: - raise NotImplementedError(f"Algebraic states and mapping cannot be use together for now.") + raise NotImplementedError( + f"Algebraic states and mapping cannot be use together for now." + ) name_ref = [f"reference_{i}" for i in range(n_references)] - nlp.variable_mappings[name] = BiMapping(list(range(n_references)), list(range(n_references))) + nlp.variable_mappings[name] = BiMapping( + list(range(n_references)), list(range(n_references)) + ) ConfigureProblem.configure_new_variable( name, name_ref, @@ -1695,7 +1972,9 @@ def configure_stochastic_ref(ocp, nlp, n_references: int): ) @staticmethod - def configure_stochastic_m(ocp, nlp, n_noised_states: int, n_collocation_points: int = 1): + def configure_stochastic_m( + ocp, nlp, n_noised_states: int, n_collocation_points: int = 1 + ): """ Configure the helper matrix M (from Gillis 2013 : https://doi.org/10.1109/CDC.2013.6761121). @@ -1707,11 +1986,15 @@ def configure_stochastic_m(ocp, nlp, n_noised_states: int, n_collocation_points: name = "m" if name in nlp.variable_mappings: - raise NotImplementedError(f"Algebraic states and mapping cannot be use together for now.") + raise NotImplementedError( + f"Algebraic states and mapping cannot be use together for now." + ) name_m = [] for name_1 in [f"X_{i}" for i in range(n_noised_states)]: - for name_2 in [f"X_{i}" for i in range(n_noised_states * n_collocation_points)]: + for name_2 in [ + f"X_{i}" for i in range(n_noised_states * n_collocation_points) + ]: name_m += [name_1 + "_&_" + name_2] nlp.variable_mappings[name] = BiMapping( list(range(n_noised_states * n_noised_states * n_collocation_points)), @@ -1729,7 +2012,9 @@ def configure_stochastic_m(ocp, nlp, n_noised_states: int, n_collocation_points: ) @staticmethod - def configure_tau(ocp, nlp, as_states: bool, as_controls: bool, fatigue: FatigueList = None): + def configure_tau( + ocp, nlp, as_states: bool, as_controls: bool, fatigue: FatigueList = None + ): """ Configure the generalized forces @@ -1749,7 +2034,14 @@ def configure_tau(ocp, nlp, as_states: bool, as_controls: bool, fatigue: Fatigue name_tau = ConfigureProblem._get_kinematics_based_names(nlp, name) axes_idx = ConfigureProblem._apply_phase_mapping(ocp, nlp, name) ConfigureProblem.configure_new_variable( - name, name_tau, ocp, nlp, as_states, as_controls, fatigue=fatigue, axes_idx=axes_idx + name, + name_tau, + ocp, + nlp, + as_states, + as_controls, + fatigue=fatigue, + axes_idx=axes_idx, ) @staticmethod @@ -1792,7 +2084,9 @@ def configure_taudot(ocp, nlp, as_states: bool, as_controls: bool): name = "taudot" name_taudot = ConfigureProblem._get_kinematics_based_names(nlp, name) axes_idx = ConfigureProblem._apply_phase_mapping(ocp, nlp, name) - ConfigureProblem.configure_new_variable(name, name_taudot, ocp, nlp, as_states, as_controls, axes_idx=axes_idx) + ConfigureProblem.configure_new_variable( + name, name_taudot, ocp, nlp, as_states, as_controls, axes_idx=axes_idx + ) @staticmethod def configure_contact_forces(ocp, nlp, as_states: bool, as_controls: bool): @@ -1812,10 +2106,22 @@ def configure_contact_forces(ocp, nlp, as_states: bool, as_controls: bool): name_contact_forces = [] for i in range(nlp.model.nb_rigid_contacts): name_contact_forces.extend( - [f"Seg{i}_FX", f"Seg{i}_FY", f"Seg{i}_FZ", f"Seg{i}_CX", f"Seg{i}_CY", f"Seg{i}_CZ"] + [ + f"Seg{i}_FX", + f"Seg{i}_FY", + f"Seg{i}_FZ", + f"Seg{i}_CX", + f"Seg{i}_CY", + f"Seg{i}_CZ", + ] ) ConfigureProblem.configure_new_variable( - "translational_forces", name_contact_forces, ocp, nlp, as_states, as_controls + "translational_forces", + name_contact_forces, + ocp, + nlp, + as_states, + as_controls, ) @staticmethod @@ -1833,7 +2139,11 @@ def configure_soft_contact_forces(ocp, nlp, as_states: bool, as_controls: bool): If the generalized force derivatives should be a control """ name_soft_contact_forces = [] - component_list = ["fx", "fy", "fz"] # TODO: find a better place to hold this or define it in biorbd ? + component_list = [ + "fx", + "fy", + "fz", + ] # TODO: find a better place to hold this or define it in biorbd ? for ii in range(nlp.model.nb_soft_contacts): name_soft_contact_forces.extend( [ @@ -1843,11 +2153,18 @@ def configure_soft_contact_forces(ocp, nlp, as_states: bool, as_controls: bool): ] ) ConfigureProblem.configure_new_variable( - "forces_in_global", name_soft_contact_forces, ocp, nlp, as_states, as_controls + "forces_in_global", + name_soft_contact_forces, + ocp, + nlp, + as_states, + as_controls, ) @staticmethod - def configure_muscles(ocp, nlp, as_states: bool, as_controls: bool, fatigue: FatigueList = None): + def configure_muscles( + ocp, nlp, as_states: bool, as_controls: bool, fatigue: FatigueList = None + ): """ Configure the muscles @@ -1897,19 +2214,29 @@ def _apply_phase_mapping(ocp, nlp, name: str) -> BiMapping | None: if nlp.phase_mapping: if name in nlp.variable_mappings.keys(): double_mapping_to_first = ( - nlp.variable_mappings[name].to_first.map(nlp.phase_mapping.to_first.map_idx).T.tolist()[0] + nlp.variable_mappings[name] + .to_first.map(nlp.phase_mapping.to_first.map_idx) + .T.tolist()[0] ) - double_mapping_to_first = [int(double_mapping_to_first[i]) for i in range(len(double_mapping_to_first))] + double_mapping_to_first = [ + int(double_mapping_to_first[i]) + for i in range(len(double_mapping_to_first)) + ] double_mapping_to_second = ( - nlp.variable_mappings[name].to_second.map(nlp.phase_mapping.to_second.map_idx).T.tolist()[0] + nlp.variable_mappings[name] + .to_second.map(nlp.phase_mapping.to_second.map_idx) + .T.tolist()[0] ) double_mapping_to_second = [ - int(double_mapping_to_second[i]) for i in range(len(double_mapping_to_second)) + int(double_mapping_to_second[i]) + for i in range(len(double_mapping_to_second)) ] else: double_mapping_to_first = nlp.phase_mapping.to_first.map_idx double_mapping_to_second = nlp.phase_mapping.to_second.map_idx - axes_idx = BiMapping(to_first=double_mapping_to_first, to_second=double_mapping_to_second) + axes_idx = BiMapping( + to_first=double_mapping_to_first, to_second=double_mapping_to_second + ) else: axes_idx = None return axes_idx @@ -1921,9 +2248,13 @@ class DynamicsFcn(FcnEnum): """ TORQUE_DRIVEN = (ConfigureProblem.torque_driven,) - TORQUE_DRIVEN_FREE_FLOATING_BASE = (ConfigureProblem.torque_driven_free_floating_base,) + TORQUE_DRIVEN_FREE_FLOATING_BASE = ( + ConfigureProblem.torque_driven_free_floating_base, + ) STOCHASTIC_TORQUE_DRIVEN = (ConfigureProblem.stochastic_torque_driven,) - STOCHASTIC_TORQUE_DRIVEN_FREE_FLOATING_BASE = (ConfigureProblem.stochastic_torque_driven_free_floating_base,) + STOCHASTIC_TORQUE_DRIVEN_FREE_FLOATING_BASE = ( + ConfigureProblem.stochastic_torque_driven_free_floating_base, + ) TORQUE_DERIVATIVE_DRIVEN = (ConfigureProblem.torque_derivative_driven,) TORQUE_ACTIVATIONS_DRIVEN = (ConfigureProblem.torque_activations_driven,) JOINTS_ACCELERATION_DRIVEN = (ConfigureProblem.joints_acceleration_driven,) @@ -2026,7 +2357,9 @@ class DynamicsList(UniquePerPhaseOptionList): Print the DynamicsList to the console """ - def add(self, dynamics_type: Callable | Dynamics | DynamicsFcn, **extra_parameters: Any): + def add( + self, dynamics_type: Callable | Dynamics | DynamicsFcn, **extra_parameters: Any + ): """ Add a new Dynamics to the list @@ -2042,7 +2375,9 @@ def add(self, dynamics_type: Callable | Dynamics | DynamicsFcn, **extra_paramete self.copy(dynamics_type) else: - super(DynamicsList, self)._add(dynamics_type=dynamics_type, option_type=Dynamics, **extra_parameters) + super(DynamicsList, self)._add( + dynamics_type=dynamics_type, option_type=Dynamics, **extra_parameters + ) def print(self): """ @@ -2051,14 +2386,19 @@ def print(self): raise NotImplementedError("Printing of DynamicsList is not ready yet") -def _check_numerical_timeseries_format(numerical_timeseries: np.ndarray, n_shooting: int, phase_idx: int): +def _check_numerical_timeseries_format( + numerical_timeseries: np.ndarray, n_shooting: int, phase_idx: int +): """Check if the numerical_data_timeseries is of the right format""" if type(numerical_timeseries) is not np.ndarray: raise RuntimeError( f"Phase {phase_idx} has numerical_data_timeseries of type {type(numerical_timeseries)} " f"but it should be of type np.ndarray" ) - if numerical_timeseries is not None and numerical_timeseries.shape[2] != n_shooting + 1: + if ( + numerical_timeseries is not None + and numerical_timeseries.shape[2] != n_shooting + 1 + ): raise RuntimeError( f"Phase {phase_idx} has {n_shooting}+1 shooting points but the numerical_data_timeseries " f"has {numerical_timeseries.shape[2]} shooting points." @@ -2083,6 +2423,10 @@ def _check_soft_contacts_dynamics( ) -def _check_contacts_in_biorbd_model(with_contact: bool, nb_contacts: int, phase_idx: int): +def _check_contacts_in_biorbd_model( + with_contact: bool, nb_contacts: int, phase_idx: int +): if with_contact and nb_contacts == 0: - raise ValueError(f"No contact defined in the .bioMod of phase {phase_idx}, set with_contact to False.") + raise ValueError( + f"No contact defined in the .bioMod of phase {phase_idx}, set with_contact to False." + ) diff --git a/bioptim/dynamics/dynamics_functions.py b/bioptim/dynamics/dynamics_functions.py index fb1848421..814d43105 100644 --- a/bioptim/dynamics/dynamics_functions.py +++ b/bioptim/dynamics/dynamics_functions.py @@ -78,7 +78,13 @@ def custom( """ return nlp.dynamics_type.dynamic_function( - time, states, controls, parameters, algebraic_states, numerical_timeseries, nlp + time, + states, + controls, + parameters, + algebraic_states, + numerical_timeseries, + nlp, ) @staticmethod @@ -138,13 +144,25 @@ def torque_driven( dq = DynamicsFunctions.compute_qdot(nlp, q, qdot) tau = DynamicsFunctions.__get_fatigable_tau(nlp, states, controls, fatigue) - tau = tau + nlp.model.passive_joint_torque()(q, qdot, nlp.parameters.cx) if with_passive_torque else tau - tau = tau + nlp.model.ligament_joint_torque()(q, qdot, nlp.parameters.cx) if with_ligament else tau + tau = ( + tau + nlp.model.passive_joint_torque()(q, qdot, nlp.parameters.cx) + if with_passive_torque + else tau + ) + tau = ( + tau + nlp.model.ligament_joint_torque()(q, qdot, nlp.parameters.cx) + if with_ligament + else tau + ) tau = tau - nlp.model.friction_coefficients @ qdot if with_friction else tau - external_forces = nlp.get_external_forces(states, controls, algebraic_states, numerical_timeseries) + external_forces = nlp.get_external_forces( + states, controls, algebraic_states, numerical_timeseries + ) - ddq = DynamicsFunctions.forward_dynamics(nlp, q, qdot, tau, with_contact, external_forces) + ddq = DynamicsFunctions.forward_dynamics( + nlp, q, qdot, tau, with_contact, external_forces + ) dxdt = nlp.cx(nlp.states.shape, ddq.shape[1]) dxdt[nlp.states["q"].index, :] = horzcat(*[dq for _ in range(ddq.shape[1])]) dxdt[nlp.states["qdot"].index, :] = ddq @@ -156,8 +174,12 @@ def torque_driven( # TODO: contacts and fatigue to be handled with implicit dynamics if nlp.ode_solver.defects_type == DefectType.IMPLICIT: if not with_contact and fatigue is None: - qddot = DynamicsFunctions.get(nlp.states_dot["qddot"], nlp.states_dot.scaled.cx) - tau_id = DynamicsFunctions.inverse_dynamics(nlp, q, qdot, qddot, with_contact, external_forces) + qddot = DynamicsFunctions.get( + nlp.states_dot["qddot"], nlp.states_dot.scaled.cx + ) + tau_id = DynamicsFunctions.inverse_dynamics( + nlp, q, qdot, qddot, with_contact, external_forces + ) defects = nlp.cx(dq.shape[0] + tau_id.shape[0], tau_id.shape[1]) dq_defects = [] @@ -167,7 +189,9 @@ def torque_driven( - DynamicsFunctions.compute_qdot( nlp, q, - DynamicsFunctions.get(nlp.states_dot.scaled["qdot"], nlp.states_dot.scaled.cx), + DynamicsFunctions.get( + nlp.states_dot.scaled["qdot"], nlp.states_dot.scaled.cx + ), ) ) defects[: dq.shape[0], :] = horzcat(*dq_defects) @@ -234,12 +258,21 @@ def torque_driven_free_floating_base( n_qdot = nlp.model.nb_qdot tau_joints = ( - tau_joints + nlp.model.passive_joint_torque()(q_full, qdot_full, nlp.parameters.cx) + tau_joints + + nlp.model.passive_joint_torque()(q_full, qdot_full, nlp.parameters.cx) if with_passive_torque else tau_joints ) - tau_joints = tau_joints + nlp.model.ligament_joint_torque()(q_full, qdot_full) if with_ligament else tau_joints - tau_joints = tau_joints - nlp.model.friction_coefficients @ qdot_joints if with_friction else tau_joints + tau_joints = ( + tau_joints + nlp.model.ligament_joint_torque()(q_full, qdot_full) + if with_ligament + else tau_joints + ) + tau_joints = ( + tau_joints - nlp.model.friction_coefficients @ qdot_joints + if with_friction + else tau_joints + ) tau_full = vertcat(nlp.cx.zeros(nlp.model.nb_root), tau_joints) @@ -298,7 +331,9 @@ def stochastic_torque_driven( qdot = DynamicsFunctions.get(nlp.states["qdot"], states) tau = DynamicsFunctions.get(nlp.controls["tau"], controls) motor_noise = DynamicsFunctions.get(nlp.parameters["motor_noise"], parameters) - sensory_noise = DynamicsFunctions.get(nlp.parameters["sensory_noise"], parameters) + sensory_noise = DynamicsFunctions.get( + nlp.parameters["sensory_noise"], parameters + ) tau += nlp.model.compute_torques_from_noise_and_feedback( nlp=nlp, @@ -314,7 +349,9 @@ def stochastic_torque_driven( tau = tau - nlp.model.friction_coefficients @ qdot if with_friction else tau dq = DynamicsFunctions.compute_qdot(nlp, q, qdot) - ddq = DynamicsFunctions.forward_dynamics(nlp, q, qdot, tau, with_contact, external_forces=None) + ddq = DynamicsFunctions.forward_dynamics( + nlp, q, qdot, tau, with_contact, external_forces=None + ) dxdt = nlp.cx(nlp.states.shape, ddq.shape[1]) dxdt[nlp.states["q"].index, :] = horzcat(*[dq for _ in range(ddq.shape[1])]) dxdt[nlp.states["qdot"].index, :] = ddq @@ -366,7 +403,9 @@ def stochastic_torque_driven_free_floating_base( qdot_joints = DynamicsFunctions.get(nlp.states["qdot_joints"], states) tau_joints = DynamicsFunctions.get(nlp.controls["tau_joints"], controls) motor_noise = DynamicsFunctions.get(nlp.parameters["motor_noise"], parameters) - sensory_noise = DynamicsFunctions.get(nlp.parameters["sensory_noise"], parameters) + sensory_noise = DynamicsFunctions.get( + nlp.parameters["sensory_noise"], parameters + ) q_full = vertcat(q_roots, q_joints) qdot_full = vertcat(qdot_roots, qdot_joints) @@ -382,7 +421,11 @@ def stochastic_torque_driven_free_floating_base( motor_noise=motor_noise, sensory_noise=sensory_noise, ) - tau_joints = (tau_joints - nlp.model.friction_coefficients @ qdot_joints) if with_friction else tau_joints + tau_joints = ( + (tau_joints - nlp.model.friction_coefficients @ qdot_joints) + if with_friction + else tau_joints + ) tau_full = vertcat(nlp.cx.zeros(nlp.model.nb_root), tau_joints) @@ -397,7 +440,9 @@ def stochastic_torque_driven_free_floating_base( return DynamicsEvaluation(dxdt=dxdt, defects=None) @staticmethod - def __get_fatigable_tau(nlp, states: MX | SX, controls: MX | SX, fatigue: FatigueList) -> MX | SX: + def __get_fatigable_tau( + nlp, states: MX | SX, controls: MX | SX, fatigue: FatigueList + ) -> MX | SX: """ Apply the forward dynamics including (or not) the torque fatigue @@ -416,7 +461,9 @@ def __get_fatigable_tau(nlp, states: MX | SX, controls: MX | SX, fatigue: Fatigu ------- The generalized accelerations """ - tau_var, tau_cx = (nlp.controls, controls) if "tau" in nlp.controls else (nlp.states, states) + tau_var, tau_cx = ( + (nlp.controls, controls) if "tau" in nlp.controls else (nlp.states, states) + ) tau = nlp.get_var_from_states_or_controls("tau", states, controls) if fatigue is not None and "tau" in fatigue: tau_fatigue = fatigue["tau"] @@ -425,19 +472,30 @@ def __get_fatigable_tau(nlp, states: MX | SX, controls: MX | SX, fatigue: Fatigu # Only homogeneous state_only is implemented yet n_state_only = sum([t.models.state_only for t in tau_fatigue]) if 0 < n_state_only < len(fatigue["tau"]): - raise NotImplementedError("fatigue list without homogeneous state_only flag is not supported yet") - apply_to_joint_dynamics = sum([t.models.apply_to_joint_dynamics for t in tau_fatigue]) + raise NotImplementedError( + "fatigue list without homogeneous state_only flag is not supported yet" + ) + apply_to_joint_dynamics = sum( + [t.models.apply_to_joint_dynamics for t in tau_fatigue] + ) if 0 < n_state_only < len(fatigue["tau"]): raise NotImplementedError( "fatigue list without homogeneous apply_to_joint_dynamics flag is not supported yet" ) if apply_to_joint_dynamics != 0: - raise NotImplementedError("apply_to_joint_dynamics is not implemented for joint torque") + raise NotImplementedError( + "apply_to_joint_dynamics is not implemented for joint torque" + ) if not tau_fatigue[0].models.split_controls and "tau" in nlp.controls: pass elif tau_fatigue[0].models.state_only: - tau = sum([DynamicsFunctions.get(tau_var[f"tau_{suffix}"], tau_cx) for suffix in tau_suffix]) + tau = sum( + [ + DynamicsFunctions.get(tau_var[f"tau_{suffix}"], tau_cx) + for suffix in tau_suffix + ] + ) else: tau = nlp.cx() for i, t in enumerate(tau_fatigue): @@ -445,7 +503,10 @@ def __get_fatigable_tau(nlp, states: MX | SX, controls: MX | SX, fatigue: Fatigu for suffix in tau_suffix: model = t.models.models[suffix] tau_tp += ( - DynamicsFunctions.get(nlp.states[f"tau_{suffix}_{model.dynamics_suffix()}"], states)[i] + DynamicsFunctions.get( + nlp.states[f"tau_{suffix}_{model.dynamics_suffix()}"], + states, + )[i] * model.scaling ) tau = vertcat(tau, tau_tp) @@ -511,10 +572,14 @@ def torque_activations_driven( if with_ligament: tau += nlp.model.ligament_joint_torque()(q, qdot, nlp.parameters.cx) - external_forces = nlp.get_external_forces(states, controls, algebraic_states, numerical_timeseries) + external_forces = nlp.get_external_forces( + states, controls, algebraic_states, numerical_timeseries + ) dq = DynamicsFunctions.compute_qdot(nlp, q, qdot) - ddq = DynamicsFunctions.forward_dynamics(nlp, q, qdot, tau, with_contact, external_forces) + ddq = DynamicsFunctions.forward_dynamics( + nlp, q, qdot, tau, with_contact, external_forces + ) dq = horzcat(*[dq for _ in range(ddq.shape[1])]) @@ -569,14 +634,26 @@ def torque_derivative_driven( qdot = DynamicsFunctions.get(nlp.states["qdot"], states) tau = DynamicsFunctions.get(nlp.states["tau"], states) - tau = tau + nlp.model.passive_joint_torque()(q, qdot, nlp.parameters.cx) if with_passive_torque else tau - tau = tau + nlp.model.ligament_joint_torque()(q, qdot, nlp.parameters.cx) if with_ligament else tau + tau = ( + tau + nlp.model.passive_joint_torque()(q, qdot, nlp.parameters.cx) + if with_passive_torque + else tau + ) + tau = ( + tau + nlp.model.ligament_joint_torque()(q, qdot, nlp.parameters.cx) + if with_ligament + else tau + ) dq = DynamicsFunctions.compute_qdot(nlp, q, qdot) dtau = DynamicsFunctions.get(nlp.controls["taudot"], controls) - external_forces = nlp.get_external_forces(states, controls, algebraic_states, numerical_timeseries) - ddq = DynamicsFunctions.forward_dynamics(nlp, q, qdot, tau, with_contact, external_forces) + external_forces = nlp.get_external_forces( + states, controls, algebraic_states, numerical_timeseries + ) + ddq = DynamicsFunctions.forward_dynamics( + nlp, q, qdot, tau, with_contact, external_forces + ) dxdt = nlp.cx(nlp.states.shape, ddq.shape[1]) dxdt[nlp.states["q"].index, :] = horzcat(*[dq for _ in range(ddq.shape[1])]) dxdt[nlp.states["qdot"].index, :] = ddq @@ -629,12 +706,24 @@ def forces_from_torque_driven( q = nlp.get_var_from_states_or_controls("q", states, controls) qdot = nlp.get_var_from_states_or_controls("qdot", states, controls) tau = nlp.get_var_from_states_or_controls("tau", states, controls) - tau = tau + nlp.model.passive_joint_torque()(q, qdot, nlp.parameters.cx) if with_passive_torque else tau - tau = tau + nlp.model.ligament_joint_torque()(q, qdot, nlp.parameters.cx) if with_ligament else tau + tau = ( + tau + nlp.model.passive_joint_torque()(q, qdot, nlp.parameters.cx) + if with_passive_torque + else tau + ) + tau = ( + tau + nlp.model.ligament_joint_torque()(q, qdot, nlp.parameters.cx) + if with_ligament + else tau + ) - external_forces = nlp.get_external_forces(states, controls, algebraic_states, numerical_timeseries) + external_forces = nlp.get_external_forces( + states, controls, algebraic_states, numerical_timeseries + ) - return nlp.model.contact_forces()(q, qdot, tau, external_forces, nlp.parameters.cx) + return nlp.model.contact_forces()( + q, qdot, tau, external_forces, nlp.parameters.cx + ) @staticmethod def forces_from_torque_activation_driven( @@ -681,11 +770,23 @@ def forces_from_torque_activation_driven( qdot = nlp.get_var_from_states_or_controls("qdot", states, controls) tau_activations = nlp.get_var_from_states_or_controls("tau", states, controls) tau = nlp.model.torque()(tau_activations, q, qdot, nlp.parameters.cx) - tau = tau + nlp.model.passive_joint_torque()(q, qdot, nlp.parameters.cx) if with_passive_torque else tau - tau = tau + nlp.model.ligament_joint_torque()(q, qdot, nlp.parameters.cx) if with_ligament else tau + tau = ( + tau + nlp.model.passive_joint_torque()(q, qdot, nlp.parameters.cx) + if with_passive_torque + else tau + ) + tau = ( + tau + nlp.model.ligament_joint_torque()(q, qdot, nlp.parameters.cx) + if with_ligament + else tau + ) - external_forces = nlp.get_external_forces(states, controls, algebraic_states, numerical_timeseries) - return nlp.model.contact_forces()(q, qdot, tau, external_forces, nlp.parameters.cx) + external_forces = nlp.get_external_forces( + states, controls, algebraic_states, numerical_timeseries + ) + return nlp.model.contact_forces()( + q, qdot, tau, external_forces, nlp.parameters.cx + ) @staticmethod def muscles_driven( @@ -699,6 +800,7 @@ def muscles_driven( with_contact: bool, with_passive_torque: bool = False, with_ligament: bool = False, + with_friction: bool = False, with_residual_torque: bool = False, fatigue=None, ) -> DynamicsEvaluation: @@ -727,6 +829,8 @@ def muscles_driven( If the dynamic with passive torque should be used with_ligament: bool If the dynamic with ligament should be used + with_friction: bool + If the dynamic with friction should be used fatigue: FatigueDynamicsList To define fatigue elements with_residual_torque: bool @@ -741,9 +845,13 @@ def muscles_driven( q = nlp.get_var_from_states_or_controls("q", states, controls) qdot = nlp.get_var_from_states_or_controls("qdot", states, controls) residual_tau = ( - DynamicsFunctions.__get_fatigable_tau(nlp, states, controls, fatigue) if with_residual_torque else None + DynamicsFunctions.__get_fatigable_tau(nlp, states, controls, fatigue) + if with_residual_torque + else None + ) + mus_activations = nlp.get_var_from_states_or_controls( + "muscles", states, controls ) - mus_activations = nlp.get_var_from_states_or_controls("muscles", states, controls) fatigue_states = None if fatigue is not None and "muscles" in fatigue: mus_fatigue = fatigue["muscles"] @@ -755,7 +863,9 @@ def muscles_driven( raise NotImplementedError( f"{fatigue_name} list without homogeneous state_only flag is not supported yet" ) - apply_to_joint_dynamics = sum([m.models.apply_to_joint_dynamics for m in mus_fatigue]) + apply_to_joint_dynamics = sum( + [m.models.apply_to_joint_dynamics for m in mus_fatigue] + ) if 0 < apply_to_joint_dynamics < len(fatigue["muscles"]): raise NotImplementedError( f"{fatigue_name} list without homogeneous apply_to_joint_dynamics flag is not supported yet" @@ -772,20 +882,39 @@ def muscles_driven( raise ValueError(f"{fatigue_name} must be of all same types") if n_state_only == 0: - mus_activations = DynamicsFunctions.get(nlp.states[f"muscles_{dyn_suffix}"], states) + mus_activations = DynamicsFunctions.get( + nlp.states[f"muscles_{dyn_suffix}"], states + ) if apply_to_joint_dynamics > 0: - fatigue_states = DynamicsFunctions.get(nlp.states[f"muscles_{fatigue_suffix}"], states) - muscles_tau = DynamicsFunctions.compute_tau_from_muscle(nlp, q, qdot, mus_activations, fatigue_states) + fatigue_states = DynamicsFunctions.get( + nlp.states[f"muscles_{fatigue_suffix}"], states + ) + muscles_tau = DynamicsFunctions.compute_tau_from_muscle( + nlp, q, qdot, mus_activations, fatigue_states + ) tau = muscles_tau + residual_tau if residual_tau is not None else muscles_tau - tau = tau + nlp.model.passive_joint_torque()(q, qdot, nlp.parameters.cx) if with_passive_torque else tau - tau = tau + nlp.model.ligament_joint_torque()(q, qdot, nlp.parameters.cx) if with_ligament else tau + tau = ( + tau + nlp.model.passive_joint_torque()(q, qdot, nlp.parameters.cx) + if with_passive_torque + else tau + ) + tau = ( + tau + nlp.model.ligament_joint_torque()(q, qdot, nlp.parameters.cx) + if with_ligament + else tau + ) + tau = tau - nlp.model.friction_coefficients @ qdot if with_friction else tau dq = DynamicsFunctions.compute_qdot(nlp, q, qdot) - external_forces = nlp.get_external_forces(states, controls, algebraic_states, numerical_timeseries) - ddq = DynamicsFunctions.forward_dynamics(nlp, q, qdot, tau, with_contact, external_forces) + external_forces = nlp.get_external_forces( + states, controls, algebraic_states, numerical_timeseries + ) + ddq = DynamicsFunctions.forward_dynamics( + nlp, q, qdot, tau, with_contact, external_forces + ) dxdt = nlp.cx(nlp.states.shape, ddq.shape[1]) dxdt[nlp.states["q"].index, :] = horzcat(*[dq for _ in range(ddq.shape[1])]) dxdt[nlp.states["qdot"].index, :] = ddq @@ -793,8 +922,12 @@ def muscles_driven( has_excitation = True if "muscles" in nlp.states else False if has_excitation: mus_excitations = DynamicsFunctions.get(nlp.controls["muscles"], controls) - dmus = DynamicsFunctions.compute_muscle_dot(nlp, mus_excitations, mus_activations) - dxdt[nlp.states["muscles"].index, :] = horzcat(*[dmus for _ in range(ddq.shape[1])]) + dmus = DynamicsFunctions.compute_muscle_dot( + nlp, mus_excitations, mus_activations + ) + dxdt[nlp.states["muscles"].index, :] = horzcat( + *[dmus for _ in range(ddq.shape[1])] + ) if fatigue is not None and "muscles" in fatigue: dxdt = fatigue["muscles"].dynamics(dxdt, nlp, states, controls) @@ -845,16 +978,36 @@ def forces_from_muscle_driven( q = nlp.get_var_from_states_or_controls("q", states, controls) qdot = nlp.get_var_from_states_or_controls("qdot", states, controls) - residual_tau = nlp.get_var_from_states_or_controls("tau", states, controls) if "tau" in nlp.controls else None - mus_activations = nlp.get_var_from_states_or_controls("muscles", states, controls) - muscles_tau = DynamicsFunctions.compute_tau_from_muscle(nlp, q, qdot, mus_activations) + residual_tau = ( + nlp.get_var_from_states_or_controls("tau", states, controls) + if "tau" in nlp.controls + else None + ) + mus_activations = nlp.get_var_from_states_or_controls( + "muscles", states, controls + ) + muscles_tau = DynamicsFunctions.compute_tau_from_muscle( + nlp, q, qdot, mus_activations + ) tau = muscles_tau + residual_tau if residual_tau is not None else muscles_tau - tau = tau + nlp.model.passive_joint_torque()(q, qdot, nlp.parameters.cx) if with_passive_torque else tau - tau = tau + nlp.model.ligament_joint_torque()(q, qdot, nlp.parameters.cx) if with_ligament else tau + tau = ( + tau + nlp.model.passive_joint_torque()(q, qdot, nlp.parameters.cx) + if with_passive_torque + else tau + ) + tau = ( + tau + nlp.model.ligament_joint_torque()(q, qdot, nlp.parameters.cx) + if with_ligament + else tau + ) - external_forces = nlp.get_external_forces(states, controls, algebraic_states, numerical_timeseries) - return nlp.model.contact_forces()(q, qdot, tau, external_forces, nlp.parameters.cx) + external_forces = nlp.get_external_forces( + states, controls, algebraic_states, numerical_timeseries + ) + return nlp.model.contact_forces()( + q, qdot, tau, external_forces, nlp.parameters.cx + ) @staticmethod def joints_acceleration_driven( @@ -895,7 +1048,9 @@ def joints_acceleration_driven( qdot = nlp.get_var_from_states_or_controls("qdot", states, controls) qddot_joints = nlp.get_var_from_states_or_controls("qddot", states, controls) - qddot_root = nlp.model.forward_dynamics_free_floating_base()(q, qdot, qddot_joints, nlp.parameters.cx) + qddot_root = nlp.model.forward_dynamics_free_floating_base()( + q, qdot, qddot_joints, nlp.parameters.cx + ) qddot_reordered = nlp.model.reorder_qddot_root_joints(qddot_root, qddot_joints) qdot_mapped = nlp.variable_mappings["qdot"].to_first.map(qdot) @@ -946,15 +1101,25 @@ def compute_qdot(nlp, q: MX | SX, qdot: MX | SX): elif "q_roots" and "q_joints" in nlp.states: mapping = BiMapping( to_first=list(nlp.states["q_roots"].mapping.to_first.map_idx) - + [i + nlp.model.nb_root for i in nlp.states["q_joints"].mapping.to_first.map_idx], + + [ + i + nlp.model.nb_root + for i in nlp.states["q_joints"].mapping.to_first.map_idx + ], to_second=list(nlp.states["q_roots"].mapping.to_second.map_idx) - + [i + nlp.model.nb_root for i in nlp.states["q_joints"].mapping.to_second.map_idx], + + [ + i + nlp.model.nb_root + for i in nlp.states["q_joints"].mapping.to_second.map_idx + ], ) elif q in nlp.controls: mapping = nlp.controls["q"].mapping else: - raise RuntimeError("Your q key combination was not found in states or controls") - return mapping.to_first.map(nlp.model.reshape_qdot()(q, qdot, nlp.parameters.cx)) + raise RuntimeError( + "Your q key combination was not found in states or controls" + ) + return mapping.to_first.map( + nlp.model.reshape_qdot()(q, qdot, nlp.parameters.cx) + ) @staticmethod def forward_dynamics( @@ -992,7 +1157,9 @@ def forward_dynamics( elif "qdot" in nlp.controls: qdot_var_mapping = nlp.controls["qdot"].mapping.to_first else: - qdot_var_mapping = BiMapping([i for i in range(qdot.shape[0])], [i for i in range(qdot.shape[0])]).to_first + qdot_var_mapping = BiMapping( + [i for i in range(qdot.shape[0])], [i for i in range(qdot.shape[0])] + ).to_first external_forces = [] if external_forces is None else external_forces qddot = nlp.model.forward_dynamics(with_contact=with_contact)( @@ -1036,7 +1203,9 @@ def inverse_dynamics( Torques in tau """ if external_forces is None: - tau = nlp.model.inverse_dynamics(with_contact=with_contact)(q, qdot, qddot, [], nlp.parameters.cx) + tau = nlp.model.inverse_dynamics(with_contact=with_contact)( + q, qdot, qddot, [], nlp.parameters.cx + ) else: tau = nlp.model.inverse_dynamics(with_contact=with_contact)( q, qdot, qddot, external_forces, nlp.parameters.cx @@ -1044,7 +1213,9 @@ def inverse_dynamics( return tau # We ignore on purpose the mapping to keep zeros in the defects of the dynamic. @staticmethod - def compute_muscle_dot(nlp, muscle_excitations: MX | SX, muscle_activations: MX | SX): + def compute_muscle_dot( + nlp, muscle_excitations: MX | SX, muscle_activations: MX | SX + ): """ Easy accessor to derivative of muscle activations @@ -1062,7 +1233,9 @@ def compute_muscle_dot(nlp, muscle_excitations: MX | SX, muscle_activations: MX The derivative of muscle activations """ - return nlp.model.muscle_activation_dot()(muscle_excitations, muscle_activations, nlp.parameters.cx) + return nlp.model.muscle_activation_dot()( + muscle_excitations, muscle_activations, nlp.parameters.cx + ) @staticmethod def compute_tau_from_muscle( @@ -1096,7 +1269,9 @@ def compute_tau_from_muscle( activations = type(q)() for k in range(len(nlp.controls["muscles"])): if fatigue_states is not None: - activations = vertcat(activations, muscle_activations[k] * (1 - fatigue_states[k])) + activations = vertcat( + activations, muscle_activations[k] * (1 - fatigue_states[k]) + ) else: activations = vertcat(activations, muscle_activations[k]) return nlp.model.muscle_joint_torque()(activations, q, qdot, nlp.parameters.cx) From b2d28d518560ce3521a92317ae2fb6a15de15093 Mon Sep 17 00:00:00 2001 From: p-shg Date: Fri, 31 Jan 2025 10:03:32 +0100 Subject: [PATCH 2/2] Fixed formatting --- bioptim/dynamics/configure_problem.py | 602 ++++++------------------- bioptim/dynamics/dynamics_functions.py | 301 +++---------- 2 files changed, 196 insertions(+), 707 deletions(-) diff --git a/bioptim/dynamics/configure_problem.py b/bioptim/dynamics/configure_problem.py index e246f5151..ff068da2d 100644 --- a/bioptim/dynamics/configure_problem.py +++ b/bioptim/dynamics/configure_problem.py @@ -95,42 +95,23 @@ def _get_kinematics_based_names(nlp, var_type: str) -> list[str]: The list of str to display on figures """ - idx = ( - nlp.phase_mapping.to_first.map_idx - if nlp.phase_mapping - else range(nlp.model.nb_q) - ) + idx = nlp.phase_mapping.to_first.map_idx if nlp.phase_mapping else range(nlp.model.nb_q) if nlp.model.nb_quaternions == 0: new_names = [nlp.model.name_dof[i] for i in idx] else: new_names = [] for i in nlp.phase_mapping.to_first.map_idx: - if ( - nlp.model.name_dof[i][-4:-1] == "Rot" - or nlp.model.name_dof[i][-6:-1] == "Trans" - ): + if nlp.model.name_dof[i][-4:-1] == "Rot" or nlp.model.name_dof[i][-6:-1] == "Trans": new_names += [nlp.model.name_dof[i]] else: if nlp.model.name_dof[i][-5:] != "QuatW": if var_type == "qdot": - new_names += [ - nlp.model.name_dof[i][:-5] - + "omega" - + nlp.model.name_dof[i][-1] - ] + new_names += [nlp.model.name_dof[i][:-5] + "omega" + nlp.model.name_dof[i][-1]] elif var_type == "qddot": - new_names += [ - nlp.model.name_dof[i][:-5] - + "omegadot" - + nlp.model.name_dof[i][-1] - ] + new_names += [nlp.model.name_dof[i][:-5] + "omegadot" + nlp.model.name_dof[i][-1]] elif var_type == "qdddot": - new_names += [ - nlp.model.name_dof[i][:-5] - + "omegaddot" - + nlp.model.name_dof[i][-1] - ] + new_names += [nlp.model.name_dof[i][:-5] + "omegaddot" + nlp.model.name_dof[i][-1]] elif var_type == "tau" or var_type == "taudot": new_names += [nlp.model.name_dof[i]] @@ -208,21 +189,13 @@ def torque_driven( A list of values to pass to the dynamics at each node. Experimental external forces should be included here. """ - _check_contacts_in_biorbd_model( - with_contact, nlp.model.nb_contacts, nlp.phase_idx - ) - _check_soft_contacts_dynamics( - soft_contacts_dynamics, nlp.model.nb_soft_contacts, nlp.phase_idx - ) + _check_contacts_in_biorbd_model(with_contact, nlp.model.nb_contacts, nlp.phase_idx) + _check_soft_contacts_dynamics(soft_contacts_dynamics, nlp.model.nb_soft_contacts, nlp.phase_idx) # Declared rigidbody states and controls ConfigureProblem.configure_q(ocp, nlp, as_states=True, as_controls=False) - ConfigureProblem.configure_qdot( - ocp, nlp, as_states=True, as_controls=False, as_states_dot=True - ) - ConfigureProblem.configure_tau( - ocp, nlp, as_states=False, as_controls=True, fatigue=fatigue - ) + ConfigureProblem.configure_qdot(ocp, nlp, as_states=True, as_controls=False, as_states_dot=True) + ConfigureProblem.configure_tau(ocp, nlp, as_states=False, as_controls=True, fatigue=fatigue) ConfigureProblem.configure_qddot(ocp, nlp, False, False, True) # Declared soft contacts controls @@ -231,9 +204,7 @@ def torque_driven( # Configure the actual ODE of the dynamics if nlp.dynamics_type.dynamic_function: - ConfigureProblem.configure_dynamics_function( - ocp, nlp, DynamicsFunctions.custom - ) + ConfigureProblem.configure_dynamics_function(ocp, nlp, DynamicsFunctions.custom) else: ConfigureProblem.configure_dynamics_function( ocp, @@ -248,9 +219,7 @@ def torque_driven( # Configure the contact forces if with_contact: - ConfigureProblem.configure_contact_function( - ocp, nlp, DynamicsFunctions.forces_from_torque_driven - ) + ConfigureProblem.configure_contact_function(ocp, nlp, DynamicsFunctions.forces_from_torque_driven) # Configure the soft contact forces ConfigureProblem.configure_soft_contact_function(ocp, nlp) @@ -377,9 +346,7 @@ def torque_driven_free_floating_base( # Configure the actual ODE of the dynamics if nlp.dynamics_type.dynamic_function: - ConfigureProblem.configure_dynamics_function( - ocp, nlp, DynamicsFunctions.custom - ) + ConfigureProblem.configure_dynamics_function(ocp, nlp, DynamicsFunctions.custom) else: ConfigureProblem.configure_dynamics_function( ocp, @@ -426,30 +393,19 @@ def stochastic_torque_driven( n_noised_tau = len(nlp.model.motor_noise_mapping["tau"].to_first.map_idx) else: n_noised_tau = nlp.model.nb_tau - n_noise = ( - nlp.model.motor_noise_magnitude.shape[0] - + nlp.model.sensory_noise_magnitude.shape[0] - ) + n_noise = nlp.model.motor_noise_magnitude.shape[0] + nlp.model.sensory_noise_magnitude.shape[0] n_noised_states = 2 * n_noised_tau # Algebraic states variables ConfigureProblem.configure_stochastic_k( - ocp, - nlp, - n_noised_controls=n_noised_tau, - n_references=nlp.model.n_references, - ) - ConfigureProblem.configure_stochastic_ref( - ocp, nlp, n_references=nlp.model.n_references + ocp, nlp, n_noised_controls=n_noised_tau, n_references=nlp.model.n_references ) + ConfigureProblem.configure_stochastic_ref(ocp, nlp, n_references=nlp.model.n_references) n_collocation_points = 1 if isinstance(problem_type, SocpType.COLLOCATION): n_collocation_points += problem_type.polynomial_degree ConfigureProblem.configure_stochastic_m( - ocp, - nlp, - n_noised_states=n_noised_states, - n_collocation_points=n_collocation_points, + ocp, nlp, n_noised_states=n_noised_states, n_collocation_points=n_collocation_points ) if isinstance(problem_type, SocpType.TRAPEZOIDAL_EXPLICIT): @@ -462,21 +418,13 @@ def stochastic_torque_driven( ) else: if with_cholesky: - ConfigureProblem.configure_stochastic_cholesky_cov( - ocp, nlp, n_noised_states=n_noised_states - ) + ConfigureProblem.configure_stochastic_cholesky_cov(ocp, nlp, n_noised_states=n_noised_states) else: - ConfigureProblem.configure_stochastic_cov_implicit( - ocp, nlp, n_noised_states=n_noised_states - ) + ConfigureProblem.configure_stochastic_cov_implicit(ocp, nlp, n_noised_states=n_noised_states) if isinstance(problem_type, SocpType.TRAPEZOIDAL_IMPLICIT): - ConfigureProblem.configure_stochastic_a( - ocp, nlp, n_noised_states=n_noised_states - ) - ConfigureProblem.configure_stochastic_c( - ocp, nlp, n_noised_states=n_noised_states, n_noise=n_noise - ) + ConfigureProblem.configure_stochastic_a(ocp, nlp, n_noised_states=n_noised_states) + ConfigureProblem.configure_stochastic_c(ocp, nlp, n_noised_states=n_noised_states, n_noise=n_noise) ConfigureProblem.torque_driven( ocp=ocp, @@ -523,30 +471,19 @@ def stochastic_torque_driven_free_floating_base( A list of values to pass to the dynamics at each node. """ n_noised_tau = nlp.model.n_noised_controls - n_noise = ( - nlp.model.motor_noise_magnitude.shape[0] - + nlp.model.sensory_noise_magnitude.shape[0] - ) + n_noise = nlp.model.motor_noise_magnitude.shape[0] + nlp.model.sensory_noise_magnitude.shape[0] n_noised_states = nlp.model.n_noised_states # Stochastic variables ConfigureProblem.configure_stochastic_k( - ocp, - nlp, - n_noised_controls=n_noised_tau, - n_references=nlp.model.n_references, - ) - ConfigureProblem.configure_stochastic_ref( - ocp, nlp, n_references=nlp.model.n_references + ocp, nlp, n_noised_controls=n_noised_tau, n_references=nlp.model.n_references ) + ConfigureProblem.configure_stochastic_ref(ocp, nlp, n_references=nlp.model.n_references) n_collocation_points = 1 if isinstance(problem_type, SocpType.COLLOCATION): n_collocation_points += problem_type.polynomial_degree ConfigureProblem.configure_stochastic_m( - ocp, - nlp, - n_noised_states=n_noised_states, - n_collocation_points=n_collocation_points, + ocp, nlp, n_noised_states=n_noised_states, n_collocation_points=n_collocation_points ) if isinstance(problem_type, SocpType.TRAPEZOIDAL_EXPLICIT): @@ -559,21 +496,13 @@ def stochastic_torque_driven_free_floating_base( ) else: if with_cholesky: - ConfigureProblem.configure_stochastic_cholesky_cov( - ocp, nlp, n_noised_states=n_noised_states - ) + ConfigureProblem.configure_stochastic_cholesky_cov(ocp, nlp, n_noised_states=n_noised_states) else: - ConfigureProblem.configure_stochastic_cov_implicit( - ocp, nlp, n_noised_states=n_noised_states - ) + ConfigureProblem.configure_stochastic_cov_implicit(ocp, nlp, n_noised_states=n_noised_states) if isinstance(problem_type, SocpType.TRAPEZOIDAL_IMPLICIT): - ConfigureProblem.configure_stochastic_a( - ocp, nlp, n_noised_states=n_noised_states - ) - ConfigureProblem.configure_stochastic_c( - ocp, nlp, n_noised_states=n_noised_states, n_noise=n_noise - ) + ConfigureProblem.configure_stochastic_a(ocp, nlp, n_noised_states=n_noised_states) + ConfigureProblem.configure_stochastic_c(ocp, nlp, n_noised_states=n_noised_states, n_noise=n_noise) ConfigureProblem.torque_driven_free_floating_base( ocp=ocp, @@ -619,13 +548,9 @@ def torque_derivative_driven( A list of values to pass to the dynamics at each node. Experimental external forces should be included here. """ - _check_contacts_in_biorbd_model( - with_contact, nlp.model.nb_contacts, nlp.phase_idx - ) + _check_contacts_in_biorbd_model(with_contact, nlp.model.nb_contacts, nlp.phase_idx) - _check_soft_contacts_dynamics( - soft_contacts_dynamics, nlp.model.nb_soft_contacts, nlp.phase_idx - ) + _check_soft_contacts_dynamics(soft_contacts_dynamics, nlp.model.nb_soft_contacts, nlp.phase_idx) ConfigureProblem.configure_q(ocp, nlp, True, False) ConfigureProblem.configure_qdot(ocp, nlp, True, False) @@ -636,9 +561,7 @@ def torque_derivative_driven( ConfigureProblem.configure_soft_contact_forces(ocp, nlp, False, True) if nlp.dynamics_type.dynamic_function: - ConfigureProblem.configure_dynamics_function( - ocp, nlp, DynamicsFunctions.custom - ) + ConfigureProblem.configure_dynamics_function(ocp, nlp, DynamicsFunctions.custom) else: ConfigureProblem.configure_dynamics_function( ocp, @@ -698,9 +621,7 @@ def torque_activations_driven( A list of values to pass to the dynamics at each node. Experimental external forces should be included here. """ - _check_contacts_in_biorbd_model( - with_contact, nlp.model.nb_contacts, nlp.phase_idx - ) + _check_contacts_in_biorbd_model(with_contact, nlp.model.nb_contacts, nlp.phase_idx) ConfigureProblem.configure_q(ocp, nlp, True, False) ConfigureProblem.configure_qdot(ocp, nlp, True, False) @@ -710,9 +631,7 @@ def torque_activations_driven( ConfigureProblem.configure_residual_tau(ocp, nlp, False, True) if nlp.dynamics_type.dynamic_function: - ConfigureProblem.configure_dynamics_function( - ocp, nlp, DynamicsFunctions.custom - ) + ConfigureProblem.configure_dynamics_function(ocp, nlp, DynamicsFunctions.custom) else: ConfigureProblem.configure_dynamics_function( ocp, @@ -750,9 +669,7 @@ def joints_acceleration_driven( A list of values to pass to the dynamics at each node. Experimental external forces should be included here. """ ConfigureProblem.configure_q(ocp, nlp, as_states=True, as_controls=False) - ConfigureProblem.configure_qdot( - ocp, nlp, as_states=True, as_controls=False, as_states_dot=True - ) + ConfigureProblem.configure_qdot(ocp, nlp, as_states=True, as_controls=False, as_states_dot=True) # Configure qddot joints nb_root = nlp.model.nb_root if not nb_root > 0: @@ -769,9 +686,7 @@ def joints_acceleration_driven( as_states_dot=True, ) - name_qddot_joints = [ - str(i + nb_root) for i in range(nlp.model.nb_qddot - nb_root) - ] + name_qddot_joints = [str(i + nb_root) for i in range(nlp.model.nb_qddot - nb_root)] ConfigureProblem.configure_new_variable( "qddot_joints", name_qddot_joints, @@ -782,9 +697,7 @@ def joints_acceleration_driven( as_states_dot=True, ) - ConfigureProblem.configure_dynamics_function( - ocp, nlp, DynamicsFunctions.joints_acceleration_driven - ) + ConfigureProblem.configure_dynamics_function(ocp, nlp, DynamicsFunctions.joints_acceleration_driven) @staticmethod def muscle_driven( @@ -829,14 +742,10 @@ def muscle_driven( numerical_data_timeseries: dict[str, np.ndarray] A list of values to pass to the dynamics at each node. Experimental external forces should be included here. """ - _check_contacts_in_biorbd_model( - with_contact, nlp.model.nb_contacts, nlp.phase_idx - ) + _check_contacts_in_biorbd_model(with_contact, nlp.model.nb_contacts, nlp.phase_idx) if fatigue is not None and "tau" in fatigue and not with_residual_torque: - raise RuntimeError( - "Residual torques need to be used to apply fatigue on torques" - ) + raise RuntimeError("Residual torques need to be used to apply fatigue on torques") ConfigureProblem.configure_q(ocp, nlp, True, False) ConfigureProblem.configure_qdot(ocp, nlp, True, False, True) @@ -844,14 +753,10 @@ def muscle_driven( if with_residual_torque: ConfigureProblem.configure_tau(ocp, nlp, False, True, fatigue=fatigue) - ConfigureProblem.configure_muscles( - ocp, nlp, with_excitations, True, fatigue=fatigue - ) + ConfigureProblem.configure_muscles(ocp, nlp, with_excitations, True, fatigue=fatigue) if nlp.dynamics_type.dynamic_function: - ConfigureProblem.configure_dynamics_function( - ocp, nlp, DynamicsFunctions.custom - ) + ConfigureProblem.configure_dynamics_function(ocp, nlp, DynamicsFunctions.custom) else: ConfigureProblem.configure_dynamics_function( ocp, @@ -874,9 +779,7 @@ def muscle_driven( ConfigureProblem.configure_soft_contact_function(ocp, nlp) @staticmethod - def holonomic_torque_driven( - ocp, nlp, numerical_data_timeseries: dict[str, np.ndarray] = None - ): + def holonomic_torque_driven(ocp, nlp, numerical_data_timeseries: dict[str, np.ndarray] = None): """ Tell the program which variables are states and controls. @@ -924,18 +827,12 @@ def holonomic_torque_driven( # extra plots ConfigureProblem.configure_qv(ocp, nlp, nlp.model.compute_q_v) ConfigureProblem.configure_qdotv(ocp, nlp, nlp.model._compute_qdot_v) - ConfigureProblem.configure_lagrange_multipliers_function( - ocp, nlp, nlp.model.compute_the_lagrangian_multipliers - ) + ConfigureProblem.configure_lagrange_multipliers_function(ocp, nlp, nlp.model.compute_the_lagrangian_multipliers) - ConfigureProblem.configure_dynamics_function( - ocp, nlp, DynamicsFunctions.holonomic_torque_driven - ) + ConfigureProblem.configure_dynamics_function(ocp, nlp, DynamicsFunctions.holonomic_torque_driven) @staticmethod - def configure_lagrange_multipliers_function( - ocp, nlp, dyn_func: Callable, **extra_params - ): + def configure_lagrange_multipliers_function(ocp, nlp, dyn_func: Callable, **extra_params): """ Configure the contact points @@ -962,12 +859,8 @@ def configure_lagrange_multipliers_function( ], [ dyn_func()( - nlp.get_var_from_states_or_controls( - "q_u", nlp.states.scaled.cx, nlp.controls.scaled.cx - ), - nlp.get_var_from_states_or_controls( - "qdot_u", nlp.states.scaled.cx, nlp.controls.scaled.cx - ), + nlp.get_var_from_states_or_controls("q_u", nlp.states.scaled.cx, nlp.controls.scaled.cx), + nlp.get_var_from_states_or_controls("qdot_u", nlp.states.scaled.cx, nlp.controls.scaled.cx), DM.zeros(nlp.model.nb_dependent_joints, 1), DynamicsFunctions.get(nlp.controls["tau"], nlp.controls.scaled.cx), ) @@ -978,39 +871,20 @@ def configure_lagrange_multipliers_function( all_multipliers_names = [] for nlp_i in ocp.nlp: - if hasattr( - nlp_i.model, "has_holonomic_constraints" - ): # making sure we have a HolonomicBiorbdModel - nlp_i_multipliers_names = [ - nlp_i.model.name_dof[i] for i in nlp_i.model.dependent_joint_index - ] + if hasattr(nlp_i.model, "has_holonomic_constraints"): # making sure we have a HolonomicBiorbdModel + nlp_i_multipliers_names = [nlp_i.model.name_dof[i] for i in nlp_i.model.dependent_joint_index] all_multipliers_names.extend( - [ - name - for name in nlp_i_multipliers_names - if name not in all_multipliers_names - ] + [name for name in nlp_i_multipliers_names if name not in all_multipliers_names] ) - all_multipliers_names = [ - f"lagrange_multiplier_{name}" for name in all_multipliers_names - ] + all_multipliers_names = [f"lagrange_multiplier_{name}" for name in all_multipliers_names] all_multipliers_names_in_phase = [ - f"lagrange_multiplier_{nlp.model.name_dof[i]}" - for i in nlp.model.dependent_joint_index + f"lagrange_multiplier_{nlp.model.name_dof[i]}" for i in nlp.model.dependent_joint_index ] axes_idx = BiMapping( - to_first=[ - i - for i, c in enumerate(all_multipliers_names) - if c in all_multipliers_names_in_phase - ], - to_second=[ - i - for i, c in enumerate(all_multipliers_names) - if c in all_multipliers_names_in_phase - ], + to_first=[i for i, c in enumerate(all_multipliers_names) if c in all_multipliers_names_in_phase], + to_second=[i for i, c in enumerate(all_multipliers_names) if c in all_multipliers_names_in_phase], ) nlp.plot["lagrange_multipliers"] = CustomPlot( @@ -1050,9 +924,7 @@ def configure_qv(ocp, nlp, dyn_func: Callable, **extra_params): ], [ dyn_func()( - nlp.get_var_from_states_or_controls( - "q_u", nlp.states.cx, nlp.controls.cx - ), + nlp.get_var_from_states_or_controls("q_u", nlp.states.cx, nlp.controls.cx), DM.zeros(nlp.model.nb_dependent_joints, 1), ) ], @@ -1062,34 +934,16 @@ def configure_qv(ocp, nlp, dyn_func: Callable, **extra_params): all_multipliers_names = [] for nlp_i in ocp.nlp: - if hasattr( - nlp_i.model, "has_holonomic_constraints" - ): # making sure we have a HolonomicBiorbdModel - nlp_i_multipliers_names = [ - nlp_i.model.name_dof[i] for i in nlp_i.model.dependent_joint_index - ] + if hasattr(nlp_i.model, "has_holonomic_constraints"): # making sure we have a HolonomicBiorbdModel + nlp_i_multipliers_names = [nlp_i.model.name_dof[i] for i in nlp_i.model.dependent_joint_index] all_multipliers_names.extend( - [ - name - for name in nlp_i_multipliers_names - if name not in all_multipliers_names - ] + [name for name in nlp_i_multipliers_names if name not in all_multipliers_names] ) - all_multipliers_names_in_phase = [ - nlp.model.name_dof[i] for i in nlp.model.dependent_joint_index - ] + all_multipliers_names_in_phase = [nlp.model.name_dof[i] for i in nlp.model.dependent_joint_index] axes_idx = BiMapping( - to_first=[ - i - for i, c in enumerate(all_multipliers_names) - if c in all_multipliers_names_in_phase - ], - to_second=[ - i - for i, c in enumerate(all_multipliers_names) - if c in all_multipliers_names_in_phase - ], + to_first=[i for i, c in enumerate(all_multipliers_names) if c in all_multipliers_names_in_phase], + to_second=[i for i, c in enumerate(all_multipliers_names) if c in all_multipliers_names_in_phase], ) nlp.plot["q_v"] = CustomPlot( @@ -1129,12 +983,8 @@ def configure_qdotv(ocp, nlp, dyn_func: Callable, **extra_params): ], [ dyn_func()( - nlp.get_var_from_states_or_controls( - "q_u", nlp.states.scaled.cx, nlp.controls.scaled.cx - ), - nlp.get_var_from_states_or_controls( - "qdot_u", nlp.states.scaled.cx, nlp.controls.scaled.cx - ), + nlp.get_var_from_states_or_controls("q_u", nlp.states.scaled.cx, nlp.controls.scaled.cx), + nlp.get_var_from_states_or_controls("qdot_u", nlp.states.scaled.cx, nlp.controls.scaled.cx), DM.zeros(nlp.model.nb_dependent_joints, 1), ) ], @@ -1144,34 +994,16 @@ def configure_qdotv(ocp, nlp, dyn_func: Callable, **extra_params): all_multipliers_names = [] for nlp_i in ocp.nlp: - if hasattr( - nlp_i.model, "has_holonomic_constraints" - ): # making sure we have a HolonomicBiorbdModel - nlp_i_multipliers_names = [ - nlp_i.model.name_dof[i] for i in nlp_i.model.dependent_joint_index - ] + if hasattr(nlp_i.model, "has_holonomic_constraints"): # making sure we have a HolonomicBiorbdModel + nlp_i_multipliers_names = [nlp_i.model.name_dof[i] for i in nlp_i.model.dependent_joint_index] all_multipliers_names.extend( - [ - name - for name in nlp_i_multipliers_names - if name not in all_multipliers_names - ] + [name for name in nlp_i_multipliers_names if name not in all_multipliers_names] ) - all_multipliers_names_in_phase = [ - nlp.model.name_dof[i] for i in nlp.model.dependent_joint_index - ] + all_multipliers_names_in_phase = [nlp.model.name_dof[i] for i in nlp.model.dependent_joint_index] axes_idx = BiMapping( - to_first=[ - i - for i, c in enumerate(all_multipliers_names) - if c in all_multipliers_names_in_phase - ], - to_second=[ - i - for i, c in enumerate(all_multipliers_names) - if c in all_multipliers_names_in_phase - ], + to_first=[i for i, c in enumerate(all_multipliers_names) if c in all_multipliers_names_in_phase], + to_second=[i for i, c in enumerate(all_multipliers_names) if c in all_multipliers_names_in_phase], ) nlp.plot["qdot_v"] = CustomPlot( @@ -1345,37 +1177,19 @@ def configure_contact_function(ocp, nlp, dyn_func: Callable, **extra_params): all_contact_names = [] for elt in ocp.nlp: - all_contact_names.extend( - [ - name - for name in elt.model.contact_names - if name not in all_contact_names - ] - ) + all_contact_names.extend([name for name in elt.model.contact_names if name not in all_contact_names]) if "contact_forces" in nlp.plot_mapping: contact_names_in_phase = [name for name in nlp.model.contact_names] axes_idx = BiMapping( to_first=nlp.plot_mapping["contact_forces"].map_idx, - to_second=[ - i - for i, c in enumerate(all_contact_names) - if c in contact_names_in_phase - ], + to_second=[i for i, c in enumerate(all_contact_names) if c in contact_names_in_phase], ) else: contact_names_in_phase = [name for name in nlp.model.contact_names] axes_idx = BiMapping( - to_first=[ - i - for i, c in enumerate(all_contact_names) - if c in contact_names_in_phase - ], - to_second=[ - i - for i, c in enumerate(all_contact_names) - if c in contact_names_in_phase - ], + to_first=[i for i, c in enumerate(all_contact_names) if c in contact_names_in_phase], + to_second=[i for i, c in enumerate(all_contact_names) if c in contact_names_in_phase], ) nlp.plot["contact_forces"] = CustomPlot( @@ -1439,11 +1253,7 @@ def configure_soft_contact_function(ocp, nlp): ] phase_mappings = BiMapping( to_first=nlp.plot_mapping["soft_contact_forces"].map_idx, - to_second=[ - i - for i, c in enumerate(all_soft_contact_names) - if c in soft_contact_names_in_phase - ], + to_second=[i for i, c in enumerate(all_soft_contact_names) if c in soft_contact_names_in_phase], ) else: soft_contact_names_in_phase = [ @@ -1452,33 +1262,16 @@ def configure_soft_contact_function(ocp, nlp): if nlp.model.soft_contact_names[i_sc] not in all_soft_contact_names ] phase_mappings = BiMapping( - to_first=[ - i - for i, c in enumerate(all_soft_contact_names) - if c in soft_contact_names_in_phase - ], - to_second=[ - i - for i, c in enumerate(all_soft_contact_names) - if c in soft_contact_names_in_phase - ], - ) - nlp.plot[f"soft_contact_forces_{nlp.model.soft_contact_names[i_sc]}"] = ( - CustomPlot( - lambda t0, phases_dt, node_idx, x, u, p, a, d: nlp.soft_contact_forces_func( - np.concatenate([t0, t0 + phases_dt[nlp.phase_idx]]), - x, - u, - p, - a, - d, - )[ - (i_sc * 6) : ((i_sc + 1) * 6), : - ], - plot_type=PlotType.INTEGRATED, - axes_idx=phase_mappings, - legend=all_soft_contact_names, + to_first=[i for i, c in enumerate(all_soft_contact_names) if c in soft_contact_names_in_phase], + to_second=[i for i, c in enumerate(all_soft_contact_names) if c in soft_contact_names_in_phase], ) + nlp.plot[f"soft_contact_forces_{nlp.model.soft_contact_names[i_sc]}"] = CustomPlot( + lambda t0, phases_dt, node_idx, x, u, p, a, d: nlp.soft_contact_forces_func( + np.concatenate([t0, t0 + phases_dt[nlp.phase_idx]]), x, u, p, a, d + )[(i_sc * 6) : ((i_sc + 1) * 6), :], + plot_type=PlotType.INTEGRATED, + axes_idx=phase_mappings, + legend=all_soft_contact_names, ) @staticmethod @@ -1574,11 +1367,7 @@ def configure_integrated_value( # TODO: compute values at collocation points # but for now only cx_start can be used - n_cx = ( - nlp.ode_solver.n_cx - 1 - if isinstance(nlp.ode_solver, OdeSolver.COLLOCATION) - else 3 - ) + n_cx = nlp.ode_solver.n_cx - 1 if isinstance(nlp.ode_solver, OdeSolver.COLLOCATION) else 3 if n_cx < 3: n_cx = 3 @@ -1592,13 +1381,8 @@ def configure_integrated_value( mapping=dummy_mapping, node_index=0, ) - for node_index in range( - 1, nlp.ns + 1 - ): # cannot use phase_dynamics == PhaseDynamics.SHARED_DURING_THE_PHASE - cx_scaled_next = [ - nlp.integrated_value_functions[name](nlp, node_index) - for _ in range(n_cx) - ] + for node_index in range(1, nlp.ns + 1): # cannot use phase_dynamics == PhaseDynamics.SHARED_DURING_THE_PHASE + cx_scaled_next = [nlp.integrated_value_functions[name](nlp, node_index) for _ in range(n_cx)] nlp.integrated_values.append( name, cx_scaled_next_formatted, @@ -1608,9 +1392,7 @@ def configure_integrated_value( ) @staticmethod - def configure_q( - ocp, nlp, as_states: bool, as_controls: bool, as_states_dot: bool = False - ): + def configure_q(ocp, nlp, as_states: bool, as_controls: bool, as_states_dot: bool = False): """ Configure the generalized coordinates @@ -1629,20 +1411,11 @@ def configure_q( name_q = nlp.model.name_dof axes_idx = ConfigureProblem._apply_phase_mapping(ocp, nlp, name) ConfigureProblem.configure_new_variable( - name, - name_q, - ocp, - nlp, - as_states, - as_controls, - as_states_dot, - axes_idx=axes_idx, + name, name_q, ocp, nlp, as_states, as_controls, as_states_dot, axes_idx=axes_idx ) @staticmethod - def configure_qdot( - ocp, nlp, as_states: bool, as_controls: bool, as_states_dot: bool = False - ): + def configure_qdot(ocp, nlp, as_states: bool, as_controls: bool, as_states_dot: bool = False): """ Configure the generalized velocities @@ -1662,20 +1435,11 @@ def configure_qdot( name_qdot = ConfigureProblem._get_kinematics_based_names(nlp, name) axes_idx = ConfigureProblem._apply_phase_mapping(ocp, nlp, name) ConfigureProblem.configure_new_variable( - name, - name_qdot, - ocp, - nlp, - as_states, - as_controls, - as_states_dot, - axes_idx=axes_idx, + name, name_qdot, ocp, nlp, as_states, as_controls, as_states_dot, axes_idx=axes_idx ) @staticmethod - def configure_qddot( - ocp, nlp, as_states: bool, as_controls: bool, as_states_dot: bool = False - ): + def configure_qddot(ocp, nlp, as_states: bool, as_controls: bool, as_states_dot: bool = False): """ Configure the generalized accelerations @@ -1695,14 +1459,7 @@ def configure_qddot( name_qddot = ConfigureProblem._get_kinematics_based_names(nlp, name) axes_idx = ConfigureProblem._apply_phase_mapping(ocp, nlp, name) ConfigureProblem.configure_new_variable( - name, - name_qddot, - ocp, - nlp, - as_states, - as_controls, - as_states_dot, - axes_idx=axes_idx, + name, name_qddot, ocp, nlp, as_states, as_controls, as_states_dot, axes_idx=axes_idx ) @staticmethod @@ -1723,9 +1480,7 @@ def configure_qdddot(ocp, nlp, as_states: bool, as_controls: bool): name = "qdddot" name_qdddot = ConfigureProblem._get_kinematics_based_names(nlp, name) axes_idx = ConfigureProblem._apply_phase_mapping(ocp, nlp, name) - ConfigureProblem.configure_new_variable( - name, name_qdddot, ocp, nlp, as_states, as_controls, axes_idx=axes_idx - ) + ConfigureProblem.configure_new_variable(name, name_qdddot, ocp, nlp, as_states, as_controls, axes_idx=axes_idx) @staticmethod def configure_stochastic_k(ocp, nlp, n_noised_controls: int, n_references: int): @@ -1739,9 +1494,7 @@ def configure_stochastic_k(ocp, nlp, n_noised_controls: int, n_references: int): name = "k" if name in nlp.variable_mappings: - raise NotImplementedError( - f"Algebraic states and mapping cannot be use together for now." - ) + raise NotImplementedError(f"Algebraic states and mapping cannot be use together for now.") name_k = [] control_names = [f"control_{i}" for i in range(n_noised_controls)] @@ -1750,8 +1503,7 @@ def configure_stochastic_k(ocp, nlp, n_noised_controls: int, n_references: int): for name_2 in ref_names: name_k += [name_1 + "_&_" + name_2] nlp.variable_mappings[name] = BiMapping( - list(range(len(control_names) * len(ref_names))), - list(range(len(control_names) * len(ref_names))), + list(range(len(control_names) * len(ref_names))), list(range(len(control_names) * len(ref_names))) ) ConfigureProblem.configure_new_variable( name, @@ -1776,17 +1528,14 @@ def configure_stochastic_c(ocp, nlp, n_noised_states: int, n_noise: int): name = "c" if name in nlp.variable_mappings: - raise NotImplementedError( - f"Algebraic states variables and mapping cannot be use together for now." - ) + raise NotImplementedError(f"Algebraic states variables and mapping cannot be use together for now.") name_c = [] for name_1 in [f"X_{i}" for i in range(n_noised_states)]: for name_2 in [f"X_{i}" for i in range(n_noise)]: name_c += [name_1 + "_&_" + name_2] nlp.variable_mappings[name] = BiMapping( - list(range(n_noised_states * n_noise)), - list(range(n_noised_states * n_noise)), + list(range(n_noised_states * n_noise)), list(range(n_noised_states * n_noise)) ) ConfigureProblem.configure_new_variable( @@ -1813,17 +1562,13 @@ def configure_stochastic_a(ocp, nlp, n_noised_states: int): name = "a" if name in nlp.variable_mappings: - raise NotImplementedError( - f"Algebraic states and mapping cannot be use together for now." - ) + raise NotImplementedError(f"Algebraic states and mapping cannot be use together for now.") name_a = [] for name_1 in [f"X_{i}" for i in range(n_noised_states)]: for name_2 in [f"X_{i}" for i in range(n_noised_states)]: name_a += [name_1 + "_&_" + name_2] - nlp.variable_mappings[name] = BiMapping( - list(range(n_noised_states**2)), list(range(n_noised_states**2)) - ) + nlp.variable_mappings[name] = BiMapping(list(range(n_noised_states**2)), list(range(n_noised_states**2))) ConfigureProblem.configure_new_variable( name, @@ -1838,9 +1583,7 @@ def configure_stochastic_a(ocp, nlp, n_noised_states: int): ) @staticmethod - def configure_stochastic_cov_explicit( - ocp, nlp, n_noised_states: int, initial_matrix: DM - ): + def configure_stochastic_cov_explicit(ocp, nlp, n_noised_states: int, initial_matrix: DM): """ Configure the covariance matrix P representing the motor noise. Parameters @@ -1851,17 +1594,13 @@ def configure_stochastic_cov_explicit( name = "cov" if name in nlp.variable_mappings: - raise NotImplementedError( - f"Algebraic states and mapping cannot be use together for now." - ) + raise NotImplementedError(f"Algebraic states and mapping cannot be use together for now.") name_cov = [] for name_1 in [f"X_{i}" for i in range(n_noised_states)]: for name_2 in [f"X_{i}" for i in range(n_noised_states)]: name_cov += [name_1 + "_&_" + name_2] - nlp.variable_mappings[name] = BiMapping( - list(range(n_noised_states**2)), list(range(n_noised_states**2)) - ) + nlp.variable_mappings[name] = BiMapping(list(range(n_noised_states**2)), list(range(n_noised_states**2))) ConfigureProblem.configure_integrated_value( name, name_cov, @@ -1882,17 +1621,13 @@ def configure_stochastic_cov_implicit(ocp, nlp, n_noised_states: int): name = "cov" if name in nlp.variable_mappings: - raise NotImplementedError( - f"Algebraic states and mapping cannot be use together for now." - ) + raise NotImplementedError(f"Algebraic states and mapping cannot be use together for now.") name_cov = [] for name_1 in [f"X_{i}" for i in range(n_noised_states)]: for name_2 in [f"X_{i}" for i in range(n_noised_states)]: name_cov += [name_1 + "_&_" + name_2] - nlp.variable_mappings[name] = BiMapping( - list(range(n_noised_states**2)), list(range(n_noised_states**2)) - ) + nlp.variable_mappings[name] = BiMapping(list(range(n_noised_states**2)), list(range(n_noised_states**2))) ConfigureProblem.configure_new_variable( name, name_cov, @@ -1917,17 +1652,13 @@ def configure_stochastic_cholesky_cov(ocp, nlp, n_noised_states: int): name = "cholesky_cov" if name in nlp.variable_mappings: - raise NotImplementedError( - f"Algebraic states and mapping cannot be use together for now." - ) + raise NotImplementedError(f"Algebraic states and mapping cannot be use together for now.") name_cov = [] for nb_1, name_1 in enumerate([f"X_{i}" for i in range(n_noised_states)]): for name_2 in [f"X_{i}" for i in range(nb_1 + 1)]: name_cov += [name_1 + "_&_" + name_2] - nlp.variable_mappings[name] = BiMapping( - list(range(len(name_cov))), list(range(len(name_cov))) - ) + nlp.variable_mappings[name] = BiMapping(list(range(len(name_cov))), list(range(len(name_cov)))) ConfigureProblem.configure_new_variable( name, name_cov, @@ -1952,14 +1683,10 @@ def configure_stochastic_ref(ocp, nlp, n_references: int): name = "ref" if name in nlp.variable_mappings: - raise NotImplementedError( - f"Algebraic states and mapping cannot be use together for now." - ) + raise NotImplementedError(f"Algebraic states and mapping cannot be use together for now.") name_ref = [f"reference_{i}" for i in range(n_references)] - nlp.variable_mappings[name] = BiMapping( - list(range(n_references)), list(range(n_references)) - ) + nlp.variable_mappings[name] = BiMapping(list(range(n_references)), list(range(n_references))) ConfigureProblem.configure_new_variable( name, name_ref, @@ -1972,9 +1699,7 @@ def configure_stochastic_ref(ocp, nlp, n_references: int): ) @staticmethod - def configure_stochastic_m( - ocp, nlp, n_noised_states: int, n_collocation_points: int = 1 - ): + def configure_stochastic_m(ocp, nlp, n_noised_states: int, n_collocation_points: int = 1): """ Configure the helper matrix M (from Gillis 2013 : https://doi.org/10.1109/CDC.2013.6761121). @@ -1986,15 +1711,11 @@ def configure_stochastic_m( name = "m" if name in nlp.variable_mappings: - raise NotImplementedError( - f"Algebraic states and mapping cannot be use together for now." - ) + raise NotImplementedError(f"Algebraic states and mapping cannot be use together for now.") name_m = [] for name_1 in [f"X_{i}" for i in range(n_noised_states)]: - for name_2 in [ - f"X_{i}" for i in range(n_noised_states * n_collocation_points) - ]: + for name_2 in [f"X_{i}" for i in range(n_noised_states * n_collocation_points)]: name_m += [name_1 + "_&_" + name_2] nlp.variable_mappings[name] = BiMapping( list(range(n_noised_states * n_noised_states * n_collocation_points)), @@ -2012,9 +1733,7 @@ def configure_stochastic_m( ) @staticmethod - def configure_tau( - ocp, nlp, as_states: bool, as_controls: bool, fatigue: FatigueList = None - ): + def configure_tau(ocp, nlp, as_states: bool, as_controls: bool, fatigue: FatigueList = None): """ Configure the generalized forces @@ -2034,14 +1753,7 @@ def configure_tau( name_tau = ConfigureProblem._get_kinematics_based_names(nlp, name) axes_idx = ConfigureProblem._apply_phase_mapping(ocp, nlp, name) ConfigureProblem.configure_new_variable( - name, - name_tau, - ocp, - nlp, - as_states, - as_controls, - fatigue=fatigue, - axes_idx=axes_idx, + name, name_tau, ocp, nlp, as_states, as_controls, fatigue=fatigue, axes_idx=axes_idx ) @staticmethod @@ -2084,9 +1796,7 @@ def configure_taudot(ocp, nlp, as_states: bool, as_controls: bool): name = "taudot" name_taudot = ConfigureProblem._get_kinematics_based_names(nlp, name) axes_idx = ConfigureProblem._apply_phase_mapping(ocp, nlp, name) - ConfigureProblem.configure_new_variable( - name, name_taudot, ocp, nlp, as_states, as_controls, axes_idx=axes_idx - ) + ConfigureProblem.configure_new_variable(name, name_taudot, ocp, nlp, as_states, as_controls, axes_idx=axes_idx) @staticmethod def configure_contact_forces(ocp, nlp, as_states: bool, as_controls: bool): @@ -2106,22 +1816,10 @@ def configure_contact_forces(ocp, nlp, as_states: bool, as_controls: bool): name_contact_forces = [] for i in range(nlp.model.nb_rigid_contacts): name_contact_forces.extend( - [ - f"Seg{i}_FX", - f"Seg{i}_FY", - f"Seg{i}_FZ", - f"Seg{i}_CX", - f"Seg{i}_CY", - f"Seg{i}_CZ", - ] + [f"Seg{i}_FX", f"Seg{i}_FY", f"Seg{i}_FZ", f"Seg{i}_CX", f"Seg{i}_CY", f"Seg{i}_CZ"] ) ConfigureProblem.configure_new_variable( - "translational_forces", - name_contact_forces, - ocp, - nlp, - as_states, - as_controls, + "translational_forces", name_contact_forces, ocp, nlp, as_states, as_controls ) @staticmethod @@ -2139,11 +1837,7 @@ def configure_soft_contact_forces(ocp, nlp, as_states: bool, as_controls: bool): If the generalized force derivatives should be a control """ name_soft_contact_forces = [] - component_list = [ - "fx", - "fy", - "fz", - ] # TODO: find a better place to hold this or define it in biorbd ? + component_list = ["fx", "fy", "fz"] # TODO: find a better place to hold this or define it in biorbd ? for ii in range(nlp.model.nb_soft_contacts): name_soft_contact_forces.extend( [ @@ -2153,18 +1847,11 @@ def configure_soft_contact_forces(ocp, nlp, as_states: bool, as_controls: bool): ] ) ConfigureProblem.configure_new_variable( - "forces_in_global", - name_soft_contact_forces, - ocp, - nlp, - as_states, - as_controls, + "forces_in_global", name_soft_contact_forces, ocp, nlp, as_states, as_controls ) @staticmethod - def configure_muscles( - ocp, nlp, as_states: bool, as_controls: bool, fatigue: FatigueList = None - ): + def configure_muscles(ocp, nlp, as_states: bool, as_controls: bool, fatigue: FatigueList = None): """ Configure the muscles @@ -2214,29 +1901,19 @@ def _apply_phase_mapping(ocp, nlp, name: str) -> BiMapping | None: if nlp.phase_mapping: if name in nlp.variable_mappings.keys(): double_mapping_to_first = ( - nlp.variable_mappings[name] - .to_first.map(nlp.phase_mapping.to_first.map_idx) - .T.tolist()[0] + nlp.variable_mappings[name].to_first.map(nlp.phase_mapping.to_first.map_idx).T.tolist()[0] ) - double_mapping_to_first = [ - int(double_mapping_to_first[i]) - for i in range(len(double_mapping_to_first)) - ] + double_mapping_to_first = [int(double_mapping_to_first[i]) for i in range(len(double_mapping_to_first))] double_mapping_to_second = ( - nlp.variable_mappings[name] - .to_second.map(nlp.phase_mapping.to_second.map_idx) - .T.tolist()[0] + nlp.variable_mappings[name].to_second.map(nlp.phase_mapping.to_second.map_idx).T.tolist()[0] ) double_mapping_to_second = [ - int(double_mapping_to_second[i]) - for i in range(len(double_mapping_to_second)) + int(double_mapping_to_second[i]) for i in range(len(double_mapping_to_second)) ] else: double_mapping_to_first = nlp.phase_mapping.to_first.map_idx double_mapping_to_second = nlp.phase_mapping.to_second.map_idx - axes_idx = BiMapping( - to_first=double_mapping_to_first, to_second=double_mapping_to_second - ) + axes_idx = BiMapping(to_first=double_mapping_to_first, to_second=double_mapping_to_second) else: axes_idx = None return axes_idx @@ -2248,13 +1925,9 @@ class DynamicsFcn(FcnEnum): """ TORQUE_DRIVEN = (ConfigureProblem.torque_driven,) - TORQUE_DRIVEN_FREE_FLOATING_BASE = ( - ConfigureProblem.torque_driven_free_floating_base, - ) + TORQUE_DRIVEN_FREE_FLOATING_BASE = (ConfigureProblem.torque_driven_free_floating_base,) STOCHASTIC_TORQUE_DRIVEN = (ConfigureProblem.stochastic_torque_driven,) - STOCHASTIC_TORQUE_DRIVEN_FREE_FLOATING_BASE = ( - ConfigureProblem.stochastic_torque_driven_free_floating_base, - ) + STOCHASTIC_TORQUE_DRIVEN_FREE_FLOATING_BASE = (ConfigureProblem.stochastic_torque_driven_free_floating_base,) TORQUE_DERIVATIVE_DRIVEN = (ConfigureProblem.torque_derivative_driven,) TORQUE_ACTIVATIONS_DRIVEN = (ConfigureProblem.torque_activations_driven,) JOINTS_ACCELERATION_DRIVEN = (ConfigureProblem.joints_acceleration_driven,) @@ -2357,9 +2030,7 @@ class DynamicsList(UniquePerPhaseOptionList): Print the DynamicsList to the console """ - def add( - self, dynamics_type: Callable | Dynamics | DynamicsFcn, **extra_parameters: Any - ): + def add(self, dynamics_type: Callable | Dynamics | DynamicsFcn, **extra_parameters: Any): """ Add a new Dynamics to the list @@ -2375,9 +2046,7 @@ def add( self.copy(dynamics_type) else: - super(DynamicsList, self)._add( - dynamics_type=dynamics_type, option_type=Dynamics, **extra_parameters - ) + super(DynamicsList, self)._add(dynamics_type=dynamics_type, option_type=Dynamics, **extra_parameters) def print(self): """ @@ -2386,19 +2055,14 @@ def print(self): raise NotImplementedError("Printing of DynamicsList is not ready yet") -def _check_numerical_timeseries_format( - numerical_timeseries: np.ndarray, n_shooting: int, phase_idx: int -): +def _check_numerical_timeseries_format(numerical_timeseries: np.ndarray, n_shooting: int, phase_idx: int): """Check if the numerical_data_timeseries is of the right format""" if type(numerical_timeseries) is not np.ndarray: raise RuntimeError( f"Phase {phase_idx} has numerical_data_timeseries of type {type(numerical_timeseries)} " f"but it should be of type np.ndarray" ) - if ( - numerical_timeseries is not None - and numerical_timeseries.shape[2] != n_shooting + 1 - ): + if numerical_timeseries is not None and numerical_timeseries.shape[2] != n_shooting + 1: raise RuntimeError( f"Phase {phase_idx} has {n_shooting}+1 shooting points but the numerical_data_timeseries " f"has {numerical_timeseries.shape[2]} shooting points." @@ -2423,10 +2087,6 @@ def _check_soft_contacts_dynamics( ) -def _check_contacts_in_biorbd_model( - with_contact: bool, nb_contacts: int, phase_idx: int -): +def _check_contacts_in_biorbd_model(with_contact: bool, nb_contacts: int, phase_idx: int): if with_contact and nb_contacts == 0: - raise ValueError( - f"No contact defined in the .bioMod of phase {phase_idx}, set with_contact to False." - ) + raise ValueError(f"No contact defined in the .bioMod of phase {phase_idx}, set with_contact to False.") diff --git a/bioptim/dynamics/dynamics_functions.py b/bioptim/dynamics/dynamics_functions.py index 814d43105..128dbd921 100644 --- a/bioptim/dynamics/dynamics_functions.py +++ b/bioptim/dynamics/dynamics_functions.py @@ -78,13 +78,7 @@ def custom( """ return nlp.dynamics_type.dynamic_function( - time, - states, - controls, - parameters, - algebraic_states, - numerical_timeseries, - nlp, + time, states, controls, parameters, algebraic_states, numerical_timeseries, nlp ) @staticmethod @@ -144,25 +138,13 @@ def torque_driven( dq = DynamicsFunctions.compute_qdot(nlp, q, qdot) tau = DynamicsFunctions.__get_fatigable_tau(nlp, states, controls, fatigue) - tau = ( - tau + nlp.model.passive_joint_torque()(q, qdot, nlp.parameters.cx) - if with_passive_torque - else tau - ) - tau = ( - tau + nlp.model.ligament_joint_torque()(q, qdot, nlp.parameters.cx) - if with_ligament - else tau - ) + tau = tau + nlp.model.passive_joint_torque()(q, qdot, nlp.parameters.cx) if with_passive_torque else tau + tau = tau + nlp.model.ligament_joint_torque()(q, qdot, nlp.parameters.cx) if with_ligament else tau tau = tau - nlp.model.friction_coefficients @ qdot if with_friction else tau - external_forces = nlp.get_external_forces( - states, controls, algebraic_states, numerical_timeseries - ) + external_forces = nlp.get_external_forces(states, controls, algebraic_states, numerical_timeseries) - ddq = DynamicsFunctions.forward_dynamics( - nlp, q, qdot, tau, with_contact, external_forces - ) + ddq = DynamicsFunctions.forward_dynamics(nlp, q, qdot, tau, with_contact, external_forces) dxdt = nlp.cx(nlp.states.shape, ddq.shape[1]) dxdt[nlp.states["q"].index, :] = horzcat(*[dq for _ in range(ddq.shape[1])]) dxdt[nlp.states["qdot"].index, :] = ddq @@ -174,12 +156,8 @@ def torque_driven( # TODO: contacts and fatigue to be handled with implicit dynamics if nlp.ode_solver.defects_type == DefectType.IMPLICIT: if not with_contact and fatigue is None: - qddot = DynamicsFunctions.get( - nlp.states_dot["qddot"], nlp.states_dot.scaled.cx - ) - tau_id = DynamicsFunctions.inverse_dynamics( - nlp, q, qdot, qddot, with_contact, external_forces - ) + qddot = DynamicsFunctions.get(nlp.states_dot["qddot"], nlp.states_dot.scaled.cx) + tau_id = DynamicsFunctions.inverse_dynamics(nlp, q, qdot, qddot, with_contact, external_forces) defects = nlp.cx(dq.shape[0] + tau_id.shape[0], tau_id.shape[1]) dq_defects = [] @@ -189,9 +167,7 @@ def torque_driven( - DynamicsFunctions.compute_qdot( nlp, q, - DynamicsFunctions.get( - nlp.states_dot.scaled["qdot"], nlp.states_dot.scaled.cx - ), + DynamicsFunctions.get(nlp.states_dot.scaled["qdot"], nlp.states_dot.scaled.cx), ) ) defects[: dq.shape[0], :] = horzcat(*dq_defects) @@ -258,21 +234,12 @@ def torque_driven_free_floating_base( n_qdot = nlp.model.nb_qdot tau_joints = ( - tau_joints - + nlp.model.passive_joint_torque()(q_full, qdot_full, nlp.parameters.cx) + tau_joints + nlp.model.passive_joint_torque()(q_full, qdot_full, nlp.parameters.cx) if with_passive_torque else tau_joints ) - tau_joints = ( - tau_joints + nlp.model.ligament_joint_torque()(q_full, qdot_full) - if with_ligament - else tau_joints - ) - tau_joints = ( - tau_joints - nlp.model.friction_coefficients @ qdot_joints - if with_friction - else tau_joints - ) + tau_joints = tau_joints + nlp.model.ligament_joint_torque()(q_full, qdot_full) if with_ligament else tau_joints + tau_joints = tau_joints - nlp.model.friction_coefficients @ qdot_joints if with_friction else tau_joints tau_full = vertcat(nlp.cx.zeros(nlp.model.nb_root), tau_joints) @@ -331,9 +298,7 @@ def stochastic_torque_driven( qdot = DynamicsFunctions.get(nlp.states["qdot"], states) tau = DynamicsFunctions.get(nlp.controls["tau"], controls) motor_noise = DynamicsFunctions.get(nlp.parameters["motor_noise"], parameters) - sensory_noise = DynamicsFunctions.get( - nlp.parameters["sensory_noise"], parameters - ) + sensory_noise = DynamicsFunctions.get(nlp.parameters["sensory_noise"], parameters) tau += nlp.model.compute_torques_from_noise_and_feedback( nlp=nlp, @@ -349,9 +314,7 @@ def stochastic_torque_driven( tau = tau - nlp.model.friction_coefficients @ qdot if with_friction else tau dq = DynamicsFunctions.compute_qdot(nlp, q, qdot) - ddq = DynamicsFunctions.forward_dynamics( - nlp, q, qdot, tau, with_contact, external_forces=None - ) + ddq = DynamicsFunctions.forward_dynamics(nlp, q, qdot, tau, with_contact, external_forces=None) dxdt = nlp.cx(nlp.states.shape, ddq.shape[1]) dxdt[nlp.states["q"].index, :] = horzcat(*[dq for _ in range(ddq.shape[1])]) dxdt[nlp.states["qdot"].index, :] = ddq @@ -403,9 +366,7 @@ def stochastic_torque_driven_free_floating_base( qdot_joints = DynamicsFunctions.get(nlp.states["qdot_joints"], states) tau_joints = DynamicsFunctions.get(nlp.controls["tau_joints"], controls) motor_noise = DynamicsFunctions.get(nlp.parameters["motor_noise"], parameters) - sensory_noise = DynamicsFunctions.get( - nlp.parameters["sensory_noise"], parameters - ) + sensory_noise = DynamicsFunctions.get(nlp.parameters["sensory_noise"], parameters) q_full = vertcat(q_roots, q_joints) qdot_full = vertcat(qdot_roots, qdot_joints) @@ -421,11 +382,7 @@ def stochastic_torque_driven_free_floating_base( motor_noise=motor_noise, sensory_noise=sensory_noise, ) - tau_joints = ( - (tau_joints - nlp.model.friction_coefficients @ qdot_joints) - if with_friction - else tau_joints - ) + tau_joints = (tau_joints - nlp.model.friction_coefficients @ qdot_joints) if with_friction else tau_joints tau_full = vertcat(nlp.cx.zeros(nlp.model.nb_root), tau_joints) @@ -440,9 +397,7 @@ def stochastic_torque_driven_free_floating_base( return DynamicsEvaluation(dxdt=dxdt, defects=None) @staticmethod - def __get_fatigable_tau( - nlp, states: MX | SX, controls: MX | SX, fatigue: FatigueList - ) -> MX | SX: + def __get_fatigable_tau(nlp, states: MX | SX, controls: MX | SX, fatigue: FatigueList) -> MX | SX: """ Apply the forward dynamics including (or not) the torque fatigue @@ -461,9 +416,7 @@ def __get_fatigable_tau( ------- The generalized accelerations """ - tau_var, tau_cx = ( - (nlp.controls, controls) if "tau" in nlp.controls else (nlp.states, states) - ) + tau_var, tau_cx = (nlp.controls, controls) if "tau" in nlp.controls else (nlp.states, states) tau = nlp.get_var_from_states_or_controls("tau", states, controls) if fatigue is not None and "tau" in fatigue: tau_fatigue = fatigue["tau"] @@ -472,30 +425,19 @@ def __get_fatigable_tau( # Only homogeneous state_only is implemented yet n_state_only = sum([t.models.state_only for t in tau_fatigue]) if 0 < n_state_only < len(fatigue["tau"]): - raise NotImplementedError( - "fatigue list without homogeneous state_only flag is not supported yet" - ) - apply_to_joint_dynamics = sum( - [t.models.apply_to_joint_dynamics for t in tau_fatigue] - ) + raise NotImplementedError("fatigue list without homogeneous state_only flag is not supported yet") + apply_to_joint_dynamics = sum([t.models.apply_to_joint_dynamics for t in tau_fatigue]) if 0 < n_state_only < len(fatigue["tau"]): raise NotImplementedError( "fatigue list without homogeneous apply_to_joint_dynamics flag is not supported yet" ) if apply_to_joint_dynamics != 0: - raise NotImplementedError( - "apply_to_joint_dynamics is not implemented for joint torque" - ) + raise NotImplementedError("apply_to_joint_dynamics is not implemented for joint torque") if not tau_fatigue[0].models.split_controls and "tau" in nlp.controls: pass elif tau_fatigue[0].models.state_only: - tau = sum( - [ - DynamicsFunctions.get(tau_var[f"tau_{suffix}"], tau_cx) - for suffix in tau_suffix - ] - ) + tau = sum([DynamicsFunctions.get(tau_var[f"tau_{suffix}"], tau_cx) for suffix in tau_suffix]) else: tau = nlp.cx() for i, t in enumerate(tau_fatigue): @@ -503,10 +445,7 @@ def __get_fatigable_tau( for suffix in tau_suffix: model = t.models.models[suffix] tau_tp += ( - DynamicsFunctions.get( - nlp.states[f"tau_{suffix}_{model.dynamics_suffix()}"], - states, - )[i] + DynamicsFunctions.get(nlp.states[f"tau_{suffix}_{model.dynamics_suffix()}"], states)[i] * model.scaling ) tau = vertcat(tau, tau_tp) @@ -572,14 +511,10 @@ def torque_activations_driven( if with_ligament: tau += nlp.model.ligament_joint_torque()(q, qdot, nlp.parameters.cx) - external_forces = nlp.get_external_forces( - states, controls, algebraic_states, numerical_timeseries - ) + external_forces = nlp.get_external_forces(states, controls, algebraic_states, numerical_timeseries) dq = DynamicsFunctions.compute_qdot(nlp, q, qdot) - ddq = DynamicsFunctions.forward_dynamics( - nlp, q, qdot, tau, with_contact, external_forces - ) + ddq = DynamicsFunctions.forward_dynamics(nlp, q, qdot, tau, with_contact, external_forces) dq = horzcat(*[dq for _ in range(ddq.shape[1])]) @@ -634,26 +569,14 @@ def torque_derivative_driven( qdot = DynamicsFunctions.get(nlp.states["qdot"], states) tau = DynamicsFunctions.get(nlp.states["tau"], states) - tau = ( - tau + nlp.model.passive_joint_torque()(q, qdot, nlp.parameters.cx) - if with_passive_torque - else tau - ) - tau = ( - tau + nlp.model.ligament_joint_torque()(q, qdot, nlp.parameters.cx) - if with_ligament - else tau - ) + tau = tau + nlp.model.passive_joint_torque()(q, qdot, nlp.parameters.cx) if with_passive_torque else tau + tau = tau + nlp.model.ligament_joint_torque()(q, qdot, nlp.parameters.cx) if with_ligament else tau dq = DynamicsFunctions.compute_qdot(nlp, q, qdot) dtau = DynamicsFunctions.get(nlp.controls["taudot"], controls) - external_forces = nlp.get_external_forces( - states, controls, algebraic_states, numerical_timeseries - ) - ddq = DynamicsFunctions.forward_dynamics( - nlp, q, qdot, tau, with_contact, external_forces - ) + external_forces = nlp.get_external_forces(states, controls, algebraic_states, numerical_timeseries) + ddq = DynamicsFunctions.forward_dynamics(nlp, q, qdot, tau, with_contact, external_forces) dxdt = nlp.cx(nlp.states.shape, ddq.shape[1]) dxdt[nlp.states["q"].index, :] = horzcat(*[dq for _ in range(ddq.shape[1])]) dxdt[nlp.states["qdot"].index, :] = ddq @@ -706,24 +629,12 @@ def forces_from_torque_driven( q = nlp.get_var_from_states_or_controls("q", states, controls) qdot = nlp.get_var_from_states_or_controls("qdot", states, controls) tau = nlp.get_var_from_states_or_controls("tau", states, controls) - tau = ( - tau + nlp.model.passive_joint_torque()(q, qdot, nlp.parameters.cx) - if with_passive_torque - else tau - ) - tau = ( - tau + nlp.model.ligament_joint_torque()(q, qdot, nlp.parameters.cx) - if with_ligament - else tau - ) + tau = tau + nlp.model.passive_joint_torque()(q, qdot, nlp.parameters.cx) if with_passive_torque else tau + tau = tau + nlp.model.ligament_joint_torque()(q, qdot, nlp.parameters.cx) if with_ligament else tau - external_forces = nlp.get_external_forces( - states, controls, algebraic_states, numerical_timeseries - ) + external_forces = nlp.get_external_forces(states, controls, algebraic_states, numerical_timeseries) - return nlp.model.contact_forces()( - q, qdot, tau, external_forces, nlp.parameters.cx - ) + return nlp.model.contact_forces()(q, qdot, tau, external_forces, nlp.parameters.cx) @staticmethod def forces_from_torque_activation_driven( @@ -770,23 +681,11 @@ def forces_from_torque_activation_driven( qdot = nlp.get_var_from_states_or_controls("qdot", states, controls) tau_activations = nlp.get_var_from_states_or_controls("tau", states, controls) tau = nlp.model.torque()(tau_activations, q, qdot, nlp.parameters.cx) - tau = ( - tau + nlp.model.passive_joint_torque()(q, qdot, nlp.parameters.cx) - if with_passive_torque - else tau - ) - tau = ( - tau + nlp.model.ligament_joint_torque()(q, qdot, nlp.parameters.cx) - if with_ligament - else tau - ) + tau = tau + nlp.model.passive_joint_torque()(q, qdot, nlp.parameters.cx) if with_passive_torque else tau + tau = tau + nlp.model.ligament_joint_torque()(q, qdot, nlp.parameters.cx) if with_ligament else tau - external_forces = nlp.get_external_forces( - states, controls, algebraic_states, numerical_timeseries - ) - return nlp.model.contact_forces()( - q, qdot, tau, external_forces, nlp.parameters.cx - ) + external_forces = nlp.get_external_forces(states, controls, algebraic_states, numerical_timeseries) + return nlp.model.contact_forces()(q, qdot, tau, external_forces, nlp.parameters.cx) @staticmethod def muscles_driven( @@ -845,13 +744,9 @@ def muscles_driven( q = nlp.get_var_from_states_or_controls("q", states, controls) qdot = nlp.get_var_from_states_or_controls("qdot", states, controls) residual_tau = ( - DynamicsFunctions.__get_fatigable_tau(nlp, states, controls, fatigue) - if with_residual_torque - else None - ) - mus_activations = nlp.get_var_from_states_or_controls( - "muscles", states, controls + DynamicsFunctions.__get_fatigable_tau(nlp, states, controls, fatigue) if with_residual_torque else None ) + mus_activations = nlp.get_var_from_states_or_controls("muscles", states, controls) fatigue_states = None if fatigue is not None and "muscles" in fatigue: mus_fatigue = fatigue["muscles"] @@ -863,9 +758,7 @@ def muscles_driven( raise NotImplementedError( f"{fatigue_name} list without homogeneous state_only flag is not supported yet" ) - apply_to_joint_dynamics = sum( - [m.models.apply_to_joint_dynamics for m in mus_fatigue] - ) + apply_to_joint_dynamics = sum([m.models.apply_to_joint_dynamics for m in mus_fatigue]) if 0 < apply_to_joint_dynamics < len(fatigue["muscles"]): raise NotImplementedError( f"{fatigue_name} list without homogeneous apply_to_joint_dynamics flag is not supported yet" @@ -882,39 +775,21 @@ def muscles_driven( raise ValueError(f"{fatigue_name} must be of all same types") if n_state_only == 0: - mus_activations = DynamicsFunctions.get( - nlp.states[f"muscles_{dyn_suffix}"], states - ) + mus_activations = DynamicsFunctions.get(nlp.states[f"muscles_{dyn_suffix}"], states) if apply_to_joint_dynamics > 0: - fatigue_states = DynamicsFunctions.get( - nlp.states[f"muscles_{fatigue_suffix}"], states - ) - muscles_tau = DynamicsFunctions.compute_tau_from_muscle( - nlp, q, qdot, mus_activations, fatigue_states - ) + fatigue_states = DynamicsFunctions.get(nlp.states[f"muscles_{fatigue_suffix}"], states) + muscles_tau = DynamicsFunctions.compute_tau_from_muscle(nlp, q, qdot, mus_activations, fatigue_states) tau = muscles_tau + residual_tau if residual_tau is not None else muscles_tau - tau = ( - tau + nlp.model.passive_joint_torque()(q, qdot, nlp.parameters.cx) - if with_passive_torque - else tau - ) - tau = ( - tau + nlp.model.ligament_joint_torque()(q, qdot, nlp.parameters.cx) - if with_ligament - else tau - ) + tau = tau + nlp.model.passive_joint_torque()(q, qdot, nlp.parameters.cx) if with_passive_torque else tau + tau = tau + nlp.model.ligament_joint_torque()(q, qdot, nlp.parameters.cx) if with_ligament else tau tau = tau - nlp.model.friction_coefficients @ qdot if with_friction else tau dq = DynamicsFunctions.compute_qdot(nlp, q, qdot) - external_forces = nlp.get_external_forces( - states, controls, algebraic_states, numerical_timeseries - ) - ddq = DynamicsFunctions.forward_dynamics( - nlp, q, qdot, tau, with_contact, external_forces - ) + external_forces = nlp.get_external_forces(states, controls, algebraic_states, numerical_timeseries) + ddq = DynamicsFunctions.forward_dynamics(nlp, q, qdot, tau, with_contact, external_forces) dxdt = nlp.cx(nlp.states.shape, ddq.shape[1]) dxdt[nlp.states["q"].index, :] = horzcat(*[dq for _ in range(ddq.shape[1])]) dxdt[nlp.states["qdot"].index, :] = ddq @@ -922,12 +797,8 @@ def muscles_driven( has_excitation = True if "muscles" in nlp.states else False if has_excitation: mus_excitations = DynamicsFunctions.get(nlp.controls["muscles"], controls) - dmus = DynamicsFunctions.compute_muscle_dot( - nlp, mus_excitations, mus_activations - ) - dxdt[nlp.states["muscles"].index, :] = horzcat( - *[dmus for _ in range(ddq.shape[1])] - ) + dmus = DynamicsFunctions.compute_muscle_dot(nlp, mus_excitations, mus_activations) + dxdt[nlp.states["muscles"].index, :] = horzcat(*[dmus for _ in range(ddq.shape[1])]) if fatigue is not None and "muscles" in fatigue: dxdt = fatigue["muscles"].dynamics(dxdt, nlp, states, controls) @@ -978,36 +849,16 @@ def forces_from_muscle_driven( q = nlp.get_var_from_states_or_controls("q", states, controls) qdot = nlp.get_var_from_states_or_controls("qdot", states, controls) - residual_tau = ( - nlp.get_var_from_states_or_controls("tau", states, controls) - if "tau" in nlp.controls - else None - ) - mus_activations = nlp.get_var_from_states_or_controls( - "muscles", states, controls - ) - muscles_tau = DynamicsFunctions.compute_tau_from_muscle( - nlp, q, qdot, mus_activations - ) + residual_tau = nlp.get_var_from_states_or_controls("tau", states, controls) if "tau" in nlp.controls else None + mus_activations = nlp.get_var_from_states_or_controls("muscles", states, controls) + muscles_tau = DynamicsFunctions.compute_tau_from_muscle(nlp, q, qdot, mus_activations) tau = muscles_tau + residual_tau if residual_tau is not None else muscles_tau - tau = ( - tau + nlp.model.passive_joint_torque()(q, qdot, nlp.parameters.cx) - if with_passive_torque - else tau - ) - tau = ( - tau + nlp.model.ligament_joint_torque()(q, qdot, nlp.parameters.cx) - if with_ligament - else tau - ) + tau = tau + nlp.model.passive_joint_torque()(q, qdot, nlp.parameters.cx) if with_passive_torque else tau + tau = tau + nlp.model.ligament_joint_torque()(q, qdot, nlp.parameters.cx) if with_ligament else tau - external_forces = nlp.get_external_forces( - states, controls, algebraic_states, numerical_timeseries - ) - return nlp.model.contact_forces()( - q, qdot, tau, external_forces, nlp.parameters.cx - ) + external_forces = nlp.get_external_forces(states, controls, algebraic_states, numerical_timeseries) + return nlp.model.contact_forces()(q, qdot, tau, external_forces, nlp.parameters.cx) @staticmethod def joints_acceleration_driven( @@ -1048,9 +899,7 @@ def joints_acceleration_driven( qdot = nlp.get_var_from_states_or_controls("qdot", states, controls) qddot_joints = nlp.get_var_from_states_or_controls("qddot", states, controls) - qddot_root = nlp.model.forward_dynamics_free_floating_base()( - q, qdot, qddot_joints, nlp.parameters.cx - ) + qddot_root = nlp.model.forward_dynamics_free_floating_base()(q, qdot, qddot_joints, nlp.parameters.cx) qddot_reordered = nlp.model.reorder_qddot_root_joints(qddot_root, qddot_joints) qdot_mapped = nlp.variable_mappings["qdot"].to_first.map(qdot) @@ -1101,25 +950,15 @@ def compute_qdot(nlp, q: MX | SX, qdot: MX | SX): elif "q_roots" and "q_joints" in nlp.states: mapping = BiMapping( to_first=list(nlp.states["q_roots"].mapping.to_first.map_idx) - + [ - i + nlp.model.nb_root - for i in nlp.states["q_joints"].mapping.to_first.map_idx - ], + + [i + nlp.model.nb_root for i in nlp.states["q_joints"].mapping.to_first.map_idx], to_second=list(nlp.states["q_roots"].mapping.to_second.map_idx) - + [ - i + nlp.model.nb_root - for i in nlp.states["q_joints"].mapping.to_second.map_idx - ], + + [i + nlp.model.nb_root for i in nlp.states["q_joints"].mapping.to_second.map_idx], ) elif q in nlp.controls: mapping = nlp.controls["q"].mapping else: - raise RuntimeError( - "Your q key combination was not found in states or controls" - ) - return mapping.to_first.map( - nlp.model.reshape_qdot()(q, qdot, nlp.parameters.cx) - ) + raise RuntimeError("Your q key combination was not found in states or controls") + return mapping.to_first.map(nlp.model.reshape_qdot()(q, qdot, nlp.parameters.cx)) @staticmethod def forward_dynamics( @@ -1157,9 +996,7 @@ def forward_dynamics( elif "qdot" in nlp.controls: qdot_var_mapping = nlp.controls["qdot"].mapping.to_first else: - qdot_var_mapping = BiMapping( - [i for i in range(qdot.shape[0])], [i for i in range(qdot.shape[0])] - ).to_first + qdot_var_mapping = BiMapping([i for i in range(qdot.shape[0])], [i for i in range(qdot.shape[0])]).to_first external_forces = [] if external_forces is None else external_forces qddot = nlp.model.forward_dynamics(with_contact=with_contact)( @@ -1203,9 +1040,7 @@ def inverse_dynamics( Torques in tau """ if external_forces is None: - tau = nlp.model.inverse_dynamics(with_contact=with_contact)( - q, qdot, qddot, [], nlp.parameters.cx - ) + tau = nlp.model.inverse_dynamics(with_contact=with_contact)(q, qdot, qddot, [], nlp.parameters.cx) else: tau = nlp.model.inverse_dynamics(with_contact=with_contact)( q, qdot, qddot, external_forces, nlp.parameters.cx @@ -1213,9 +1048,7 @@ def inverse_dynamics( return tau # We ignore on purpose the mapping to keep zeros in the defects of the dynamic. @staticmethod - def compute_muscle_dot( - nlp, muscle_excitations: MX | SX, muscle_activations: MX | SX - ): + def compute_muscle_dot(nlp, muscle_excitations: MX | SX, muscle_activations: MX | SX): """ Easy accessor to derivative of muscle activations @@ -1233,9 +1066,7 @@ def compute_muscle_dot( The derivative of muscle activations """ - return nlp.model.muscle_activation_dot()( - muscle_excitations, muscle_activations, nlp.parameters.cx - ) + return nlp.model.muscle_activation_dot()(muscle_excitations, muscle_activations, nlp.parameters.cx) @staticmethod def compute_tau_from_muscle( @@ -1269,9 +1100,7 @@ def compute_tau_from_muscle( activations = type(q)() for k in range(len(nlp.controls["muscles"])): if fatigue_states is not None: - activations = vertcat( - activations, muscle_activations[k] * (1 - fatigue_states[k]) - ) + activations = vertcat(activations, muscle_activations[k] * (1 - fatigue_states[k])) else: activations = vertcat(activations, muscle_activations[k]) return nlp.model.muscle_joint_torque()(activations, q, qdot, nlp.parameters.cx)