From e7fec490058d90e19c9409980de3d2f80af2fa03 Mon Sep 17 00:00:00 2001 From: Charbie Date: Fri, 21 Jun 2024 06:35:13 -0400 Subject: [PATCH] initial thoughts --- bioptim/dynamics/configure_problem.py | 76 +++++++++---------- bioptim/dynamics/dynamics_functions.py | 20 +++-- bioptim/models/biorbd/biorbd_model.py | 42 ++++++---- bioptim/optimization/optimization_variable.py | 12 +-- bioptim/optimization/parameters.py | 4 +- 5 files changed, 85 insertions(+), 69 deletions(-) diff --git a/bioptim/dynamics/configure_problem.py b/bioptim/dynamics/configure_problem.py index 7cb60d5df..142688061 100644 --- a/bioptim/dynamics/configure_problem.py +++ b/bioptim/dynamics/configure_problem.py @@ -294,21 +294,21 @@ def torque_driven( external_forces=external_forces, ) - # Configure the contact forces - if with_contact: - ConfigureProblem.configure_contact_function( - ocp, nlp, DynamicsFunctions.forces_from_torque_driven, external_forces=external_forces - ) - # Configure the soft contact forces - ConfigureProblem.configure_soft_contact_function(ocp, nlp) - # Algebraic constraints of soft contact forces if needed - if soft_contacts_dynamics == SoftContactDynamics.CONSTRAINT: - ocp.implicit_constraints.add( - ImplicitConstraintFcn.SOFT_CONTACTS_EQUALS_SOFT_CONTACTS_DYNAMICS, - node=Node.ALL_SHOOTING, - penalty_type=ConstraintType.IMPLICIT, - phase=nlp.phase_idx, - ) + # # Configure the contact forces + # if with_contact: + # ConfigureProblem.configure_contact_function( + # ocp, nlp, DynamicsFunctions.forces_from_torque_driven, external_forces=external_forces + # ) + # # Configure the soft contact forces + # ConfigureProblem.configure_soft_contact_function(ocp, nlp) + # # Algebraic constraints of soft contact forces if needed + # if soft_contacts_dynamics == SoftContactDynamics.CONSTRAINT: + # ocp.implicit_constraints.add( + # ImplicitConstraintFcn.SOFT_CONTACTS_EQUALS_SOFT_CONTACTS_DYNAMICS, + # node=Node.ALL_SHOOTING, + # penalty_type=ConstraintType.IMPLICIT, + # phase=nlp.phase_idx, + # ) @staticmethod def torque_driven_free_floating_base( @@ -963,12 +963,12 @@ def configure_dynamics_function(ocp, nlp, dyn_func, **extra_params): DynamicsFunctions.apply_parameters(nlp) dynamics_eval = dyn_func( - nlp.time_mx, - nlp.states.scaled.mx_reduced, - nlp.controls.scaled.mx_reduced, - nlp.parameters.scaled.mx_reduced, - nlp.algebraic_states.scaled.mx_reduced, - nlp.numerical_timeseries.mx, + nlp.time_cx, + nlp.states.scaled.cx, + nlp.controls.scaled.cx, + nlp.parameters.scaled.cx, + nlp.algebraic_states.scaled.cx, + nlp.numerical_timeseries.cx, nlp, **extra_params, ) @@ -976,17 +976,17 @@ def configure_dynamics_function(ocp, nlp, dyn_func, **extra_params): if isinstance(dynamics_dxdt, (list, tuple)): dynamics_dxdt = vertcat(*dynamics_dxdt) - time_span_sym = vertcat(nlp.time_mx, nlp.dt_mx) + time_span_sym = vertcat(nlp.time_cx, nlp.dt) if nlp.dynamics_func is None: nlp.dynamics_func = Function( "ForwardDyn", [ time_span_sym, - nlp.states.scaled.mx_reduced, - nlp.controls.scaled.mx_reduced, - nlp.parameters.scaled.mx_reduced, - nlp.algebraic_states.scaled.mx_reduced, - nlp.numerical_timeseries.mx, + nlp.states.scaled.cx, + nlp.controls.scaled.cx, + nlp.parameters.scaled.cx, + nlp.algebraic_states.scaled.cx, + nlp.numerical_timeseries.cx, ], [dynamics_dxdt], ["t_span", "x", "u", "p", "a", "d"], @@ -1011,12 +1011,12 @@ def configure_dynamics_function(ocp, nlp, dyn_func, **extra_params): "DynamicsDefects", [ time_span_sym, - nlp.states.scaled.mx_reduced, - nlp.controls.scaled.mx_reduced, - nlp.parameters.scaled.mx_reduced, - nlp.algebraic_states.scaled.mx_reduced, - nlp.numerical_timeseries.mx, - nlp.states_dot.scaled.mx_reduced, + nlp.states.scaled.cx, + nlp.controls.scaled.cx, + nlp.parameters.scaled.cx, + nlp.algebraic_states.scaled.cx, + nlp.numerical_timeseries.cx, + nlp.states_dot.scaled.cx, ], [dynamics_eval.defects], ["t_span", "x", "u", "p", "a", "d", "xdot"], @@ -1040,11 +1040,11 @@ def configure_dynamics_function(ocp, nlp, dyn_func, **extra_params): "ForwardDyn", [ time_span_sym, - nlp.states.scaled.mx_reduced, - nlp.controls.scaled.mx_reduced, - nlp.parameters.scaled.mx_reduced, - nlp.algebraic_states.scaled.mx_reduced, - nlp.numerical_timeseries.mx, + nlp.states.scaled.cx, + nlp.controls.scaled.cx, + nlp.parameters.scaled.cx, + nlp.algebraic_states.scaled.cx, + nlp.numerical_timeseries.cx, ], [dynamics_dxdt], ["t_span", "x", "u", "p", "a", "d"], diff --git a/bioptim/dynamics/dynamics_functions.py b/bioptim/dynamics/dynamics_functions.py index 6b17fadad..b58f3b91f 100644 --- a/bioptim/dynamics/dynamics_functions.py +++ b/bioptim/dynamics/dynamics_functions.py @@ -146,10 +146,13 @@ def torque_driven( dq = DynamicsFunctions.compute_qdot(nlp, q, qdot) - tau = DynamicsFunctions.__get_fatigable_tau(nlp, states, controls, fatigue) - tau = tau + nlp.model.passive_joint_torque(q, qdot) if with_passive_torque else tau - tau = tau + nlp.model.ligament_joint_torque(q, qdot) if with_ligament else tau - tau = tau - nlp.model.friction_coefficients @ qdot if with_friction else tau + # Charbie: to remove + tau = DynamicsFunctions.get(nlp.controls["tau"], controls) + # + # tau = DynamicsFunctions.__get_fatigable_tau(nlp, states, controls, fatigue) + # tau = tau + nlp.model.passive_joint_torque(q, qdot) if with_passive_torque else tau + # tau = tau + nlp.model.ligament_joint_torque(q, qdot) if with_ligament else tau + # tau = tau - nlp.model.friction_coefficients @ qdot if with_friction else tau if ( rigidbody_dynamics == RigidBodyDynamics.DAE_INVERSE_DYNAMICS @@ -169,7 +172,7 @@ def torque_driven( dxdt[nlp.states["qddot"].index, :] = DynamicsFunctions.get(nlp.controls["qdddot"], controls) else: ddq = DynamicsFunctions.forward_dynamics(nlp, q, qdot, tau, with_contact, external_forces) - dxdt = MX(nlp.states.shape, ddq.shape[1]) + dxdt = nlp.cx(nlp.states.shape, ddq.shape[1]) dxdt[nlp.states["q"].index, :] = horzcat(*[dq for _ in range(ddq.shape[1])]) dxdt[nlp.states["qdot"].index, :] = ddq @@ -1090,7 +1093,7 @@ def compute_qdot(nlp, q: MX | SX, qdot: MX | SX): mapping = nlp.controls["q"].mapping else: raise RuntimeError("Your q key combination was not found in states or controls") - return mapping.to_first.map(nlp.model.reshape_qdot(q, qdot)) + return mapping.to_first.map(nlp.model.reshape_qdot()(q, qdot)) @staticmethod def forward_dynamics( @@ -1122,6 +1125,7 @@ def forward_dynamics( ------- The derivative of qdot """ + # Get the mapping of the output if "qdot" in nlp.states: qdot_var_mapping = nlp.states["qdot"].mapping.to_first elif "qdot" in nlp.controls: @@ -1133,14 +1137,14 @@ def forward_dynamics( if with_contact: qddot = nlp.model.constrained_forward_dynamics(q, qdot, tau) else: - qddot = nlp.model.forward_dynamics(q, qdot, tau) + qddot = nlp.model.forward_dynamics()(q, qdot, tau) return qdot_var_mapping.map(qddot) else: if with_contact: qddot = nlp.model.constrained_forward_dynamics(q, qdot, tau, external_forces) else: - qddot = nlp.model.forward_dynamics(q, qdot, tau, external_forces) + qddot = nlp.model.forward_dynamics()(q, qdot, tau, external_forces) return qdot_var_mapping.map(qddot) @staticmethod diff --git a/bioptim/models/biorbd/biorbd_model.py b/bioptim/models/biorbd/biorbd_model.py index 3586d3c14..efd77a85e 100644 --- a/bioptim/models/biorbd/biorbd_model.py +++ b/bioptim/models/biorbd/biorbd_model.py @@ -6,7 +6,7 @@ GeneralizedTorque, GeneralizedAcceleration, ) -from casadi import SX, MX, vertcat, horzcat, norm_fro +from casadi import SX, MX, vertcat, horzcat, norm_fro, Function from typing import Callable, Any from ..utils import _var_mapping, bounds_from_ranges @@ -208,14 +208,20 @@ def angular_momentum(self, q, qdot) -> MX: qdot_biorbd = GeneralizedVelocity(qdot) return self.model.angularMomentum(q_biorbd, qdot_biorbd, True).to_mx() - def reshape_qdot(self, q, qdot, k_stab=1) -> MX: - self.check_q_size(q) - self.check_qdot_size(qdot) - return self.model.computeQdot( - GeneralizedCoordinates(q), - GeneralizedCoordinates(qdot), # mistake in biorbd + def reshape_qdot(self, k_stab=1) -> MX: + q_mx_reduced = MX.sym("q_mx_reduced", self.nb_q, 1) + qdot_mx_reduced = MX.sym("qdot_mx_reduced", self.nb_qdot, 1) + biorbd_return = self.model.computeQdot( + GeneralizedCoordinates(q_mx_reduced), + GeneralizedCoordinates(qdot_mx_reduced), # mistake in biorbd k_stab, ).to_mx() + casadi_fun = Function( + "reshape_qdot", + [q_mx_reduced, qdot_mx_reduced], + [biorbd_return], + ) + return casadi_fun def segment_angular_velocity(self, q, qdot, idx) -> MX: """ @@ -345,16 +351,22 @@ def _dispatch_forces(self, external_forces: MX, translational_forces: MX): return external_forces_set - def forward_dynamics(self, q, qdot, tau, external_forces=None, translational_forces=None) -> MX: - self.check_q_size(q) - self.check_qdot_size(qdot) - self.check_tau_size(tau) + def forward_dynamics(self, external_forces=None, translational_forces=None) -> MX: + q_mx_reduced = MX.sym("q_mx_reduced", self.nb_q, 1) + qdot_mx_reduced = MX.sym("qdot_mx_reduced", self.nb_qdot, 1) + tau_mx_reduced = MX.sym("tau_mx_reduced", self.nb_tau, 1) external_forces_set = self._dispatch_forces(external_forces, translational_forces) - q_biorbd = GeneralizedCoordinates(q) - qdot_biorbd = GeneralizedVelocity(qdot) - tau_biorbd = GeneralizedTorque(tau) - return self.model.ForwardDynamics(q_biorbd, qdot_biorbd, tau_biorbd, external_forces_set).to_mx() + q_biorbd = GeneralizedCoordinates(q_mx_reduced) + qdot_biorbd = GeneralizedVelocity(qdot_mx_reduced) + tau_biorbd = GeneralizedTorque(tau_mx_reduced) + biorbd_return = self.model.ForwardDynamics(q_biorbd, qdot_biorbd, tau_biorbd, external_forces_set).to_mx() + casadi_fun = Function( + "forward_dynamics", + [q_mx_reduced, qdot_mx_reduced, tau_mx_reduced], + [biorbd_return], + ) + return casadi_fun def constrained_forward_dynamics(self, q, qdot, tau, external_forces=None, translational_forces=None) -> MX: external_forces_set = self._dispatch_forces(external_forces, translational_forces) diff --git a/bioptim/optimization/optimization_variable.py b/bioptim/optimization/optimization_variable.py index a9a7a8578..9b5f812c2 100644 --- a/bioptim/optimization/optimization_variable.py +++ b/bioptim/optimization/optimization_variable.py @@ -189,7 +189,7 @@ def __init__(self, cx_constructor, phase_dynamics): self._cx_mid: MX | SX | np.ndarray = np.array([]) self._cx_end: MX | SX | np.ndarray = np.array([]) self._cx_intermediates: list = [] - self.mx_reduced: MX = MX.sym("var", 0, 0) + # self.mx_reduced: MX = MX.sym("var", 0, 0) self.cx_constructor = cx_constructor self._current_cx_to_get = 0 self.phase_dynamics = phase_dynamics @@ -310,7 +310,7 @@ def append(self, name: str, cx: list, mx: MX, bimapping: BiMapping): else: self._cx_intermediates[i] = vertcat(self._cx_intermediates[i], c) - self.mx_reduced = vertcat(self.mx_reduced, MX.sym("var", cx[0].shape[0])) + # self.mx_reduced = vertcat(self.mx_reduced, MX.sym("var", cx[0].shape[0])) self.elements.append(OptimizationVariable(name, mx, cx, index, bimapping, parent_list=self)) def append_from_scaled( @@ -346,7 +346,7 @@ def append_from_scaled( else: self._cx_intermediates[i] = vertcat(self._cx_intermediates[i], c) - self.mx_reduced = scaled_optimization_variable.mx_reduced + # self.mx_reduced = scaled_optimization_variable.mx_reduced var = scaled_optimization_variable[name] self.elements.append(OptimizationVariable(name, var.mx, cx, var.index, var.mapping, self)) @@ -564,9 +564,9 @@ def shape(self): def mx(self): return self.unscaled.mx - @property - def mx_reduced(self): - return self.unscaled.mx_reduced + # @property + # def mx_reduced(self): + # return self.unscaled.mx_reduced @property def cx(self): diff --git a/bioptim/optimization/parameters.py b/bioptim/optimization/parameters.py index f6df2e3a7..c8bf19254 100644 --- a/bioptim/optimization/parameters.py +++ b/bioptim/optimization/parameters.py @@ -170,7 +170,7 @@ def add( self.scaling.add(key=name, scaling=scaling) index = range(self._cx_start.shape[0], self._cx_start.shape[0] + cx[0].shape[0]) self._cx_start = vertcat(self._cx_start, cx[0]) - self.mx_reduced = vertcat(self.mx_reduced, MX.sym("var", cx[0].shape[0])) + # self.mx_reduced = vertcat(self.mx_reduced, MX.sym("var", cx[0].shape[0])) mx = MX.sym(name, size) self.elements.append( Parameter( @@ -224,7 +224,7 @@ def to_unscaled( unscaled_parameter._cx_start = vertcat( unscaled_parameter._cx_start, element.cx_start * element.scaling.scaling ) - unscaled_parameter.mx_reduced = vertcat(unscaled_parameter.mx_reduced, element.mx * element.scaling.scaling) + # unscaled_parameter.mx_reduced = vertcat(unscaled_parameter.cx, element.cx * element.scaling.scaling) return unscaled_parameter