diff --git a/bioptim/__init__.py b/bioptim/__init__.py index b194666d8..4e064c78c 100644 --- a/bioptim/__init__.py +++ b/bioptim/__init__.py @@ -172,7 +172,8 @@ from .dynamics.fatigue.effort_perception import EffortPerception, TauEffortPerception from .dynamics.ode_solver import OdeSolver, OdeSolverBase from .interfaces.solver_options import Solver -from .interfaces.biorbd_model import BiorbdModel, MultiBiorbdModel +from .interfaces.biorbd_model import BiorbdModel +from .interfaces.multi_biorbd_model import MultiBiorbdModel from .interfaces.biomodel import BioModel from .interfaces.holonomic_biomodel import HolonomicBioModel from .interfaces.variational_biomodel import VariationalBioModel diff --git a/bioptim/interfaces/biomodel.py b/bioptim/interfaces/biomodel.py index 22e08130c..808e7e830 100644 --- a/bioptim/interfaces/biomodel.py +++ b/bioptim/interfaces/biomodel.py @@ -21,7 +21,7 @@ def serialize(self) -> tuple[Callable, dict]: @property def friction_coefficients(self) -> MX: - """Get the coeffient of friction to apply to specified elements in the dymamics""" + """Get the coefficient of friction to apply to specified elements in the dynamics""" return MX() @property @@ -80,7 +80,7 @@ def segments(self) -> tuple: """Get all segments""" return () - def homogeneous_matrices_in_global(self, q, reference_idx, inverse=False) -> tuple: + def homogeneous_matrices_in_global(self, q, segment_id, inverse=False) -> tuple: """ Get the homogeneous matrices of all segments in the world frame, such as: P_R0 = T_R0_R1 * P_R1 @@ -89,9 +89,9 @@ def homogeneous_matrices_in_global(self, q, reference_idx, inverse=False) -> tup P_R1 the position of any point P in the segment R1 frame. """ - def homogeneous_matrices_in_child(self, *args) -> tuple: + def homogeneous_matrices_in_child(self, segment_id) -> MX: """ - Get the homogeneous matrices of all segments in their parent frame, + Get the homogeneous matrices of one segment in its parent frame, such as: P_R1 = T_R1_R2 * P_R2 with P_R1 the position of any point P in the segment R1 frame, with P_R2 the position of any point P in the segment R2 frame, diff --git a/bioptim/interfaces/biorbd_model.py b/bioptim/interfaces/biorbd_model.py index b70547b4e..8fd844673 100644 --- a/bioptim/interfaces/biorbd_model.py +++ b/bioptim/interfaces/biorbd_model.py @@ -171,15 +171,13 @@ def nb_root(self) -> int: def segments(self) -> tuple[biorbd.Segment]: return self.model.segments() - def homogeneous_matrices_in_global(self, q, reference_index, inverse=False): - val = self.model.globalJCS(GeneralizedCoordinates(q), reference_index) - if inverse: - return val.transpose() - else: - return val + def homogeneous_matrices_in_global(self, q, segment_id, inverse=False) -> biorbd.RotoTrans: + # Todo: one of the last ouput of BiorbdModel which is not a MX but a biorbd object + rt_matrix = self.model.globalJCS(GeneralizedCoordinates(q), segment_id) + return rt_matrix.transpose() if inverse else rt_matrix - def homogeneous_matrices_in_child(self, *args): - return self.model.localJCS(*args) + def homogeneous_matrices_in_child(self, segment_id) -> MX: + return self.model.localJCS(segment_id).to_mx() @property def mass(self) -> MX: @@ -614,664 +612,3 @@ def animate( return None else: return all_bioviz - - -class MultiBiorbdModel: - """ - This class allows to define multiple biorbd models for the same phase. - - - Attributes - ---------- - models : list[BiorbdModel] - The list of biorbd models to be handled in the optimal control program. - extra_models : list[BiorbdModel] - A list of extra biorbd models stored in the class for further use. - - Methods - ------- - variable_index() - Get the index of the variables in the global vector for a given model index. - nb_models() - Get the number of models. - nb_extra_models() - Get the number of extra models. - - """ - - def __init__( - self, - bio_model: tuple[str | biorbd.Model | BiorbdModel, ...], - extra_bio_models: tuple[str | biorbd.Model | BiorbdModel, ...] = (), - ): - self.models = [] - if not isinstance(bio_model, tuple): - raise ValueError("The models must be a 'str', 'biorbd.Model', 'bioptim.BiorbdModel'" " or a tuple of those") - - for model in bio_model: - if isinstance(model, str): - self.models.append(BiorbdModel(model)) - elif isinstance(model, biorbd.Model): - self.models.append(BiorbdModel(model)) - elif isinstance(model, BiorbdModel): - self.models.append(model) - else: - raise ValueError("The models should be of type 'str', 'biorbd.Model' or 'bioptim.BiorbdModel'") - - if not isinstance(extra_bio_models, tuple): - raise ValueError("The models must be a 'str', 'biorbd.Model', 'bioptim.BiorbdModel'" " or a tuple of those") - - self.extra_models = [] - for model in extra_bio_models: - if isinstance(model, str): - self.extra_models.append(BiorbdModel(model)) - elif isinstance(model, biorbd.Model): - self.extra_models.append(BiorbdModel(model)) - elif isinstance(model, BiorbdModel): - self.extra_models.append(model) - else: - raise ValueError("The models should be of type 'str', 'biorbd.Model' or 'bioptim.BiorbdModel'") - - def __getitem__(self, index): - return self.models[index] - - def deep_copy(self, *args): - raise NotImplementedError("Deep copy is not implemented yet for MultiBiorbdModel class") - - @property - def path(self) -> (list[str], list[str]): - return [model.path for model in self.models], [model.path for model in self.extra_models] - - def copy(self): - return MultiBiorbdModel(tuple(self.path[0]), tuple(self.path[1])) - - def serialize(self) -> tuple[Callable, dict]: - return MultiBiorbdModel, dict(bio_model=tuple(self.path[0]), extra_bio_models=tuple(self.path[1])) - - def variable_index(self, variable: str, model_index: int) -> range: - """ - Get the index of the variables in the global vector for a given model index - - Parameters - ---------- - variable: str - The variable to get the index from such as 'q', 'qdot', 'qddot', 'tau', 'contact' - model_index: int - The index of the model to get the index from - - Returns - ------- - range - The index of the variable in the global vector - """ - if variable == "q": - current_idx = 0 - for model in self.models[:model_index]: - current_idx += model.nb_q - return range(current_idx, current_idx + self.models[model_index].nb_q) - elif variable == "qdot": - current_idx = 0 - for model in self.models[:model_index]: - current_idx += model.nb_qdot - return range(current_idx, current_idx + self.models[model_index].nb_qdot) - elif variable == "qddot": - current_idx = 0 - for model in self.models[:model_index]: - current_idx += model.nb_qddot - return range(current_idx, current_idx + self.models[model_index].nb_qddot) - elif variable == "qddot_joints": - current_idx = 0 - for model in self.models[:model_index]: - current_idx += model.nb_qddot - model.nb_root - return range( - current_idx, current_idx + self.models[model_index].nb_qddot - self.models[model_index].nb_root - ) - elif variable == "qddot_root": - current_idx = 0 - for model in self.models[:model_index]: - current_idx += model.nb_root - return range(current_idx, current_idx + self.models[model_index].nb_root) - elif variable == "tau": - current_idx = 0 - for model in self.models[:model_index]: - current_idx += model.nb_tau - return range(current_idx, current_idx + self.models[model_index].nb_tau) - elif variable == "contact": - current_idx = 0 - for model in self.models[:model_index]: - current_idx += model.nb_rigid_contacts - return range(current_idx, current_idx + self.models[model_index].nb_rigid_contacts) - - @property - def nb_models(self) -> int: - """ - Get the number of models - - Returns - ------- - int - The number of models - """ - return len(self.models) - - @property - def nb_extra_models(self) -> int: - """ - Get the number of extra models - - Returns - ------- - int - The number of extra models - """ - return len(self.extra_models) - - @property - def gravity(self) -> MX: - return vertcat(*(model.gravity for model in self.models)) - - def set_gravity(self, new_gravity) -> None: - for model in self.models: - model.set_gravity(new_gravity) - return - - @property - def nb_tau(self) -> int: - return sum(model.nb_tau for model in self.models) - - @property - def nb_segments(self) -> int: - return sum(model.nb_segments for model in self.models) - - def segment_index(self, name) -> int: - raise NotImplementedError("segment_index is not implemented for MultiBiorbdModel") - - @property - def nb_quaternions(self) -> int: - return sum(model.nb_quaternions for model in self.models) - - @property - def nb_q(self) -> int: - return sum(model.nb_q for model in self.models) - - @property - def nb_qdot(self) -> int: - return sum(model.nb_qdot for model in self.models) - - @property - def nb_qddot(self) -> int: - return sum(model.nb_qddot for model in self.models) - - @property - def nb_root(self) -> int: - return sum(model.nb_root for model in self.models) - - @property - def segments(self) -> tuple[biorbd.Segment, ...]: - out = () - for model in self.models: - out += model.segments - return out - - def homogeneous_matrices_in_global(self, q, reference_index, inverse=False): - raise NotImplementedError("homogeneous_matrices_in_global is not implemented for MultiBiorbdModel") - - def homogeneous_matrices_in_child(self, *args) -> tuple: - raise NotImplementedError("homogeneous_matrices_in_child is not implemented for MultiBiorbdModel") - - @property - def mass(self) -> MX: - return vertcat(*(model.mass for model in self.models)) - - def center_of_mass(self, q) -> MX: - out = MX() - for i, model in enumerate(self.models): - q_model = q[self.variable_index("q", i)] - out = vertcat(out, model.center_of_mass(q_model)) - return out - - def center_of_mass_velocity(self, q, qdot) -> MX: - out = MX() - for i, model in enumerate(self.models): - q_model = q[self.variable_index("q", i)] - qdot_model = qdot[self.variable_index("qdot", i)] - out = vertcat( - out, - model.center_of_mass_velocity(q_model, qdot_model), - ) - return out - - def center_of_mass_acceleration(self, q, qdot, qddot) -> MX: - out = MX() - for i, model in enumerate(self.models): - q_model = q[self.variable_index("q", i)] - qdot_model = qdot[self.variable_index("qdot", i)] - qddot_model = qddot[self.variable_index("qddot", i)] - out = vertcat( - out, - model.center_of_mass_acceleration(q_model, qdot_model, qddot_model), - ) - return out - - def mass_matrix(self, q) -> list[MX]: - out = [] - for i, model in enumerate(self.models): - q_model = q[self.variable_index("q", i)] - out += [model.mass_matrix(q_model)] - return out - - def non_linear_effects(self, q, qdot) -> list[MX]: - out = [] - for i, model in enumerate(self.models): - q_model = q[self.variable_index("q", i)] - qdot_model = qdot[self.variable_index("qdot", i)] - out += [model.non_linear_effects(q_model, qdot_model)] - return out - - def angular_momentum(self, q, qdot) -> MX: - out = MX() - for i, model in enumerate(self.models): - q_model = q[self.variable_index("q", i)] - qdot_model = qdot[self.variable_index("qdot", i)] - out = vertcat( - out, - model.angular_momentum(q_model, qdot_model), - ) - return out - - def reshape_qdot(self, q, qdot, k_stab=1) -> MX: - out = MX() - for i, model in enumerate(self.models): - q_model = q[self.variable_index("q", i)] - qdot_model = qdot[self.variable_index("qdot", i)] - out = vertcat( - out, - model.reshape_qdot(q_model, qdot_model, k_stab), - ) - return out - - def segment_angular_velocity(self, q, qdot, idx) -> MX: - out = MX() - for i, model in enumerate(self.models): - q_model = q[self.variable_index("q", i)] - qdot_model = qdot[self.variable_index("qdot", i)] - out = vertcat( - out, - model.segment_angular_velocity(q_model, qdot_model, idx), - ) - return out - - @property - def name_dof(self) -> tuple[str, ...]: - return tuple([dof for model in self.models for dof in model.name_dof]) - - @property - def contact_names(self) -> tuple[str, ...]: - return tuple([contact for model in self.models for contact in model.contact_names]) - - @property - def nb_soft_contacts(self) -> int: - return sum(model.nb_soft_contacts for model in self.models) - - @property - def soft_contact_names(self) -> tuple[str, ...]: - return tuple([contact for model in self.models for contact in model.soft_contact_names]) - - def soft_contact(self, soft_contact_index, *args): - current_number_of_soft_contacts = 0 - out = [] - for model in self.models: - if soft_contact_index < current_number_of_soft_contacts + model.nb_soft_contacts: - out = model.soft_contact(soft_contact_index - current_number_of_soft_contacts, *args) - break - current_number_of_soft_contacts += model.nb_soft_contacts - return out - - @property - def muscle_names(self) -> tuple[str, ...]: - return tuple([muscle for model in self.models for muscle in model.muscle_names]) - - @property - def nb_muscles(self) -> int: - return sum(model.nb_muscles for model in self.models) - - def torque(self, tau_activations, q, qdot) -> MX: - out = MX() - for i, model in enumerate(self.models): - tau_activations_model = tau_activations[self.variable_index("tau", i)] - q_model = q[self.variable_index("q", i)] - qdot_model = qdot[self.variable_index("qdot", i)] - out = vertcat( - out, - model.torque( - tau_activations_model, - q_model, - qdot_model, - ), - ) - return out - - def forward_dynamics_free_floating_base(self, q, qdot, qddot_joints) -> MX: - out = MX() - for i, model in enumerate(self.models): - q_model = q[self.variable_index("q", i)] - qdot_model = qdot[self.variable_index("qdot", i)] - qddot_joints_model = qddot_joints[self.variable_index("qddot_joints", i)] - out = vertcat( - out, - model.forward_dynamics_free_floating_base( - q_model, - qdot_model, - qddot_joints_model, - ), - ) - return out - - def reorder_qddot_root_joints(self, qddot_root, qddot_joints): - out = MX() - for i, model in enumerate(self.models): - qddot_root_model = qddot_root[self.variable_index("qddot_root", i)] - qddot_joints_model = qddot_joints[self.variable_index("qddot_joints", i)] - out = vertcat( - out, - model.reorder_qddot_root_joints(qddot_root_model, qddot_joints_model), - ) - - return out - - def forward_dynamics(self, q, qdot, tau, external_forces=None, f_contacts=None) -> MX: - if f_contacts is not None or external_forces is not None: - raise NotImplementedError( - "External forces and contact forces are not implemented yet for MultiBiorbdModel." - ) - - out = MX() - for i, model in enumerate(self.models): - q_model = q[self.variable_index("q", i)] - qdot_model = qdot[self.variable_index("qdot", i)] - tau_model = tau[self.variable_index("tau", i)] - out = vertcat( - out, - model.forward_dynamics( - q_model, - qdot_model, - tau_model, - external_forces, - f_contacts, - ), - ) - return out - - def constrained_forward_dynamics(self, q, qdot, tau, external_forces=None) -> MX: - if external_forces is not None: - raise NotImplementedError("External forces are not implemented yet for MultiBiorbdModel.") - - out = MX() - for i, model in enumerate(self.models): - q_model = q[self.variable_index("q", i)] - qdot_model = qdot[self.variable_index("qdot", i)] - tau_model = tau[self.variable_index("qddot", i)] # Due to a bug in biorbd - out = vertcat( - out, - model.constrained_forward_dynamics( - q_model, - qdot_model, - tau_model, - external_forces, - ), - ) - return out - - def inverse_dynamics(self, q, qdot, qddot, external_forces=None, f_contacts=None) -> MX: - if f_contacts is not None or external_forces is not None: - raise NotImplementedError( - "External forces and contact forces are not implemented yet for MultiBiorbdModel." - ) - - out = MX() - for i, model in enumerate(self.models): - q_model = q[self.variable_index("q", i)] - qdot_model = qdot[self.variable_index("qdot", i)] - qddot_model = qddot[self.variable_index("qddot", i)] - out = vertcat( - out, - model.inverse_dynamics( - q_model, - qdot_model, - qddot_model, - external_forces, - f_contacts, - ), - ) - return out - - def contact_forces_from_constrained_forward_dynamics(self, q, qdot, tau, external_forces=None) -> MX: - if external_forces is not None: - raise NotImplementedError("External forces are not implemented yet for MultiBiorbdModel.") - out = MX() - for i, model in enumerate(self.models): - q_model = q[self.variable_index("q", i)] - qdot_model = qdot[self.variable_index("qdot", i)] - tau_model = tau[self.variable_index("qddot", i)] # Due to a bug in biorbd - out = vertcat( - out, - model.contact_forces_from_constrained_forward_dynamics( - q_model, - qdot_model, - tau_model, - external_forces, - ), - ) - return out - - def qdot_from_impact(self, q, qdot_pre_impact) -> MX: - out = MX() - for i, model in enumerate(self.models): - q_model = q[self.variable_index("q", i)] - qdot_pre_impact_model = qdot_pre_impact[self.variable_index("qdot", i)] - out = vertcat( - out, - model.qdot_from_impact( - q_model, - qdot_pre_impact_model, - ), - ) - return out - - def muscle_activation_dot(self, muscle_excitations) -> MX: - out = MX() - for model in self.models: - muscle_states = model.model.stateSet() # still call from Biorbd - for k in range(model.nb_muscles): - muscle_states[k].setExcitation(muscle_excitations[k]) - out = vertcat(out, model.model.activationDot(muscle_states).to_mx()) - return out - - def muscle_joint_torque(self, activations, q, qdot) -> MX: - out = MX() - for model in self.models: - muscles_states = model.model.stateSet() # still call from Biorbd - for k in range(model.nb_muscles): - muscles_states[k].setActivation(activations[k]) - q_biorbd = GeneralizedCoordinates(q) - qdot_biorbd = GeneralizedVelocity(qdot) - out = vertcat(out, model.model.muscularJointTorque(muscles_states, q_biorbd, qdot_biorbd).to_mx()) - return out - - def markers(self, q) -> Any | list[MX]: - out = [] - for i, model in enumerate(self.models): - q_model = q[self.variable_index("q", i)] - out.append(model.markers(q_model)) - return [item for sublist in out for item in sublist] - - @property - def nb_markers(self) -> int: - return sum(model.nb_markers for model in self.models) - - def marker_index(self, name): - raise NotImplementedError("marker_index is not implemented yet for MultiBiorbdModel") - - def marker(self, q, index, reference_segment_index=None) -> MX: - raise NotImplementedError("marker is not implemented yet for MultiBiorbdModel") - - @property - def nb_rigid_contacts(self) -> int: - """ - Returns the number of rigid contacts. - Example: - First contact with axis YZ - Second contact with axis Z - nb_rigid_contacts = 2 - """ - return sum(model.nb_rigid_contacts for model in self.models) - - @property - def nb_contacts(self) -> int: - """ - Returns the number of contact index. - Example: - First contact with axis YZ - Second contact with axis Z - nb_contacts = 3 - """ - return sum(model.nb_contacts for model in self.models) - - def rigid_contact_index(self, contact_index) -> tuple: - """ - Returns the axis index of this specific rigid contact. - Example: - First contact with axis YZ - Second contact with axis Z - rigid_contact_index(0) = (1, 2) - """ - - model_selected = None - for i, model in enumerate(self.models): - if contact_index in self.variable_index("contact", i): - model_selected = model - # Note: may not work if the contact_index is not in the first model - return model_selected.rigid_contact_index(contact_index) - - def marker_velocities(self, q, qdot, reference_index=None) -> list[MX]: - if reference_index is not None: - raise RuntimeError("marker_velocities is not implemented yet with reference_index for MultiBiorbdModel") - - out = [] - for i, model in enumerate(self.models): - q_model = q[self.variable_index("q", i)] - qdot_model = qdot[self.variable_index("qdot", i)] - out.extend( - model.marker_velocities(q_model, qdot_model, reference_index), - ) - return out - - def tau_max(self, q, qdot) -> tuple[MX, MX]: - out_max = MX() - out_min = MX() - for i, model in enumerate(self.models): - q_model = q[self.variable_index("q", i)] - qdot_model = qdot[self.variable_index("qdot", i)] - torque_max, torque_min = model.tau_max(q_model, qdot_model) - out_max = vertcat(out_max, torque_max) - out_min = vertcat(out_min, torque_min) - return out_max, out_min - - def rigid_contact_acceleration(self, q, qdot, qddot, contact_index, contact_axis) -> MX: - model_selected = None - model_idx = -1 - for i, model in enumerate(self.models): - if contact_index in self.variable_index("contact", i): - model_selected = model - model_idx = i - q_model = q[self.variable_index("q", model_idx)] - qdot_model = qdot[self.variable_index("qdot", model_idx)] - qddot_model = qddot[self.variable_index("qddot", model_idx)] - return model_selected.rigid_contact_acceleration(q_model, qdot_model, qddot_model, contact_index, contact_axis) - - @property - def nb_dof(self) -> int: - return sum(model.nb_dof for model in self.models) - - @property - def marker_names(self) -> tuple[str, ...]: - return tuple([name for model in self.models for name in model.marker_names]) - - def soft_contact_forces(self, q, qdot) -> MX: - out = MX() - for i, model in enumerate(self.models): - q_model = q[self.variable_index("q", i)] - qdot_model = qdot[self.variable_index("qdot", i)] - soft_contact_forces = model.soft_contact_forces(q_model, qdot_model) - out = vertcat(out, soft_contact_forces) - return out - - def reshape_fext_to_fcontact(self, fext: MX) -> biorbd.VecBiorbdVector: - raise NotImplementedError("reshape_fext_to_fcontact is not implemented yet for MultiBiorbdModel") - - def normalize_state_quaternions(self, x: MX | SX) -> MX | SX: - all_q_normalized = MX() - for i, model in enumerate(self.models): - q_model = x[self.variable_index("q", i)] # quaternions are only in q - q_normalized = model.normalize_state_quaternions(q_model) - all_q_normalized = vertcat(all_q_normalized, q_normalized) - idx_first_qdot = self.nb_q # assuming x = [q, qdot] - x_normalized = vertcat(all_q_normalized, x[idx_first_qdot:]) - - return x_normalized - - def contact_forces(self, q, qdot, tau, external_forces: list = None) -> MX: - if external_forces is not None: - raise NotImplementedError("contact_forces is not implemented yet with external_forces for MultiBiorbdModel") - - out = MX() - for i, model in enumerate(self.models): - q_model = q[self.variable_index("q", i)] - qdot_model = qdot[self.variable_index("qdot", i)] - tau_model = tau[self.variable_index("tau", i)] - - contact_forces = model.contact_forces(q_model, qdot_model, tau_model, external_forces) - out = vertcat(out, contact_forces) - - return out - - def passive_joint_torque(self, q, qdot) -> MX: - out = MX() - for i, model in enumerate(self.models): - q_model = q[self.variable_index("q", i)] - qdot_model = qdot[self.variable_index("qdot", i)] - out = vertcat(out, model.passive_joint_torque(q_model, qdot_model)) - return out - - def ligament_joint_torque(self, q, qdot) -> MX: - out = MX() - for i, model in enumerate(self.models): - q_model = q[self.variable_index("q", i)] - qdot_model = qdot[self.variable_index("qdot", i)] - out = vertcat(out, model.ligament_joint_torque(q_model, qdot_model)) - return out - - def ranges_from_model(self, variable: str): - return [the_range for model in self.models for the_range in model.ranges_from_model(variable)] - - def bounds_from_ranges(self, variables: str | list[str, ...], mapping: BiMapping | BiMappingList = None) -> Bounds: - return bounds_from_ranges(self, variables, mapping) - - def _q_mapping(self, mapping: BiMapping = None) -> dict: - return _q_mapping(self, mapping) - - def _qdot_mapping(self, mapping: BiMapping = None) -> dict: - return _qdot_mapping(self, mapping) - - def _qddot_mapping(self, mapping: BiMapping = None) -> dict: - return _qddot_mapping(self, mapping) - - def lagrangian(self): - raise NotImplementedError("lagrangian is not implemented yet for MultiBiorbdModel") - - def partitioned_forward_dynamics(self, q_u, qdot_u, tau, external_forces=None, f_contacts=None, q_v_init=None): - raise NotImplementedError("partitioned_forward_dynamics is not implemented yet for MultiBiorbdModel") - - @staticmethod - def animate(solution: Any, show_now: bool = True, tracked_markers: list = None, **kwargs: Any) -> None | list: - raise NotImplementedError("animate is not implemented yet for MultiBiorbdModel") diff --git a/bioptim/interfaces/multi_biorbd_model.py b/bioptim/interfaces/multi_biorbd_model.py new file mode 100644 index 000000000..d683f5670 --- /dev/null +++ b/bioptim/interfaces/multi_biorbd_model.py @@ -0,0 +1,791 @@ +from typing import Callable, Any + +import biorbd_casadi as biorbd +from biorbd_casadi import ( + GeneralizedCoordinates, + GeneralizedVelocity, + GeneralizedTorque, + GeneralizedAcceleration, +) +from casadi import SX, MX, vertcat, horzcat, norm_fro +import numpy as np + +from ..misc.utils import check_version +from ..limits.path_conditions import Bounds +from ..misc.mapping import BiMapping, BiMappingList +from .biorbd_model import _q_mapping, _qdot_mapping, _qddot_mapping, BiorbdModel, bounds_from_ranges + + +class MultiBiorbdModel: + """ + This class allows to define multiple biorbd models for the same phase. + + + Attributes + ---------- + models : list[BiorbdModel] + The list of biorbd models to be handled in the optimal control program. + extra_models : list[BiorbdModel] + A list of extra biorbd models stored in the class for further use. + + Methods + ------- + variable_index() + Get the index of the variables in the global vector for a given model index. + nb_models() + Get the number of models. + nb_extra_models() + Get the number of extra models. + + """ + + def __init__( + self, + bio_model: tuple[str | biorbd.Model | BiorbdModel, ...], + extra_bio_models: tuple[str | biorbd.Model | BiorbdModel, ...] = (), + ): + self.models = [] + if not isinstance(bio_model, tuple): + raise ValueError("The models must be a 'str', 'biorbd.Model', 'bioptim.BiorbdModel'" " or a tuple of those") + + for model in bio_model: + if isinstance(model, str): + self.models.append(BiorbdModel(model)) + elif isinstance(model, biorbd.Model): + self.models.append(BiorbdModel(model)) + elif isinstance(model, BiorbdModel): + self.models.append(model) + else: + raise ValueError("The models should be of type 'str', 'biorbd.Model' or 'bioptim.BiorbdModel'") + + if not isinstance(extra_bio_models, tuple): + raise ValueError("The models must be a 'str', 'biorbd.Model', 'bioptim.BiorbdModel'" " or a tuple of those") + + self.extra_models = [] + for model in extra_bio_models: + if isinstance(model, str): + self.extra_models.append(BiorbdModel(model)) + elif isinstance(model, biorbd.Model): + self.extra_models.append(BiorbdModel(model)) + elif isinstance(model, BiorbdModel): + self.extra_models.append(model) + else: + raise ValueError("The models should be of type 'str', 'biorbd.Model' or 'bioptim.BiorbdModel'") + + def __getitem__(self, index): + return self.models[index] + + def deep_copy(self, *args): + raise NotImplementedError("Deep copy is not implemented yet for MultiBiorbdModel class") + + @property + def path(self) -> (list[str], list[str]): + return [model.path for model in self.models], [model.path for model in self.extra_models] + + def copy(self): + return MultiBiorbdModel(tuple(self.path[0]), tuple(self.path[1])) + + def serialize(self) -> tuple[Callable, dict]: + return MultiBiorbdModel, dict(bio_model=tuple(self.path[0]), extra_bio_models=tuple(self.path[1])) + + def variable_index(self, variable: str, model_index: int) -> range: + """ + Get the index of the variables in the global vector for a given model index + + Parameters + ---------- + variable: str + The variable to get the index from such as 'q', 'qdot', 'qddot', 'tau', 'contact', 'markers' + model_index: int + The index of the model to get the index from + + Returns + ------- + range + The index of the variable in the global vector + """ + current_idx = 0 + + if variable == "q": + for model in self.models[:model_index]: + current_idx += model.nb_q + return range(current_idx, current_idx + self.models[model_index].nb_q) + + elif variable == "qdot": + for model in self.models[:model_index]: + current_idx += model.nb_qdot + return range(current_idx, current_idx + self.models[model_index].nb_qdot) + + elif variable == "qddot": + for model in self.models[:model_index]: + current_idx += model.nb_qddot + return range(current_idx, current_idx + self.models[model_index].nb_qddot) + + elif variable == "qddot_joints": + for model in self.models[:model_index]: + current_idx += model.nb_qddot - model.nb_root + return range( + current_idx, current_idx + self.models[model_index].nb_qddot - self.models[model_index].nb_root + ) + + elif variable == "qddot_root": + for model in self.models[:model_index]: + current_idx += model.nb_root + return range(current_idx, current_idx + self.models[model_index].nb_root) + + elif variable == "tau": + for model in self.models[:model_index]: + current_idx += model.nb_tau + return range(current_idx, current_idx + self.models[model_index].nb_tau) + + elif variable == "contact": + for model in self.models[:model_index]: + current_idx += model.nb_rigid_contacts + return range(current_idx, current_idx + self.models[model_index].nb_rigid_contacts) + + elif variable == "markers": + for model in self.models[:model_index]: + current_idx += model.nb_markers + return range(current_idx, current_idx + self.models[model_index].nb_markers) + + elif variable == "segment": + for model in self.models[:model_index]: + current_idx += model.nb_segments + return range(current_idx, current_idx + self.models[model_index].nb_segments) + + else: + raise ValueError( + "The variable must be 'q', 'qdot', 'qddot', 'tau', 'contact' or 'markers'" f" and {variable} was sent." + ) + + def global_variable_id(self, variable: str, model_index: int, model_variable_id: int) -> int: + """ + Get the id of the variable in the global vector for a given model index + + Parameters + ---------- + variable: str + The variable to get the index from such as 'q', 'qdot', 'qddot', 'tau', 'contact', 'markers' + model_index: int + The index of the model to get the index from + model_variable_id: int + The id of the variable in the model vector + + Returns + ------- + int + The id of the variable in the global vector + """ + return self.variable_index(variable, model_index)[model_variable_id] + + def local_variable_id(self, variable: str, global_index: int) -> tuple[int, int]: + """ + Get the id of the variable in the local vector and the model index for a given index of the global vector + + Parameters + ---------- + variable: str + The variable to get the index from such as 'q', 'qdot', 'qddot', 'tau', 'contact', 'markers' + global_index: int + The index of the variable in the global vector + + Returns + ------- + tuple(int, int) + The id of the variable in the local vector and the model index + """ + + for model_id, model in enumerate(self.models): + if global_index in self.variable_index(variable, model_id): + return global_index - self.variable_index(variable, model_id)[0], model_id + + @property + def nb_models(self) -> int: + """ + Get the number of models + + Returns + ------- + int + The number of models + """ + return len(self.models) + + @property + def nb_extra_models(self) -> int: + """ + Get the number of extra models + + Returns + ------- + int + The number of extra models + """ + return len(self.extra_models) + + @property + def gravity(self) -> MX: + return vertcat(*(model.gravity for model in self.models)) + + def set_gravity(self, new_gravity) -> None: + for model in self.models: + model.set_gravity(new_gravity) + return + + @property + def nb_tau(self) -> int: + return sum(model.nb_tau for model in self.models) + + @property + def nb_segments(self) -> int: + return sum(model.nb_segments for model in self.models) + + def segment_index(self, name) -> int: + raise NotImplementedError("segment_index is not implemented for MultiBiorbdModel") + + @property + def nb_quaternions(self) -> int: + return sum(model.nb_quaternions for model in self.models) + + @property + def nb_q(self) -> int: + return sum(model.nb_q for model in self.models) + + @property + def nb_qdot(self) -> int: + return sum(model.nb_qdot for model in self.models) + + @property + def nb_qddot(self) -> int: + return sum(model.nb_qddot for model in self.models) + + @property + def nb_root(self) -> int: + return sum(model.nb_root for model in self.models) + + @property + def segments(self) -> tuple[biorbd.Segment, ...]: + out = () + for model in self.models: + out += model.segments + return out + + def homogeneous_matrices_in_global(self, q, segment_id, inverse=False) -> biorbd.RotoTrans: + local_segment_id, model_id = self.local_variable_id("segment", segment_id) + q_model = q[self.variable_index("q", model_id)] + return self.models[model_id].homogeneous_matrices_in_global(q_model, local_segment_id, inverse) + + def homogeneous_matrices_in_child(self, segment_id) -> MX: + local_id, model_id = self.local_variable_id("segment", segment_id) + return self.models[model_id].homogeneous_matrices_in_child(local_id) + + @property + def mass(self) -> MX: + return vertcat(*(model.mass for model in self.models)) + + def center_of_mass(self, q) -> MX: + out = MX() + for i, model in enumerate(self.models): + q_model = q[self.variable_index("q", i)] + out = vertcat(out, model.center_of_mass(q_model)) + return out + + def center_of_mass_velocity(self, q, qdot) -> MX: + out = MX() + for i, model in enumerate(self.models): + q_model = q[self.variable_index("q", i)] + qdot_model = qdot[self.variable_index("qdot", i)] + out = vertcat( + out, + model.center_of_mass_velocity(q_model, qdot_model), + ) + return out + + def center_of_mass_acceleration(self, q, qdot, qddot) -> MX: + out = MX() + for i, model in enumerate(self.models): + q_model = q[self.variable_index("q", i)] + qdot_model = qdot[self.variable_index("qdot", i)] + qddot_model = qddot[self.variable_index("qddot", i)] + out = vertcat( + out, + model.center_of_mass_acceleration(q_model, qdot_model, qddot_model), + ) + return out + + def mass_matrix(self, q) -> list[MX]: + out = [] + for i, model in enumerate(self.models): + q_model = q[self.variable_index("q", i)] + out += [model.mass_matrix(q_model)] + return out + + def non_linear_effects(self, q, qdot) -> list[MX]: + out = [] + for i, model in enumerate(self.models): + q_model = q[self.variable_index("q", i)] + qdot_model = qdot[self.variable_index("qdot", i)] + out += [model.non_linear_effects(q_model, qdot_model)] + return out + + def angular_momentum(self, q, qdot) -> MX: + out = MX() + for i, model in enumerate(self.models): + q_model = q[self.variable_index("q", i)] + qdot_model = qdot[self.variable_index("qdot", i)] + out = vertcat( + out, + model.angular_momentum(q_model, qdot_model), + ) + return out + + def reshape_qdot(self, q, qdot, k_stab=1) -> MX: + out = MX() + for i, model in enumerate(self.models): + q_model = q[self.variable_index("q", i)] + qdot_model = qdot[self.variable_index("qdot", i)] + out = vertcat( + out, + model.reshape_qdot(q_model, qdot_model, k_stab), + ) + return out + + def segment_angular_velocity(self, q, qdot, idx) -> MX: + out = MX() + for i, model in enumerate(self.models): + q_model = q[self.variable_index("q", i)] + qdot_model = qdot[self.variable_index("qdot", i)] + out = vertcat( + out, + model.segment_angular_velocity(q_model, qdot_model, idx), + ) + return out + + @property + def name_dof(self) -> tuple[str, ...]: + return tuple([dof for model in self.models for dof in model.name_dof]) + + @property + def contact_names(self) -> tuple[str, ...]: + return tuple([contact for model in self.models for contact in model.contact_names]) + + @property + def nb_soft_contacts(self) -> int: + return sum(model.nb_soft_contacts for model in self.models) + + @property + def soft_contact_names(self) -> tuple[str, ...]: + return tuple([contact for model in self.models for contact in model.soft_contact_names]) + + def soft_contact(self, soft_contact_index, *args): + current_number_of_soft_contacts = 0 + out = [] + for model in self.models: + if soft_contact_index < current_number_of_soft_contacts + model.nb_soft_contacts: + out = model.soft_contact(soft_contact_index - current_number_of_soft_contacts, *args) + break + current_number_of_soft_contacts += model.nb_soft_contacts + return out + + @property + def muscle_names(self) -> tuple[str, ...]: + return tuple([muscle for model in self.models for muscle in model.muscle_names]) + + @property + def nb_muscles(self) -> int: + return sum(model.nb_muscles for model in self.models) + + def torque(self, tau_activations, q, qdot) -> MX: + out = MX() + for i, model in enumerate(self.models): + tau_activations_model = tau_activations[self.variable_index("tau", i)] + q_model = q[self.variable_index("q", i)] + qdot_model = qdot[self.variable_index("qdot", i)] + out = vertcat( + out, + model.torque( + tau_activations_model, + q_model, + qdot_model, + ), + ) + return out + + def forward_dynamics_free_floating_base(self, q, qdot, qddot_joints) -> MX: + out = MX() + for i, model in enumerate(self.models): + q_model = q[self.variable_index("q", i)] + qdot_model = qdot[self.variable_index("qdot", i)] + qddot_joints_model = qddot_joints[self.variable_index("qddot_joints", i)] + out = vertcat( + out, + model.forward_dynamics_free_floating_base( + q_model, + qdot_model, + qddot_joints_model, + ), + ) + return out + + def reorder_qddot_root_joints(self, qddot_root, qddot_joints): + out = MX() + for i, model in enumerate(self.models): + qddot_root_model = qddot_root[self.variable_index("qddot_root", i)] + qddot_joints_model = qddot_joints[self.variable_index("qddot_joints", i)] + out = vertcat( + out, + model.reorder_qddot_root_joints(qddot_root_model, qddot_joints_model), + ) + + return out + + def forward_dynamics(self, q, qdot, tau, external_forces=None, f_contacts=None) -> MX: + if f_contacts is not None or external_forces is not None: + raise NotImplementedError( + "External forces and contact forces are not implemented yet for MultiBiorbdModel." + ) + + out = MX() + for i, model in enumerate(self.models): + q_model = q[self.variable_index("q", i)] + qdot_model = qdot[self.variable_index("qdot", i)] + tau_model = tau[self.variable_index("tau", i)] + out = vertcat( + out, + model.forward_dynamics( + q_model, + qdot_model, + tau_model, + external_forces, + f_contacts, + ), + ) + return out + + def constrained_forward_dynamics(self, q, qdot, tau, external_forces=None) -> MX: + if external_forces is not None: + raise NotImplementedError("External forces are not implemented yet for MultiBiorbdModel.") + + out = MX() + for i, model in enumerate(self.models): + q_model = q[self.variable_index("q", i)] + qdot_model = qdot[self.variable_index("qdot", i)] + tau_model = tau[self.variable_index("qddot", i)] # Due to a bug in biorbd + out = vertcat( + out, + model.constrained_forward_dynamics( + q_model, + qdot_model, + tau_model, + external_forces, + ), + ) + return out + + def inverse_dynamics(self, q, qdot, qddot, external_forces=None, f_contacts=None) -> MX: + if f_contacts is not None or external_forces is not None: + raise NotImplementedError( + "External forces and contact forces are not implemented yet for MultiBiorbdModel." + ) + + out = MX() + for i, model in enumerate(self.models): + q_model = q[self.variable_index("q", i)] + qdot_model = qdot[self.variable_index("qdot", i)] + qddot_model = qddot[self.variable_index("qddot", i)] + out = vertcat( + out, + model.inverse_dynamics( + q_model, + qdot_model, + qddot_model, + external_forces, + f_contacts, + ), + ) + return out + + def contact_forces_from_constrained_forward_dynamics(self, q, qdot, tau, external_forces=None) -> MX: + if external_forces is not None: + raise NotImplementedError("External forces are not implemented yet for MultiBiorbdModel.") + out = MX() + for i, model in enumerate(self.models): + q_model = q[self.variable_index("q", i)] + qdot_model = qdot[self.variable_index("qdot", i)] + tau_model = tau[self.variable_index("qddot", i)] # Due to a bug in biorbd + out = vertcat( + out, + model.contact_forces_from_constrained_forward_dynamics( + q_model, + qdot_model, + tau_model, + external_forces, + ), + ) + return out + + def qdot_from_impact(self, q, qdot_pre_impact) -> MX: + out = MX() + for i, model in enumerate(self.models): + q_model = q[self.variable_index("q", i)] + qdot_pre_impact_model = qdot_pre_impact[self.variable_index("qdot", i)] + out = vertcat( + out, + model.qdot_from_impact( + q_model, + qdot_pre_impact_model, + ), + ) + return out + + def muscle_activation_dot(self, muscle_excitations) -> MX: + out = MX() + for model in self.models: + muscle_states = model.model.stateSet() # still call from Biorbd + for k in range(model.nb_muscles): + muscle_states[k].setExcitation(muscle_excitations[k]) + out = vertcat(out, model.model.activationDot(muscle_states).to_mx()) + return out + + def muscle_joint_torque(self, activations, q, qdot) -> MX: + out = MX() + for model in self.models: + muscles_states = model.model.stateSet() # still call from Biorbd + for k in range(model.nb_muscles): + muscles_states[k].setActivation(activations[k]) + q_biorbd = GeneralizedCoordinates(q) + qdot_biorbd = GeneralizedVelocity(qdot) + out = vertcat(out, model.model.muscularJointTorque(muscles_states, q_biorbd, qdot_biorbd).to_mx()) + return out + + def markers(self, q) -> Any | list[MX]: + out = [] + for i, model in enumerate(self.models): + q_model = q[self.variable_index("q", i)] + out.append(model.markers(q_model)) + return [item for sublist in out for item in sublist] + + @property + def nb_markers(self) -> int: + return sum(model.nb_markers for model in self.models) + + def marker_index(self, name): + for i, model in enumerate(self.models): + if name in model.marker_names: + marker_id = biorbd.marker_index(model.model, name) + return self.variable_index("markers", model_index=i)[marker_id] + + raise ValueError(f"{name} is not in the MultiBiorbdModel") + + def marker(self, q, index, reference_segment_index=None) -> MX: + local_marker_id, model_id = self.local_variable_id("markers", index) + q_model = q[self.variable_index("q", model_id)] + + return self.models[model_id].marker(q_model, local_marker_id, reference_segment_index) + + @property + def nb_rigid_contacts(self) -> int: + """ + Returns the number of rigid contacts. + Example: + First contact with axis YZ + Second contact with axis Z + nb_rigid_contacts = 2 + """ + return sum(model.nb_rigid_contacts for model in self.models) + + @property + def nb_contacts(self) -> int: + """ + Returns the number of contact index. + Example: + First contact with axis YZ + Second contact with axis Z + nb_contacts = 3 + """ + return sum(model.nb_contacts for model in self.models) + + def rigid_contact_index(self, contact_index) -> tuple: + """ + Returns the axis index of this specific rigid contact. + Example: + First contact with axis YZ + Second contact with axis Z + rigid_contact_index(0) = (1, 2) + """ + + model_selected = None + for i, model in enumerate(self.models): + if contact_index in self.variable_index("contact", i): + model_selected = model + # Note: may not work if the contact_index is not in the first model + return model_selected.rigid_contact_index(contact_index) + + def marker_velocities(self, q, qdot, reference_index=None) -> list[MX]: + if reference_index is not None: + raise RuntimeError("marker_velocities is not implemented yet with reference_index for MultiBiorbdModel") + + out = [] + for i, model in enumerate(self.models): + q_model = q[self.variable_index("q", i)] + qdot_model = qdot[self.variable_index("qdot", i)] + out.extend( + model.marker_velocities(q_model, qdot_model, reference_index), + ) + return out + + def tau_max(self, q, qdot) -> tuple[MX, MX]: + out_max = MX() + out_min = MX() + for i, model in enumerate(self.models): + q_model = q[self.variable_index("q", i)] + qdot_model = qdot[self.variable_index("qdot", i)] + torque_max, torque_min = model.tau_max(q_model, qdot_model) + out_max = vertcat(out_max, torque_max) + out_min = vertcat(out_min, torque_min) + return out_max, out_min + + def rigid_contact_acceleration(self, q, qdot, qddot, contact_index, contact_axis) -> MX: + model_selected = None + model_idx = -1 + for i, model in enumerate(self.models): + if contact_index in self.variable_index("contact", i): + model_selected = model + model_idx = i + q_model = q[self.variable_index("q", model_idx)] + qdot_model = qdot[self.variable_index("qdot", model_idx)] + qddot_model = qddot[self.variable_index("qddot", model_idx)] + return model_selected.rigid_contact_acceleration(q_model, qdot_model, qddot_model, contact_index, contact_axis) + + @property + def nb_dof(self) -> int: + return sum(model.nb_dof for model in self.models) + + @property + def marker_names(self) -> tuple[str, ...]: + return tuple([name for model in self.models for name in model.marker_names]) + + def soft_contact_forces(self, q, qdot) -> MX: + out = MX() + for i, model in enumerate(self.models): + q_model = q[self.variable_index("q", i)] + qdot_model = qdot[self.variable_index("qdot", i)] + soft_contact_forces = model.soft_contact_forces(q_model, qdot_model) + out = vertcat(out, soft_contact_forces) + return out + + def reshape_fext_to_fcontact(self, fext: MX) -> biorbd.VecBiorbdVector: + raise NotImplementedError("reshape_fext_to_fcontact is not implemented yet for MultiBiorbdModel") + + def normalize_state_quaternions(self, x: MX | SX) -> MX | SX: + all_q_normalized = MX() + for i, model in enumerate(self.models): + q_model = x[self.variable_index("q", i)] # quaternions are only in q + q_normalized = model.normalize_state_quaternions(q_model) + all_q_normalized = vertcat(all_q_normalized, q_normalized) + idx_first_qdot = self.nb_q # assuming x = [q, qdot] + x_normalized = vertcat(all_q_normalized, x[idx_first_qdot:]) + + return x_normalized + + def contact_forces(self, q, qdot, tau, external_forces: list = None) -> MX: + if external_forces is not None: + raise NotImplementedError("contact_forces is not implemented yet with external_forces for MultiBiorbdModel") + + out = MX() + for i, model in enumerate(self.models): + q_model = q[self.variable_index("q", i)] + qdot_model = qdot[self.variable_index("qdot", i)] + tau_model = tau[self.variable_index("tau", i)] + + contact_forces = model.contact_forces(q_model, qdot_model, tau_model, external_forces) + out = vertcat(out, contact_forces) + + return out + + def passive_joint_torque(self, q, qdot) -> MX: + out = MX() + for i, model in enumerate(self.models): + q_model = q[self.variable_index("q", i)] + qdot_model = qdot[self.variable_index("qdot", i)] + out = vertcat(out, model.passive_joint_torque(q_model, qdot_model)) + return out + + def ligament_joint_torque(self, q, qdot) -> MX: + out = MX() + for i, model in enumerate(self.models): + q_model = q[self.variable_index("q", i)] + qdot_model = qdot[self.variable_index("qdot", i)] + out = vertcat(out, model.ligament_joint_torque(q_model, qdot_model)) + return out + + def ranges_from_model(self, variable: str): + return [the_range for model in self.models for the_range in model.ranges_from_model(variable)] + + def bounds_from_ranges(self, variables: str | list[str, ...], mapping: BiMapping | BiMappingList = None) -> Bounds: + return bounds_from_ranges(self, variables, mapping) + + def _q_mapping(self, mapping: BiMapping = None) -> dict: + return _q_mapping(self, mapping) + + def _qdot_mapping(self, mapping: BiMapping = None) -> dict: + return _qdot_mapping(self, mapping) + + def _qddot_mapping(self, mapping: BiMapping = None) -> dict: + return _qddot_mapping(self, mapping) + + def lagrangian(self): + raise NotImplementedError("lagrangian is not implemented yet for MultiBiorbdModel") + + def partitioned_forward_dynamics(self, q_u, qdot_u, tau, external_forces=None, f_contacts=None, q_v_init=None): + raise NotImplementedError("partitioned_forward_dynamics is not implemented yet for MultiBiorbdModel") + + @staticmethod + def animate(solution: Any, show_now: bool = True, tracked_markers: list = None, **kwargs: Any) -> None | list: + try: + import bioviz + except ModuleNotFoundError: + raise RuntimeError("bioviz must be install to animate the model") + + check_version(bioviz, "2.3.0", "2.4.0") + + states = solution.states + if not isinstance(states, (list, tuple)): + states = [states] + + if tracked_markers is None: + tracked_markers = [None] * len(states) + + all_bioviz = [] + for idx_phase, data in enumerate(states): + # This calls each of the function that modify the internal dynamic model based on the parameters + nlp = solution.ocp.nlp[idx_phase] + + if isinstance(nlp.model, MultiBiorbdModel): + if nlp.model.nb_models > 1: + raise NotImplementedError( + f"Animation is only implemented for MultiBiorbdModel with 1 model." + f" There are {nlp.model.nb_models} models in the phase {idx_phase}." + ) + else: + model = nlp.model.models[0] + + biorbd_model: BiorbdModel = model + + all_bioviz.append(bioviz.Viz(biorbd_model.path, **kwargs)) + all_bioviz[-1].load_movement(solution.ocp.nlp[idx_phase].variable_mappings["q"].to_second.map(data["q"])) + + if tracked_markers[idx_phase] is not None: + all_bioviz[-1].load_experimental_markers(tracked_markers[idx_phase]) + + if show_now: + b_is_visible = [True] * len(all_bioviz) + while sum(b_is_visible): + for i, b in enumerate(all_bioviz): + if b.vtk_window.is_active: + b.update() + else: + b_is_visible[i] = False + return None + else: + return all_bioviz diff --git a/bioptim/optimization/solution.py b/bioptim/optimization/solution.py index 8e88d3c46..2ea0c5f63 100644 --- a/bioptim/optimization/solution.py +++ b/bioptim/optimization/solution.py @@ -1831,13 +1831,16 @@ def animate( # assuming that all the models or the same type. self._check_models_comes_from_same_super_class() - self.ocp.nlp[0].model.animate( + + all_bioviz = self.ocp.nlp[0].model.animate( solution=data_to_animate, show_now=show_now, tracked_markers=tracked_markers, **kwargs, ) + return all_bioviz + def _check_models_comes_from_same_super_class(self): """Check that all the models comes from the same super class""" for i, nlp in enumerate(self.ocp.nlp): diff --git a/tests/shard1/test_biorbd_multi_model.py b/tests/shard1/test_biorbd_multi_model.py index bee616db5..f5c0563fc 100644 --- a/tests/shard1/test_biorbd_multi_model.py +++ b/tests/shard1/test_biorbd_multi_model.py @@ -8,6 +8,7 @@ BiMappingList, BoundsList, ) +from ..utils import TestUtils def test_biorbd_model_import(): @@ -47,25 +48,40 @@ def test_biorbd_model(): ), ) - nb_q = models.nb_q - nb_qdot = models.nb_qdot - nb_qddot = models.nb_qddot - nb_root = models.nb_root - nb_tau = models.nb_tau - nb_quaternions = models.nb_quaternions - nb_segments = models.nb_segments - nb_muscles = models.nb_muscles - nb_soft_contacts = models.nb_soft_contacts - nb_markers = models.nb_markers - nb_rigid_contacts = models.nb_rigid_contacts - nb_contacts = models.nb_contacts - nb_dof = models.nb_dof - - name_dof = models.name_dof - contact_names = models.contact_names - soft_contact_names = models.soft_contact_names - marker_names = models.marker_names - muscle_names = models.muscle_names + assert models.nb_q == 6 + assert models.nb_qdot == 6 + assert models.nb_qddot == 6 + assert models.nb_root == 2 + assert models.nb_tau == 6 + assert models.nb_quaternions == 0 + assert models.nb_segments == 6 + assert models.nb_muscles == 0 + assert models.nb_soft_contacts == 0 + assert models.nb_markers == 12 + assert models.nb_rigid_contacts == 0 + assert models.nb_contacts == 0 + assert models.nb_dof == 6 + assert models.nb_models == 2 + assert models.nb_extra_models == 3 + + assert models.name_dof == ("Seg1_RotX", "Seg2_RotX", "Seg3_RotX", "Seg1_RotX", "Seg2_RotX", "Seg3_RotX") + assert models.contact_names == () + assert models.soft_contact_names == () + assert models.marker_names == ( + "marker_1", + "marker_2", + "marker_3", + "marker_4", + "marker_5", + "marker_6", + "marker_1", + "marker_2", + "marker_3", + "marker_4", + "marker_5", + "marker_6", + ) + assert models.muscle_names == () variable_mappings = BiMappingList() variable_mappings.add("q", to_second=[0, 1, 2, 3, 4, 5], to_first=[0, 1, 2, 3, 4, 5]) @@ -73,332 +89,302 @@ def test_biorbd_model(): variable_mappings.add("qddot", to_second=[0, 1, 2, 3, 4, 5], to_first=[0, 1, 2, 3, 4, 5]) variable_mappings.add("tau", to_second=[None, 0, 1, None, 0, 2], to_first=[1, 2, 5]) - np.random.seed(42) - q = MX(np.random.random((nb_q,))) - qdot = MX(np.random.random((nb_qdot,))) - qddot = MX(np.random.random((nb_qddot,))) - tau = MX(np.random.random((nb_tau,))) - qddot_joints = MX(np.random.random((nb_tau - nb_root,))) - f_ext = MX(6) - f_contact = MX(6) - muscle_excitations = MX(np.random.random((nb_muscles,))) - # model_deep_copied = models.deep_copy() # TODO: Fix deep copy models.copy() models.serialize() + models.set_gravity(np.array([0, 0, -3])) - model_gravity_modified = Function("Gravity", [], [models.gravity])()["o0"] + TestUtils.assert_equal(models.gravity, np.array([0, 0, -3, 0, 0, -3])) + models.set_gravity(np.array([0, 0, -9.81])) with pytest.raises(NotImplementedError, match="segment_index is not implemented for MultiBiorbdModel"): segment_index = models.segment_index("Seg1") - segments = models.segments - - with pytest.raises( - NotImplementedError, match="homogeneous_matrices_in_global is not implemented for MultiBiorbdModel" - ): - Function("RT_parent", [], [models.homogeneous_matrices_in_global(q[:3], 0, 0).to_mx()])()["o0"] - - with pytest.raises( - NotImplementedError, match="homogeneous_matrices_in_child is not implemented for MultiBiorbdModel" - ): - Function("RT_child", [], [models.homogeneous_matrices_in_child(0)[0].to_mx()])()["o0"] - mass = Function("Mass", [], [models.mass])()["o0"] - center_of_mass = Function("CoM", [], [models.center_of_mass(q)])()["o0"] - center_of_mass_velocity = Function("CoMdot", [], [models.center_of_mass_velocity(q, qdot)])()["o0"] - center_of_mass_acceleration = Function("CoMddot", [], [models.center_of_mass_acceleration(q, qdot, qdot)])()["o0"] - mass_matrix = Function("MassMatrix", [], models.mass_matrix(q))() - non_linear_effects = Function("NonLinearEffect", [], models.non_linear_effects(q, qdot))() - angular_momentum = Function("AngMom", [], [models.angular_momentum(q, qdot)])()["o0"] - reshape_qdot = Function("GetQdot", [], [models.reshape_qdot(q, qdot, 1)])()["o0"] - segment_angular_velocity = Function("SegmentAngMom", [], [models.segment_angular_velocity(q, qdot, 0)])()["o0"] - soft_contact = Function("SoftContact", [], [models.soft_contact(0, 0)])()[ - "o0" - ] # TODO: Fix soft contact (biorbd call error) - - with pytest.raises(RuntimeError, match="Close the actuator model before calling torqueMax"): - Function("TorqueFromActivation", [], [models.torque(tau, q, qdot)])()[ - "o0" - ] # TODO: Fix torque (Close the actuator model before calling torqueMax) - - forward_dynamics_free_floating_base = Function( - "RootForwardDynamics", [], [models.forward_dynamics_free_floating_base(q, qdot, qddot_joints)] - )()["o0"] - forward_dynamics = Function("ForwardDynamics", [], [models.forward_dynamics(q, qdot, tau)])()["o0"] - - constrained_forward_dynamics = Function( - "ConstrainedForwardDynamics", [], [models.constrained_forward_dynamics(q, qdot, tau, None)] - )()["o0"] - with pytest.raises(NotImplementedError, match="External forces are not implemented yet for MultiBiorbdModel."): - Function("ConstrainedForwardDynamics", [], [models.constrained_forward_dynamics(q, qdot, tau, f_ext)])()["o0"] - - inverse_dynamics = Function("InverseDynamics", [], [models.inverse_dynamics(q, qdot, tau)])()["o0"] - - contact_forces_from_constrained_dynamics = Function( - "ContactForcesFromDynamics", [], [models.contact_forces_from_constrained_forward_dynamics(q, qdot, tau, None)] - )()["o0"] - with pytest.raises(NotImplementedError, match="External forces are not implemented yet for MultiBiorbdModel."): - Function( - "ContactForcesFromDynamics", - [], - [models.contact_forces_from_constrained_forward_dynamics(q, qdot, tau, f_ext)], - )()["o0"] - - qdot_from_impact = Function("QdotFromImpact", [], [models.qdot_from_impact(q, qdot)])()["o0"] - muscle_activation_dot = Function("MusActivationdot", [], [models.muscle_activation_dot(muscle_excitations)])()["o0"] - muscle_joint_torque = Function("MusTau", [], [models.muscle_joint_torque(muscle_excitations, q, qdot)])()["o0"] - markers = Function("Markers", [], [models.markers(q)[0]])()["o0"] - - with pytest.raises(NotImplementedError, match="marker is not implemented yet for MultiBiorbdModel"): - Function("Marker", [], [models.marker(q[:3], index=0)])()["o0"] + assert len(models.segments) == 6 + assert isinstance(models.segments, tuple) + assert isinstance(models.segments[0], biorbd.biorbd.Segment) + + TestUtils.assert_equal( + # one of the last ouput of BiorbdModel which is not a MX but a biorbd object + models.homogeneous_matrices_in_global(np.array([1, 2, 3]), 0, 0).to_mx(), + np.array( + [ + [1.0, 0.0, 0.0, 0.0], + [0.0, 0.540302, -0.841471, 0.0], + [0.0, 0.841471, 0.540302, 0.0], + [0.0, 0.0, 0.0, 1.0], + ] + ), + ) + TestUtils.assert_equal( + models.homogeneous_matrices_in_child(4), + np.array([[1.0, 0.0, 0.0, 0.0], [0.0, 1.0, 0.0, 0.0], [0.0, 0.0, 1.0, -1.0], [0.0, 0.0, 0.0, 1.0]]), + ) - with pytest.raises(NotImplementedError, match="marker_index is not implemented yet for MultiBiorbdModel"): - models.marker_index("marker_3") + TestUtils.assert_equal(models.mass, np.array([3, 3])) + TestUtils.assert_equal( + models.center_of_mass(np.array([1, 2, 3, 4, 5, 6])), + np.array([-5.000000e-04, 8.433844e-01, -1.764446e-01, -5.000000e-04, -3.232674e-01, 1.485815e00]), + ) + TestUtils.assert_equal( + models.center_of_mass_velocity(np.array([1, 2.1, 3, 4.1, 5, 6.1]), np.array([1, 2.1, 3, 4.1, 5, 6])), + np.array([0.0, 0.425434, 0.638069, 0.0, -12.293126, 0.369492]), + ) + TestUtils.assert_equal( + models.center_of_mass_acceleration( + np.array([1, 2.1, 3, 4.1, 5, 6.1]), + np.array([1, 2.1, 3, 4.1, 5, 6]), + np.array([0.1, 0.2, 0.3, 0.4, 0.5, 0.6]), + ), + np.array([0.0, 0.481652, 6.105858, 0.0, -33.566971, -126.795179]), + ) - list_markers_velocities = models.marker_velocities(q, qdot) - assert isinstance(list_markers_velocities, list) - marker_velocities = Function( - "Markerdot", [], [list_markers_velocities[i] for i in range(len(list_markers_velocities))] - )()["o0"] + mass_matrices = models.mass_matrix(np.array([1, 2.1, 3, 4.1, 5, 6.1])) + assert len(mass_matrices) == 2 + TestUtils.assert_equal( + mass_matrices[0], + np.array( + [ + [2.711080e00, 3.783457e-01, 4.243336e-01], + [3.783457e-01, 9.999424e-01, -2.881681e-05], + [4.243336e-01, -2.881681e-05, 9.543311e-01], + ] + ), + ) + TestUtils.assert_equal( + mass_matrices[1], + np.array([[9.313616, 5.580191, 2.063886], [5.580191, 4.791997, 1.895999], [2.063886, 1.895999, 0.945231]]), + ) - with pytest.raises(RuntimeError, match="All dof must have their actuators set"): - Function("TauMax", [], [models.tau_max(q, qdot)])()[ - "o0" - ] # TODO: add an actuator model (AnaisFarr will do it when her PR will be merged) + nonlinear_effects = models.non_linear_effects(np.array([1, 2.1, 3, 4.1, 5, 6.1]), np.array([1, 2.1, 3, 4.1, 5, 6])) + assert len(nonlinear_effects) == 2 + TestUtils.assert_equal( + nonlinear_effects[0], + np.array([38.817401, -1.960653, -1.259441]), + ) + TestUtils.assert_equal( + nonlinear_effects[1], + np.array([322.060726, -22.147881, -20.660836]), + ) - # TODO: add a model with a contact to test this function - # rigid_contact_acceleration = models.rigid_contact_acceleration(q, qdot, qddot, 0, 0) + TestUtils.assert_equal( + models.angular_momentum(np.array([1, 2.1, 3, 4.1, 5, 6.1]), np.array([1, 2.1, 3, 4.1, 5, 6])), + np.array([3.001448e00, 0.000000e00, -2.168404e-19, 2.514126e01, 3.252607e-19, 0.000000e00]), + decimal=5, + ) - soft_contact_forces = Function("SoftContactForces", [], [models.soft_contact_forces(q, qdot)])()["o0"] + TestUtils.assert_equal( + models.reshape_qdot(np.array([1, 2.1, 3, 4.1, 5, 6.1]), np.array([1, 2.1, 3, 4.1, 5, 6]), 1), + np.array([1.0, 2.1, 3.0, 4.1, 5.0, 6.0]), + ) - reshape_fext_to_fcontact = Function("Fext_to_Fcontact", [], [models.models[0].reshape_fext_to_fcontact(f_ext)])()[ - "o0" - ] + TestUtils.assert_equal( + models.segment_angular_velocity(np.array([1, 2.1, 3, 4.1, 5, 6.1]), np.array([1, 2.1, 3, 4.1, 5, 6]), 1), + np.array([3.1, 0.0, 0.0, 9.1, 0.0, 0.0]), + ) - with pytest.raises( - NotImplementedError, match="reshape_fext_to_fcontact is not implemented yet for MultiBiorbdModel" - ): - Function("Fext_to_Fcontact", [], [models.reshape_fext_to_fcontact(f_ext)])()["o0"] + assert models.soft_contact(0, 0) == [] # TODO: Fix soft contact (biorbd call error) - normalize_state_quaternions = Function( - "normalize_state_quaternions", [], [models.normalize_state_quaternions(vertcat(q, qdot))] - )()["o0"] + with pytest.raises(RuntimeError, match="Close the actuator model before calling torqueMax"): + models.torque( + np.array([3.1, 0.0, 0.0, 9.1, 0.0, 0.0]), + np.array([3.1, 0.0, 0.0, 9.1, 0.0, 0.0]), + np.array([3.1, 0.0, 0.0, 9.1, 0.0, 0.0]), + ) # TODO: Fix torque (Close the actuator model before calling torqueMax) + + TestUtils.assert_equal( + models.forward_dynamics_free_floating_base( + np.array([1, 2.1, 3, 4.1, 5, 6.1]), + np.array([1, 2.1, 3, 4.1, 5, 6]), + np.array([3.1, 0.0, 0.0, 9.1, 0.0, 0.0]), + ), + np.array([-14.750679, -36.596107]), + ) - get_quaternion_idx = models.models[0].get_quaternion_idx() - get_quaternion_idx = models.models[1].get_quaternion_idx() + TestUtils.assert_equal( + models.forward_dynamics( + np.array([1, 2.1, 3, 4.1, 5, 6.1]), + np.array([1, 2.1, 3, 4.1, 5, 6]), + np.array([3.1, 1, 2, 9.1, 1, 2]), + ), + np.array([-16.092093, 9.049853, 10.570878, -121.783105, 154.820759, -20.664452]), + ) - contact_forces = Function("Fcontact", [], [models.contact_forces(q, qdot, tau, None)])()["o0"] - with pytest.raises( - NotImplementedError, match="contact_forces is not implemented yet with external_forces for MultiBiorbdModel" - ): - contact_forces = Function("Fcontact", [], [models.contact_forces(q, qdot, tau, f_ext)])()["o0"] + TestUtils.assert_equal( + models.constrained_forward_dynamics( + np.array([1, 2.1, 3, 4.1, 5, 6.1]), + np.array([1, 2.1, 3, 4.1, 5, 6]), + np.array([3.1, 1, 2, 9.1, 1, 2]), + ), + np.array([-16.092093, 9.049853, 10.570878, -121.783105, 154.820759, -20.664452]), + ) - passive_joint_torque = Function("PassiveTau", [], [models.passive_joint_torque(q, qdot)])()["o0"] - q_mapping = models._q_mapping(variable_mappings) - qdot_mapping = models._q_mapping(variable_mappings) - qddot_mapping = models._q_mapping(variable_mappings) - tau_mapping = models._q_mapping(variable_mappings) - bounds_from_ranges = BoundsList() - bounds_from_ranges["q"] = models.bounds_from_ranges("q", variable_mappings) - bounds_from_ranges["qdot"] = models.bounds_from_ranges("qdot", variable_mappings) + with pytest.raises(NotImplementedError, match="External forces are not implemented yet for MultiBiorbdModel."): + models.constrained_forward_dynamics( + np.array([1, 2.1, 3, 4.1, 5, 6.1]), + np.array([1, 2.1, 3, 4.1, 5, 6]), + np.array([3.1, 1, 2, 9.1, 1, 2]), + np.array([3.1, 1, 2, 9.1, 1, 2]), + ) - np.testing.assert_equal(nb_q, 6) - np.testing.assert_equal(nb_qdot, 6) - np.testing.assert_equal(nb_qddot, 6) - np.testing.assert_equal(nb_root, 2) - np.testing.assert_equal(nb_tau, 6) - np.testing.assert_equal(nb_quaternions, 0) - np.testing.assert_equal(nb_segments, 6) - np.testing.assert_equal(nb_muscles, 0) - np.testing.assert_equal(nb_soft_contacts, 0) - np.testing.assert_equal(nb_markers, 12) - np.testing.assert_equal(nb_rigid_contacts, 0) - np.testing.assert_equal(nb_contacts, 0) - np.testing.assert_equal(nb_dof, 6) - - np.testing.assert_equal(name_dof, ("Seg1_RotX", "Seg2_RotX", "Seg3_RotX", "Seg1_RotX", "Seg2_RotX", "Seg3_RotX")) - np.testing.assert_equal(contact_names, ()) - np.testing.assert_equal(soft_contact_names, ()) - np.testing.assert_equal( - marker_names, - ( - "marker_1", - "marker_2", - "marker_3", - "marker_4", - "marker_5", - "marker_6", - "marker_1", - "marker_2", - "marker_3", - "marker_4", - "marker_5", - "marker_6", + TestUtils.assert_equal( + models.inverse_dynamics( + np.array([1, 2.1, 3, 4.1, 5, 6.1]), + np.array([1, 2.1, 3, 4.1, 5, 6]), + np.array([3.1, 1, 2, 9.1, 1, 2]), ), + np.array([4.844876e01, 2.121037e-01, 1.964626e00, 4.165226e02, 3.721585e01, 1.906986e00]), + decimal=5, ) - np.testing.assert_equal(muscle_names, ()) - - for i in range(model_gravity_modified.shape[0]): - np.testing.assert_almost_equal(model_gravity_modified[i], DM(np.array([0, 0, -3, 0, 0, -3])[i])) - for i in range(mass.shape[0]): - np.testing.assert_almost_equal(mass[i], DM(np.array([3, 3][i]))) + TestUtils.assert_equal( + models.contact_forces_from_constrained_forward_dynamics( + np.array([1, 2.1, 3, 4.1, 5, 6.1]), + np.array([1, 2.1, 3, 4.1, 5, 6]), + np.array([3.1, 1, 2, 9.1, 1, 2]), + None, + ), + np.array([0.0, 0.0]), + ) - for i in range(center_of_mass.shape[0]): - np.testing.assert_almost_equal( - center_of_mass[i], DM(np.array([-0.0005, 1.28949, -0.875209, -0.0005, 1.30214, -1.43631])[i]), decimal=5 + with pytest.raises(NotImplementedError, match="External forces are not implemented yet for MultiBiorbdModel."): + models.contact_forces_from_constrained_forward_dynamics( + np.array([1, 2.1, 3, 4.1, 5, 6.1]), + np.array([1, 2.1, 3, 4.1, 5, 6]), + np.array([3.1, 1, 2, 9.1, 1, 2]), + np.array([3.1, 1, 2, 9.1, 1, 2]), ) - for i in range(center_of_mass_velocity.shape[0]): - np.testing.assert_almost_equal( - center_of_mass_velocity[i], DM(np.array([0, -0.0792032, 1.02386, 0, 1.20171, 1.19432])[i]), decimal=5 - ) + TestUtils.assert_equal( + models.qdot_from_impact( + np.array([1, 2.1, 3, 4.1, 5, 6.1]), + np.array([1, 2.1, 3, 4.1, 5, 6]), + ), + np.array([1.0, 2.1, 3.0, 4.1, 5.0, 6.0]), + ) - for i in range(center_of_mass_acceleration.shape[0]): - np.testing.assert_almost_equal( - center_of_mass_acceleration[i], DM(np.array([0, -1.2543, 0.750037, 0, -0.097267, 2.34977])[i]), decimal=5 - ) + TestUtils.assert_equal( + models.muscle_activation_dot( + np.array([1, 2.1, 3, 4.1, 5, 6.1]), + ), + np.array([], dtype=np.float64), + ) - for i in range(3): - for j in range(3): - np.testing.assert_almost_equal( - mass_matrix["o0"][i, j], - DM( - np.array([[8.99991, 5.14091, 1.44319], [5.14091, 4.23625, 1.61812], [1.44319, 1.61812, 0.954331]])[ - i, j - ] - ), - decimal=5, - ) - np.testing.assert_almost_equal( - mass_matrix["o1"][i, j], - DM( - np.array([[13.3131, 7.56109, 2.76416], [7.56109, 4.75431, 1.87716], [2.76416, 1.87716, 0.945231]])[ - i, j - ] - ), - decimal=5, - ) - - for j in range(3): - np.testing.assert_almost_equal( - non_linear_effects["o0"][j], DM(np.array([7.01844, 7.1652, 3.02573])[j]), decimal=5 - ) - np.testing.assert_almost_equal( - non_linear_effects["o1"][j], DM(np.array([10.3449, 6.41134, 2.68226])[j]), decimal=5 - ) + TestUtils.assert_equal( + models.muscle_joint_torque( + np.array([1, 2.1, 3, 4.1, 5, 6.1]), + np.array([1, 2.1, 3, 4.1, 5, 6.1]), + np.array([1, 2.1, 3, 4.1, 5, 6]), + ), + np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0]), + ) - for i in range(angular_momentum.shape[0]): - np.testing.assert_almost_equal( - angular_momentum[i], - DM(np.array([2.09041, -2.1684e-19, -1.0842e-19, 2.41968, 2.1684e-19, -2.1684e-19])[i]), - decimal=5, - ) + TestUtils.assert_equal( + models.markers( + np.array([1, 2.1, 3, 4.1, 5, 6.1]), + )[0], + np.array([0.0, 0.0, 0.0]), + ) - for i in range(reshape_qdot.shape[0]): - np.testing.assert_almost_equal( - reshape_qdot[i], - DM(np.array([0.0580836, 0.866176, 0.601115, 0.708073, 0.0205845, 0.96991])[i]), - decimal=5, - ) + TestUtils.assert_equal( + models.marker( + np.array([1, 2.1, 3, 4.1, 5, 6.1]), + index=1, + ), + np.array([0.0, 0.841471, -0.540302]), + ) - for i in range(segment_angular_velocity.shape[0]): - np.testing.assert_almost_equal( - segment_angular_velocity[i], DM(np.array([0.0580836, 0, 0, 0.708073, 0, 0])[i]), decimal=5 - ) + assert models.marker_index("marker_3") == 2 - np.testing.assert_equal(soft_contact.shape, (0, 0)) + markers_velocities = models.marker_velocities( + np.array([1, 2.1, 3, 4.1, 5, 6.1]), + np.array([1, 2.1, 3, 4.1, 5, 6]), + ) + assert isinstance(markers_velocities, list) - for i in range(forward_dynamics_free_floating_base.shape[0]): - np.testing.assert_almost_equal( - forward_dynamics_free_floating_base[i], DM(np.array([-1.16626, -0.997219])[i]), decimal=5 - ) + TestUtils.assert_equal( + markers_velocities[0], + np.array([0.0, 0.0, 0.0]), + ) + TestUtils.assert_equal( + markers_velocities[1], + np.array([0.0, 0.540302, 0.841471]), + ) - for i in range(forward_dynamics.shape[0]): - np.testing.assert_almost_equal( - forward_dynamics[i], - DM(np.array([1.06535, -3.78598, 2.27888, -0.348228, -0.778017, 0.113326])[i]), - decimal=5, - ) + with pytest.raises(RuntimeError, match="All dof must have their actuators set"): + models.tau_max( + np.array([1, 2.1, 3, 4.1, 5, 6.1]), + np.array([1, 2.1, 3, 4.1, 5, 6]), + ) # TODO: add an actuator model (AnaisFarr will do it when her PR will be merged) - for i in range(constrained_forward_dynamics.shape[0]): - np.testing.assert_almost_equal( - constrained_forward_dynamics[i], - DM(np.array([1.06535, -3.78598, 2.27888, -0.348228, -0.778017, 0.113326])[i]), - decimal=5, - ) + # TODO: add a model with a contact to test this function + # rigid_contact_acceleration = models.rigid_contact_acceleration(q, qdot, qddot, 0, 0) - for i in range(inverse_dynamics.shape[0]): - np.testing.assert_almost_equal( - inverse_dynamics[i], DM(np.array([13.2861, 11.6096, 4.70426, 15.4236, 9.54273, 3.96254])[i]), decimal=4 - ) + TestUtils.assert_equal( + models.soft_contact_forces( + np.array([1, 2.1, 3, 4.1, 5, 6.1]), + np.array([1, 2.1, 3, 4.1, 5, 6]), + ), + np.array([], dtype=np.float64), + ) - for i in range(contact_forces_from_constrained_dynamics.shape[0]): - np.testing.assert_almost_equal( - contact_forces_from_constrained_dynamics[i], - DM(np.array([0, 0])[i]), - decimal=5, - ) + with pytest.raises( + NotImplementedError, match="reshape_fext_to_fcontact is not implemented yet for MultiBiorbdModel" + ): + models.reshape_fext_to_fcontact(np.array([1, 2.1, 3, 4.1, 5, 6.1])) + + # this test doesn't properly test the function, but it's the best we can do for now + # we should add a quaternion to the model to test it + # anyway it's tested elsewhere. + TestUtils.assert_equal( + models.normalize_state_quaternions( + np.array([1, 2.1, 3, 4.1, 5, 6.1, 1, 2.1, 3, 4.1, 5, 6.6]), + ), + np.array([1.0, 2.1, 3.0, 4.1, 5.0, 6.1, 1.0, 2.1, 3.0, 4.1, 5.0, 6.6]), + ) - for i in range(qdot_from_impact.shape[0]): - np.testing.assert_almost_equal( - qdot_from_impact[i], - DM(np.array([0.0580836, 0.866176, 0.601115, 0.708073, 0.0205845, 0.96991])[i]), - decimal=5, - ) + TestUtils.assert_equal( + models.contact_forces( + np.array([1, 2.1, 3, 4.1, 5, 6.1]), + np.array([1, 2.1, 3, 4.1, 5, 6]), + np.array([3.1, 1, 2, 9.1, 1, 2]), + None, + ), + np.array([0.0, 0.0]), + ) - np.testing.assert_equal(muscle_activation_dot.shape, (0, 1)) - - for i in range(muscle_joint_torque.shape[0]): - np.testing.assert_almost_equal(muscle_joint_torque[i], DM(np.zeros((6,))[i]), decimal=5) - - for i in range(markers.shape[0]): - np.testing.assert_almost_equal(markers[i], DM(np.zeros((3,))[i]), decimal=5) - - for i in range(marker_velocities.shape[0]): - np.testing.assert_almost_equal(marker_velocities[i], DM(np.zeros((6,))[i]), decimal=5) - - np.testing.assert_equal(soft_contact_forces.shape, (0, 1)) - - np.testing.assert_equal(reshape_fext_to_fcontact.shape, (0, 0)) - - for i in range(normalize_state_quaternions.shape[0]): - np.testing.assert_almost_equal( - normalize_state_quaternions[i], - DM( - np.array( - [ - 0.37454, - 0.950714, - 0.731994, - 0.598658, - 0.156019, - 0.155995, - 0.0580836, - 0.866176, - 0.601115, - 0.708073, - 0.0205845, - 0.96991, - ] - )[i] - ), - decimal=5, + with pytest.raises( + NotImplementedError, match="contact_forces is not implemented yet with external_forces for MultiBiorbdModel" + ): + models.contact_forces( + np.array([1, 2.1, 3, 4.1, 5, 6.1]), + np.array([1, 2.1, 3, 4.1, 5, 6]), + np.array([3.1, 1, 2, 9.1, 1, 2]), + np.array([3.1, 1, 2, 9.1, 1, 2]), ) - np.testing.assert_equal(get_quaternion_idx, []) - - for i in range(contact_forces.shape[0]): - np.testing.assert_almost_equal(contact_forces[i], DM(np.array([0, 0])[i]), decimal=5) + TestUtils.assert_equal( + models.passive_joint_torque( + np.array([1, 2.1, 3, 4.1, 5, 6.1]), + np.array([1, 2.1, 3, 4.1, 5, 6]), + ), + np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0]), + ) - for i in range(passive_joint_torque.shape[0]): - np.testing.assert_almost_equal(passive_joint_torque[i], DM(np.zeros((6,))[i]), decimal=5) + q_mapping = models._q_mapping(variable_mappings) + qdot_mapping = models._q_mapping(variable_mappings) + qddot_mapping = models._q_mapping(variable_mappings) + tau_mapping = models._q_mapping(variable_mappings) np.testing.assert_equal(q_mapping["q"].to_first.map_idx, [0, 1, 2, 3, 4, 5]) np.testing.assert_equal(qdot_mapping["qdot"].to_first.map_idx, [0, 1, 2, 3, 4, 5]) np.testing.assert_equal(qddot_mapping["qddot"].to_first.map_idx, [0, 1, 2, 3, 4, 5]) np.testing.assert_equal(tau_mapping["tau"].to_first.map_idx, [1, 2, 5]) + bounds_from_ranges = BoundsList() + bounds_from_ranges["q"] = models.bounds_from_ranges("q", variable_mappings) + bounds_from_ranges["qdot"] = models.bounds_from_ranges("qdot", variable_mappings) + for key in bounds_from_ranges.keys(): if key == "q": expected = [ @@ -423,5 +409,71 @@ def test_biorbd_model(): np.testing.assert_almost_equal(bounds_from_ranges[key].min, DM(np.array(expected)), decimal=5) - assert models.nb_models == 2 - assert models.nb_extra_models == 3 + assert models.variable_index("q", 0) == range(0, 3) + assert models.variable_index("qdot", 1) == range(3, 6) + assert models.variable_index("tau", 0) == range(0, 3) + assert models.variable_index("qddot", 1) == range(3, 6) + assert models.variable_index("qddot_joints", 0) == range(0, 2) + assert models.variable_index("qddot_root", 1) == range(1, 2) + assert models.variable_index("contact", 0) == range(0, 0) + assert models.variable_index("markers", 0) == range(0, 6) + + variable_name = "wrong" + with pytest.raises( + ValueError, + match=f"The variable must be 'q', 'qdot', 'qddot', 'tau', 'contact' or 'markers'" + f" and {variable_name} was sent.", + ): + models.variable_index(variable_name, 0) + + assert models.global_variable_id("q", 0, 1) == 1 + assert models.global_variable_id("qdot", 1, 0) == 3 + assert models.global_variable_id("tau", 0, 1) == 1 + assert models.global_variable_id("qddot", 1, 0) == 3 + assert models.global_variable_id("qddot_joints", 0, 1) == 1 + assert models.global_variable_id("qddot_root", 1, 0) == 1 + + with pytest.raises(IndexError, match="range object index out of range"): + models.global_variable_id("contact", 0, 1) + + assert models.global_variable_id("markers", 0, 1) == 1 + + local_id, model_id = models.local_variable_id("q", 2) + assert local_id == 2 + assert model_id == 0 + + local_id, model_id = models.local_variable_id("qdot", 5) + assert local_id == 2 + assert model_id == 1 + + local_id, model_id = models.local_variable_id("tau", 2) + assert local_id == 2 + assert model_id == 0 + + local_id, model_id = models.local_variable_id("qddot", 5) + assert local_id == 2 + assert model_id == 1 + + local_id, model_id = models.local_variable_id("qddot_joints", 2) + assert local_id == 0 + assert model_id == 1 + + local_id, model_id = models.local_variable_id("qddot_root", 1) + assert local_id == 0 + assert model_id == 1 + + local_id, model_id = models.local_variable_id("markers", 2) + assert local_id == 2 + assert model_id == 0 + + local_id, model_id = models.local_variable_id("segment", 2) + assert local_id == 2 + assert model_id == 0 + + variable_name = "wrong" + with pytest.raises( + ValueError, + match=f"The variable must be 'q', 'qdot', 'qddot', 'tau', 'contact' or 'markers'" + f" and {variable_name} was sent.", + ): + models.local_variable_id(variable_name, 0)