From 22ef2297bab38d31e94c37cf99c9ce33a6f2899b Mon Sep 17 00:00:00 2001 From: eve-mac Date: Fri, 20 Sep 2024 09:58:35 -0400 Subject: [PATCH] multi_model --- .../getting_started/custom_dynamics.py | 2 +- .../muscle_excitations_tracker.py | 14 +- bioptim/limits/constraints.py | 3 +- bioptim/models/biorbd/biorbd_model.py | 2 +- bioptim/models/biorbd/multi_biorbd_model.py | 662 +++++++++++------- bioptim/models/holonomic_constraints.py | 6 +- tests/shard1/test__global_plots.py | 2 +- tests/shard1/test_biorbd_model_size.py | 67 +- tests/shard1/test_biorbd_multi_model.py | 144 ++-- 9 files changed, 512 insertions(+), 390 deletions(-) diff --git a/bioptim/examples/getting_started/custom_dynamics.py b/bioptim/examples/getting_started/custom_dynamics.py index 69faeb168..d5a6f8e0a 100644 --- a/bioptim/examples/getting_started/custom_dynamics.py +++ b/bioptim/examples/getting_started/custom_dynamics.py @@ -74,7 +74,7 @@ def custom_dynamics( # You can directly call biorbd function (as for ddq) or call bioptim accessor (as for dq) dq = DynamicsFunctions.compute_qdot(nlp, q, qdot) * my_additional_factor - ddq = nlp.model.forward_dynamics(q, qdot, tau, []) + ddq = nlp.model.forward_dynamics()(q, qdot, tau, []) # the user has to choose if want to return the explicit dynamics dx/dt = f(x,u,p) # as the first argument of DynamicsEvaluation or diff --git a/bioptim/examples/muscle_driven_ocp/muscle_excitations_tracker.py b/bioptim/examples/muscle_driven_ocp/muscle_excitations_tracker.py index 659915f34..4f2235fa9 100644 --- a/bioptim/examples/muscle_driven_ocp/muscle_excitations_tracker.py +++ b/bioptim/examples/muscle_driven_ocp/muscle_excitations_tracker.py @@ -12,7 +12,7 @@ import biorbd_casadi as biorbd import numpy as np -from casadi import MX, vertcat +from casadi import MX, SX, vertcat, horzcat from matplotlib import pyplot as plt from scipy.integrate import solve_ivp @@ -44,6 +44,7 @@ def generate_data( n_shooting: int, use_residual_torque: bool = True, phase_dynamics: PhaseDynamics = PhaseDynamics.SHARED_DURING_THE_PHASE, + use_sx: bool = False, ) -> tuple: """ Generate random data. If np.random.seed is defined before, it will always return the same results @@ -91,7 +92,8 @@ def generate_data( symbolic_controls = vertcat(*(symbolic_tau, symbolic_mus_controls)) symbolic_parameters = MX.sym("u", 0, 0) - nlp = NonLinearProgram(phase_dynamics=phase_dynamics) + nlp = NonLinearProgram(phase_dynamics=phase_dynamics, use_sx=use_sx) + nlp.cx = SX if use_sx else MX nlp.model = bio_model nlp.variable_mappings = { "q": BiMapping(range(n_q), range(n_q)), @@ -100,8 +102,6 @@ def generate_data( "tau": BiMapping(range(n_tau), range(n_tau)), "muscles": BiMapping(range(n_mus), range(n_mus)), } - markers_func = biorbd.to_casadi_func("ForwardKin", bio_model.markers, symbolic_q) - nlp.states = OptimizationVariableContainer(phase_dynamics=phase_dynamics) nlp.states_dot = OptimizationVariableContainer(phase_dynamics=phase_dynamics) nlp.controls = OptimizationVariableContainer(phase_dynamics=phase_dynamics) @@ -202,7 +202,7 @@ def dyn_interface(t, x, u): def add_to_data(i, q): X[:, i] = q - markers[:, :, i] = markers_func(q[:n_q]) + markers[:, :, i] = horzcat(*bio_model.markers()(q[:n_q])) x_init = np.array([0.0] * n_q + [0.0] * n_qdot + [0.5] * n_mus) add_to_data(0, x_init) @@ -374,10 +374,8 @@ def main(): n_frames = q.shape[1] markers = np.ndarray((3, n_mark, q.shape[1])) - symbolic_states = MX.sym("x", n_q, 1) - markers_func = biorbd.to_casadi_func("ForwardKin", bio_model.markers, symbolic_states) for i in range(n_frames): - markers[:, :, i] = markers_func(q[:, i]) + markers[:, :, i] = horzcat(*bio_model.markers()(q[:, i])) plt.figure("Markers") n_steps_ode = ocp.nlp[0].ode_solver.steps + 1 if ocp.nlp[0].ode_solver.is_direct_collocation else 1 diff --git a/bioptim/limits/constraints.py b/bioptim/limits/constraints.py index 8d8db49f4..82e7460dc 100644 --- a/bioptim/limits/constraints.py +++ b/bioptim/limits/constraints.py @@ -414,7 +414,8 @@ def qddot_equals_forward_dynamics( tau = tau + controller.model.ligament_joint_torque(q, qdot) if with_ligament else tau qddot = controller.controls["qddot"].cx if "qddot" in controller.controls else controller.states["qddot"].cx - qddot_fd = controller.model.forward_dynamics(with_contact=with_contact)(q, qdot, tau) + # @Ipuch: no f_ext allowed ? + qddot_fd = controller.model.forward_dynamics(with_contact=with_contact)(q, qdot, tau, []) return qddot - qddot_fd @staticmethod diff --git a/bioptim/models/biorbd/biorbd_model.py b/bioptim/models/biorbd/biorbd_model.py index 0538998ec..800dfee01 100644 --- a/bioptim/models/biorbd/biorbd_model.py +++ b/bioptim/models/biorbd/biorbd_model.py @@ -221,7 +221,7 @@ def center_of_mass_acceleration(self) -> Function: biorbd_return = self.model.CoMddot(q_biorbd, qdot_biorbd, qddot_biorbd, True).to_mx() casadi_fun = Function( "center_of_mass_acceleration", - [self.q, self.qdot, self.tau], + [self.q, self.qdot, self.qddot], [biorbd_return], ) return casadi_fun diff --git a/bioptim/models/biorbd/multi_biorbd_model.py b/bioptim/models/biorbd/multi_biorbd_model.py index a32e0f2e6..91c526c66 100644 --- a/bioptim/models/biorbd/multi_biorbd_model.py +++ b/bioptim/models/biorbd/multi_biorbd_model.py @@ -1,6 +1,6 @@ import biorbd_casadi as biorbd -from casadi import SX, MX, vertcat, Function -from typing import Callable, Any +from casadi import MX, vertcat, Function, horzcat +from typing import Callable from .biorbd_model import BiorbdModel from ..utils import _var_mapping @@ -65,6 +65,15 @@ def __init__( else: raise ValueError("The models should be of type 'str', 'biorbd.Model' or 'bioptim.BiorbdModel'") + # Declaration of MX variables of the right shape for the creation of CasADi Functions + self.q = MX.sym("q_mx", self.nb_q, 1) + self.qdot = MX.sym("qdot_mx", self.nb_qdot, 1) + self.qddot = MX.sym("qddot_mx", self.nb_qddot, 1) + self.qddot_roots = MX.sym("qddot_roots_mx", self.nb_root, 1) + self.qddot_joints = MX.sym("qddot_joints_mx", self.nb_qddot - self.nb_root, 1) + self.tau = MX.sym("tau_mx", self.nb_tau, 1) + self.muscle = MX.sym("muscle_mx", self.nb_muscles, 1) + def __getitem__(self, index): return self.models[index] @@ -217,10 +226,17 @@ def nb_extra_models(self) -> int: return len(self.extra_models) @property - def gravity(self) -> MX: - return vertcat(*(model.gravity for model in self.models)) + def gravity(self) -> Function: + biorbd_return = vertcat(*(model.gravity()['o0'] for model in self.models)) + casadi_fun = Function( + "gravity", + [MX()], + [biorbd_return], + ) + return casadi_fun def set_gravity(self, new_gravity) -> None: + # All models have the same gravity, but it could be changed if needed for model in self.models: model.set_gravity(new_gravity) return @@ -264,114 +280,149 @@ def segments(self) -> tuple[biorbd.Segment, ...]: out += (seg,) return out - def biorbd_homogeneous_matrices_in_global(self, segment_idx, inverse=False) -> biorbd.RotoTrans: - # Charbie todo remove + def homogeneous_matrices_in_global(self, segment_idx, inverse=False) -> Function: local_segment_id, model_id = self.local_variable_id("segment", segment_idx) q_model = self.models[model_id].q - return self.models[model_id].homogeneous_matrices_in_global(q_model, local_segment_id, inverse) - - def homogeneous_matrices_in_global(self, segment_idx, inverse=False) -> MX: - local_segment_id, model_id = self.local_variable_id("segment", segment_idx) - q_model = self.models[model_id].q - return self.models[model_id].homogeneous_matrices_in_global(local_segment_id, inverse)(q_model) - - def homogeneous_matrices_in_child(self, segment_id) -> MX: + biorbd_return = self.models[model_id].homogeneous_matrices_in_global(local_segment_id, inverse)(q_model) + casadi_fun = Function( + "homogeneous_matrices_in_global", + [self.models[model_id].q], + [biorbd_return], + ) + return casadi_fun + + def homogeneous_matrices_in_child(self, segment_id) -> Function: local_id, model_id = self.local_variable_id("segment", segment_id) - return self.models[model_id].homogeneous_matrices_in_child(local_id) + casadi_fun = self.models[model_id].homogeneous_matrices_in_child(local_id) + return casadi_fun @property - def mass(self) -> MX: - return vertcat(*(model.mass for model in self.models)) - - def center_of_mass(self) -> MX: - out = MX() + def mass(self) -> Function: + biorbd_return = vertcat(*(model.mass()['o0'] for model in self.models)) + casadi_fun = Function( + "mass", + [MX()], + [biorbd_return], + ) + return casadi_fun + + def center_of_mass(self) -> Function: + biorbd_return = MX() for i, model in enumerate(self.models): - q_model = model.q[self.variable_index("q", i)] - out = model.center_of_mass if i == 0 else vertcat(out, model.center_of_mass()(q_model)) - return out - - def center_of_mass_velocity(self) -> MX: - out = MX() + q_model = self.q[self.variable_index("q", i)] + biorbd_return = vertcat(biorbd_return, model.center_of_mass()(q_model)) + casadi_fun = Function( + "center_of_mass", + [self.q], + [biorbd_return], + ) + return casadi_fun + + def center_of_mass_velocity(self) -> Function: + biorbd_return = MX() for i, model in enumerate(self.models): - q_model = model.q[self.variable_index("q", i)] - qdot_model = model.qdot[self.variable_index("qdot", i)] - out = ( - model.center_of_mass_velocity()(q_model, qdot_model) - if i == 0 - else vertcat( - out, + q_model = self.q[self.variable_index("q", i)] + qdot_model = self.qdot[self.variable_index("qdot", i)] + biorbd_return = vertcat( + biorbd_return, model.center_of_mass_velocity()(q_model, qdot_model), ) - ) - return out - - def center_of_mass_acceleration(self) -> MX: - out = MX() + casadi_fun = Function( + "center_of_mass_velocity", + [self.q, self.qdot], + [biorbd_return], + ) + return casadi_fun + + def center_of_mass_acceleration(self) -> Function: + biorbd_return = MX() for i, model in enumerate(self.models): - q_model = model.q[self.variable_index("q", i)] - qdot_model = model.qdot[self.variable_index("qdot", i)] - qddot_model = model.qddot[self.variable_index("qddot", i)] - out = ( - model.center_of_mass_acceleration()(q_model, qdot_model, qddot_model) - if i == 0 - else vertcat( - out, + q_model = self.q[self.variable_index("q", i)] + qdot_model = self.qdot[self.variable_index("qdot", i)] + qddot_model = self.qddot[self.variable_index("qddot", i)] + biorbd_return = vertcat( + biorbd_return, model.center_of_mass_acceleration()(q_model, qdot_model, qddot_model), ) - ) - return out - - def mass_matrix(self) -> list[MX]: - out = [] + casadi_fun = Function( + "center_of_mass_acceleration", + [self.q, self.qdot, self.qddot], + [biorbd_return], + ) + return casadi_fun + + def mass_matrix(self) -> Function: + biorbd_return = [] for i, model in enumerate(self.models): - q_model = model.q[self.variable_index("q", i)] - out += [model.mass_matrix()(q_model)] - return out - - def non_linear_effects(self) -> list[MX]: - out = [] + q_model = self.q[self.variable_index("q", i)] + biorbd_return += [model.mass_matrix()(q_model)] + casadi_fun = Function( + "mass_matrix", + [self.q], + biorbd_return, + ) + return casadi_fun + + def non_linear_effects(self) -> Function: + biorbd_return = [] for i, model in enumerate(self.models): - q_model = model.q[self.variable_index("q", i)] - qdot_model = model.qdot[self.variable_index("qdot", i)] - out += [model.non_linear_effects()(q_model, qdot_model)] - return out - - def angular_momentum(self) -> MX: + q_model = self.q[self.variable_index("q", i)] + qdot_model = self.qdot[self.variable_index("qdot", i)] + biorbd_return += [model.non_linear_effects()(q_model, qdot_model)] + casadi_fun = Function( + "non_linear_effects", + [self.q, self.qdot], + biorbd_return, + ) + return casadi_fun + + def angular_momentum(self) -> Function: + biorbd_return = MX() for i, model in enumerate(self.models): q_model = self.q[self.variable_index("q", i)] qdot_model = self.qdot[self.variable_index("qdot", i)] - out = ( - model.angular_momentum()(q_model, qdot_model) - if i == 0 - else vertcat( - out, + biorbd_return = vertcat( + biorbd_return, model.angular_momentum()(q_model, qdot_model), ) - ) - return out - - # Charbie ici ... todo - def reshape_qdot(self, q, qdot, k_stab=1) -> MX: - out = MX() + casadi_fun = Function( + "angular_momentum", + [self.q, self.qdot], + [biorbd_return], + ) + return casadi_fun + + def reshape_qdot(self, k_stab=1) -> Function: + biorbd_return = 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), + q_model = self.q[self.variable_index("q", i)] + qdot_model = self.qdot[self.variable_index("qdot", i)] + biorbd_return = vertcat( + biorbd_return, + model.reshape_qdot(k_stab)(q_model, qdot_model), ) - return out - - def segment_angular_velocity(self, q, qdot, idx) -> MX: - out = MX() + casadi_fun = Function( + "reshape_qdot", + [self.q, self.qdot], + [biorbd_return], + ) + return casadi_fun + + def segment_angular_velocity(self, idx) -> Function: + biorbd_return = 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), + q_model = self.q[self.variable_index("q", i)] + qdot_model = self.qdot[self.variable_index("qdot", i)] + biorbd_return = vertcat( + biorbd_return, + model.segment_angular_velocity(idx)(q_model, qdot_model), ) - return out + casadi_fun = Function( + "segment_angular_velocity", + [self.q, self.qdot], + [biorbd_return], + ) + return casadi_fun @property def name_dof(self) -> tuple[str, ...]: @@ -390,6 +441,7 @@ 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): + # What does that function return? current_number_of_soft_contacts = 0 out = [] for model in self.models: @@ -407,167 +459,199 @@ def muscle_names(self) -> tuple[str, ...]: 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() + def torque(self) -> Function: + biorbd_return = 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 = self.muscle[self.variable_index("tau", i)] + q_model = self.q[self.variable_index("q", i)] + qdot_model = self.qdot[self.variable_index("qdot", i)] + biorbd_return = vertcat( + biorbd_return, + model.torque()( tau_activations_model, q_model, qdot_model, ), ) - return out + casadi_fun = Function( + "torque", + [self.muscle, self.q, self.qdot], + [biorbd_return], + ) + return casadi_fun - def forward_dynamics_free_floating_base(self, q, qdot, qddot_joints) -> MX: - out = MX() + + def forward_dynamics_free_floating_base(self) -> Function: + biorbd_return = 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 = self.q[self.variable_index("q", i)] + qdot_model = self.qdot[self.variable_index("qdot", i)] + qddot_joints_model = self.qddot_joints[self.variable_index("qddot_joints", i)] + biorbd_return = vertcat( + biorbd_return, + 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() + casadi_fun = Function( + "forward_dynamics_free_floating_base", + [self.q, self.qdot, self.qddot_joints], + [biorbd_return], + ) + return casadi_fun + + def reorder_qddot_root_joints(self): + biorbd_return = 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, + qddot_root_model = self.qddot_roots[self.variable_index("qddot_root", i)] + qddot_joints_model = self.qddot_joints[self.variable_index("qddot_joints", i)] + biorbd_return = vertcat( + biorbd_return, model.reorder_qddot_root_joints(qddot_root_model, qddot_joints_model), ) - return out + casadi_fun = Function( + "reorder_qddot_root_joints", + [self.qddot_roots, self.qddot_joints], + [biorbd_return], + ) + return casadi_fun - def forward_dynamics(self, with_contact, 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)] - if with_contact: - out = vertcat( - out, - model.constrained_forward_dynamics( - q_model, - qdot_model, - tau_model, - external_forces, - f_contacts, - ), - ) - else: - out = vertcat( - out, - model.forward_dynamics( - q_model, - qdot_model, - tau_model, - external_forces, - f_contacts, - ), - ) - return out + def forward_dynamics(self, with_contact) -> Function: + """External forces and contact forces are not implemented yet for MultiBiorbdModel.""" - 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." + biorbd_return = MX() + for i, model in enumerate(self.models): + q_model = self.q[self.variable_index("q", i)] + qdot_model = self.qdot[self.variable_index("qdot", i)] + tau_model = self.tau[self.variable_index("tau", i)] + biorbd_return = vertcat( + biorbd_return, + model.forward_dynamics(with_contact=with_contact)( + q_model, + qdot_model, + tau_model, + [], + ), ) + casadi_fun = Function( + "forward_dynamics", + [self.q, self.qdot, self.tau], + [biorbd_return], + ) + return casadi_fun - out = MX() + def inverse_dynamics(self) -> Function: + """External forces and contact forces are not implemented yet for MultiBiorbdModel.""" + + biorbd_return = 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 = self.q[self.variable_index("q", i)] + qdot_model = self.qdot[self.variable_index("qdot", i)] + qddot_model = self.qddot[self.variable_index("qddot", i)] + biorbd_return = vertcat( + biorbd_return, + 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() + casadi_fun = Function( + "inverse_dynamics", + [self.q, self.qdot, self.qddot], + [biorbd_return], + ) + return casadi_fun + + def contact_forces_from_constrained_forward_dynamics(self) -> Function: + """External forces are not implemented yet for MultiBiorbdModel.""" + biorbd_return = 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 = self.q[self.variable_index("q", i)] + qdot_model = self.qdot[self.variable_index("qdot", i)] + tau_model = self.tau[self.variable_index("qddot", i)] # Due to a bug in biorbd + biorbd_return = vertcat( + biorbd_return, + 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() + casadi_fun = Function( + "contact_forces_from_constrained_forward_dynamics", + [self.q, self.qdot, self.tau], + [biorbd_return], + ) + return casadi_fun + + def qdot_from_impact(self) -> Function: + biorbd_return = 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 = self.q[self.variable_index("q", i)] + qdot_pre_impact_model = self.qdot[self.variable_index("qdot", i)] + biorbd_return = vertcat( + biorbd_return, + model.qdot_from_impact()( q_model, qdot_pre_impact_model, ), ) - return out - - def muscle_activation_dot(self, muscle_excitations) -> MX: - out = MX() + casadi_fun = Function( + "qdot_from_impact", + [self.q, self.qdot], + [biorbd_return], + ) + return casadi_fun + + def muscle_activation_dot(self) -> Function: + biorbd_return = 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() + muscle_states[k].setExcitation(self.muscle[k]) + biorbd_return = vertcat(biorbd_return, model.model.activationDot(muscle_states).to_mx()) + casadi_fun = Function( + "muscle_activation_dot", + [self.muscle], + [biorbd_return], + ) + return casadi_fun + + def muscle_joint_torque(self) -> Function: + biorbd_return = MX() for i, model in enumerate(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_model = q[self.variable_index("q", i)] - qdot_model = qdot[self.variable_index("qdot", i)] - out = vertcat(out, model.model.muscularJointTorque(muscles_states, q_model, qdot_model).to_mx()) - return out - - def markers(self, q) -> Any | list[MX]: - out = [] + muscles_states[k].setActivation(self.activations[k]) + q_model = self.q[self.variable_index("q", i)] + qdot_model = self.qdot[self.variable_index("qdot", i)] + biorbd_return = vertcat(biorbd_return, model.model.muscularJointTorque(muscles_states, q_model, qdot_model).to_mx()) + casadi_fun = Function( + "muscle_joint_torque", + [self.muscle, self.q, self.qdot], + [biorbd_return], + ) + return casadi_fun + + def markers(self) -> Function: + biorbd_return = [] 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] + q_model = self.q[self.variable_index("q", i)] + biorbd_return += [model.markers()(q_model)] + biorbd_return = [item for sublist in biorbd_return for item in sublist] + casadi_fun = Function( + "markers", + [self.q], + [horzcat(*biorbd_return)], + ) + return casadi_fun @property def nb_markers(self) -> int: @@ -581,11 +665,16 @@ def marker_index(self, name): raise ValueError(f"{name} is not in the MultiBiorbdModel") - def marker(self, q, index, reference_segment_index=None) -> MX: + def marker(self, index, reference_segment_index=None) -> Function: 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) + q_model = self.q[self.variable_index("q", model_id)] + biorbd_return = self.models[model_id].marker(local_marker_id, reference_segment_index)(q_model) + casadi_fun = Function( + "marker", + [self.q], + [biorbd_return], + ) + return casadi_fun @property def nb_rigid_contacts(self) -> int: @@ -625,41 +714,56 @@ def rigid_contact_index(self, contact_index) -> tuple: # 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]: + def marker_velocities(self, reference_index=None) -> Function: if reference_index is not None: raise RuntimeError("marker_velocities is not implemented yet with reference_index for MultiBiorbdModel") - out = [] + biorbd_return = [] 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]: + q_model = self.q[self.variable_index("q", i)] + qdot_model = self.qdot[self.variable_index("qdot", i)] + biorbd_return += [model.marker_velocities(reference_index)(q_model, qdot_model)] + biorbd_return = [item for sublist in biorbd_return for item in sublist] + casadi_fun = Function( + "marker_velocities", + [self.q, self.qdot], + [horzcat(*biorbd_return)], + ) + return casadi_fun + + def tau_max(self) -> Function: 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) + q_model = self.q[self.variable_index("q", i)] + qdot_model = self.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: + casadi_fun = Function( + "tau_max", + [self.q, self.qdot], + [out_max, out_min], + ) + return casadi_fun + + def rigid_contact_acceleration(self, contact_index, contact_axis) -> Function: 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) + q_model = self.q[self.variable_index("q", model_idx)] + qdot_model = self.qdot[self.variable_index("qdot", model_idx)] + qddot_model = self.qddot[self.variable_index("qddot", model_idx)] + biorbd_return = model_selected.rigid_contact_acceleration(contact_index, contact_axis)(q_model, qdot_model, qddot_model) + casadi_fun = Function( + "rigid_contact_acceleration", + [self.q, self.qdot, self.qddot], + [biorbd_return], + ) + return casadi_fun @property def nb_dof(self) -> int: @@ -669,59 +773,79 @@ def nb_dof(self) -> int: 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() + def soft_contact_forces(self, q, qdot) -> Function: + biorbd_return = 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: + soft_contact_forces = model.soft_contact_forces()(q_model, qdot_model) + biorbd_return = vertcat(biorbd_return, soft_contact_forces) + casadi_fun = Function( + "soft_contact_forces", + [self.q, self.qdot], + [biorbd_return], + ) + return casadi_fun + + def reshape_fext_to_fcontact(self): 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() + def normalize_state_quaternions(self) -> Function: + biorbd_return = 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() + q_model = self.q[self.variable_index("q", i)] + q_normalized = model.normalize_state_quaternions()(q_model) + biorbd_return = vertcat(biorbd_return, q_normalized) + casadi_fun = Function( + "normalize_state_quaternions", + [self.q], + [biorbd_return], + ) + return casadi_fun + + def contact_forces(self) -> Function: + """external_forces is not implemented yet for MultiBiorbdModel""" + biorbd_return = 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)] + q_model = self.q[self.variable_index("q", i)] + qdot_model = self.qdot[self.variable_index("qdot", i)] + tau_model = self.tau[self.variable_index("tau", i)] - contact_forces = model.contact_forces(q_model, qdot_model, tau_model, external_forces) - out = vertcat(out, contact_forces) + contact_forces = model.contact_forces()(q_model, qdot_model, tau_model, []) + biorbd_return = vertcat(biorbd_return, contact_forces) + casadi_fun = Function( + "contact_forces", + [self.q, self.qdot, self.tau], + [biorbd_return], + ) + return casadi_fun - return out - def passive_joint_torque(self, q, qdot) -> MX: - out = MX() + def passive_joint_torque(self) -> Function: + biorbd_return = 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() + q_model = self.q[self.variable_index("q", i)] + qdot_model = self.qdot[self.variable_index("qdot", i)] + biorbd_return = vertcat(biorbd_return, model.passive_joint_torque()(q_model, qdot_model)) + casadi_fun = Function( + "passive_joint_torque", + [self.q, self.qdot], + [biorbd_return], + ) + return casadi_fun + + def ligament_joint_torque(self) -> Function: + biorbd_return = 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 + q_model = self.q[self.variable_index("q", i)] + qdot_model = self.qdot[self.variable_index("qdot", i)] + biorbd_return = vertcat(biorbd_return, model.ligament_joint_torque()(q_model, qdot_model)) + casadi_fun = Function( + "ligament_joint_torque", + [self.q, self.qdot], + [biorbd_return], + ) + return casadi_fun def ranges_from_model(self, variable: str): return [the_range for model in self.models for the_range in model.ranges_from_model(variable)] diff --git a/bioptim/models/holonomic_constraints.py b/bioptim/models/holonomic_constraints.py index 248a97e3e..486baa928 100644 --- a/bioptim/models/holonomic_constraints.py +++ b/bioptim/models/holonomic_constraints.py @@ -52,16 +52,16 @@ def superimpose_markers( q_ddot_sym = MX.sym("q_ddot", biorbd_model.nb_qdot, 1) # symbolic markers in global frame - marker_1_sym = biorbd_model.marker(q_sym, index=biorbd_model.marker_index(marker_1)) + marker_1_sym = biorbd_model.marker(index=biorbd_model.marker_index(marker_1))(q_sym) if marker_2 is not None: - marker_2_sym = biorbd_model.marker(q_sym, index=biorbd_model.marker_index(marker_2)) + marker_2_sym = biorbd_model.marker(index=biorbd_model.marker_index(marker_2))(q_sym) else: marker_2_sym = MX([0, 0, 0]) # if local frame is provided, the markers are expressed in the same local frame if local_frame_index is not None: - jcs_t = biorbd_model.homogeneous_matrices_in_global(q_sym, local_frame_index, inverse=True) + jcs_t = biorbd_model.homogeneous_matrices_in_global(segment_idx=local_frame_index, inverse=True)(q_sym) marker_1_sym = (jcs_t @ vertcat(marker_1_sym, 1))[:3] marker_2_sym = (jcs_t @ vertcat(marker_2_sym, 1))[:3] diff --git a/tests/shard1/test__global_plots.py b/tests/shard1/test__global_plots.py index 12f9601fe..6845a8c54 100644 --- a/tests/shard1/test__global_plots.py +++ b/tests/shard1/test__global_plots.py @@ -117,7 +117,7 @@ def test_plot_merged_graphs(phase_dynamics): # Generate random data to fit np.random.seed(42) - t, markers_ref, x_ref, muscle_excitations_ref = ocp_module.generate_data(bio_model, final_time, n_shooting) + t, markers_ref, x_ref, muscle_excitations_ref = ocp_module.generate_data(bio_model, final_time, n_shooting, use_sx=False) bio_model = BiorbdModel(model_path) # To prevent from free variable, the model must be reloaded ocp = ocp_module.prepare_ocp( diff --git a/tests/shard1/test_biorbd_model_size.py b/tests/shard1/test_biorbd_model_size.py index 6ee312a99..1523ee2c2 100644 --- a/tests/shard1/test_biorbd_model_size.py +++ b/tests/shard1/test_biorbd_model_size.py @@ -10,6 +10,7 @@ def model(): return +# Pariterre: Can I remove this file all together? def generate_q_vectors(model): q_valid = MX([0.1] * model.nb_q) @@ -66,7 +67,7 @@ def test_center_of_mass_valid_and_too_large_q_input(model): q_valid, q_too_large = generate_q_vectors(model) # q valid - model.center_of_mass(q_valid) + model.center_of_mass()(q_valid) # q not valid with pytest.raises(ValueError, match="Length of q size should be: 4, but got: 5"): model.center_of_mass(q_too_large) @@ -80,13 +81,13 @@ def test_center_of_mass_velocity_valid_and_too_large_q_or_qdot_input(model): q_valid, qdot_valid, q_too_large, qdot_too_large = generate_q_and_qdot_vectors(model) # q and qdot valid - model.center_of_mass_velocity(q_valid, qdot_valid) + model.center_of_mass_velocity()(q_valid, qdot_valid) # qdot valid but q not valid with pytest.raises(ValueError, match="Length of q size should be: 4, but got: 5"): - model.center_of_mass_velocity(q_too_large, qdot_valid) + model.center_of_mass_velocity()(q_too_large, qdot_valid) # q valid but qdot not valid with pytest.raises(ValueError, match="Length of qdot size should be: 4, but got: 5"): - model.center_of_mass_velocity(q_valid, qdot_too_large) + model.center_of_mass_velocity()(q_valid, qdot_too_large) def test_center_of_mass_acceleration_valid_and_too_large_q_or_qdot_or_qddot_input( @@ -106,16 +107,16 @@ def test_center_of_mass_acceleration_valid_and_too_large_q_or_qdot_or_qddot_inpu ) = generate_q_qdot_qddot_vectors(model) # q, qdot and qddot valid - model.center_of_mass_acceleration(q_valid, qdot_valid, qddot_valid) + model.center_of_mass_acceleration()(q_valid, qdot_valid, qddot_valid) # qdot and qddot valid but q not valid with pytest.raises(ValueError, match="Length of q size should be: 4, but got: 5"): - model.center_of_mass_acceleration(q_too_large, qdot_valid, qddot_valid) + model.center_of_mass_acceleration()(q_too_large, qdot_valid, qddot_valid) # q and qddot valid but qdot not valid with pytest.raises(ValueError, match="Length of qdot size should be: 4, but got: 5"): - model.center_of_mass_acceleration(q_valid, qdot_too_large, qddot_valid) + model.center_of_mass_acceleration()(q_valid, qdot_too_large, qddot_valid) # q and qdot valid but qddot not valid with pytest.raises(ValueError, match="Length of qddot size should be: 4, but got: 5"): - model.center_of_mass_acceleration(q_valid, qdot_valid, qddot_too_large) + model.center_of_mass_acceleration()(q_valid, qdot_valid, qddot_too_large) def test_body_rotation_rate_valid_and_too_large_q_or_qdot_input(model): @@ -126,13 +127,13 @@ def test_body_rotation_rate_valid_and_too_large_q_or_qdot_input(model): q_valid, qdot_valid, q_too_large, qdot_too_large = generate_q_and_qdot_vectors(model) # q and qdot valid - model.body_rotation_rate(q_valid, qdot_valid) + model.body_rotation_rate()(q_valid, qdot_valid) # qdot valid but q not valid with pytest.raises(ValueError, match="Length of q size should be: 4, but got: 5"): - model.body_rotation_rate(q_too_large, qdot_valid) + model.body_rotation_rate()(q_too_large, qdot_valid) # q valid but qdot not valid with pytest.raises(ValueError, match="Length of qdot size should be: 4, but got: 5"): - model.body_rotation_rate(q_valid, qdot_too_large) + model.body_rotation_rate()(q_valid, qdot_too_large) def test_mass_matrix_valid_and_too_large_q_input(model): @@ -143,10 +144,10 @@ def test_mass_matrix_valid_and_too_large_q_input(model): q_valid, q_too_large = generate_q_vectors(model) # q valid - model.mass_matrix(q_valid) + model.mass_matrix()(q_valid) # q not valid with pytest.raises(ValueError, match="Length of q size should be: 4, but got: 5"): - model.mass_matrix(q_too_large) + model.mass_matrix()(q_too_large) def test_non_linear_effects_valid_and_too_large_q_or_qdot_input(model): @@ -157,13 +158,13 @@ def test_non_linear_effects_valid_and_too_large_q_or_qdot_input(model): q_valid, qdot_valid, q_too_large, qdot_too_large = generate_q_and_qdot_vectors(model) # q and qdot valid - model.non_linear_effects(q_valid, qdot_valid) + model.non_linear_effects()(q_valid, qdot_valid) # qdot valid but q not valid with pytest.raises(ValueError, match="Length of q size should be: 4, but got: 5"): - model.non_linear_effects(q_too_large, qdot_valid) + model.non_linear_effects()(q_too_large, qdot_valid) # q valid but qdot not valid with pytest.raises(ValueError, match="Length of qdot size should be: 4, but got: 5"): - model.non_linear_effects(q_valid, qdot_too_large) + model.non_linear_effects()(q_valid, qdot_too_large) def test_angular_momentum_valid_and_too_large_q_or_qdot_input(model): @@ -174,13 +175,13 @@ def test_angular_momentum_valid_and_too_large_q_or_qdot_input(model): q_valid, qdot_valid, q_too_large, qdot_too_large = generate_q_and_qdot_vectors(model) # q and qdot valid - model.angular_momentum(q_valid, qdot_valid) + model.angular_momentum()(q_valid, qdot_valid) # qdot valid but q not valid with pytest.raises(ValueError, match="Length of q size should be: 4, but got: 5"): - model.angular_momentum(q_too_large, qdot_valid) + model.angular_momentum()(q_too_large, qdot_valid) # q valid but qdot not valid with pytest.raises(ValueError, match="Length of qdot size should be: 4, but got: 5"): - model.angular_momentum(q_valid, qdot_too_large) + model.angular_momentum()(q_valid, qdot_too_large) def test_reshape_qdot_valid_and_too_large_q_or_qdot_input(model): @@ -191,13 +192,13 @@ def test_reshape_qdot_valid_and_too_large_q_or_qdot_input(model): q_valid, qdot_valid, q_too_large, qdot_too_large = generate_q_and_qdot_vectors(model) # q and qdot valid - model.reshape_qdot(q_valid, qdot_valid) + model.reshape_qdot()(q_valid, qdot_valid) # qdot valid but q not valid with pytest.raises(ValueError, match="Length of q size should be: 4, but got: 5"): - model.reshape_qdot(q_too_large, qdot_valid) + model.reshape_qdot()(q_too_large, qdot_valid) # q valid but qdot not valid with pytest.raises(ValueError, match="Length of qdot size should be: 4, but got: 5"): - model.reshape_qdot(q_valid, qdot_too_large) + model.reshape_qdot()(q_valid, qdot_too_large) def test_segment_angular_velocity_valid_and_too_large_q_or_qdot_input(model): @@ -208,13 +209,13 @@ def test_segment_angular_velocity_valid_and_too_large_q_or_qdot_input(model): q_valid, qdot_valid, q_too_large, qdot_too_large = generate_q_and_qdot_vectors(model) idx = 1 # q and qdot valid - model.segment_angular_velocity(q_valid, qdot_valid, idx) + model.segment_angular_velocity(idx)(q_valid, qdot_valid) # qdot valid but q not valid with pytest.raises(ValueError, match="Length of q size should be: 4, but got: 5"): - model.segment_angular_velocity(q_too_large, qdot_valid, idx) + model.segment_angular_velocity(idx)(q_too_large, qdot_valid) # q valid but qdot not valid with pytest.raises(ValueError, match="Length of qdot size should be: 4, but got: 5"): - model.segment_angular_velocity(q_valid, qdot_too_large, idx) + model.segment_angular_velocity(idx)(q_valid, qdot_too_large) def test_forward_dynamics_free_floating_base_valid_and_too_large_q_or_qdot_or_qddot_joints_input( @@ -234,16 +235,16 @@ def test_forward_dynamics_free_floating_base_valid_and_too_large_q_or_qdot_or_qd ) = generate_q_qdot_qddot_vectors(model, root_dynamics=True) # q, qdot and qddot_joints valid - model.forward_dynamics_free_floating_base(q_valid, qdot_valid, qddot_joints_valid) + model.forward_dynamics_free_floating_base()(q_valid, qdot_valid, qddot_joints_valid) # qdot and qddot_joints valid but q not valid with pytest.raises(ValueError, match="Length of q size should be: 4, but got: 5"): - model.forward_dynamics_free_floating_base(q_too_large, qdot_valid, qddot_joints_valid) + model.forward_dynamics_free_floating_base()(q_too_large, qdot_valid, qddot_joints_valid) # q and qddot_joints valid but qdot not valid with pytest.raises(ValueError, match="Length of qdot size should be: 4, but got: 5"): - model.forward_dynamics_free_floating_base(q_valid, qdot_too_large, qddot_joints_valid) + model.forward_dynamics_free_floating_base()(q_valid, qdot_too_large, qddot_joints_valid) # q and qdot valid but qddot_joints not valid with pytest.raises(ValueError, match="Length of qddot_joints size should be: 1, but got: 5"): - model.forward_dynamics_free_floating_base(q_valid, qdot_valid, qddot_joints_too_large) + model.forward_dynamics_free_floating_base()(q_valid, qdot_valid, qddot_joints_too_large) def test_forward_dynamics_valid_and_too_large_q_or_qdot_or_tau_activations_input(model): @@ -255,16 +256,16 @@ def test_forward_dynamics_valid_and_too_large_q_or_qdot_or_tau_activations_input tau_valid, tau_too_large = generate_tau_activations_vectors(model) # q, qdot and tau valid - model.forward_dynamics(q_valid, qdot_valid, tau_valid) + model.forward_dynamics()(q_valid, qdot_valid, tau_valid, []) # qdot and tau valid but q not valid with pytest.raises(ValueError, match="Length of q size should be: 4, but got: 5"): - model.forward_dynamics(q_too_large, qdot_valid, tau_valid) + model.forward_dynamics()(q_too_large, qdot_valid, tau_valid, []) # q and tau valid but qdot not valid with pytest.raises(ValueError, match="Length of qdot size should be: 4, but got: 5"): - model.forward_dynamics(q_valid, qdot_too_large, tau_valid) + model.forward_dynamics()(q_valid, qdot_too_large, tau_valid, []) # q and qdot valid but tau not valid with pytest.raises(ValueError, match="Length of tau size should be: 4, but got: 5"): - model.forward_dynamics(q_valid, qdot_valid, tau_too_large) + model.forward_dynamics()(q_valid, qdot_valid, tau_too_large, []) def test_constrained_forward_dynamics_valid_and_too_large_q_or_qdot_or_tau_input(model): diff --git a/tests/shard1/test_biorbd_multi_model.py b/tests/shard1/test_biorbd_multi_model.py index 815d69fee..58bad292f 100644 --- a/tests/shard1/test_biorbd_multi_model.py +++ b/tests/shard1/test_biorbd_multi_model.py @@ -95,7 +95,7 @@ def test_biorbd_model(): models.serialize() models.set_gravity(np.array([0, 0, -3])) - TestUtils.assert_equal(models.gravity, np.array([0, 0, -3, 0, 0, -3])) + TestUtils.assert_equal(models.gravity()['o0'], np.array([0, 0, -3, 0, 0, -3]).reshape(6, 1)) models.set_gravity(np.array([0, 0, -9.81])) with pytest.raises(NotImplementedError, match="segment_index is not implemented for MultiBiorbdModel"): @@ -106,8 +106,7 @@ def test_biorbd_model(): 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.biorbd_homogeneous_matrices_in_global(np.array([1, 2, 3]), 0, 0), + models.homogeneous_matrices_in_global(0, 0)(np.array([1, 2, 3])), np.array( [ [1.0, 0.0, 0.0, 0.0], @@ -118,29 +117,29 @@ def test_biorbd_model(): ), ) TestUtils.assert_equal( - models.homogeneous_matrices_in_child(4), + models.homogeneous_matrices_in_child(4)()["o0"], 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]]), ) - TestUtils.assert_equal(models.mass, np.array([3, 3])) + TestUtils.assert_equal(models.mass()["o0"], np.array([3, 3]).reshape(2, 1)) 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]), + 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]).reshape(6, 1), ) 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]), + 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]).reshape(6, 1), ) TestUtils.assert_equal( - models.center_of_mass_acceleration( + 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]), + np.array([0.0, 0.481652, 6.105858, 0.0, -33.566971, -126.795179]).reshape(6, 1), ) - mass_matrices = models.mass_matrix(np.array([1, 2.1, 3, 4.1, 5, 6.1])) + 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], @@ -157,58 +156,58 @@ def test_biorbd_model(): np.array([[9.313616, 5.580191, 2.063886], [5.580191, 4.791997, 1.895999], [2.063886, 1.895999, 0.945231]]), ) - 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])) + 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]), + np.array([38.817401, -1.960653, -1.259441]).reshape(3, 1) ) TestUtils.assert_equal( nonlinear_effects[1], - np.array([322.060726, -22.147881, -20.660836]), + np.array([322.060726, -22.147881, -20.660836]).reshape(3, 1) ) 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]), + 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]).reshape(6, 1), decimal=5, ) 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]), + models.reshape_qdot(1)(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]).reshape(6, 1), ) 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]), + models.segment_angular_velocity(1)(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]).reshape(6, 1), ) assert models.soft_contact(0, 0) == [] # TODO: Fix soft contact (biorbd call error) - 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) + # 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( + 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([3.1, 0.0, 9.1, 0.0]), ), - np.array([-14.750679, -36.596107]), + np.array([-14.750679, -40.031762]).reshape(2, 1) ) TestUtils.assert_equal( - models.forward_dynamics( + models.forward_dynamics(with_contact=False)( 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]), + np.array([-16.092093, 9.049853, 10.570878, -121.783105, 154.820759, -20.664452]).reshape(6, 1) ) TestUtils.assert_equal( @@ -217,10 +216,11 @@ def test_biorbd_model(): 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]), + np.array([-16.092093, 9.049853, 10.570878, -121.783105, 154.820759, -20.664452]).reshape(6, 1) ) - with pytest.raises(NotImplementedError, match="External forces are not implemented yet for MultiBiorbdModel."): + with pytest.raises(RuntimeError, match="Incorrect number of inputs: Expected 3, got 4"): + # Because external_forces are not implemented models.forward_dynamics(with_contact=True)( np.array([1, 2.1, 3, 4.1, 5, 6.1]), np.array([1, 2.1, 3, 4.1, 5, 6]), @@ -229,27 +229,27 @@ def test_biorbd_model(): ) TestUtils.assert_equal( - models.inverse_dynamics( + 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]), + np.array([4.844876e01, 2.121037e-01, 1.964626e00, 4.165226e02, 3.721585e01, 1.906986e00]).reshape(6, 1), decimal=5, ) TestUtils.assert_equal( - models.contact_forces_from_constrained_forward_dynamics( + 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]), + np.array([0.0, 0.0]).reshape(2, 1) ) - with pytest.raises(NotImplementedError, match="External forces are not implemented yet for MultiBiorbdModel."): - models.contact_forces_from_constrained_forward_dynamics( + with pytest.raises(RuntimeError, match="Incorrect number of inputs: Expected 3, got 4"): + # Because external_forces are not implemented + 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]), @@ -257,63 +257,61 @@ def test_biorbd_model(): ) TestUtils.assert_equal( - models.qdot_from_impact( + 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]), + np.array([1.0, 2.1, 3.0, 4.1, 5.0, 6.0]).reshape(6, 1) ) TestUtils.assert_equal( - models.muscle_activation_dot( - np.array([1, 2.1, 3, 4.1, 5, 6.1]), + models.muscle_activation_dot()( + [], # There is no muscle in the models ), - np.array([], dtype=np.float64), + np.array([], dtype=np.float64).reshape(0, 1), ) TestUtils.assert_equal( - models.muscle_joint_torque( - np.array([1, 2.1, 3, 4.1, 5, 6.1]), + models.muscle_joint_torque()( + np.array([]), # There is no muscle in the models 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]), + np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0]).reshape(6, 1), ) TestUtils.assert_equal( - models.markers( + models.markers()( np.array([1, 2.1, 3, 4.1, 5, 6.1]), - )[0], - np.array([0.0, 0.0, 0.0]), + )[:, 0], + np.array([0.0, 0.0, 0.0]).reshape(3, 1), ) TestUtils.assert_equal( - models.marker( + models.marker(index=1)( np.array([1, 2.1, 3, 4.1, 5, 6.1]), - index=1, ), - np.array([0.0, 0.841471, -0.540302]), + np.array([0.0, 0.841471, -0.540302]).reshape(3, 1) ) assert models.marker_index("marker_3") == 2 - markers_velocities = models.marker_velocities( + 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) TestUtils.assert_equal( - markers_velocities[0], - np.array([0.0, 0.0, 0.0]), + markers_velocities[:, 0], + np.array([0.0, 0.0, 0.0]).reshape(3, 1), ) TestUtils.assert_equal( - markers_velocities[1], - np.array([0.0, 0.540302, 0.841471]), + markers_velocities[:, 1], + np.array([0.0, 0.540302, 0.841471]).reshape(3, 1), ) with pytest.raises(RuntimeError, match="All dof must have their actuators set"): - models.tau_max( + 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) @@ -332,32 +330,32 @@ def test_biorbd_model(): 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])) + 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]), + models.normalize_state_quaternions()( + np.array([1, 2.1, 3, 4.1, 5, 6.1]), ), - 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]), + np.array([1.0, 2.1, 3.0, 4.1, 5.0, 6.1]).reshape(6, 1), ) TestUtils.assert_equal( - models.contact_forces( + 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.array([0.0, 0.0]).reshape(2, 1) ) + # Because external_forces are not implemented yet with pytest.raises( - NotImplementedError, match="contact_forces is not implemented yet with external_forces for MultiBiorbdModel" + RuntimeError, match="Incorrect number of inputs: Expected 3, got 4" ): - models.contact_forces( + 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]), @@ -365,11 +363,11 @@ def test_biorbd_model(): ) TestUtils.assert_equal( - models.passive_joint_torque( + 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]), + np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0]).reshape(6, 1), ) q_mapping = models._var_mapping("q", None, variable_mappings)