Skip to content

Commit

Permalink
initial thoughts
Browse files Browse the repository at this point in the history
  • Loading branch information
@EveCharbie
EveCharbie committed Jun 21, 2024
1 parent db03d6d commit e7fec49
Show file tree
Hide file tree
Showing 5 changed files with 85 additions and 69 deletions.
76 changes: 38 additions & 38 deletions bioptim/dynamics/configure_problem.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down Expand Up @@ -963,30 +963,30 @@ 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,
)
dynamics_dxdt = dynamics_eval.dxdt
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"],
Expand All @@ -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"],
Expand All @@ -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"],
Expand Down
20 changes: 12 additions & 8 deletions bioptim/dynamics/dynamics_functions.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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

Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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:
Expand All @@ -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
Expand Down
42 changes: 27 additions & 15 deletions bioptim/models/biorbd/biorbd_model.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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:
"""
Expand Down Expand Up @@ -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)
Expand Down
12 changes: 6 additions & 6 deletions bioptim/optimization/optimization_variable.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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))

Expand Down Expand Up @@ -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):
Expand Down
4 changes: 2 additions & 2 deletions bioptim/optimization/parameters.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down Expand Up @@ -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

Expand Down

0 comments on commit e7fec49

Please sign in to comment.