[RTR] Implicit defect type for muscles

bioptim/dynamics/configure_problem.py

@@ -495,7 +495,8 @@ def muscle_driven(
             raise NotImplementedError("MUSCLE_DRIVEN cannot be used with this enum RigidBodyDynamics yet")
 
         ConfigureProblem.configure_q(nlp, True, False)
-        ConfigureProblem.configure_qdot(nlp, True, False)
+        ConfigureProblem.configure_qdot(nlp, True, False, True)
+        ConfigureProblem.configure_qddot(nlp, False, False, True)
         if with_torque:
             ConfigureProblem.configure_tau(nlp, False, True, fatigue=fatigue)
         ConfigureProblem.configure_muscles(nlp, with_excitations, True, fatigue=fatigue)

bioptim/dynamics/dynamics_functions.py

@@ -486,7 +486,25 @@ def muscles_driven(
         if fatigue is not None and "muscles" in fatigue:
             dxdt = fatigue["muscles"].dynamics(dxdt, nlp, states, controls)
 
-        return DynamicsEvaluation(dxdt=dxdt, defects=None)
+        defects = None
+        # todo: contacts and fatigue to be handled with implicit dynamics
+        if not with_contact and fatigue is None:
+            qddot = DynamicsFunctions.get(nlp.states_dot["qddot"], nlp.states_dot.mx_reduced)
+            tau_id = DynamicsFunctions.inverse_dynamics(nlp, q, qdot, qddot, with_contact)
+            defects = MX(dq.shape[0] + tau_id.shape[0], tau_id.shape[1])
+
+            dq_defects = []
+            for _ in range(tau_id.shape[1]):
+                dq_defects.append(
+                    dq
+                    - DynamicsFunctions.compute_qdot(
+                        nlp, q, DynamicsFunctions.get(nlp.states_dot["qdot"], nlp.states_dot.mx_reduced)
+                    )
+                )
+            defects[: dq.shape[0], :] = horzcat(*dq_defects)
+            defects[dq.shape[0] :, :] = tau - tau_id
+
+        return DynamicsEvaluation(dxdt=dxdt, defects=defects)
 
     @staticmethod
     def forces_from_muscle_driven(states: MX.sym, controls: MX.sym, parameters: MX.sym, nlp) -> MX:

tests/test_global_torque_driven_ocp.py

@@ -6,7 +6,7 @@
 
 import numpy as np
 import biorbd_casadi as biorbd
-from bioptim import OdeSolver, ConstraintList, ConstraintFcn, Node
+from bioptim import OdeSolver, ConstraintList, ConstraintFcn, Node, DefectType
 
 from .utils import TestUtils
 
@@ -287,6 +287,88 @@ def test_track_marker_2D_pendulum(ode_solver):
     TestUtils.simulate(sol)
 
 
+@pytest.mark.parametrize("ode_solver", [OdeSolver.IRK, OdeSolver.COLLOCATION])
+@pytest.mark.parametrize("defects_type", [DefectType.EXPLICIT, DefectType.IMPLICIT])
+def test_track_marker_2D_pendulum(ode_solver, defects_type):
+    # Load muscle_activations_contact_tracker
+    from bioptim.examples.torque_driven_ocp import track_markers_2D_pendulum as ocp_module
+
+    bioptim_folder = os.path.dirname(ocp_module.__file__)
+
+    ode_solver = ode_solver()
+
+    # Define the problem
+    model_path = bioptim_folder + "/models/pendulum.bioMod"
+    biorbd_model = biorbd.Model(model_path)
+
+    final_time = 2
+    n_shooting = 30
+
+    # Generate data to fit
+    np.random.seed(42)
+    markers_ref = np.random.rand(3, 2, n_shooting + 1)
+    tau_ref = np.random.rand(2, n_shooting)
+
+    if isinstance(ode_solver, OdeSolver.IRK):
+        tau_ref = tau_ref * 5
+
+    ocp = ocp_module.prepare_ocp(biorbd_model, final_time, n_shooting, markers_ref, tau_ref, ode_solver=ode_solver)
+    sol = ocp.solve()
+
+    # Check constraints
+    g = np.array(sol.constraints)
+
+    # Check some of the results
+    q, qdot, tau = sol.states["q"], sol.states["qdot"], sol.controls["tau"]
+
+    if isinstance(ode_solver, OdeSolver.IRK):
+        np.testing.assert_equal(g.shape, (n_shooting * 4, 1))
+        np.testing.assert_almost_equal(g, np.zeros((n_shooting * 4, 1)))
+
+        # Check objective function value
+        f = np.array(sol.cost)
+        np.testing.assert_equal(f.shape, (1, 1))
+        np.testing.assert_almost_equal(f[0, 0], 290.6751231)
+
+        # initial and final position
+        np.testing.assert_almost_equal(q[:, 0], np.array((0, 0)))
+        np.testing.assert_almost_equal(q[:, -1], np.array((0.64142484, 2.85371719)))
+
+        # initial and final velocities
+        np.testing.assert_almost_equal(qdot[:, 0], np.array((0, 0)))
+        np.testing.assert_almost_equal(qdot[:, -1], np.array((3.46921861, 3.24168308)))
+
+        # initial and final controls
+        np.testing.assert_almost_equal(tau[:, 0], np.array((9.11770196, -13.83677175)))
+        np.testing.assert_almost_equal(tau[:, -2], np.array((1.16836132, 4.77230548)))
+
+    else:
+        np.testing.assert_equal(g.shape, (n_shooting * 4 * 5, 1))
+        np.testing.assert_almost_equal(g, np.zeros((n_shooting * 4 * 5, 1)))
+
+        # Check objective function value
+        f = np.array(sol.cost)
+        np.testing.assert_equal(f.shape, (1, 1))
+        np.testing.assert_almost_equal(f[0, 0], 281.8462122624288)
+
+        # initial and final position
+        np.testing.assert_almost_equal(q[:, 0], np.array((0, 0)))
+        np.testing.assert_almost_equal(q[:, -1], np.array((0.8390514, 3.3819348)))
+
+        # initial and final velocities
+        np.testing.assert_almost_equal(qdot[:, 0], np.array((0, 0)))
+        np.testing.assert_almost_equal(qdot[:, -1], np.array((3.2598235, 3.8800289)))
+
+        # initial and final controls
+        np.testing.assert_almost_equal(tau[:, 0], np.array((6.8532419, -12.1810791)))
+        np.testing.assert_almost_equal(tau[:, -2], np.array((0.1290981, 0.9345706)))
+
+    # save and load
+    TestUtils.save_and_load(sol, ocp, False)
+
+    # simulate
+    TestUtils.simulate(sol)
+
 def test_trampo_quaternions():
     # Load trampo_quaternion
     from bioptim.examples.torque_driven_ocp import trampo_quaternions as ocp_module