[RTM when modify] spring load - explore more objective function - tests will come later

bioptim/examples/torque_driven_ocp/spring_load.py

@@ -6,22 +6,106 @@
 
 import platform
 
-from casadi import MX, vertcat
+from casadi import MX, vertcat, sign
 import numpy as np
 from bioptim import (
     BiorbdModel,
     OptimalControlProgram,
     Dynamics,
     ConfigureProblem,
-    Objective,
+    ObjectiveList,
     DynamicsFunctions,
     ObjectiveFcn,
     BoundsList,
     NonLinearProgram,
     Solver,
     DynamicsEvaluation,
     PhaseDynamics,
+    SolutionMerge,
 )
+from matplotlib import pyplot as plt
+
+# scenarios are based on a Mayer term (at Tf)
+# 0: maximize upward speed - expected kinematics: negative torque to get as low as possible and release
+# 1: maximize downward speed - expected kinematics: positive torque to get as high as possible and release
+# 2: minimize quadratic speed - expected kinematics: no torque no move
+# 3: maximize quadratic speed - as in 1
+# 4-7 same as 0-3 but for COM
+
+scenarios = {
+    0: {
+        "label": "max qdot(T)",
+        "quad": False,
+        "sign": -1,
+        "tau_min": -100,
+        "tau_max": 0,
+        "check_tau": -1,
+        "check_qdot(T)": 1,
+    },
+    1: {
+        "label": "min qdot(T)",
+        "quad": False,
+        "sign": 1,
+        "tau_min": 0,
+        "tau_max": 100,
+        "check_tau": 1,
+        "check_qdot(T)": -1,
+    },
+    2: {
+        "label": "min qdot(T)**2",
+        "quad": True,
+        "sign": 1,
+        "tau_min": -100,
+        "tau_max": 100,
+        "check_tau": 1,
+        "check_qdot(T)": 1,
+    },
+    3: {
+        "label": "max qdot(T)**2",
+        "quad": True,
+        "sign": -1,
+        "tau_min": -100,
+        "tau_max": 0,
+        "check_tau": -1,
+        "check_qdot(T)": 1,
+    },
+    4: {
+        "label": "max COMdot(T)",
+        "quad": False,
+        "sign": -1,
+        "tau_min": -100,
+        "tau_max": 0,
+        "check_tau": -1,
+        "check_qdot(T)": 1,
+    },
+    5: {
+        "label": "min COMdot(T)",
+        "quad": False,
+        "sign": 1,
+        "tau_min": 0,
+        "tau_max": 100,
+        "check_tau": 1,
+        "check_qdot(T)": -1,
+    },
+    6: {
+        "label": "min COMdot(T)**2",
+        "quad": True,
+        "sign": 1,
+        "tau_min": -100,
+        "tau_max": 100,
+        "check_tau": 1,
+        "check_qdot(T)": 1,
+    },
+    7: {
+        "label": "max COMdot(T)**2",
+        "quad": True,
+        "sign": -1,
+        "tau_min": -100,
+        "tau_max": 0,
+        "check_tau": -1,
+        "check_qdot(T)": 1,
+    },
+}
 
 
 def custom_dynamic(
@@ -55,7 +139,8 @@
     tau = DynamicsFunctions.get(nlp.controls["tau"], controls)
 
     force_vector = MX.zeros(6)
-    force_vector[5] = 100 * q[0] ** 2
+    stiffness = 100
+    force_vector[5] = -sign(q[0]) * stiffness * q[0] ** 2  # traction-compression spring
 
     qddot = nlp.model.forward_dynamics(q, qdot, tau, [["Point", force_vector]])
 
@@ -83,13 +168,41 @@
     biorbd_model_path: str = "models/mass_point.bioMod",
     phase_dynamics: PhaseDynamics = PhaseDynamics.SHARED_DURING_THE_PHASE,
     expand_dynamics: bool = True,
+    phase_time: float = 0.5,
+    n_shooting: float = 30,
+    scenario=1,
 ):
     # BioModel path
     m = BiorbdModel(biorbd_model_path)
     m.set_gravity(np.array((0, 0, 0)))
 
+    weight = 1
+
     # Add objective functions (high upward velocity at end point)
-    objective_functions = Objective(ObjectiveFcn.Mayer.MINIMIZE_STATE, key="qdot", index=0, weight=-1)
+    objective_functions = ObjectiveList()
+
+    if "qdot" in scenarios[scenario]["label"]:
+        objective_functions.add(
+            ObjectiveFcn.Mayer.MINIMIZE_STATE,
+            key="qdot",
+            index=0,
+            weight=weight * scenarios[scenario]["sign"],
+            quadratic=scenarios[scenario]["quad"],
+        )
+    elif "COMdot" in scenarios[scenario]["label"]:
+        objective_functions.add(
+            ObjectiveFcn.Mayer.MINIMIZE_COM_VELOCITY,
+            index=2,
+            weight=weight * scenarios[scenario]["sign"],
+            quadratic=scenarios[scenario]["quad"],
+        )
+
+    objective_functions.add(
+        ObjectiveFcn.Lagrange.MINIMIZE_CONTROL,
+        key="tau",
+        weight=1e-5,
+        quadratic=True,
+    )
 
     # Dynamics
     dynamics = Dynamics(
@@ -108,27 +221,54 @@
 
     # Define control path constraint
     u_bounds = BoundsList()
-    u_bounds["tau"] = [-100] * m.nb_tau, [0] * m.nb_tau
+    u_bounds["tau"] = scenarios[scenario]["tau_min"] * m.nb_tau, scenarios[scenario]["tau_max"] * m.nb_tau
 
     return OptimalControlProgram(
         m,
         dynamics,
-        n_shooting=30,
-        phase_time=0.5,
+        n_shooting=n_shooting,
+        phase_time=phase_time,
         x_bounds=x_bounds,
         u_bounds=u_bounds,
         objective_functions=objective_functions,
     )
 
 
 def main():
-    ocp = prepare_ocp()
+    phase_time = 0.5
+    n_shooting = 30
+    fig, axs = plt.subplots(1, 3)
+
+    for scenario in range(8):  # in [1]: #
+        print(scenarios[scenario]["label"])
+        ocp = prepare_ocp(phase_time=phase_time, n_shooting=n_shooting, scenario=scenario)
+
+        ocp.print(to_console=True, to_graph=False)
+
+        # --- Solve the program --- #
+        sol = ocp.solve(Solver.IPOPT(show_online_optim=platform.system() == "Linux"))
+        q = sol.decision_states(to_merge=SolutionMerge.NODES)["q"]
+        qdot = sol.decision_states(to_merge=SolutionMerge.NODES)["qdot"]
+        tau = sol.decision_controls(to_merge=SolutionMerge.NODES)["tau"]
+        time = np.linspace(0, phase_time, n_shooting + 1)
+        eps = 1e-6
+
+        axs[0].plot(time, q.flatten(), label=scenarios[scenario]["label"])
+        axs[0].set_title("q")
+
+        axs[1].plot(time, qdot.flatten(), label=scenarios[scenario]["label"])
+        axs[1].set_title("qdot")
+
+        axs[2].step(time, np.hstack([tau.flatten(), np.nan]), label=scenarios[scenario]["label"])
+        axs[2].set_title("tau")
 
-    # --- Solve the program --- #
-    sol = ocp.solve(Solver.IPOPT(show_online_optim=platform.system() == "Linux"))
+        # --- Show results --- #
+        sol.print_cost()
+        # sol.graphs()
+        # sol.animate()
 
-    # --- Show results --- #
-    sol.animate()
+        axs[2].legend()
+    plt.show()
 
 
 if __name__ == "__main__":

tests/shard1/test_all_examples.py → tests/shard1/test_prepare_all_examples.py

@@ -165,7 +165,17 @@ def test__getting_started__custom_plotting():
     )
 
 
-# todo: Add example_continuity_as_objective.py?
+def test__getting_started__example_continuity_as_objective():
+    from bioptim.examples.getting_started import example_continuity_as_objective as ocp_module
+
+    bioptim_folder = os.path.dirname(ocp_module.__file__)
+
+    ocp_module.prepare_ocp_first_pass(
+        biorbd_model_path=bioptim_folder + "/models/pendulum_maze.bioMod",
+        final_time=1,
+        n_shooting=100,
+        state_continuity_weight=1_000_000,
+    )
 
 
 def test__getting_started__example_cyclic_movement():
@@ -194,7 +204,7 @@ def test__getting_started__example_external_forces():
     )
 
 
-# todo: Add example_external_forces.py?
+# todo: Add example_implicit_dynamics.py?
 
 
 def test__getting_started__example_inequality_constraint():
@@ -216,7 +226,7 @@ def test__getting_started__example_inequality_constraint():
     )
 
 
-# todo: Add example_joint_acceleratio_driven.py?
+# todo: Add example_joint_acceleration_driven.py?
 
 
 def test__getting_started__example_mapping():
@@ -423,6 +433,7 @@ def test__optimal_time_ocp__time_constraint():
     )
 
 
+## torque_driven_ocp_folder
 def test__symmetrical_torque_driven_ocp__symmetry_by_constraint():
     from bioptim.examples.symmetrical_torque_driven_ocp import (
         symmetry_by_constraint as ocp_module,
@@ -522,10 +533,12 @@ def test__torque_driven_ocp__spring_load():
 
     bioptim_folder = os.path.dirname(ocp_module.__file__)
 
-    ocp_module.prepare_ocp(
-        biorbd_model_path=bioptim_folder + "/models/mass_point.bioMod",
-        expand_dynamics=False,
-    )
+    for scenario in range(8):
+        ocp_module.prepare_ocp(
+            biorbd_model_path=bioptim_folder + "/models/mass_point.bioMod",
+            expand_dynamics=False,
+            scenario=scenario,
+        )
 
 
 def test__track__optimal_estimation():