[RTR] Alignment contraint examples and bug fixes

bioptim/dynamics/configure_variables.py

@@ -2,11 +2,11 @@
 from typing import Callable, Any
 
 import numpy as np
-from casadi import DM, vertcat, Function
+from casadi import DM, vertcat, Function, horzcat
 
 from .configure_new_variable import NewVariableConfiguration
 from .fatigue.fatigue_dynamics import FatigueList
-from ..misc.enums import PlotType, ContactType
+from ..misc.enums import PlotType, ContactType, ControlType
 from ..misc.fcn_enum import FcnEnum
 from ..misc.mapping import BiMapping, Mapping
 from ..models.protocols.stochastic_biomodel import StochasticBioModel
@@ -1204,20 +1204,31 @@ def configure_qv(ocp, nlp, **extra_params) -> None:
 
         time_span_sym = vertcat(nlp.time_cx, nlp.dt)
 
-        nlp.q_v_function = Function(
+        sym_qv = (
+            nlp.controls.scaled.cx
+            if nlp.control_type == ControlType.CONSTANT
+            else nlp.cx.sym("linear_continuous_compatible_controls", nlp.controls.scaled.cx.shape[0], 2)
+        )
+        qv_init_w_algebraic = (
+            nlp.algebraic_states["q_v"].cx
+            if "q_v" in nlp.algebraic_states.keys()
+            else DM.zeros(nlp.model.nb_dependent_joints, 1)
+        )
+
+        q_v_plot_function = Function(
             "qv_function",
             [
                 time_span_sym,
-                nlp.states.cx,
-                nlp.controls.cx,
-                nlp.parameters.cx,
-                nlp.algebraic_states.cx,
+                nlp.states.scaled.cx,
+                sym_qv,
+                nlp.parameters.scaled.cx,
+                nlp.algebraic_states.scaled.cx,
                 nlp.numerical_timeseries.cx,
             ],
             [
                 nlp.model.compute_q_v()(
                     nlp.states["q_u"].cx,
-                    DM.zeros(nlp.model.nb_dependent_joints, 1),
+                    qv_init_w_algebraic,
                 )
             ],
             ["t_span", "x", "u", "p", "a", "d"],
@@ -1239,7 +1250,7 @@ def configure_qv(ocp, nlp, **extra_params) -> None:
         )
 
         nlp.plot["q_v"] = CustomPlot(
-            lambda t0, phases_dt, node_idx, x, u, p, a, d: nlp.q_v_function(
+            lambda t0, phases_dt, node_idx, x, u, p, a, d: q_v_plot_function(
                 np.concatenate([t0, t0 + phases_dt[nlp.phase_idx]]), x, u, p, a, d
             ),
             plot_type=PlotType.INTEGRATED,
@@ -1261,12 +1272,24 @@ def configure_qdotv(ocp, nlp, **extra_params) -> None:
         """
 
         time_span_sym = vertcat(nlp.time_cx, nlp.dt)
-        nlp.q_v_function = Function(
+
+        sym_qv = (
+            nlp.controls.scaled.cx
+            if nlp.control_type == ControlType.CONSTANT
+            else nlp.cx.sym("new_control", nlp.controls.scaled.cx.shape[0], 2)
+        )
+        qv_init_w_algebraic = (
+            nlp.algebraic_states["q_v"].cx
+            if "q_v" in nlp.algebraic_states.keys()
+            else DM.zeros(nlp.model.nb_dependent_joints, 1)
+        )
+
+        qdot_v_plot_function = Function(
             "qdot_v_function",
             [
                 time_span_sym,
                 nlp.states.scaled.cx,
-                nlp.controls.scaled.cx,
+                sym_qv,
                 nlp.parameters.scaled.cx,
                 nlp.algebraic_states.scaled.cx,
                 nlp.numerical_timeseries.cx,
@@ -1275,7 +1298,7 @@ def configure_qdotv(ocp, nlp, **extra_params) -> None:
                 nlp.model._compute_qdot_v()(
                     nlp.states.scaled["q_u"].cx,
                     nlp.states.scaled["qdot_u"].cx,
-                    DM.zeros(nlp.model.nb_dependent_joints, 1),
+                    qv_init_w_algebraic,
                 )
             ],
             ["t_span", "x", "u", "p", "a", "d"],
@@ -1297,7 +1320,7 @@ def configure_qdotv(ocp, nlp, **extra_params) -> None:
         )
 
         nlp.plot["qdot_v"] = CustomPlot(
-            lambda t0, phases_dt, node_idx, x, u, p, a, d: nlp.q_v_function(
+            lambda t0, phases_dt, node_idx, x, u, p, a, d: qdot_v_plot_function(
                 np.concatenate([t0, t0 + phases_dt[nlp.phase_idx]]), x, u, p, a, d
             ),
             plot_type=PlotType.INTEGRATED,
@@ -1319,12 +1342,25 @@ def configure_lagrange_multipliers_function(ocp, nlp: NpArrayDictOptional, **ext
         """
 
         time_span_sym = vertcat(nlp.time_cx, nlp.dt)
-        nlp.lagrange_multipliers_function = Function(
+
+        sym_qv = (
+            nlp.controls.scaled.cx
+            if nlp.control_type == ControlType.CONSTANT
+            else nlp.cx.sym("linear_continuous_compatible_controls", nlp.controls.scaled.cx.shape[0], 2)
+        )
+        ctrl_cx = sym_qv[:, 1] if nlp.control_type == ControlType.LINEAR_CONTINUOUS else nlp.controls.scaled.cx
+        qv_init_w_algebraic = (
+            nlp.algebraic_states["q_v"].cx
+            if "q_v" in nlp.algebraic_states.keys()
+            else DM.zeros(nlp.model.nb_dependent_joints, 1)
+        )
+
+        lagrange_multipliers_plot_function = Function(
             "lagrange_multipliers_function",
             [
                 time_span_sym,
                 nlp.states.scaled.cx,
-                nlp.controls.scaled.cx,
+                sym_qv,
                 nlp.parameters.scaled.cx,
                 nlp.algebraic_states.scaled.cx,
                 nlp.numerical_timeseries.cx,
@@ -1333,8 +1369,11 @@ def configure_lagrange_multipliers_function(ocp, nlp: NpArrayDictOptional, **ext
                 nlp.model.compute_the_lagrangian_multipliers()(
                     nlp.states.scaled["q_u"].cx,
                     nlp.states.scaled["qdot_u"].cx,
-                    DM.zeros(nlp.model.nb_dependent_joints, 1),
-                    DynamicsFunctions.get(nlp.controls["tau"], nlp.controls.scaled.cx),
+                    qv_init_w_algebraic,
+                    DynamicsFunctions.get(
+                        nlp.controls["tau"],
+                        ctrl_cx,
+                    ),
                 )
             ],
             ["t_span", "x", "u", "p", "a", "d"],
@@ -1360,7 +1399,7 @@ def configure_lagrange_multipliers_function(ocp, nlp: NpArrayDictOptional, **ext
         )
 
         nlp.plot["lagrange_multipliers"] = CustomPlot(
-            lambda t0, phases_dt, node_idx, x, u, p, a, d: nlp.lagrange_multipliers_function(
+            lambda t0, phases_dt, node_idx, x, u, p, a, d: lagrange_multipliers_plot_function(
                 np.concatenate([t0, t0 + phases_dt[nlp.phase_idx]]), x, u, p, a, d
             ),
             plot_type=PlotType.INTEGRATED,

bioptim/examples/models/two_cubes_lagrange2D_6DOF.bioMod

@@ -0,0 +1,160 @@
+version 1.0
+
+gravity 0 0 -9.81
+
+segment root
+
+    marker root_origin
+        parent root
+        position 0 0 0
+    endmarker
+
+endsegment
+
+segment cube0
+    parent root
+    RTinMatrix    0
+    RT pi/7 -pi/8 -pi/6  xyz 0 0 0
+    mass 1.0
+    com 0 0 0
+    inertia   1 0 0
+	            0 1 0
+	            0 0 1
+    mesh 0 -1 -1
+    mesh 0 0 -1
+    mesh 0 0 0
+    mesh 0 -1 0
+    mesh 0 -1 -1
+    mesh 1 -1 -1
+    mesh 1 0 -1
+    mesh 0 0 -1
+    mesh 1 0 -1
+    mesh 1 0 0
+    mesh 0 0 0
+    mesh 1 0 0
+    mesh 1 -1 0
+    mesh 0 -1 0
+    mesh 1 -1 0
+    mesh 1 -1 -1
+
+    rotations	xyz
+    translations xyz
+    rangesQ
+	    -4*pi 4*pi
+	    -4*pi 4*pi
+	    -4*pi 4*pi
+        -100 100
+        -100 100
+        -100 100
+endsegment
+
+    marker cube0_1
+        parent cube0
+        position 0 0 0
+    endmarker
+
+
+segment cube1
+    parent root
+    RTinMatrix    0
+    RT 0 0 0 xyz 0 2 0
+    mass 1.0
+    com 0.1 0.2 0.3
+    inertia   1 0 0
+	            0 1 0
+	            0 0 1
+    mesh 0 -1 -1
+    mesh 0 0 -1
+    mesh 0 0 0
+    mesh 0 -1 0
+    mesh 0 -1 -1
+    mesh 1 -1 -1
+    mesh 1 0 -1
+    mesh 0 0 -1
+    mesh 1 0 -1
+    mesh 1 0 0
+    mesh 0 0 0
+    mesh 1 0 0
+    mesh 1 -1 0
+    mesh 0 -1 0
+    mesh 1 -1 0
+    mesh 1 -1 -1
+
+    rotations	xyz
+    translations xyz
+
+    rangesQ
+	    -4*pi 4*pi
+	    -4*pi 4*pi
+	    -4*pi 4*pi
+        -100 100
+        -100 100
+        -100 100
+endsegment
+
+    marker cube1_1
+        parent cube1
+        position 0 0 0
+    endmarker
+
+segment new_frame
+    parent cube1
+    RTinMatrix    0
+    RT 1 0 0 xyz 0 0 0
+endsegment
+
+segment new_frame2
+    parent cube1
+    RTinMatrix    1
+    RT
+		1.000000	0.000000	0.000000	0.000000
+		0.000000	0.000000	-1.000000	0.000000
+		0.000000	1.000000	0.000000	0.000000
+		0.000000	0.000000	0.000000	1.000000
+endsegment
+
+segment new_frame3
+    parent cube1
+    RTinMatrix    1
+    RT
+		0.7500000	-0.4330127	0.5000000	0
+		0.6495190	0.6250000	-0.4330127	0
+		-0.1250000	0.6495190	0.7500000	0
+		0.000000	0.000000	0.000000	1.000000
+endsegment
+
+segment new_frame4
+    parent cube1
+    RTinMatrix    1
+    RT
+		1	0	0	0
+		0	-0.5000000	-0.8660254	0
+		0	0.8660254	-0.5000000	0
+		0.000000	0.000000	0.000000	1.000000
+endsegment
+
+segment new_frame5
+    parent cube1
+    RTinMatrix    1
+    RT
+		-0.000000	-0.070711	-0.997497	0
+		-0.124294	0.989762	-0.070163	0
+		0.992245	0.123983	-0.008789	0
+		0.000000	0.000000	0.000000	1.000000
+endsegment
+
+segment new_frame6
+    parent cube1
+    RTinMatrix    1
+    RT
+		0	1	0	0
+		0	0	1	0
+		1	0	0	0
+		0.000000	0.000000	0.000000	1.000000
+endsegment
+
+segment new_frame7
+    parent cube1
+    RTinMatrix    0
+    RT 0 pi/2 pi/2 xyz 0 0 0
+endsegment

bioptim/examples/models/two_cubes_lagrange2D_outofplane.bioMod

@@ -0,0 +1,76 @@
+version 1.0
+
+gravity 0 0 0.0001
+
+segment root
+
+    marker root_origin
+        parent root
+        position 0 0 0
+    endmarker
+
+endsegment
+
+segment cube0
+    parent root
+    RTinMatrix    0
+    RT pi/7 -pi/8 -pi/6  xyz 0 0 0
+    mass 1.0
+    com 0 0 0
+    inertia   1 0 0
+	            0 1 0
+	            0 0 1
+    mesh 0 -1 -1
+    mesh 0 0 -1
+    mesh 0 0 0
+    mesh 0 -1 0
+    mesh 0 -1 -1
+    mesh 1 -1 -1
+    mesh 1 0 -1
+    mesh 0 0 -1
+    mesh 1 0 -1
+    mesh 1 0 0
+    mesh 0 0 0
+    mesh 1 0 0
+    mesh 1 -1 0
+    mesh 0 -1 0
+    mesh 1 -1 0
+    mesh 1 -1 -1
+
+    rotations	x
+    rangesQ
+	    -4*pi 4*pi
+endsegment
+
+segment cube1
+    parent root
+    RTinMatrix    0
+    RT 0 0 0 xyz 0 2 0
+    mass 1.0
+    com 0 0 0
+    inertia   1 0 0
+	            0 1 0
+	            0 0 1
+    mesh 0 -1 -1
+    mesh 0 0 -1
+    mesh 0 0 0
+    mesh 0 -1 0
+    mesh 0 -1 -1
+    mesh 1 -1 -1
+    mesh 1 0 -1
+    mesh 0 0 -1
+    mesh 1 0 -1
+    mesh 1 0 0
+    mesh 0 0 0
+    mesh 1 0 0
+    mesh 1 -1 0
+    mesh 0 -1 0
+    mesh 1 -1 0
+    mesh 1 -1 -1
+
+    rotations	xyz
+    rangesQ
+	    -4*pi 4*pi
+	    -4*pi 4*pi
+	    -4*pi 4*pi
+endsegment