diff --git a/bioptim/dynamics/configure_problem.py b/bioptim/dynamics/configure_problem.py index 1314c5a03..1f43c09fd 100644 --- a/bioptim/dynamics/configure_problem.py +++ b/bioptim/dynamics/configure_problem.py @@ -128,6 +128,7 @@ def initialize(ocp, nlp): A reference to the ocp nlp: NonLinearProgram A reference to the phase + A reference to the phase """ nlp.dynamics_type.type(ocp, nlp, **nlp.dynamics_type.params) @@ -613,7 +614,7 @@ def holonomic_torque_driven(ocp, nlp): ConfigureProblem.configure_dynamics_function(ocp, nlp, DynamicsFunctions.holonomic_torque_driven, expand=False) @staticmethod - def configure_dynamics_function(ocp, nlp, dyn_func, t: MX | SX = None, expand: bool = True, **extra_params): + def configure_dynamics_function(ocp, nlp, dyn_func, expand: bool = True, **extra_params): """ Configure the dynamics of the system @@ -633,6 +634,7 @@ def configure_dynamics_function(ocp, nlp, dyn_func, t: MX | SX = None, expand: b DynamicsFunctions.apply_parameters(nlp.parameters.mx, nlp) dynamics_eval = dyn_func( + nlp.time.mx, nlp.states.scaled.mx_reduced, nlp.controls.scaled.mx_reduced, nlp.parameters.mx, @@ -648,13 +650,14 @@ def configure_dynamics_function(ocp, nlp, dyn_func, t: MX | SX = None, expand: b nlp.dynamics_func = Function( "ForwardDyn", [ + nlp.time.mx, nlp.states.scaled.mx_reduced, nlp.controls.scaled.mx_reduced, nlp.parameters.mx, nlp.stochastic_variables.mx, ], [dynamics_dxdt], - ["x", "u", "p", "s"], + ["t", "x", "u", "p", "s"], ["xdot"], ) @@ -665,6 +668,7 @@ def configure_dynamics_function(ocp, nlp, dyn_func, t: MX | SX = None, expand: b nlp.implicit_dynamics_func = Function( "DynamicsDefects", [ + nlp.time.mx, nlp.states.scaled.mx_reduced, nlp.controls.scaled.mx_reduced, nlp.parameters.mx, @@ -672,7 +676,7 @@ def configure_dynamics_function(ocp, nlp, dyn_func, t: MX | SX = None, expand: b nlp.states_dot.scaled.mx_reduced, ], [dynamics_eval.defects], - ["x", "u", "p", "s", "xdot"], + ["t", "x", "u", "p", "s", "xdot"], ["defects"], ) if expand: @@ -696,6 +700,7 @@ def configure_contact_function(ocp, nlp, dyn_func: Callable, **extra_params): nlp.contact_forces_func = Function( "contact_forces_func", [ + nlp.time.mx, nlp.states.scaled.mx_reduced, nlp.controls.scaled.mx_reduced, nlp.parameters.mx, @@ -703,6 +708,7 @@ def configure_contact_function(ocp, nlp, dyn_func: Callable, **extra_params): ], [ dyn_func( + nlp.time.mx, nlp.states.scaled.mx_reduced, nlp.controls.scaled.mx_reduced, nlp.parameters.mx, @@ -711,7 +717,7 @@ def configure_contact_function(ocp, nlp, dyn_func: Callable, **extra_params): **extra_params, ) ], - ["x", "u", "p", "s"], + ["t", "x", "u", "p", "s"], ["contact_forces"], ).expand() @@ -759,9 +765,9 @@ def configure_soft_contact_function(ocp, nlp): ) nlp.soft_contact_forces_func = Function( "soft_contact_forces_func", - [nlp.states.mx_reduced, nlp.controls.mx_reduced, nlp.parameters.mx], + [nlp.time.mx, nlp.states.mx_reduced, nlp.controls.mx_reduced, nlp.parameters.mx], [global_soft_contact_force_func], - ["x", "u", "p"], + ["t", "x", "u", "p"], ["soft_contact_forces"], ).expand() @@ -796,7 +802,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 t, x, u, p, s: nlp.soft_contact_forces_func(x, u, p, s)[(i_sc * 6) : ((i_sc + 1) * 6), :], + lambda t, x, u, p, s: nlp.soft_contact_forces_func(t, x, u, p, s)[(i_sc * 6) : ((i_sc + 1) * 6), :], plot_type=PlotType.INTEGRATED, axes_idx=phase_mappings, legend=all_soft_contact_names, @@ -917,6 +923,7 @@ def configure_new_variable( as_controls: bool, as_states_dot: bool = False, as_stochastic: bool = False, + as_time: bool = False, fatigue: FatigueList = None, combine_name: str = None, combine_state_control_plot: bool = False, @@ -1037,6 +1044,7 @@ def define_cx_unscaled(_cx_scaled: list, scaling: np.ndarray) -> list: nlp.u_scaling.add(name, scaling=np.ones(len(nlp.variable_mappings[name].to_first.map_idx))) # Use of states[0] and controls[0] is permitted since ocp.assume_phase_dynamics is True + mx_time = [] mx_states = [] if not copy_states else [ocp.nlp[nlp.use_states_from_phase_idx].states[0][name].mx] mx_states_dot = ( [] if not copy_states_dot else [ocp.nlp[nlp.use_states_dot_from_phase_idx].states_dot[0][name].mx] @@ -1057,8 +1065,10 @@ def define_cx_unscaled(_cx_scaled: list, scaling: np.ndarray) -> list: if not copy_controls: mx_controls.append(MX.sym(var_name, 1, 1)) + mx_time.append(MX.sym(var_name, 1, 1)) mx_stochastic.append(MX.sym(var_name, 1, 1)) + mx_time = vertcat(*mx_time) mx_states = vertcat(*mx_states) mx_states_dot = vertcat(*mx_states_dot) mx_controls = vertcat(*mx_controls) @@ -1159,6 +1169,16 @@ def define_cx_unscaled(_cx_scaled: list, scaling: np.ndarray) -> list: name, cx_scaled[0], cx_scaled[0], mx_stochastic, nlp.variable_mappings[name], node_index ) + if as_time: + for node_index in range((0 if ocp.assume_phase_dynamics else nlp.ns) + 1): + n_cx = nlp.ode_solver.polynomial_degree + 1 if isinstance(nlp.ode_solver, OdeSolver.COLLOCATION) else 3 + if n_cx < 3: + n_cx = 3 + cx_scaled = define_cx_scaled(n_col=n_cx, n_shooting=1, initial_node=node_index) + nlp.time.append( + name, cx_scaled[0], cx_scaled[0], mx_time, nlp.variable_mappings[name], node_index + ) + @staticmethod def configure_integrated_value( name: str, diff --git a/bioptim/dynamics/dynamics_functions.py b/bioptim/dynamics/dynamics_functions.py index 1e63d6967..1fae1b8fc 100644 --- a/bioptim/dynamics/dynamics_functions.py +++ b/bioptim/dynamics/dynamics_functions.py @@ -43,7 +43,7 @@ class DynamicsFunctions: @staticmethod def custom( - states: MX.sym, controls: MX.sym, parameters: MX.sym, stochastic_variables: MX.sym, nlp + time: MX.sym, states: MX.sym, controls: MX.sym, parameters: MX.sym, stochastic_variables: MX.sym, nlp ) -> DynamicsEvaluation: """ Interface to custom dynamic function provided by the user. @@ -67,10 +67,11 @@ def custom( The defects of the implicit dynamics """ - return nlp.dynamics_type.dynamic_function(states, controls, parameters, stochastic_variables, nlp) + return nlp.dynamics_type.dynamic_function(time, states, controls, parameters, stochastic_variables, nlp) @staticmethod def torque_driven( + time: MX.sym, states: MX.sym, controls: MX.sym, parameters: MX.sym, @@ -227,6 +228,7 @@ def __get_fatigable_tau(nlp: NonLinearProgram, states: MX, controls: MX, fatigue @staticmethod def torque_activations_driven( + time: MX.sym, states: MX.sym, controls: MX.sym, parameters: MX.sym, @@ -286,6 +288,7 @@ def torque_activations_driven( @staticmethod def torque_derivative_driven( + time: MX.sym, states: MX.sym, controls: MX.sym, parameters: MX.sym, @@ -357,6 +360,7 @@ def torque_derivative_driven( @staticmethod def forces_from_torque_driven( + time: MX.sym, states: MX.sym, controls: MX.sym, parameters: MX.sym, @@ -405,6 +409,7 @@ def forces_from_torque_driven( @staticmethod def forces_from_torque_activation_driven( + time: MX.sym, states: MX.sym, controls: MX.sym, parameters: MX.sym, @@ -451,6 +456,7 @@ def forces_from_torque_activation_driven( @staticmethod def muscles_driven( + time: MX.sym, states: MX.sym, controls: MX.sym, parameters: MX.sym, @@ -587,6 +593,7 @@ def muscles_driven( @staticmethod def forces_from_muscle_driven( + time: MX.sym, states: MX.sym, controls: MX.sym, parameters: MX.sym, @@ -634,6 +641,7 @@ def forces_from_muscle_driven( @staticmethod def joints_acceleration_driven( + time: MX.sym, states: MX.sym, controls: MX.sym, parameters: MX.sym, @@ -901,6 +909,7 @@ def compute_tau_from_muscle( @staticmethod def holonomic_torque_driven( + time: MX | SX, states: MX | SX, controls: MX | SX, parameters: MX | SX, diff --git a/bioptim/dynamics/integrator.py b/bioptim/dynamics/integrator.py index 309273d72..b63b5424b 100644 --- a/bioptim/dynamics/integrator.py +++ b/bioptim/dynamics/integrator.py @@ -19,6 +19,8 @@ class Integrator: The index of the degrees of freedom to integrate cx: MX | SX The CasADi type the integration should be built from + t_sym: MX | SX + The state variables x_sym: MX | SX The state variables u_sym: MX | SX @@ -50,7 +52,7 @@ class Integrator: Get the multithreaded CasADi graph of the integration get_u(self, u: np.ndarray, dt_norm: float) -> np.ndarray Get the control at a given time - dxdt(self, h: float, states: MX | SX, controls: MX | SX, params: MX | SX, stochastic_variables: MX | SX) -> tuple[SX, list[SX]] + dxdt(self, h: float, time: MX | SX, states: MX | SX, controls: MX | SX, params: MX | SX, stochastic_variables: MX | SX) -> tuple[SX, list[SX]] The dynamics of the system _finish_init(self) Prepare the CasADi function from dxdt @@ -71,6 +73,7 @@ def __init__(self, ode: dict, ode_opt: dict): self.t_span = ode_opt["t0"], ode_opt["tf"] self.idx = ode_opt["idx"] self.cx = ode_opt["cx"] + self.t_sym = ode["time"] self.x_sym = ode["x_scaled"] self.u_sym = [] if ode_opt["control_type"] is ControlType.NONE else ode["p_scaled"] self.param_sym = ode_opt["param"].cx @@ -83,7 +86,6 @@ def __init__(self, ode: dict, ode_opt: dict): self.control_type = ode_opt["control_type"] self.step_time = self.t_span[1] - self.t_span[0] self.h = self.step_time - self.t = ode["time"] self.function = None def __call__(self, *args, **kwargs): @@ -123,6 +125,7 @@ def get_u(self, u: np.ndarray, dt_norm: float) -> np.ndarray: return u elif self.control_type == ControlType.LINEAR_CONTINUOUS: return u[:, 0] + (u[:, 1] - u[:, 0]) * dt_norm + # TODO LINEAR CONTINUOUS error because time is SX(0x1) elif self.control_type == ControlType.NONE: return np.ndarray((0,)) else: @@ -131,6 +134,7 @@ def get_u(self, u: np.ndarray, dt_norm: float) -> np.ndarray: def dxdt( self, h: float, + time: MX | SX, states: MX | SX, controls: MX | SX, params: MX | SX, @@ -144,6 +148,8 @@ def dxdt( ---------- h: float The time step + time: MX | SX + The time of the system states: MX | SX The states of the system controls: MX | SX @@ -169,11 +175,11 @@ def _finish_init(self): self.function = Function( "integrator", - [self.x_sym, self.u_sym, self.param_sym, self.stochastic_variables_sym], + [self.t_sym, self.x_sym, self.u_sym, self.param_sym, self.stochastic_variables_sym], self.dxdt( - self.h, self.x_sym, self.u_sym, self.param_sym, self.param_scaling, self.stochastic_variables_sym + self.h, self.t_sym, self.x_sym, self.u_sym, self.param_sym, self.param_scaling, self.stochastic_variables_sym ), - ["x0", "p", "params", "s"], + ["t", "x0", "p", "params", "s"], ["xf", "xall"], ) @@ -195,7 +201,7 @@ class RK(Integrator): ------- next_x(self, h: float, t: float, x_prev: MX | SX, u: MX | SX, p: MX | SX, s: MX | SX) Compute the next integrated state (abstract) - dxdt(self, h: float, states: MX | SX, controls: MX | SX, params: MX | SX, stochastic_variables: MX | SX) -> tuple[SX, list[SX]] + dxdt(self, h: float, time: MX | SX, states: MX | SX, controls: MX | SX, params: MX | SX, stochastic_variables: MX | SX) -> tuple[SX, list[SX]] The dynamics of the system """ @@ -243,6 +249,7 @@ def next_x(self, h: float, t: float, x_prev: MX | SX, u: MX | SX, p: MX | SX, s: def dxdt( self, h: float, + time: MX | SX, states: MX | SX, controls: MX | SX, params: MX | SX, @@ -256,6 +263,8 @@ def dxdt( ---------- h: float The time step + time: MX | SX + The time of the system states: MX | SX The states of the system controls: MX | SX @@ -271,7 +280,7 @@ def dxdt( ------- The derivative of the states """ - + t = time u = controls x = self.cx(states.shape[0], self.n_step + 1) p = params * param_scaling @@ -279,7 +288,8 @@ def dxdt( s = stochastic_variables for i in range(1, self.n_step + 1): - t_norm_init = self.t if 't' in self.fun.name_in() else (i - 1) / self.n_step + # t_norm_init = t if 't' in self.fun.name_in() else (i - 1) / self.n_step + t_norm_init = t x[:, i] = self.next_x(h, t_norm_init, x[:, i - 1], u, p, s) if self.model.nb_quaternions > 0: x[:, i] = self.model.normalize_state_quaternions(x[:, i]) @@ -334,7 +344,7 @@ def next_x(self, h: float, t: float, x_prev: MX | SX, u: MX | SX, p: MX | SX, s: The next integrate states """ - return x_prev + h * self.fun(x_prev, self.get_u(u, t), p, s)[:, self.idx] + return x_prev + h * self.fun(t, x_prev, self.get_u(u, t), p, s)[:, self.idx] class RK2(RK): @@ -383,8 +393,8 @@ def next_x(self, h: float, t: float, x_prev: MX | SX, u: MX | SX, p: MX | SX, s: ------- The next integrate states """ - k1 = self.fun(x_prev, self.get_u(u, t), p, s)[:, self.idx] - return x_prev + h * self.fun(x_prev + h / 2 * k1, self.get_u(u, t + self.h_norm / 2), p, s)[:, self.idx] + k1 = self.fun(t, x_prev, self.get_u(u, t), p, s)[:, self.idx] + return x_prev + h * self.fun(t, x_prev + h / 2 * k1, self.get_u(u, t + self.h_norm / 2), p, s)[:, self.idx] class RK4(RK): @@ -434,10 +444,15 @@ def next_x(self, h: float, t: float, x_prev: MX | SX, u: MX | SX, p: MX | SX, s: The next integrate states """ - k1 = self.fun(x_prev, self.get_u(u, t), p, s, t)[:, self.idx] if 't' in self.fun.name_in() else self.fun(x_prev, self.get_u(u, t), p, s)[:, self.idx] - k2 = self.fun(x_prev + h / 2 * k1, self.get_u(u, t + self.h_norm / 2), p, s, t + self.h_norm / 2)[:,self.idx] if 't' in self.fun.name_in() else self.fun(x_prev + h / 2 * k1,self.get_u(u, t + self.h_norm / 2), p, s)[:,self.idx] - k3 = self.fun(x_prev + h / 2 * k2, self.get_u(u, t + self.h_norm / 2), p, s, t + self.h_norm / 2)[:,self.idx] if 't' in self.fun.name_in() else self.fun(x_prev + h / 2 * k2,self.get_u(u, t + self.h_norm / 2), p, s)[:,self.idx] - k4 = self.fun(x_prev + h * k3, self.get_u(u, t + self.h_norm), p, s, t + self.h_norm)[:,self.idx] if 't' in self.fun.name_in() else self.fun(x_prev + h * k3, self.get_u(u, t + self.h_norm), p, s)[:, self.idx] + # k1 = self.fun(x_prev, self.get_u(u, t), p, s, t)[:, self.idx] if 't' in self.fun.name_in() else self.fun(x_prev, self.get_u(u, t), p, s)[:, self.idx] + # k2 = self.fun(x_prev + h / 2 * k1, self.get_u(u, t + self.h_norm / 2), p, s, t + self.h_norm / 2)[:,self.idx] if 't' in self.fun.name_in() else self.fun(x_prev + h / 2 * k1,self.get_u(u, t + self.h_norm / 2), p, s)[:,self.idx] + # k3 = self.fun(x_prev + h / 2 * k2, self.get_u(u, t + self.h_norm / 2), p, s, t + self.h_norm / 2)[:,self.idx] if 't' in self.fun.name_in() else self.fun(x_prev + h / 2 * k2,self.get_u(u, t + self.h_norm / 2), p, s)[:,self.idx] + # k4 = self.fun(x_prev + h * k3, self.get_u(u, t + self.h_norm), p, s, t + self.h_norm)[:,self.idx] if 't' in self.fun.name_in() else self.fun(x_prev + h * k3, self.get_u(u, t + self.h_norm), p, s)[:, self.idx] + + k1 = self.fun(t, x_prev, self.get_u(u, t), p, s)[:, self.idx] + k2 = self.fun(t + self.h_norm / 2, x_prev + h / 2 * k1, self.get_u(u, t + self.h_norm / 2), p, s)[:, self.idx] + k3 = self.fun(t + self.h_norm / 2, x_prev + h / 2 * k2, self.get_u(u, t + self.h_norm / 2), p, s)[:, self.idx] + k4 = self.fun(t + self.h_norm, x_prev + h * k3, self.get_u(u, t + self.h_norm), p, s)[:, self.idx] return x_prev + h / 6 * (k1 + 2 * k2 + 2 * k3 + k4) @@ -490,34 +505,34 @@ def next_x(self, h: float, t: float, x_prev: MX | SX, u: MX | SX, p: MX | SX, s: The next integrate states """ - k1 = self.fun(x_prev, self.get_u(u, t), p, s)[:, self.idx] - k2 = self.fun(x_prev + (h * 4 / 27) * k1, self.get_u(u, t + self.h_norm * (4 / 27)), p, s)[:, self.idx] - k3 = self.fun(x_prev + (h / 18) * (k1 + 3 * k2), self.get_u(u, t + self.h_norm * (2 / 9)), p, s)[:, self.idx] - k4 = self.fun(x_prev + (h / 12) * (k1 + 3 * k3), self.get_u(u, t + self.h_norm * (1 / 3)), p, s)[:, self.idx] - k5 = self.fun(x_prev + (h / 8) * (k1 + 3 * k4), self.get_u(u, t + self.h_norm * (1 / 2)), p, s)[:, self.idx] + k1 = self.fun(t, x_prev, self.get_u(u, t), p, s)[:, self.idx] + k2 = self.fun(t, x_prev + (h * 4 / 27) * k1, self.get_u(u, t + self.h_norm * (4 / 27)), p, s)[:, self.idx] + k3 = self.fun(t, x_prev + (h / 18) * (k1 + 3 * k2), self.get_u(u, t + self.h_norm * (2 / 9)), p, s)[:, self.idx] + k4 = self.fun(t, x_prev + (h / 12) * (k1 + 3 * k3), self.get_u(u, t + self.h_norm * (1 / 3)), p, s)[:, self.idx] + k5 = self.fun(t, x_prev + (h / 8) * (k1 + 3 * k4), self.get_u(u, t + self.h_norm * (1 / 2)), p, s)[:, self.idx] k6 = self.fun( - x_prev + (h / 54) * (13 * k1 - 27 * k3 + 42 * k4 + 8 * k5), self.get_u(u, t + self.h_norm * (2 / 3)), p, s + t, x_prev + (h / 54) * (13 * k1 - 27 * k3 + 42 * k4 + 8 * k5), self.get_u(u, t + self.h_norm * (2 / 3)), p, s )[:, self.idx] k7 = self.fun( - x_prev + (h / 4320) * (389 * k1 - 54 * k3 + 966 * k4 - 824 * k5 + 243 * k6), + t, x_prev + (h / 4320) * (389 * k1 - 54 * k3 + 966 * k4 - 824 * k5 + 243 * k6), self.get_u(u, t + self.h_norm * (1 / 6)), p, s, )[:, self.idx] k8 = self.fun( - x_prev + (h / 20) * (-234 * k1 + 81 * k3 - 1164 * k4 + 656 * k5 - 122 * k6 + 800 * k7), + t, x_prev + (h / 20) * (-234 * k1 + 81 * k3 - 1164 * k4 + 656 * k5 - 122 * k6 + 800 * k7), self.get_u(u, t + self.h_norm), p, s, )[:, self.idx] k9 = self.fun( - x_prev + (h / 288) * (-127 * k1 + 18 * k3 - 678 * k4 + 456 * k5 - 9 * k6 + 576 * k7 + 4 * k8), + t, x_prev + (h / 288) * (-127 * k1 + 18 * k3 - 678 * k4 + 456 * k5 - 9 * k6 + 576 * k7 + 4 * k8), self.get_u(u, t + self.h_norm * (5 / 6)), p, s, )[:, self.idx] k10 = self.fun( - x_prev + t, x_prev + (h / 820) * (1481 * k1 - 81 * k3 + 7104 * k4 - 3376 * k5 + 72 * k6 - 5040 * k7 - 60 * k8 + 720 * k9), self.get_u(u, t + self.h_norm), p, @@ -540,7 +555,7 @@ class COLLOCATION(Integrator): ------- get_u(self, u: np.ndarray, dt_norm: float) -> np.ndarray Get the control at a given time - dxdt(self, h: float, states: MX | SX, controls: MX | SX, params: MX | SX, stochastic_variables: MX | SX) -> tuple[SX, list[SX]] + dxdt(self, h: float, time: MX | SX, states: MX | SX, controls: MX | SX, params: MX | SX, stochastic_variables: MX | SX) -> tuple[SX, list[SX]] The dynamics of the system """ @@ -624,6 +639,7 @@ def get_u(self, u: np.ndarray, dt_norm: float) -> np.ndarray: def dxdt( self, h: float, + time: MX | SX, states: MX | SX, controls: MX | SX, params: MX | SX, @@ -637,6 +653,8 @@ def dxdt( ---------- h: float The time step + time: MX | SX + The time of the system states: MX | SX The states of the system controls: MX | SX @@ -664,12 +682,13 @@ def dxdt( if self.defects_type == DefectType.EXPLICIT: f_j = self.fun( - states[j], self.get_u(controls, self.step_time[j]), params * param_scaling, stochastic_variables + time, states[j], self.get_u(controls, self.step_time[j]), params * param_scaling, stochastic_variables )[:, self.idx] defects.append(h * f_j - xp_j) elif self.defects_type == DefectType.IMPLICIT: defects.append( self.implicit_fun( + time, states[j], self.get_u(controls, self.step_time[j]), params * param_scaling, @@ -694,11 +713,11 @@ def _finish_init(self): self.function = Function( "integrator", - [horzcat(*self.x_sym), self.u_sym, self.param_sym, self.stochastic_variables_sym], + [self.t_sym, horzcat(*self.x_sym), self.u_sym, self.param_sym, self.stochastic_variables_sym], self.dxdt( - self.h, self.x_sym, self.u_sym, self.param_sym, self.param_scaling, self.stochastic_variables_sym + self.h, self.t_sym, self.x_sym, self.u_sym, self.param_sym, self.param_scaling, self.stochastic_variables_sym ), - ["x0", "p", "params", "s"], + ["t", "x0", "p", "params", "s"], ["xf", "xall", "defects"], ) @@ -730,6 +749,7 @@ def __init__(self, ode: dict, ode_opt: dict): def dxdt( self, h: float, + time: MX | SX, states: MX | SX, controls: MX | SX, params: MX | SX, @@ -743,6 +763,8 @@ def dxdt( ---------- h: float The time step + time: MX | SX + The time of the system states: MX | SX The states of the system controls: MX | SX @@ -760,16 +782,16 @@ def dxdt( """ nx = states[0].shape[0] - _, _, defect = super(IRK, self).dxdt(h, states, controls, params, param_scaling, stochastic_variables) + _, _, defect = super(IRK, self).dxdt(h, time, states, controls, params, param_scaling, stochastic_variables) # Root-finding function, implicitly defines x_collocation_points as a function of x0 and p vfcn = Function( - "vfcn", [vertcat(*states[1:]), states[0], controls, params, stochastic_variables], [defect] + "vfcn", [vertcat(*states[1:]), time, states[0], controls, params, stochastic_variables], [defect] ).expand() # Create a implicit function instance to solve the system of equations ifcn = rootfinder("ifcn", "newton", vfcn) - x_irk_points = ifcn(self.cx(), states[0], controls, params, stochastic_variables) + x_irk_points = ifcn(self.cx(), time, states[0], controls, params, stochastic_variables) x = [states[0] if r == 0 else x_irk_points[(r - 1) * nx : r * nx] for r in range(self.degree + 1)] # Get an expression for the state at the end of the finite element @@ -786,11 +808,11 @@ def _finish_init(self): self.function = Function( "integrator", - [self.x_sym[0], self.u_sym, self.param_sym, self.stochastic_variables_sym], + [self.t_sym, self.x_sym[0], self.u_sym, self.param_sym, self.stochastic_variables_sym], self.dxdt( - self.h, self.x_sym, self.u_sym, self.param_sym, self.param_scaling, self.stochastic_variables_sym + self.h, self.t_sym, self.x_sym, self.u_sym, self.param_sym, self.param_scaling, self.stochastic_variables_sym ), - ["x0", "p", "params", "s"], + ["t", "x0", "p", "params", "s"], ["xf", "xall"], ) diff --git a/bioptim/dynamics/ode_solver.py b/bioptim/dynamics/ode_solver.py index 5ba3b81b0..a5f43e42a 100644 --- a/bioptim/dynamics/ode_solver.py +++ b/bioptim/dynamics/ode_solver.py @@ -119,7 +119,7 @@ def integrator(self, ocp, nlp, node_index: int, t: float | MX | SX) -> list: ------- A list of integrators """ - + nlp.time.node_index = node_index nlp.states.node_index = node_index nlp.states_dot.node_index = node_index nlp.controls.node_index = node_index @@ -138,6 +138,7 @@ def integrator(self, ocp, nlp, node_index: int, t: float | MX | SX) -> list: } ode = { + "time": nlp.time.cx_start, "x_unscaled": nlp.states.cx_start, "x_scaled": nlp.states.scaled.cx_start, "p_unscaled": nlp.controls.cx_start @@ -149,7 +150,6 @@ def integrator(self, ocp, nlp, node_index: int, t: float | MX | SX) -> list: "stochastic_variables": nlp.stochastic_variables.cx_start, "ode": nlp.dynamics_func, "implicit_ode": nlp.implicit_dynamics_func, - "time": t } if ode["ode"].size2_out("xdot") != 1: @@ -289,7 +289,7 @@ def integrator(self, ocp, nlp, node_index: int, t: float | MX | SX) -> list: ------- A list of integrators """ - + nlp.time.node_index = node_index nlp.states.node_index = node_index nlp.states_dot.node_index = node_index nlp.controls.node_index = node_index @@ -305,6 +305,7 @@ def integrator(self, ocp, nlp, node_index: int, t: float | MX | SX) -> list: ) ode = { + "time": nlp.time.cx_start, "x_unscaled": [nlp.states.cx_start] + nlp.states.cx_intermediates_list, "x_scaled": [nlp.states.scaled.cx_start] + nlp.states.scaled.cx_intermediates_list, "p_unscaled": nlp.controls.cx_start, @@ -312,7 +313,7 @@ def integrator(self, ocp, nlp, node_index: int, t: float | MX | SX) -> list: "stochastic_variables": nlp.stochastic_variables.cx_start, "ode": nlp.dynamics_func, "implicit_ode": nlp.implicit_dynamics_func, - "time": t + # "time": t } ode_opt = { "t0": 0, @@ -422,7 +423,7 @@ def integrator(self, ocp, nlp, node_index: int, t: float | MX | SX) -> list: ------- A list of integrators """ - + nlp.time.node_index = node_index nlp.states.node_index = node_index nlp.states_dot.node_index = node_index nlp.controls.node_index = node_index @@ -438,15 +439,16 @@ def integrator(self, ocp, nlp, node_index: int, t: float | MX | SX) -> list: raise RuntimeError("CVODES cannot be used with piece-wise linear controls (only RK4)") ode = { + "time": nlp.time.cx_start, "x": nlp.states.scaled.cx_start, "p": nlp.controls.scaled.cx_start, "ode": nlp.dynamics_func( + nlp.time.scaled.cx_start, nlp.states.scaled.cx_start, nlp.controls.scaled.cx_start, nlp.parameters.cx, nlp.stochastic_variables.cx_start, ), - "time": t, } ode_opt = {"t0": 0, "tf": nlp.dt} @@ -456,6 +458,7 @@ def integrator(self, ocp, nlp, node_index: int, t: float | MX | SX) -> list: Function( "integrator", [ + nlp.time, nlp.states.scaled.cx_start, nlp.controls.scaled.cx_start, nlp.parameters.cx, @@ -466,7 +469,7 @@ def integrator(self, ocp, nlp, node_index: int, t: float | MX | SX) -> list: nlp.states.scaled.cx_start, nlp.controls.scaled.cx_start, ), - ["x0", "p", "params", "s"], + ["t", "x0", "p", "params", "s"], ["xf", "xall"], ) ] diff --git a/bioptim/examples/custom_model/custom_package/custom_dynamics.py b/bioptim/examples/custom_model/custom_package/custom_dynamics.py index fa2f5bc00..4e1aba873 100644 --- a/bioptim/examples/custom_model/custom_package/custom_dynamics.py +++ b/bioptim/examples/custom_model/custom_package/custom_dynamics.py @@ -14,6 +14,7 @@ def custom_dynamics( + time: MX, states: MX, controls: MX, parameters: MX, diff --git a/bioptim/examples/getting_started/custom_dynamics.py b/bioptim/examples/getting_started/custom_dynamics.py index 15cc0bdd7..19f2d5bff 100644 --- a/bioptim/examples/getting_started/custom_dynamics.py +++ b/bioptim/examples/getting_started/custom_dynamics.py @@ -32,6 +32,7 @@ def custom_dynamic( + time: MX | SX, states: MX | SX, controls: MX | SX, parameters: MX | SX, @@ -45,6 +46,8 @@ def custom_dynamic( Parameters ---------- + time: MX | SX + The time of the system states: MX | SX The state of the system controls: MX | SX diff --git a/bioptim/examples/muscle_driven_ocp/muscle_activations_tracker.py b/bioptim/examples/muscle_driven_ocp/muscle_activations_tracker.py index 8c285c3cd..7cff3679c 100644 --- a/bioptim/examples/muscle_driven_ocp/muscle_activations_tracker.py +++ b/bioptim/examples/muscle_driven_ocp/muscle_activations_tracker.py @@ -164,6 +164,7 @@ def generate_data( dynamics_func = biorbd.to_casadi_func( "ForwardDyn", dyn_func( + time=MX(), states=symbolic_states, controls=symbolic_controls, parameters=symbolic_parameters, diff --git a/bioptim/examples/muscle_driven_ocp/muscle_excitations_tracker.py b/bioptim/examples/muscle_driven_ocp/muscle_excitations_tracker.py index 26bcbd9d9..2bb8cd843 100644 --- a/bioptim/examples/muscle_driven_ocp/muscle_excitations_tracker.py +++ b/bioptim/examples/muscle_driven_ocp/muscle_excitations_tracker.py @@ -82,6 +82,7 @@ def generate_data( symbolic_tau = MX.sym("tau", n_tau, 1) symbolic_mus_controls = MX.sym("mus", n_mus, 1) + symbolic_time = MX.sym("t", 0, 0) symbolic_states = vertcat(*(symbolic_q, symbolic_qdot, symbolic_mus_states)) symbolic_controls = vertcat(*(symbolic_tau, symbolic_mus_controls)) @@ -166,6 +167,7 @@ def generate_data( dynamics_func = biorbd.to_casadi_func( "ForwardDyn", DynamicsFunctions.muscles_driven( + time=symbolic_time, states=symbolic_states, controls=symbolic_controls, parameters=symbolic_parameters, @@ -174,6 +176,7 @@ def generate_data( with_contact=False, rigidbody_dynamics=RigidBodyDynamics.ODE, ).dxdt, + symbolic_time, symbolic_states, symbolic_controls, symbolic_parameters, @@ -183,7 +186,7 @@ def generate_data( def dyn_interface(t, x, u): u = np.concatenate([np.zeros(n_tau), u]) - return np.array(dynamics_func(x, u, [])[:, 0]).squeeze() + return np.array(dynamics_func(t, x, u, [])[:, 0]).squeeze() # Generate some muscle excitations U = np.random.rand(n_shooting, n_mus).T diff --git a/bioptim/examples/torque_driven_ocp/spring_load.py b/bioptim/examples/torque_driven_ocp/spring_load.py index 3c0f843d0..6fdfe4930 100644 --- a/bioptim/examples/torque_driven_ocp/spring_load.py +++ b/bioptim/examples/torque_driven_ocp/spring_load.py @@ -24,13 +24,15 @@ def custom_dynamic( - states: MX, controls: MX, parameters: MX, stochastic_variables: MX, nlp: NonLinearProgram + time: MX, states: MX, controls: MX, parameters: MX, stochastic_variables: MX, nlp: NonLinearProgram ) -> DynamicsEvaluation: """ The dynamics of the system using an external force (see custom_dynamics for more explanation) Parameters ---------- + time: MX + The current time of the system states: MX The current states of the system controls: MX diff --git a/bioptim/gui/check_conditioning.py b/bioptim/gui/check_conditioning.py index a54b0df0d..5dde8673c 100644 --- a/bioptim/gui/check_conditioning.py +++ b/bioptim/gui/check_conditioning.py @@ -59,6 +59,7 @@ def jacobian_hessian_constraints(): list_constraints = [] for constraints in nlp.g: node_index = constraints.node_idx[0] # TODO deal with assume_phase_dynamics=False + nlp.time.node_index = node_index nlp.states.node_index = node_index nlp.states_dot.node_index = node_index nlp.controls.node_index = node_index @@ -72,13 +73,14 @@ def jacobian_hessian_constraints(): ): # depends if there are parameters if nlp.parameters.shape == 0: - vertcat_obj = vertcat(*nlp.X_scaled, *nlp.U_scaled, nlp.parameters.cx, *nlp.S) + vertcat_obj = vertcat(*nlp.T_scaled, *nlp.X_scaled, *nlp.U_scaled, nlp.parameters.cx, *nlp.S) else: - vertcat_obj = vertcat(*nlp.X_scaled, *nlp.U_scaled, *[nlp.parameters.cx, *nlp.S]) + vertcat_obj = vertcat(*nlp.T_scaled, *nlp.X_scaled, *nlp.U_scaled, *[nlp.parameters.cx, *nlp.S]) list_constraints.append( jacobian( constraints.function[constraints.node_idx[0]]( + nlp.time.cx_start, nlp.states.cx_start, nlp.controls.cx_start, nlp.parameters.cx, @@ -92,9 +94,9 @@ def jacobian_hessian_constraints(): # depends if there are parameters if nlp.parameters.shape == 0: - vertcat_obj = vertcat(*nlp.X_scaled, *nlp.U_scaled, nlp.parameters.cx, *nlp.S) + vertcat_obj = vertcat(*nlp.T_scaled, *nlp.X_scaled, *nlp.U_scaled, nlp.parameters.cx, *nlp.S) else: - vertcat_obj = vertcat(*nlp.X_scaled, *nlp.U_scaled, *[nlp.parameters.cx], *nlp.S) + vertcat_obj = vertcat(*nlp.T_scaled, *nlp.X_scaled, *nlp.U_scaled, *[nlp.parameters.cx], *nlp.S) jac_func = Function( "jacobian", @@ -102,16 +104,23 @@ def jacobian_hessian_constraints(): [jacobian_cas], ) + nb_t_init = sum([nlp.t_init[key].shape[0] for key in nlp.t_init.keys()]) nb_x_init = sum([nlp.x_init[key].shape[0] for key in nlp.x_init.keys()]) nb_u_init = sum([nlp.u_init[key].shape[0] for key in nlp.u_init.keys()]) nb_s_init = sum([nlp.s_init[key].shape[0] for key in nlp.s_init.keys()]) # evaluate jac_func at X_init, U_init, considering the parameters + t_init = np.zeros((len(nlp.T), nb_t_init)) x_init = np.zeros((len(nlp.X), nb_x_init)) u_init = np.zeros((len(nlp.U), nb_u_init)) param_init = np.array([ocp.parameter_init[key].shape[0] for key in ocp.parameter_init.keys()]) s_init = np.zeros((len(nlp.S), nb_s_init)) + for key in nlp.time.keys(): + nlp.t_init[key].check_and_adjust_dimensions(len(nlp.time[key]), nlp.ns + 1) + for node_index in range(nlp.ns + 1): + nlp.time.node_index = node_index + t_init[node_index, nlp.time[key].index] = np.array(nlp.t_init[key].init.evaluate_at(node_index)) for key in nlp.states.keys(): nlp.x_init[key].check_and_adjust_dimensions(len(nlp.states[key]), nlp.ns + 1) for node_index in range(nlp.ns + 1): @@ -130,12 +139,13 @@ def jacobian_hessian_constraints(): nlp.s_init[key].init.evaluate_at(node_index) ) + t_init = t_init.reshape((t_init.size, 1)) x_init = x_init.reshape((x_init.size, 1)) u_init = u_init.reshape((u_init.size, 1)) param_init = param_init.reshape((param_init.size, 1)) s_init = s_init.reshape((s_init.size, 1)) - jacobian_matrix = np.array(jac_func(np.vstack((x_init, u_init, param_init, s_init)))) + jacobian_matrix = np.array(jac_func(np.vstack((t_init, x_init, u_init, param_init, s_init)))) jacobian_list.append(jacobian_matrix) @@ -152,6 +162,7 @@ def jacobian_hessian_constraints(): list_norm = [] for constraints in nlp.g: node_index = constraints.node_idx[0] # TODO deal with assume_phase_dynamics=False + nlp.time.node_index = node_index nlp.states.node_index = node_index nlp.states_dot.node_index = node_index nlp.controls.node_index = node_index @@ -160,19 +171,20 @@ def jacobian_hessian_constraints(): for axis in range( 0, constraints.function[node_index]( - nlp.states.cx_start, nlp.controls.cx_start, nlp.parameters.cx, nlp.stochastic_variables.cx_start + nlp.time.cx_start, nlp.states.cx_start, nlp.controls.cx_start, nlp.parameters.cx, nlp.stochastic_variables.cx_start ).shape[0], ): # find all equality constraints if constraints.bounds.min[axis][0] == constraints.bounds.max[axis][0]: # parameters if nlp.parameters.shape == 0: - vertcat_obj = vertcat(*nlp.X_scaled, *nlp.U_scaled, nlp.parameters.cx, *nlp.S) + vertcat_obj = vertcat(*nlp.T_scaled, *nlp.X_scaled, *nlp.U_scaled, nlp.parameters.cx, *nlp.S) else: - vertcat_obj = vertcat(*nlp.X_scaled, *nlp.U_scaled, *[nlp.parameters.cx, *nlp.S]) + vertcat_obj = vertcat(*nlp.T_scaled, *nlp.X_scaled, *nlp.U_scaled, *[nlp.parameters.cx, *nlp.S]) hessian_cas = hessian( constraints.function[node_index]( + nlp.time.cx_start, nlp.states.cx_start, nlp.controls.cx_start, nlp.parameters.cx, @@ -189,7 +201,7 @@ def jacobian_hessian_constraints(): [hessian_cas], ) - hessian_matrix = np.array(hes_func(np.vstack((x_init, u_init, param_init, s_init)))) + hessian_matrix = np.array(hes_func(np.vstack((t_init, x_init, u_init, param_init, s_init)))) # append hessian list list_hessian.append(hessian_matrix) @@ -304,6 +316,7 @@ def hessian_objective(): objective = 0 node_index = obj.node_idx[0] # TODO deal with assume_phase_dynamics=False + nlp.time.node_index = node_index nlp.states.node_index = node_index nlp.states_dot.node_index = node_index nlp.controls.node_index = node_index @@ -313,12 +326,15 @@ def hessian_objective(): if obj.multinode_penalty or obj.transition: phase = ocp.nlp[phase - 1] nlp_post = nlp + time_pre = phase.time_cx_end + time_post = nlp_post.time_cx_start states_pre = phase.states.cx_end states_post = nlp_post.states.cx_start controls_pre = phase.controls.cx_end controls_post = nlp_post.controls.cx_start stochastic_pre = phase.stochastic_cx_end stochastic_post = nlp_post.stochastic_cx_start + time_cx = vertcat(time_pre, time_post) state_cx = vertcat(states_pre, states_post) control_cx = vertcat(controls_pre, controls_post) stochastic_cx = vertcat(stochastic_pre, stochastic_post) @@ -328,16 +344,19 @@ def hessian_objective(): state_cx = horzcat(*([nlp.states.cx_start] + nlp.states.cx_intermediates_list)) else: state_cx = nlp.states.cx_start + time_cx = nlp.time.cx_start control_cx = nlp.controls.cx_start stochastic_cx = nlp.stochastic_variables.cx_start if obj.explicit_derivative: if obj.derivative: raise RuntimeError("derivative and explicit_derivative cannot be simultaneously true") + time_cx = horzcat(time_cx, nlp.time_cx_end) state_cx = horzcat(state_cx, nlp.states.cx_end) control_cx = horzcat(control_cx, nlp.controls.cx_end) stochastic_cx = horzcat(stochastic_cx, nlp.stochastic_cx_end) if obj.derivative: + time_cx = horzcat(nlp.time.cx_end, nlp.time.cx_start) state_cx = horzcat(nlp.states.cx_end, nlp.states.cx_start) control_cx = horzcat(nlp.controls.cx_end, nlp.controls.cx_start) stochastic_cx = horzcat(nlp.stochastic.cx_end, nlp.stochastic.cx_start) @@ -349,6 +368,11 @@ def hessian_objective(): ) if is_trapezoidal: + time_cx = ( + horzcat(nlp.time.cx_start, nlp.time.cx_end) + if obj.integration_rule == QuadratureRule.APPROXIMATE_TRAPEZOIDAL + else nlp.time.cx_start + ) state_cx = ( horzcat(nlp.states.cx_start, nlp.states.cx_end) if obj.integration_rule == QuadratureRule.APPROXIMATE_TRAPEZOIDAL @@ -367,6 +391,7 @@ def hessian_objective(): if obj.target is None: p = obj.weighted_function[node_index]( + time_cx, state_cx, control_cx, nlp.parameters.cx, @@ -377,6 +402,7 @@ def hessian_objective(): ) else: p = obj.weighted_function[node_index]( + time_cx, state_cx, control_cx, nlp.parameters.cx, @@ -392,9 +418,9 @@ def hessian_objective(): # create function to build the hessian # parameters if nlp.parameters.shape == 0: - vertcat_obj = vertcat(*nlp.X_scaled, *nlp.U_scaled, nlp.parameters.cx, *nlp.S) + vertcat_obj = vertcat(*nlp.T_scaled, *nlp.X_scaled, *nlp.U_scaled, nlp.parameters.cx, *nlp.S) else: - vertcat_obj = vertcat(*nlp.X_scaled, *nlp.U_scaled, *[nlp.parameters.cx], *nlp.S) + vertcat_obj = vertcat(*nlp.T_scaled, *nlp.X_scaled, *nlp.U_scaled, *[nlp.parameters.cx], *nlp.S) hessian_cas = hessian(objective, vertcat_obj)[0] @@ -404,16 +430,23 @@ def hessian_objective(): [hessian_cas], ) + nb_t_init = sum([nlp.t_init[key].shape[0] for key in nlp.t_init.keys()]) nb_x_init = sum([nlp.x_init[key].shape[0] for key in nlp.x_init.keys()]) nb_u_init = sum([nlp.u_init[key].shape[0] for key in nlp.u_init.keys()]) nb_s_init = sum([nlp.s_init[key].shape[0] for key in nlp.s_init.keys()]) # evaluate jac_func at X_init, U_init, considering the parameters + t_init = np.zeros((len(nlp.T), nb_t_init)) x_init = np.zeros((len(nlp.X), nb_x_init)) u_init = np.zeros((len(nlp.U), nb_u_init)) param_init = np.array([nlp.x_init[key].shape[0] for key in ocp.parameter_init.keys()]) s_init = np.zeros((len(nlp.S), nb_s_init)) + for key in nlp.time.keys(): + nlp.t_init[key].check_and_adjust_dimensions(len(nlp.time[key]), nlp.ns + 1) + for node_index in range(nlp.ns + 1): + nlp.time.node_index = node_index + t_init[node_index, nlp.time[key].index] = np.array(nlp.t_init[key].init.evaluate_at(node_index)) for key in nlp.states.keys(): nlp.x_init[key].check_and_adjust_dimensions(len(nlp.states[key]), nlp.ns + 1) for node_index in range(nlp.ns + 1): @@ -432,12 +465,13 @@ def hessian_objective(): nlp.s_init[key].init.evaluate_at(node_index) ) + t_init = t_init.reshape((t_init.size, 1)) x_init = x_init.reshape((x_init.size, 1)) u_init = u_init.reshape((u_init.size, 1)) param_init = param_init.reshape((param_init.size, 1)) s_init = s_init.reshape((s_init.size, 1)) - hessian_obj_matrix = np.array(hes_func(np.vstack((x_init, u_init, param_init, s_init)))) + hessian_obj_matrix = np.array(hes_func(np.vstack((t_init, x_init, u_init, param_init, s_init)))) hessian_obj_list.append(hessian_obj_matrix) # Convexity checking (positive semi-definite hessian) diff --git a/bioptim/gui/graph.py b/bioptim/gui/graph.py index 97d1ce2fc..f34fc1aca 100644 --- a/bioptim/gui/graph.py +++ b/bioptim/gui/graph.py @@ -279,6 +279,7 @@ def print(self): """ for phase_idx in range(self.ocp.n_phases): # We only need to use the first index since the bounds are not depend on the dynamics + self.ocp.nlp[phase_idx].time.node_index = 0 self.ocp.nlp[phase_idx].states.node_index = 0 self.ocp.nlp[phase_idx].states_dot.node_index = 0 self.ocp.nlp[phase_idx].controls.node_index = 0 diff --git a/bioptim/gui/plot.py b/bioptim/gui/plot.py index f48207070..b619d4e36 100644 --- a/bioptim/gui/plot.py +++ b/bioptim/gui/plot.py @@ -22,7 +22,7 @@ class CustomPlot: Attributes ---------- - function: Callable[states, controls, parameters, stochastic_variables] + function: Callable[time, states, controls, parameters, stochastic_variables] The function to call to update the graph type: PlotType Type of plot to use @@ -66,7 +66,7 @@ def __init__( """ Parameters ---------- - update_function: Callable[states, controls, parameters, stochastic_variables] + update_function: Callable[time, states, controls, parameters, stochastic_variables] The function to call to update the graph plot_type: PlotType Type of plot to use @@ -340,12 +340,14 @@ def legend_without_duplicate_labels(ax): node_index = 0 # TODO deal with assume_phase_dynamics=False if nlp.plot[key].node_idx is not None: node_index = nlp.plot[key].node_idx[0] + nlp.time.node_index = node_index nlp.states.node_index = node_index nlp.states_dot.node_index = node_index nlp.controls.node_index = node_index nlp.stochastic_variables.node_index = node_index # If multi-node penalties = None, stays zero + size_t = 0 size_x = 0 size_u = 0 size_p = 0 @@ -354,16 +356,19 @@ def legend_without_duplicate_labels(ax): casadi_function = nlp.plot[key].parameters["penalty"].weighted_function_non_threaded[0] if nlp.plot[key].parameters["penalty"].multinode_penalty: if casadi_function is not None: - size_x = len(casadi_function.nominal_in(0)) - size_u = len(casadi_function.nominal_in(1)) - size_p = len(casadi_function.nominal_in(2)) - size_s = len(casadi_function.nominal_in(3)) + size_t = len(casadi_function.nominal_in(0)) + size_x = len(casadi_function.nominal_in(1)) + size_u = len(casadi_function.nominal_in(2)) + size_p = len(casadi_function.nominal_in(3)) + size_s = len(casadi_function.nominal_in(4)) else: + size_t = nlp.time.shape size_x = nlp.states.shape size_u = nlp.controls.shape size_p = nlp.parameters.shape size_s = nlp.stochastic_variables.shape else: + size_t = nlp.time.shape size_x = nlp.states.shape size_u = nlp.controls.shape size_p = nlp.parameters.shape @@ -373,6 +378,7 @@ def legend_without_duplicate_labels(ax): nlp.plot[key] .function( node_index, + np.zeros((size_t, 1)), np.zeros((size_x, 1)), np.zeros((size_u, 1)), np.zeros((size_p, 1)), @@ -707,7 +713,12 @@ def update_data(self, v: dict): if self.integrator != SolutionIntegrator.OCP else nlp.ode_solver.steps + 1 ) - + if isinstance(data_time, dict): + n_elements = data_time[list(data_time.keys())[0]].shape[1] + elif isinstance(data_time, list): + n_elements = data_time[i][list(data_time[i].keys())[0]].shape[1] + else: + raise RuntimeError("Invalid data_time type") if isinstance(data_states, dict): n_elements = data_states[list(data_states.keys())[0]].shape[1] elif isinstance(data_states, list): @@ -774,6 +785,7 @@ def update_data(self, v: dict): val_tempo = self.plot_func[key][i].function( idx, + time[:, idx: idx + 1 + 1], state[:, step_size * idx : step_size * (idx + 1) + x_mod], control[:, idx : idx + u_mod + 1], data_params_in_dyn, @@ -809,6 +821,22 @@ def update_data(self, v: dict): y = np.array([np.nan]) penalty: MultinodeConstraint = self.plot_func[key][i].parameters["penalty"] + t_phase = np.ndarray((0, len(penalty.nodes_phase))) + if sol.ocp.n_phases == 1 and isinstance(data_time, dict): + data_time = [data_time] + for time_key in data_time[i].keys(): + t_phase_tp = np.ndarray((data_time[i][time_key].shape[0], 0)) + for tp in range(len(penalty.nodes_phase)): + phase_tp = penalty.nodes_phase[tp] + node_idx_tp = penalty.all_nodes_index[tp] + t_phase_tp = np.hstack( + ( + t_phase_tp, + data_time[phase_tp][time_key][:, node_idx_tp][:, np.newaxis], + ) + ) + t_phase = np.vstack((t_phase, t_phase_tp)) + x_phase = np.ndarray((0, len(penalty.nodes_phase))) if sol.ocp.n_phases == 1 and isinstance(data_states, dict): data_states = [data_states] @@ -862,6 +890,7 @@ def update_data(self, v: dict): val_tempo = self.plot_func[key][i].function( self.plot_func[key][i].node_idx[0], + t_phase, x_phase, u_phase, data_params_in_dyn, @@ -878,6 +907,12 @@ def update_data(self, v: dict): raise NotImplementedError("See todo!") val = self.plot_func[key][i].function( node_idx, + np.hstack( + ( + data_time[node_idx]["all"][:, -1], + data_time[node_idx + 1]["all"][:, 0], + ) + ), np.hstack( ( data_states[node_idx]["all"][:, -1], @@ -900,6 +935,7 @@ def update_data(self, v: dict): **self.plot_func[key][i].parameters, ) else: + time_tp = stochastic[:, node_idx: node_idx + 1 + 1] if ( self.plot_func[key][i].label == "CONTINUITY" and nlp.ode_solver.is_direct_collocation @@ -915,6 +951,7 @@ def update_data(self, v: dict): stochastic_tp = stochastic[:, node_idx : node_idx + 1 + 1] val = self.plot_func[key][i].function( node_idx, + time_tp, states, control_tp, data_params_in_dyn, @@ -933,6 +970,7 @@ def update_data(self, v: dict): for i_node, node_idx in enumerate(self.plot_func[key][i].node_idx): val_tempo = self.plot_func[key][i].function( node_idx, + time[:, node_idx: node_idx + 1 + 1], state[:, node_idx * step_size : (node_idx + 1) * step_size + 1 : step_size], control[:, node_idx : node_idx + 1 + 1], data_params_in_dyn, @@ -960,6 +998,7 @@ def update_data(self, v: dict): val_tempo = self.plot_func[key][i].function( nodes, + time[:, node_idx: node_idx + 1 + 1], state[:, ::step_size], control, data_params_in_dyn, diff --git a/bioptim/interfaces/interface_utils.py b/bioptim/interfaces/interface_utils.py index dc52e67d7..a865b9001 100644 --- a/bioptim/interfaces/interface_utils.py +++ b/bioptim/interfaces/interface_utils.py @@ -221,6 +221,7 @@ def get_control_modificator(index): ) # Make an exception to the fact that U is not available for the last node + _t = ocp.cx() _x = ocp.cx() _u = ocp.cx() _s = ocp.cx() @@ -239,10 +240,12 @@ def get_control_modificator(index): if ocp.assume_phase_dynamics or index_i < len(nlp_i.U_scaled) else [] ) + _t_tp = nlp_i.T[index_i] _s_tp = nlp_i.S[index_i] # 0th column since this constraint can only be applied to a single point. This is to account for # the COLLOCATION which will have multiple column, but are not intended to be used here + _t = vertcat(_t, _t_tp) _x = vertcat(_x, _x_tp[:, 0]) _u = vertcat(_u, _u_tp) _s = vertcat(_s, _s_tp) @@ -254,6 +257,7 @@ def get_control_modificator(index): else: _x = nlp.X_scaled[_idx] _u = nlp.U_scaled[_idx][:, 0] if _idx < len(nlp.U_scaled) else [] + _t = nlp.T[_idx] _s = nlp.S[_idx] else: if is_unscaled: @@ -265,6 +269,7 @@ def get_control_modificator(index): _u = [] else: _u = nlp.U_scaled[_idx][:, 0] if _idx < len(nlp.U_scaled) else [] + _t = nlp.T[_idx][:, 0] _s = nlp.S[_idx][:, 0] if _penalty.derivative or _penalty.explicit_derivative: @@ -275,8 +280,10 @@ def get_control_modificator(index): else: x = nlp.X_scaled[_idx + 1][:, 0] u = nlp.U_scaled[_idx + 1][:, 0] if _idx + 1 < len(nlp.U_scaled) else [] + t = nlp.T[_idx + 1][:, 0] s = nlp.S[_idx + 1][:, 0] + _t = horzcat(_t, t) _x = horzcat(_x, x) _u = horzcat(_u, u) _s = horzcat(_s, s) @@ -293,6 +300,8 @@ def get_control_modificator(index): else: u = nlp.U_scaled[_idx + 1][:, 0] if _idx + 1 < len(nlp.U_scaled) else [] _u = horzcat(_u, u) + t = nlp.T[_idx + 1][:, 0] + _t = horzcat(_t, t) s = nlp.S[_idx + 1][:, 0] _s = horzcat(_s, s) @@ -303,7 +312,7 @@ def get_control_modificator(index): else: u = nlp.U_scaled[_idx + 1][:, 0] if _idx + 1 < len(nlp.U_scaled) else [] _u = horzcat(_u, u) - return _x, _u, _s + return _t, _x, _u, _s param = interface.ocp.cx(interface.ocp.parameters.cx) out = interface.ocp.cx() @@ -316,11 +325,13 @@ def get_control_modificator(index): raise NotImplementedError("multi_thread penalty with target shape != [n x m] is not implemented yet") target = penalty.target[0] if penalty.target is not None else [] + t = nlp.cx() x = nlp.cx() u = nlp.cx() s = nlp.cx() for idx in penalty.node_idx: - x_tp, u_tp, s_tp = get_x_and_u_at_idx(penalty, idx, is_unscaled) + t_tp, x_tp, u_tp, s_tp = get_x_and_u_at_idx(penalty, idx, is_unscaled) + t = horzcat(t, t_tp) x = horzcat(x, x_tp) u = horzcat(u, u_tp) s = horzcat(s, s_tp) @@ -330,7 +341,7 @@ def get_control_modificator(index): u = horzcat(u, u[:, -1]) # We can call penalty.weighted_function[0] since multi-thread declares all the node at [0] - p = reshape(penalty.weighted_function[0](x, u, param, s, penalty.weight, target, penalty.dt), -1, 1) + p = reshape(penalty.weighted_function[0](t, x, u, param, s, penalty.weight, target, penalty.dt), -1, 1) else: p = interface.ocp.cx() @@ -351,12 +362,13 @@ def get_control_modificator(index): continue if not nlp: + t = [] x = [] u = [] s = [] else: - x, u, s = get_x_and_u_at_idx(penalty, idx, is_unscaled) - p = vertcat(p, penalty.weighted_function[idx](x, u, param, s, penalty.weight, target, penalty.dt)) + t, x, u, s = get_x_and_u_at_idx(penalty, idx, is_unscaled) + p = vertcat(p, penalty.weighted_function[idx](t, x, u, param, s, penalty.weight, target, penalty.dt)) out = vertcat(out, sum2(p)) return out diff --git a/bioptim/interfaces/solve_ivp_interface.py b/bioptim/interfaces/solve_ivp_interface.py index a51f8e15f..f2d734f61 100644 --- a/bioptim/interfaces/solve_ivp_interface.py +++ b/bioptim/interfaces/solve_ivp_interface.py @@ -8,6 +8,7 @@ def solve_ivp_interface( dynamics_func: Callable, t_eval: np.ndarray | List[float], + t: np.ndarray, x0: np.ndarray, u: np.ndarray, params: np.ndarray, @@ -25,6 +26,8 @@ def solve_ivp_interface( function that computes the dynamics of the system t_eval : np.ndarray | List[float] array of times t the controls u are evaluated at + t : np.ndarray + array of time x0 : np.ndarray array of initial conditions u : np.ndarray @@ -66,6 +69,7 @@ def solve_ivp_interface( y = run_solve_ivp( dynamics_func=dynamics_func, t_eval=t_eval_step, + t=t, x0=x0i, u=ui, s=s, @@ -113,6 +117,7 @@ def solve_ivp_interface( y = run_solve_ivp( dynamics_func=dynamics_func, t_eval=t_eval_step, + t=t, x0=x0i, u=ui, s=s, @@ -132,6 +137,7 @@ def solve_ivp_interface( def run_solve_ivp( dynamics_func: Callable, t_eval: np.ndarray | List[float], + t: np.ndarray, x0: np.ndarray, u: np.ndarray, params: np.ndarray, @@ -149,6 +155,8 @@ def run_solve_ivp( function that computes the dynamics of the system t_eval : np.ndarray | List[float] array of times t the controls u are evaluated at + t : np.ndarray + array of time x0 : np.ndarray array of initial conditions u : np.ndarray @@ -176,7 +184,7 @@ def run_solve_ivp( t_span = [t_eval[0], t_eval[-1]] integrated_sol = solve_ivp( - lambda t, x: np.array(dynamics_func(x, control_function(t), params, s))[:, 0], + lambda t, x: np.array(dynamics_func(t, x, control_function(t), params, s))[:, 0], t_span=t_span, y0=x0, t_eval=np.array(t_eval, dtype=np.float64), # prevent error with dtype=object @@ -310,6 +318,7 @@ def previous_t_except_the_last_one(t: float, t_eval: np.ndarray | List[float]) - def solve_ivp_bioptim_interface( dynamics_func: list[Callable], keep_intermediate_points: bool, + t: np.ndarray, x0: np.ndarray, u: np.ndarray, params: np.ndarray, @@ -327,6 +336,8 @@ def solve_ivp_bioptim_interface( function that computes the dynamics of the system keep_intermediate_points : bool whether to keep the intermediate points or not + t : np.ndarray + array of time x0 : np.ndarray array of initial conditions u : np.ndarray @@ -366,7 +377,7 @@ def solve_ivp_bioptim_interface( np.array([], dtype=np.float64).reshape(x0i.shape[0], 0) if keep_intermediate_points else x0i, # x0 or None - np.array(func(x0=x0i, p=u_controls, params=params / param_scaling, s=s)[dynamics_output]), + np.array(func(t=t, x0=x0i, p=u_controls, params=params / param_scaling, s=s)[dynamics_output]), ), # xf or xall axis=1, ) diff --git a/bioptim/limits/constraints.py b/bioptim/limits/constraints.py index 53b4a38a7..535a2b5f8 100644 --- a/bioptim/limits/constraints.py +++ b/bioptim/limits/constraints.py @@ -244,6 +244,7 @@ def non_slipping( constraint.max_bound = np.array([np.inf, np.inf]) contact = controller.get_nlp.contact_forces_func( + controller.time.cx_start, controller.states.cx_start, controller.controls.cx_start, controller.parameters.cx, @@ -655,6 +656,7 @@ def stochastic_dg_dx_implicit( sensory_noise = MX.sym("sensory_noise", sensory_noise_magnitude.shape[0], 1) dx = dynamics( + controller.time.cx_start, vertcat(q_root, q_joints, qdot_root, qdot_joints), controller.controls.cx_start, controller.parameters.cx_start, @@ -671,6 +673,7 @@ def stochastic_dg_dx_implicit( DF_DX_fun = Function( "DF_DX_fun", [ + controller.time.cx_start, q_root, q_joints, qdot_root, @@ -685,6 +688,7 @@ def stochastic_dg_dx_implicit( ) DF_DX = DF_DX_fun( + controller.time.cx_start, controller.states["q"].cx_start[:nb_root], controller.states["q"].cx_start[nb_root:], controller.states["qdot"].cx_start[:nb_root], diff --git a/bioptim/limits/penalty.py b/bioptim/limits/penalty.py index 0ad585853..720bd1b3c 100644 --- a/bioptim/limits/penalty.py +++ b/bioptim/limits/penalty.py @@ -444,6 +444,7 @@ def minimize_qddot(penalty: PenaltyOption, controller: PenaltyController): if "qddot" not in controller.states and "qddot" not in controller.controls: return controller.dynamics( + controller.time.cx_start, controller.states.cx_start, controller.controls.cx_start, controller.parameters.cx, @@ -554,6 +555,7 @@ def minimize_com_acceleration(penalty: PenaltyOption, controller: PenaltyControl controller.states["q"].mx, controller.states["qdot"].mx, controller.dynamics( + controller.time.mx, controller.states.mx, controller.controls.mx, controller.parameters.mx, @@ -663,6 +665,7 @@ def minimize_contact_forces( penalty.quadratic = True if penalty.quadratic is None else penalty.quadratic contact_force = controller.get_nlp.contact_forces_func( + controller.time.cx_start, controller.states.cx_start, controller.controls.cx_start, controller.parameters.cx, @@ -988,16 +991,16 @@ def continuity(penalty: PenaltyOption, controller: PenaltyController | list): continuity = controller.states.cx_end if controller.get_nlp.ode_solver.is_direct_collocation: cx = horzcat(*([controller.states.cx_start] + controller.states.cx_intermediates_list)) - continuity -= controller.integrate(x0=cx, p=u, params=controller.parameters.cx)["xf"] + continuity -= controller.integrate(x0=cx, p=u, params=controller.parameters.cx, t = controller.ocp.node_time(phase_idx=controller.get_nlp.phase_idx, node_idx=controller.node_index))["xf"] continuity = vertcat( continuity, - controller.integrate(x0=cx, p=u, params=controller.parameters.cx)["defects"], + controller.integrate(x0=cx, p=u, params=controller.parameters.cx, t = controller.ocp.node_time(phase_idx=controller.get_nlp.phase_idx, node_idx=controller.node_index))["defects"], ) penalty.integrate = True else: - continuity -= controller.integrate(x0=controller.states.cx_start, p=u, params=controller.parameters.cx)[ + continuity -= controller.integrate(x0=controller.states.cx_start, p=u, params=controller.parameters.cx, t = controller.ocp.node_time(phase_idx=controller.get_nlp.phase_idx, node_idx=controller.node_index))[ "xf" ] # continuity -= controller.integrate( @@ -1005,6 +1008,7 @@ def continuity(penalty: PenaltyOption, controller: PenaltyController | list): # x0=controller.states.cx_start, p=u, params=controller.parameters.cx)[ # "xf"] + penalty.explicit_derivative = True penalty.multi_thread = True diff --git a/bioptim/limits/penalty_controller.py b/bioptim/limits/penalty_controller.py index b2bcb2b1d..0eee0cc5f 100644 --- a/bioptim/limits/penalty_controller.py +++ b/bioptim/limits/penalty_controller.py @@ -110,6 +110,20 @@ def mx_to_cx(self): def model(self): return self._nlp.model + @property + def time(self) -> OptimizationVariableList: + """ + Return the time associated with the current node index + + Returns + ------- + The time at node node_index + """ + self._nlp.time.node_index = self.node_index + out = self._nlp.time.unscaled + out.current_cx_to_get = self.cx_index_to_get + return out + @property def states(self) -> OptimizationVariableList: """ @@ -187,6 +201,23 @@ def integrate(self): def dynamics(self): return self._nlp.dynamics_func + @property + def time_scaled(self) -> OptimizationVariableList: + """ + Return the scaled states associated with the current node index. + + Warning: Most of the time, the user does not want that states but the normal `states`, that said, it can + sometime be useful for very limited number of use case. + + Returns + ------- + The scaled states at node node_index + """ + self._nlp.time.node_index = self.node_index + out = self._nlp.time.scaled + out.current_cx_to_get = self.cx_index_to_get + return out + @property def states_scaled(self) -> OptimizationVariableList: """ diff --git a/bioptim/limits/penalty_option.py b/bioptim/limits/penalty_option.py index f7e0a6f32..e252b0ede 100644 --- a/bioptim/limits/penalty_option.py +++ b/bioptim/limits/penalty_option.py @@ -436,10 +436,12 @@ def get_u(u: MX | SX, dt: MX | SX): for ctrl in controllers: self.all_nodes_index.extend(ctrl.t) + time_cx_scaled = ocp.cx() state_cx_scaled = ocp.cx() control_cx_scaled = ocp.cx() stochastic_cx_scaled = ocp.cx() for ctrl in controllers: + time_cx_scaled = vertcat(time_cx_scaled, ctrl.time_scaled.cx) state_cx_scaled = vertcat(state_cx_scaled, ctrl.states_scaled.cx) control_cx_scaled = vertcat(control_cx_scaled, ctrl.controls_scaled.cx) stochastic_cx_scaled = vertcat(stochastic_cx_scaled, ctrl.stochastic_variables.unscaled.cx) @@ -455,10 +457,12 @@ def get_u(u: MX | SX, dt: MX | SX): else: state_cx_scaled = controller.states_scaled.cx_start control_cx_scaled = controller.controls_scaled.cx_start + time_cx_scaled = controller.time.cx_start stochastic_cx_scaled = controller.stochastic_variables.cx_start if self.explicit_derivative: if self.derivative: raise RuntimeError("derivative and explicit_derivative cannot be simultaneously true") + time_cx_scaled = horzcat(time_cx_scaled, controller.time_scaled.cx_end) state_cx_scaled = horzcat(state_cx_scaled, controller.states_scaled.cx_end) control_cx_scaled = horzcat(control_cx_scaled, controller.controls_scaled.cx_end) stochastic_cx_scaled = horzcat(stochastic_cx_scaled, controller.stochastic_variables.cx_end) @@ -478,28 +482,32 @@ def get_u(u: MX | SX, dt: MX | SX): # Do not use nlp.add_casadi_func because all functions must be registered sub_fcn = fcn[self.rows, self.cols] self.function[node] = controller.to_casadi_func( - name, sub_fcn, state_cx_scaled, control_cx_scaled, param_cx, stochastic_cx_scaled, expand=self.expand + name, sub_fcn, time_cx_scaled, state_cx_scaled, control_cx_scaled, param_cx, stochastic_cx_scaled, expand=self.expand ) self.function_non_threaded[node] = self.function[node] if self.derivative: + time_cx_scaled = horzcat(controller.time_scaled.cx_end, controller.time_scaled.cx_start) state_cx_scaled = horzcat(controller.states_scaled.cx_end, controller.states_scaled.cx_start) control_cx_scaled = horzcat(controller.controls_scaled.cx_end, controller.controls_scaled.cx_start) self.function[node] = biorbd.to_casadi_func( f"{name}", # TODO: Charbie -> this is Flase, add stochastic_variables for start, mid AND end self.function[node]( + controller.time_scaled.cx_end, controller.states_scaled.cx_end, controller.controls_scaled.cx_end, param_cx, controller.stochastic_variables.cx_start, ) - self.function[node]( + controller.time_scaled.cx_start, controller.states_scaled.cx_start, controller.controls_scaled.cx_start, param_cx, controller.stochastic_variables.cx_start, ), + time_cx_scaled, state_cx_scaled, control_cx_scaled, param_cx, @@ -522,6 +530,14 @@ def get_u(u: MX | SX, dt: MX | SX): exponent = 2 if self.quadratic and self.weight else 1 if is_trapezoidal: + # Hypothesis: the function is continuous on time + # it neglects the discontinuities at the beginning of the optimization + time_cx_scaled = ( + horzcat(controller.time_scaled.cx_start, controller.time_scaled.cx_end) + if self.integration_rule == QuadratureRule.APPROXIMATE_TRAPEZOIDAL + else controller.time_scaled.cx_start + ) + # Hypothesis: the function is continuous on states # it neglects the discontinuities at the beginning of the optimization state_cx_scaled = ( @@ -551,7 +567,7 @@ def get_u(u: MX | SX, dt: MX | SX): state_cx_end_scaled = ( controller.states_scaled.cx_end if self.integration_rule == QuadratureRule.APPROXIMATE_TRAPEZOIDAL - else controller.integrate(x0=state_cx, p=control_cx_end, params=controller.parameters.cx)["xf"] + else controller.integrate(t=controller.ocp.node_time(phase_idx=controller.get_nlp.phase_idx, node_idx=controller.node_index), x0=state_cx, p=control_cx_end, params=controller.parameters.cx)["xf"] ) stochastic_cx_scaled = ( @@ -565,6 +581,7 @@ def get_u(u: MX | SX, dt: MX | SX): ( ( self.function[node]( + controller.time_scaled.cx_start, controller.states_scaled.cx_start, controller.controls_scaled.cx_start, param_cx, @@ -574,12 +591,13 @@ def get_u(u: MX | SX, dt: MX | SX): ) ** exponent + ( - self.function[node](state_cx_end_scaled, control_cx_end_scaled, param_cx, stochastic_cx_scaled) + self.function[node](time_cx_scaled, state_cx_end_scaled, control_cx_end_scaled, param_cx, stochastic_cx_scaled) - target_cx[:, 1] ) ** exponent ) / 2, + time_cx_scaled, state_cx_scaled, control_cx_scaled, param_cx, @@ -588,11 +606,11 @@ def get_u(u: MX | SX, dt: MX | SX): dt_cx, ) modified_fcn = modified_function( - state_cx_scaled, control_cx_scaled, param_cx, stochastic_cx_scaled, target_cx, dt_cx + time_cx_scaled, state_cx_scaled, control_cx_scaled, param_cx, stochastic_cx_scaled, target_cx, dt_cx ) else: modified_fcn = ( - self.function[node](state_cx_scaled, control_cx_scaled, param_cx, stochastic_cx_scaled) - target_cx + self.function[node](time_cx_scaled, state_cx_scaled, control_cx_scaled, param_cx, stochastic_cx_scaled) - target_cx ) ** exponent # for the future bioptim adventurer: here lies the reason that a constraint must have weight = 0. @@ -601,7 +619,7 @@ def get_u(u: MX | SX, dt: MX | SX): # Do not use nlp.add_casadi_func because all of them must be registered self.weighted_function[node] = Function( name, - [state_cx_scaled, control_cx_scaled, param_cx, stochastic_cx_scaled, weight_cx, target_cx, dt_cx], + [time_cx_scaled, state_cx_scaled, control_cx_scaled, param_cx, stochastic_cx_scaled, weight_cx, target_cx, dt_cx], [modified_fcn], ) self.weighted_function_non_threaded[node] = self.weighted_function[node] diff --git a/bioptim/optimization/non_linear_program.py b/bioptim/optimization/non_linear_program.py index cd7facff1..b2d02b14e 100644 --- a/bioptim/optimization/non_linear_program.py +++ b/bioptim/optimization/non_linear_program.py @@ -87,6 +87,10 @@ class NonLinearProgram: The casadi variables for the integration at each node of the phase controls: OptimizationVariableContainer A list of all the control variables + t_bounds = Bounds() + The bounds for the time + t_init = InitialGuess() + The initial guess for the time x_bounds = Bounds() The bounds for the states x_init = InitialGuess() @@ -146,6 +150,8 @@ def __init__(self, assume_phase_dynamics): self.phase_mapping = None self.plot = {} self.plot_mapping = {} + self.T = None + self.T_scaled = None self.t0 = None self.tf = None self.t_initial_guess = None @@ -155,9 +161,12 @@ def __init__(self, assume_phase_dynamics): self.U_scaled = None self.u_scaling = None self.U = None + self.use_time_from_phase_idx = NodeMapping() self.use_states_from_phase_idx = NodeMapping() self.use_controls_from_phase_idx = NodeMapping() self.use_states_dot_from_phase_idx = NodeMapping() + self.t_bounds = BoundsList() + self.t_init = InitialGuessList() self.x_bounds = BoundsList() self.x_init = InitialGuessList() self.X_scaled = None @@ -167,6 +176,7 @@ def __init__(self, assume_phase_dynamics): self.s_init = InitialGuessList() self.S = None self.assume_phase_dynamics = assume_phase_dynamics + self.time = OptimizationVariableContainer(assume_phase_dynamics) self.states = OptimizationVariableContainer(assume_phase_dynamics) self.states_dot = OptimizationVariableContainer(assume_phase_dynamics) self.controls = OptimizationVariableContainer(assume_phase_dynamics) @@ -191,6 +201,7 @@ def initialize(self, cx: Callable = None): self.g_internal = [] self.casadi_func = {} + self.time.initialize_from_shooting(n_shooting=self.ns + 1, cx=self.cx) self.states.initialize_from_shooting(n_shooting=self.ns + 1, cx=self.cx) self.states_dot.initialize_from_shooting(n_shooting=self.ns + 1, cx=self.cx) self.controls.initialize_from_shooting(n_shooting=self.ns + 1, cx=self.cx) diff --git a/bioptim/optimization/optimal_control_program.py b/bioptim/optimization/optimal_control_program.py index 98092ad70..fd704daf8 100644 --- a/bioptim/optimization/optimal_control_program.py +++ b/bioptim/optimization/optimal_control_program.py @@ -151,9 +151,11 @@ def __init__( dynamics: Dynamics | DynamicsList, n_shooting: int | list | tuple, phase_time: int | float | list | tuple, + t_bounds: BoundsList = None, x_bounds: BoundsList = None, u_bounds: BoundsList = None, s_bounds: BoundsList = None, + t_init: InitialGuessList | None = None, x_init: InitialGuessList | None = None, u_init: InitialGuessList | None = None, s_init: InitialGuessList | None = None, @@ -195,12 +197,16 @@ def __init__( The number of shooting point of the phases phase_time: int | float | list | tuple The phase time of the phases + t_init: InitialGuess | InitialGuessList + The initial guesses for the time x_init: InitialGuess | InitialGuessList The initial guesses for the states u_init: InitialGuess | InitialGuessList The initial guesses for the controls s_init: InitialGuess | InitialGuessList The initial guesses for the stochastic variables + t_bounds: Bounds | BoundsList + The bounds for the time x_bounds: Bounds | BoundsList The bounds for the states u_bounds: Bounds | BoundsList @@ -263,9 +269,11 @@ def __init__( dynamics, n_shooting, phase_time, + t_init, x_init, u_init, s_init, + t_bounds, x_bounds, u_bounds, s_bounds, @@ -305,9 +313,11 @@ def __init__( n_threads, n_shooting, phase_time, + t_bounds, x_bounds, u_bounds, s_bounds, + t_init, x_init, u_init, s_init, @@ -366,9 +376,11 @@ def set_original_values( dynamics, n_shooting, phase_time, + t_init, x_init, u_init, s_init, + t_bounds, x_bounds, u_bounds, s_bounds, @@ -408,9 +420,11 @@ def set_original_values( "dynamics": dynamics, "n_shooting": n_shooting, "phase_time": phase_time, + "t_init": t_init, "x_init": x_init, "u_init": u_init, "s_init": s_init, + "t_bounds": t_bounds, "x_bounds": x_bounds, "u_bounds": u_bounds, "s_bounds": s_bounds, @@ -448,9 +462,11 @@ def check_arguments_and_build_nlp( n_threads, n_shooting, phase_time, + t_bounds, x_bounds, u_bounds, s_bounds, + t_init, x_init, u_init, s_init, @@ -498,6 +514,11 @@ def check_arguments_and_build_nlp( else: raise RuntimeError("phase_time should be a number or a list of number") + if t_bounds is None: + t_bounds = BoundsList() + elif not isinstance(t_bounds, BoundsList): + raise RuntimeError("t_bounds should be built from a BoundsList") + if x_bounds is None: x_bounds = BoundsList() elif not isinstance(x_bounds, BoundsList): @@ -513,6 +534,11 @@ def check_arguments_and_build_nlp( elif not isinstance(s_bounds, BoundsList): raise RuntimeError("s_bounds should be built from a BoundsList") + if t_init is None: + t_init = InitialGuessList() + if not isinstance(t_init, InitialGuessList): + raise RuntimeError("t_init should be built from a InitialGuessList") + if x_init is None: x_init = InitialGuessList() if not isinstance(x_init, InitialGuessList): @@ -528,10 +554,12 @@ def check_arguments_and_build_nlp( if not isinstance(s_init, InitialGuessList): raise RuntimeError("s_init should be built from a InitialGuessList") + t_bounds = self._prepare_option_dict_for_phase("t_bounds", t_bounds, BoundsList) x_bounds = self._prepare_option_dict_for_phase("x_bounds", x_bounds, BoundsList) u_bounds = self._prepare_option_dict_for_phase("u_bounds", u_bounds, BoundsList) s_bounds = self._prepare_option_dict_for_phase("s_bounds", s_bounds, BoundsList) + t_init = self._prepare_option_dict_for_phase("t_init", t_init, InitialGuessList) x_init = self._prepare_option_dict_for_phase("x_init", x_init, InitialGuessList) u_init = self._prepare_option_dict_for_phase("u_init", u_init, InitialGuessList) s_init = self._prepare_option_dict_for_phase("s_init", s_init, InitialGuessList) @@ -724,8 +752,8 @@ def check_arguments_and_build_nlp( self.parameter_bounds = BoundsList() self.parameter_init = InitialGuessList() - self.update_bounds(x_bounds, u_bounds, parameter_bounds, s_bounds) - self.update_initial_guess(x_init, u_init, parameter_init, s_init) + self.update_bounds(t_bounds, x_bounds, u_bounds, parameter_bounds, s_bounds) + self.update_initial_guess(t_init, x_init, u_init, parameter_init, s_init) # Define the actual NLP problem OptimizationVectorHelper.declare_ocp_shooting_points(self) @@ -1091,6 +1119,7 @@ def _declare_parameters(self, new_parameters: ParameterList): def update_bounds( self, + t_bounds: BoundsList = None, x_bounds: BoundsList = None, u_bounds: BoundsList = None, parameter_bounds: BoundsList = None, @@ -1101,6 +1130,8 @@ def update_bounds( Parameters ---------- + t_bounds: BoundsList + The time bounds to add x_bounds: BoundsList The state bounds to add u_bounds: BoundsList @@ -1111,6 +1142,13 @@ def update_bounds( The stochastic variable bounds to add """ for i in range(self.n_phases): + if t_bounds is not None: + if not isinstance(t_bounds, BoundsList): + raise RuntimeError("t_bounds should be built from a BoundsList") + origin_phase = 0 if len(t_bounds) == 1 else i + for key in t_bounds[origin_phase].keys(): + self.nlp[i].t_bounds.add(key, t_bounds[origin_phase][key], phase=0) + if x_bounds is not None: if not isinstance(x_bounds, BoundsList): raise RuntimeError("x_bounds should be built from a BoundsList") @@ -1148,6 +1186,7 @@ def update_bounds( def update_initial_guess( self, + t_init: InitialGuessList = None, x_init: InitialGuessList = None, u_init: InitialGuessList = None, parameter_init: InitialGuessList = None, @@ -1158,6 +1197,8 @@ def update_initial_guess( Parameters ---------- + t_init: BoundsList + The time initial guess to add x_init: BoundsList The state initial guess to add u_init: BoundsList @@ -1169,6 +1210,18 @@ def update_initial_guess( """ for i in range(self.n_phases): + if t_init: + if not isinstance(t_init, InitialGuessList): + raise RuntimeError("t_init should be built from a InitialGuessList") + origin_phase = 0 if len(t_init) == 1 else i + for key in t_init[origin_phase].keys(): + if ( + not self.nlp[i].ode_solver.is_direct_collocation + and x_init[origin_phase].type == InterpolationType.ALL_POINTS + ): + raise ValueError("InterpolationType.ALL_POINTS must only be used with direct collocation") + self.nlp[i].t_init.add(key, t_init[origin_phase][key], phase=0) + if x_init: if not isinstance(x_init, InitialGuessList): raise RuntimeError("x_init should be built from a InitialGuessList") @@ -1381,7 +1434,7 @@ def compute_penalty_values(t, x, u, p, s, penalty, dt: int | Callable): if penalty.transition or penalty.multinode_penalty: out.append( penalty.weighted_function_non_threaded[t]( - x.reshape((-1, 1)), u.reshape((-1, 1)), p, s.reshape((-1, 1)), penalty.weight, _target, 1 + t, x.reshape((-1, 1)), u.reshape((-1, 1)), p, s.reshape((-1, 1)), penalty.weight, _target, 1 ) ) # dt=1 because multinode penalties behave like Mayer functions @@ -1395,19 +1448,19 @@ def compute_penalty_values(t, x, u, p, s, penalty, dt: int | Callable): (x.shape[0], int(penalty.weighted_function_non_threaded[t].nnz_in(0) / x.shape[0])) ) - out.append(penalty.weighted_function_non_threaded[t](state_value, u, p, s, penalty.weight, _target, dt)) + out.append(penalty.weighted_function_non_threaded[t](t, state_value, u, p, s, penalty.weight, _target, dt)) elif ( penalty.integration_rule == QuadratureRule.APPROXIMATE_TRAPEZOIDAL or penalty.integration_rule == QuadratureRule.TRAPEZOIDAL ): out = [ penalty.weighted_function_non_threaded[t]( - x[:, [i, i + 1]], u[:, i], p, s, penalty.weight, _target, dt + t, x[:, [i, i + 1]], u[:, i], p, s, penalty.weight, _target, dt ) for i in range(x.shape[1] - 1) ] else: - out.append(penalty.weighted_function_non_threaded[t](x, u, p, s, penalty.weight, _target, dt)) + out.append(penalty.weighted_function_non_threaded[t](t, x, u, p, s, penalty.weight, _target, dt)) return sum1(horzcat(*out)) def add_penalty(_penalties): diff --git a/bioptim/optimization/optimization_vector.py b/bioptim/optimization/optimization_vector.py index e62384d08..219cdf32d 100644 --- a/bioptim/optimization/optimization_vector.py +++ b/bioptim/optimization/optimization_vector.py @@ -33,12 +33,16 @@ def declare_ocp_shooting_points(ocp): """ Declare all the casadi variables with the right size to be used during a specific phase """ + t = [] + t_scaled = [] x = [] x_scaled = [] u = [] u_scaled = [] s = [] for nlp in ocp.nlp: + t.append([]) + t_scaled.append([]) x.append([]) x_scaled.append([]) u.append([]) @@ -90,8 +94,15 @@ def declare_ocp_shooting_points(ocp): nlp.cx.sym("S_" + str(nlp.phase_idx) + "_" + str(k), nlp.stochastic_variables.shape, 1) ) + t[nlp.phase_idx].append( + nlp.cx.sym("T_" + str(nlp.phase_idx) + "_" + str(k), nlp.time.scaled.shape, 1) + ) + OptimizationVectorHelper._set_node_index(nlp, 0) + nlp.T_scaled = t_scaled[nlp.phase_idx] + nlp.T = t[nlp.phase_idx] + nlp.X_scaled = x_scaled[nlp.phase_idx] nlp.X = x[nlp.phase_idx] @@ -109,7 +120,7 @@ def vector(ocp): ------- The vector of all variables """ - + t_scaled = [] x_scaled = [] u_scaled = [] s = [] @@ -118,10 +129,11 @@ def vector(ocp): x_scaled += [x.reshape((-1, 1)) for x in nlp.X_scaled] else: x_scaled += nlp.X_scaled + t_scaled += nlp.T_scaled u_scaled += nlp.U_scaled s += nlp.S - return vertcat(*x_scaled, *u_scaled, ocp.parameters.cx, *s) + return vertcat(*t_scaled, *x_scaled, *u_scaled, ocp.parameters.cx, *s) @staticmethod def bounds_vectors(ocp) -> tuple[np.ndarray, np.ndarray]: @@ -262,6 +274,29 @@ def init_vector(ocp): """ v_init = np.ndarray((0, 1)) + # For time + for i_phase in range(len(ocp.nlp)): + nlp = ocp.nlp[i_phase] + OptimizationVectorHelper._set_node_index(nlp, 0) + + for key in nlp.time.keys(): + if key in nlp.t_init.keys(): + nlp.t_init[key].check_and_adjust_dimensions(nlp.time[key].cx.shape[0], nlp.ns) + + for k in range(nlp.ns + 1): + OptimizationVectorHelper._set_node_index(nlp, k) + collapsed_values = np.ndarray((nlp.time.shape, 1)) + for key in nlp.time: + if key in nlp.t_init.keys(): + value = nlp.t_init[key].init.evaluate_at(shooting_point=k) + else: + value = 0 + + # Organize the stochastic variables according to the correct indices + collapsed_values[nlp.time[key].index, 0] = value + + v_init = np.concatenate((v_init, np.reshape(collapsed_values.T, (-1, 1)))) + # For states for i_phase in range(len(ocp.nlp)): current_nlp = ocp.nlp[i_phase] @@ -416,6 +451,7 @@ def to_dictionaries(ocp, data: np.ndarray | DM) -> tuple: v_array = np.array(data).squeeze() + data_time = [] data_states = [] data_controls = [] data_stochastic_variables = [] @@ -424,6 +460,7 @@ def to_dictionaries(ocp, data: np.ndarray | DM) -> tuple: nlp.controls.node_index = 0 n_points = nlp.ns * (1 if nlp.ode_solver.is_direct_shooting else (nlp.ode_solver.polynomial_degree + 1)) + 1 + data_time.append({key: np.ndarray((nlp.time[key].shape, n_points)) for key in nlp.time}) data_states.append({key: np.ndarray((nlp.states[key].shape, n_points)) for key in nlp.states}) data_controls.append( { @@ -444,6 +481,23 @@ def to_dictionaries(ocp, data: np.ndarray | DM) -> tuple: ) data_parameters = {key: np.ndarray((0, 1)) for key in ocp.parameters.keys()} + # For time + offset = 0 + p_idx = 0 + for p in range(ocp.n_phases): + nlp = ocp.nlp[p] + nx = nlp.time.shape + + if nlp.use_time_from_phase_idx == nlp.phase_idx: + repeat = (nlp.ode_solver.polynomial_degree + 1) if nlp.ode_solver.is_direct_collocation else 1 + for k in range((nlp.ns * repeat) + 1): + nlp.time.node_index = k // repeat + x_array = v_array[offset : offset + nx].reshape((nlp.time.scaled.shape, -1), order="F") + for key in nlp.time: + data_time[p_idx][key][:, k : k + 1] = x_array[nlp.time[key].index, :] + offset += nx + p_idx += 1 + # For states offset = 0 p_idx = 0 @@ -509,7 +563,7 @@ def to_dictionaries(ocp, data: np.ndarray | DM) -> tuple: offset += nstochastic p_idx += 1 - return data_states, data_controls, data_parameters, data_stochastic_variables + return data_time, data_states, data_controls, data_parameters, data_stochastic_variables @staticmethod def _nb_points(nlp, interpolation_type): diff --git a/bioptim/optimization/solution.py b/bioptim/optimization/solution.py index bef2a2ac0..89e54b4ec 100644 --- a/bioptim/optimization/solution.py +++ b/bioptim/optimization/solution.py @@ -66,6 +66,8 @@ class Solution: The number of iterations that were required to solve the program status: int Optimization success status (Ipopt: 0=Succeeded, 1=Failed) + _time: list + The data structure that holds the time _states: list The data structure that holds the states _controls: list @@ -215,9 +217,11 @@ def __init__(self, nlp: NonLinearProgram): self.tf = nlp.tf self.phase_idx = nlp.phase_idx + self.use_time_from_phase_idx = nlp.use_time_from_phase_idx self.use_states_from_phase_idx = nlp.use_states_from_phase_idx self.use_controls_from_phase_idx = nlp.use_controls_from_phase_idx self.model = nlp.model + self.time = nlp.time self.states = nlp.states self.states_dot = nlp.states_dot self.controls = nlp.controls @@ -316,6 +320,7 @@ def __init__(self, ocp, sol: dict | list | tuple | np.ndarray | DM | None): self._time_vector = None # Extract the data now for further use + self._time = {} self._states = {} self._controls = {} self.parameters = {} @@ -323,7 +328,7 @@ def __init__(self, ocp, sol: dict | list | tuple | np.ndarray | DM | None): self._integrated_values = {} self.phase_time = [] - def get_integrated_values(states, controls, parameters, stochastic_variables): + def get_integrated_values(time, states, controls, parameters, stochastic_variables): integrated_values_num = [{} for _ in self.ocp.nlp] for i_phase, nlp in enumerate(self.ocp.nlp): nlp.states.node_index = 0 @@ -433,6 +438,7 @@ def init_from_dict(_sol: dict): # Extract the data now for further use ( + self._time["scaled"], self._states["scaled"], self._controls["scaled"], self.parameters, @@ -445,7 +451,7 @@ def init_from_dict(_sol: dict): self.phase_time = OptimizationVectorHelper.extract_phase_time(self.ocp, self.vector) self._time_vector = self._generate_time() self._integrated_values = get_integrated_values( - self._states["unscaled"], self._controls["unscaled"], self.parameters, self._stochastic_variables + self._time, self._states["unscaled"], self._controls["unscaled"], self.parameters, self._stochastic_variables ) def init_from_initial_guess(_sol: list): @@ -535,6 +541,7 @@ def init_from_initial_guess(_sol: list): self.vector = np.concatenate((self.vector, ss[key].init.evaluate_at(i)[:, np.newaxis])) ( + self._time["scaled"], self._states["scaled"], self._controls["scaled"], self.parameters, @@ -547,7 +554,7 @@ def init_from_initial_guess(_sol: list): self.phase_time = OptimizationVectorHelper.extract_phase_time(self.ocp, self.vector) self._time_vector = self._generate_time() self._integrated_values = get_integrated_values( - self._states["unscaled"], self._controls["unscaled"], self.parameters, self._stochastic_variables + self._time, self._states["unscaled"], self._controls["unscaled"], self.parameters, self._stochastic_variables ) def init_from_vector(_sol: np.ndarray | DM): @@ -562,6 +569,7 @@ def init_from_vector(_sol: np.ndarray | DM): self.vector = _sol ( + self._time["scaled"], self._states["scaled"], self._controls["scaled"], self.parameters, @@ -573,7 +581,7 @@ def init_from_vector(_sol: np.ndarray | DM): self._complete_control() self.phase_time = OptimizationVectorHelper.extract_phase_time(self.ocp, self.vector) self._integrated_values = get_integrated_values( - self._states["unscaled"], self._controls["unscaled"], self.parameters, self._stochastic_variables + self._time, self._states["unscaled"], self._controls["unscaled"], self.parameters, self._stochastic_variables ) if isinstance(sol, dict): @@ -663,8 +671,9 @@ def copy(self, skip_data: bool = False) -> Any: if skip_data: new._states["unscaled"], new._controls["unscaled"] = [], [] - new._states["scaled"], new._controls["scaled"], new.parameters, new._stochastic_variables = [], [], {}, [] + self._time["scaled"], new._states["scaled"], new._controls["scaled"], new.parameters, new._stochastic_variables = [], [], [], {}, [] else: + new._time["scaled"] = deepcopy(self._time["scaled"]) new._states["scaled"] = deepcopy(self._states["scaled"]) new._controls["scaled"] = deepcopy(self._controls["scaled"]) new.parameters = deepcopy(self.parameters) @@ -1202,11 +1211,15 @@ def __perform_integration( for p, (nlp, t_eval) in enumerate(zip(self.ocp.nlp, out._time_vector)): self.ocp.nlp[p].controls.node_index = 0 - states_phase_idx = self.ocp.nlp[p].use_states_from_phase_idx controls_phase_idx = self.ocp.nlp[p].use_controls_from_phase_idx param_scaling = nlp.parameters.scaling x0 = self._get_first_frame_states(out, shooting_type, phase=p) + if self.ocp.nlp[p].time.keys(): + t = np.concatenate([self._time[p][key] for key in self.ocp.nlp[p].time]) + else: + t = np.array([]) + u = ( np.array([]) if nlp.control_type == ControlType.NONE @@ -1225,6 +1238,7 @@ def __perform_integration( integrated_sol = solve_ivp_bioptim_interface( dynamics_func=nlp.dynamics, keep_intermediate_points=keep_intermediate_points, + t=t, x0=x0, u=u, s=s, @@ -1238,6 +1252,7 @@ def __perform_integration( dynamics_func=nlp.dynamics_func, keep_intermediate_points=keep_intermediate_points, t_eval=t_eval[:-1] if shooting_type == Shooting.MULTIPLE else t_eval, + t=t, x0=x0, u=u, s=s, @@ -1291,7 +1306,7 @@ def interpolate(self, n_frames: int | list | tuple) -> Any: t_all.append(np.linspace(sum(out.phase_time[: p + 1]), sum(out.phase_time[: p + 2]), out.ns[p] + 1)) if isinstance(n_frames, int): - _, data_states, _, _, out.phase_time, out.ns = self._merge_phases(skip_controls=True) + time, _, data_states, _, _, out.phase_time, out.ns = self._merge_phases(skip_controls=True) t_all = [np.concatenate((np.concatenate([_t[:-1] for _t in t_all]), [t_all[-1][-1]]))] n_frames = [n_frames] @@ -1630,12 +1645,14 @@ def _get_penalty_cost(self, nlp, penalty): ) for idx in penalty.node_idx: + t = [] x = [] u = [] s = [] target = [] if nlp is not None: if penalty.multinode_penalty or penalty.transition: + t = np.array(()) x = np.array(()) u = np.array(()) s = np.array(()) @@ -1643,6 +1660,8 @@ def _get_penalty_cost(self, nlp, penalty): node_idx = penalty.multinode_idx[i] phase_idx = penalty.nodes_phase[i] + for key in nlp.time: + t = np.concatenate((t, self._time["scaled"][phase_idx][key][:, node_idx])) for key in nlp.states: x = np.concatenate((x, self._states["scaled"][phase_idx][key][:, node_idx])) for key in nlp.controls: @@ -1723,8 +1742,8 @@ def _get_penalty_cost(self, nlp, penalty): # casadi function if not self.ocp.assume_phase_dynamics and ((isinstance(u, list) and u != []) or isinstance(u, np.ndarray)): u = u[:, ~np.isnan(np.sum(u, axis=0))] - val.append(penalty.function[idx](x, u, p, s)) - val_weighted.append(penalty.weighted_function[idx](x, u, p, s, penalty.weight, target, dt)) + val.append(penalty.function[idx](t, x, u, p, s)) + val_weighted.append(penalty.weighted_function[idx](t, x, u, p, s, penalty.weight, target, dt)) val = np.nansum(val) val_weighted = np.nansum(val_weighted) diff --git a/bioptim/optimization/stochastic_optimal_control_program.py b/bioptim/optimization/stochastic_optimal_control_program.py index 145286ee1..196783805 100644 --- a/bioptim/optimization/stochastic_optimal_control_program.py +++ b/bioptim/optimization/stochastic_optimal_control_program.py @@ -36,9 +36,11 @@ def __init__( dynamics: Dynamics | DynamicsList, n_shooting: int | list | tuple, phase_time: int | float | list | tuple, + t_bounds: BoundsList = None, x_bounds: BoundsList = None, u_bounds: BoundsList = None, s_bounds: BoundsList = None, + t_init: InitialGuessList | None = None, x_init: InitialGuessList | None = None, u_init: InitialGuessList | None = None, s_init: InitialGuessList | None = None, @@ -100,9 +102,11 @@ def __init__( dynamics, n_shooting, phase_time, + t_init, x_init, u_init, s_init, + t_bounds, x_bounds, u_bounds, s_bounds, @@ -142,9 +146,11 @@ def __init__( n_threads, n_shooting, phase_time, + t_bounds, x_bounds, u_bounds, s_bounds, + t_init, x_init, u_init, s_init, diff --git a/tests/test__global_plots.py b/tests/test__global_plots.py index 7b11459b7..6b3628593 100644 --- a/tests/test__global_plots.py +++ b/tests/test__global_plots.py @@ -217,6 +217,7 @@ def override_penalty(pen: list[PenaltyOption]): for cmp, p in enumerate(pen): if p: for node_index in p.node_idx: + nlp.time.node_index = node_index nlp.states.node_index = node_index nlp.states_dot.node_index = node_index nlp.controls.node_index = node_index @@ -232,23 +233,24 @@ def override_penalty(pen: list[PenaltyOption]): .replace(".", "_") .replace("__", "_") ) - x = MX.sym("x", *p.weighted_function[node_index].sparsity_in("i0").shape) - u = MX.sym("u", *p.weighted_function[node_index].sparsity_in("i1").shape) - if p.weighted_function[node_index].sparsity_in("i1").shape == (0, 0): + t = MX.sym("t", *p.weighted_function[node_index].sparsity_in("i0").shape) + x = MX.sym("x", *p.weighted_function[node_index].sparsity_in("i1").shape) + u = MX.sym("u", *p.weighted_function[node_index].sparsity_in("i2").shape) + if p.weighted_function[node_index].sparsity_in("i3").shape == (0, 0): u = MX.sym("u", 3, 1) - param = MX.sym("param", *p.weighted_function[node_index].sparsity_in("i2").shape) - s = MX.sym("s", *p.weighted_function[node_index].sparsity_in("i3").shape) - weight = MX.sym("weight", *p.weighted_function[node_index].sparsity_in("i4").shape) - target = MX.sym("target", *p.weighted_function[node_index].sparsity_in("i5").shape) - dt = MX.sym("dt", *p.weighted_function[node_index].sparsity_in("i6").shape) + param = MX.sym("param", *p.weighted_function[node_index].sparsity_in("i3").shape) + s = MX.sym("s", *p.weighted_function[node_index].sparsity_in("i4").shape) + weight = MX.sym("weight", *p.weighted_function[node_index].sparsity_in("i5").shape) + target = MX.sym("target", *p.weighted_function[node_index].sparsity_in("i6").shape) + dt = MX.sym("dt", *p.weighted_function[node_index].sparsity_in("i7").shape) p.function[node_index] = Function( - name, [x, u, param, s], [np.array([range(cmp, len(p.rows) + cmp)]).T] + name, [t, x, u, param, s], [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, - [x, u, param, s, weight, target, dt], + [t, x, u, param, s, weight, target, dt], [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/test_controltype_none.py b/tests/test_controltype_none.py index 7a9879435..6666d7e5c 100644 --- a/tests/test_controltype_none.py +++ b/tests/test_controltype_none.py @@ -77,15 +77,15 @@ def system_dynamics( def custom_dynamics( self, + time: MX, states: MX, controls: MX, parameters: MX, nlp: NonLinearProgram, - t=None, ) -> DynamicsEvaluation: t_phase = nlp.parameters.mx[-1] return DynamicsEvaluation( - dxdt=nlp.model.system_dynamics(a=states[0], b=states[1], c=states[2], t=t, t_phase=t_phase), + dxdt=nlp.model.system_dynamics(a=states[0], b=states[1], c=states[2], t=time, t_phase=t_phase), defects=None, ) @@ -97,34 +97,25 @@ def custom_configure_dynamics_function(self, ocp, nlp, **extra_params): nlp.parameters = ocp.parameters DynamicsFunctions.apply_parameters(nlp.parameters.mx, nlp) - # Gets the t0 time for the current phase - t0_phase_in_ocp = sum1(nlp.parameters.mx[0 : nlp.phase_idx]) - # Gets every time node for the current phase - - dynamics_eval_horzcat = np.array(()) - for i in range(nlp.ns): - t_node_in_phase = nlp.parameters.mx[nlp.phase_idx] / (nlp.ns + 1) * i - t_node_in_ocp = t0_phase_in_ocp + t_node_in_phase - extra_params["t"] = t_node_in_ocp - - dynamics_eval = self.custom_dynamics( - nlp.states.scaled.mx_reduced, nlp.controls.scaled.mx_reduced, nlp.parameters.mx, nlp, **extra_params - ) + dynamics_eval = self.custom_dynamics( + nlp.time.scaled.mx_reduced, nlp.states.scaled.mx_reduced, nlp.controls.scaled.mx_reduced, nlp.parameters.mx, nlp, **extra_params + ) - dynamics_eval_horzcat = ( - horzcat(dynamics_eval.dxdt) if i == 0 else horzcat(dynamics_eval_horzcat, dynamics_eval.dxdt) - ) + dynamics_dxdt = dynamics_eval.dxdt + if isinstance(dynamics_dxdt, (list, tuple)): + dynamics_dxdt = vertcat(*dynamics_dxdt) nlp.dynamics_func = Function( "ForwardDyn", [ + nlp.time.scaled.mx_reduced, nlp.states.scaled.mx_reduced, nlp.controls.scaled.mx_reduced, nlp.parameters.mx, nlp.stochastic_variables.mx, ], - [dynamics_eval_horzcat], - ["x", "u", "p", "s"], + [dynamics_dxdt], + ["t", "x", "u", "p", "s"], ["xdot"], ) @@ -165,9 +156,9 @@ def declare_variables(self, ocp: OptimalControlProgram, nlp: NonLinearProgram): as_states_dot=False, ) - t = MX.sym("t") # t needs a symbolic value to start computing in custom_configure_dynamics_function - self.custom_configure_dynamics_function(ocp, nlp, t=t) - + # t = MX.sym("t") # t needs a symbolic value to start computing in custom_configure_dynamics_function + # self.custom_configure_dynamics_function(ocp, nlp, t=t) + self.custom_configure_dynamics_function(ocp, nlp) def prepare_ocp( n_phase: int, @@ -253,6 +244,7 @@ def test_main_control_type_none(use_sx, assume_phase_dynamics): """ # TODO It seems assume_phase_dynamics=True is broken + # I THINK IT'S NORMAL AS THIS FUN CHANGES AT EACH NODE # number of stimulation corresponding to phases n = 10 diff --git a/tests/test_dynamics.py b/tests/test_dynamics.py index 5cb1c7b20..b141afea3 100644 --- a/tests/test_dynamics.py +++ b/tests/test_dynamics.py @@ -70,9 +70,11 @@ def test_torque_driven(with_contact, with_external_force, cx, rigidbody_dynamics ) phase_index = [i for i in range(ocp.n_phases)] NonLinearProgram.add(ocp, "phase_idx", phase_index, False) + use_time_from_phase_idx = [i for i in range(ocp.n_phases)] use_states_from_phase_idx = [i for i in range(ocp.n_phases)] use_states_dot_from_phase_idx = [i for i in range(ocp.n_phases)] use_controls_from_phase_idx = [i for i in range(ocp.n_phases)] + NonLinearProgram.add(ocp, "use_time_from_phase_idx", use_time_from_phase_idx, False) NonLinearProgram.add(ocp, "use_states_from_phase_idx", use_states_from_phase_idx, False) NonLinearProgram.add(ocp, "use_states_dot_from_phase_idx", use_states_dot_from_phase_idx, False) NonLinearProgram.add(ocp, "use_controls_from_phase_idx", use_controls_from_phase_idx, False) @@ -88,14 +90,15 @@ def test_torque_driven(with_contact, with_external_force, cx, rigidbody_dynamics ConfigureProblem.initialize(ocp, nlp) # Test the results + time = np.random.rand(nlp.time.shape, nlp.ns) states = np.random.rand(nlp.states.shape, nlp.ns) controls = np.random.rand(nlp.controls.shape, nlp.ns) params = np.random.rand(nlp.parameters.shape, nlp.ns) stochastic_variables = np.random.rand(nlp.stochastic_variables.shape, nlp.ns) - x_out = np.array(nlp.dynamics_func(states, controls, params, stochastic_variables)) + x_out = np.array(nlp.dynamics_func(time, states, controls, params, stochastic_variables)) if rigidbody_dynamics == RigidBodyDynamics.ODE: if with_contact: - contact_out = np.array(nlp.contact_forces_func(states, controls, params, stochastic_variables)) + contact_out = np.array(nlp.contact_forces_func(time, states, controls, params, stochastic_variables)) if with_external_force: np.testing.assert_almost_equal( x_out[:, 0], @@ -132,7 +135,7 @@ def test_torque_driven(with_contact, with_external_force, cx, rigidbody_dynamics ) elif rigidbody_dynamics == RigidBodyDynamics.DAE_FORWARD_DYNAMICS: if with_contact: - contact_out = np.array(nlp.contact_forces_func(states, controls, params, stochastic_variables)) + contact_out = np.array(nlp.contact_forces_func(time, states, controls, params, stochastic_variables)) if with_external_force: np.testing.assert_almost_equal( x_out[:, 0], [0.8631034, 0.3251833, 0.1195942, 0.4937956, 0.8074402, 0.4271078, 0.417411, 0.3232029] @@ -157,7 +160,7 @@ def test_torque_driven(with_contact, with_external_force, cx, rigidbody_dynamics ) elif rigidbody_dynamics == RigidBodyDynamics.DAE_INVERSE_DYNAMICS: if with_contact: - contact_out = np.array(nlp.contact_forces_func(states, controls, params, stochastic_variables)) + contact_out = np.array(nlp.contact_forces_func(time, states, controls, params, stochastic_variables)) if with_external_force: np.testing.assert_almost_equal( x_out[:, 0], [0.8631034, 0.3251833, 0.1195942, 0.4937956, 0.8074402, 0.4271078, 0.417411, 0.3232029] @@ -216,9 +219,11 @@ def test_torque_driven_implicit(with_contact, cx, assume_phase_dynamics): ) phase_index = [i for i in range(ocp.n_phases)] NonLinearProgram.add(ocp, "phase_idx", phase_index, False) + use_time_from_phase_idx = [i for i in range(ocp.n_phases)] use_states_from_phase_idx = [i for i in range(ocp.n_phases)] use_states_dot_from_phase_idx = [i for i in range(ocp.n_phases)] use_controls_from_phase_idx = [i for i in range(ocp.n_phases)] + NonLinearProgram.add(ocp, "use_time_from_phase_idx", use_time_from_phase_idx, False) NonLinearProgram.add(ocp, "use_states_from_phase_idx", use_states_from_phase_idx, False) NonLinearProgram.add(ocp, "use_states_dot_from_phase_idx", use_states_dot_from_phase_idx, False) NonLinearProgram.add(ocp, "use_controls_from_phase_idx", use_controls_from_phase_idx, False) @@ -228,14 +233,15 @@ def test_torque_driven_implicit(with_contact, cx, assume_phase_dynamics): # Test the results np.random.seed(42) + time = np.random.rand(nlp.time.shape, nlp.ns) states = np.random.rand(nlp.states.shape, nlp.ns) controls = np.random.rand(nlp.controls.shape, nlp.ns) params = np.random.rand(nlp.parameters.shape, nlp.ns) stochastic_variables = np.random.rand(nlp.stochastic_variables.shape, nlp.ns) - x_out = np.array(nlp.dynamics_func(states, controls, params, stochastic_variables)) + x_out = np.array(nlp.dynamics_func(time, states, controls, params, stochastic_variables)) if with_contact: - contact_out = np.array(nlp.contact_forces_func(states, controls, params, stochastic_variables)) + contact_out = np.array(nlp.contact_forces_func(time, states, controls, params, stochastic_variables)) np.testing.assert_almost_equal( x_out[:, 0], [0.6118529, 0.785176, 0.6075449, 0.8083973, 0.3886773, 0.5426961, 0.7722448, 0.7290072] ) @@ -281,9 +287,11 @@ def test_torque_driven_soft_contacts_dynamics(with_contact, cx, implicit_contact phase_index = [i for i in range(ocp.n_phases)] NonLinearProgram.add(ocp, "phase_idx", phase_index, False) + use_time_from_phase_idx = [i for i in range(ocp.n_phases)] use_states_from_phase_idx = [i for i in range(ocp.n_phases)] use_states_dot_from_phase_idx = [i for i in range(ocp.n_phases)] use_controls_from_phase_idx = [i for i in range(ocp.n_phases)] + NonLinearProgram.add(ocp, "use_time_from_phase_idx", use_time_from_phase_idx, False) NonLinearProgram.add(ocp, "use_states_from_phase_idx", use_states_from_phase_idx, False) NonLinearProgram.add(ocp, "use_states_dot_from_phase_idx", use_states_dot_from_phase_idx, False) NonLinearProgram.add(ocp, "use_controls_from_phase_idx", use_controls_from_phase_idx, False) @@ -293,14 +301,15 @@ def test_torque_driven_soft_contacts_dynamics(with_contact, cx, implicit_contact # Test the results np.random.seed(42) + time = np.random.rand(nlp.time.shape, nlp.ns) states = np.random.rand(nlp.states.shape, nlp.ns) controls = np.random.rand(nlp.controls.shape, nlp.ns) params = np.random.rand(nlp.parameters.shape, nlp.ns) stochastic_variables = np.random.rand(nlp.stochastic_variables.shape, nlp.ns) - x_out = np.array(nlp.dynamics_func(states, controls, params, stochastic_variables)) + x_out = np.array(nlp.dynamics_func(time, states, controls, params, stochastic_variables)) if with_contact: - contact_out = np.array(nlp.contact_forces_func(states, controls, params, stochastic_variables)) + contact_out = np.array(nlp.contact_forces_func(time, states, controls, params, stochastic_variables)) np.testing.assert_almost_equal( x_out[:, 0], [0.6118529, 0.785176, 0.6075449, 0.8083973, -0.3214905, -0.1912131, 0.6507164, -0.2359716] ) @@ -348,9 +357,11 @@ def test_torque_derivative_driven(with_contact, with_external_force, cx, assume_ phase_index = [i for i in range(ocp.n_phases)] NonLinearProgram.add(ocp, "phase_idx", phase_index, False) + use_time_from_phase_idx = [i for i in range(ocp.n_phases)] use_states_from_phase_idx = [i for i in range(ocp.n_phases)] use_states_dot_from_phase_idx = [i for i in range(ocp.n_phases)] use_controls_from_phase_idx = [i for i in range(ocp.n_phases)] + NonLinearProgram.add(ocp, "use_time_from_phase_idx", use_time_from_phase_idx, False) NonLinearProgram.add(ocp, "use_states_from_phase_idx", use_states_from_phase_idx, False) NonLinearProgram.add(ocp, "use_states_dot_from_phase_idx", use_states_dot_from_phase_idx, False) NonLinearProgram.add(ocp, "use_controls_from_phase_idx", use_controls_from_phase_idx, False) @@ -366,14 +377,15 @@ def test_torque_derivative_driven(with_contact, with_external_force, cx, assume_ ConfigureProblem.initialize(ocp, nlp) # Test the results + time = np.random.rand(nlp.time.shape, nlp.ns) states = np.random.rand(nlp.states.shape, nlp.ns) controls = np.random.rand(nlp.controls.shape, nlp.ns) params = np.random.rand(nlp.parameters.shape, nlp.ns) stochastic_variables = np.random.rand(nlp.stochastic_variables.shape, nlp.ns) - x_out = np.array(nlp.dynamics_func(states, controls, params, stochastic_variables)) + x_out = np.array(nlp.dynamics_func(time, states, controls, params, stochastic_variables)) if with_contact: - contact_out = np.array(nlp.contact_forces_func(states, controls, params, stochastic_variables)) + contact_out = np.array(nlp.contact_forces_func(time, states, controls, params, stochastic_variables)) if with_external_force: np.testing.assert_almost_equal( x_out[:, 0], @@ -485,9 +497,11 @@ def test_torque_derivative_driven_implicit(with_contact, cx, assume_phase_dynami ) phase_index = [i for i in range(ocp.n_phases)] NonLinearProgram.add(ocp, "phase_idx", phase_index, False) + use_time_from_phase_idx = [i for i in range(ocp.n_phases)] use_states_from_phase_idx = [i for i in range(ocp.n_phases)] use_states_dot_from_phase_idx = [i for i in range(ocp.n_phases)] use_controls_from_phase_idx = [i for i in range(ocp.n_phases)] + NonLinearProgram.add(ocp, "use_time_from_phase_idx", use_time_from_phase_idx, False) NonLinearProgram.add(ocp, "use_states_from_phase_idx", use_states_from_phase_idx, False) NonLinearProgram.add(ocp, "use_states_dot_from_phase_idx", use_states_dot_from_phase_idx, False) NonLinearProgram.add(ocp, "use_controls_from_phase_idx", use_controls_from_phase_idx, False) @@ -497,14 +511,15 @@ def test_torque_derivative_driven_implicit(with_contact, cx, assume_phase_dynami # Test the results np.random.seed(42) + time = np.random.rand(nlp.time.shape, nlp.ns) states = np.random.rand(nlp.states.shape, nlp.ns) controls = np.random.rand(nlp.controls.shape, nlp.ns) params = np.random.rand(nlp.parameters.shape, nlp.ns) stochastic_variables = np.random.rand(nlp.stochastic_variables.shape, nlp.ns) - x_out = np.array(nlp.dynamics_func(states, controls, params, stochastic_variables)) + x_out = np.array(nlp.dynamics_func(time, states, controls, params, stochastic_variables)) if with_contact: - contact_out = np.array(nlp.contact_forces_func(states, controls, params, stochastic_variables)) + contact_out = np.array(nlp.contact_forces_func(time, states, controls, params, stochastic_variables)) np.testing.assert_almost_equal( x_out[:, 0], [ @@ -584,9 +599,11 @@ def test_torque_derivative_driven_soft_contacts_dynamics(with_contact, cx, impli phase_index = [i for i in range(ocp.n_phases)] NonLinearProgram.add(ocp, "phase_idx", phase_index, False) + use_time_from_phase_idx = [i for i in range(ocp.n_phases)] use_states_from_phase_idx = [i for i in range(ocp.n_phases)] use_states_dot_from_phase_idx = [i for i in range(ocp.n_phases)] use_controls_from_phase_idx = [i for i in range(ocp.n_phases)] + NonLinearProgram.add(ocp, "use_time_from_phase_idx", use_time_from_phase_idx, False) NonLinearProgram.add(ocp, "use_states_from_phase_idx", use_states_from_phase_idx, False) NonLinearProgram.add(ocp, "use_states_dot_from_phase_idx", use_states_dot_from_phase_idx, False) NonLinearProgram.add(ocp, "use_controls_from_phase_idx", use_controls_from_phase_idx, False) @@ -596,14 +613,15 @@ def test_torque_derivative_driven_soft_contacts_dynamics(with_contact, cx, impli # Test the results np.random.seed(42) + time = np.random.rand(nlp.time.shape, nlp.ns) states = np.random.rand(nlp.states.shape, nlp.ns) controls = np.random.rand(nlp.controls.shape, nlp.ns) params = np.random.rand(nlp.parameters.shape, nlp.ns) stochastic_variables = np.random.rand(nlp.stochastic_variables.shape, nlp.ns) - x_out = np.array(nlp.dynamics_func(states, controls, params, stochastic_variables)) + x_out = np.array(nlp.dynamics_func(time, states, controls, params, stochastic_variables)) if with_contact: - contact_out = np.array(nlp.contact_forces_func(states, controls, params, stochastic_variables)) + contact_out = np.array(nlp.contact_forces_func(time, states, controls, params, stochastic_variables)) np.testing.assert_almost_equal( x_out[:, 0], [ @@ -670,9 +688,11 @@ def test_soft_contacts_dynamics_errors(dynamics, assume_phase_dynamics): ) phase_index = [i for i in range(ocp.n_phases)] NonLinearProgram.add(ocp, "phase_idx", phase_index, False) + use_time_from_phase_idx = [i for i in range(ocp.n_phases)] use_states_from_phase_idx = [i for i in range(ocp.n_phases)] use_states_dot_from_phase_idx = [i for i in range(ocp.n_phases)] use_controls_from_phase_idx = [i for i in range(ocp.n_phases)] + NonLinearProgram.add(ocp, "use_time_from_phase_idx", use_time_from_phase_idx, False) NonLinearProgram.add(ocp, "use_states_from_phase_idx", use_states_from_phase_idx, False) NonLinearProgram.add(ocp, "use_states_dot_from_phase_idx", use_states_dot_from_phase_idx, False) NonLinearProgram.add(ocp, "use_controls_from_phase_idx", use_controls_from_phase_idx, False) @@ -712,9 +732,11 @@ def test_implicit_dynamics_errors(dynamics, assume_phase_dynamics): ) phase_index = [i for i in range(ocp.n_phases)] NonLinearProgram.add(ocp, "phase_idx", phase_index, False) + use_time_from_phase_idx = [i for i in range(ocp.n_phases)] use_states_from_phase_idx = [i for i in range(ocp.n_phases)] use_states_dot_from_phase_idx = [i for i in range(ocp.n_phases)] use_controls_from_phase_idx = [i for i in range(ocp.n_phases)] + NonLinearProgram.add(ocp, "use_time_from_phase_idx", use_time_from_phase_idx, False) NonLinearProgram.add(ocp, "use_states_from_phase_idx", use_states_from_phase_idx, False) NonLinearProgram.add(ocp, "use_states_dot_from_phase_idx", use_states_dot_from_phase_idx, False) NonLinearProgram.add(ocp, "use_controls_from_phase_idx", use_controls_from_phase_idx, False) @@ -757,9 +779,11 @@ def test_torque_activation_driven(with_contact, with_external_force, cx, assume_ ) phase_index = [i for i in range(ocp.n_phases)] NonLinearProgram.add(ocp, "phase_idx", phase_index, False) + use_time_from_phase_idx = [i for i in range(ocp.n_phases)] use_states_from_phase_idx = [i for i in range(ocp.n_phases)] use_states_dot_from_phase_idx = [i for i in range(ocp.n_phases)] use_controls_from_phase_idx = [i for i in range(ocp.n_phases)] + NonLinearProgram.add(ocp, "use_time_from_phase_idx", use_time_from_phase_idx, False) NonLinearProgram.add(ocp, "use_states_from_phase_idx", use_states_from_phase_idx, False) NonLinearProgram.add(ocp, "use_states_dot_from_phase_idx", use_states_dot_from_phase_idx, False) NonLinearProgram.add(ocp, "use_controls_from_phase_idx", use_controls_from_phase_idx, False) @@ -775,14 +799,15 @@ def test_torque_activation_driven(with_contact, with_external_force, cx, assume_ ConfigureProblem.initialize(ocp, nlp) # Test the results + time = np.random.rand(nlp.time.shape, nlp.ns) states = np.random.rand(nlp.states.shape, nlp.ns) controls = np.random.rand(nlp.controls.shape, nlp.ns) params = np.random.rand(nlp.parameters.shape, nlp.ns) stochastic_variables = np.random.rand(nlp.stochastic_variables.shape, nlp.ns) - x_out = np.array(nlp.dynamics_func(states, controls, params, stochastic_variables)) + x_out = np.array(nlp.dynamics_func(time, states, controls, params, stochastic_variables)) if with_contact: - contact_out = np.array(nlp.contact_forces_func(states, controls, params, stochastic_variables)) + contact_out = np.array(nlp.contact_forces_func(time, states, controls, params, stochastic_variables)) if with_external_force: np.testing.assert_almost_equal( x_out[:, 0], @@ -863,9 +888,11 @@ def test_torque_activation_driven_with_residual_torque( ) phase_index = [i for i in range(ocp.n_phases)] NonLinearProgram.add(ocp, "phase_idx", phase_index, False) + use_time_from_phase_idx = [i for i in range(ocp.n_phases)] use_states_from_phase_idx = [i for i in range(ocp.n_phases)] use_states_dot_from_phase_idx = [i for i in range(ocp.n_phases)] use_controls_from_phase_idx = [i for i in range(ocp.n_phases)] + NonLinearProgram.add(ocp, "use_time_from_phase_idx", use_time_from_phase_idx, False) NonLinearProgram.add(ocp, "use_states_from_phase_idx", use_states_from_phase_idx, False) NonLinearProgram.add(ocp, "use_states_dot_from_phase_idx", use_states_dot_from_phase_idx, False) NonLinearProgram.add(ocp, "use_controls_from_phase_idx", use_controls_from_phase_idx, False) @@ -881,11 +908,12 @@ def test_torque_activation_driven_with_residual_torque( ConfigureProblem.initialize(ocp, nlp) # Test the results + time = np.random.rand(nlp.time.shape, nlp.ns) states = np.random.rand(nlp.states.shape, nlp.ns) controls = np.random.rand(nlp.controls.shape, nlp.ns) params = np.random.rand(nlp.parameters.shape, nlp.ns) stochastic_variables = np.random.rand(nlp.stochastic_variables.shape, nlp.ns) - x_out = np.array(nlp.dynamics_func(states, controls, params, stochastic_variables)) + x_out = np.array(nlp.dynamics_func(time, states, controls, params, stochastic_variables)) if with_residual_torque: if with_external_force: @@ -990,9 +1018,11 @@ def test_muscle_driven( ) phase_index = [i for i in range(ocp.n_phases)] NonLinearProgram.add(ocp, "phase_idx", phase_index, False) + use_time_from_phase_idx = [i for i in range(ocp.n_phases)] use_states_from_phase_idx = [i for i in range(ocp.n_phases)] use_states_dot_from_phase_idx = [i for i in range(ocp.n_phases)] use_controls_from_phase_idx = [i for i in range(ocp.n_phases)] + NonLinearProgram.add(ocp, "use_time_from_phase_idx", use_time_from_phase_idx, False) NonLinearProgram.add(ocp, "use_states_from_phase_idx", use_states_from_phase_idx, False) NonLinearProgram.add(ocp, "use_states_dot_from_phase_idx", use_states_dot_from_phase_idx, False) NonLinearProgram.add(ocp, "use_controls_from_phase_idx", use_controls_from_phase_idx, False) @@ -1008,11 +1038,12 @@ def test_muscle_driven( ConfigureProblem.initialize(ocp, nlp) # Test the results + time = np.random.rand(nlp.time.shape, nlp.ns) states = np.random.rand(nlp.states.shape, nlp.ns) controls = np.random.rand(nlp.controls.shape, nlp.ns) params = np.random.rand(nlp.parameters.shape, nlp.ns) stochastic_variables = np.random.rand(nlp.stochastic_variables.shape, nlp.ns) - x_out = np.array(nlp.dynamics_func(states, controls, params, stochastic_variables)) + x_out = np.array(nlp.dynamics_func(time, states, controls, params, stochastic_variables)) if with_contact: # Warning this test is a bit bogus, there since the model does not have contacts if rigidbody_dynamics == RigidBodyDynamics.DAE_INVERSE_DYNAMICS: @@ -1496,9 +1527,11 @@ def test_joints_acceleration_driven(cx, rigid_body_dynamics, assume_phase_dynami np.random.seed(42) phase_index = [i for i in range(ocp.n_phases)] NonLinearProgram.add(ocp, "phase_idx", phase_index, False) + use_time_from_phase_idx = [i for i in range(ocp.n_phases)] use_states_from_phase_idx = [i for i in range(ocp.n_phases)] use_states_dot_from_phase_idx = [i for i in range(ocp.n_phases)] use_controls_from_phase_idx = [i for i in range(ocp.n_phases)] + NonLinearProgram.add(ocp, "use_time_from_phase_idx", use_time_from_phase_idx, False) NonLinearProgram.add(ocp, "use_states_from_phase_idx", use_states_from_phase_idx, False) NonLinearProgram.add(ocp, "use_states_dot_from_phase_idx", use_states_dot_from_phase_idx, False) NonLinearProgram.add(ocp, "use_controls_from_phase_idx", use_controls_from_phase_idx, False) @@ -1511,11 +1544,12 @@ def test_joints_acceleration_driven(cx, rigid_body_dynamics, assume_phase_dynami ConfigureProblem.initialize(ocp, nlp) # Test the results + time = np.random.rand(nlp.time.shape, nlp.ns) states = np.random.rand(nlp.states.shape, nlp.ns) controls = np.random.rand(nlp.controls.shape, nlp.ns) params = np.random.rand(nlp.parameters.shape, nlp.ns) stochastic_variables = np.random.rand(nlp.stochastic_variables.shape, nlp.ns) - x_out = np.array(nlp.dynamics_func(states, controls, params, stochastic_variables)) + x_out = np.array(nlp.dynamics_func(time, states, controls, params, stochastic_variables)) # obtained using Ipuch reference implementation. [https://github.com/Ipuch/OnDynamicsForSomersaults] np.testing.assert_almost_equal(x_out[:, 0], [0.02058449, 0.18340451, -2.95556261, 0.61185289]) @@ -1525,7 +1559,7 @@ def test_joints_acceleration_driven(cx, rigid_body_dynamics, assume_phase_dynami @pytest.mark.parametrize("with_contact", [False, True]) def test_custom_dynamics(with_contact, assume_phase_dynamics): def custom_dynamic( - states, controls, parameters, stochastic_variables, nlp, with_contact=False + time, states, controls, parameters, stochastic_variables, nlp, with_contact=False ) -> DynamicsEvaluation: q = DynamicsFunctions.get(nlp.states["q"], states) qdot = DynamicsFunctions.get(nlp.states["qdot"], states) @@ -1566,9 +1600,11 @@ def configure(ocp, nlp, with_contact=None): ) phase_index = [i for i in range(ocp.n_phases)] NonLinearProgram.add(ocp, "phase_idx", phase_index, False) + use_time_from_phase_idx = [i for i in range(ocp.n_phases)] use_states_from_phase_idx = [i for i in range(ocp.n_phases)] use_states_dot_from_phase_idx = [i for i in range(ocp.n_phases)] use_controls_from_phase_idx = [i for i in range(ocp.n_phases)] + NonLinearProgram.add(ocp, "use_time_from_phase_idx", use_time_from_phase_idx, False) NonLinearProgram.add(ocp, "use_states_from_phase_idx", use_states_from_phase_idx, False) NonLinearProgram.add(ocp, "use_states_dot_from_phase_idx", use_states_dot_from_phase_idx, False) NonLinearProgram.add(ocp, "use_controls_from_phase_idx", use_controls_from_phase_idx, False) @@ -1579,14 +1615,15 @@ def configure(ocp, nlp, with_contact=None): ConfigureProblem.initialize(ocp, nlp) # Test the results + time = np.random.rand(nlp.time.shape, nlp.ns) states = np.random.rand(nlp.states.shape, nlp.ns) controls = np.random.rand(nlp.controls.shape, nlp.ns) params = np.random.rand(nlp.parameters.shape, nlp.ns) stochastic_variables = np.random.rand(nlp.stochastic_variables.shape, nlp.ns) - x_out = np.array(nlp.dynamics_func(states, controls, params, stochastic_variables)) + x_out = np.array(nlp.dynamics_func(time, states, controls, params, stochastic_variables)) if with_contact: - contact_out = np.array(nlp.contact_forces_func(states, controls, params, stochastic_variables)) + contact_out = np.array(nlp.contact_forces_func(time, states, controls, params, stochastic_variables)) np.testing.assert_almost_equal( x_out[:, 0], [0.6118529, 0.785176, 0.6075449, 0.8083973, -0.3214905, -0.1912131, 0.6507164, -0.2359716] ) diff --git a/tests/test_ligaments.py b/tests/test_ligaments.py index f65ba6e34..d2dc26ade 100644 --- a/tests/test_ligaments.py +++ b/tests/test_ligaments.py @@ -57,9 +57,11 @@ def test_torque_driven_with_ligament(with_ligament, cx, assume_phase_dynamics): ) phase_index = [i for i in range(ocp.n_phases)] NonLinearProgram.add(ocp, "phase_idx", phase_index, False) + use_time_from_phase_idx = [i for i in range(ocp.n_phases)] use_states_from_phase_idx = [i for i in range(ocp.n_phases)] use_states_dot_from_phase_idx = [i for i in range(ocp.n_phases)] use_controls_from_phase_idx = [i for i in range(ocp.n_phases)] + NonLinearProgram.add(ocp, "use_time_from_phase_idx", use_time_from_phase_idx, False) NonLinearProgram.add(ocp, "use_states_from_phase_idx", use_states_from_phase_idx, False) NonLinearProgram.add(ocp, "use_states_dot_from_phase_idx", use_states_dot_from_phase_idx, False) NonLinearProgram.add(ocp, "use_controls_from_phase_idx", use_controls_from_phase_idx, False) @@ -70,11 +72,12 @@ def test_torque_driven_with_ligament(with_ligament, cx, assume_phase_dynamics): ConfigureProblem.initialize(ocp, nlp) # Test the results + time = np.random.rand(nlp.time.shape, nlp.ns) states = np.random.rand(nlp.states.shape, nlp.ns) controls = np.random.rand(nlp.controls.shape, nlp.ns) params = np.random.rand(nlp.parameters.shape, nlp.ns) stochastic_variables = np.random.rand(nlp.stochastic_variables.shape, nlp.ns) - x_out = np.array(nlp.dynamics_func(states, controls, params, stochastic_variables)) + x_out = np.array(nlp.dynamics_func(time, states, controls, params, stochastic_variables)) if with_ligament: np.testing.assert_almost_equal( x_out[:, 0], @@ -117,9 +120,11 @@ def test_torque_derivative_driven_with_ligament(with_ligament, cx, assume_phase_ phase_index = [i for i in range(ocp.n_phases)] NonLinearProgram.add(ocp, "phase_idx", phase_index, False) + use_time_from_phase_idx = [i for i in range(ocp.n_phases)] use_states_from_phase_idx = [i for i in range(ocp.n_phases)] use_states_dot_from_phase_idx = [i for i in range(ocp.n_phases)] use_controls_from_phase_idx = [i for i in range(ocp.n_phases)] + NonLinearProgram.add(ocp, "use_time_from_phase_idx", use_time_from_phase_idx, False) NonLinearProgram.add(ocp, "use_states_from_phase_idx", use_states_from_phase_idx, False) NonLinearProgram.add(ocp, "use_states_dot_from_phase_idx", use_states_dot_from_phase_idx, False) NonLinearProgram.add(ocp, "use_controls_from_phase_idx", use_controls_from_phase_idx, False) @@ -130,11 +135,12 @@ def test_torque_derivative_driven_with_ligament(with_ligament, cx, assume_phase_ ConfigureProblem.initialize(ocp, nlp) # Test the results + time = np.random.rand(nlp.time.shape, nlp.ns) states = np.random.rand(nlp.states.shape, nlp.ns) controls = np.random.rand(nlp.controls.shape, nlp.ns) params = np.random.rand(nlp.parameters.shape, nlp.ns) stochastic_variables = np.random.rand(nlp.stochastic_variables.shape, nlp.ns) - x_out = np.array(nlp.dynamics_func(states, controls, params, stochastic_variables)) + x_out = np.array(nlp.dynamics_func(time, states, controls, params, stochastic_variables)) if with_ligament: np.testing.assert_almost_equal( x_out[:, 0], @@ -174,9 +180,11 @@ def test_torque_activation_driven_with_ligament(with_ligament, cx, assume_phase_ ) phase_index = [i for i in range(ocp.n_phases)] NonLinearProgram.add(ocp, "phase_idx", phase_index, False) + use_time_from_phase_idx = [i for i in range(ocp.n_phases)] use_states_from_phase_idx = [i for i in range(ocp.n_phases)] use_states_dot_from_phase_idx = [i for i in range(ocp.n_phases)] use_controls_from_phase_idx = [i for i in range(ocp.n_phases)] + NonLinearProgram.add(ocp, "use_time_from_phase_idx", use_time_from_phase_idx, False) NonLinearProgram.add(ocp, "use_states_from_phase_idx", use_states_from_phase_idx, False) NonLinearProgram.add(ocp, "use_states_dot_from_phase_idx", use_states_dot_from_phase_idx, False) NonLinearProgram.add(ocp, "use_controls_from_phase_idx", use_controls_from_phase_idx, False) @@ -186,11 +194,12 @@ def test_torque_activation_driven_with_ligament(with_ligament, cx, assume_phase_ ConfigureProblem.initialize(ocp, nlp) # Test the results + time = np.random.rand(nlp.time.shape, nlp.ns) states = np.random.rand(nlp.states.shape, nlp.ns) controls = np.random.rand(nlp.controls.shape, nlp.ns) params = np.random.rand(nlp.parameters.shape, nlp.ns) stochastic_variables = np.random.rand(nlp.stochastic_variables.shape, nlp.ns) - x_out = np.array(nlp.dynamics_func(states, controls, params, stochastic_variables)) + x_out = np.array(nlp.dynamics_func(time, states, controls, params, stochastic_variables)) if with_ligament: np.testing.assert_almost_equal( x_out[:, 0], @@ -237,9 +246,11 @@ def test_muscle_driven_with_ligament(with_ligament, cx, assume_phase_dynamics): ) phase_index = [i for i in range(ocp.n_phases)] NonLinearProgram.add(ocp, "phase_idx", phase_index, False) + use_time_from_phase_idx = [i for i in range(ocp.n_phases)] use_states_from_phase_idx = [i for i in range(ocp.n_phases)] use_states_dot_from_phase_idx = [i for i in range(ocp.n_phases)] use_controls_from_phase_idx = [i for i in range(ocp.n_phases)] + NonLinearProgram.add(ocp, "use_time_from_phase_idx", use_time_from_phase_idx, False) NonLinearProgram.add(ocp, "use_states_from_phase_idx", use_states_from_phase_idx, False) NonLinearProgram.add(ocp, "use_states_dot_from_phase_idx", use_states_dot_from_phase_idx, False) NonLinearProgram.add(ocp, "use_controls_from_phase_idx", use_controls_from_phase_idx, False) @@ -250,11 +261,12 @@ def test_muscle_driven_with_ligament(with_ligament, cx, assume_phase_dynamics): ConfigureProblem.initialize(ocp, nlp) # Test the results + time = np.random.rand(nlp.time.shape, nlp.ns) states = np.random.rand(nlp.states.shape, nlp.ns) controls = np.random.rand(nlp.controls.shape, nlp.ns) params = np.random.rand(nlp.parameters.shape, nlp.ns) stochastic_variables = np.random.rand(nlp.stochastic_variables.shape, nlp.ns) - x_out = np.array(nlp.dynamics_func(states, controls, params, stochastic_variables)) + x_out = np.array(nlp.dynamics_func(time, states, controls, params, stochastic_variables)) if with_ligament: np.testing.assert_almost_equal( diff --git a/tests/test_passive_torque.py b/tests/test_passive_torque.py index 59eb75bf2..f458dac79 100644 --- a/tests/test_passive_torque.py +++ b/tests/test_passive_torque.py @@ -62,9 +62,11 @@ def test_torque_driven_with_passive_torque(with_passive_torque, cx, rigidbody_dy ) phase_index = [i for i in range(ocp.n_phases)] NonLinearProgram.add(ocp, "phase_idx", phase_index, False) + use_time_from_phase_idx = [i for i in range(ocp.n_phases)] use_states_from_phase_idx = [i for i in range(ocp.n_phases)] use_states_dot_from_phase_idx = [i for i in range(ocp.n_phases)] use_controls_from_phase_idx = [i for i in range(ocp.n_phases)] + NonLinearProgram.add(ocp, "use_time_from_phase_idx", use_time_from_phase_idx, False) NonLinearProgram.add(ocp, "use_states_from_phase_idx", use_states_from_phase_idx, False) NonLinearProgram.add(ocp, "use_states_dot_from_phase_idx", use_states_dot_from_phase_idx, False) NonLinearProgram.add(ocp, "use_controls_from_phase_idx", use_controls_from_phase_idx, False) @@ -75,11 +77,12 @@ def test_torque_driven_with_passive_torque(with_passive_torque, cx, rigidbody_dy ConfigureProblem.initialize(ocp, nlp) # Test the results + time = np.random.rand(nlp.time.shape, nlp.ns) states = np.random.rand(nlp.states.shape, nlp.ns) controls = np.random.rand(nlp.controls.shape, nlp.ns) params = np.random.rand(nlp.parameters.shape, nlp.ns) stochastic_variables = np.random.rand(nlp.stochastic_variables.shape, nlp.ns) - x_out = np.array(nlp.dynamics_func(states, controls, params, stochastic_variables)) + x_out = np.array(nlp.dynamics_func(time, states, controls, params, stochastic_variables)) if rigidbody_dynamics == RigidBodyDynamics.ODE: if with_passive_torque: np.testing.assert_almost_equal( @@ -143,9 +146,11 @@ def test_torque_derivative_driven_with_passive_torque(with_passive_torque, cx, a phase_index = [i for i in range(ocp.n_phases)] NonLinearProgram.add(ocp, "phase_idx", phase_index, False) + use_time_from_phase_idx = [i for i in range(ocp.n_phases)] use_states_from_phase_idx = [i for i in range(ocp.n_phases)] use_states_dot_from_phase_idx = [i for i in range(ocp.n_phases)] use_controls_from_phase_idx = [i for i in range(ocp.n_phases)] + NonLinearProgram.add(ocp, "use_time_from_phase_idx", use_time_from_phase_idx, False) NonLinearProgram.add(ocp, "use_states_from_phase_idx", use_states_from_phase_idx, False) NonLinearProgram.add(ocp, "use_states_dot_from_phase_idx", use_states_dot_from_phase_idx, False) NonLinearProgram.add(ocp, "use_controls_from_phase_idx", use_controls_from_phase_idx, False) @@ -156,11 +161,12 @@ def test_torque_derivative_driven_with_passive_torque(with_passive_torque, cx, a ConfigureProblem.initialize(ocp, nlp) # Test the results + time = np.random.rand(nlp.time.shape, nlp.ns) states = np.random.rand(nlp.states.shape, nlp.ns) controls = np.random.rand(nlp.controls.shape, nlp.ns) params = np.random.rand(nlp.parameters.shape, nlp.ns) stochastic_variables = np.random.rand(nlp.stochastic_variables.shape, nlp.ns) - x_out = np.array(nlp.dynamics_func(states, controls, params, stochastic_variables)) + x_out = np.array(nlp.dynamics_func(time, states, controls, params, stochastic_variables)) if with_passive_torque: np.testing.assert_almost_equal( x_out[:, 0], @@ -236,9 +242,11 @@ def test_torque_activation_driven_with_passive_torque( ) phase_index = [i for i in range(ocp.n_phases)] NonLinearProgram.add(ocp, "phase_idx", phase_index, False) + use_time_from_phase_idx = [i for i in range(ocp.n_phases)] use_states_from_phase_idx = [i for i in range(ocp.n_phases)] use_states_dot_from_phase_idx = [i for i in range(ocp.n_phases)] use_controls_from_phase_idx = [i for i in range(ocp.n_phases)] + NonLinearProgram.add(ocp, "use_time_from_phase_idx", use_time_from_phase_idx, False) NonLinearProgram.add(ocp, "use_states_from_phase_idx", use_states_from_phase_idx, False) NonLinearProgram.add(ocp, "use_states_dot_from_phase_idx", use_states_dot_from_phase_idx, False) NonLinearProgram.add(ocp, "use_controls_from_phase_idx", use_controls_from_phase_idx, False) @@ -249,11 +257,12 @@ def test_torque_activation_driven_with_passive_torque( ConfigureProblem.initialize(ocp, nlp) # Test the results + time = np.random.rand(nlp.time.shape, nlp.ns) states = np.random.rand(nlp.states.shape, nlp.ns) controls = np.random.rand(nlp.controls.shape, nlp.ns) params = np.random.rand(nlp.parameters.shape, nlp.ns) stochastic_variables = np.random.rand(nlp.stochastic_variables.shape, nlp.ns) - x_out = np.array(nlp.dynamics_func(states, controls, params, stochastic_variables)) + x_out = np.array(nlp.dynamics_func(time, states, controls, params, stochastic_variables)) if with_residual_torque: if with_passive_torque: np.testing.assert_almost_equal( @@ -350,9 +359,11 @@ def test_muscle_driven_with_passive_torque(with_passive_torque, rigidbody_dynami ) phase_index = [i for i in range(ocp.n_phases)] NonLinearProgram.add(ocp, "phase_idx", phase_index, False) + use_time_from_phase_idx = [i for i in range(ocp.n_phases)] use_states_from_phase_idx = [i for i in range(ocp.n_phases)] use_states_dot_from_phase_idx = [i for i in range(ocp.n_phases)] use_controls_from_phase_idx = [i for i in range(ocp.n_phases)] + NonLinearProgram.add(ocp, "use_time_from_phase_idx", use_time_from_phase_idx, False) NonLinearProgram.add(ocp, "use_states_from_phase_idx", use_states_from_phase_idx, False) NonLinearProgram.add(ocp, "use_states_dot_from_phase_idx", use_states_dot_from_phase_idx, False) NonLinearProgram.add(ocp, "use_controls_from_phase_idx", use_controls_from_phase_idx, False) @@ -365,11 +376,12 @@ def test_muscle_driven_with_passive_torque(with_passive_torque, rigidbody_dynami ConfigureProblem.initialize(ocp, nlp) # Test the results + time = np.random.rand(nlp.time.shape, nlp.ns) states = np.random.rand(nlp.states.shape, nlp.ns) controls = np.random.rand(nlp.controls.shape, nlp.ns) params = np.random.rand(nlp.parameters.shape, nlp.ns) stochastic_variables = np.random.rand(nlp.stochastic_variables.shape, nlp.ns) - x_out = np.array(nlp.dynamics_func(states, controls, params, stochastic_variables)) + x_out = np.array(nlp.dynamics_func(time, states, controls, params, stochastic_variables)) if rigidbody_dynamics == RigidBodyDynamics.DAE_INVERSE_DYNAMICS: if with_passive_torque: diff --git a/tests/test_penalty.py b/tests/test_penalty.py index 5471df168..b79a32a42 100644 --- a/tests/test_penalty.py +++ b/tests/test_penalty.py @@ -75,6 +75,7 @@ def get_penalty_value(ocp, penalty, t, x, u, p, s): if isinstance(val, float): return val + time = ocp.nlp[0].time.cx_start if ocp.nlp[0].time.cx_start.shape != (0, 0) else ocp.cx(0, 0) states = ocp.nlp[0].states.cx_start if ocp.nlp[0].states.cx_start.shape != (0, 0) else ocp.cx(0, 0) controls = ocp.nlp[0].controls.cx_start if ocp.nlp[0].controls.cx_start.shape != (0, 0) else ocp.cx(0, 0) parameters = ocp.nlp[0].parameters.cx if ocp.nlp[0].parameters.cx.shape != (0, 0) else ocp.cx(0, 0) @@ -83,8 +84,8 @@ def get_penalty_value(ocp, penalty, t, x, u, p, s): if ocp.nlp[0].stochastic_variables.cx_start.shape != (0, 0) else ocp.cx(0, 0) ) - return ocp.nlp[0].to_casadi_func("penalty", val, states, controls, parameters, stochastic_variables)( - x[0], u[0], p, s + return ocp.nlp[0].to_casadi_func("penalty", val, time, states, controls, parameters, stochastic_variables)( + t, x[0], u[0], p, s )