diff --git a/bioptim/limits/constraints.py b/bioptim/limits/constraints.py index 96dddf22e..e0df2aa8b 100644 --- a/bioptim/limits/constraints.py +++ b/bioptim/limits/constraints.py @@ -286,6 +286,8 @@ class ConstraintFcn(Enum): TRACK_CONTACT_FORCES = (PenaltyType.TRACK_CONTACT_FORCES,) ALIGN_SEGMENT_WITH_CUSTOM_RT = (PenaltyType.ALIGN_SEGMENT_WITH_CUSTOM_RT,) ALIGN_MARKER_WITH_SEGMENT_AXIS = (PenaltyType.ALIGN_MARKER_WITH_SEGMENT_AXIS,) + COM_POSITION = (PenaltyType.MINIMIZE_COM_POSITION,) + COM_VELOCITY = (PenaltyType.MINIMIZE_COM_VELOCITY,) CUSTOM = (PenaltyType.CUSTOM,) CONTACT_FORCE = (ConstraintFunction.Functions.contact_force,) NON_SLIPPING = (ConstraintFunction.Functions.non_slipping,) diff --git a/bioptim/limits/objective_functions.py b/bioptim/limits/objective_functions.py index cdee30675..f900e7f6a 100644 --- a/bioptim/limits/objective_functions.py +++ b/bioptim/limits/objective_functions.py @@ -334,6 +334,8 @@ class Lagrange(Enum): TRACK_ALL_CONTROLS = (PenaltyType.TRACK_ALL_CONTROLS,) MINIMIZE_CONTACT_FORCES = (PenaltyType.MINIMIZE_CONTACT_FORCES,) TRACK_CONTACT_FORCES = (PenaltyType.TRACK_CONTACT_FORCES,) + MINIMIZE_COM_POSITION = (PenaltyType.MINIMIZE_COM_POSITION,) + MINIMIZE_COM_VELOCITY = (PenaltyType.MINIMIZE_COM_VELOCITY,) ALIGN_SEGMENT_WITH_CUSTOM_RT = (PenaltyType.ALIGN_SEGMENT_WITH_CUSTOM_RT,) ALIGN_MARKER_WITH_SEGMENT_AXIS = (PenaltyType.ALIGN_MARKER_WITH_SEGMENT_AXIS,) CUSTOM = (PenaltyType.CUSTOM,) @@ -368,6 +370,8 @@ class Mayer(Enum): MINIMIZE_CONTACT_FORCES = (PenaltyType.MINIMIZE_CONTACT_FORCES,) TRACK_CONTACT_FORCES = (PenaltyType.TRACK_CONTACT_FORCES,) MINIMIZE_PREDICTED_COM_HEIGHT = (PenaltyType.MINIMIZE_PREDICTED_COM_HEIGHT,) + MINIMIZE_COM_POSITION = (PenaltyType.MINIMIZE_COM_POSITION,) + MINIMIZE_COM_VELOCITY = (PenaltyType.MINIMIZE_COM_VELOCITY,) ALIGN_SEGMENT_WITH_CUSTOM_RT = (PenaltyType.ALIGN_SEGMENT_WITH_CUSTOM_RT,) ALIGN_MARKER_WITH_SEGMENT_AXIS = (PenaltyType.ALIGN_MARKER_WITH_SEGMENT_AXIS,) CUSTOM = (PenaltyType.CUSTOM,) diff --git a/bioptim/limits/penalty.py b/bioptim/limits/penalty.py index 787590b3f..5c031f9aa 100644 --- a/bioptim/limits/penalty.py +++ b/bioptim/limits/penalty.py @@ -346,6 +346,61 @@ def minimize_predicted_com_height(penalty, ocp, nlp, t, x, u, p): CoM_height = (CoM_dot[2] * CoM_dot[2]) / (2 * -g) + CoM[2] penalty.type.get_type().add_to_penalty(ocp, nlp, CoM_height, penalty) + @staticmethod + def minimize_com_position(penalty, ocp, nlp, t, x, u, p, axis=None): + """ + Adds the objective that the position of the center of mass of the model should be minimized. + If no axis is specified, the squared-norm of the CoM's position is mimimized. + Otherwise, the projection of the CoM's position on the specified axis is minimized. + """ + + target = None + if penalty.target is not None: + target = PenaltyFunctionAbstract._check_and_fill_tracking_data_size(penalty.target, (1, len(x))) + + PenaltyFunctionAbstract._add_to_casadi_func(nlp, "biorbd_CoM", nlp.model.CoM, nlp.q) + for i, v in enumerate(x): + q = nlp.mapping["q"].expand.map(v[: nlp.shape["q"]]) + CoM = nlp.casadi_func["biorbd_CoM"](q) + + if axis == None: + CoM_proj = CoM + elif not isinstance(axis, Axe): + raise RuntimeError("axis must be a bioptim.Axe") + else: + CoM_proj = CoM[axis] + + penalty.sliced_target = target[:, i] if target is not None else None + penalty.type.get_type().add_to_penalty(ocp, nlp, CoM_proj, penalty) + + @staticmethod + def minimize_com_velocity(penalty, ocp, nlp, t, x, u, p, axis=None): + """ + Adds the objective that the velocity of the center of mass of the model should be minimized. + If no axis is specified, the squared-norm of the CoM's velocity is mimimized. + Otherwise, the projection of the CoM's velocity on the specified axis is minimized. + """ + + target = None + if penalty.target is not None: + target = PenaltyFunctionAbstract._check_and_fill_tracking_data_size(penalty.target, (1, len(x))) + + PenaltyFunctionAbstract._add_to_casadi_func(nlp, "biorbd_CoMdot", nlp.model.CoMdot, nlp.q, nlp.q_dot) + for i, v in enumerate(x): + q = nlp.mapping["q"].expand.map(v[: nlp.shape["q"]]) + q_dot = nlp.mapping["q_dot"].expand.map(v[nlp.shape["q"] :]) + CoM_dot = nlp.casadi_func["biorbd_CoMdot"](q, q_dot) + + if axis == None: + CoM_dot_proj = CoM_dot[0] ** 2 + CoM_dot[1] ** 2 + CoM_dot[2] ** 2 + elif not isinstance(axis, Axe): + raise RuntimeError("axis must be a bioptim.Axe") + else: + CoM_dot_proj = CoM_dot[axis] + + penalty.sliced_target = target[:, i] if target is not None else None + penalty.type.get_type().add_to_penalty(ocp, nlp, CoM_dot_proj, penalty) + @staticmethod def minimize_contact_forces(penalty, ocp, nlp, t, x, u, p): """ @@ -528,6 +583,8 @@ def _parameter_modifier(penalty_function, parameters): or penalty_function == PenaltyType.ALIGN_SEGMENT_WITH_CUSTOM_RT or penalty_function == PenaltyType.ALIGN_MARKER_WITH_SEGMENT_AXIS or penalty_function == PenaltyType.MINIMIZE_TORQUE_DERIVATIVE + or penalty_function == PenaltyType.MINIMIZE_COM_POSITION + or penalty_function == PenaltyType.MINIMIZE_COM_VELOCITY ): parameters.quadratic = True else: @@ -745,6 +802,8 @@ class PenaltyType(Enum): MINIMIZE_CONTACT_FORCES = PenaltyFunctionAbstract.Functions.minimize_contact_forces TRACK_CONTACT_FORCES = MINIMIZE_CONTACT_FORCES MINIMIZE_PREDICTED_COM_HEIGHT = PenaltyFunctionAbstract.Functions.minimize_predicted_com_height + MINIMIZE_COM_POSITION = PenaltyFunctionAbstract.Functions.minimize_com_position + MINIMIZE_COM_VELOCITY = PenaltyFunctionAbstract.Functions.minimize_com_velocity ALIGN_SEGMENT_WITH_CUSTOM_RT = PenaltyFunctionAbstract.Functions.align_segment_with_custom_rt ALIGN_MARKER_WITH_SEGMENT_AXIS = PenaltyFunctionAbstract.Functions.align_marker_with_segment_axis CUSTOM = PenaltyFunctionAbstract.Functions.custom diff --git a/examples/torque_driven_with_contact/maximize_predicted_height_CoM.py b/examples/torque_driven_with_contact/maximize_predicted_height_CoM.py index 1296646de..d7df7abc8 100644 --- a/examples/torque_driven_with_contact/maximize_predicted_height_CoM.py +++ b/examples/torque_driven_with_contact/maximize_predicted_height_CoM.py @@ -1,4 +1,5 @@ import biorbd +import numpy as np from bioptim import ( OptimalControlProgram, @@ -13,10 +14,22 @@ InitialGuessList, ShowResult, OdeSolver, + Axe, + ConstraintList, + ConstraintFcn, + Node, ) -def prepare_ocp(model_path, phase_time, number_shooting_points, use_actuators=False, ode_solver=OdeSolver.RK4): +def prepare_ocp( + model_path, + phase_time, + number_shooting_points, + use_actuators=False, + ode_solver=OdeSolver.RK4, + objective_name="MINIMIZE_PREDICTED_COM_HEIGHT", + com_constraints=False, +): # --- Options --- # # Model path biorbd_model = biorbd.Model(model_path) @@ -30,7 +43,12 @@ def prepare_ocp(model_path, phase_time, number_shooting_points, use_actuators=Fa # Add objective functions objective_functions = ObjectiveList() - objective_functions.add(ObjectiveFcn.Mayer.MINIMIZE_PREDICTED_COM_HEIGHT, weight=-1) + if objective_name == "MINIMIZE_PREDICTED_COM_HEIGHT": + objective_functions.add(ObjectiveFcn.Mayer.MINIMIZE_PREDICTED_COM_HEIGHT, weight=-1) + elif objective_name == "MINIMIZE_COM_POSITION": + objective_functions.add(ObjectiveFcn.Lagrange.MINIMIZE_COM_POSITION, axis=Axe.Z, weight=-1) + elif objective_name == "MINIMIZE_COM_VELOCITY": + objective_functions.add(ObjectiveFcn.Lagrange.MINIMIZE_COM_VELOCITY, axis=Axe.Z, weight=-1) objective_functions.add(ObjectiveFcn.Lagrange.MINIMIZE_TORQUE, weight=1 / 100) # Dynamics @@ -40,6 +58,19 @@ def prepare_ocp(model_path, phase_time, number_shooting_points, use_actuators=Fa else: dynamics.add(DynamicsFcn.TORQUE_DRIVEN_WITH_CONTACT) + # Constraints + constraints = ConstraintList() + if com_constraints: + constraints.add( + ConstraintFcn.COM_VELOCITY, + node=Node.ALL, + min_bound=np.array([-100, -100, -100]), + max_bound=np.array([100, 100, 100]), + ) + constraints.add( + ConstraintFcn.COM_POSITION, node=Node.ALL, min_bound=np.array([-1, -1, -1]), max_bound=np.array([1, 1, 1]) + ) + # Path constraint nb_q = biorbd_model.nbQ() nb_qdot = nb_q @@ -72,6 +103,7 @@ def prepare_ocp(model_path, phase_time, number_shooting_points, use_actuators=Fa x_bounds, u_bounds, objective_functions, + constraints=constraints, tau_mapping=tau_mapping, ode_solver=ode_solver, ) @@ -81,7 +113,14 @@ def prepare_ocp(model_path, phase_time, number_shooting_points, use_actuators=Fa model_path = "2segments_4dof_2contacts.bioMod" t = 0.5 ns = 20 - ocp = prepare_ocp(model_path=model_path, phase_time=t, number_shooting_points=ns, use_actuators=False) + ocp = prepare_ocp( + model_path=model_path, + phase_time=t, + number_shooting_points=ns, + use_actuators=False, + objective_name="MINIMIZE_COM_VELOCITY", + com_constraints=True, + ) # --- Solve the program --- # sol = ocp.solve(show_online_optim=True) diff --git a/tests/test_global_torque_driven_with_contact_ocp.py b/tests/test_global_torque_driven_with_contact_ocp.py index a1d819d46..0e0fe4055 100644 --- a/tests/test_global_torque_driven_with_contact_ocp.py +++ b/tests/test_global_torque_driven_with_contact_ocp.py @@ -16,7 +16,11 @@ @pytest.mark.parametrize("ode_solver", [OdeSolver.RK4, OdeSolver.RK8, OdeSolver.IRK]) -def test_maximize_predicted_height_CoM(ode_solver): +@pytest.mark.parametrize( + "objective_name", ["MINIMIZE_PREDICTED_COM_HEIGHT", "MINIMIZE_COM_POSITION", "MINIMIZE_COM_VELOCITY"] +) +@pytest.mark.parametrize("com_constraints", [False, True]) +def test_maximize_predicted_height_CoM(ode_solver, objective_name, com_constraints): PROJECT_FOLDER = Path(__file__).parent / ".." spec = importlib.util.spec_from_file_location( "maximize_predicted_height_CoM", @@ -31,32 +35,64 @@ def test_maximize_predicted_height_CoM(ode_solver): number_shooting_points=20, use_actuators=False, ode_solver=ode_solver, + objective_name=objective_name, + com_constraints=com_constraints, ) sol = ocp.solve() - # Check objective function value f = np.array(sol["f"]) np.testing.assert_equal(f.shape, (1, 1)) - np.testing.assert_almost_equal(f[0, 0], 0.7592028279017864) # Check constraints g = np.array(sol["g"]) - np.testing.assert_equal(g.shape, (160, 1)) - np.testing.assert_almost_equal(g, np.zeros((160, 1))) + if com_constraints: + np.testing.assert_equal(g.shape, (244, 1)) + else: + np.testing.assert_equal(g.shape, (160, 1)) + np.testing.assert_almost_equal(g, np.zeros((160, 1))) # Check some of the results states, controls = Data.get_data(ocp, sol["x"]) q, qdot, tau = states["q"], states["q_dot"], controls["tau"] - # initial and final position + # initial position np.testing.assert_almost_equal(q[:, 0], np.array((0.0, 0.0, -0.5, 0.5))) - np.testing.assert_almost_equal(q[:, -1], np.array((0.1189651, -0.0904378, -0.7999996, 0.7999996))) - # initial and final velocities + # initial velocities np.testing.assert_almost_equal(qdot[:, 0], np.array((0, 0, 0, 0))) - np.testing.assert_almost_equal(qdot[:, -1], np.array((1.2636414, -1.3010929, -3.6274687, 3.6274687))) - # initial and final controls - np.testing.assert_almost_equal(tau[:, 0], np.array((-22.1218282))) - np.testing.assert_almost_equal(tau[:, -1], np.array(0.2653957)) + + if objective_name == "MINIMIZE_PREDICTED_COM_HEIGHT": + # Check objective function value + np.testing.assert_almost_equal(f[0, 0], 0.7592028279017864) + + # final position + np.testing.assert_almost_equal(q[:, -1], np.array((0.1189651, -0.0904378, -0.7999996, 0.7999996))) + # final velocities + np.testing.assert_almost_equal(qdot[:, -1], np.array((1.2636414, -1.3010929, -3.6274687, 3.6274687))) + # initial and final controls + np.testing.assert_almost_equal(tau[:, 0], np.array((-22.1218282))) + np.testing.assert_almost_equal(tau[:, -1], np.array(0.2653957)) + elif objective_name == "MINIMIZE_COM_POSITION": + # Check objective function value + np.testing.assert_almost_equal(f[0, 0], 0.458575464873056) + + # final position + np.testing.assert_almost_equal(q[:, -1], np.array((0.1189651, -0.0904378, -0.7999996, 0.7999996))) + # final velocities + np.testing.assert_almost_equal(qdot[:, -1], np.array((1.24525494, -1.28216182, -3.57468814, 3.57468814))) + # initial and final controls + np.testing.assert_almost_equal(tau[:, 0], np.array((-21.96213697))) + np.testing.assert_almost_equal(tau[:, -1], np.array(-0.22120207)) + elif objective_name == "MINIMIZE_COM_VELOCITY": + # Check objective function value + np.testing.assert_almost_equal(f[0, 0], 0.4709888694097001) + + # final position + np.testing.assert_almost_equal(q[:, -1], np.array((0.1189652, -0.09043785, -0.79999979, 0.79999979))) + # final velocities + np.testing.assert_almost_equal(qdot[:, -1], np.array((1.26103572, -1.29841047, -3.61998944, 3.61998944))) + # initial and final controls + np.testing.assert_almost_equal(tau[:, 0], np.array((-22.18008227))) + np.testing.assert_almost_equal(tau[:, -1], np.array(-0.02280469)) # save and load TestUtils.save_and_load(sol, ocp, False) diff --git a/tests/test_penalty.py b/tests/test_penalty.py index e422c4cbc..edf1eb23b 100644 --- a/tests/test_penalty.py +++ b/tests/test_penalty.py @@ -689,6 +689,39 @@ def test_penalty_minimize_predicted_com_height(value): ) +@pytest.mark.parametrize("penalty_origin", [ObjectiveFcn.Lagrange, ObjectiveFcn.Mayer, ConstraintFcn]) +@pytest.mark.parametrize("value", [0.1, -10]) +def test_penalty_minimize_com_position(value, penalty_origin): + ocp = prepare_test_ocp() + x = [DM.ones((12, 1)) * value] + if "COM_POSITION" in penalty_origin._member_names_: + penalty_type = penalty_origin.COM_POSITION + else: + penalty_type = penalty_origin.MINIMIZE_COM_POSITION + + if isinstance(penalty_type, (ObjectiveFcn.Lagrange, ObjectiveFcn.Mayer)): + penalty = Objective(penalty_type) + else: + penalty = Constraint(penalty_type) + + penalty_type.value[0](penalty, ocp, ocp.nlp[0], [], x, [], []) + + if isinstance(penalty_type, (ObjectiveFcn.Lagrange, ObjectiveFcn.Mayer)): + res = ocp.nlp[0].J[0][0]["val"] + else: + res = ocp.nlp[0].g[0][0]["val"] + + expected = np.array([[0.05], [0.05], [0.05]]) + if value == -10: + expected = np.array([[-5], [0.05], [-5]]) + + np.testing.assert_almost_equal(res, expected) + + if isinstance(penalty_type, ConstraintFcn): + np.testing.assert_almost_equal(ocp.nlp[0].g[0][0]["bounds"].min, np.array(0)) + np.testing.assert_almost_equal(ocp.nlp[0].g[0][0]["bounds"].max, np.array(0)) + + @pytest.mark.parametrize("penalty_origin", [ObjectiveFcn.Lagrange, ObjectiveFcn.Mayer, ConstraintFcn]) @pytest.mark.parametrize("value", [0.1, -10]) def test_penalty_align_segment_with_custom_rt(penalty_origin, value):