blacked

pyomeca · Oct 28, 2024 · 978c2d2 · 978c2d2
1 parent 1517dbe
commit 978c2d2
Show file tree

Hide file tree

Showing 6 changed files with 54 additions and 61 deletions.
diff --git a/bioptim/dynamics/dynamics_functions.py b/bioptim/dynamics/dynamics_functions.py
@@ -1136,7 +1136,9 @@ def holonomic_torque_driven(
         q_u = DynamicsFunctions.get(nlp.states["q_u"], states)
         qdot_u = DynamicsFunctions.get(nlp.states["qdot_u"], states)
         tau = DynamicsFunctions.get(nlp.controls["tau"], controls)
-        q_v_init = DM.zeros(nlp.model.nb_dependent_joints)  # @ipuch: I'm not sure of this, where q_v_init is supposed to be if not zeros?
+        q_v_init = DM.zeros(
+            nlp.model.nb_dependent_joints
+        )  # @ipuch: I'm not sure of this, where q_v_init is supposed to be if not zeros?
         qddot_u = nlp.model.partitioned_forward_dynamics()(q_u, qdot_u, q_v_init, tau, nlp.parameters.cx)
 
         return DynamicsEvaluation(dxdt=vertcat(qdot_u, qddot_u), defects=None)
diff --git a/bioptim/examples/holonomic_constraints/two_pendulums.py b/bioptim/examples/holonomic_constraints/two_pendulums.py
@@ -71,7 +71,13 @@ def compute_all_states(sol, bio_model: HolonomicBiorbdModel):
             .squeeze()
         )
         qddot[:, i] = bio_model.compute_qddot()(q[:, i], qdot[:, i], qddot_u_i).toarray().squeeze()
-        lambdas[:, i] = bio_model.compute_the_lagrangian_multipliers()(states["q_u"][:, i][:, np.newaxis], states["qdot_u"][:, i], q_v_init[:, i], tau[:, i]).toarray().squeeze()
+        lambdas[:, i] = (
+            bio_model.compute_the_lagrangian_multipliers()(
+                states["q_u"][:, i][:, np.newaxis], states["qdot_u"][:, i], q_v_init[:, i], tau[:, i]
+            )
+            .toarray()
+            .squeeze()
+        )
 
     return q, qdot, qddot, lambdas
 

diff --git a/bioptim/misc/enums.py b/bioptim/misc/enums.py
@@ -233,4 +233,3 @@ class ExternalForceType(Enum):
     FORCE = "force"
     TORQUE = "torque"
     TORQUE_AND_FORCE = "torque_and_force"
-
diff --git a/bioptim/models/biorbd/holonomic_biorbd_model.py b/bioptim/models/biorbd/holonomic_biorbd_model.py
@@ -42,7 +42,6 @@ def __init__(
         self.qddot_v = MX.sym("qddot_v_mx", self.nb_dependent_joints, 1)
         self.q_v_init = MX.sym("q_v_init_mx", self.nb_dependent_joints, 1)
 
-
     def set_newton_tol(self, newton_tol: float):
         self._newton_tol = newton_tol
 
@@ -178,9 +177,9 @@ def constrained_forward_dynamics(self) -> Function:
 
         biais = -self.holonomic_constraints_jacobian(self.qdot) @ self.qdot
         if self.stabilization:
-            biais -= self.alpha * self.holonomic_constraints(self.q) + self.beta * self.holonomic_constraints_derivative(
-                self.q, self.qdot
-            )
+            biais -= self.alpha * self.holonomic_constraints(
+                self.q
+            ) + self.beta * self.holonomic_constraints_derivative(self.q, self.qdot)
 
         tau_augmented = vertcat(tau_augmented, biais)
 
@@ -293,11 +292,13 @@ def partitioned_forward_dynamics(self) -> Function:
             modified_generalized_forces - second_term @ self.biais_vector(q, qdot) - modified_non_linear_effect
         )
 
-        casadi_fun = Function("partitioned_forward_dynamics",
-                              [self.q_u, self.qdot_u, self.q_v_init, self.tau],
-                              [qddot_u],
-                              ["q_u", "qdot_u", "q_v_init", "tau"],
-                              ["qddot_u"])
+        casadi_fun = Function(
+            "partitioned_forward_dynamics",
+            [self.q_u, self.qdot_u, self.q_v_init, self.tau],
+            [qddot_u],
+            ["q_u", "qdot_u", "q_v_init", "tau"],
+            ["qddot_u"],
+        )
 
         return casadi_fun
 
@@ -384,51 +385,33 @@ def compute_q_v(self) -> Function:
         ifcn = rootfinder("ifcn", "newton", residuals, opts)
         v_opt = ifcn(*ifcn_input)
 
-        casadi_fun = Function("compute_q_v",
-                              [self.q_u, self.q_v_init],
-                              [v_opt],
-                              ["q_u", "q_v_init"],
-                              ["q_v"])
+        casadi_fun = Function("compute_q_v", [self.q_u, self.q_v_init], [v_opt], ["q_u", "q_v_init"], ["q_v"])
         return casadi_fun
 
     def compute_q(self) -> Function:
         q_v = self.compute_q_v()(self.q_u, self.q_v_init)
         biorbd_return = self.state_from_partition(self.q_u, q_v)
-        casadi_fun = Function("compute_q",
-                                [self.q_u, self.q_v_init],
-                                [biorbd_return],
-                                ["q_u", "q_v_init"],
-                                ["q"])
+        casadi_fun = Function("compute_q", [self.q_u, self.q_v_init], [biorbd_return], ["q_u", "q_v_init"], ["q"])
         return casadi_fun
 
     def compute_qdot_v(self) -> Function:
         coupling_matrix_vu = self.coupling_matrix(self.q)
         biorbd_return = coupling_matrix_vu @ self.qdot_u
-        casadi_fun = Function("compute_qdot_v",
-                                [self.q, self.qdot_u],
-                                [biorbd_return],
-                                ["q", "qdot_u"],
-                                ["qdot_v"])
+        casadi_fun = Function("compute_qdot_v", [self.q, self.qdot_u], [biorbd_return], ["q", "qdot_u"], ["qdot_v"])
         return casadi_fun
 
     def _compute_qdot_v(self) -> Function:
         q = self.compute_q()(self.q_u, self.q_v_init)
         biorbd_return = self.compute_qdot_v()(q, self.qdot_u)
-        casadi_fun = Function("compute_qdot_v",
-                                [self.q, self.qdot_u, self.q_v_init],
-                                [biorbd_return],
-                                ["q", "qdot_u"],
-                                ["qdot_v"])
+        casadi_fun = Function(
+            "compute_qdot_v", [self.q, self.qdot_u, self.q_v_init], [biorbd_return], ["q", "qdot_u"], ["qdot_v"]
+        )
         return casadi_fun
 
     def compute_qdot(self) -> Function:
         qdot_v = self.compute_qdot_v()(self.q, self.qdot_u)
         biorbd_return = self.state_from_partition(self.qdot_u, qdot_v)
-        casadi_fun = Function("compute_qdot",
-                                [self.q, self.qdot_u],
-                                [biorbd_return],
-                                ["q", "qdot_u"],
-                                ["qdot"])
+        casadi_fun = Function("compute_qdot", [self.q, self.qdot_u], [biorbd_return], ["q", "qdot_u"], ["qdot"])
         return casadi_fun
 
     def compute_qddot_v(self) -> Function:
@@ -442,11 +425,9 @@ def compute_qddot_v(self) -> Function:
         """
         coupling_matrix_vu = self.coupling_matrix(self.q)
         biorbd_return = coupling_matrix_vu @ self.qddot_u + self.biais_vector(self.q, self.qdot)
-        casadi_fun = Function("compute_qddot_v",
-                                [self.q, self.qdot, self.qddot_u],
-                                [biorbd_return],
-                                ["q", "qdot", "qddot_u"],
-                                ["qddot_v"])
+        casadi_fun = Function(
+            "compute_qddot_v", [self.q, self.qdot, self.qddot_u], [biorbd_return], ["q", "qdot", "qddot_u"], ["qddot_v"]
+        )
         return casadi_fun
 
     def compute_qddot(self) -> Function:
@@ -460,11 +441,9 @@ def compute_qddot(self) -> Function:
         """
         qddot_v = self.compute_qddot_v()(self.q, self.qdot, self.qddot_u)
         biorbd_return = self.state_from_partition(self.qddot_u, qddot_v)
-        casadi_fun = Function("compute_qddot",
-                                [self.q, self.qdot, self.qddot_u],
-                                [biorbd_return],
-                                ["q", "qdot", "qddot_u"],
-                                ["qddot"])
+        casadi_fun = Function(
+            "compute_qddot", [self.q, self.qdot, self.qddot_u], [biorbd_return], ["q", "qdot", "qddot_u"], ["qddot"]
+        )
         return casadi_fun
 
     def compute_the_lagrangian_multipliers(self) -> Function:
@@ -481,11 +460,13 @@ def compute_the_lagrangian_multipliers(self) -> Function:
         qddot = self.compute_qddot()(q, qdot, qddot_u)
 
         biorbd_return = self._compute_the_lagrangian_multipliers()(q, qdot, qddot, self.tau)
-        casadi_fun = Function("compute_the_lagrangian_multipliers",
-                                [self.q_u, self.qdot_u, self.q_v_init, self.tau],
-                                [biorbd_return],
-                                ["q_u", "qdot_u", "q_v_init", "tau"],
-                                ["lambda"])
+        casadi_fun = Function(
+            "compute_the_lagrangian_multipliers",
+            [self.q_u, self.qdot_u, self.q_v_init, self.tau],
+            [biorbd_return],
+            ["q_u", "qdot_u", "q_v_init", "tau"],
+            ["lambda"],
+        )
         return casadi_fun
 
     def _compute_the_lagrangian_multipliers(self) -> Function:
@@ -518,10 +499,11 @@ def _compute_the_lagrangian_multipliers(self) -> Function:
         biorbd_return = partitioned_constraints_jacobian_v_t_inv @ (
             m_vu @ qddot_u + m_vv @ qddot_v + non_linear_effect_v - partitioned_tau_v
         )
-        casadi_fun = Function("_compute_the_lagrangian_multipliers",
-                                [self.q, self.qdot, self.qddot, self.tau],
-                                [biorbd_return],
-                                ["q", "qdot", "qddot", "tau"],
-                                ["lambda"])
+        casadi_fun = Function(
+            "_compute_the_lagrangian_multipliers",
+            [self.q, self.qdot, self.qddot, self.tau],
+            [biorbd_return],
+            ["q", "qdot", "qddot", "tau"],
+            ["lambda"],
+        )
         return casadi_fun
-
diff --git a/bioptim/models/biorbd/variational_biorbd_model.py b/bioptim/models/biorbd/variational_biorbd_model.py
@@ -29,7 +29,6 @@ def __init__(
         self.control_type = control_type
         self.control_discrete_approximation = control_discrete_approximation
 
-
     def discrete_lagrangian(
         self,
         q1: MX | SX | DM,

diff --git a/tests/shard1/test_biorbd_model_holonomic.py b/tests/shard1/test_biorbd_model_holonomic.py
@@ -103,7 +103,9 @@ def test_model_holonomic():
     )
     TestUtils.assert_equal(model.holonomic_constraints_derivative(q, q_dot), [-7.65383105, -0.44473154])
     TestUtils.assert_equal(model.holonomic_constraints_double_derivative(q, q_dot, q_ddot), [10.23374996, -11.73729905])
-    TestUtils.assert_equal(model.constrained_forward_dynamics()(q, q_dot, tau, []), [-5.18551845, -3.01921376, 25.79451813])
+    TestUtils.assert_equal(
+        model.constrained_forward_dynamics()(q, q_dot, tau, []), [-5.18551845, -3.01921376, 25.79451813]
+    )
     TestUtils.assert_equal(
         model.partitioned_mass_matrix(q),
         [
@@ -146,7 +148,10 @@ def test_model_holonomic():
     )
     TestUtils.assert_equal(
         model.compute_the_lagrangian_multipliers()(
-            MX(np.zeros(model.nb_independent_joints)), MX(np.ones(model.nb_independent_joints) * 0.001), MX(np.zeros(model.nb_dependent_joints)), tau
+            MX(np.zeros(model.nb_independent_joints)),
+            MX(np.ones(model.nb_independent_joints) * 0.001),
+            MX(np.zeros(model.nb_dependent_joints)),
+            tau,
         ),
         [np.nan, np.nan],
         expand=False,