diff --git a/bioptim/limits/constraints.py b/bioptim/limits/constraints.py index 5cede291f..104f88dd4 100644 --- a/bioptim/limits/constraints.py +++ b/bioptim/limits/constraints.py @@ -514,9 +514,9 @@ def implicit_marker_acceleration( ) var = [] - var.extend([controller.states[key] for key in controller.states]) - var.extend([controller.controls[key] for key in controller.controls]) - var.extend([controller.parameters[key] for key in controller.parameters]) + var.extend([controller.states[key] for key in controller.states.keys()]) + var.extend([controller.controls[key] for key in controller.controls.keys()]) + var.extend([controller.parameters[key] for key in controller.parameters.keys()]) return controller.mx_to_cx("contact_acceleration", contact_acceleration, *var) diff --git a/bioptim/limits/multinode_constraint.py b/bioptim/limits/multinode_constraint.py index fffb4ab1a..2645ce101 100644 --- a/bioptim/limits/multinode_constraint.py +++ b/bioptim/limits/multinode_constraint.py @@ -102,6 +102,7 @@ class MultinodeConstraintFcn(FcnEnum): COM_EQUALITY = (MultinodeConstraintFunctions.Functions.com_equality,) COM_VELOCITY_EQUALITY = (MultinodeConstraintFunctions.Functions.com_velocity_equality,) TIME_CONSTRAINT = (MultinodeConstraintFunctions.Functions.time_equality,) + TRACK_TOTAL_TIME = (MultinodeConstraintFunctions.Functions.track_total_time,) STOCHASTIC_HELPER_MATRIX_EXPLICIT = (MultinodeConstraintFunctions.Functions.stochastic_helper_matrix_explicit,) STOCHASTIC_HELPER_MATRIX_IMPLICIT = (MultinodeConstraintFunctions.Functions.stochastic_helper_matrix_implicit,) STOCHASTIC_COVARIANCE_MATRIX_CONTINUITY_IMPLICIT = ( diff --git a/bioptim/limits/multinode_penalty.py b/bioptim/limits/multinode_penalty.py index 6da53489a..b82f42a46 100644 --- a/bioptim/limits/multinode_penalty.py +++ b/bioptim/limits/multinode_penalty.py @@ -321,22 +321,7 @@ def time_equality(penalty, controllers: list[PenaltyController, PenaltyControlle MultinodePenaltyFunctions.Functions._prepare_controller_cx(controllers) - def get_time_parameter_idx(controller: PenaltyController, i_phase): - time_idx = None - for i in range(controller.parameters.cx.shape[0]): - param_name = controller.parameters.cx[i].name() - if param_name == "time_phase_" + str(controller.phase_idx): - time_idx = controller.phase_idx - if time_idx is None: - raise RuntimeError( - f"Time penalty can't be established since the {i_phase}th phase has no time parameter. " - f"\nTime parameter can be added with : " - f"\nobjective_functions.add(ObjectiveFcn.[Mayer or Lagrange].MINIMIZE_TIME) or " - f"\nwith constraints.add(ConstraintFcn.TIME_CONSTRAINT)." - ) - return time_idx - - time_idx = [get_time_parameter_idx(controller, i) for i, controller in enumerate(controllers)] + time_idx = [controller.get_time_parameter_idx() for i, controller in enumerate(controllers)] time_0 = controllers[0].parameters.cx[time_idx[0]] out = controllers[0].cx.zeros((1, 1)) @@ -346,6 +331,34 @@ def get_time_parameter_idx(controller: PenaltyController, i_phase): return out + @staticmethod + def track_total_time(penalty, controllers: list[PenaltyController, PenaltyController]): + """ + The total duration of the phases must be equal to a defined duration + + Parameters + ---------- + penalty : MultinodePenalty + A reference to the phase penalty + controllers: list[PenaltyController, PenaltyController] + The penalty node elements + + Returns + ------- + The difference between the duration of the phases + """ + + MultinodePenaltyFunctions.Functions._prepare_controller_cx(controllers) + + time_idx = [controller.get_time_parameter_idx() for i, controller in enumerate(controllers)] + + time = controllers[0].parameters.cx[time_idx[0]] + for i in range(1, len(controllers)): + time_i = controllers[i].parameters.cx[time_idx[i]] + time += time_i + + return time + @staticmethod def stochastic_helper_matrix_explicit( penalty, diff --git a/bioptim/limits/penalty.py b/bioptim/limits/penalty.py index c87e1118a..8c82091ef 100644 --- a/bioptim/limits/penalty.py +++ b/bioptim/limits/penalty.py @@ -1464,6 +1464,7 @@ def _get_markers_acceleration(controller: PenaltyController, markers, CoM=False) "com_ddot" if CoM else "markers_acceleration", markers, controller.time, + controller.parameters, controller.states["q"], controller.states["qdot"], last_param, diff --git a/bioptim/limits/penalty_controller.py b/bioptim/limits/penalty_controller.py index f0299b5cb..827799700 100644 --- a/bioptim/limits/penalty_controller.py +++ b/bioptim/limits/penalty_controller.py @@ -321,3 +321,18 @@ def parameters_scaled(self) -> OptimizationVariableList: The scaled parameters """ return self._nlp.parameters.scaled + + def get_time_parameter_idx(self): + time_idx = None + for i in range(self.parameters.cx.shape[0]): + param_name = self.parameters.cx[i].name() + if param_name == "time_phase_" + str(self.phase_idx): + time_idx = self.phase_idx + if time_idx is None: + raise RuntimeError( + f"Time penalty can't be established since the {self.phase_idx}th phase has no time parameter. " + f"\nTime parameter can be added with : " + f"\nobjective_functions.add(ObjectiveFcn.[Mayer or Lagrange].MINIMIZE_TIME) or " + f"\nwith constraints.add(ConstraintFcn.TIME_CONSTRAINT)." + ) + return time_idx diff --git a/bioptim/optimization/non_linear_program.py b/bioptim/optimization/non_linear_program.py index 5043cdfdd..42ad86310 100644 --- a/bioptim/optimization/non_linear_program.py +++ b/bioptim/optimization/non_linear_program.py @@ -339,11 +339,15 @@ def mx_to_cx(name: str, symbolic_expression: SX | MX | Callable, *all_param: Any cx_types = OptimizationVariable, OptimizationVariableList, Parameter, ParameterList mx = [var.mx if isinstance(var, cx_types) else var for var in all_param] - cx = [ - var.mapping.to_second.map(var.cx_start) if hasattr(var, "mapping") else var.cx_start - for var in all_param - if isinstance(var, cx_types) - ] + cx = [] + for var in all_param: + if hasattr(var, "mapping"): + cx += [var.mapping.to_second.map(var.cx_start)] + elif hasattr(var, "cx_start"): + cx += [var.cx_start] + else: + cx += [var.cx] # This is a temporary hack until parameters are included as OptimizationVariables + return NonLinearProgram.to_casadi_func(name, symbolic_expression, *mx)(*cx) @staticmethod diff --git a/tests/shard4/test_penalty.py b/tests/shard4/test_penalty.py index 12682b849..a8705e533 100644 --- a/tests/shard4/test_penalty.py +++ b/tests/shard4/test_penalty.py @@ -11,6 +11,10 @@ Axis, ConstraintFcn, Constraint, + MultinodeConstraintFcn, + MultinodeConstraintList, + MultinodeConstraint, + MultinodeObjective, Node, RigidBodyDynamics, ControlType, @@ -65,11 +69,14 @@ def prepare_test_ocp( dynamics = DynamicsList() dynamics.add(DynamicsFcn.TORQUE_DRIVEN, expand_dynamics=True, phase_dynamics=phase_dynamics) + objective_functions = Objective(ObjectiveFcn.Mayer.MINIMIZE_TIME) + ocp = OptimalControlProgram( bio_model, dynamics, 10, 1.0, + objective_functions=objective_functions, use_sx=use_sx, ) @@ -79,7 +86,13 @@ def prepare_test_ocp( def get_penalty_value(ocp, penalty, t, x, u, p, s): - val = penalty.type(penalty, PenaltyController(ocp, ocp.nlp[0], t, x, u, [], [], [], s, [], 0), **penalty.params) + if isinstance(penalty, MultinodeConstraint) or isinstance(penalty, MultinodeObjective): + controller = [ + PenaltyController(ocp, ocp.nlp[0], t, x, u, [], [], p, s, [], 0) for i in range(len(penalty.nodes_phase)) + ] + else: + controller = PenaltyController(ocp, ocp.nlp[0], t, x, u, [], [], p, s, [], 0) + val = penalty.type(penalty, controller, **penalty.params) # changed only this one if isinstance(val, float): return val @@ -117,7 +130,7 @@ def test_penalty_minimize_time(penalty_origin, value, phase_dynamics): t = [0] x = [DM.ones((8, 1)) * value] u = [0] - p = [] + p = [1] s = [] penalty_type = penalty_origin.MINIMIZE_TIME @@ -365,7 +378,7 @@ def test_penalty_minimize_markers_acceleration(penalty_origin, implicit, value, t = [0] x = [DM.ones((8, 1)) * value] u = [0] - p = [] + p = [0] s = [] penalty_type = penalty_origin.MINIMIZE_MARKERS_ACCELERATION @@ -1086,14 +1099,46 @@ def test_penalty_time_constraint(value, phase_dynamics): t = [0] x = [0] u = [0] - p = [] + p = [0] s = [] penalty_type = ConstraintFcn.TIME_CONSTRAINT penalty = Constraint(penalty_type) res = get_penalty_value(ocp, penalty, t, x, u, p, s) - np.testing.assert_almost_equal(res, np.array([])) + np.testing.assert_almost_equal(res, np.array(0)) + + +@pytest.mark.parametrize("phase_dynamics", [PhaseDynamics.SHARED_DURING_THE_PHASE, PhaseDynamics.ONE_PER_NODE]) +@pytest.mark.parametrize("value", [0.1, -10]) +def test_penalty_constraint_total_time(value, phase_dynamics): + ocp = prepare_test_ocp(phase_dynamics=phase_dynamics) + t = [0] + x = [DM.ones((8, 1)) * value] + u = [0] + p = [0.1] + s = [] + + penalty_type = MultinodeConstraintFcn.TRACK_TOTAL_TIME + penalty = MultinodeConstraintList() + penalty.add( + penalty_type, + min_bound=0.01, + max_bound=20, + nodes_phase=(0, 1), + nodes=(Node.END, Node.END), + ) + + penalty_type( + penalty[0], + [ + PenaltyController(ocp, ocp.nlp[0], [], [], [], [], [], p, s, [], 0), + PenaltyController(ocp, ocp.nlp[0], [], [], [], [], [], p, s, [], 0), + ], + ) + res = get_penalty_value(ocp, penalty[0], t, x, u, p, s) + + np.testing.assert_almost_equal(res, np.array(0.2)) @pytest.mark.parametrize("phase_dynamics", [PhaseDynamics.SHARED_DURING_THE_PHASE, PhaseDynamics.ONE_PER_NODE])