blacked

pyomeca · Sep 20, 2024 · 37b215c · 37b215c
1 parent 224758f
commit 37b215c
Show file tree

Hide file tree

Showing 5 changed files with 96 additions and 48 deletions.
diff --git a/bioptim/dynamics/dynamics_functions.py b/bioptim/dynamics/dynamics_functions.py
@@ -260,7 +260,9 @@ def torque_driven_free_floating_base(
         n_qdot = nlp.model.nb_qdot
 
         tau_joints = (
-            tau_joints + nlp.model.passive_joint_torque()(q_full, qdot_full, nlp.parameters.cx) if with_passive_torque else tau_joints
+            tau_joints + nlp.model.passive_joint_torque()(q_full, qdot_full, nlp.parameters.cx)
+            if with_passive_torque
+            else tau_joints
         )
         tau_joints = tau_joints + nlp.model.ligament_joint_torque()(q_full, qdot_full) if with_ligament else tau_joints
         tau_joints = tau_joints - nlp.model.friction_coefficients @ qdot_joints if with_friction else tau_joints
@@ -1116,7 +1118,9 @@ def forward_dynamics(
             qddot = nlp.model.forward_dynamics(with_contact=with_contact)(q, qdot, tau, [], nlp.parameters.cx)
             return qdot_var_mapping.map(qddot)
         else:
-            qddot = nlp.model.forward_dynamics(with_contact=with_contact)(q, qdot, tau, external_forces, nlp.parameters.cx)
+            qddot = nlp.model.forward_dynamics(with_contact=with_contact)(
+                q, qdot, tau, external_forces, nlp.parameters.cx
+            )
             return qdot_var_mapping.map(qddot)
 
     @staticmethod
@@ -1162,7 +1166,9 @@ def inverse_dynamics(
                 tau_shape = nlp.model.nb_tau
             tau = nlp.cx(tau_shape, nlp.ns)
             for i in range(external_forces.shape[1]):
-                tau[:, i] = nlp.model.inverse_dynamics(with_contact=with_contact)(q, qdot, qddot, external_forces[:, i], nlp.parameters.cx)
+                tau[:, i] = nlp.model.inverse_dynamics(with_contact=with_contact)(
+                    q, qdot, qddot, external_forces[:, i], nlp.parameters.cx
+                )
         return tau  # We ignore on purpose the mapping to keep zeros in the defects of the dynamic.
 
     @staticmethod

diff --git a/bioptim/limits/constraints.py b/bioptim/limits/constraints.py
@@ -411,11 +411,17 @@ def qddot_equals_forward_dynamics(
             passive_torque = controller.model.passive_joint_torque()(q, qdot, controller.parameters.cx)
             tau = controller.states["tau"].cx if "tau" in controller.states else controller.tau.cx
             tau = tau + passive_torque if with_passive_torque else tau
-            tau = tau + controller.model.ligament_joint_torque(q, qdot, controller.parameters.cx) if with_ligament else tau
+            tau = (
+                tau + controller.model.ligament_joint_torque(q, qdot, controller.parameters.cx)
+                if with_ligament
+                else tau
+            )
 
             qddot = controller.controls["qddot"].cx if "qddot" in controller.controls else controller.states["qddot"].cx
             # @Ipuch: no f_ext allowed ?
-            qddot_fd = controller.model.forward_dynamics(with_contact=with_contact)(q, qdot, tau, [], controller.parameters.cx)
+            qddot_fd = controller.model.forward_dynamics(with_contact=with_contact)(
+                q, qdot, tau, [], controller.parameters.cx
+            )
             return qddot - qddot_fd
 
         @staticmethod
@@ -453,7 +459,11 @@ def tau_equals_inverse_dynamics(
             qddot = controller.states["qddot"].cx if "qddot" in controller.states else controller.controls["qddot"].cx
             passive_torque = controller.model.passive_joint_torque()(q, qdot, controller.parameters.cx)
             tau = tau + passive_torque if with_passive_torque else tau
-            tau = tau + controller.model.ligament_joint_torque()(q, qdot, controller.parameters.cx) if with_ligament else tau
+            tau = (
+                tau + controller.model.ligament_joint_torque()(q, qdot, controller.parameters.cx)
+                if with_ligament
+                else tau
+            )
 
             if "fext" in controller.controls:
                 f_ext = controller.controls["fext"].cx
@@ -464,7 +474,9 @@ def tau_equals_inverse_dynamics(
             else:
                 raise ValueError("External forces must be provided")
 
-            tau_id = controller.model.inverse_dynamics(with_contact=with_contact)(q, qdot, qddot, f_ext, controller.parameters.cx)
+            tau_id = controller.model.inverse_dynamics(with_contact=with_contact)(
+                q, qdot, qddot, f_ext, controller.parameters.cx
+            )
 
             var = []
             var.extend([controller.states[key] for key in controller.states])
@@ -534,7 +546,11 @@ def tau_from_muscle_equal_inverse_dynamics(
                 muscles_states[k].setActivation(muscle_activations[k])
             muscle_tau = controller.model.muscle_joint_torque(muscles_states, q, qdot, controller.parameters.cx)
             muscle_tau = muscle_tau + passive_torque if with_passive_torque else muscle_tau
-            muscle_tau = muscle_tau + controller.model.ligament_joint_torque(q, qdot, controller.parameters.cx) if with_ligament else muscle_tau
+            muscle_tau = (
+                muscle_tau + controller.model.ligament_joint_torque(q, qdot, controller.parameters.cx)
+                if with_ligament
+                else muscle_tau
+            )
             qddot = controller.states["qddot"].cx if "qddot" in controller.states else controller.controls["qddot"].cx
 
             if controller.get_nlp.numerical_timeseries:

diff --git a/bioptim/limits/penalty.py b/bioptim/limits/penalty.py
@@ -124,7 +124,9 @@ def minimize_power(
             if key_control == "tau":
                 return controls * controller.qdot.cx_start
             elif key_control == "muscles":
-                muscles_dot = controller.model.muscle_velocity()(controller.q.cx, controller.qdot.cx, controller.parameters.cx)
+                muscles_dot = controller.model.muscle_velocity()(
+                    controller.q.cx, controller.qdot.cx, controller.parameters.cx
+                )
 
                 return controls * muscles_dot
 
@@ -295,7 +297,9 @@ def minimize_markers(
             jcs_t = (
                 CX_eye(4)
                 if reference_jcs is None
-                else controller.model.homogeneous_matrices_in_global(reference_jcs, inverse=True)(q.cx, controller.parameters.cx)
+                else controller.model.homogeneous_matrices_in_global(reference_jcs, inverse=True)(
+                    q.cx, controller.parameters.cx
+                )
             )
 
             markers = []
@@ -426,9 +430,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.parameters.cx) - controller.model.marker(
-                first_marker_idx
-            )(controller.q.cx, controller.parameters.cx)
+            diff_markers = controller.model.marker(second_marker_idx)(
+                controller.q.cx, controller.parameters.cx
+            ) - controller.model.marker(first_marker_idx)(controller.q.cx, controller.parameters.cx)
 
             return diff_markers
 
@@ -470,7 +474,9 @@ def superimpose_markers_velocity(
             PenaltyFunctionAbstract.set_axes_rows(penalty, axes)
             penalty.quadratic = True if penalty.quadratic is None else penalty.quadratic
 
-            marker_velocity = controller.model.marker_velocities()(controller.q.cx, controller.qdot.cx, controller.parameters.cx)
+            marker_velocity = controller.model.marker_velocities()(
+                controller.q.cx, controller.qdot.cx, controller.parameters.cx
+            )
             marker_1 = marker_velocity[first_marker_idx][:]
             marker_2 = marker_velocity[second_marker_idx][:]
 
@@ -600,7 +606,9 @@ def minimize_predicted_com_height(_: PenaltyOption, controller: PenaltyControlle
 
             g = controller.model.gravity()["o0"][2]
             com = controller.model.center_of_mass()(controller.q.cx, controller.parameters.cx)
-            com_dot = controller.model.center_of_mass_velocity()(controller.q.cx, controller.qdot.cx, controller.parameters.cx)
+            com_dot = controller.model.center_of_mass_velocity()(
+                controller.q.cx, controller.qdot.cx, controller.parameters.cx
+            )
             com_height = (com_dot[2] * com_dot[2]) / (2 * -g) + com[2]
             return com_height
 
@@ -648,7 +656,9 @@ def minimize_com_velocity(penalty: PenaltyOption, controller: PenaltyController,
             PenaltyFunctionAbstract.set_axes_rows(penalty, axes)
             penalty.quadratic = True if penalty.quadratic is None else penalty.quadratic
 
-            return controller.model.center_of_mass_velocity()(controller.q.cx, controller.qdot.cx, controller.parameters.cx)
+            return controller.model.center_of_mass_velocity()(
+                controller.q.cx, controller.qdot.cx, controller.parameters.cx
+            )
 
         @staticmethod
         def minimize_com_acceleration(penalty: PenaltyOption, controller: PenaltyController, axes: tuple | list = None):
@@ -673,7 +683,9 @@ 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, controller.parameters.cx)
+            marker = controller.model.center_of_mass_acceleration()(
+                controller.q.cx, controller.qdot.cx, qddot, controller.parameters.cx
+            )
 
             return marker
 
@@ -718,7 +730,9 @@ def minimize_linear_momentum(penalty: PenaltyOption, controller: PenaltyControll
             PenaltyFunctionAbstract.set_axes_rows(penalty, axes)
             penalty.quadratic = True if penalty.quadratic is None else penalty.quadratic
 
-            com_velocity = controller.model.center_of_mass_velocity()(controller.q.cx, controller.qdot.cx, controller.parameters.cx)
+            com_velocity = controller.model.center_of_mass_velocity()(
+                controller.q.cx, controller.qdot.cx, controller.parameters.cx
+            )
             mass = controller.model.mass()["o0"]
             linear_momentum_cx = com_velocity * mass
             return linear_momentum_cx
@@ -892,7 +906,9 @@ def track_segment_with_custom_rt(
                 raise NotImplementedError(
                     "The track_segment_with_custom_rt penalty can only be called with a BiorbdModel"
                 )
-            r_seg_transposed = controller.model.homogeneous_matrices_in_global(segment_index)(controller.q.cx, controller.parameters.cx)[:3, :3].T
+            r_seg_transposed = controller.model.homogeneous_matrices_in_global(segment_index)(
+                controller.q.cx, controller.parameters.cx
+            )[:3, :3].T
             r_rt = controller.model.rt(rt_idx)(controller.q.cx, controller.parameters.cx)[:3, :3]
             # @Pariterre: why was this sequence is fixed?
             # @Pariterre: this is suspicious and it breaks the tests!
@@ -985,7 +1001,9 @@ def minimize_segment_rotation(
             if not isinstance(controller.model, BiorbdModel):
                 raise NotImplementedError("The minimize_segment_rotation penalty can only be called with a BiorbdModel")
 
-            jcs_segment = controller.model.homogeneous_matrices_in_global(segment_idx)(controller.q.cx, controller.parameters.cx)[:3, :3]
+            jcs_segment = controller.model.homogeneous_matrices_in_global(segment_idx)(
+                controller.q.cx, controller.parameters.cx
+            )[:3, :3]
             angles_segment = controller.model.rotation_matrix_to_euler_angles(sequence)(jcs_segment)
 
             if axes is None:
@@ -1030,7 +1048,9 @@ def minimize_segment_velocity(
                     "The minimize_segments_velocity penalty can only be called with a BiorbdModel"
                 )
             model: BiorbdModel = controller.model
-            segment_angular_velocity = model.segment_angular_velocity(segment_idx)(controller.q.cx, controller.qdot.cx, controller.parameters.cx)
+            segment_angular_velocity = model.segment_angular_velocity(segment_idx)(
+                controller.q.cx, controller.qdot.cx, controller.parameters.cx
+            )
 
             if axes is None:
                 axes = [Axis.X, Axis.Y, Axis.Z]
@@ -1097,10 +1117,18 @@ def track_vector_orientations_from_markers(
                 else vector_1_marker_1
             )
 
-            vector_0_marker_0_position = controller.model.marker(vector_0_marker_0_idx)(controller.q.cx, controller.parameters.cx)
-            vector_0_marker_1_position = controller.model.marker(vector_0_marker_1_idx)(controller.q.cx, controller.parameters.cx)
-            vector_1_marker_0_position = controller.model.marker(vector_1_marker_0_idx)(controller.q.cx, controller.parameters.cx)
-            vector_1_marker_1_position = controller.model.marker(vector_1_marker_1_idx)(controller.q.cx, controller.parameters.cx)
+            vector_0_marker_0_position = controller.model.marker(vector_0_marker_0_idx)(
+                controller.q.cx, controller.parameters.cx
+            )
+            vector_0_marker_1_position = controller.model.marker(vector_0_marker_1_idx)(
+                controller.q.cx, controller.parameters.cx
+            )
+            vector_1_marker_0_position = controller.model.marker(vector_1_marker_0_idx)(
+                controller.q.cx, controller.parameters.cx
+            )
+            vector_1_marker_1_position = controller.model.marker(vector_1_marker_1_idx)(
+                controller.q.cx, controller.parameters.cx
+            )
 
             vector_0 = vector_0_marker_1_position - vector_0_marker_0_position
             vector_1 = vector_1_marker_1_position - vector_1_marker_0_position

diff --git a/bioptim/models/biorbd/multi_biorbd_model.py b/bioptim/models/biorbd/multi_biorbd_model.py
@@ -230,7 +230,7 @@ def nb_extra_models(self) -> int:
 
     @property
     def gravity(self) -> Function:
-        biorbd_return = vertcat(*(model.gravity()['o0'] for model in self.models))
+        biorbd_return = vertcat(*(model.gravity()["o0"] for model in self.models))
         casadi_fun = Function(
             "gravity",
             [MX()],
@@ -301,7 +301,7 @@ def homogeneous_matrices_in_child(self, segment_id) -> Function:
 
     @property
     def mass(self) -> Function:
-        biorbd_return = vertcat(*(model.mass()['o0'] for model in self.models))
+        biorbd_return = vertcat(*(model.mass()["o0"] for model in self.models))
         casadi_fun = Function(
             "mass",
             [MX()],
@@ -327,9 +327,9 @@ def center_of_mass_velocity(self) -> Function:
             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),
-                )
+                biorbd_return,
+                model.center_of_mass_velocity()(q_model, qdot_model),
+            )
         casadi_fun = Function(
             "center_of_mass_velocity",
             [self.q, self.qdot],
@@ -344,9 +344,9 @@ def center_of_mass_acceleration(self) -> Function:
             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),
-                )
+                biorbd_return,
+                model.center_of_mass_acceleration()(q_model, qdot_model, qddot_model),
+            )
         casadi_fun = Function(
             "center_of_mass_acceleration",
             [self.q, self.qdot, self.qddot],
@@ -385,9 +385,9 @@ def angular_momentum(self) -> Function:
             q_model = self.q[self.variable_index("q", i)]
             qdot_model = self.qdot[self.variable_index("qdot", i)]
             biorbd_return = vertcat(
-                    biorbd_return,
-                    model.angular_momentum()(q_model, qdot_model),
-                )
+                biorbd_return,
+                model.angular_momentum()(q_model, qdot_model),
+            )
         casadi_fun = Function(
             "angular_momentum",
             [self.q, self.qdot],
@@ -483,7 +483,6 @@ def torque(self) -> Function:
         )
         return casadi_fun
 
-
     def forward_dynamics_free_floating_base(self) -> Function:
         biorbd_return = MX()
         for i, model in enumerate(self.models):
@@ -522,7 +521,6 @@ def reorder_qddot_root_joints(self):
         )
         return casadi_fun
 
-
     def forward_dynamics(self, with_contact) -> Function:
         """External forces and contact forces are not implemented yet for MultiBiorbdModel."""
 
@@ -557,12 +555,7 @@ def inverse_dynamics(self) -> Function:
             qddot_model = self.qddot[self.variable_index("qddot", i)]
             biorbd_return = vertcat(
                 biorbd_return,
-                model.inverse_dynamics()(
-                    q_model,
-                    qdot_model,
-                    qddot_model,
-                    []
-                ),
+                model.inverse_dynamics()(q_model, qdot_model, qddot_model, []),
             )
         casadi_fun = Function(
             "inverse_dynamics",
@@ -635,7 +628,9 @@ def muscle_joint_torque(self) -> Function:
                 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())
+            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],
@@ -760,7 +755,9 @@ def rigid_contact_acceleration(self, contact_index, contact_axis) -> Function:
         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)
+        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],
@@ -823,7 +820,6 @@ def contact_forces(self) -> Function:
         )
         return casadi_fun
 
-
     def passive_joint_torque(self) -> Function:
         biorbd_return = MX()
         for i, model in enumerate(self.models):

diff --git a/tests/shard1/test__global_plots.py b/tests/shard1/test__global_plots.py
@@ -117,7 +117,9 @@ 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, use_sx=False)
+    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(