[RTR] FIX: example which had an unexpected behavior, and a new that new implementation should pass

bioptim/examples/getting_started/example_continuity_as_objective.py

@@ -15,8 +15,9 @@
 
 import platform
 
-from casadi import sqrt
 import numpy as np
+from casadi import sqrt
+
 from bioptim import (
     BiorbdModel,
     OptimalControlProgram,
@@ -51,15 +52,16 @@ def out_of_sphere(controller: PenaltyController, y, z):
 
 
 def prepare_ocp_first_pass(
-    biorbd_model_path: str,
-    final_time: float,
-    n_shooting: int,
-    state_continuity_weight: float,
-    ode_solver: OdeSolverBase = OdeSolver.RK4(),
-    use_sx: bool = True,
-    n_threads: int = 1,
-    phase_dynamics: PhaseDynamics = PhaseDynamics.SHARED_DURING_THE_PHASE,
-    expand_dynamics: bool = True,
+        biorbd_model_path: str,
+        final_time: float,
+        n_shooting: int,
+        state_continuity_weight: float,
+        ode_solver: OdeSolverBase = OdeSolver.RK4(),
+        use_sx: bool = True,
+        n_threads: int = 1,
+        phase_dynamics: PhaseDynamics = PhaseDynamics.SHARED_DURING_THE_PHASE,
+        expand_dynamics: bool = True,
+        minimize_time: bool = True,
 ) -> OptimalControlProgram:
     """
     The initialization of an ocp
@@ -89,6 +91,8 @@ def prepare_ocp_first_pass(
         If the dynamics function should be expanded. Please note, this will solve the problem faster, but will slow down
         the declaration of the OCP, so it is a trade-off. Also depending on the solver, it may or may not work
         (for instance IRK is not compatible with expanded dynamics)
+    minimize_time: bool
+        If the time should be minimized
 
     Returns
     -------
@@ -100,7 +104,8 @@ def prepare_ocp_first_pass(
     # Add objective functions
     objective_functions = ObjectiveList()
     objective_functions.add(ObjectiveFcn.Lagrange.MINIMIZE_CONTROL, weight=1, key="tau")
-    objective_functions.add(ObjectiveFcn.Mayer.MINIMIZE_TIME, weight=100)
+    if minimize_time:
+        objective_functions.add(ObjectiveFcn.Mayer.MINIMIZE_TIME, weight=100 / n_shooting)
 
     # Dynamics
     dynamics = Dynamics(
@@ -122,7 +127,7 @@ def prepare_ocp_first_pass(
     n_qdot = bio_model.nb_qdot
     x_init = InitialGuessList()
     x_init["q"] = [0] * n_q
-    x_init["qdot"] = [0] * n_q
+    x_init["qdot"] = [0] * n_qdot
     x_init.add_noise(bounds=x_bounds, magnitude=0.001, n_shooting=n_shooting + 1)
 
     # Define control path constraint
@@ -164,12 +169,13 @@ def prepare_ocp_first_pass(
 
 
 def prepare_ocp_second_pass(
-    biorbd_model_path: str,
-    ns: int,
-    solution: Solution,
-    ode_solver: OdeSolverBase = OdeSolver.RK4(),
-    use_sx: bool = True,
-    n_threads: int = 1,
+        biorbd_model_path: str,
+        n_shooting: int,
+        solution: Solution,
+        ode_solver: OdeSolverBase = OdeSolver.RK4(),
+        use_sx: bool = True,
+        n_threads: int = 1,
+        minimize_time: bool = True,
 ) -> OptimalControlProgram:
     """
     The initialization of an ocp
@@ -178,12 +184,18 @@ def prepare_ocp_second_pass(
     ----------
     biorbd_model_path: str
         The path to the biorbd model
+    n_shooting: int
+        The number of shooting points to define int the direct multiple shooting program
+    solution: Solution
+        The first pass solution
     ode_solver: OdeSolverBase = OdeSolver.RK4()
         Which type of OdeSolver to use
     use_sx: bool
         If the SX variable should be used instead of MX (can be extensive on RAM)
     n_threads: int
         The number of threads to use in the paralleling (1 = no parallel computing)
+    minimize_time: bool
+        If the time should be minimized
 
     Returns
     -------
@@ -195,7 +207,8 @@ def prepare_ocp_second_pass(
     # Add objective functions
     objective_functions = ObjectiveList()
     objective_functions.add(ObjectiveFcn.Lagrange.MINIMIZE_CONTROL, weight=1, key="tau")
-    objective_functions.add(ObjectiveFcn.Mayer.MINIMIZE_TIME, weight=100)
+    if minimize_time:
+        objective_functions.add(ObjectiveFcn.Mayer.MINIMIZE_TIME, weight=100 / n_shooting)
 
     # Dynamics
     dynamics = Dynamics(DynamicsFcn.TORQUE_DRIVEN)
@@ -208,10 +221,15 @@ def prepare_ocp_second_pass(
     x_bounds["qdot"][:, 0] = 0
 
     # Initial guess
-    states = solution.decision_states(to_merge=SolutionMerge.NODES)
     x_init = InitialGuessList()
-    x_init.add("q", states["q"], interpolation=InterpolationType.EACH_FRAME)
-    x_init.add("qdot", states["qdot"], interpolation=InterpolationType.EACH_FRAME)
+    x_init.add(
+        "q", solution.decision_states(to_merge=SolutionMerge.NODES)["q"], interpolation=InterpolationType.EACH_FRAME
+    )
+    x_init.add(
+        "qdot",
+        solution.decision_states(to_merge=SolutionMerge.NODES)["qdot"],
+        interpolation=InterpolationType.EACH_FRAME,
+    )
 
     # Define control path constraint
     n_tau = bio_model.nb_tau
@@ -220,25 +238,28 @@ def prepare_ocp_second_pass(
     u_bounds["tau"] = [tau_min] * n_tau, [tau_max] * n_tau
     u_bounds["tau"][1, :] = 0  # Prevent the model from actively rotate
 
-    controls = solution.decision_controls(to_merge=SolutionMerge.NODES)
     u_init = InitialGuessList()
-    u_init.add("tau", controls["tau"], interpolation=InterpolationType.EACH_FRAME)
+    u_init.add(
+        "tau",
+        solution.decision_controls(to_merge=SolutionMerge.NODES)["tau"],
+        interpolation=InterpolationType.EACH_FRAME,
+    )
 
     constraints = ConstraintList()
     constraints.add(ConstraintFcn.SUPERIMPOSE_MARKERS, node=Node.END, first_marker="marker_2", second_marker="target_2")
     constraints.add(out_of_sphere, y=-0.45, z=0, min_bound=0.35, max_bound=np.inf, node=Node.ALL_SHOOTING)
     constraints.add(out_of_sphere, y=0.05, z=0, min_bound=0.35, max_bound=np.inf, node=Node.ALL_SHOOTING)
-    # HERE (referenced in first pass)
     constraints.add(out_of_sphere, y=0.55, z=-0.85, min_bound=0.35, max_bound=np.inf, node=Node.ALL_SHOOTING)
     constraints.add(out_of_sphere, y=0.75, z=0.2, min_bound=0.35, max_bound=np.inf, node=Node.ALL_SHOOTING)
     constraints.add(out_of_sphere, y=1.4, z=0.5, min_bound=0.35, max_bound=np.inf, node=Node.ALL_SHOOTING)
     constraints.add(out_of_sphere, y=2, z=1.2, min_bound=0.35, max_bound=np.inf, node=Node.ALL_SHOOTING)
 
     final_time = float(solution.decision_time(to_merge=SolutionMerge.NODES)[-1, 0])
+
     return OptimalControlProgram(
         bio_model,
         dynamics,
-        ns,
+        n_shooting,
         final_time,
         x_init=x_init,
         u_init=u_init,
@@ -260,19 +281,21 @@ def main():
     # --- First pass --- #
     # --- Prepare the ocp --- #
     np.random.seed(123456)
-    n_shooting = 500
     ocp_first = prepare_ocp_first_pass(
         biorbd_model_path="models/pendulum_maze.bioMod",
         final_time=5,
-        n_shooting=n_shooting,
+        n_shooting=500,
         # change the weight to observe the impact on the continuity of the solution
         # or comment to see how the constrained program would fare
         state_continuity_weight=1_000_000,
         n_threads=3,
     )
     # ocp_first.print(to_console=True)
 
-    solver_first = Solver.IPOPT(show_online_optim=platform.system() == "Linux", show_options=dict(show_bounds=True))
+    solver_first = Solver.IPOPT(
+        show_online_optim=platform.system() == "Linux",
+        show_options=dict(show_bounds=True),
+    )
     # change maximum iterations to affect the initial solution
     # it doesn't mather if it exits before the optimal solution, only that there is an initial guess
     solver_first.set_maximum_iterations(500)
@@ -286,11 +309,14 @@ def main():
 
     # # --- Second pass ---#
     # # --- Prepare the ocp --- #
-    solver_second = Solver.IPOPT(show_online_optim=platform.system() == "Linux", show_options=dict(show_bounds=True))
+    solver_second = Solver.IPOPT(
+        show_online_optim=platform.system() == "Linux",
+        show_options=dict(show_bounds=True),
+    )
     solver_second.set_maximum_iterations(10000)
 
     ocp_second = prepare_ocp_second_pass(
-        biorbd_model_path="models/pendulum_maze.bioMod", ns=n_shooting, solution=sol_first, n_threads=3
+        biorbd_model_path="models/pendulum_maze.bioMod", n_shooting=500, solution=sol_first, n_threads=3
     )
 
     # Custom plots

tests/shard1/test_continuity_as_objective.py

@@ -1,26 +1,144 @@
-from bioptim import PhaseDynamics, SolutionMerge
 import numpy as np
 import pytest
 
+from bioptim import PhaseDynamics, SolutionMerge
 from tests.utils import TestUtils
 
 
-@pytest.mark.parametrize("phase_dynamics", [PhaseDynamics.SHARED_DURING_THE_PHASE, PhaseDynamics.ONE_PER_NODE])
+@pytest.mark.parametrize(
+    "phase_dynamics",
+    [PhaseDynamics.SHARED_DURING_THE_PHASE, PhaseDynamics.ONE_PER_NODE],
+)
 def test_continuity_as_objective(phase_dynamics):
-    from bioptim.examples.getting_started import example_continuity_as_objective as ocp_module
+    from bioptim.examples.getting_started import (
+        example_continuity_as_objective as ocp_module,
+    )
 
     np.random.seed(42)
     model_path = TestUtils.bioptim_folder() + "/examples/getting_started/models/pendulum_maze.bioMod"
 
+    # first pass
     ocp = ocp_module.prepare_ocp_first_pass(
         biorbd_model_path=model_path,
+        n_threads=1,
         final_time=1,
         n_shooting=3,
-        state_continuity_weight=1000000,
+        state_continuity_weight=1000,
         phase_dynamics=phase_dynamics,
         expand_dynamics=True,
+        minimize_time=False,  # we only want to test the continuity as objective here
     )
     sol = ocp.solve()
 
-    expected = np.array([-0.1376, 2.9976372])
-    np.testing.assert_almost_equal(sol.decision_states(to_merge=SolutionMerge.NODES)["q"][:, -1], expected)
+    expected_q = [[0.0, -0.1820716, 0.0502083, -0.1376], [0.0, 0.2059882, -0.3885045, 2.9976372]]
+    expected_qdot = [[0.0, 0.13105439, -3.43794783, -23.6570729], [0.0, -0.66178869, 3.07970721, -19.12526049]]
+    expected_controls = [[-1.49607534, -0.24541618, -19.12881238], [0.0, 0.0, 0.0]]
+    expected_iterations = range(300, 600)  # 436 on my laptop @ipuch
+
+    np.testing.assert_almost_equal(sol.decision_states(to_merge=SolutionMerge.NODES)["q"], expected_q)
+    np.testing.assert_almost_equal(sol.decision_states(to_merge=SolutionMerge.NODES)["qdot"], expected_qdot)
+    np.testing.assert_almost_equal(sol.decision_controls(to_merge=SolutionMerge.NODES)["tau"], expected_controls)
+    assert sol.iterations in expected_iterations
+
+    # second pass
+    ocp_second_pass = ocp_module.prepare_ocp_second_pass(
+        biorbd_model_path=model_path,
+        n_threads=1,
+        n_shooting=3,
+        solution=sol,
+        minimize_time=False,
+    )
+    sol_second_pass = ocp_second_pass.solve()
+
+    expected_q = [[0.0, -0.12359617, 0.21522375, -0.1376], [0.0, 0.06184961, -0.37118107, 2.9976372]]
+    expected_qdot = [[0.0, 0.14235975, -1.10526128, -4.21797828], [0.0, -0.55992744, 1.17407988, -1.06473819]]
+    expected_tau = [[-1.16548046, 1.10283517, -27.94121882], [0.0, 0.0, 0.0]]
+
+    expected_vector = [
+        0.333333,
+        0,
+        0,
+        0,
+        0,
+        -0.123596,
+        0.0618496,
+        0.14236,
+        -0.559927,
+        0.215224,
+        -0.371181,
+        -1.10526,
+        1.17408,
+        -0.1376,
+        2.99764,
+        -4.21798,
+        -1.06474,
+        -1.16548,
+        0,
+        1.10284,
+        0,
+        -27.9412,
+        0,
+    ]
+    expected_constraints = [
+        1.90126e-15,
+        1.99146e-15,
+        -3.33067e-16,
+        -1.11022e-16,
+        6.10623e-16,
+        5.60663e-15,
+        1.55431e-15,
+        -1.33227e-15,
+        -1.47937e-14,
+        -5.77316e-15,
+        -1.1573e-12,
+        -5.08926e-13,
+        0,
+        -1.30451e-15,
+        -4.04121e-14,
+        1.08612,
+        1.05124,
+        0.991253,
+        0.954385,
+        0.949236,
+        0.9216,
+        0.492353,
+        0.554696,
+        0.620095,
+        1.34023,
+        1.36917,
+        1.38148,
+        1.97149,
+        2.0114,
+        2.03747,
+        2.89311,
+        2.93228,
+        2.95656,
+    ]
+
+    expected_cost = 261.095
+    expected_detailed_cost = [
+        {
+            "name": "Lagrange.MINIMIZE_CONTROL",
+            "cost_value_weighted": 261.0954331500721,
+            "cost_value": -28.003864114406223,
+        }
+    ]
+    expected_iterations = range(5, 35)  # 20 on my laptop @ipuch
+
+    np.testing.assert_almost_equal(sol_second_pass.decision_states(to_merge=SolutionMerge.NODES)["q"], expected_q)
+    np.testing.assert_almost_equal(sol_second_pass.decision_states(to_merge=SolutionMerge.NODES)["qdot"], expected_qdot)
+    np.testing.assert_almost_equal(sol_second_pass.decision_controls(to_merge=SolutionMerge.NODES)["tau"], expected_tau)
+
+    np.testing.assert_almost_equal(sol_second_pass.vector, np.array([expected_vector]).T, decimal=4)
+    np.testing.assert_almost_equal(sol_second_pass.constraints, np.array([expected_constraints]).T, decimal=4)
+    np.testing.assert_almost_equal(float(sol_second_pass.cost), expected_cost, decimal=2)
+
+    assert sol_second_pass.detailed_cost[0]["name"] == expected_detailed_cost[0]["name"]
+    np.testing.assert_almost_equal(
+        sol_second_pass.detailed_cost[0]["cost_value_weighted"], expected_detailed_cost[0]["cost_value_weighted"]
+    )
+    np.testing.assert_almost_equal(
+        sol_second_pass.detailed_cost[0]["cost_value"], expected_detailed_cost[0]["cost_value"]
+    )
+
+    assert sol_second_pass.iterations in expected_iterations