From 5f27ca8e5316ad03684eed24cc5cbed23edd9115 Mon Sep 17 00:00:00 2001 From: Charbie Date: Fri, 29 Mar 2024 18:11:35 -0400 Subject: [PATCH 1/7] fixed some tests --- bioptim/dynamics/configure_problem.py | 6 ++- bioptim/dynamics/dynamics_functions.py | 42 ++++++++++++++++++- .../custom_package/custom_dynamics.py | 1 + .../getting_started/custom_dynamics.py | 1 + bioptim/limits/penalty.py | 1 + 5 files changed, 47 insertions(+), 4 deletions(-) diff --git a/bioptim/dynamics/configure_problem.py b/bioptim/dynamics/configure_problem.py index aaff5c4d1..582992010 100644 --- a/bioptim/dynamics/configure_problem.py +++ b/bioptim/dynamics/configure_problem.py @@ -149,7 +149,7 @@ def custom(ocp, A reference to the phase """ - nlp.dynamics_type.configure(ocp, nlp, dynamics_constants_used_at_each_nodes, **extra_params) + nlp.dynamics_type.configure(ocp, nlp, **extra_params) @staticmethod def torque_driven( @@ -1077,6 +1077,7 @@ def configure_contact_function(ocp, nlp, dyn_func: Callable, **extra_params): nlp.controls.scaled.mx_reduced, nlp.parameters.scaled.mx_reduced, nlp.algebraic_states.scaled.mx_reduced, + nlp.dynamics_constants.mx, ], [ dyn_func( @@ -1085,11 +1086,12 @@ def configure_contact_function(ocp, nlp, dyn_func: Callable, **extra_params): nlp.controls.scaled.mx_reduced, nlp.parameters.scaled.mx_reduced, nlp.algebraic_states.scaled.mx_reduced, + nlp.dynamics_constants.mx, nlp, **extra_params, ) ], - ["t_span", "x", "u", "p", "a"], + ["t_span", "x", "u", "p", "a", "dynamics_constants"], ["contact_forces"], ).expand() diff --git a/bioptim/dynamics/dynamics_functions.py b/bioptim/dynamics/dynamics_functions.py index 3543a996c..695f25de9 100644 --- a/bioptim/dynamics/dynamics_functions.py +++ b/bioptim/dynamics/dynamics_functions.py @@ -44,7 +44,7 @@ class DynamicsFunctions: @staticmethod def custom( - time: MX.sym, states: MX.sym, controls: MX.sym, parameters: MX.sym, algebraic_states: MX.sym, nlp + time: MX.sym, states: MX.sym, controls: MX.sym, parameters: MX.sym, algebraic_states: MX.sym, dynamics_constants: MX.sym, nlp ) -> DynamicsEvaluation: """ Interface to custom dynamic function provided by the user. @@ -61,6 +61,8 @@ def custom( The parameters of the system algebraic_states: MX.sym The algebraic_states of the system + dynamics_constants: MX.sym + The dynamics constants of the system nlp: NonLinearProgram The definition of the system @@ -72,7 +74,7 @@ def custom( The defects of the implicit dynamics """ - return nlp.dynamics_type.dynamic_function(time, states, controls, parameters, algebraic_states, nlp) + return nlp.dynamics_type.dynamic_function(time, states, controls, parameters, algebraic_states, dynamics_constants, nlp) @staticmethod def torque_driven( @@ -81,6 +83,7 @@ def torque_driven( controls: MX.sym, parameters: MX.sym, algebraic_states: MX.sym, + dynamics_constants: MX.sym, nlp, with_contact: bool, with_passive_torque: bool, @@ -105,6 +108,8 @@ def torque_driven( The parameters of the system algebraic_states: MX.sym The algebraic states of the system + dynamics_constants: MX.sym + The dynamics constants of the system nlp: NonLinearProgram The definition of the system with_contact: bool @@ -196,6 +201,7 @@ def torque_driven_free_floating_base( controls: MX.sym, parameters: MX.sym, algebraic_states: MX.sym, + dynamics_constants: MX.sym, nlp, with_passive_torque: bool, with_ligament: bool, @@ -216,6 +222,8 @@ def torque_driven_free_floating_base( The parameters of the system algebraic_states: MX.sym The algebraic states of the system + dynamics_constants: MX.sym + The dynamics constants of the system nlp: NonLinearProgram The definition of the system with_passive_torque: bool @@ -267,6 +275,7 @@ def stochastic_torque_driven( controls: MX.sym, parameters: MX.sym, algebraic_states: MX.sym, + dynamics_constants: MX.sym, nlp, with_contact: bool, with_friction: bool, @@ -286,6 +295,8 @@ def stochastic_torque_driven( The parameters of the system algebraic_states: MX.sym The algebraic states variables of the system + dynamics_constants: MX.sym + The dynamics constants of the system nlp: NonLinearProgram The definition of the system with_contact: bool @@ -332,6 +343,7 @@ def stochastic_torque_driven_free_floating_base( controls: MX.sym, parameters: MX.sym, algebraic_states: MX.sym, + dynamics_constants: MX.sym, nlp, with_friction: bool, ) -> DynamicsEvaluation: @@ -350,6 +362,8 @@ def stochastic_torque_driven_free_floating_base( The parameters of the system algebraic_states: MX.sym The algebraic states of the system + dynamics_constants: MX.sym + The dynamics constants of the system nlp: NonLinearProgram The definition of the system with_friction: bool @@ -457,6 +471,7 @@ def torque_activations_driven( controls: MX.sym, parameters: MX.sym, algebraic_states: MX.sym, + dynamics_constants: MX.sym, nlp, with_contact: bool, with_passive_torque: bool, @@ -479,6 +494,8 @@ def torque_activations_driven( The parameters of the system algebraic_states: MX.sym The algebraic states of the system + dynamics_constants: MX.sym + The dynamics constants of the system nlp: NonLinearProgram The definition of the system with_contact: bool @@ -524,6 +541,7 @@ def torque_derivative_driven( controls: MX.sym, parameters: MX.sym, algebraic_states: MX.sym, + dynamics_constants: MX.sym, nlp, rigidbody_dynamics: RigidBodyDynamics, with_contact: bool, @@ -546,6 +564,8 @@ def torque_derivative_driven( The parameters of the system algebraic_states: MX.sym The algebraic states of the system + dynamics_constants: MX.sym + The dynamics constants of the system nlp: NonLinearProgram The definition of the system rigidbody_dynamics: RigidBodyDynamics @@ -603,6 +623,7 @@ def forces_from_torque_driven( controls: MX.sym, parameters: MX.sym, algebraic_states: MX.sym, + dynamics_constants: MX.sym, nlp, with_passive_torque: bool = False, with_ligament: bool = False, @@ -623,6 +644,8 @@ def forces_from_torque_driven( The parameters of the system algebraic_states: MX.sym The algebraic states of the system + dynamics_constants: MX.sym + The dynamics constants of the system nlp: NonLinearProgram The definition of the system with_passive_torque: bool @@ -653,6 +676,7 @@ def forces_from_torque_activation_driven( controls: MX.sym, parameters: MX.sym, algebraic_states: MX.sym, + dynamics_constants: MX.sym, nlp, with_passive_torque: bool = False, with_ligament: bool = False, @@ -673,6 +697,8 @@ def forces_from_torque_activation_driven( The parameters of the system algebraic_states: MX.sym The algebraic states of the system + dynamics_constants: MX.sym + The dynamics constants of the system nlp: NonLinearProgram The definition of the system with_passive_torque: bool @@ -703,6 +729,7 @@ def muscles_driven( controls: MX.sym, parameters: MX.sym, algebraic_states: MX.sym, + dynamics_constants: MX.sym, nlp, with_contact: bool, with_passive_torque: bool = False, @@ -727,6 +754,8 @@ def muscles_driven( The parameters of the system algebraic_states: MX.sym The algebraic states of the system + dynamics_constants: MX.sym + The dynamics constants of the system nlp: NonLinearProgram The definition of the system with_contact: bool @@ -848,6 +877,7 @@ def forces_from_muscle_driven( controls: MX.sym, parameters: MX.sym, algebraic_states: MX.sym, + dynamics_constants: MX.sym, nlp, with_passive_torque: bool = False, with_ligament: bool = False, @@ -868,6 +898,8 @@ def forces_from_muscle_driven( The parameters of the system algebraic_states: MX.sym The algebraic states of the system + dynamics_constants: MX.sym + The dynamics constants of the system nlp: NonLinearProgram The definition of the system with_passive_torque: bool @@ -901,6 +933,7 @@ def joints_acceleration_driven( controls: MX.sym, parameters: MX.sym, algebraic_states: MX.sym, + dynamics_constants: MX.sym, nlp, rigidbody_dynamics: RigidBodyDynamics = RigidBodyDynamics.ODE, ) -> DynamicsEvaluation: @@ -919,6 +952,8 @@ def joints_acceleration_driven( The parameters of the system algebraic_states: MX.sym The algebraic states of the system + dynamics_constants: MX.sym + The dynamics constants of the system nlp: NonLinearProgram The definition of the system rigidbody_dynamics: RigidBodyDynamics @@ -1210,6 +1245,7 @@ def holonomic_torque_driven( controls: MX.sym, parameters: MX.sym, algebraic_states: MX.sym, + dynamics_constants: MX.sym, nlp, external_forces: list = None, ) -> DynamicsEvaluation: @@ -1228,6 +1264,8 @@ def holonomic_torque_driven( The parameters acting on the system algebraic_states: MX.sym The algebraic states of the system + dynamics_constants: MX.sym + The dynamics constants of the system nlp: NonLinearProgram A reference to the phase external_forces: list[Any] diff --git a/bioptim/examples/custom_model/custom_package/custom_dynamics.py b/bioptim/examples/custom_model/custom_package/custom_dynamics.py index 6d50ea7e3..fb806b509 100644 --- a/bioptim/examples/custom_model/custom_package/custom_dynamics.py +++ b/bioptim/examples/custom_model/custom_package/custom_dynamics.py @@ -20,6 +20,7 @@ def custom_dynamics( controls: MX, parameters: MX, algebraic_states: MX, + dynamics_constants: MX, nlp: NonLinearProgram, ) -> DynamicsEvaluation: """ diff --git a/bioptim/examples/getting_started/custom_dynamics.py b/bioptim/examples/getting_started/custom_dynamics.py index e886cbfac..6eca18237 100644 --- a/bioptim/examples/getting_started/custom_dynamics.py +++ b/bioptim/examples/getting_started/custom_dynamics.py @@ -38,6 +38,7 @@ def custom_dynamics( controls: MX | SX, parameters: MX | SX, algebraic_states: MX | SX, + dynamics_constants: MX | SX, nlp: NonLinearProgram, my_additional_factor=1, ) -> DynamicsEvaluation: diff --git a/bioptim/limits/penalty.py b/bioptim/limits/penalty.py index 15edf5aab..8f4f9aa83 100644 --- a/bioptim/limits/penalty.py +++ b/bioptim/limits/penalty.py @@ -792,6 +792,7 @@ def minimize_contact_forces( controller.controls.cx_start, controller.parameters.cx, controller.algebraic_states.cx_start, + controller.dynamics_constants.cx, ) return contact_force From afae924090a19a9ed7ee7b6bab95de8d66bd88ea Mon Sep 17 00:00:00 2001 From: Charbie Date: Mon, 1 Apr 2024 11:17:00 -0400 Subject: [PATCH 2/7] plots working again --- bioptim/dynamics/configure_new_variable.py | 16 +++++------ bioptim/dynamics/configure_problem.py | 4 +-- bioptim/examples/getting_started/pendulum.py | 5 ++-- bioptim/gui/plot.py | 28 +++++++++++++++++--- bioptim/limits/constraints.py | 1 + 5 files changed, 39 insertions(+), 15 deletions(-) diff --git a/bioptim/dynamics/configure_new_variable.py b/bioptim/dynamics/configure_new_variable.py index 422a09c7d..dcc2ced85 100644 --- a/bioptim/dynamics/configure_new_variable.py +++ b/bioptim/dynamics/configure_new_variable.py @@ -418,7 +418,7 @@ def _declare_cx_and_plot(self): ) if not self.skip_plot: self.nlp.plot[f"{self.name}_states"] = CustomPlot( - lambda t0, phases_dt, node_idx, x, u, p, a: ( + lambda t0, phases_dt, node_idx, x, u, p, a, d: ( x[self.nlp.states.key_index(self.name), :] if x.any() else np.ndarray((cx[0][0].shape[0], 1)) * np.nan @@ -455,7 +455,7 @@ def _declare_cx_and_plot(self): plot_type = PlotType.PLOT if self.nlp.control_type == ControlType.LINEAR_CONTINUOUS else PlotType.STEP if not self.skip_plot: self.nlp.plot[f"{self.name}_controls"] = CustomPlot( - lambda t0, phases_dt, node_idx, x, u, p, a: ( + lambda t0, phases_dt, node_idx, x, u, p, a, d: ( u[self.nlp.controls.key_index(self.name), :] if u.any() else np.ndarray((cx[0][0].shape[0], 1)) * np.nan @@ -521,7 +521,7 @@ def _declare_cx_and_plot(self): if not self.skip_plot: all_variables_in_one_subplot = True if self.name in ["m", "cov", "k"] else False self.nlp.plot[f"{self.name}_algebraic"] = CustomPlot( - lambda t0, phases_dt, node_idx, x, u, p, a: ( + lambda t0, phases_dt, node_idx, x, u, p, a, d: ( a[self.nlp.algebraic_states.key_index(self.name), :] if a.any() else np.ndarray((cx[0][0].shape[0], 1)) * np.nan @@ -591,7 +591,7 @@ def _manage_fatigue_to_new_variable( legend = [f"{name}_{i}" for i in name_elements] fatigue_plot_name = f"fatigue_{name}" nlp.plot[fatigue_plot_name] = CustomPlot( - lambda t0, phases_dt, node_idx, x, u, p, a: ( + lambda t0, phases_dt, node_idx, x, u, p, a, d: ( x[:n_elements, :] if x.any() else np.ndarray((len(name_elements), 1)) ) * np.nan, @@ -601,7 +601,7 @@ def _manage_fatigue_to_new_variable( ) control_plot_name = f"{name}_controls" if not multi_interface and split_controls else f"{name}" nlp.plot[control_plot_name] = CustomPlot( - lambda t0, phases_dt, node_idx, x, u, p, a: ( + lambda t0, phases_dt, node_idx, x, u, p, a, d: ( u[:n_elements, :] if u.any() else np.ndarray((len(name_elements), 1)) ) * np.nan, @@ -621,7 +621,7 @@ def _manage_fatigue_to_new_variable( var_names_with_suffix[-1], name_elements, ocp, nlp, as_states, as_controls, skip_plot=True ) nlp.plot[f"{var_names_with_suffix[-1]}_controls"] = CustomPlot( - lambda t0, phases_dt, node_idx, x, u, p, a, key: ( + lambda t0, phases_dt, node_idx, x, u, p, a, d, key: ( u[nlp.controls.key_index(key), :] if u.any() else np.ndarray((len(name_elements), 1)) * np.nan ), plot_type=PlotType.STEP, @@ -632,7 +632,7 @@ def _manage_fatigue_to_new_variable( elif i == 0: NewVariableConfiguration(f"{name}", name_elements, ocp, nlp, as_states, as_controls, skip_plot=True) nlp.plot[f"{name}_controls"] = CustomPlot( - lambda t0, phases_dt, node_idx, x, u, p, a, key: ( + lambda t0, phases_dt, node_idx, x, u, p, a, d, key: ( u[nlp.controls.key_index(key), :] if u.any() else np.ndarray((len(name_elements), 1)) * np.nan ), plot_type=PlotType.STEP, @@ -645,7 +645,7 @@ def _manage_fatigue_to_new_variable( name_tp = f"{var_names_with_suffix[-1]}_{params}" NewVariableConfiguration(name_tp, name_elements, ocp, nlp, True, False, skip_plot=True) nlp.plot[name_tp] = CustomPlot( - lambda t0, phases_dt, node_idx, x, u, p, a, key, mod: ( + lambda t0, phases_dt, node_idx, x, u, p, a, d, key, mod: ( mod * x[nlp.states.key_index(key), :] if x.any() else np.ndarray((len(name_elements), 1)) * np.nan ), plot_type=PlotType.INTEGRATED, diff --git a/bioptim/dynamics/configure_problem.py b/bioptim/dynamics/configure_problem.py index 582992010..5ed98035a 100644 --- a/bioptim/dynamics/configure_problem.py +++ b/bioptim/dynamics/configure_problem.py @@ -1113,7 +1113,7 @@ def configure_contact_function(ocp, nlp, dyn_func: Callable, **extra_params): ) nlp.plot["contact_forces"] = CustomPlot( - lambda t0, phases_dt, node_idx, x, u, p, a: nlp.contact_forces_func( + lambda t0, phases_dt, node_idx, x, u, p, a, d: nlp.contact_forces_func( [t0, t0 + phases_dt[nlp.phase_idx]], x, u, p, a ), plot_type=PlotType.INTEGRATED, @@ -1186,7 +1186,7 @@ def configure_soft_contact_function(ocp, nlp): to_second=[i for i, c in enumerate(all_soft_contact_names) if c in soft_contact_names_in_phase], ) nlp.plot[f"soft_contact_forces_{nlp.model.soft_contact_names[i_sc]}"] = CustomPlot( - lambda t0, phases_dt, node_idx, x, u, p, a: nlp.soft_contact_forces_func( + lambda t0, phases_dt, node_idx, x, u, p, a, d: nlp.soft_contact_forces_func( [t0, t0 + phases_dt[nlp.phase_idx]], x, u, p, a )[(i_sc * 6) : ((i_sc + 1) * 6), :], plot_type=PlotType.INTEGRATED, diff --git a/bioptim/examples/getting_started/pendulum.py b/bioptim/examples/getting_started/pendulum.py index 50eeaaa13..36b5b30c8 100644 --- a/bioptim/examples/getting_started/pendulum.py +++ b/bioptim/examples/getting_started/pendulum.py @@ -133,7 +133,7 @@ def main(): """ # --- Prepare the ocp --- # - ocp = prepare_ocp(biorbd_model_path="models/pendulum.bioMod", final_time=1, n_shooting=400, n_threads=2) + ocp = prepare_ocp(biorbd_model_path="models/pendulum.bioMod", final_time=1, n_shooting=40, n_threads=2) # Custom plots ocp.add_plot_penalty(CostType.ALL) # This will display the objectives and constraints at the current iteration @@ -147,7 +147,8 @@ def main(): ocp.print(to_console=False, to_graph=False) # --- Solve the ocp. Please note that online graphics only works with the Linux operating system --- # - sol = ocp.solve(Solver.IPOPT(show_online_optim=platform.system() == "Linux")) + sol = ocp.solve(Solver.IPOPT(show_online_optim=False)) + sol.graphs() sol.print_cost() diff --git a/bioptim/gui/plot.py b/bioptim/gui/plot.py index c5c37787c..a05a0ea4c 100644 --- a/bioptim/gui/plot.py +++ b/bioptim/gui/plot.py @@ -350,6 +350,7 @@ def legend_without_duplicate_labels(ax): size_u = nlp.controls.shape size_p = nlp.parameters.shape size_a = nlp.algebraic_states.shape + size_d = nlp.dynamics_constants.shape if "penalty" in nlp.plot[key].parameters: penalty = nlp.plot[key].parameters["penalty"] @@ -362,6 +363,7 @@ def legend_without_duplicate_labels(ax): size_u = casadi_function.size_in("u")[0] size_p = casadi_function.size_in("p")[0] size_a = casadi_function.size_in("a")[0] + size_d = casadi_function.size_in("dynamics_constants")[0] size = ( nlp.plot[key] @@ -373,6 +375,7 @@ def legend_without_duplicate_labels(ax): np.zeros((size_u, 1)), # controls np.zeros((size_p, 1)), # parameters np.zeros((size_a, 1)), # algebraic_states + np.zeros((size_d, 1)), # dynamics_constants **nlp.plot[key].parameters, # parameters ) .shape[0] @@ -711,11 +714,19 @@ def update_data( self._update_xdata(time_stepwise) for nlp in self.ocp.nlp: + from ..interfaces.interface_utils import _get_dynamics_constants phase_idx = nlp.phase_idx x_decision = data_states_decision[phase_idx] x_stepwise = data_states_stepwise[phase_idx] u = data_controls[phase_idx] a = data_algebraic_states[phase_idx] + d = [] + for n_idx in range(nlp.ns + 1): + d_tp = _get_dynamics_constants(self.ocp, phase_idx, n_idx, 0) + if d_tp.shape == (0, 0): + d += [np.array([])] + else: + d += [np.array(d_tp)] for key in self.variable_sizes[phase_idx]: y_data = self._compute_y_from_plot_func( @@ -728,6 +739,7 @@ def update_data( u, p, a, + d, ) if y_data is None: continue @@ -749,7 +761,7 @@ def update_data( update_conditioning_plots(args["x"], self.ocp) def _compute_y_from_plot_func( - self, custom_plot: CustomPlot, phase_idx, time_stepwise, dt, x_decision, x_stepwise, u, p, a + self, custom_plot: CustomPlot, phase_idx, time_stepwise, dt, x_decision, x_stepwise, u, p, a, d ): """ Compute the y data from the plot function @@ -774,11 +786,14 @@ def _compute_y_from_plot_func( The parameters of the current phase a The algebraic states of the current phase + d + The dynamics constants of the current phase Returns ------- The y data """ + from ..interfaces.interface_utils import _get_dynamics_constants if not custom_plot: return None @@ -809,7 +824,13 @@ def _compute_y_from_plot_func( idx, lambda p_idx, n_idx, sn_idx: a[n_idx][:, sn_idx] if n_idx < len(a) else np.ndarray((0, 1)), ) - + d_node = PenaltyHelpers.dynamics_constants( + penalty, + idx, + lambda p_idx, n_idx, sn_idx: _get_dynamics_constants(self.ocp, p_idx, n_idx, sn_idx) + ) + if d_node.shape == (0, 0): + d_node = DM(0, 1) else: t0 = time_stepwise[phase_idx][node_idx][0] @@ -817,8 +838,9 @@ def _compute_y_from_plot_func( u_node = u[node_idx] if node_idx < len(u) else np.ndarray((0, 1)) p_node = p a_node = a[node_idx] + d_node = d[node_idx] - tp = custom_plot.function(t0, dt, node_idx, x_node, u_node, p_node, a_node, **custom_plot.parameters) + tp = custom_plot.function(t0, dt, node_idx, x_node, u_node, p_node, a_node, d_node, **custom_plot.parameters) y_tp = np.ndarray((max(custom_plot.phase_mappings.to_first.map_idx) + 1, tp.shape[1])) * np.nan for ctr, axe_index in enumerate(custom_plot.phase_mappings.to_first.map_idx): diff --git a/bioptim/limits/constraints.py b/bioptim/limits/constraints.py index 76468cf04..0ec9968b2 100644 --- a/bioptim/limits/constraints.py +++ b/bioptim/limits/constraints.py @@ -260,6 +260,7 @@ def non_slipping( controller.controls.cx_start, controller.parameters.cx, controller.algebraic_states.cx_start, + controller.dynamics_constants.cx, ) normal_contact_force_squared = sum1(contact[normal_component_idx, 0]) ** 2 if len(tangential_component_idx) == 1: From 775990eb7968b6b35074afed088a79cd7be3f404 Mon Sep 17 00:00:00 2001 From: Charbie Date: Mon, 1 Apr 2024 11:21:20 -0400 Subject: [PATCH 3/7] reverted changes on pendulum --- bioptim/examples/getting_started/pendulum.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/bioptim/examples/getting_started/pendulum.py b/bioptim/examples/getting_started/pendulum.py index 36b5b30c8..50eeaaa13 100644 --- a/bioptim/examples/getting_started/pendulum.py +++ b/bioptim/examples/getting_started/pendulum.py @@ -133,7 +133,7 @@ def main(): """ # --- Prepare the ocp --- # - ocp = prepare_ocp(biorbd_model_path="models/pendulum.bioMod", final_time=1, n_shooting=40, n_threads=2) + ocp = prepare_ocp(biorbd_model_path="models/pendulum.bioMod", final_time=1, n_shooting=400, n_threads=2) # Custom plots ocp.add_plot_penalty(CostType.ALL) # This will display the objectives and constraints at the current iteration @@ -147,8 +147,7 @@ def main(): ocp.print(to_console=False, to_graph=False) # --- Solve the ocp. Please note that online graphics only works with the Linux operating system --- # - sol = ocp.solve(Solver.IPOPT(show_online_optim=False)) - sol.graphs() + sol = ocp.solve(Solver.IPOPT(show_online_optim=platform.system() == "Linux")) sol.print_cost() From fa1905d7da0226053eb4b0ff42768162926520f0 Mon Sep 17 00:00:00 2001 From: Charbie Date: Mon, 1 Apr 2024 11:56:28 -0400 Subject: [PATCH 4/7] fixed some tests --- bioptim/dynamics/configure_problem.py | 4 +- bioptim/dynamics/dynamics_functions.py | 10 +- .../getting_started/custom_plotting.py | 6 +- .../example_pendulum_time_dependent.py | 7 +- .../track_markers_2D_pendulum.py | 8 +- bioptim/interfaces/acados_interface.py | 9 +- bioptim/limits/constraints.py | 3 +- bioptim/limits/penalty_controller.py | 4 + bioptim/limits/penalty_option.py | 10 +- tests/shard1/test__global_plots.py | 7 +- tests/shard4/test_penalty.py | 133 ++++++++++++------ tests/shard6/test_time_dependent_problems.py | 5 +- 12 files changed, 129 insertions(+), 77 deletions(-) diff --git a/bioptim/dynamics/configure_problem.py b/bioptim/dynamics/configure_problem.py index 5ed98035a..b5d9de993 100644 --- a/bioptim/dynamics/configure_problem.py +++ b/bioptim/dynamics/configure_problem.py @@ -1114,7 +1114,7 @@ def configure_contact_function(ocp, nlp, dyn_func: Callable, **extra_params): nlp.plot["contact_forces"] = CustomPlot( lambda t0, phases_dt, node_idx, x, u, p, a, d: nlp.contact_forces_func( - [t0, t0 + phases_dt[nlp.phase_idx]], x, u, p, a + [t0, t0 + phases_dt[nlp.phase_idx]], x, u, p, a, d ), plot_type=PlotType.INTEGRATED, axes_idx=axes_idx, @@ -1187,7 +1187,7 @@ def configure_soft_contact_function(ocp, nlp): ) nlp.plot[f"soft_contact_forces_{nlp.model.soft_contact_names[i_sc]}"] = CustomPlot( lambda t0, phases_dt, node_idx, x, u, p, a, d: nlp.soft_contact_forces_func( - [t0, t0 + phases_dt[nlp.phase_idx]], x, u, p, a + [t0, t0 + phases_dt[nlp.phase_idx]], x, u, p, a, d )[(i_sc * 6) : ((i_sc + 1) * 6), :], plot_type=PlotType.INTEGRATED, axes_idx=phase_mappings, diff --git a/bioptim/dynamics/dynamics_functions.py b/bioptim/dynamics/dynamics_functions.py index 695f25de9..6e6b9b0f7 100644 --- a/bioptim/dynamics/dynamics_functions.py +++ b/bioptim/dynamics/dynamics_functions.py @@ -175,7 +175,7 @@ def torque_driven( ): if not with_contact and fatigue is None: qddot = DynamicsFunctions.get(nlp.states_dot["qddot"], nlp.states_dot.scaled.mx_reduced) - tau_id = DynamicsFunctions.inverse_dynamics(nlp, q, qdot, qddot, with_contact) + tau_id = DynamicsFunctions.inverse_dynamics(nlp, q, qdot, qddot, with_contact, external_forces) defects = MX(dq.shape[0] + tau_id.shape[0], tau_id.shape[1]) dq_defects = [] @@ -1166,7 +1166,7 @@ def inverse_dynamics( Torques in tau """ - if nlp.external_forces is None: + if external_forces is None: tau = nlp.model.inverse_dynamics(q, qdot, qddot) else: if "tau" in nlp.states: @@ -1176,9 +1176,7 @@ def inverse_dynamics( else: tau_shape = nlp.model.nb_tau tau = MX(tau_shape, nlp.ns) - # Todo: Should be added to pass f_ext in controls (as a symoblic value) - # this would avoid to create multiple equations of motions per node - for i, f_ext in enumerate(nlp.external_forces): + for i, f_ext in enumerate(external_forces): tau[:, i] = nlp.model.inverse_dynamics(q, qdot, qddot, f_ext) return tau # We ignore on purpose the mapping to keep zeros in the defects of the dynamic. @@ -1250,7 +1248,7 @@ def holonomic_torque_driven( external_forces: list = None, ) -> DynamicsEvaluation: """ - The custom dynamics function that provides the derivative of the states: dxdt = f(t, x, u, p, a) + The custom dynamics function that provides the derivative of the states: dxdt = f(t, x, u, p, a, d) Parameters ---------- diff --git a/bioptim/examples/getting_started/custom_plotting.py b/bioptim/examples/getting_started/custom_plotting.py index 1238a1e52..98b92b733 100644 --- a/bioptim/examples/getting_started/custom_plotting.py +++ b/bioptim/examples/getting_started/custom_plotting.py @@ -95,18 +95,18 @@ def prepare_ocp( # Add my lovely new plot ocp.add_plot( "My New Extra Plot", - lambda t0, phases_dt, node_idx, x, u, p, a: custom_plot_callback(x, [0, 1, 3]), + lambda t0, phases_dt, node_idx, x, u, p, a, d: custom_plot_callback(x, [0, 1, 3]), plot_type=PlotType.PLOT, ) ocp.add_plot( # This one combines to the previous one as they have the same name "My New Extra Plot", - lambda t0, phases_dt, node_idx, x, u, p, a: custom_plot_callback(x, [1, 3]), + lambda t0, phases_dt, node_idx, x, u, p, a, d: custom_plot_callback(x, [1, 3]), plot_type=PlotType.STEP, axes_idx=[1, 2], ) ocp.add_plot( "My Second New Extra Plot", - lambda t0, phases_dt, node_idx, x, u, p, a: custom_plot_callback(x, [2, 1]), + lambda t0, phases_dt, node_idx, x, u, p, a, d: custom_plot_callback(x, [2, 1]), plot_type=PlotType.INTEGRATED, axes_idx=[0, 2], ) diff --git a/bioptim/examples/torque_driven_ocp/example_pendulum_time_dependent.py b/bioptim/examples/torque_driven_ocp/example_pendulum_time_dependent.py index c23e857c2..85e8b7e12 100644 --- a/bioptim/examples/torque_driven_ocp/example_pendulum_time_dependent.py +++ b/bioptim/examples/torque_driven_ocp/example_pendulum_time_dependent.py @@ -39,10 +39,11 @@ def time_dependent_dynamic( controls: MX | SX, parameters: MX | SX, algebraic_states: MX | SX, + dynamics_constants: MX | SX, nlp: NonLinearProgram, ) -> DynamicsEvaluation: """ - The custom dynamics function that provides the derivative of the states: dxdt = f(t, x, u, p, a) + The custom dynamics function that provides the derivative of the states: dxdt = f(t, x, u, p, a, d) Parameters ---------- @@ -55,7 +56,9 @@ def time_dependent_dynamic( parameters: MX | SX The parameters acting on the system algebraic_states: MX | SX - The Algebraic states variables of the system + The algebraic states variables of the system + dynamics_constants: MX | SX + The dynamics constants of the system nlp: NonLinearProgram A reference to the phase diff --git a/bioptim/examples/torque_driven_ocp/track_markers_2D_pendulum.py b/bioptim/examples/torque_driven_ocp/track_markers_2D_pendulum.py index bc364a17e..9837d4d45 100644 --- a/bioptim/examples/torque_driven_ocp/track_markers_2D_pendulum.py +++ b/bioptim/examples/torque_driven_ocp/track_markers_2D_pendulum.py @@ -191,14 +191,14 @@ def main(): ocp.add_plot( "Markers plot coordinates", - update_function=lambda t0, phases_dt, node_idx, x, u, p, a: get_markers_pos(x, 0, markers_fun, n_q), + update_function=lambda t0, phases_dt, node_idx, x, u, p, a, d: get_markers_pos(x, 0, markers_fun, n_q), linestyle=".-", plot_type=PlotType.STEP, color=marker_color[0], ) ocp.add_plot( "Markers plot coordinates", - update_function=lambda t0, phases_dt, node_idx, x, u, p, a: get_markers_pos(x, 1, markers_fun, n_q), + update_function=lambda t0, phases_dt, node_idx, x, u, p, a, d: get_markers_pos(x, 1, markers_fun, n_q), linestyle=".-", plot_type=PlotType.STEP, color=marker_color[1], @@ -206,14 +206,14 @@ def main(): ocp.add_plot( "Markers plot coordinates", - update_function=lambda t0, phases_dt, node_idx, x, u, p, a: markers_ref[:, 0, :], + update_function=lambda t0, phases_dt, node_idx, x, u, p, a, d: markers_ref[:, 0, :], plot_type=PlotType.PLOT, color=marker_color[0], legend=title_markers, ) ocp.add_plot( "Markers plot coordinates", - update_function=lambda t0, phases_dt, node_idx, x, u, p, a: markers_ref[:, 1, :], + update_function=lambda t0, phases_dt, node_idx, x, u, p, a, d: markers_ref[:, 1, :], plot_type=PlotType.PLOT, color=marker_color[1], legend=title_markers, diff --git a/bioptim/interfaces/acados_interface.py b/bioptim/interfaces/acados_interface.py index 746f14ebc..3d96fb271 100644 --- a/bioptim/interfaces/acados_interface.py +++ b/bioptim/interfaces/acados_interface.py @@ -167,6 +167,7 @@ def __acados_export_model(self, ocp): p_sym = ocp.nlp[0].parameters.scaled.cx a = ocp.nlp[0].algebraic_states.cx_start a_sym = ocp.nlp[0].algebraic_states.scaled.cx_start + d = ocp.nlp[0].dynamics_constants.cx_start if ocp.parameters: for key in ocp.parameters: @@ -180,7 +181,7 @@ def __acados_export_model(self, ocp): x_sym = vertcat(p_sym, x_sym) x_dot_sym = SX.sym("x_dot", x_sym.shape[0], x_sym.shape[1]) - f_expl = vertcat([0] * self.nparams, ocp.nlp[0].dynamics_func(t, x, u, p, a)) + f_expl = vertcat([0] * self.nparams, ocp.nlp[0].dynamics_func(t, x, u, p, a, d)) f_impl = x_dot_sym - f_expl self.acados_model.f_impl_expr = f_impl @@ -495,14 +496,14 @@ def _adjust_dim(): else: raise RuntimeError(f"{objectives.type.name} is an incompatible objective term with LINEAR_LS cost type") - def add_nonlinear_ls_lagrange(acados, objectives, t, dt, x, u, p, a): + def add_nonlinear_ls_lagrange(acados, objectives, t, dt, x, u, p, a, d): if objectives.function[0].size_in("x")[0] == x.shape[0] * 2: x = vertcat(x, x) if objectives.function[0].size_in("u")[0] == u.shape[0] * 2: u = vertcat(u, u) acados.lagrange_costs = vertcat( - acados.lagrange_costs, objectives.function[0](t, dt, x, u, p, a).reshape((-1, 1)) + acados.lagrange_costs, objectives.function[0](t, dt, x, u, p, a, d).reshape((-1, 1)) ) acados.W = linalg.block_diag(acados.W, np.diag([objectives.weight] * objectives.function[0].numel_out())) @@ -512,7 +513,7 @@ def add_nonlinear_ls_lagrange(acados, objectives, t, dt, x, u, p, a): else: acados.y_ref.append([np.zeros((objectives.function[0].numel_out(), 1)) for _ in node_idx]) - def add_nonlinear_ls_mayer(acados, objectives, t, dt, x, u, p, a, node=None): + def add_nonlinear_ls_mayer(acados, objectives, t, dt, x, u, p, a, d, node=None): if objectives.node[0] not in [Node.INTERMEDIATES, Node.PENULTIMATE, Node.END]: acados.W_0 = linalg.block_diag( acados.W_0, np.diag([objectives.weight] * objectives.function[0].numel_out()) diff --git a/bioptim/limits/constraints.py b/bioptim/limits/constraints.py index 0ec9968b2..e59d6773c 100644 --- a/bioptim/limits/constraints.py +++ b/bioptim/limits/constraints.py @@ -928,8 +928,9 @@ def collocation_jacobians(penalty, controller): vertcat(controller.t_span.cx), horzcat(controller.states_scaled.cx, horzcat(*controller.states_scaled.cx_intermediates_list)), controller.controls_scaled.cx, - controller.parameters.cx, # TODO: fix parameter scaling + controller.parameters_scaled.cx, controller.algebraic_states_scaled.cx, + controller.dynamics_constants.cx, ) initial_defect = controller.states_scaled.cx_start - controller.states_scaled.cx_intermediates_list[0] diff --git a/bioptim/limits/penalty_controller.py b/bioptim/limits/penalty_controller.py index f32cc6684..6a4a66058 100644 --- a/bioptim/limits/penalty_controller.py +++ b/bioptim/limits/penalty_controller.py @@ -30,6 +30,7 @@ def __init__( p: MX | SX | list, a: list, a_scaled: list, + d: list, node_index: int = None, ): """ @@ -55,6 +56,8 @@ def __init__( References to the algebraic_states variables a_scaled: list References to the scaled algebraic_states variables + d: list + References to the dynamics constants node_index: int Current node index if nlp.phase_dynamics is SHARED_DURING_THE_PHASE, then node_index is expected to be set to 0 @@ -70,6 +73,7 @@ def __init__( self.a = a self.a_scaled = a_scaled self.p = vertcat(p) if p is not None else p + self.d = d self.node_index = node_index self.cx_index_to_get = 0 diff --git a/bioptim/limits/penalty_option.py b/bioptim/limits/penalty_option.py index b04e1f95d..50c0dae13 100644 --- a/bioptim/limits/penalty_option.py +++ b/bioptim/limits/penalty_option.py @@ -300,11 +300,11 @@ def _check_target_dimensions(self, controller: PenaltyController): def transform_penalty_to_stochastic(self, controller: PenaltyController, fcn, state_cx_scaled): """ Transform the penalty fcn into the variation of fcn depending on the noise: - fcn = fcn(x, u, p, a) becomes d/dx(fcn) * covariance * d/dx(fcn).T + fcn = fcn(x, u, p, a, d) becomes d/dx(fcn) * covariance * d/dx(fcn).T - Please note that this is usually used to add a buffer around an equality constraint h(x, u, p, a) = 0 + Please note that this is usually used to add a buffer around an equality constraint h(x, u, p, a, d) = 0 transforming it into an inequality constraint of the form: - h(x, u, p, a) + sqrt(dh/dx * covariance * dh/dx.T) <= 0 + h(x, u, p, a, d) + sqrt(dh/dx * covariance * dh/dx.T) <= 0 Here, we chose a different implementation to avoid the discontinuity of the sqrt, we instead decompose the two terms, meaning that you have to declare the constraint h=0 and the "variation of h"=buffer ** 2 with @@ -525,7 +525,7 @@ def _set_penalty_function(self, controllers: list[PenaltyController], fcn: MX | name, [time, phases_dt, x, u, p, a, dynamics_constants], [(func_at_start + func_at_end) / 2], - ["t", "dt", "x", "u", "p", "a"], + ["t", "dt", "x", "u", "p", "a", "dynamics_constants"], ["val"], ) elif self.derivative: @@ -828,7 +828,7 @@ def _finish_add_target_to_plot(self, controller: PenaltyController): """ - def plot_function(t0, phases_dt, node_idx, x, u, p, a, penalty=None): + def plot_function(t0, phases_dt, node_idx, x, u, p, a, d, penalty=None): if isinstance(node_idx, (list, tuple)): return self.target_to_plot[:, [self.node_idx.index(idx) for idx in node_idx]] else: diff --git a/tests/shard1/test__global_plots.py b/tests/shard1/test__global_plots.py index f99bc6707..b9110ae84 100644 --- a/tests/shard1/test__global_plots.py +++ b/tests/shard1/test__global_plots.py @@ -158,7 +158,7 @@ def test_add_new_plot(phase_dynamics): sol = ocp.solve(solver) # Test 1 - Working plot - ocp.add_plot("My New Plot", lambda t0, phases_dt, node_idx, x, u, p, a: x[0:2, :]) + ocp.add_plot("My New Plot", lambda t0, phases_dt, node_idx, x, u, p, a, d: x[0:2, :]) sol.graphs(automatically_organize=False) # Add the plot of objectives and constraints to this mess @@ -258,16 +258,17 @@ def override_penalty(pen: list[PenaltyOption]): u = MX.sym("u", 3, 1) param = MX.sym("param", *p.weighted_function[node_index].size_in("p")) a = MX.sym("a", *p.weighted_function[node_index].size_in("a")) + d = MX.sym("dynamics_constants", *p.weighted_function[node_index].size_in("dynamics_constants")) weight = MX.sym("weight", *p.weighted_function[node_index].size_in("weight")) target = MX.sym("target", *p.weighted_function[node_index].size_in("target")) p.function[node_index] = Function( - name, [t, phases_dt, x, u, param, a], [np.array([range(cmp, len(p.rows) + cmp)]).T] + name, [t, phases_dt, x, u, param, a, d], [np.array([range(cmp, len(p.rows) + cmp)]).T] ) p.function_non_threaded[node_index] = p.function[node_index] p.weighted_function[node_index] = Function( name, - [t, phases_dt, x, u, param, a, weight, target], + [t, phases_dt, x, u, param, a, d, weight, target], [np.array([range(cmp + 1, len(p.rows) + cmp + 1)]).T], ) p.weighted_function_non_threaded[node_index] = p.weighted_function[node_index] diff --git a/tests/shard4/test_penalty.py b/tests/shard4/test_penalty.py index 4b3bd8316..df1f54931 100644 --- a/tests/shard4/test_penalty.py +++ b/tests/shard4/test_penalty.py @@ -86,13 +86,13 @@ def prepare_test_ocp( return ocp -def get_penalty_value(ocp, penalty, t, phases_dt, x, u, p, a): +def get_penalty_value(ocp, penalty, t, phases_dt, x, u, p, a, d): if isinstance(penalty, MultinodeConstraint) or isinstance(penalty, MultinodeObjective): controller = [ - PenaltyController(ocp, ocp.nlp[0], t, x, u, [], [], p, a, [], 0) for i in range(len(penalty.nodes_phase)) + PenaltyController(ocp, ocp.nlp[0], t, x, u, [], [], p, a, [], d, 0) for i in range(len(penalty.nodes_phase)) ] else: - controller = PenaltyController(ocp, ocp.nlp[0], t, x, u, [], [], p, a, [], 0) + controller = PenaltyController(ocp, ocp.nlp[0], t, x, u, [], [], p, a, [], d, 0) val = penalty.type(penalty, controller, **penalty.extra_parameters) # changed only this one if isinstance(val, float): @@ -133,11 +133,12 @@ def test_penalty_minimize_time(penalty_origin, value, phase_dynamics): u = [0] p = [1] a = [] + d = [] penalty_type = penalty_origin.MINIMIZE_TIME penalty = Objective(penalty_type) - penalty_type(penalty, PenaltyController(ocp, ocp.nlp[0], [], [], [], [], [], p, a, [], 0)) - res = get_penalty_value(ocp, penalty, t, phases_dt, x, u, p, a) + penalty_type(penalty, PenaltyController(ocp, ocp.nlp[0], [], [], [], [], [], p, a, [], d, 0)) + res = get_penalty_value(ocp, penalty, t, phases_dt, x, u, p, a, d) if penalty_origin == ObjectiveFcn.Lagrange: np.testing.assert_almost_equal(res, np.array(1)) @@ -156,9 +157,10 @@ def test_penalty_minimize_state(penalty_origin, value, phase_dynamics): u = [0] p = [] a = [] + d = [] penalty = Objective(penalty_origin.MINIMIZE_STATE, key="qdot") - res = get_penalty_value(ocp, penalty, t, phases_dt, x, u, p, a) + res = get_penalty_value(ocp, penalty, t, phases_dt, x, u, p, a, d) np.testing.assert_almost_equal(res, np.array([[value]] * 4)) @@ -173,8 +175,10 @@ def test_penalty_minimize_joint_power(penalty_origin, value, phase_dynamics): u = [1] p = [] a = [] + d = [] + penalty = Objective(penalty_origin.MINIMIZE_POWER, key_control="tau") - res = get_penalty_value(ocp, penalty, t, phases_dt, x, u, p, a) + res = get_penalty_value(ocp, penalty, t, phases_dt, x, u, p, a, d) np.testing.assert_almost_equal(res, np.array([[value]] * 4)) @@ -189,9 +193,10 @@ def test_penalty_minimize_muscle_power(penalty_origin, value, phase_dynamics): u = [DM.ones((8, 1)) * value] p = [] a = [] + d = [] penalty = Objective(penalty_origin.MINIMIZE_POWER, key_control="muscles") - res = get_penalty_value(ocp, penalty, t, phases_dt, x, u, p, a) + res = get_penalty_value(ocp, penalty, t, phases_dt, x, u, p, a, d) if value == 0.1: np.testing.assert_almost_equal( res, np.array([[0.00475812, -0.00505504, -0.000717714, 0.00215864, 0.00215864, -0.00159915]]).T @@ -213,6 +218,7 @@ def test_penalty_minimize_qddot(penalty_origin, value, phase_dynamics): u = [DM.ones((4, 1)) * value] p = [] a = [] + d = [] if penalty_origin == ConstraintFcn: with pytest.raises(AttributeError, match="MINIMIZE_QDDOT"): @@ -221,7 +227,7 @@ def test_penalty_minimize_qddot(penalty_origin, value, phase_dynamics): else: penalty_type = penalty_origin.MINIMIZE_QDDOT penalty = Objective(penalty_type) - res = get_penalty_value(ocp, penalty, t, phases_dt, x, u, p, a).T + res = get_penalty_value(ocp, penalty, t, phases_dt, x, u, p, a, d).T np.testing.assert_almost_equal(res, [[value, -9.81 + value, value, value]]) @@ -237,13 +243,14 @@ def test_penalty_track_state(penalty_origin, value, phase_dynamics): u = [0] p = [] a = [] + d = [] penalty_type = penalty_origin.TRACK_STATE if isinstance(penalty_type, (ObjectiveFcn.Lagrange, ObjectiveFcn.Mayer)): penalty = Objective(penalty_type, key="qdot", target=np.ones((4, 1)) * value) else: penalty = Constraint(penalty_type, key="qdot", target=np.ones((4, 1)) * value) - res = get_penalty_value(ocp, penalty, t, phases_dt, x, u, p, a) + res = get_penalty_value(ocp, penalty, t, phases_dt, x, u, p, a, d) np.testing.assert_almost_equal(res, [[value]] * 4) @@ -258,12 +265,14 @@ def test_penalty_track_joint_power(penalty_origin, value, phase_dynamics): u = [1] p = [] a = [] + d = [] + penalty_type = penalty_origin.TRACK_POWER if isinstance(penalty_type, (ObjectiveFcn.Lagrange, ObjectiveFcn.Mayer)): penalty = Objective(penalty_type, key_control="tau") else: penalty = Constraint(penalty_type, key_control="tau") - res = get_penalty_value(ocp, penalty, t, phases_dt, x, u, p, a) + res = get_penalty_value(ocp, penalty, t, phases_dt, x, u, p, a, d) np.testing.assert_almost_equal(res, [[value]] * 4) @@ -278,10 +287,11 @@ def test_penalty_minimize_markers(penalty_origin, value, phase_dynamics): u = [0] p = [] a = [] + d = [] penalty_type = penalty_origin.MINIMIZE_MARKERS penalty = Objective(penalty_type) - res = get_penalty_value(ocp, penalty, t, phases_dt, x, u, p, a) + res = get_penalty_value(ocp, penalty, t, phases_dt, x, u, p, a, d) expected = np.array( [ @@ -313,6 +323,7 @@ def test_penalty_track_markers(penalty_origin, value, phase_dynamics): u = [0] p = [] a = [] + d = [] penalty_type = penalty_origin.TRACK_MARKERS @@ -320,7 +331,7 @@ def test_penalty_track_markers(penalty_origin, value, phase_dynamics): penalty = Objective(penalty_type, target=np.ones((3, 7, 1)) * value) else: penalty = Constraint(penalty_type, target=np.ones((3, 7, 1)) * value) - res = get_penalty_value(ocp, penalty, t, phases_dt, x, u, p, a) + res = get_penalty_value(ocp, penalty, t, phases_dt, x, u, p, a, d) expected = np.array( [ @@ -352,10 +363,11 @@ def test_penalty_minimize_markers_velocity(penalty_origin, value, phase_dynamics u = [0] p = [] a = [] + d = [] penalty_type = penalty_origin.MINIMIZE_MARKERS_VELOCITY penalty = Objective(penalty_type) - res = get_penalty_value(ocp, penalty, t, phases_dt, x, u, p, a) + res = get_penalty_value(ocp, penalty, t, phases_dt, x, u, p, a, d) if value == 0.1: np.testing.assert_almost_equal( @@ -394,6 +406,8 @@ def test_penalty_minimize_markers_acceleration(penalty_origin, implicit, value, u = [0] p = [0] a = [] + d = [] + penalty_type = penalty_origin.MINIMIZE_MARKERS_ACCELERATION if isinstance(penalty_type, (ObjectiveFcn.Lagrange, ObjectiveFcn.Mayer)): @@ -402,7 +416,7 @@ def test_penalty_minimize_markers_acceleration(penalty_origin, implicit, value, penalty = Constraint(penalty_type) if not implicit: - res = get_penalty_value(ocp, penalty, t, phases_dt, x, u, p, a) + res = get_penalty_value(ocp, penalty, t, phases_dt, x, u, p, a, d) expected = np.array( [ @@ -422,7 +436,7 @@ def test_penalty_minimize_markers_acceleration(penalty_origin, implicit, value, np.testing.assert_almost_equal(res, expected, decimal=5) else: - res = get_penalty_value(ocp, penalty, t, phases_dt, x, u, p, a) + res = get_penalty_value(ocp, penalty, t, phases_dt, x, u, p, a, d) expected = np.array( [ @@ -457,6 +471,7 @@ def test_penalty_track_markers_velocity(penalty_origin, value, phase_dynamics): u = [0] p = [] a = [] + d = [] penalty_type = penalty_origin.TRACK_MARKERS_VELOCITY @@ -464,7 +479,7 @@ def test_penalty_track_markers_velocity(penalty_origin, value, phase_dynamics): penalty = Objective(penalty_type, target=np.ones((3, 7, 1)) * value) else: penalty = Constraint(penalty_type, target=np.ones((3, 7, 1)) * value) - res = get_penalty_value(ocp, penalty, t, phases_dt, x, u, p, a) + res = get_penalty_value(ocp, penalty, t, phases_dt, x, u, p, a, d) if value == 0.1: np.testing.assert_almost_equal( @@ -564,6 +579,7 @@ def test_penalty_track_super_impose_marker(penalty_origin, value, phase_dynamics u = [0] p = [] a = [] + d = [] penalty_type = penalty_origin.SUPERIMPOSE_MARKERS @@ -571,7 +587,7 @@ def test_penalty_track_super_impose_marker(penalty_origin, value, phase_dynamics penalty = Objective(penalty_type, first_marker=0, second_marker=1) else: penalty = Constraint(penalty_type, first_marker=0, second_marker=1) - res = get_penalty_value(ocp, penalty, t, phases_dt, x, u, p, a) + res = get_penalty_value(ocp, penalty, t, phases_dt, x, u, p, a, d) expected = [[0.8951707, 0, -1.0948376]] if value == 0.1 else [[-1.3830926, 0, 0.2950504]] np.testing.assert_almost_equal(res.T, expected) @@ -588,6 +604,7 @@ def test_penalty_track_super_impose_marker_velocity(penalty_origin, value, phase u = [0] p = [] a = [] + d = [] penalty_type = penalty_origin.SUPERIMPOSE_MARKERS_VELOCITY @@ -595,7 +612,7 @@ def test_penalty_track_super_impose_marker_velocity(penalty_origin, value, phase penalty = Objective(penalty_type, first_marker=0, second_marker=1) else: penalty = Constraint(penalty_type, first_marker=0, second_marker=1) - res = get_penalty_value(ocp, penalty, t, phases_dt, x, u, p, a) + res = get_penalty_value(ocp, penalty, t, phases_dt, x, u, p, a, d) expected = [[-0.1094838, 0.0, -0.0895171]] if value == 0.1 else [[-2.9505042, 0.0, -13.8309264]] np.testing.assert_almost_equal(res.T, expected) @@ -613,6 +630,7 @@ def test_penalty_proportional_state(penalty_origin, value, value_intercept, phas u = [0] p = [] a = [] + d = [] penalty_type = penalty_origin.PROPORTIONAL_STATE @@ -636,7 +654,7 @@ def test_penalty_proportional_state(penalty_origin, value, value_intercept, phas first_dof_intercept=value_intercept, second_dof_intercept=value_intercept, ) - res = get_penalty_value(ocp, penalty, t, phases_dt, x, u, p, a) + res = get_penalty_value(ocp, penalty, t, phases_dt, x, u, p, a, d) if value_intercept == 0.0: np.testing.assert_almost_equal(res, -value) @@ -658,6 +676,7 @@ def test_penalty_proportional_control(penalty_origin, value, phase_dynamics): u = [DM.ones((4, 1)) * value] p = [] a = [] + d = [] penalty_type = penalty_origin.PROPORTIONAL_CONTROL @@ -669,7 +688,7 @@ def test_penalty_proportional_control(penalty_origin, value, phase_dynamics): penalty = Objective(penalty_type, key="tau", first_dof=first, second_dof=second, coef=coef) else: penalty = Constraint(penalty_type, key="tau", first_dof=first, second_dof=second, coef=coef) - res = get_penalty_value(ocp, penalty, t, phases_dt, x, u, p, a) + res = get_penalty_value(ocp, penalty, t, phases_dt, x, u, p, a, d) np.testing.assert_almost_equal(res, np.array(u[0][first] - coef * u[0][second])) @@ -685,9 +704,10 @@ def test_penalty_minimize_torque(penalty_origin, value, phase_dynamics): u = [DM.ones((4, 1)) * value] p = [] a = [] + d = [] penalty = Objective(penalty_origin.MINIMIZE_CONTROL, key="tau") - res = get_penalty_value(ocp, penalty, t, phases_dt, x, u, p, a) + res = get_penalty_value(ocp, penalty, t, phases_dt, x, u, p, a, d) np.testing.assert_almost_equal(res, np.array([[value, value, value, value]]).T) @@ -703,6 +723,7 @@ def test_penalty_track_torque(penalty_origin, value, phase_dynamics): u = [DM.ones((4, 1)) * value] p = [] a = [] + d = [] penalty_type = penalty_origin.TRACK_CONTROL @@ -710,7 +731,7 @@ def test_penalty_track_torque(penalty_origin, value, phase_dynamics): penalty = Objective(penalty_type, key="tau", target=np.ones((4, 1)) * value) else: penalty = Constraint(penalty_type, key="tau", target=np.ones((4, 1)) * value) - res = get_penalty_value(ocp, penalty, t, phases_dt, x, u, p, a) + res = get_penalty_value(ocp, penalty, t, phases_dt, x, u, p, a, d) np.testing.assert_almost_equal(res, np.array([[value, value, value, value]]).T) @@ -726,10 +747,11 @@ def test_penalty_minimize_muscles_control(penalty_origin, value, phase_dynamics) u = [DM.ones((8, 1)) * value] p = [] a = [] + d = [] penalty_type = penalty_origin.MINIMIZE_CONTROL penalty = Objective(penalty_type, key="muscles") - res = get_penalty_value(ocp, penalty, t, phases_dt, x, u, p, a) + res = get_penalty_value(ocp, penalty, t, phases_dt, x, u, p, a, d) np.testing.assert_almost_equal(res, np.array([[value, value, value, value, value, value]]).T) @@ -745,10 +767,11 @@ def test_penalty_minimize_contact_forces(penalty_origin, value, phase_dynamics): u = [DM.ones((4, 1)) * value] p = [] a = [] + d = [] penalty_type = penalty_origin.MINIMIZE_CONTACT_FORCES penalty = Objective(penalty_type) - res = get_penalty_value(ocp, penalty, t, phases_dt, x, u, p, a) + res = get_penalty_value(ocp, penalty, t, phases_dt, x, u, p, a, d) if value == 0.1: np.testing.assert_almost_equal(res, np.array([[-9.6680105, 127.2360329, 5.0905995]]).T) @@ -767,6 +790,7 @@ def test_penalty_track_contact_forces(penalty_origin, value, phase_dynamics): u = [DM.ones((4, 1)) * value] p = [] a = [] + d = [] penalty_type = penalty_origin.TRACK_CONTACT_FORCES @@ -774,7 +798,7 @@ def test_penalty_track_contact_forces(penalty_origin, value, phase_dynamics): penalty = Objective(penalty_type, target=np.ones((1, 1)) * value, index=0) else: penalty = Constraint(penalty_type, target=np.ones((1, 1)) * value, index=0) - res = get_penalty_value(ocp, penalty, t, phases_dt, x, u, p, a) + res = get_penalty_value(ocp, penalty, t, phases_dt, x, u, p, a, d) if value == 0.1: np.testing.assert_almost_equal(res.T, [[-9.6680105, 127.2360329, 5.0905995]]) @@ -792,10 +816,11 @@ def test_penalty_minimize_predicted_com_height(value, phase_dynamics): u = [0] p = [] a = [] + d = [] penalty_type = ObjectiveFcn.Mayer.MINIMIZE_PREDICTED_COM_HEIGHT penalty = Objective(penalty_type) - res = get_penalty_value(ocp, penalty, t, phases_dt, x, u, p, a) + res = get_penalty_value(ocp, penalty, t, phases_dt, x, u, p, a, d) expected = np.array(0.0501274 if value == 0.1 else -3.72579) np.testing.assert_almost_equal(res, expected) @@ -812,6 +837,7 @@ def test_penalty_minimize_com_position(value, penalty_origin, phase_dynamics): u = [0] p = [] a = [] + d = [] if "TRACK_COM_POSITION" in penalty_origin._member_names_: penalty_type = penalty_origin.TRACK_COM_POSITION @@ -822,7 +848,7 @@ def test_penalty_minimize_com_position(value, penalty_origin, phase_dynamics): penalty = Objective(penalty_type) else: penalty = Constraint(penalty_type) - res = get_penalty_value(ocp, penalty, t, phases_dt, x, u, p, a) + res = get_penalty_value(ocp, penalty, t, phases_dt, x, u, p, a, d) expected = np.array([[0.05], [0.05], [0.05]]) if value == -10: @@ -842,6 +868,7 @@ def test_penalty_minimize_angular_momentum(value, penalty_origin, phase_dynamics u = [0] p = [] a = [] + d = [] penalty_type = penalty_origin.MINIMIZE_ANGULAR_MOMENTUM @@ -849,7 +876,7 @@ def test_penalty_minimize_angular_momentum(value, penalty_origin, phase_dynamics penalty = Objective(penalty_type) else: penalty = Constraint(penalty_type) - res = get_penalty_value(ocp, penalty, t, phases_dt, x, u, p, a) + res = get_penalty_value(ocp, penalty, t, phases_dt, x, u, p, a, d) expected = np.array([[-0.005], [0.2], [0.005]]) if value == -10: @@ -870,6 +897,7 @@ def test_penalty_minimize_linear_momentum(value, penalty_origin, use_sx, phase_d u = [0] p = [] a = [] + d = [] penalty_type = penalty_origin.MINIMIZE_LINEAR_MOMENTUM @@ -877,7 +905,7 @@ def test_penalty_minimize_linear_momentum(value, penalty_origin, use_sx, phase_d penalty = Objective(penalty_type) else: penalty = Constraint(penalty_type) - res = get_penalty_value(ocp, penalty, t, phases_dt, x, u, p, a) + res = get_penalty_value(ocp, penalty, t, phases_dt, x, u, p, a, d) expected = np.array([[0.1], [0], [0.1]]) if value == -10: @@ -898,6 +926,7 @@ def test_penalty_minimize_comddot(value, penalty_origin, implicit, phase_dynamic u = [0] p = [] a = [] + d = [] penalty_type = penalty_origin.MINIMIZE_COM_ACCELERATION @@ -907,7 +936,7 @@ def test_penalty_minimize_comddot(value, penalty_origin, implicit, phase_dynamic penalty = Constraint(penalty_type) if not implicit: - res = get_penalty_value(ocp, penalty, t, phases_dt, x, u, p, a) + res = get_penalty_value(ocp, penalty, t, phases_dt, x, u, p, a, d) expected = np.array([[0.0], [-0.7168803], [-0.0740871]]) if value == -10: @@ -935,6 +964,7 @@ def test_penalty_track_segment_with_custom_rt(penalty_origin, value, phase_dynam u = [0] p = [] a = [] + d = [] penalty_type = penalty_origin.TRACK_SEGMENT_WITH_CUSTOM_RT @@ -942,7 +972,7 @@ def test_penalty_track_segment_with_custom_rt(penalty_origin, value, phase_dynam penalty = Objective(penalty_type, segment="ground", rt=0) else: penalty = Constraint(penalty_type, segment="ground", rt=0) - res = get_penalty_value(ocp, penalty, t, phases_dt, x, u, p, a) + res = get_penalty_value(ocp, penalty, t, phases_dt, x, u, p, a, d) expected = np.array([[0], [0.1], [0]]) if value == -10: @@ -962,6 +992,7 @@ def test_penalty_track_marker_with_segment_axis(penalty_origin, value, phase_dyn u = [0] p = [] a = [] + d = [] penalty_type = penalty_origin.TRACK_MARKER_WITH_SEGMENT_AXIS @@ -969,7 +1000,7 @@ def test_penalty_track_marker_with_segment_axis(penalty_origin, value, phase_dyn penalty = Objective(penalty_type, marker="m0", segment="ground", axis=Axis.X) else: penalty = Constraint(penalty_type, marker="m0", segment="ground", axis=Axis.X) - res = get_penalty_value(ocp, penalty, t, phases_dt, x, u, p, a) + res = get_penalty_value(ocp, penalty, t, phases_dt, x, u, p, a, d) expected = [[value, 0, value]] np.testing.assert_almost_equal(res.T, expected) @@ -986,6 +1017,7 @@ def test_penalty_minimize_segment_rotation(penalty_origin, value, phase_dynamics u = [0] p = [] a = [] + d = [] if penalty_origin == ObjectiveFcn.Lagrange or penalty_origin == ObjectiveFcn.Mayer: penalty_type = penalty_origin.MINIMIZE_SEGMENT_ROTATION @@ -993,7 +1025,7 @@ def test_penalty_minimize_segment_rotation(penalty_origin, value, phase_dynamics else: penalty_type = penalty_origin.TRACK_SEGMENT_ROTATION penalty = Constraint(penalty_type, segment=2) - res = get_penalty_value(ocp, penalty, t, phases_dt, x, u, p, a) + res = get_penalty_value(ocp, penalty, t, phases_dt, x, u, p, a, d) expected = [[0, value, 0]] if value == 0.1 else [[3.1415927, 0.575222, 3.1415927]] np.testing.assert_almost_equal(res.T, expected) @@ -1010,6 +1042,7 @@ def test_penalty_minimize_segment_velocity(penalty_origin, value, phase_dynamics u = [0] p = [] a = [] + d = [] if penalty_origin == ObjectiveFcn.Lagrange or penalty_origin == ObjectiveFcn.Mayer: penalty_type = penalty_origin.MINIMIZE_SEGMENT_VELOCITY @@ -1017,7 +1050,7 @@ def test_penalty_minimize_segment_velocity(penalty_origin, value, phase_dynamics else: penalty_type = penalty_origin.TRACK_SEGMENT_VELOCITY penalty = Constraint(penalty_type, segment=2) - res = get_penalty_value(ocp, penalty, t, phases_dt, x, u, p, a) + res = get_penalty_value(ocp, penalty, t, phases_dt, x, u, p, a, d) expected = [[0, value, 0]] np.testing.assert_almost_equal(res.T, expected) @@ -1034,6 +1067,7 @@ def test_penalty_minimize_vector_orientation(penalty_origin, value, phase_dynami u = [0] p = [] a = [] + d = [] penalty_type = penalty_origin.TRACK_VECTOR_ORIENTATIONS_FROM_MARKERS @@ -1054,7 +1088,7 @@ def test_penalty_minimize_vector_orientation(penalty_origin, value, phase_dynami vector_1_marker_1="m6", ) - res = get_penalty_value(ocp, penalty, t, phases_dt, x, u, p, a) + res = get_penalty_value(ocp, penalty, t, phases_dt, x, u, p, a, d) if value == 0.1: np.testing.assert_almost_equal(float(res), 0.09999999999999999) @@ -1073,10 +1107,11 @@ def test_penalty_contact_force_inequality(penalty_origin, value, phase_dynamics) u = [DM.ones((4, 1)) * value] p = [] a = [] + d = [] penalty_type = penalty_origin.TRACK_CONTACT_FORCES penalty = Constraint(penalty_type, contact_index=0) - res = get_penalty_value(ocp, penalty, t, phases_dt, x, u, p, a) + res = get_penalty_value(ocp, penalty, t, phases_dt, x, u, p, a, d) expected = [[-9.6680105, 127.2360329, 5.0905995]] if value == 0.1 else [[25.6627161, 462.7973306, -94.0182191]] np.testing.assert_almost_equal(res.T, expected) @@ -1092,12 +1127,13 @@ def test_penalty_non_slipping(value, phase_dynamics): u = [DM.ones((4, 1)) * value] p = [] a = [] + d = [] penalty_type = ConstraintFcn.NON_SLIPPING penalty = Constraint( penalty_type, tangential_component_idx=0, normal_component_idx=1, static_friction_coefficient=2 ) - res = get_penalty_value(ocp, penalty, t, phases_dt, x, u, p, a) + res = get_penalty_value(ocp, penalty, t, phases_dt, x, u, p, a, d) expected = [[64662.56185612, 64849.5027121]] if value == 0.1 else [[856066.90177734, 857384.05177395]] np.testing.assert_almost_equal(res.T, expected) @@ -1114,15 +1150,16 @@ def test_tau_max_from_actuators(value, threshold, phase_dynamics): u = [DM.ones((3, 1)) * value, DM.ones((3, 1)) * value] p = [] a = [] + d = [] penalty_type = ConstraintFcn.TORQUE_MAX_FROM_Q_AND_QDOT penalty = Constraint(penalty_type, min_torque=threshold) if threshold and threshold < 0: with pytest.raises(ValueError, match="min_torque cannot be negative in tau_max_from_actuators"): - get_penalty_value(ocp, penalty, t, phases_dt, x, u, p, a) + get_penalty_value(ocp, penalty, t, phases_dt, x, u, p, a, d) return else: - res = get_penalty_value(ocp, penalty, t, phases_dt, x, u, p, a) + res = get_penalty_value(ocp, penalty, t, phases_dt, x, u, p, a, d) if threshold: np.testing.assert_almost_equal(res, np.repeat([value + threshold, value - threshold], 3)[:, np.newaxis]) @@ -1140,10 +1177,11 @@ def test_penalty_time_constraint(value, phase_dynamics): u = [0] p = [0] a = [] + d = [] penalty_type = ConstraintFcn.TIME_CONSTRAINT penalty = Constraint(penalty_type) - res = get_penalty_value(ocp, penalty, t, phases_dt, x, u, p, a) + res = get_penalty_value(ocp, penalty, t, phases_dt, x, u, p, a, d) np.testing.assert_almost_equal(res, np.array(0.05) * ocp.nlp[0].ns) @@ -1158,6 +1196,7 @@ def test_penalty_constraint_total_time(value, phase_dynamics): u = [0] p = [0.1] a = [] + d = [] penalty_type = MultinodeConstraintFcn.TRACK_TOTAL_TIME penalty = MultinodeConstraintList() @@ -1173,11 +1212,11 @@ def test_penalty_constraint_total_time(value, phase_dynamics): penalty_type( penalty[0], [ - PenaltyController(ocp, ocp.nlp[0], [], [], [], [], [], p, a, [], 0), - PenaltyController(ocp, ocp.nlp[0], [], [], [], [], [], p, a, [], 0), + PenaltyController(ocp, ocp.nlp[0], [], [], [], [], [], p, a, [], d, 0), + PenaltyController(ocp, ocp.nlp[0], [], [], [], [], [], p, a, [], d, 0), ], ) - res = get_penalty_value(ocp, penalty[0], t, phases_dt, x, u, p, a) + res = get_penalty_value(ocp, penalty[0], t, phases_dt, x, u, p, a, d) np.testing.assert_almost_equal(res, np.array(0.05) * ocp.nlp[0].ns * 2) @@ -1197,6 +1236,7 @@ def custom(controller: PenaltyController, mult): u = [0] p = [] a = [] + d = [] penalty_type = penalty_origin.CUSTOM @@ -1205,7 +1245,7 @@ def custom(controller: PenaltyController, mult): penalty = Objective(custom, index=0, mult=mult, custom_type=penalty_origin) else: penalty = Constraint(custom, index=0, mult=mult) - res = get_penalty_value(ocp, penalty, t, phases_dt, x, u, p, a) + res = get_penalty_value(ocp, penalty, t, phases_dt, x, u, p, a, d) np.testing.assert_almost_equal(res, [[value * mult]] * 4) @@ -1278,9 +1318,10 @@ def custom_with_bounds(controller: PenaltyController): u = [0] p = [] a = [] + d = [] penalty = Constraint(custom_with_bounds) - res = get_penalty_value(ocp, penalty, t, phases_dt, x, u, p, a) + res = get_penalty_value(ocp, penalty, t, phases_dt, x, u, p, a, d) np.testing.assert_almost_equal(res, [[value]] * 4) np.testing.assert_almost_equal(penalty.min_bound, -10) diff --git a/tests/shard6/test_time_dependent_problems.py b/tests/shard6/test_time_dependent_problems.py index 684d2dda0..07a5e38f8 100644 --- a/tests/shard6/test_time_dependent_problems.py +++ b/tests/shard6/test_time_dependent_problems.py @@ -38,10 +38,11 @@ def time_dynamic( controls: MX | SX, parameters: MX | SX, algebraic_states: MX | SX, + dynamics_constants: MX | SX, nlp: NonLinearProgram, ) -> DynamicsEvaluation: """ - The custom dynamics function that provides the derivative of the states: dxdt = f(t, x, u, p, a) + The custom dynamics function that provides the derivative of the states: dxdt = f(t, x, u, p, a, d) Parameters ---------- @@ -55,6 +56,8 @@ def time_dynamic( The parameters acting on the system algebraic_states: MX | SX The Algebraic states variables of the system + dynamics_constants: MX | SX + The dynamics constants of the system nlp: NonLinearProgram A reference to the phase From 9f899097dfd0c71b76d7857d63a0a5227e582308 Mon Sep 17 00:00:00 2001 From: Charbie Date: Mon, 1 Apr 2024 11:57:12 -0400 Subject: [PATCH 5/7] blacked --- bioptim/dynamics/configure_problem.py | 5 +- bioptim/dynamics/dynamics_functions.py | 12 +- bioptim/dynamics/ode_solver.py | 11 +- bioptim/gui/plot.py | 9 +- bioptim/interfaces/solve_ivp_interface.py | 4 +- bioptim/optimization/solution/solution.py | 4 +- tests/shard1/test_dynamics.py | 374 +++++++++++----------- 7 files changed, 222 insertions(+), 197 deletions(-) diff --git a/bioptim/dynamics/configure_problem.py b/bioptim/dynamics/configure_problem.py index b5d9de993..f028e905f 100644 --- a/bioptim/dynamics/configure_problem.py +++ b/bioptim/dynamics/configure_problem.py @@ -134,10 +134,7 @@ def initialize(ocp, nlp): nlp.dynamics_type.type(ocp, nlp, **nlp.dynamics_type.extra_parameters) @staticmethod - def custom(ocp, - nlp, - dynamics_constants_used_at_each_nodes: dict[list] = {}, - **extra_params): + def custom(ocp, nlp, dynamics_constants_used_at_each_nodes: dict[list] = {}, **extra_params): """ Call the user-defined dynamics configuration function diff --git a/bioptim/dynamics/dynamics_functions.py b/bioptim/dynamics/dynamics_functions.py index 6e6b9b0f7..074f72599 100644 --- a/bioptim/dynamics/dynamics_functions.py +++ b/bioptim/dynamics/dynamics_functions.py @@ -44,7 +44,13 @@ class DynamicsFunctions: @staticmethod def custom( - time: MX.sym, states: MX.sym, controls: MX.sym, parameters: MX.sym, algebraic_states: MX.sym, dynamics_constants: MX.sym, nlp + time: MX.sym, + states: MX.sym, + controls: MX.sym, + parameters: MX.sym, + algebraic_states: MX.sym, + dynamics_constants: MX.sym, + nlp, ) -> DynamicsEvaluation: """ Interface to custom dynamic function provided by the user. @@ -74,7 +80,9 @@ def custom( The defects of the implicit dynamics """ - return nlp.dynamics_type.dynamic_function(time, states, controls, parameters, algebraic_states, dynamics_constants, nlp) + return nlp.dynamics_type.dynamic_function( + time, states, controls, parameters, algebraic_states, dynamics_constants, nlp + ) @staticmethod def torque_driven( diff --git a/bioptim/dynamics/ode_solver.py b/bioptim/dynamics/ode_solver.py index 5e0a92a25..9d1582a67 100644 --- a/bioptim/dynamics/ode_solver.py +++ b/bioptim/dynamics/ode_solver.py @@ -174,7 +174,9 @@ def d_ode(self, nlp) -> MX: """ raise RuntimeError("This method should be implemented in the child class") - def initialize_integrator(self, ocp, nlp, dynamics_index: int, node_index: int, is_extra_dynamics: bool = False, **extra_opt) -> Callable: + def initialize_integrator( + self, ocp, nlp, dynamics_index: int, node_index: int, is_extra_dynamics: bool = False, **extra_opt + ) -> Callable: """ Initialize the integrator @@ -596,7 +598,12 @@ def initialize_integrator(self, ocp, nlp, dynamics_index: int, node_index: int, "x": nlp.states.scaled.cx_start, "u": nlp.controls.scaled.cx_start, # todo: add p=parameters "ode": dynamics_func( - vertcat(*t), self.x_ode(nlp), self.p_ode(nlp), self.param_ode(nlp), self.a_ode(nlp), self.d_ode(nlp), + vertcat(*t), + self.x_ode(nlp), + self.p_ode(nlp), + self.param_ode(nlp), + self.a_ode(nlp), + self.d_ode(nlp), ), } diff --git a/bioptim/gui/plot.py b/bioptim/gui/plot.py index a05a0ea4c..bdd120df1 100644 --- a/bioptim/gui/plot.py +++ b/bioptim/gui/plot.py @@ -715,6 +715,7 @@ def update_data( for nlp in self.ocp.nlp: from ..interfaces.interface_utils import _get_dynamics_constants + phase_idx = nlp.phase_idx x_decision = data_states_decision[phase_idx] x_stepwise = data_states_stepwise[phase_idx] @@ -825,9 +826,7 @@ def _compute_y_from_plot_func( lambda p_idx, n_idx, sn_idx: a[n_idx][:, sn_idx] if n_idx < len(a) else np.ndarray((0, 1)), ) d_node = PenaltyHelpers.dynamics_constants( - penalty, - idx, - lambda p_idx, n_idx, sn_idx: _get_dynamics_constants(self.ocp, p_idx, n_idx, sn_idx) + penalty, idx, lambda p_idx, n_idx, sn_idx: _get_dynamics_constants(self.ocp, p_idx, n_idx, sn_idx) ) if d_node.shape == (0, 0): d_node = DM(0, 1) @@ -840,7 +839,9 @@ def _compute_y_from_plot_func( a_node = a[node_idx] d_node = d[node_idx] - tp = custom_plot.function(t0, dt, node_idx, x_node, u_node, p_node, a_node, d_node, **custom_plot.parameters) + tp = custom_plot.function( + t0, dt, node_idx, x_node, u_node, p_node, a_node, d_node, **custom_plot.parameters + ) y_tp = np.ndarray((max(custom_plot.phase_mappings.to_first.map_idx) + 1, tp.shape[1])) * np.nan for ctr, axe_index in enumerate(custom_plot.phase_mappings.to_first.map_idx): diff --git a/bioptim/interfaces/solve_ivp_interface.py b/bioptim/interfaces/solve_ivp_interface.py index 3a7728203..7d9d276fe 100644 --- a/bioptim/interfaces/solve_ivp_interface.py +++ b/bioptim/interfaces/solve_ivp_interface.py @@ -81,7 +81,9 @@ def solve_ivp_interface( result = _solve_ivp_scipy_interface( lambda t, x: np.array( - list_of_dynamics[node](t, x, _control_function(control_type, t, t_span, u[node]), p, a[node], d[node]) + list_of_dynamics[node]( + t, x, _control_function(control_type, t, t_span, u[node]), p, a[node], d[node] + ) )[:, 0], x0=x0i, t_span=np.array(t_span), diff --git a/bioptim/optimization/solution/solution.py b/bioptim/optimization/solution/solution.py index b3f0cbeec..ba33ede13 100644 --- a/bioptim/optimization/solution/solution.py +++ b/bioptim/optimization/solution/solution.py @@ -1007,9 +1007,7 @@ def _states_for_phase_integration( ), ) d_tp = PenaltyHelpers.dynamics_constants( - penalty, - 0, - lambda p, n, sn: _get_dynamics_constants(self.ocp, p, n, sn) + penalty, 0, lambda p, n, sn: _get_dynamics_constants(self.ocp, p, n, sn) ) d = np.array([]) if d_tp.shape == (0, 0) else np.array(d_tp) diff --git a/tests/shard1/test_dynamics.py b/tests/shard1/test_dynamics.py index 4c5580cd4..47d76c997 100644 --- a/tests/shard1/test_dynamics.py +++ b/tests/shard1/test_dynamics.py @@ -70,45 +70,45 @@ def test_torque_driven(with_contact, with_external_force, cx, rigidbody_dynamics if with_external_force: external_forces = np.zeros((6, 1, nlp.ns + 1)) external_forces[:, 0, 0] = [ - 0.374540118847362, - 0.950714306409916, - 0.731993941811405, - 0.598658484197037, - 0.156018640442437, - 0.155994520336203, - ] + 0.374540118847362, + 0.950714306409916, + 0.731993941811405, + 0.598658484197037, + 0.156018640442437, + 0.155994520336203, + ] external_forces[:, 0, 1] = [ - 0.058083612168199, - 0.866176145774935, - 0.601115011743209, - 0.708072577796045, - 0.020584494295802, - 0.969909852161994, - ] + 0.058083612168199, + 0.866176145774935, + 0.601115011743209, + 0.708072577796045, + 0.020584494295802, + 0.969909852161994, + ] external_forces[:, 0, 2] = [ - 0.832442640800422, - 0.212339110678276, - 0.181824967207101, - 0.183404509853434, - 0.304242242959538, - 0.524756431632238, - ] + 0.832442640800422, + 0.212339110678276, + 0.181824967207101, + 0.183404509853434, + 0.304242242959538, + 0.524756431632238, + ] external_forces[:, 0, 3] = [ - 0.431945018642116, - 0.291229140198042, - 0.611852894722379, - 0.139493860652042, - 0.292144648535218, - 0.366361843293692, - ] + 0.431945018642116, + 0.291229140198042, + 0.611852894722379, + 0.139493860652042, + 0.292144648535218, + 0.366361843293692, + ] external_forces[:, 0, 4] = [ - 0.456069984217036, - 0.785175961393014, - 0.19967378215836, - 0.514234438413612, - 0.592414568862042, - 0.046450412719998, - ] + 0.456069984217036, + 0.785175961393014, + 0.19967378215836, + 0.514234438413612, + 0.592414568862042, + 0.046450412719998, + ] ocp = OptimalControlProgram(nlp, use_sx=(True if cx == SX else False)) nlp.control_type = ControlType.CONSTANT @@ -121,7 +121,9 @@ def test_torque_driven(with_contact, with_external_force, cx, rigidbody_dynamics rigidbody_dynamics=rigidbody_dynamics, expand_dynamics=True, phase_dynamics=phase_dynamics, - dynamics_constants_used_at_each_nodes={"external_forces": external_forces} if external_forces is not None else {}, + dynamics_constants_used_at_each_nodes=( + {"external_forces": external_forces} if external_forces is not None else {} + ), ), False, ) @@ -416,45 +418,45 @@ def test_torque_derivative_driven(with_contact, with_external_force, cx, phase_d if with_external_force: external_forces = np.zeros((6, 1, nlp.ns + 1)) external_forces[:, 0, 0] = [ - 0.3745401188473625, - 0.9507143064099162, - 0.7319939418114051, - 0.5986584841970366, - 0.15601864044243652, - 0.15599452033620265, + 0.3745401188473625, + 0.9507143064099162, + 0.7319939418114051, + 0.5986584841970366, + 0.15601864044243652, + 0.15599452033620265, ] external_forces[:, 0, 1] = [ - 0.05808361216819946, - 0.8661761457749352, - 0.6011150117432088, - 0.7080725777960455, - 0.020584494295802447, - 0.9699098521619943, - ] + 0.05808361216819946, + 0.8661761457749352, + 0.6011150117432088, + 0.7080725777960455, + 0.020584494295802447, + 0.9699098521619943, + ] external_forces[:, 0, 2] = [ - 0.8324426408004217, - 0.21233911067827616, - 0.18182496720710062, - 0.18340450985343382, - 0.3042422429595377, - 0.5247564316322378, - ] + 0.8324426408004217, + 0.21233911067827616, + 0.18182496720710062, + 0.18340450985343382, + 0.3042422429595377, + 0.5247564316322378, + ] external_forces[:, 0, 3] = [ - 0.43194501864211576, - 0.2912291401980419, - 0.6118528947223795, - 0.13949386065204183, - 0.29214464853521815, - 0.3663618432936917, - ] + 0.43194501864211576, + 0.2912291401980419, + 0.6118528947223795, + 0.13949386065204183, + 0.29214464853521815, + 0.3663618432936917, + ] external_forces[:, 0, 4] = [ - 0.45606998421703593, - 0.7851759613930136, - 0.19967378215835974, - 0.5142344384136116, - 0.5924145688620425, - 0.046450412719997725, - ] + 0.45606998421703593, + 0.7851759613930136, + 0.19967378215835974, + 0.5142344384136116, + 0.5924145688620425, + 0.046450412719997725, + ] NonLinearProgram.add( ocp, @@ -464,7 +466,9 @@ def test_torque_derivative_driven(with_contact, with_external_force, cx, phase_d with_contact=with_contact, expand_dynamics=True, phase_dynamics=phase_dynamics, - dynamics_constants_used_at_each_nodes={"external_forces": external_forces} if external_forces is not None else {}, + dynamics_constants_used_at_each_nodes=( + {"external_forces": external_forces} if external_forces is not None else {} + ), ), False, ) @@ -891,45 +895,45 @@ def test_torque_activation_driven(with_contact, with_external_force, cx, phase_d if with_external_force: external_forces = np.zeros((6, 1, nlp.ns + 1)) external_forces[:, 0, 0] = [ - 0.3745401188473625, - 0.9507143064099162, - 0.7319939418114051, - 0.5986584841970366, - 0.15601864044243652, - 0.15599452033620265, - ] + 0.3745401188473625, + 0.9507143064099162, + 0.7319939418114051, + 0.5986584841970366, + 0.15601864044243652, + 0.15599452033620265, + ] external_forces[:, 0, 1] = [ - 0.05808361216819946, - 0.8661761457749352, - 0.6011150117432088, - 0.7080725777960455, - 0.020584494295802447, - 0.9699098521619943, - ] + 0.05808361216819946, + 0.8661761457749352, + 0.6011150117432088, + 0.7080725777960455, + 0.020584494295802447, + 0.9699098521619943, + ] external_forces[:, 0, 2] = [ - 0.8324426408004217, - 0.21233911067827616, - 0.18182496720710062, - 0.18340450985343382, - 0.3042422429595377, - 0.5247564316322378, - ] + 0.8324426408004217, + 0.21233911067827616, + 0.18182496720710062, + 0.18340450985343382, + 0.3042422429595377, + 0.5247564316322378, + ] external_forces[:, 0, 3] = [ - 0.43194501864211576, - 0.2912291401980419, - 0.6118528947223795, - 0.13949386065204183, - 0.29214464853521815, - 0.3663618432936917, - ] + 0.43194501864211576, + 0.2912291401980419, + 0.6118528947223795, + 0.13949386065204183, + 0.29214464853521815, + 0.3663618432936917, + ] external_forces[:, 0, 4] = [ - 0.45606998421703593, - 0.7851759613930136, - 0.19967378215835974, - 0.5142344384136116, - 0.5924145688620425, - 0.046450412719997725, - ] + 0.45606998421703593, + 0.7851759613930136, + 0.19967378215835974, + 0.5142344384136116, + 0.5924145688620425, + 0.046450412719997725, + ] ocp = OptimalControlProgram(nlp, use_sx=(True if cx == SX else False)) nlp.control_type = ControlType.CONSTANT @@ -941,7 +945,9 @@ def test_torque_activation_driven(with_contact, with_external_force, cx, phase_d with_contact=with_contact, expand_dynamics=True, phase_dynamics=phase_dynamics, - dynamics_constants_used_at_each_nodes={"external_forces": external_forces} if external_forces is not None else {}, + dynamics_constants_used_at_each_nodes=( + {"external_forces": external_forces} if external_forces is not None else {} + ), ), False, ) @@ -1048,45 +1054,45 @@ def test_torque_activation_driven_with_residual_torque( if with_external_force: external_forces = np.zeros((6, 1, nlp.ns + 1)) external_forces[:, 0, 0] = [ - 0.3745401188473625, - 0.9507143064099162, - 0.7319939418114051, - 0.5986584841970366, - 0.15601864044243652, - 0.15599452033620265, - ] + 0.3745401188473625, + 0.9507143064099162, + 0.7319939418114051, + 0.5986584841970366, + 0.15601864044243652, + 0.15599452033620265, + ] external_forces[:, 0, 1] = [ - 0.05808361216819946, - 0.8661761457749352, - 0.6011150117432088, - 0.7080725777960455, - 0.020584494295802447, - 0.9699098521619943, - ] + 0.05808361216819946, + 0.8661761457749352, + 0.6011150117432088, + 0.7080725777960455, + 0.020584494295802447, + 0.9699098521619943, + ] external_forces[:, 0, 2] = [ - 0.8324426408004217, - 0.21233911067827616, - 0.18182496720710062, - 0.18340450985343382, - 0.3042422429595377, - 0.5247564316322378, - ] + 0.8324426408004217, + 0.21233911067827616, + 0.18182496720710062, + 0.18340450985343382, + 0.3042422429595377, + 0.5247564316322378, + ] external_forces[:, 0, 3] = [ - 0.43194501864211576, - 0.2912291401980419, - 0.6118528947223795, - 0.13949386065204183, - 0.29214464853521815, - 0.3663618432936917, - ] + 0.43194501864211576, + 0.2912291401980419, + 0.6118528947223795, + 0.13949386065204183, + 0.29214464853521815, + 0.3663618432936917, + ] external_forces[:, 0, 4] = [ - 0.45606998421703593, - 0.7851759613930136, - 0.19967378215835974, - 0.5142344384136116, - 0.5924145688620425, - 0.046450412719997725, - ] + 0.45606998421703593, + 0.7851759613930136, + 0.19967378215835974, + 0.5142344384136116, + 0.5924145688620425, + 0.046450412719997725, + ] ocp = OptimalControlProgram(nlp, use_sx=(True if cx == SX else False)) nlp.control_type = ControlType.CONSTANT @@ -1098,7 +1104,9 @@ def test_torque_activation_driven_with_residual_torque( with_residual_torque=with_residual_torque, expand_dynamics=True, phase_dynamics=phase_dynamics, - dynamics_constants_used_at_each_nodes={"external_forces": external_forces} if external_forces is not None else {}, + dynamics_constants_used_at_each_nodes=( + {"external_forces": external_forces} if external_forces is not None else {} + ), ), False, ) @@ -1251,8 +1259,10 @@ def test_muscle_driven( ): # Prepare the program nlp = NonLinearProgram(phase_dynamics=phase_dynamics) - nlp.model = BiorbdModel(TestUtils.bioptim_folder() + "/examples/muscle_driven_ocp/models/arm26_with_contact.bioMod", - segments_to_apply_external_forces=["r_ulna_radius_hand_rotation1"]) + nlp.model = BiorbdModel( + TestUtils.bioptim_folder() + "/examples/muscle_driven_ocp/models/arm26_with_contact.bioMod", + segments_to_apply_external_forces=["r_ulna_radius_hand_rotation1"], + ) nlp.ns = 5 nlp.cx = cx nlp.time_mx = MX.sym("time", 1, 1) @@ -1271,45 +1281,45 @@ def test_muscle_driven( if with_external_force: external_forces = np.zeros((6, 1, nlp.ns + 1)) external_forces[:, 0, 0] = [ - 0.3745401188473625, - 0.9507143064099162, - 0.7319939418114051, - 0.5986584841970366, - 0.15601864044243652, - 0.15599452033620265, - ] + 0.3745401188473625, + 0.9507143064099162, + 0.7319939418114051, + 0.5986584841970366, + 0.15601864044243652, + 0.15599452033620265, + ] external_forces[:, 0, 1] = [ - 0.05808361216819946, - 0.8661761457749352, - 0.6011150117432088, - 0.7080725777960455, - 0.020584494295802447, - 0.9699098521619943, - ] + 0.05808361216819946, + 0.8661761457749352, + 0.6011150117432088, + 0.7080725777960455, + 0.020584494295802447, + 0.9699098521619943, + ] external_forces[:, 0, 2] = [ - 0.8324426408004217, - 0.21233911067827616, - 0.18182496720710062, - 0.18340450985343382, - 0.3042422429595377, - 0.5247564316322378, - ] + 0.8324426408004217, + 0.21233911067827616, + 0.18182496720710062, + 0.18340450985343382, + 0.3042422429595377, + 0.5247564316322378, + ] external_forces[:, 0, 3] = [ - 0.43194501864211576, - 0.2912291401980419, - 0.6118528947223795, - 0.13949386065204183, - 0.29214464853521815, - 0.3663618432936917, - ] + 0.43194501864211576, + 0.2912291401980419, + 0.6118528947223795, + 0.13949386065204183, + 0.29214464853521815, + 0.3663618432936917, + ] external_forces[:, 0, 4] = [ - 0.45606998421703593, - 0.7851759613930136, - 0.19967378215835974, - 0.5142344384136116, - 0.5924145688620425, - 0.046450412719997725, - ] + 0.45606998421703593, + 0.7851759613930136, + 0.19967378215835974, + 0.5142344384136116, + 0.5924145688620425, + 0.046450412719997725, + ] ocp = OptimalControlProgram(nlp, use_sx=(True if cx == SX else False)) nlp.control_type = ControlType.CONSTANT @@ -1324,7 +1334,9 @@ def test_muscle_driven( rigidbody_dynamics=rigidbody_dynamics, expand_dynamics=True, phase_dynamics=phase_dynamics, - dynamics_constants_used_at_each_nodes={"external_forces": external_forces} if external_forces is not None else {}, + dynamics_constants_used_at_each_nodes=( + {"external_forces": external_forces} if external_forces is not None else {} + ), ), False, ) From 29083ab93c5f264f1fddb3f94efdf721f0adbcef Mon Sep 17 00:00:00 2001 From: Charbie Date: Mon, 1 Apr 2024 12:31:55 -0400 Subject: [PATCH 6/7] woups --- bioptim/limits/penalty_option.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/bioptim/limits/penalty_option.py b/bioptim/limits/penalty_option.py index 50c0dae13..16caa7c7e 100644 --- a/bioptim/limits/penalty_option.py +++ b/bioptim/limits/penalty_option.py @@ -1002,4 +1002,5 @@ def get_penalty_controller(self, ocp, nlp) -> PenaltyController: if nlp.A is not None and (not isinstance(nlp.A, list) or nlp.A != []): a = [nlp.A[idx] for idx in t_idx] a_scaled = [nlp.A_scaled[idx] for idx in t_idx] - return PenaltyController(ocp, nlp, t_idx, x, u, x_scaled, u_scaled, nlp.parameters.cx, a, a_scaled) + d = [nlp.dynamics_constants for idx in t_idx] + return PenaltyController(ocp, nlp, t_idx, x, u, x_scaled, u_scaled, nlp.parameters.cx, a, a_scaled, d) From 21cd216574220d1f0e2178e453712ddfabf8e2cb Mon Sep 17 00:00:00 2001 From: Charbie Date: Mon, 1 Apr 2024 21:52:29 -0400 Subject: [PATCH 7/7] fixed some tests --- bioptim/dynamics/configure_problem.py | 34 +++++++++++++++---- bioptim/dynamics/dynamics_functions.py | 1 + .../getting_started/custom_dynamics.py | 2 +- .../muscle_excitations_tracker.py | 1 + .../arm_reaching_muscle_driven.py | 33 ++++++++++++------ ...arm_reaching_torque_driven_collocations.py | 1 + .../arm_reaching_torque_driven_explicit.py | 22 +++++++----- .../arm_reaching_torque_driven_implicit.py | 2 ++ .../stochastic_optimal_control/common.py | 4 +-- bioptim/interfaces/acados_interface.py | 1 + bioptim/limits/constraints.py | 28 +++++++++++---- bioptim/limits/multinode_penalty.py | 8 +++++ bioptim/limits/penalty.py | 3 ++ bioptim/limits/penalty_helpers.py | 21 +++++++----- bioptim/limits/penalty_option.py | 6 ++-- .../models/biorbd/stochastic_biorbd_model.py | 4 +-- bioptim/optimization/solution/solution.py | 16 +++++++-- .../stochastic_optimal_control_program.py | 1 + tests/shard4/test_penalty.py | 6 ++-- 19 files changed, 143 insertions(+), 51 deletions(-) diff --git a/bioptim/dynamics/configure_problem.py b/bioptim/dynamics/configure_problem.py index f028e905f..8b586c911 100644 --- a/bioptim/dynamics/configure_problem.py +++ b/bioptim/dynamics/configure_problem.py @@ -314,6 +314,7 @@ def torque_driven_free_floating_base( with_passive_torque: bool = False, with_ligament: bool = False, with_friction: bool = False, + dynamics_constants_used_at_each_nodes: dict[list] = {}, ): """ Configure the dynamics for a torque driven program with a free floating base. @@ -333,6 +334,8 @@ def torque_driven_free_floating_base( If the dynamic with ligament should be used with_friction: bool If the dynamic with joint friction should be used (friction = coefficients * qdot) + dynamics_constants_used_at_each_nodes: dict[np.ndarray] + A list of values to pass to the dynamics at each node. """ nb_q = nlp.model.nb_q @@ -437,6 +440,7 @@ def stochastic_torque_driven( with_friction: bool = False, with_cholesky: bool = False, initial_matrix: DM = None, + dynamics_constants_used_at_each_nodes: dict[list] = {}, ): """ Configure the dynamics for a torque driven stochastic program (states are q and qdot, controls are tau) @@ -451,6 +455,12 @@ def stochastic_torque_driven( If the dynamic with contact should be used with_friction: bool If the dynamic with joint friction should be used (friction = coefficient * qdot) + with_cholesky: bool + If the Cholesky decomposition should be used for the covariance matrix. + initial_matrix: DM + The initial value for the covariance matrix + dynamics_constants_used_at_each_nodes: dict[np.ndarray] + A list of values to pass to the dynamics at each node. Experimental external forces should be included here. """ if "tau" in nlp.model.motor_noise_mapping: @@ -510,10 +520,10 @@ def stochastic_torque_driven_free_floating_base( ocp, nlp, problem_type, - with_contact: bool = False, with_friction: bool = False, with_cholesky: bool = False, initial_matrix: DM = None, + dynamics_constants_used_at_each_nodes: dict[list] = {}, ): """ Configure the dynamics for a stochastic torque driven program with a free floating base. @@ -525,10 +535,14 @@ def stochastic_torque_driven_free_floating_base( A reference to the ocp nlp: NonLinearProgram A reference to the phase - with_contact: bool - If the dynamic with contact should be used with_friction: bool If the dynamic with joint friction should be used (friction = coefficient * qdot) + with_cholesky: bool + If the Cholesky decomposition should be used for the covariance matrix. + initial_matrix: DM + The initial value for the covariance matrix + dynamics_constants_used_at_each_nodes: dict[np.ndarray] + A list of values to pass to the dynamics at each node. """ n_noised_tau = nlp.model.n_noised_controls n_noise = nlp.model.motor_noise_magnitude.shape[0] + nlp.model.sensory_noise_magnitude.shape[0] @@ -751,7 +765,11 @@ def torque_activations_driven( ConfigureProblem.configure_soft_contact_function(ocp, nlp) @staticmethod - def joints_acceleration_driven(ocp, nlp, rigidbody_dynamics: RigidBodyDynamics = RigidBodyDynamics.ODE): + def joints_acceleration_driven(ocp, + nlp, + rigidbody_dynamics: RigidBodyDynamics = RigidBodyDynamics.ODE, + dynamics_constants_used_at_each_nodes: dict[list] = {} + ): """ Configure the dynamics for a joints acceleration driven program (states are q and qdot, controls are qddot_joints) @@ -764,7 +782,8 @@ def joints_acceleration_driven(ocp, nlp, rigidbody_dynamics: RigidBodyDynamics = A reference to the phase rigidbody_dynamics: RigidBodyDynamics which rigidbody dynamics should be used - + dynamics_constants_used_at_each_nodes: dict[np.ndarray] + A list of values to pass to the dynamics at each node. Experimental external forces should be included here. """ if rigidbody_dynamics != RigidBodyDynamics.ODE: raise NotImplementedError("Implicit dynamics not implemented yet.") @@ -902,7 +921,10 @@ def muscle_driven( ConfigureProblem.configure_soft_contact_function(ocp, nlp) @staticmethod - def holonomic_torque_driven(ocp, nlp): + def holonomic_torque_driven(ocp, + nlp, + dynamics_constants_used_at_each_nodes: dict[list] = {} + ): """ Tell the program which variables are states and controls. diff --git a/bioptim/dynamics/dynamics_functions.py b/bioptim/dynamics/dynamics_functions.py index 074f72599..62eb5b5e9 100644 --- a/bioptim/dynamics/dynamics_functions.py +++ b/bioptim/dynamics/dynamics_functions.py @@ -331,6 +331,7 @@ def stochastic_torque_driven( controls=controls, parameters=parameters, algebraic_states=algebraic_states, + dynamics_constants=dynamics_constants, sensory_noise=sensory_noise, motor_noise=motor_noise, ) diff --git a/bioptim/examples/getting_started/custom_dynamics.py b/bioptim/examples/getting_started/custom_dynamics.py index 6eca18237..7e6227447 100644 --- a/bioptim/examples/getting_started/custom_dynamics.py +++ b/bioptim/examples/getting_started/custom_dynamics.py @@ -82,7 +82,7 @@ def custom_dynamics( return DynamicsEvaluation(dxdt=vertcat(dq, ddq), defects=None) -def custom_configure(ocp: OptimalControlProgram, nlp: NonLinearProgram, my_additional_factor=1): +def custom_configure(ocp: OptimalControlProgram, nlp: NonLinearProgram, dynamics_constants_used_at_each_nodes, my_additional_factor=1): """ Tell the program which variables are states and controls. The user is expected to use the ConfigureProblem.configure_xxx functions. diff --git a/bioptim/examples/muscle_driven_ocp/muscle_excitations_tracker.py b/bioptim/examples/muscle_driven_ocp/muscle_excitations_tracker.py index ac39550cf..86fd5cf55 100644 --- a/bioptim/examples/muscle_driven_ocp/muscle_excitations_tracker.py +++ b/bioptim/examples/muscle_driven_ocp/muscle_excitations_tracker.py @@ -175,6 +175,7 @@ def generate_data( controls=symbolic_controls, parameters=symbolic_parameters, algebraic_states=MX(), + dynamics_constants=MX(), nlp=nlp, with_contact=False, rigidbody_dynamics=RigidBodyDynamics.ODE, diff --git a/bioptim/examples/stochastic_optimal_control/arm_reaching_muscle_driven.py b/bioptim/examples/stochastic_optimal_control/arm_reaching_muscle_driven.py index 3180b03a7..6a034aae2 100644 --- a/bioptim/examples/stochastic_optimal_control/arm_reaching_muscle_driven.py +++ b/bioptim/examples/stochastic_optimal_control/arm_reaching_muscle_driven.py @@ -53,6 +53,7 @@ def sensory_reference( controls: cas.MX | cas.SX, parameters: cas.MX | cas.SX, algebraic_states: cas.MX | cas.SX, + dynamics_constants: cas.MX | cas.SX, nlp: NonLinearProgram, ): """ @@ -70,6 +71,7 @@ def stochastic_forward_dynamics( controls: cas.MX | cas.SX, parameters: cas.MX | cas.SX, algebraic_states: cas.MX | cas.SX, + dynamics_constants: cas.MX | cas.SX, nlp: NonLinearProgram, force_field_magnitude, with_noise, @@ -92,7 +94,7 @@ def stochastic_forward_dynamics( k = DynamicsFunctions.get(nlp.algebraic_states["k"], algebraic_states) k_matrix = StochasticBioModel.reshape_to_matrix(k, nlp.model.matrix_shape_k) - hand_pos_velo = nlp.model.sensory_reference(time, states, controls, parameters, algebraic_states, nlp) + hand_pos_velo = nlp.model.sensory_reference(time, states, controls, parameters, algebraic_states, dynamics_constants, nlp) mus_excitations_fb += nlp.model.get_excitation_with_feedback(k_matrix, hand_pos_velo, ref, sensory_noise) noise_torque = motor_noise @@ -131,7 +133,8 @@ def stochastic_forward_dynamics( return DynamicsEvaluation(dxdt=cas.vertcat(dq_computed, dqdot_computed, dactivations_computed), defects=None) -def configure_stochastic_optimal_control_problem(ocp: OptimalControlProgram, nlp: NonLinearProgram): +def configure_stochastic_optimal_control_problem(ocp: OptimalControlProgram, + nlp: NonLinearProgram): ConfigureProblem.configure_q(ocp, nlp, True, False, False) ConfigureProblem.configure_qdot(ocp, nlp, True, False, True) ConfigureProblem.configure_qddot(ocp, nlp, False, False, True) @@ -148,15 +151,15 @@ def configure_stochastic_optimal_control_problem(ocp: OptimalControlProgram, nlp ConfigureProblem.configure_dynamics_function( ocp, nlp, - dyn_func=lambda time, states, controls, parameters, algebraic_states, nlp: nlp.dynamics_type.dynamic_function( - time, states, controls, parameters, algebraic_states, nlp, with_noise=False + dyn_func=lambda time, states, controls, parameters, algebraic_states, dynamics_constants, nlp: nlp.dynamics_type.dynamic_function( + time, states, controls, parameters, algebraic_states, dynamics_constants, nlp, with_noise=False ), ) ConfigureProblem.configure_dynamics_function( ocp, nlp, - dyn_func=lambda time, states, controls, parameters, algebraic_states, nlp: nlp.dynamics_type.dynamic_function( - time, states, controls, parameters, algebraic_states, nlp, with_noise=True + dyn_func=lambda time, states, controls, parameters, algebraic_states, dynamics_constants, nlp: nlp.dynamics_type.dynamic_function( + time, states, controls, parameters, algebraic_states, dynamics_constants, nlp, with_noise=True ), ) @@ -198,14 +201,15 @@ def get_cov_mat(nlp, node_index): nlp.controls.mx, nlp.parameters.mx, nlp.algebraic_states.mx, + nlp.dynamics_constants.mx, nlp, force_field_magnitude=nlp.model.force_field_magnitude, with_noise=True, ) dx.dxdt = cas.Function( - "tp", [nlp.states.mx, nlp.controls.mx, nlp.parameters.mx, nlp.algebraic_states.mx], [dx.dxdt] - )(nlp.states.cx, nlp.controls.cx, nlp.parameters.cx, nlp.algebraic_states.cx) + "tp", [nlp.states.mx, nlp.controls.mx, nlp.parameters.mx, nlp.algebraic_states.mx, nlp.dynamics_constants.mx], [dx.dxdt] + )(nlp.states.cx, nlp.controls.cx, nlp.parameters.cx, nlp.algebraic_states.cx, nlp.dynamics_constants.cx) ddx_dwm = cas.jacobian(dx.dxdt, cas.vertcat(sensory_noise, motor_noise)) dg_dw = -ddx_dwm * dt @@ -220,7 +224,7 @@ def get_cov_mat(nlp, node_index): func_eval = cas.Function( "p_next", - [dt, nlp.states.cx, nlp.controls.cx, nlp.parameters.cx, nlp.algebraic_states.cx, cov_sym], + [dt, nlp.states.cx, nlp.controls.cx, nlp.parameters.cx, nlp.algebraic_states.cx, nlp.dynamics_constants.cx, cov_sym], [p_next], )( nlp.dt, @@ -228,6 +232,7 @@ def get_cov_mat(nlp, node_index): nlp.controls.cx, parameters, nlp.algebraic_states.cx, + nlp.dynamics_constants.cx, nlp.integrated_values["cov"].cx, ) p_vector = StochasticBioModel.reshape_to_vector(func_eval) @@ -303,6 +308,7 @@ def expected_feedback_effort(controllers: list[PenaltyController], sensory_noise controllers[0].controls.cx, controllers[0].parameters.cx, controllers[0].algebraic_states.cx, + controllers[0].dynamics_constants.cx, controllers[0].get_nlp, ) trace_k_sensor_k = cas.trace(k_matrix @ sensory_noise_matrix @ k_matrix.T) @@ -335,6 +341,7 @@ def zero_acceleration(controller: PenaltyController, force_field_magnitude: floa controller.controls.cx, controller.parameters.cx, controller.algebraic_states.cx, + controller.dynamics_constants.cx, controller.get_nlp, force_field_magnitude=force_field_magnitude, with_noise=False, @@ -468,18 +475,20 @@ def prepare_socp( dynamics = DynamicsList() dynamics.add( configure_stochastic_optimal_control_problem, - dynamic_function=lambda time, states, controls, parameters, algebraic_states, nlp, with_noise: stochastic_forward_dynamics( + dynamic_function=lambda time, states, controls, parameters, algebraic_states, dynamincs_constants, nlp, with_noise: stochastic_forward_dynamics( time, states, controls, parameters, algebraic_states, + dynamincs_constants, nlp, force_field_magnitude=force_field_magnitude, with_noise=with_noise, ), expand_dynamics=False, phase_dynamics=PhaseDynamics.ONE_PER_NODE, + dynamics_constants_used_at_each_nodes={}, ) n_muscles = 6 @@ -708,6 +717,7 @@ def main(): controls = socp.nlp[0].controls.cx parameters = socp.nlp[0].parameters.cx algebraic_states = socp.nlp[0].algebraic_states.cx + dynamics_constants = socp.nlp[0].dynamics_constants.cx nlp = socp.nlp[0] out = stochastic_forward_dynamics( cas.vertcat(time, time + dt), @@ -715,11 +725,12 @@ def main(): controls, parameters, algebraic_states, + dynamics_constants, nlp, force_field_magnitude=force_field_magnitude, with_noise=True, ) - dyn_fun = cas.Function("dyn_fun", [dt, time, states, controls, parameters, algebraic_states], [out.dxdt]) + dyn_fun = cas.Function("dyn_fun", [dt, time, states, controls, parameters, algebraic_states, dynamics_constants], [out.dxdt]) fig, axs = plt.subplots(3, 2) n_simulations = 30 diff --git a/bioptim/examples/stochastic_optimal_control/arm_reaching_torque_driven_collocations.py b/bioptim/examples/stochastic_optimal_control/arm_reaching_torque_driven_collocations.py index b373cf926..7fce8f177 100644 --- a/bioptim/examples/stochastic_optimal_control/arm_reaching_torque_driven_collocations.py +++ b/bioptim/examples/stochastic_optimal_control/arm_reaching_torque_driven_collocations.py @@ -39,6 +39,7 @@ def sensory_reference( controls: cas.MX | cas.SX, parameters: cas.MX | cas.SX, algebraic_states: cas.MX | cas.SX, + dynamics_constants: cas.MX | cas.SX, nlp: NonLinearProgram, ): """ diff --git a/bioptim/examples/stochastic_optimal_control/arm_reaching_torque_driven_explicit.py b/bioptim/examples/stochastic_optimal_control/arm_reaching_torque_driven_explicit.py index 2296636dd..1dbd97639 100644 --- a/bioptim/examples/stochastic_optimal_control/arm_reaching_torque_driven_explicit.py +++ b/bioptim/examples/stochastic_optimal_control/arm_reaching_torque_driven_explicit.py @@ -49,6 +49,7 @@ def stochastic_forward_dynamics( controls: cas.MX | cas.SX, parameters: cas.MX | cas.SX, algebraic_states: cas.MX | cas.SX, + dynamics_constants: cas.MX | cas.SX, nlp: NonLinearProgram, with_noise: bool, ) -> DynamicsEvaluation: @@ -76,7 +77,7 @@ def stochastic_forward_dynamics( qdddot = DynamicsFunctions.get(nlp.controls["qdddot"], controls) dqdot_constraint = dynamics_torque_driven_with_feedbacks( - time, states, controls, parameters, algebraic_states, nlp, with_noise=with_noise + time, states, controls, parameters, algebraic_states, dynamics_constants, nlp, with_noise=with_noise ) defects = cas.vertcat(dqdot_constraint - qddot) @@ -104,19 +105,20 @@ def configure_stochastic_optimal_control_problem(ocp: OptimalControlProgram, nlp ConfigureProblem.configure_dynamics_function( ocp, nlp, - dyn_func=lambda time, states, controls, parameters, algebraic_states, nlp: nlp.dynamics_type.dynamic_function( - time, states, controls, parameters, algebraic_states, nlp, with_noise=False + dyn_func=lambda time, states, controls, parameters, algebraic_states, dynamics_constants, nlp: nlp.dynamics_type.dynamic_function( + time, states, controls, parameters, algebraic_states, dynamics_constants, nlp, with_noise=False ), ) ConfigureProblem.configure_dynamics_function( ocp, nlp, - dyn_func=lambda time, states, controls, parameters, algebraic_states, nlp: nlp.dynamics_type.dynamic_function( + dyn_func=lambda time, states, controls, parameters, algebraic_states, dynamics_constants, nlp: nlp.dynamics_type.dynamic_function( time, states, controls, parameters, algebraic_states, + dynamics_constants, nlp, with_noise=True, ), @@ -142,6 +144,7 @@ def sensory_reference( controls: cas.MX | cas.SX, parameters: cas.MX | cas.SX, algebraic_states: cas.MX | cas.SX, + dynamics_constants: cas.MX | cas.SX, nlp: NonLinearProgram, ): """ @@ -186,13 +189,13 @@ def get_cov_mat(nlp, node_index, use_sx): cov_matrix = StochasticBioModel.reshape_to_matrix(cov_sym, nlp.model.matrix_shape_cov) dx = stochastic_forward_dynamics( - nlp.time_cx, nlp.states.mx, nlp.controls.mx, nlp.parameters.mx, nlp.algebraic_states.mx, nlp, with_noise=True + nlp.time_cx, nlp.states.mx, nlp.controls.mx, nlp.parameters.mx, nlp.algebraic_states.mx, nlp.dynamics_constants.mx, nlp, with_noise=True ) dx.dxdt = cas.Function( "tp", - [nlp.states.mx, nlp.controls.mx, nlp.parameters.mx, nlp.algebraic_states.mx], + [nlp.states.mx, nlp.controls.mx, nlp.parameters.mx, nlp.algebraic_states.mx, nlp.dynamics_constants.mx], [dx.dxdt], - )(nlp.states.cx, nlp.controls.cx, nlp.parameters.cx, nlp.algebraic_states.cx) + )(nlp.states.cx, nlp.controls.cx, nlp.parameters.cx, nlp.algebraic_states.cx, nlp.dynamics_constants.cx) ddx_dwm = cas.jacobian(dx.dxdt, cas.vertcat(sensory_noise, motor_noise)) dg_dw = -ddx_dwm * dt @@ -201,7 +204,7 @@ def get_cov_mat(nlp, node_index, use_sx): p_next = M_matrix @ (dg_dx @ cov_matrix @ dg_dx.T + dg_dw @ sigma_w @ dg_dw.T) @ M_matrix.T func = cas.Function( - "p_next", [dt, nlp.states.cx, nlp.controls.cx, nlp.parameters.cx, nlp.algebraic_states.cx, cov_sym], [p_next] + "p_next", [dt, nlp.states.cx, nlp.controls.cx, nlp.parameters.cx, nlp.algebraic_states.cx, nlp.dynamics_constants.cx, cov_sym], [p_next] ) nlp.states.node_index = node_index - 1 @@ -219,6 +222,7 @@ def get_cov_mat(nlp, node_index, use_sx): nlp.controls.cx, parameters, nlp.algebraic_states.cx, + nlp.dynamics_constants.cx, nlp.integrated_values["cov"].cx, ) p_vector = StochasticBioModel.reshape_to_vector(func_eval) @@ -290,6 +294,7 @@ def expected_feedback_effort(controllers: list[PenaltyController]) -> cas.MX: controllers[0].controls.cx, controllers[0].parameters.cx, controllers[0].algebraic_states, + controllers[0].dynamics_constants.cx, controllers[0].get_nlp, ) e_fb = k_matrix @ ((estimated_ref - ref) + controllers[0].model.sensory_noise_magnitude) @@ -434,6 +439,7 @@ def prepare_socp( dynamic_function=stochastic_forward_dynamics, expand_dynamics=False, phase_dynamics=PhaseDynamics.ONE_PER_NODE, + dynamics_constants_used_at_each_nodes={}, ) states_min = np.ones((n_states, n_shooting + 1)) * -cas.inf diff --git a/bioptim/examples/stochastic_optimal_control/arm_reaching_torque_driven_implicit.py b/bioptim/examples/stochastic_optimal_control/arm_reaching_torque_driven_implicit.py index d912fa891..44c3f5826 100644 --- a/bioptim/examples/stochastic_optimal_control/arm_reaching_torque_driven_implicit.py +++ b/bioptim/examples/stochastic_optimal_control/arm_reaching_torque_driven_implicit.py @@ -50,6 +50,7 @@ def sensory_reference( controls: cas.MX | cas.SX, parameters: cas.MX | cas.SX, algebraic_states: cas.MX | cas.SX, + dynamics_constants: cas.MX | cas.SX, nlp: NonLinearProgram, ): """ @@ -188,6 +189,7 @@ def prepare_socp( expand_dynamics=False, phase_dynamics=PhaseDynamics.ONE_PER_NODE, with_friction=True, + dynamics_constants_used_at_each_nodes={}, ) states_min = np.ones((n_states, n_shooting + 1)) * -cas.inf diff --git a/bioptim/examples/stochastic_optimal_control/common.py b/bioptim/examples/stochastic_optimal_control/common.py index 9307b899f..8de94ee59 100644 --- a/bioptim/examples/stochastic_optimal_control/common.py +++ b/bioptim/examples/stochastic_optimal_control/common.py @@ -10,7 +10,7 @@ import matplotlib.transforms as transforms -def dynamics_torque_driven_with_feedbacks(time, states, controls, parameters, algebraic_states, nlp, with_noise): +def dynamics_torque_driven_with_feedbacks(time, states, controls, parameters, algebraic_states, dynamics_constants, nlp, with_noise): q = DynamicsFunctions.get(nlp.states["q"], states) qdot = DynamicsFunctions.get(nlp.states["qdot"], states) tau = DynamicsFunctions.get(nlp.controls["tau"], controls) @@ -24,7 +24,7 @@ def dynamics_torque_driven_with_feedbacks(time, states, controls, parameters, al motor_noise = DynamicsFunctions.get(nlp.parameters["motor_noise"], parameters) sensory_noise = DynamicsFunctions.get(nlp.parameters["sensory_noise"], parameters) - end_effector = nlp.model.sensory_reference(time, states, controls, parameters, algebraic_states, nlp) + end_effector = nlp.model.sensory_reference(time, states, controls, parameters, algebraic_states, dynamics_constants, nlp) tau_feedback = get_excitation_with_feedback(k_matrix, end_effector, ref, sensory_noise) diff --git a/bioptim/interfaces/acados_interface.py b/bioptim/interfaces/acados_interface.py index 3d96fb271..c30a9be71 100644 --- a/bioptim/interfaces/acados_interface.py +++ b/bioptim/interfaces/acados_interface.py @@ -296,6 +296,7 @@ def __set_constraints(self, ocp): u = nlp.controls.cx_start p = nlp.parameters.cx a = nlp.algebraic_states.cx_start + d = nlp.dynamics_constants.cx for g, G in enumerate(nlp.g): if not G: diff --git a/bioptim/limits/constraints.py b/bioptim/limits/constraints.py index e59d6773c..46c86c46b 100644 --- a/bioptim/limits/constraints.py +++ b/bioptim/limits/constraints.py @@ -697,6 +697,7 @@ def stochastic_df_dx_implicit( qdot_joints = MX.sym("qdot_joints", nu, 1) tau_joints = MX.sym("tau_joints", nu, 1) algebraic_states_sym = MX.sym("algebraic_states_sym", controller.algebraic_states.shape, 1) + dynamics_constants_sym = MX.sym("dynamics_constants_sym", controller.dynamics_constants.shape, 1) dx = controller.extra_dynamics(0)( controller.t_span.mx, @@ -704,6 +705,7 @@ def stochastic_df_dx_implicit( tau_joints, controller.parameters.mx, algebraic_states_sym, + dynamics_constants_sym, ) non_root_index = list(range(nb_root, nb_root + nu)) + list( @@ -720,6 +722,7 @@ def stochastic_df_dx_implicit( tau_joints, controller.parameters.mx, algebraic_states_sym, + dynamics_constants_sym, ], [jacobian(dx[non_root_index], vertcat(q_joints, qdot_joints))], ) @@ -737,6 +740,7 @@ def stochastic_df_dx_implicit( controller.controls.cx, parameters, controller.algebraic_states.cx, + controller.dynamics_constants.cx, ) CX_eye = SX_eye if controller.ocp.cx == SX else MX_eye @@ -776,6 +780,7 @@ def stochastic_helper_matrix_collocation( controller.controls.cx_start, parameters, controller.algebraic_states.cx_start, + controller.dynamics_constants.cx ) return StochasticBioModel.reshape_to_vector(constraint) @@ -814,6 +819,7 @@ def stochastic_covariance_matrix_continuity_collocation( controller.controls.cx_start, parameters, controller.algebraic_states.cx_start, + controller.dynamics_constants.cx, ) cov_implicit_defect = cov_matrix_next - cov_next_computed @@ -842,6 +848,7 @@ def stochastic_mean_sensory_input_equals_reference( controls=controller.controls.mx, parameters=controller.parameters.mx, algebraic_states=controller.algebraic_states.mx, + dynamics_constants=controller.dynamics_constants.mx, nlp=controller.get_nlp, ) @@ -853,6 +860,7 @@ def stochastic_mean_sensory_input_equals_reference( controller.controls.mx, controller.parameters.mx, controller.algebraic_states.mx, + controller.dynamics_constants.mx, ], [sensory_input], )( @@ -861,6 +869,7 @@ def stochastic_mean_sensory_input_equals_reference( controller.controls.cx_start, controller.parameters.cx, controller.algebraic_states.cx_start, + controller.dynamics_constants.cx, ) return sensory_input[: controller.model.n_feedbacks] - ref[: controller.model.n_feedbacks] @@ -949,8 +958,10 @@ def collocation_jacobians(penalty, controller): controller.states_scaled.cx_start, horzcat(*controller.states_scaled.cx_intermediates_list), controller.controls_scaled.cx_start, - controller.parameters.cx, # TODO: fix parameter scaling + controller.parameters_scaled.cx, controller.algebraic_states_scaled.cx_start, + controller.dynamics_constants.cx, + ], [Fdz.T - Gdz.T @ m_matrix.T], ) @@ -965,8 +976,9 @@ def collocation_jacobians(penalty, controller): controller.states_scaled.cx_start, horzcat(*controller.states_scaled.cx_intermediates_list), controller.controls_scaled.cx_start, - controller.parameters.cx, # TODO: fix parameter scaling + controller.parameters_scaled.cx, controller.algebraic_states_scaled.cx_start, + controller.dynamics_constants.cx, ], [m_matrix @ (Gdx @ cov_matrix @ Gdx.T + Gdw @ sigma_ww @ Gdw.T) @ m_matrix.T], ) @@ -980,8 +992,9 @@ def collocation_jacobians(penalty, controller): controller.states_scaled.cx_start, horzcat(*controller.states_scaled.cx_intermediates_list), controller.controls_scaled.cx_start, - controller.parameters.cx, # TODO: fix parameter scaling + controller.parameters_scaled.cx, controller.algebraic_states_scaled.cx_start, + controller.dynamics_constants.cx, ], [Gdx], ) @@ -993,8 +1006,9 @@ def collocation_jacobians(penalty, controller): controller.states_scaled.cx_start, horzcat(*controller.states_scaled.cx_intermediates_list), controller.controls_scaled.cx_start, - controller.parameters.cx, # TODO: fix parameter scaling + controller.parameters_scaled.cx, controller.algebraic_states_scaled.cx_start, + controller.dynamics_constants.cx, ], [Gdz], ) @@ -1006,8 +1020,9 @@ def collocation_jacobians(penalty, controller): controller.states_scaled.cx_start, horzcat(*controller.states_scaled.cx_intermediates_list), controller.controls_scaled.cx_start, - controller.parameters.cx, # TODO: fix parameter scaling + controller.parameters_scaled.cx, controller.algebraic_states_scaled.cx_start, + controller.dynamics_constants.cx, ], [Gdw], ) @@ -1019,8 +1034,9 @@ def collocation_jacobians(penalty, controller): controller.states_scaled.cx_start, horzcat(*controller.states_scaled.cx_intermediates_list), controller.controls_scaled.cx_start, - controller.parameters.cx, # TODO: fix parameter scaling + controller.parameters_scaled.cx, controller.algebraic_states_scaled.cx_start, + controller.dynamics_constants.cx, ], [Fdz], ) diff --git a/bioptim/limits/multinode_penalty.py b/bioptim/limits/multinode_penalty.py index 35b4d765c..bd89db052 100644 --- a/bioptim/limits/multinode_penalty.py +++ b/bioptim/limits/multinode_penalty.py @@ -397,6 +397,7 @@ def stochastic_helper_matrix_explicit( controllers[0].controls.cx, parameters, controllers[0].algebraic_states.cx, + controllers[0].dynamics_constants.cx, ) DdZ_DX_fun = Function( @@ -407,6 +408,7 @@ def stochastic_helper_matrix_explicit( controllers[0].controls.cx, controllers[0].parameters.cx, controllers[0].algebraic_states.cx, + controllers[0].dynamics_constants.cx, ], [jacobian(dx, controllers[0].states.cx)], ) @@ -421,6 +423,7 @@ def stochastic_helper_matrix_explicit( controllers[1].controls.cx, parameters, controllers[1].algebraic_states.cx, + controllers[1].dynamics_constants.cx, ) CX_eye = SX_eye if controllers[0].cx == SX else MX_eye @@ -551,6 +554,7 @@ def stochastic_df_dw_implicit( qdot_joints = MX.sym("qdot_joints", nu, 1) tau_joints = MX.sym("tau_joints", nu, 1) algebraic_states_sym = MX.sym("algebraic_states_sym", controllers[0].algebraic_states.shape, 1) + dynamics_constants_sym = MX.sym("dynamics_constants_sym", controllers[0].dynamics_constants.shape, 1) dx = controllers[0].extra_dynamics(0)( controllers[0].time.mx, @@ -558,6 +562,7 @@ def stochastic_df_dw_implicit( tau_joints, controllers[0].parameters.mx, algebraic_states_sym, + dynamics_constants_sym, ) non_root_index = list(range(nb_root, nb_root + nu)) + list( @@ -575,6 +580,7 @@ def stochastic_df_dw_implicit( tau_joints, controllers[0].parameters.mx, algebraic_states_sym, + dynamics_constants_sym, ], [ jacobian( @@ -599,6 +605,7 @@ def stochastic_df_dw_implicit( controllers[0].controls["tau"].cx, parameters, controllers[0].algebraic_states.cx, + controllers[0].dynamics_constants.cx, ) parameters = controllers[1].parameters.cx @@ -614,6 +621,7 @@ def stochastic_df_dw_implicit( controllers[1].controls.cx, parameters, controllers[1].algebraic_states.cx, + controllers[1].dynamics_constants.cx, ) out = c_matrix - (-(DF_DW + DF_DW_plus) / 2 * dt) diff --git a/bioptim/limits/penalty.py b/bioptim/limits/penalty.py index 8f4f9aa83..4b802ae29 100644 --- a/bioptim/limits/penalty.py +++ b/bioptim/limits/penalty.py @@ -205,6 +205,7 @@ def stochastic_minimize_expected_feedback_efforts(penalty: PenaltyOption, contro controls=controller.controls_scaled.mx, parameters=controller.parameters_scaled.mx, algebraic_states=controller.algebraic_states_scaled.mx, + dynamics_constants=controller.dynamics_constants.mx, sensory_noise=controller.model.sensory_noise_magnitude, motor_noise=controller.model.motor_noise_magnitude, ) @@ -219,6 +220,7 @@ def stochastic_minimize_expected_feedback_efforts(penalty: PenaltyOption, contro controller.controls_scaled.mx, controller.parameters_scaled.mx, controller.algebraic_states_scaled.mx, + controller.dynamics_constants.mx, ], [jac_e_fb_x], )( @@ -227,6 +229,7 @@ def stochastic_minimize_expected_feedback_efforts(penalty: PenaltyOption, contro controller.controls.cx_start, controller.parameters.cx_start, controller.algebraic_states.cx_start, + controller.dynamics_constants.cx_start, ) trace_jac_p_jack = trace(jac_e_fb_x_cx @ cov_matrix @ jac_e_fb_x_cx.T) diff --git a/bioptim/limits/penalty_helpers.py b/bioptim/limits/penalty_helpers.py index 1aaeafaca..09b893045 100644 --- a/bioptim/limits/penalty_helpers.py +++ b/bioptim/limits/penalty_helpers.py @@ -132,15 +132,18 @@ def parameters(penalty, index, get_parameter_decision: Callable): def dynamics_constants(penalty, index, get_dynamics_constants: Callable, is_constructing_penalty: bool = False): node = penalty.node_idx[index] if penalty.multinode_penalty: - d = [] - phases, _, subnodes = _get_multinode_indices(penalty, is_constructing_penalty) - for phase, sub in zip(phases, subnodes): - d_tp = get_dynamics_constants(phase, node, sub) - if d_tp.shape != (0, 0): - d += [_reshape_to_vector(get_dynamics_constants(phase, node, sub))] - return _vertcat(d) if len(d) > 0 else DM() - - d = get_dynamics_constants(penalty.phase, node, 0) # cx_start + d = get_dynamics_constants(penalty.nodes_phase[0], node, 0) # cx_start + # d = [] + # phases, _, subnodes = _get_multinode_indices(penalty, is_constructing_penalty) + # for phase, sub in zip(phases, subnodes): + # d_tp = get_dynamics_constants(phase, node, sub) + # if d_tp.shape != (0, 0): + # d += [_reshape_to_vector(get_dynamics_constants(phase, node, sub))] + # return _vertcat(d) if len(d) > 0 else DM() + + else: + d = get_dynamics_constants(penalty.phase, node, 0) # cx_start + if d.shape != (0, 0): d = _reshape_to_vector(d) diff --git a/bioptim/limits/penalty_option.py b/bioptim/limits/penalty_option.py index 16caa7c7e..1eb6393d0 100644 --- a/bioptim/limits/penalty_option.py +++ b/bioptim/limits/penalty_option.py @@ -675,7 +675,6 @@ def get_variable_inputs(self, controllers: list[PenaltyController]): lambda p_idx, n_idx, sn_idx: self._get_dynamics_constants(ocp, p_idx, n_idx, sn_idx), is_constructing_penalty=True, ) - # dynamics_constants = controller.dynamics_constants.cx return controller, t0, x, u, p, a, dynamics_constants @@ -780,7 +779,10 @@ def vertcat_cx_end(): def _get_dynamics_constants(self, ocp, p_idx, n_idx, sn_idx): - nlp = ocp.nlp[p_idx] + try: + nlp = ocp.nlp[p_idx] + except: + print("ici") dynamics_constants = nlp.dynamics_constants if dynamics_constants.cx_start.shape == (0, 0): diff --git a/bioptim/models/biorbd/stochastic_biorbd_model.py b/bioptim/models/biorbd/stochastic_biorbd_model.py index 0212de065..7106f1676 100644 --- a/bioptim/models/biorbd/stochastic_biorbd_model.py +++ b/bioptim/models/biorbd/stochastic_biorbd_model.py @@ -8,7 +8,7 @@ def _compute_torques_from_noise_and_feedback_default( - nlp, time, states, controls, parameters, algebraic_states, sensory_noise, motor_noise + nlp, time, states, controls, parameters, algebraic_states, dynamics_constants, sensory_noise, motor_noise ): tau_nominal = DynamicsFunctions.get(nlp.controls["tau"], controls) @@ -18,7 +18,7 @@ def _compute_torques_from_noise_and_feedback_default( k_matrix = StochasticBioModel.reshape_to_matrix(k, nlp.model.matrix_shape_k) - sensory_input = nlp.model.sensory_reference(time, states, controls, parameters, algebraic_states, nlp) + sensory_input = nlp.model.sensory_reference(time, states, controls, parameters, algebraic_states, dynamics_constants, nlp) tau_fb = k_matrix @ ((sensory_input - ref) + sensory_noise) tau_motor_noise = motor_noise diff --git a/bioptim/optimization/solution/solution.py b/bioptim/optimization/solution/solution.py index ba33ede13..b40e34a23 100644 --- a/bioptim/optimization/solution/solution.py +++ b/bioptim/optimization/solution/solution.py @@ -784,6 +784,7 @@ def integrate( ------- Return the integrated states """ + from ...interfaces.interface_utils import _get_dynamics_constants t_spans, x, u, params, a = self._prepare_integrate(integrator=integrator) @@ -843,9 +844,10 @@ def noisy_integrate( size: int = 100, ): """ - TODO: Charbie! + Integrated the states with different noise values sampled from the covariance matrix. """ from ...optimization.stochastic_optimal_control_program import StochasticOptimalControlProgram + from ...interfaces.interface_utils import _get_dynamics_constants if not isinstance(self.ocp, StochasticOptimalControlProgram): raise ValueError("This method is only available for StochasticOptimalControlProgram.") @@ -874,6 +876,14 @@ def noisy_integrate( cov_matrix = StochasticBioModel.reshape_to_matrix(a[0][0][cov_index, :], self.ocp.nlp[0].model.matrix_shape_cov) first_x = np.random.multivariate_normal(x[0][0][:, 0], cov_matrix, size=size).T for p, nlp in enumerate(self.ocp.nlp): + d = [] + for n_idx in range(nlp.ns + 1): + d_tp = _get_dynamics_constants(self.ocp, p, n_idx, 0) + if d_tp.shape == (0, 0): + d += [np.array([])] + else: + d += [np.array(d_tp)] + motor_noise = np.zeros((len(params[motor_noise_index]), nlp.ns, size)) for i in range(len(params[motor_noise_index])): motor_noise[i, :] = np.random.normal(0, params[motor_noise_index[i]], size=(nlp.ns, size)) @@ -902,7 +912,7 @@ def noisy_integrate( raise NotImplementedError("Noisy integration is not available for multiple extra dynamics.") cas_func = Function( "noised_extra_dynamics", - [nlp.time_cx, nlp.states.cx, nlp.controls.cx, parameters_cx, nlp.algebraic_states.cx], + [nlp.time_cx, nlp.states.cx, nlp.controls.cx, parameters_cx, nlp.algebraic_states.cx, nlp.dynamics_constants.cx], [ nlp.extra_dynamics_func[0]( nlp.time_cx, @@ -910,6 +920,7 @@ def noisy_integrate( nlp.controls.cx, params_this_time[node], nlp.algebraic_states.cx, + nlp.dynamics_constants.cx, ) ], ) @@ -924,6 +935,7 @@ def noisy_integrate( u=u[p], # No need to add noise on the controls, the extra_dynamics should do it for us a=a[p], p=parameters, + d=d, method=integrator, ) for i_node in range(nlp.ns + 1): diff --git a/bioptim/optimization/stochastic_optimal_control_program.py b/bioptim/optimization/stochastic_optimal_control_program.py index 908110021..f094a468a 100644 --- a/bioptim/optimization/stochastic_optimal_control_program.py +++ b/bioptim/optimization/stochastic_optimal_control_program.py @@ -531,6 +531,7 @@ def get_cov_init( p=p_guess, a=fake_algebraic_states, a_scaled=[], + d=[], node_index=0, ) _, _, Gdx, Gdz, Gdw, Fdz = ConstraintFunction.Functions.collocation_jacobians(penalty, penalty_controller) diff --git a/tests/shard4/test_penalty.py b/tests/shard4/test_penalty.py index df1f54931..67e5f4e8c 100644 --- a/tests/shard4/test_penalty.py +++ b/tests/shard4/test_penalty.py @@ -106,9 +106,11 @@ def get_penalty_value(ocp, penalty, t, phases_dt, x, u, p, a, d): algebraic_states = ( ocp.nlp[0].algebraic_states.cx_start if ocp.nlp[0].algebraic_states.cx_start.shape != (0, 0) else ocp.cx(0, 0) ) + dynamics_constants = ocp.nlp[0].dynamics_constants.cx if ocp.nlp[0].dynamics_constants.cx.shape != (0, 0) else ocp.cx(0, 0) + return ocp.nlp[0].to_casadi_func( - "penalty", val, time, phases_dt_cx, states, controls, parameters, algebraic_states - )(t, phases_dt, x[0], u[0], p, a) + "penalty", val, time, phases_dt_cx, states, controls, parameters, algebraic_states, dynamics_constants + )(t, phases_dt, x[0], u[0], p, a, d) def test_penalty_targets_shapes():