[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

@@ -140,7 +140,7 @@ def torque_driven(
             dxdt = fatigue["tau"].dynamics(dxdt, nlp, states, controls)
 
         defects = None
-        # todo: contacts and fatigue to be handled with implicit dynamics
+        # 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)
@@ -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:
@@ -694,10 +712,14 @@ def inverse_dynamics(
         Torques in tau
         """
 
-        tau_var = nlp.states["tau"] if "tau" in nlp.states else nlp.controls["tau"]
-
         if nlp.external_forces:
-            tau = MX(tau_var.mx.shape[0], nlp.ns)
+            if "tau" in nlp.states:
+                tau_shape = nlp.states["tau"].mx.shape[0]
+            elif "tau" in nlp.controls:
+                tau_shape = nlp.controls["tau"].mx.shape[0]
+            else:
+                tau_shape = nlp.model.nbGeneralizedTorque()
+            tau = MX(tau_shape, nlp.ns)
             for i, f_ext in enumerate(nlp.external_forces):
                 tau[:, i] = nlp.model.InverseDynamics(q, qdot, qddot, f_ext).to_mx()
         else:

bioptim/examples/muscle_driven_ocp/muscle_excitations_tracker.py

@@ -58,13 +58,15 @@ def generate_data(
     # Aliases
     n_q = biorbd_model.nbQ()
     n_qdot = biorbd_model.nbQdot()
+    n_qddot = biorbd_model.nbQddot()
     n_tau = biorbd_model.nbGeneralizedTorque()
     n_mus = biorbd_model.nbMuscleTotal()
     dt = final_time / n_shooting
 
     # Casadi related stuff
     symbolic_q = MX.sym("q", n_q, 1)
     symbolic_qdot = MX.sym("qdot", n_qdot, 1)
+    symbolic_qddot = MX.sym("qddot", n_qddot, 1)
     symbolic_mus_states = MX.sym("mus", n_mus, 1)
 
     symbolic_tau = MX.sym("tau", n_tau, 1)
@@ -79,6 +81,7 @@ def generate_data(
     nlp.variable_mappings = {
         "q": BiMapping(range(n_q), range(n_q)),
         "qdot": BiMapping(range(n_qdot), range(n_qdot)),
+        "qddot": BiMapping(range(n_qddot), range(n_qddot)),
         "tau": BiMapping(range(n_tau), range(n_tau)),
         "muscles": BiMapping(range(n_mus), range(n_mus)),
     }
@@ -97,6 +100,8 @@ def generate_data(
         symbolic_mus_controls,
         nlp.variable_mappings["muscles"],
     )
+    nlp.states_dot.append("qdot", [symbolic_qdot, symbolic_qdot], symbolic_qdot, nlp.variable_mappings["qdot"])
+    nlp.states_dot.append("qddot", [symbolic_qddot, symbolic_qddot], symbolic_qddot, nlp.variable_mappings["qddot"])
 
     dynamics_func = biorbd.to_casadi_func(
         "ForwardDyn",

tests/test_global_muscle_driven_ocp.py

@@ -5,7 +5,7 @@
 import pytest
 
 import numpy as np
-from bioptim import OdeSolver, Solver
+from bioptim import OdeSolver, Solver, DefectType
 
 from .utils import TestUtils
 
@@ -111,3 +111,86 @@ def test_muscle_driven_ocp(ode_solver):
 
     # simulate
     TestUtils.simulate(sol, decimal_value=5)
+
+
+@pytest.mark.parametrize("ode_solver", [OdeSolver.COLLOCATION, OdeSolver.IRK])
+def test_muscle_driven_ocp(ode_solver):
+    from bioptim.examples.muscle_driven_ocp import static_arm as ocp_module
+
+    bioptim_folder = os.path.dirname(ocp_module.__file__)
+
+    ode_solver_obj = ode_solver(defects_type=DefectType.IMPLICIT)
+    ocp = ocp_module.prepare_ocp(
+        bioptim_folder + "/models/arm26.bioMod",
+        final_time=0.1,
+        n_shooting=5,
+        weight=1,
+        ode_solver=ode_solver_obj,
+    )
+    sol = ocp.solve()
+
+    # Check objective function value
+    f = np.array(sol.cost)
+    np.testing.assert_equal(f.shape, (1, 1))
+
+    # Check constraints
+    g = np.array(sol.constraints)
+    if ode_solver == OdeSolver.COLLOCATION:
+        np.testing.assert_equal(g.shape, (20 * 5, 1))
+        np.testing.assert_almost_equal(g, np.zeros((20 * 5, 1)), decimal=5)
+    else:
+        np.testing.assert_equal(g.shape, (20, 1))
+        np.testing.assert_almost_equal(g, np.zeros((20, 1)), decimal=5)
+
+    # Check some of the results
+    q, qdot, tau, mus = sol.states["q"], sol.states["qdot"], sol.controls["tau"], sol.controls["muscles"]
+
+    if ode_solver == OdeSolver.IRK:
+        np.testing.assert_almost_equal(f[0, 0], 0.12644299285122357)
+
+        # initial and final position
+        np.testing.assert_almost_equal(q[:, 0], np.array([0.07, 1.4]))
+        np.testing.assert_almost_equal(q[:, -1], np.array([-0.19992522, 2.65885512]))
+        # initial and final velocities
+        np.testing.assert_almost_equal(qdot[:, 0], np.array([0.0, 0.0]))
+        np.testing.assert_almost_equal(qdot[:, -1], np.array([-2.31428244, 14.18136079]))
+        # initial and final controls
+        np.testing.assert_almost_equal(tau[:, 0], np.array([0.00799548, 0.02025833]))
+        np.testing.assert_almost_equal(tau[:, -2], np.array([0.00228284, 0.00281158]))
+        np.testing.assert_almost_equal(
+            mus[:, 0],
+            np.array([7.16894627e-06, 6.03295877e-01, 3.37029458e-01, 1.08379096e-05, 1.14087059e-05, 3.66744423e-01]),
+        )
+        np.testing.assert_almost_equal(
+            mus[:, -2],
+            np.array([5.46688078e-05, 6.60548530e-03, 3.77595547e-03, 4.92828831e-04, 5.09444822e-04, 9.08082070e-03]),
+        )
+
+    elif ode_solver == OdeSolver.COLLOCATION:
+        np.testing.assert_almost_equal(f[0, 0], 0.12644297341855165)
+
+        # initial and final position
+        np.testing.assert_almost_equal(q[:, 0], np.array([0.07, 1.4]))
+        np.testing.assert_almost_equal(q[:, -1], np.array([-0.19992534, 2.65884909]))
+        # initial and final velocities
+        np.testing.assert_almost_equal(qdot[:, 0], np.array([0.0, 0.0]))
+        np.testing.assert_almost_equal(qdot[:, -1], np.array([-2.3143106, 14.1812974]))
+        # initial and final controls
+        np.testing.assert_almost_equal(tau[:, 0], np.array([0.00799575, 0.02025812]))
+        np.testing.assert_almost_equal(tau[:, -2], np.array([0.00228286, 0.00281158]))
+        np.testing.assert_almost_equal(
+            mus[:, 0],
+            np.array([7.16887076e-06, 6.03293415e-01, 3.37026700e-01, 1.08380212e-05, 1.14088234e-05, 3.66740786e-01]),
+        )
+        np.testing.assert_almost_equal(
+            mus[:, -2],
+            np.array([5.4664028e-05, 6.5610959e-03, 3.7092411e-03, 4.6592962e-04, 4.8159442e-04, 9.0543847e-03]),
+        )
+    else:
+        raise ValueError("Test not ready")
+
+    # save and load
+    TestUtils.save_and_load(sol, ocp, False)
+
+    # simulate
+    TestUtils.simulate(sol, decimal_value=5)

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,89 @@ 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