Skip to content

Commit

Permalink
blacked
Browse files Browse the repository at this point in the history
  • Loading branch information
EveCharbie committed Sep 19, 2024
1 parent 40e2ae8 commit 3881cca
Show file tree
Hide file tree
Showing 4 changed files with 42 additions and 32 deletions.
28 changes: 16 additions & 12 deletions bioptim/limits/penalty.py
Original file line number Diff line number Diff line change
Expand Up @@ -341,7 +341,9 @@ def minimize_markers_velocity(
penalty.quadratic = True if penalty.quadratic is None else penalty.quadratic

# Add the penalty in the requested reference frame. None for global
markers = horzcat(*controller.model.marker_velocities(reference_index=reference_jcs)(controller.q.cx, controller.qdot.cx))
markers = horzcat(
*controller.model.marker_velocities(reference_index=reference_jcs)(controller.q.cx, controller.qdot.cx)
)

return markers

Expand Down Expand Up @@ -380,9 +382,9 @@ def minimize_markers_acceleration(
qddot = PenaltyFunctionAbstract._get_qddot(controller, "cx")

markers = horzcat(
*controller.model.marker_accelerations(reference_index=reference_jcs)(controller.q.cx,
controller.qdot.cx,
qddot)
*controller.model.marker_accelerations(reference_index=reference_jcs)(
controller.q.cx, controller.qdot.cx, qddot
)
)

return markers
Expand Down Expand Up @@ -425,9 +427,9 @@ def superimpose_markers(
PenaltyFunctionAbstract.set_axes_rows(penalty, axes)
penalty.quadratic = True if penalty.quadratic is None else penalty.quadratic

diff_markers = controller.model.marker(second_marker_idx)(controller.q.cx) - controller.model.marker(first_marker_idx)(
controller.q.cx
)
diff_markers = controller.model.marker(second_marker_idx)(controller.q.cx) - controller.model.marker(
first_marker_idx
)(controller.q.cx)

return diff_markers

Expand Down Expand Up @@ -597,7 +599,7 @@ def minimize_predicted_com_height(_: PenaltyOption, controller: PenaltyControlle
The penalty node elements
"""

g = controller.model.gravity()['o0'][2]
g = controller.model.gravity()["o0"][2]
com = controller.model.center_of_mass()(controller.q.cx)
com_dot = controller.model.center_of_mass_velocity()(controller.q.cx, controller.qdot.cx)
com_height = (com_dot[2] * com_dot[2]) / (2 * -g) + com[2]
Expand Down Expand Up @@ -672,9 +674,7 @@ def minimize_com_acceleration(penalty: PenaltyOption, controller: PenaltyControl

qddot = PenaltyFunctionAbstract._get_qddot(controller, "cx")

marker = controller.model.center_of_mass_acceleration()(controller.q.cx,
controller.qdot.cx,
qddot)
marker = controller.model.center_of_mass_acceleration()(controller.q.cx, controller.qdot.cx, qddot)

return marker

Expand Down Expand Up @@ -860,7 +860,11 @@ def minimize_soft_contact_forces(

@staticmethod
def track_segment_with_custom_rt(
penalty: PenaltyOption, controller: PenaltyController, segment: int | str, rt_idx: int, sequence: str = "zyx"
penalty: PenaltyOption,
controller: PenaltyController,
segment: int | str,
rt_idx: int,
sequence: str = "zyx",
):
"""
Minimize the difference of the euler angles extracted from the coordinate system of a segment
Expand Down
12 changes: 4 additions & 8 deletions bioptim/models/biorbd/biorbd_model.py
Original file line number Diff line number Diff line change
Expand Up @@ -132,9 +132,7 @@ def rotation_matrix_to_euler_angles(self, sequence) -> Function:
"""
r = MX.sym("r_mx", 3, 3)
# @Pariterre: is this the right order?
r_matrix = biorbd.Rotation(r[0, 0], r[0, 1], r[0, 2],
r[0, 0], r[0, 1], r[0, 2],
r[0, 0], r[0, 1], r[0, 2])
r_matrix = biorbd.Rotation(r[0, 0], r[0, 1], r[0, 2], r[0, 0], r[0, 1], r[0, 2], r[0, 0], r[0, 1], r[0, 2])
biorbd_return = biorbd.Rotation.toEulerAngles(r_matrix, sequence).to_mx()
casadi_fun = Function(
"rotation_matrix_to_euler_angles",
Expand Down Expand Up @@ -409,7 +407,7 @@ def _dispatch_forces(self, external_forces: MX):

return external_forces_set

def forward_dynamics(self, with_contact: bool=False) -> Function:
def forward_dynamics(self, with_contact: bool = False) -> Function:
external_forces_set = self._dispatch_forces(self.external_forces)

q_biorbd = GeneralizedCoordinates(self.q)
Expand All @@ -434,7 +432,7 @@ def forward_dynamics(self, with_contact: bool=False) -> Function:
)
return casadi_fun

def inverse_dynamics(self, with_contact: bool=False) -> Function:
def inverse_dynamics(self, with_contact: bool = False) -> Function:
# @ipuch: I do not understand what is happening here? Do we have f_ext or it is just the contact forces?
if with_contact:
f_ext = self.reshape_fext_to_fcontact(self.external_forces)
Expand Down Expand Up @@ -790,9 +788,7 @@ def contact_forces(self, external_forces: MX = None) -> Function:
[biorbd_return],
)
else:
biorbd_return = self.contact_forces_from_constrained_forward_dynamics()(
self.q, self.qdot, self.tau, MX()
)
biorbd_return = self.contact_forces_from_constrained_forward_dynamics()(self.q, self.qdot, self.tau, MX())
casadi_fun = Function(
"contact_forces",
[self.q, self.qdot, self.tau],
Expand Down
30 changes: 21 additions & 9 deletions bioptim/models/biorbd/multi_biorbd_model.py
Original file line number Diff line number Diff line change
Expand Up @@ -295,9 +295,13 @@ def center_of_mass_velocity(self) -> 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,
model.center_of_mass_velocity()(q_model, qdot_model),
out = (
model.center_of_mass_velocity()(q_model, qdot_model)
if i == 0
else vertcat(
out,
model.center_of_mass_velocity()(q_model, qdot_model),
)
)
return out

Expand All @@ -307,9 +311,13 @@ def center_of_mass_acceleration(self) -> MX:
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,
model.center_of_mass_acceleration()(q_model, qdot_model, qddot_model),
out = (
model.center_of_mass_acceleration()(q_model, qdot_model, qddot_model)
if i == 0
else vertcat(
out,
model.center_of_mass_acceleration()(q_model, qdot_model, qddot_model),
)
)
return out

Expand All @@ -332,9 +340,13 @@ def angular_momentum(self) -> 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,
model.angular_momentum()(q_model, qdot_model),
out = (
model.angular_momentum()(q_model, qdot_model)
if i == 0
else vertcat(
out,
model.angular_momentum()(q_model, qdot_model),
)
)
return out

Expand Down
4 changes: 1 addition & 3 deletions bioptim/models/protocols/biomodel.py
Original file line number Diff line number Diff line change
Expand Up @@ -358,9 +358,7 @@ def lagrangian(self) -> Function:
The Lagrangian.
"""

def partitioned_forward_dynamics(
self, q_u, qdot_u, tau, external_forces=None, q_v_init=None
) -> Function:
def partitioned_forward_dynamics(self, q_u, qdot_u, tau, external_forces=None, q_v_init=None) -> Function:
"""
@ipuch: I need help on how to implement this!
This is the forward dynamics of the model, but only for the independent joints
Expand Down

0 comments on commit 3881cca

Please sign in to comment.