diff --git a/bioptim/dynamics/configure_new_variable.py b/bioptim/dynamics/configure_new_variable.py index ac5bc6ad0..422a09c7d 100644 --- a/bioptim/dynamics/configure_new_variable.py +++ b/bioptim/dynamics/configure_new_variable.py @@ -518,6 +518,20 @@ def _declare_cx_and_plot(self): self.nlp.variable_mappings[self.name], node_index, ) + if not self.skip_plot: + all_variables_in_one_subplot = True if self.name in ["m", "cov", "k"] else False + self.nlp.plot[f"{self.name}_algebraic"] = CustomPlot( + lambda t0, phases_dt, node_idx, x, u, p, a: ( + a[self.nlp.algebraic_states.key_index(self.name), :] + if a.any() + else np.ndarray((cx[0][0].shape[0], 1)) * np.nan + ), + plot_type=PlotType.STEP, + axes_idx=self.axes_idx, + legend=self.legend, + combine_to=self.combine_name, + all_variables_in_one_subplot=all_variables_in_one_subplot, + ) def _manage_fatigue_to_new_variable( diff --git a/bioptim/dynamics/configure_problem.py b/bioptim/dynamics/configure_problem.py index 4ee4d3a1e..a60b4f280 100644 --- a/bioptim/dynamics/configure_problem.py +++ b/bioptim/dynamics/configure_problem.py @@ -908,31 +908,85 @@ def configure_dynamics_function(ocp, nlp, dyn_func, **extra_params): A reference to the ocp nlp: NonLinearProgram A reference to the phase - dyn_func: Callable[time, states, controls, param, algebraic_states] | tuple[Callable[time, states, controls, param, algebraic_states], ...] + dyn_func: Callable[time, states, controls, param, algebraic_states] The function to get the derivative of the states """ DynamicsFunctions.apply_parameters(nlp) - if not isinstance(dyn_func, (tuple, list)): - dyn_func = (dyn_func,) + dynamics_eval = dyn_func( + nlp.time_mx, + nlp.states.scaled.mx_reduced, + nlp.controls.scaled.mx_reduced, + nlp.parameters.scaled.mx_reduced, + nlp.algebraic_states.scaled.mx_reduced, + nlp, + **extra_params, + ) + dynamics_dxdt = dynamics_eval.dxdt + if isinstance(dynamics_dxdt, (list, tuple)): + dynamics_dxdt = vertcat(*dynamics_dxdt) - for func in dyn_func: - dynamics_eval = func( - nlp.time_mx, - nlp.states.scaled.mx_reduced, - nlp.controls.scaled.mx_reduced, - nlp.parameters.scaled.mx_reduced, - nlp.algebraic_states.scaled.mx_reduced, - nlp, - **extra_params, + time_span_sym = vertcat(nlp.time_mx, nlp.dt_mx) + if nlp.dynamics_func is None: + nlp.dynamics_func = Function( + "ForwardDyn", + [ + time_span_sym, + nlp.states.scaled.mx_reduced, + nlp.controls.scaled.mx_reduced, + nlp.parameters.scaled.mx_reduced, + nlp.algebraic_states.scaled.mx_reduced, + ], + [dynamics_dxdt], + ["t_span", "x", "u", "p", "a"], + ["xdot"], ) - dynamics_dxdt = dynamics_eval.dxdt - if isinstance(dynamics_dxdt, (list, tuple)): - dynamics_dxdt = vertcat(*dynamics_dxdt) - time_span_sym = vertcat(nlp.time_mx, nlp.dt_mx) - nlp.dynamics_func.append( + # TODO: allow expand for each dynamics independently + if nlp.dynamics_type.expand_dynamics: + try: + nlp.dynamics_func = nlp.dynamics_func.expand() + except Exception as me: + RuntimeError( + f"An error occurred while executing the 'expand()' function for the dynamic function. " + f"Please review the following casadi error message for more details.\n" + "Several factors could be causing this issue. One of the most likely is the inability to " + "use expand=True at all. In that case, try adding expand=False to the dynamics.\n" + "Original casadi error message:\n" + f"{me}" + ) + + # Only possible for regular dynamics, not for extra_dynamics + if dynamics_eval.defects is not None: + nlp.implicit_dynamics_func = Function( + "DynamicsDefects", + [ + time_span_sym, + nlp.states.scaled.mx_reduced, + nlp.controls.scaled.mx_reduced, + nlp.parameters.scaled.mx_reduced, + nlp.algebraic_states.scaled.mx_reduced, + nlp.states_dot.scaled.mx_reduced, + ], + [dynamics_eval.defects], + ["t_span", "x", "u", "p", "a", "xdot"], + ["defects"], + ) + if nlp.dynamics_type.expand_dynamics: + try: + nlp.implicit_dynamics_func = nlp.implicit_dynamics_func.expand() + except Exception as me: + RuntimeError( + f"An error occurred while executing the 'expand()' function for the dynamic function. " + f"Please review the following casadi error message for more details.\n" + "Several factors could be causing this issue. One of the most likely is the inability to " + "use expand=True at all. In that case, try adding expand=False to the dynamics.\n" + "Original casadi error message:\n" + f"{me}" + ) + else: + nlp.extra_dynamics_func.append( Function( "ForwardDyn", [ @@ -951,7 +1005,7 @@ def configure_dynamics_function(ocp, nlp, dyn_func, **extra_params): # TODO: allow expand for each dynamics independently if nlp.dynamics_type.expand_dynamics: try: - nlp.dynamics_func[-1] = nlp.dynamics_func[-1].expand() + nlp.extra_dynamics_func[-1] = nlp.dynamics_funcextra_dynamics_func[-1].expand() except Exception as me: RuntimeError( f"An error occurred while executing the 'expand()' function for the dynamic function. " @@ -962,36 +1016,6 @@ def configure_dynamics_function(ocp, nlp, dyn_func, **extra_params): f"{me}" ) - if dynamics_eval.defects is not None: - nlp.implicit_dynamics_func.append( - Function( - "DynamicsDefects", - [ - time_span_sym, - nlp.states.scaled.mx_reduced, - nlp.controls.scaled.mx_reduced, - nlp.parameters.scaled.mx_reduced, - nlp.algebraic_states.scaled.mx_reduced, - nlp.states_dot.scaled.mx_reduced, - ], - [dynamics_eval.defects], - ["t_span", "x", "u", "p", "a", "xdot"], - ["defects"], - ) - ) - if nlp.dynamics_type.expand_dynamics: - try: - nlp.implicit_dynamics_func[-1] = nlp.implicit_dynamics_func[-1].expand() - except Exception as me: - RuntimeError( - f"An error occurred while executing the 'expand()' function for the dynamic function. " - f"Please review the following casadi error message for more details.\n" - "Several factors could be causing this issue. One of the most likely is the inability to " - "use expand=True at all. In that case, try adding expand=False to the dynamics.\n" - "Original casadi error message:\n" - f"{me}" - ) - @staticmethod def configure_contact_function(ocp, nlp, dyn_func: Callable, **extra_params): """ @@ -1369,7 +1393,6 @@ def configure_stochastic_k(ocp, nlp, n_noised_controls: int, n_references: int): as_controls=False, as_states_dot=False, as_algebraic_states=True, - skip_plot=True, ) @staticmethod @@ -1493,7 +1516,6 @@ def configure_stochastic_cov_implicit(ocp, nlp, n_noised_states: int): as_controls=False, as_states_dot=False, as_algebraic_states=True, - skip_plot=True, ) @staticmethod @@ -1525,7 +1547,6 @@ def configure_stochastic_cholesky_cov(ocp, nlp, n_noised_states: int): as_controls=False, as_states_dot=False, as_algebraic_states=True, - skip_plot=True, ) @staticmethod @@ -1554,7 +1575,6 @@ def configure_stochastic_ref(ocp, nlp, n_references: int): as_controls=False, as_states_dot=False, as_algebraic_states=True, - skip_plot=True, ) @staticmethod @@ -1589,7 +1609,6 @@ def configure_stochastic_m(ocp, nlp, n_noised_states: int, n_collocation_points: as_controls=False, as_states_dot=False, as_algebraic_states=True, - skip_plot=True, ) @staticmethod diff --git a/bioptim/dynamics/ode_solver.py b/bioptim/dynamics/ode_solver.py index 43c0ada4d..7a34cc24f 100644 --- a/bioptim/dynamics/ode_solver.py +++ b/bioptim/dynamics/ode_solver.py @@ -159,7 +159,9 @@ def param_ode(self, nlp) -> MX: """ return nlp.parameters.scaled.cx_start - def initialize_integrator(self, ocp, nlp, dynamics_index: int, node_index: int, **extra_opt) -> Callable: + def initialize_integrator( + self, ocp, nlp, node_index: int, dynamics_index: int = 0, is_extra_dynamics: bool = False, **extra_opt + ) -> Callable: """ Initialize the integrator @@ -169,10 +171,12 @@ def initialize_integrator(self, ocp, nlp, dynamics_index: int, node_index: int, The Optimal control program handler nlp The NonLinearProgram handler - dynamics_index - The current dynamics to resolve (that can be referred to nlp.dynamics_func[index]) node_index The index of the node currently initialized + dynamics_index + The current extra dynamics to resolve (that can be referred to nlp.extra_dynamics_func[index]) + is_extra_dynamics + If the dynamics is an extra dynamics extra_opt Any extra options to pass to the integrator @@ -181,16 +185,20 @@ def initialize_integrator(self, ocp, nlp, dynamics_index: int, node_index: int, The initialized integrator function """ + if dynamics_index > 0 and not is_extra_dynamics: + raise RuntimeError("dynamics_index should be 0 if is_extra_dynamics is False") + nlp.states.node_index = node_index nlp.states_dot.node_index = node_index nlp.controls.node_index = node_index nlp.algebraic_states.node_index = node_index + dynamics_func = nlp.dynamics_func if not is_extra_dynamics else nlp.extra_dynamics_func[dynamics_index] ode_opt = { "model": nlp.model, "cx": nlp.cx, "control_type": nlp.control_type, "defects_type": self.defects_type, - "ode_index": node_index if nlp.dynamics_func[dynamics_index].size2_out("xdot") > 1 else 0, + "ode_index": node_index if dynamics_func.size2_out("xdot") > 1 else 0, "duplicate_starting_point": self.duplicate_starting_point, **extra_opt, } @@ -201,13 +209,8 @@ def initialize_integrator(self, ocp, nlp, dynamics_index: int, node_index: int, "u": self.p_ode(nlp), "a": self.a_ode(nlp), "param": self.param_ode(nlp), - "ode": nlp.dynamics_func[dynamics_index], - # TODO this actually checks "not nlp.implicit_dynamics_func" (or that nlp.implicit_dynamics_func == []) - "implicit_ode": ( - nlp.implicit_dynamics_func[dynamics_index] - if len(nlp.implicit_dynamics_func) > 0 - else nlp.implicit_dynamics_func - ), + "ode": dynamics_func, + "implicit_ode": nlp.implicit_dynamics_func, } return nlp.ode_solver.integrator(ode, ode_opt) @@ -235,14 +238,19 @@ def prepare_dynamic_integrator(self, ocp, nlp): # Extra dynamics extra_dynamics = [] - for i in range(1, len(nlp.dynamics_func)): - extra_dynamics += [nlp.ode_solver.initialize_integrator(ocp, nlp, dynamics_index=i, node_index=0)] + for i in range(len(nlp.extra_dynamics_func)): + extra_dynamics += [ + nlp.ode_solver.initialize_integrator(ocp, nlp, dynamics_index=i, node_index=0, is_extra_dynamics=True) + ] if nlp.phase_dynamics == PhaseDynamics.SHARED_DURING_THE_PHASE: extra_dynamics = extra_dynamics * nlp.ns else: for node_index in range(1, nlp.ns): - extra_dynamics += [nlp.ode_solver.initialize_integrator(ocp, nlp, dynamics_index=i, node_index=0)] - # TODO include this in nlp.dynamics so the index of nlp.dynamics_func and nlp.dynamics match + extra_dynamics += [ + nlp.ode_solver.initialize_integrator( + ocp, nlp, dynamics_index=i, node_index=node_index, is_extra_dynamics=True + ) + ] nlp.extra_dynamics.append(extra_dynamics) @@ -538,7 +546,9 @@ def p_ode(self, nlp): def a_ode(self, nlp): return nlp.algebraic_states.scaled.cx - def initialize_integrator(self, ocp, nlp, dynamics_index: int, node_index: int, **extra_opt): + def initialize_integrator( + self, ocp, nlp, node_index: int, dynamics_index: int = 0, is_extra_dynamics: bool = False, **extra_opt + ): raise NotImplementedError("CVODES is not yet implemented") if extra_opt: @@ -560,10 +570,11 @@ def initialize_integrator(self, ocp, nlp, dynamics_index: int, node_index: int, raise RuntimeError("CVODES cannot be used with algebraic_states variables") t = [self.t_ode(nlp)[0], self.t_ode(nlp)[1] - self.t_ode(nlp)[0]] + dynamics_func = nlp.dynamics_func if not is_extra_dynamics else nlp.extra_dynamics_func[dynamics_index] ode = { "x": nlp.states.scaled.cx_start, "u": nlp.controls.scaled.cx_start, # todo: add p=parameters - "ode": nlp.dynamics_func[dynamics_index]( + "ode": dynamics_func( vertcat(*t), self.x_ode(nlp), self.p_ode(nlp), self.param_ode(nlp), self.a_ode(nlp) ), } diff --git a/bioptim/examples/getting_started/pendulum.py b/bioptim/examples/getting_started/pendulum.py index b936f9c5f..50eeaaa13 100644 --- a/bioptim/examples/getting_started/pendulum.py +++ b/bioptim/examples/getting_started/pendulum.py @@ -136,7 +136,9 @@ def main(): ocp = prepare_ocp(biorbd_model_path="models/pendulum.bioMod", final_time=1, n_shooting=400, n_threads=2) # Custom plots - ocp.add_plot_penalty(CostType.ALL) + ocp.add_plot_penalty(CostType.ALL) # This will display the objectives and constraints at the current iteration + # ocp.add_plot_ipopt_outputs() # This will display the solver's output at the current iteration + # ocp.add_plot_check_conditioning() # This will display the conditioning of the problem at the current iteration # --- If one is interested in checking the conditioning of the problem, they can uncomment the following line --- # # ocp.check_conditioning() @@ -146,12 +148,60 @@ def main(): # --- Solve the ocp. Please note that online graphics only works with the Linux operating system --- # sol = ocp.solve(Solver.IPOPT(show_online_optim=platform.system() == "Linux")) + sol.print_cost() # --- Show the results (graph or animation) --- # - # sol.graphs(show_bounds=True) + # sol.graphs(show_bounds=True, save_name="results.png") sol.animate(n_frames=100) + # # --- Save the solution --- # + # Here is an example of how we recommend to save the solution. Please note that sol.ocp is not picklable and that sol will be loaded using the current bioptim version, not the version at the time of the generation of the results. + # import pickle + # import git + # from datetime import date + # + # # Save the version of bioptim and the date of the optimization for future reference + # repo = git.Repo(search_parent_directories=True) + # commit_id = str(repo.commit()) + # branch = str(repo.active_branch) + # tag = repo.git.describe("--tags") + # bioptim_version = repo.git.version_info + # git_date = repo.git.log("-1", "--format=%cd") + # version_dic = { + # "commit_id": commit_id, + # "git_date": git_date, + # "branch": branch, + # "tag": tag, + # "bioptim_version": bioptim_version, + # "date_of_the_optimization": date.today().strftime("%b-%d-%Y-%H-%M-%S"), + # } + # + # q = sol.decision_states(to_merge=SolutionMerge.NODES)["q"] + # qdot = sol.decision_states(to_merge=SolutionMerge.NODES)["qdot"] + # tau = sol.decision_controls(to_merge=SolutionMerge.NODES)["tau"] + # + # # Do everything you need with the solution here before we delete ocp + # integrated_sol = sol.integrate(to_merge=SolutionMerge.NODES) + # q_integrated = integrated_sol["q"] + # qdot_integrated = integrated_sol["qdot"] + # + # # Save the output of the optimization + # with open("pendulum_data.pkl", "wb") as file: + # data = {"q": q, + # "qdot": qdot, + # "tau": tau, + # "real_time_to_optimize": sol.real_time_to_optimize, + # "version": version_dic, + # "q_integrated": q_integrated, + # "qdot_integrated": qdot_integrated} + # pickle.dump(data, file) + # + # # Save the solution for future use, you will only need to do sol.ocp = prepare_ocp() to get the same solution object as above. + # with open("pendulum_sol.pkl", "wb") as file: + # del sol.ocp + # pickle.dump(sol, file) + if __name__ == "__main__": main() diff --git a/bioptim/examples/stochastic_optimal_control/arm_reaching_muscle_driven.py b/bioptim/examples/stochastic_optimal_control/arm_reaching_muscle_driven.py index bbfb4d907..3180b03a7 100644 --- a/bioptim/examples/stochastic_optimal_control/arm_reaching_muscle_driven.py +++ b/bioptim/examples/stochastic_optimal_control/arm_reaching_muscle_driven.py @@ -43,7 +43,7 @@ ControlType, ) -from bioptim.examples.stochastic_optimal_control.leuven_arm_model import LeuvenArmModel +from bioptim.examples.stochastic_optimal_control.models.leuven_arm_model import LeuvenArmModel from bioptim.examples.stochastic_optimal_control.arm_reaching_torque_driven_implicit import ExampleType diff --git a/bioptim/examples/stochastic_optimal_control/leuven_arm_model.py b/bioptim/examples/stochastic_optimal_control/models/leuven_arm_model.py similarity index 100% rename from bioptim/examples/stochastic_optimal_control/leuven_arm_model.py rename to bioptim/examples/stochastic_optimal_control/models/leuven_arm_model.py diff --git a/bioptim/examples/stochastic_optimal_control/mass_point_model.py b/bioptim/examples/stochastic_optimal_control/models/mass_point_model.py similarity index 96% rename from bioptim/examples/stochastic_optimal_control/mass_point_model.py rename to bioptim/examples/stochastic_optimal_control/models/mass_point_model.py index 4879aab79..8508222fe 100644 --- a/bioptim/examples/stochastic_optimal_control/mass_point_model.py +++ b/bioptim/examples/stochastic_optimal_control/models/mass_point_model.py @@ -91,10 +91,13 @@ def dynamics_numerical(self, states, controls, motor_noise=0): """ The dynamics from equation (22). """ + # to avoid dimension pb with solve_ivp + if states.ndim == 2: + states = states.reshape((-1,)) + q = states[: self.nb_q] qdot = states[self.nb_q :] u = controls - qddot = -self.kapa * (q - u) - self.beta * qdot * sqrt(qdot[0] ** 2 + qdot[1] ** 2 + self.c**2) + motor_noise return vertcat(qdot, qddot) diff --git a/bioptim/examples/stochastic_optimal_control/rockit_model.py b/bioptim/examples/stochastic_optimal_control/models/rockit_model.py similarity index 97% rename from bioptim/examples/stochastic_optimal_control/rockit_model.py rename to bioptim/examples/stochastic_optimal_control/models/rockit_model.py index 57bea806f..079943573 100644 --- a/bioptim/examples/stochastic_optimal_control/rockit_model.py +++ b/bioptim/examples/stochastic_optimal_control/models/rockit_model.py @@ -80,7 +80,7 @@ def dynamics_numerical(self, states, controls, motor_noise=0): qdot = states[self.nb_q :] u = controls - qddot = -0.1 * (1 - q**2) * qdot - q + u # + motor_noise + qddot = -0.1 * (1 - q**2) * qdot - q + u + motor_noise return vertcat(qdot, qddot) diff --git a/bioptim/examples/stochastic_optimal_control/obstacle_avoidance_direct_collocation.py b/bioptim/examples/stochastic_optimal_control/obstacle_avoidance_direct_collocation.py index d503c8fb3..8e76c934b 100644 --- a/bioptim/examples/stochastic_optimal_control/obstacle_avoidance_direct_collocation.py +++ b/bioptim/examples/stochastic_optimal_control/obstacle_avoidance_direct_collocation.py @@ -6,8 +6,10 @@ import matplotlib.pyplot as plt from matplotlib.patches import Ellipse +import matplotlib.cm as cm import casadi as cas import numpy as np +import pickle from bioptim import ( StochasticOptimalControlProgram, @@ -33,15 +35,267 @@ PhaseDynamics, BoundsList, SolutionMerge, + SolutionIntegrator, + Shooting, ) -from bioptim.examples.stochastic_optimal_control.mass_point_model import MassPointModel +from bioptim.examples.stochastic_optimal_control.models.mass_point_model import MassPointModel from bioptim.examples.stochastic_optimal_control.common import ( test_matrix_semi_definite_positiveness, test_eigen_values, reshape_to_matrix, ) +from scipy.integrate import solve_ivp + + +def plot_results( + sol_socp, + states, + controls, + time, + algebraic_states, + bio_model, + motor_noise_magnitude, + n_shooting, + polynomial_degree, + is_stochastic, + q_init, +): + """ + This function plots the reintegration of the optimal solution considering the motor noise. + The plot compares the covariance obtained numerically by doing 100 orbits, the covariance obtained by the optimal control problem and the covariance obtained by the noisy integration. + """ + q = states["q"] + qdot = states["qdot"] + u = controls["u"] + Tf = time[-1] + tgrid = np.linspace(0, Tf, n_shooting + 1).squeeze() + + fig, ax = plt.subplots(2, 2) + fig_comparison, ax_comparison = plt.subplots(1, 1) + for i in range(2): + a = bio_model.super_ellipse_a[i] + b = bio_model.super_ellipse_b[i] + n = bio_model.super_ellipse_n[i] + x_0 = bio_model.super_ellipse_center_x[i] + y_0 = bio_model.super_ellipse_center_y[i] + + X, Y, Z = superellipse(a, b, n, x_0, y_0) + + ax[0, 0].contourf(X, Y, Z, levels=[-1000, 0], colors=["#DA1984"], alpha=0.5, label="Obstacles") + ax_comparison.contourf(X, Y, Z, levels=[-1000, 0], colors=["#DA1984"], alpha=0.5, label="Obstacles") + + ax[0, 0].plot(q_init[0], q_init[1], "-k", label="Initial guess") + ax[0, 0].plot(q[0][0], q[1][0], "og", label="Optimal initial node") + ax[0, 0].plot(q[0], q[1], "-g", label="Optimal trajectory") + + ax[0, 1].plot(q[0], q[1], "b", label="Optimal trajectory") + ax[0, 1].plot(u[0], u[1], "r", label="Optimal controls") + for i in range(n_shooting): + if i == 0: + ax[0, 1].plot( + (u[0][i], q[0][i * (polynomial_degree + 2)]), + (u[1][i], q[1][i * (polynomial_degree + 2)]), + ":k", + label="Spring orientation", + ) + else: + ax[0, 1].plot( + (u[0][i], q[0][i * (polynomial_degree + 2)]), (u[1][i], q[1][i * (polynomial_degree + 2)]), ":k" + ) + ax[0, 1].legend() + + ax[1, 0].step(tgrid[:-1], u.T, "-.", label=["Optimal controls X", "Optimal controls Y"]) + ax[1, 0].fill_between( + tgrid[:-1], + u.T[:, 0] - motor_noise_magnitude[0], + u.T[:, 0] + motor_noise_magnitude[0], + step="pre", + alpha=0.3, + color="#1f77b4", + ) + ax[1, 0].fill_between( + tgrid[:-1], + u.T[:, 1] - motor_noise_magnitude[1], + u.T[:, 1] + motor_noise_magnitude[1], + step="pre", + alpha=0.3, + color="#ff7f0e", + ) + + ax[1, 0].plot(tgrid, q[0, :: polynomial_degree + 2], "--", label="Optimal trajectory X") + ax[1, 0].plot(tgrid, q[1, :: polynomial_degree + 2], "-", label="Optimal trajectory Y") + + ax[1, 0].set_xlabel("Time [s]") + ax[1, 0].legend() + + if is_stochastic: + m = algebraic_states["m"] + cov = algebraic_states["cov"] + + # estimate covariance using series of noisy trials + iter = 200 + np.random.seed(42) + noise = np.vstack( + [ + np.random.normal(loc=0, scale=motor_noise_magnitude[0], size=(1, n_shooting, iter)), + np.random.normal(loc=0, scale=motor_noise_magnitude[1], size=(1, n_shooting, iter)), + ] + ) + + nx = bio_model.nb_q + bio_model.nb_qdot + cov_numeric = np.zeros((nx, nx, n_shooting)) + x_mean = np.zeros((nx, n_shooting + 1)) + x_std = np.zeros((nx, n_shooting + 1)) + dt = Tf / (n_shooting) + + x_j = np.zeros((nx,)) + for i in range(n_shooting): + x_i = np.hstack([q[:, i * (polynomial_degree + 2)], qdot[:, i * (polynomial_degree + 2)]]) + new_u = np.hstack([u[:, i:], u[:, :i]]) + next_x = np.zeros((nx, iter)) + for it in range(iter): + + x_j[:] = x_i[:] + for j in range(n_shooting): + dynamics = ( + lambda t, x: bio_model.dynamics_numerical( + states=x, controls=new_u[:, j].T, motor_noise=noise[:, j, it].T + ) + .full() + .T + ) + sol_ode = solve_ivp(dynamics, t_span=[0, dt], y0=x_j, method="RK45") + x_j[:] = sol_ode.y[:, -1] + + next_x[:, it] = x_j[:] + + x_mean[:, i] = np.mean(next_x, axis=1) + x_std[:, i] = np.std(next_x, axis=1) + + cov_numeric[:, :, i] = np.cov(next_x) + if i == 0: + ax[0, 0].plot(next_x[0, :], next_x[1, :], ".r", label="Noisy integration") + else: + ax[0, 0].plot(next_x[0, :], next_x[1, :], ".r") + # We can draw the X and Y covariance just for personnal reference, but the eigen vectors of the covariance matrix do not have to be aligned with the horizontal and vertical axis + # ax[0, 0].plot([x_mean[0, i], x_mean[0, i]], x_mean[1, i] + [-x_std[1, i], x_std[1, i]], "-k", label="Numerical covariance") + # ax[0, 0].plot(x_mean[0, i] + [-x_std[0, i], x_std[0, i]], [x_mean[1, i], x_mean[1, i]], "-k") + if i == 0: + draw_cov_ellipse( + cov_numeric[:2, :2, i], x_mean[:, i], ax[0, 0], color="r", label="Numerical covariance" + ) + draw_cov_ellipse( + cov_numeric[:2, :2, i], x_mean[:, i], ax_comparison, color="r", label="Numerical covariance" + ) + else: + draw_cov_ellipse(cov_numeric[:2, :2, i], x_mean[:, i], ax[0, 0], color="r") + draw_cov_ellipse(cov_numeric[:2, :2, i], x_mean[:, i], ax_comparison, color="r") + + ax[1, 0].fill_between( + tgrid, + q[0, :: polynomial_degree + 2] - x_std[0, :], + q[0, :: polynomial_degree + 2] + x_std[0, :], + alpha=0.3, + color="#2ca02c", + ) + + ax[1, 0].fill_between( + tgrid, + q[1, :: polynomial_degree + 2] - x_std[1, :], + q[1, :: polynomial_degree + 2] + x_std[1, :], + alpha=0.3, + color="#d62728", + ) + + ax[0, 0].plot(x_mean[0, :], x_mean[1, :], "+b", label="Numerical mean") + + for i in range(n_shooting + 1): + cov_i = cov[:, i] + if not test_matrix_semi_definite_positiveness(cov_i): + print(f"Something went wrong at the {i}th node. (Semi-definiteness)") + + if not test_eigen_values(cov_i): + print(f"Something went wrong at the {i}th node. (Eigen values)") + + cov_i = reshape_to_matrix(cov_i, (bio_model.matrix_shape_cov)) + if i == 0: + draw_cov_ellipse( + cov_i[:2, :2], q[:, i * (polynomial_degree + 2)], ax[0, 0], color="y", label="Optimal covariance" + ) + draw_cov_ellipse( + cov_i[:2, :2], + q[:, i * (polynomial_degree + 2)], + ax_comparison, + color="y", + label="Optimal covariance", + ) + else: + draw_cov_ellipse(cov_i[:2, :2], q[:, i * (polynomial_degree + 2)], ax[0, 0], color="y") + draw_cov_ellipse(cov_i[:2, :2], q[:, i * (polynomial_degree + 2)], ax_comparison, color="y") + ax[0, 0].legend() + + # Integrate the nominal dynamics (as if it was deterministic) + integrated_sol = sol_socp.integrate( + shooting_type=Shooting.SINGLE, integrator=SolutionIntegrator.SCIPY_RK45, to_merge=SolutionMerge.NODES + ) + # Integrate the stochastic dynamics (considering the feedback and the motor and sensory noises) + noisy_integrated_sol = sol_socp.noisy_integrate( + integrator=SolutionIntegrator.SCIPY_RK45, to_merge=SolutionMerge.NODES + ) + + # compare with noisy integration + cov_integrated = np.zeros((2, 2, n_shooting + 1)) + mean_integrated = np.zeros((2, n_shooting + 1)) + i_node = 0 + for i in range(noisy_integrated_sol["q"][0].shape[0]): + if i == 0: + ax_comparison.plot( + noisy_integrated_sol["q"][0][i, :], + noisy_integrated_sol["q"][1][i, :], + ".", + color=cm.viridis(i / noisy_integrated_sol["q"][0].shape[0]), + alpha=0.1, + label="Noisy integration", + ) + else: + ax_comparison.plot( + noisy_integrated_sol["q"][0][i, :], + noisy_integrated_sol["q"][1][i, :], + ".", + color=cm.viridis(i / noisy_integrated_sol["q"][0].shape[0]), + alpha=0.1, + ) + if i % 7 == 0: + cov_integrated[:, :, i_node] = np.cov( + np.vstack((noisy_integrated_sol["q"][0][i, :], noisy_integrated_sol["q"][1][i, :])) + ) + mean_integrated[:, i_node] = np.mean( + np.vstack((noisy_integrated_sol["q"][0][i, :], noisy_integrated_sol["q"][1][i, :])), axis=1 + ) + + if i == 0: + draw_cov_ellipse( + cov_integrated[:2, :2, i_node], + mean_integrated[:, i_node], + ax_comparison, + color="b", + label="Noisy integration covariance", + ) + else: + draw_cov_ellipse( + cov_integrated[:2, :2, i_node], + mean_integrated[:, i_node], + ax_comparison, + color="b", + ) + i_node += 1 + ax_comparison.legend() + fig_comparison.tight_layout() + fig_comparison.savefig("comparison.png") + plt.show() + def superellipse(a=1, b=1, n=2, x_0=0, y_0=0, resolution=100): x = np.linspace(-2 * a + x_0, 2 * a + x_0, resolution) @@ -52,7 +306,7 @@ def superellipse(a=1, b=1, n=2, x_0=0, y_0=0, resolution=100): return X, Y, Z -def draw_cov_ellipse(cov, pos, ax, color="b"): +def draw_cov_ellipse(cov, pos, ax, **kwargs): """ Draw an ellipse representing the covariance at a given point. """ @@ -67,7 +321,7 @@ def eigsorted(cov): # Width and height are "full" widths, not radius width, height = 2 * np.sqrt(vals) - ellip = plt.matplotlib.patches.Ellipse(xy=pos, width=width, height=height, angle=theta, color=color, alpha=0.1) + ellip = plt.matplotlib.patches.Ellipse(xy=pos, width=width, height=height, angle=theta, alpha=0.5, **kwargs) ax.add_patch(ellip) return ellip @@ -339,7 +593,7 @@ def main(): n_shooting = 40 final_time = 4 - motor_noise_magnitude = np.array([1, 1]) + motor_noise_magnitude = np.array([1, 1]) * 1 bio_model = MassPointModel(socp_type=socp_type, motor_noise_magnitude=motor_noise_magnitude) q_init = np.zeros((bio_model.nb_q, (polynomial_degree + 2) * n_shooting + 1)) @@ -364,7 +618,7 @@ def main(): # Solver parameters solver = Solver.IPOPT(show_online_optim=False, show_options=dict(show_bounds=True)) - # solver.set_linear_solver("ma57") + solver.set_linear_solver("ma57") sol_socp = socp.solve(solver) time = sol_socp.decision_time(to_merge=SolutionMerge.NODES) @@ -372,53 +626,28 @@ def main(): controls = sol_socp.decision_controls(to_merge=SolutionMerge.NODES) algebraic_states = sol_socp.decision_algebraic_states(to_merge=SolutionMerge.NODES) - q = states["q"] - u = controls["u"] - Tf = time[-1] - tgrid = np.linspace(0, Tf, n_shooting + 1).squeeze() - - fig, ax = plt.subplots(2, 2) - for i in range(2): - a = bio_model.super_ellipse_a[i] - b = bio_model.super_ellipse_b[i] - n = bio_model.super_ellipse_n[i] - x_0 = bio_model.super_ellipse_center_x[i] - y_0 = bio_model.super_ellipse_center_y[i] - - X, Y, Z = superellipse(a, b, n, x_0, y_0) - - ax[0, 0].contourf(X, Y, Z, levels=[-1000, 0], colors=["#DA1984"], alpha=0.5) - - ax[0, 0].plot(q_init[0], q_init[1], "-k", label="Initial guess") - ax[0, 0].plot(q[0][0], q[1][0], "og") - ax[0, 0].plot(q[0], q[1], "-g", label="q") - - ax[0, 1].plot(q[0], q[1], "b") - ax[0, 1].plot(u[0], u[1], "r") - for i in range(n_shooting): - ax[0, 1].plot((u[0][i], q[0][i * (polynomial_degree + 2)]), (u[1][i], q[1][i * (polynomial_degree + 2)]), ":k") - - ax[1, 0].plot(tgrid, q[0, :: polynomial_degree + 2], "--", label="px") - ax[1, 0].plot(tgrid, q[1, :: polynomial_degree + 2], "-", label="py") - ax[1, 0].step(tgrid[:-1], u.T, "-.", label="u") - ax[1, 0].set_xlabel("t") - - if is_stochastic: - m = algebraic_states["m"] - cov = algebraic_states["cov"] - - for i in range(n_shooting + 1): - cov_i = cov[:, i] - if not test_matrix_semi_definite_positiveness(cov_i): - print(f"Something went wrong at the {i}th node. (Semi-definiteness)") - - if not test_eigen_values(cov_i): - print(f"Something went wrong at the {i}th node. (Eigen values)") - - cov_i = reshape_to_matrix(cov_i, (bio_model.matrix_shape_cov)) - draw_cov_ellipse(cov_i[:2, :2], q[:, i * (polynomial_degree + 2)], ax[0, 0], color="b") - - plt.show() + data_to_save = { + "time": time, + "states": states, + "controls": controls, + "algebraic_states": algebraic_states, + } + with open("obstacle.pkl", "wb") as file: + pickle.dump(data_to_save, file) + + plot_results( + sol_socp, + states, + controls, + time, + algebraic_states, + bio_model, + motor_noise_magnitude, + n_shooting, + polynomial_degree, + is_stochastic, + q_init, + ) if __name__ == "__main__": diff --git a/bioptim/examples/stochastic_optimal_control/rockit_matrix_lyapunov.py b/bioptim/examples/stochastic_optimal_control/rockit_matrix_lyapunov.py index 3cc882760..eaf8bdc19 100644 --- a/bioptim/examples/stochastic_optimal_control/rockit_matrix_lyapunov.py +++ b/bioptim/examples/stochastic_optimal_control/rockit_matrix_lyapunov.py @@ -28,13 +28,15 @@ OdeSolver, StochasticBioModel, PhaseDynamics, + SolutionMerge, ) -from bioptim.examples.stochastic_optimal_control.rockit_model import RockitModel +from bioptim.examples.stochastic_optimal_control.models.rockit_model import RockitModel from bioptim.examples.stochastic_optimal_control.common import ( test_matrix_semi_definite_positiveness, test_eigen_values, reshape_to_matrix, ) +from scipy.integrate import solve_ivp def configure_optimal_control_problem(ocp: OptimalControlProgram, nlp: NonLinearProgram): @@ -87,7 +89,7 @@ def bound(t): def path_constraint(controller, is_robustified: bool = False): - t = controller.time.cx + t = controller.t_span[0] q = controller.states["q"].cx sup = bound(t) if is_robustified: @@ -239,9 +241,11 @@ def main(): socp_type = SocpType.COLLOCATION(polynomial_degree=polynomial_degree, method="legendre") n_shooting = 40 final_time = 1 + motor_noise_magnitude = np.array([1]) + bio_model = RockitModel(socp_type=socp_type, motor_noise_magnitude=motor_noise_magnitude) + dt = final_time / n_shooting ts = np.arange(n_shooting + 1) * dt - motor_noise_magnitude = np.array([0]) # Solver parameters solver = Solver.IPOPT(show_online_optim=False, show_options=dict(show_bounds=True)) @@ -258,18 +262,54 @@ def main(): ) sol_socp = socp.solve(solver) - T = sol_socp.time - q = sol_socp.states["q"] - u = sol_socp.controls["u"] + states = sol_socp.decision_states() + controls = sol_socp.decision_controls() + q = np.array([item.flatten()[0] for item in states["q"]]) + qdot = np.array([item.flatten()[0] for item in states["qdot"]]) + u = np.vstack([np.array([item.flatten() for item in controls["u"]]), np.array([[np.nan]])]) + time = np.array([item.full().flatten()[0] for item in sol_socp.stepwise_time()]) # sol_ocp.graphs() plt.figure() - plt.plot(T, bound(T), "k", label="bound") - plt.plot(np.squeeze(T), np.squeeze(q), label="q") - plt.step(ts, np.squeeze(u / 40), label="u/40") + plt.plot(time, bound(time), "k", label="bound") + plt.plot(time, q, label="q") + plt.step(time, u / 40, label="u/40") if is_stochastic: - cov = sol_socp.decision_algebraic_states["cov"] + cov = sol_socp.decision_algebraic_states(to_merge=SolutionMerge.NODES)["cov"] + + # estimate covariance using series of noisy trials + iter = 200 + np.random.seed(42) + noise = np.random.normal(loc=0, scale=motor_noise_magnitude, size=(1, n_shooting, iter)) + + nx = bio_model.nb_q + bio_model.nb_qdot + cov_numeric = np.empty((nx, nx, iter)) + x_mean = np.empty((nx, iter)) + x_std = np.empty((nx, iter)) + + for i in range(n_shooting): + x_i = np.hstack([q[:, i], qdot[:, i]]) # .T + t_span = time[i : i + 2] + + next_x = np.empty((nx, iter)) + for it in range(iter): + dynamics = ( + lambda t, x: bio_model.dynamics_numerical( + states=x, controls=u[:, i].T, motor_noise=noise[:, i, it].T + ) + .full() + .T + ) + sol_ode = solve_ivp(dynamics, t_span, x_i, method="RK45") + next_x[:, it] = sol_ode.y[:, -1] + + x_mean[:, i] = np.mean(next_x, axis=1) + x_std[:, i] = np.std(next_x, axis=1) + cov_numeric[:, :, i] = np.cov(next_x) + + plt.plot(np.tile(time[i + 1], 2), x_mean[0, i] + [-x_std[0, i], x_std[0, i]], "-k") + plt.plot(np.tile(time[i + 1], iter), next_x[0, :], ".r") o = np.array([[1, 0]]) sigma = np.zeros((1, n_shooting + 1)) diff --git a/bioptim/gui/check_conditioning.py b/bioptim/gui/check_conditioning.py index 8a435db70..ff637bce1 100644 --- a/bioptim/gui/check_conditioning.py +++ b/bioptim/gui/check_conditioning.py @@ -1,556 +1,265 @@ import numpy as np -from casadi import MX, SX, Function, horzcat, vertcat, jacobian, vcat, hessian +from casadi import Function, jacobian, hessian, sum1 from matplotlib import pyplot as plt import matplotlib.colors as mcolors import matplotlib.cm as mcm -from ..misc.enums import ( - ControlType, -) -from ..misc.enums import QuadratureRule -from ..dynamics.ode_solver import OdeSolver -from ..limits.penalty_helpers import PenaltyHelpers +from ..interfaces.ipopt_interface import IpoptInterface -def check_conditioning(ocp): +def jacobian_hessian_constraints(variables_vector, all_g): """ - Visualisation of jacobian and hessian contraints and hessian objective for each phase at initial time + Returns + ------- + A list with jacobian matrix of constraints evaluates at initial time for each phase + A list with the rank of each jacobian matrix + A list with the different type of constraints + A list with norms of hessian matrix of constraints at initial time for each phase """ - def get_u(nlp, u: MX | SX, dt: MX | SX): - """ - Get the control at a given time - - Parameters - ---------- - nlp: NonlinearProgram - The nonlinear program - u: MX | SX - The control matrix - dt: MX | SX - The time a which control should be computed - - Returns - ------- - The control at a given time - """ - - if nlp.control_type in (ControlType.CONSTANT,): - return u - elif nlp.control_type == ControlType.LINEAR_CONTINUOUS: - return u[:, 0] + (u[:, 1] - u[:, 0]) * dt - else: - raise RuntimeError(f"{nlp.control_type} ControlType not implemented yet") - - def jacobian_hessian_constraints(): - """ - Returns - ------- - A list with jacobian matrix of constraints evaluates at initial time for each phase - A list with the rank of each jacobian matrix - A list with the different type of constraints - A list with norms of hessian matrix of constraints at initial time for each phase - """ - - jacobian_list = [] - jacobian_rank = [] - tick_labels_list = [] - hessian_norm_list = [] - - # JACOBIAN - phases_dt = ocp.dt_parameter.cx - for nlp in ocp.nlp: - list_constraints = [] - - for constraints in nlp.g: - node_index = constraints.node_idx[0] # TODO deal with phase_dynamics=PHASE_DYNAMICS.ONE_PER_NOE - nlp.states.node_index = node_index - nlp.states_dot.node_index = node_index - nlp.controls.node_index = node_index - nlp.algebraic_states.node_index = node_index - - if constraints.multinode_penalty: - n_phases = ocp.n_phases - for phase_idx in constraints.nodes_phase: - controllers.append(constraints.get_penalty_controller(ocp, ocp.nlp[phase_idx % n_phases])) - else: - controllers = [constraints.get_penalty_controller(ocp, nlp)] - - for axis in range(constraints.function[node_index].size_out("val")[0]): - # depends if there are parameters - if nlp.parameters.shape == 0: - vertcat_obj = vertcat([], *nlp.X_scaled, *nlp.U_scaled, nlp.parameters.cx, *nlp.A_scaled) - else: - vertcat_obj = vertcat([], *nlp.X_scaled, *nlp.U_scaled, *[nlp.parameters.cx, *nlp.A_scaled]) - - for controller in controllers: - controller.node_index = constraints.node_idx[0] - - _, t0, x, u, p, a = constraints.get_variable_inputs(controllers) - - list_constraints.append( - jacobian( - constraints.function[constraints.node_idx[0]](t0, phases_dt, x, u, p, a)[axis], - vertcat_obj, - ) - ) - - jacobian_cas = vcat(list_constraints).T - - # depends if there are parameters - if nlp.parameters.shape == 0: - vertcat_obj = vertcat([], *nlp.X_scaled, *nlp.U_scaled, nlp.parameters.cx, *nlp.A) - else: - vertcat_obj = vertcat([], *nlp.X_scaled, *nlp.U_scaled, *[nlp.parameters.cx], *nlp.A) - - jac_func = Function( - "jacobian", - [vertcat_obj], - [jacobian_cas], + # JACOBIAN + constraints_jac_func = Function( + "constraints_jacobian", + [variables_vector], + [jacobian(all_g, variables_vector)], + ) + + # HESSIAN + constraints_hess_func = [] + for i in range(all_g.shape[0]): + constraints_hess_func.append( + Function( + "constraints_hessian", + [variables_vector], + [hessian(all_g[i], variables_vector)[0]], ) + ) - 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_a_init = sum([nlp.a_init[key].shape[0] for key in nlp.a_init.keys()]) - - # evaluate jac_func at X_init, U_init, considering the parameters - time_init = np.array([], dtype=np.float64) - 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()]) - a_init = np.zeros((len(nlp.A), nb_a_init)) - - 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): - nlp.states.node_index = node_index - x_init[node_index, nlp.states[key].index] = np.array(nlp.x_init[key].init.evaluate_at(node_index)) - for key in nlp.controls.keys(): - if not key in nlp.controls: - continue - nlp.u_init[key].check_and_adjust_dimensions(len(nlp.controls[key]), nlp.ns) - for node_index in range(nlp.ns): - nlp.controls.node_index = node_index - u_init[node_index, nlp.controls[key].index] = np.array(nlp.u_init[key].init.evaluate_at(node_index)) - for key in nlp.algebraic_states.keys(): - nlp.a_init[key].check_and_adjust_dimensions(len(nlp.algebraic_states[key]), nlp.ns) - for node_index in range(nlp.ns): - nlp.algebraic_states.node_index = node_index - a_init[node_index, nlp.algebraic_states[key].index] = np.array( - nlp.a_init[key].init.evaluate_at(node_index) - ) - - time_init = time_init.reshape((time_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)) - a_init = a_init.reshape((a_init.size, 1)) - - vector_init = np.vstack((time_init, x_init, u_init, param_init, a_init)) - jacobian_matrix = np.array(jac_func(vector_init)) - jacobian_list.append(jacobian_matrix) - - # calculate jacobian rank - if jacobian_matrix.size > 0: - rank = np.linalg.matrix_rank(jacobian_matrix) - jacobian_rank.append(rank) - else: - jacobian_rank.append("No constraints") - - # HESSIAN - tick_labels = [] - list_hessian = [] - list_norm = [] - for constraints in nlp.g: - node_index = constraints.node_idx[0] # TODO deal with phase_dynamics=PhaseDynamics.ONE_PER_NODE - nlp.states.node_index = node_index - nlp.states_dot.node_index = node_index - nlp.controls.node_index = node_index - nlp.algebraic_states.node_index = node_index - - if constraints.multinode_penalty: - n_phases = ocp.n_phases - for phase_idx in constraints.nodes_phase: - controllers.append(constraints.get_penalty_controller(ocp, ocp.nlp[phase_idx % n_phases])) - else: - controllers = [constraints.get_penalty_controller(ocp, nlp)] - - for axis in range(constraints.function[node_index].size_out("val")[0]): - # find all equality constraints - if constraints.bounds.min[axis][0] == constraints.bounds.max[axis][0]: - vertcat_obj = vertcat([], *nlp.X_scaled, *nlp.U_scaled) # time, states, controls - if nlp.parameters.shape == 0: - vertcat_obj = vertcat(vertcat_obj, nlp.parameters.cx) - else: - vertcat_obj = vertcat(vertcat_obj, *[nlp.parameters.cx]) - vertcat_obj = vertcat(vertcat_obj, *nlp.A_scaled) - - for controller in controllers: - controller.node_index = constraints.node_idx[0] - _, t0, x, u, p, a = constraints.get_variable_inputs(controllers) - - hessian_cas = hessian( - constraints.function[node_index](t0, phases_dt, x, u, p, a)[axis], vertcat_obj - )[0] - - tick_labels.append(constraints.name) - - hes_func = Function( - "hessian", - [vertcat_obj], - [hessian_cas], - ) - - vector_init = np.vstack((time_init, x_init, u_init, param_init, a_init)) - hessian_matrix = np.array(hes_func(vector_init)) - list_hessian.append(hessian_matrix) - - tick_labels_list.append(tick_labels) - - # calculate norm - for hessian_idx in list_hessian: - norm = np.linalg.norm(hessian_idx) - list_norm.append(norm) - array_norm = np.array(list_norm).reshape(len(list_hessian), 1) - hessian_norm_list.append(array_norm) - - return jacobian_list, jacobian_rank, tick_labels_list, hessian_norm_list - - def check_constraints_plot(): - """ - Visualisation of jacobian matrix and hessian norm matrix - """ - jacobian_list, jacobian_rank, tick_labels_list, hessian_norm_list = jacobian_hessian_constraints() - - max_norm = [] - min_norm = [] - if hessian_norm_list[0].size != 0: - for hessian_norm in hessian_norm_list: - max_norm.append(np.ndarray.max(hessian_norm)) - min_norm.append(np.ndarray.min(hessian_norm)) - min_norm = min(min_norm) - max_norm = max(max_norm) - else: - max_norm = 0 - min_norm = 0 - - max_jac = [] - min_jac = [] - if jacobian_list[0].size != 0: - for jacobian_obj in jacobian_list: - max_jac.append(np.ndarray.max(jacobian_obj)) - min_jac.append(np.ndarray.min(jacobian_obj)) - max_jac = max(max_jac) - min_jac = min(min_jac) - else: - max_jac = 0 - min_jac = 0 - - # PLOT GENERAL - fig, axis = plt.subplots(1, 2 * ocp.n_phases) - for ax in range(ocp.n_phases): - # Jacobian plot - jacobian_list[ax][(jacobian_list[ax] == 0)] = np.nan - current_cmap = mcm.get_cmap("seismic") - # todo: The get_cmap function was deprecated. Use ``matplotlib.colormaps[name]`` or ``matplotlib.colormaps.get_cmap(obj)`` instead - current_cmap.set_bad(color="k") - norm = mcolors.TwoSlopeNorm(vmin=min_jac - 0.01, vmax=max_jac + 0.01, vcenter=0) - im = axis[ax].imshow(jacobian_list[ax], aspect="auto", cmap=current_cmap, norm=norm) - axis[ax].set_title("Jacobian constraints \n Phase " + str(ax), fontweight="bold", fontsize=8) - axis[ax].text( - 0, - jacobian_list[ax].shape[0] * 1.08, - "Matrix rank = " - + str(jacobian_rank[ax]) - + "\n Number of constraints = " - + str(jacobian_list[ax].shape[1]), - horizontalalignment="center", - fontweight="bold", - fontsize=8, - ) - axis[ax].text( - 0, - jacobian_list[ax].shape[0] * 1.1, - "The rank should be equal to the number of constraints", - horizontalalignment="center", - fontsize=6, - ) - cbar_ax = fig.add_axes([0.02, 0.4, 0.015, 0.3]) - fig.colorbar(im, cax=cbar_ax) - - # Hessian constraints plot - hessian_norm_list[ax][~(hessian_norm_list[ax] != 0).astype(bool)] = np.nan - current_cmap2 = mcm.get_cmap("seismic") - current_cmap2.set_bad(color="k") - norm2 = mcolors.TwoSlopeNorm(vmin=min_norm - 0.01, vmax=max_norm + 0.01, vcenter=0) - yticks = [] - for i in range(len(hessian_norm_list[ax])): - yticks.append(i) - - im2 = axis[ax + ocp.n_phases].imshow(hessian_norm_list[ax], aspect="auto", cmap=current_cmap2, norm=norm2) - axis[ax + ocp.n_phases].set_title( - "Hessian constraint norms \n Phase " + str(ax), fontweight="bold", fontsize=8 - ) - axis[ax + ocp.n_phases].set_xticks([0]) - axis[ax + ocp.n_phases].set_xticklabels(["Norms should be close to 0"], fontsize=8) - axis[ax + ocp.n_phases].set_yticks(yticks) - axis[ax + ocp.n_phases].set_yticklabels(tick_labels_list[ax], fontsize=6, rotation=90) - cbar_ax2 = fig.add_axes([0.95, 0.4, 0.015, 0.3]) - fig.colorbar(im2, cax=cbar_ax2) - fig.legend(["Black = 0"], loc="upper left") - plt.suptitle("Check conditioning for constraints ", color="b", fontsize=15, fontweight="bold") - - def hessian_objective(): - """ - - Returns - ------- - A list with the hessian of objectives evaluate at initial time for each phase - A list with the condition numbers of each phases - A list that indicates if the objective is convexe or not - """ - - hessian_obj_list = [] - phases_dt = ocp.dt_parameter.cx - for phase, nlp in enumerate(ocp.nlp): - for obj in nlp.J: - objective = 0 - - node_index = obj.node_idx[0] # TODO deal with phase_dynamics=PhaseDynamics.ONE_PER_NODE - nlp.states.node_index = node_index - nlp.states_dot.node_index = node_index - nlp.controls.node_index = node_index - nlp.algebraic_states.node_index = node_index - - if obj.multinode_penalty: - n_phases = ocp.n_phases - for phase_idx in obj.nodes_phase: - controllers.append(obj.get_penalty_controller(ocp, ocp.nlp[phase_idx % n_phases])) - else: - controllers = [obj.get_penalty_controller(ocp, nlp)] - - # Test every possibility - if obj.multinode_penalty: - 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 - algebraic_states_pre = phase.algebraic_states.cx_end - algebraic_states_post = nlp_post.algebraic_states.cx_start - state_cx = vertcat(states_pre, states_post) - control_cx = vertcat(controls_pre, controls_post) - algebraic_states_cx = vertcat(algebraic_states_pre, algebraic_states_post) - - else: - if obj.integrate: - state_cx = horzcat(*([nlp.states.cx_start] + nlp.states.cx_intermediates_list)) - else: - state_cx = nlp.states.cx_start - control_cx = nlp.controls.cx_start - algebraic_states_cx = nlp.algebraic_states.cx_start - if obj.explicit_derivative: - if obj.derivative: - raise RuntimeError("derivative and explicit_derivative cannot be simultaneously true") - state_cx = horzcat(state_cx, nlp.states.cx_end) - control_cx = horzcat(control_cx, nlp.controls.cx_end) - algebraic_states_cx = horzcat(algebraic_states_cx, nlp.algebraic_states.cx_end) - - if obj.derivative: - state_cx = horzcat(nlp.states.cx_end, nlp.states.cx_start) - control_cx = horzcat(nlp.controls.cx_end, nlp.controls.cx_start) - algebraic_states_cx = horzcat(nlp.algebraic_states.cx_end, nlp.algebraic_states.cx_start) - - dt_cx = nlp.cx.sym("dt", 1, 1) - is_trapezoidal = ( - obj.integration_rule == QuadratureRule.APPROXIMATE_TRAPEZOIDAL - or obj.integration_rule == QuadratureRule.TRAPEZOIDAL - ) - - if is_trapezoidal: - state_cx = ( - horzcat(nlp.states.cx_start, nlp.states.cx_end) - if obj.integration_rule == QuadratureRule.APPROXIMATE_TRAPEZOIDAL - else nlp.states.cx_start - ) - control_cx = ( - horzcat(nlp.controls.cx_start) - if nlp.control_type == ControlType.CONSTANT - else horzcat(nlp.controls.cx_start, nlp.controls.cx_end) - ) - algebraic_states_cx = ( - horzcat(nlp.algebraic_states.cx_start, nlp.algebraic_states.cx_end) - if obj.integration_rule == QuadratureRule.APPROXIMATE_TRAPEZOIDAL - else nlp.algebraic_states.cx_start - ) - - for controller in controllers: - controller.node_index = obj.node_idx[0] - _, t0, x, u, p, a = obj.get_variable_inputs(controllers) - - target = PenaltyHelpers.target(obj, obj.node_idx.index(node_index)) - - penalty = obj.weighted_function[node_index](t0, phases_dt, x, u, p, a, obj.weight, target) - - for i in range(penalty.shape[0]): - objective += penalty[i] ** 2 - - # create function to build the hessian - vertcat_obj = vertcat([], *nlp.X_scaled, *nlp.U_scaled) # time, states, controls - if nlp.parameters.shape == 0: - vertcat_obj = vertcat(vertcat_obj, nlp.parameters.cx) - else: - vertcat_obj = vertcat(vertcat_obj, *[nlp.parameters.cx]) - if vertcat(*nlp.A_scaled).shape[0] > 0: - vertcat_obj = vertcat(vertcat_obj, *nlp.A_scaled) - - hessian_cas = hessian(objective, vertcat_obj)[0] - - hes_func = Function( - "hessian", - [vertcat_obj], - [hessian_cas], - ) + return constraints_jac_func, constraints_hess_func - 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_a_init = sum([nlp.a_init[key].shape[0] for key in nlp.a_init.keys()]) - - # evaluate jac_func at X_init, U_init, considering the parameters - time_init = np.array([], dtype=np.float64) - 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()]) - a_init = np.zeros((len(nlp.A), nb_a_init)) - - 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): - nlp.states.node_index = node_index - x_init[node_index, nlp.states[key].index] = np.array(nlp.x_init[key].init.evaluate_at(node_index)) - for key in nlp.controls.keys(): - nlp.u_init[key].check_and_adjust_dimensions(len(nlp.controls[key]), nlp.ns) - for node_index in range(nlp.ns): - nlp.controls.node_index = node_index - u_init[node_index, nlp.controls[key].index] = np.array(nlp.u_init[key].init.evaluate_at(node_index)) - for key in nlp.algebraic_states.keys(): - nlp.a_init[key].check_and_adjust_dimensions(len(nlp.algebraic_states[key]), nlp.ns) - for node_index in range(nlp.ns): - nlp.algebraic_states.node_index = node_index - a_init[node_index, nlp.algebraic_states[key].index] = np.array( - nlp.a_init[key].init.evaluate_at(node_index) - ) - - time_init = time_init.reshape((time_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)) - a_init = a_init.reshape((a_init.size, 1)) - vector_init = np.vstack((time_init, x_init, u_init, param_init, a_init)) - - hessian_obj_matrix = np.array(hes_func(vector_init)) - hessian_obj_list.append(hessian_obj_matrix) - - # Convexity checking (positive semi-definite hessian) - # On R (convex), the objective is convex if and only if the hessian is positive semi definite (psd) - # And, as the hessian is symmetric (Schwarz), the hessian is psd if and only if the eigenvalues are positive - convexity = [] - condition_number = [] - for matrix in range(len(hessian_obj_list)): - eigen_values = np.linalg.eigvals(hessian_obj_list[matrix]) - ev_max = min(eigen_values) - ev_min = max(eigen_values) - if ev_min == 0: - condition_number.append(" /!\ Ev_min is 0") - if ev_min != 0: - condition_number.append(np.abs(ev_max) / np.abs(ev_min)) - convexity.append("Possible") - for ev in range(eigen_values.size): - if eigen_values[ev] < 0: - convexity[matrix] = "False" - break - - return hessian_obj_list, condition_number, convexity - - def check_objective_plot(): - """ - Visualisation of hessian objective matrix - """ - - hessian_obj_list, condition_number, convexity = hessian_objective() - - max_hes = [] - min_hes = [] - for hessian_matrix_obj in hessian_obj_list: - max_hes.append(np.ndarray.max(hessian_matrix_obj)) - min_hes.append(np.ndarray.min(hessian_matrix_obj)) - min_hes = min(min_hes) - max_hes = max(max_hes) - - # PLOT GENERAL - fig_obj, axis_obj = plt.subplots(1, ocp.n_phases) - for ax in range(ocp.n_phases): - hessian_obj_list[ax][~(hessian_obj_list[ax] != 0).astype(bool)] = np.nan - current_cmap3 = mcm.get_cmap("seismic") - current_cmap3.set_bad(color="k") - norm = mcolors.TwoSlopeNorm(vmin=min_hes - 0.01, vmax=max_hes + 0.01, vcenter=0) - if ocp.n_phases == 1: - im3 = axis_obj.imshow(hessian_obj_list[ax], cmap=current_cmap3, norm=norm) - axis_obj.set_title("Hessian objective \n Phase " + str(ax), fontweight="bold", fontsize=8) - axis_obj.text( - hessian_obj_list[ax].shape[0] / 2, - hessian_obj_list[ax].shape[0] * 1.1, - "Convexity = " + convexity[ax], - horizontalalignment="center", - fontweight="bold", - fontsize=8, - ) - axis_obj.text( - hessian_obj_list[ax].shape[0] / 2, - hessian_obj_list[ax].shape[0] * 1.2, - "|λmax|/|λmin| = Condition number = " + condition_number[ax], - horizontalalignment="center", - fontweight="bold", - fontsize=8, - ) - else: - im3 = axis_obj[ax].imshow(hessian_obj_list[ax], cmap=current_cmap3, norm=norm) - axis_obj[ax].set_title("Hessian objective \n Phase " + str(ax), fontweight="bold", fontsize=8) - axis_obj[ax].text( - hessian_obj_list[ax].shape[0] / 2, - hessian_obj_list[ax].shape[0] * 1.1, - "Convexity = " + convexity[ax], - horizontalalignment="center", - fontweight="bold", - fontsize=8, - ) - axis_obj[ax].text( - hessian_obj_list[ax].shape[0] / 2, - hessian_obj_list[ax].shape[0] * 1.2, - "|λmax|/|λmin| = Condition number = " + condition_number[ax], - horizontalalignment="center", - fontweight="bold", - fontsize=8, - ) - cbar_ax3 = fig_obj.add_axes([0.02, 0.4, 0.015, 0.3]) - fig_obj.colorbar(im3, cax=cbar_ax3) - fig_obj.legend(["Black = 0"], loc="upper left") - fig_obj.text( - 0.5, - 0.1, - "Every hessian should be convexe \n Condition numbers should be close to 0", - horizontalalignment="center", - fontsize=12, - fontweight="bold", - ) - plt.suptitle("Check conditioning for objectives", color="b", fontsize=15, fontweight="bold") - if sum(isinstance(ocp.nlp[i].ode_solver, OdeSolver.COLLOCATION) for i in range(ocp.n_phases)) > 0: - raise NotImplementedError("Conditioning check is not implemented for collocations") +def evaluate_jacobian_hessian_constraints(v, ocp): + + # JACOBIAN + constraints_jac_func = ocp.conditioning_plots["constraints_jac_func"] + evaluated_constraints_jacobian = constraints_jac_func(v) + jacobian_matrix = np.array(evaluated_constraints_jacobian) + + # Jacobian rank + if jacobian_matrix.size > 0: + jacobian_rank = np.linalg.matrix_rank(jacobian_matrix) + else: + jacobian_rank = "No constraints" + + # HESSIAN + constraints_hess_func = ocp.conditioning_plots["constraints_hess_func"] + hess_min_mean_max = np.zeros((len(constraints_hess_func), 3)) + for i in range(len(constraints_hess_func)): + evaluated_constraints_hessian = constraints_hess_func[i](v) + hess_min_mean_max[i, 0] = np.min(evaluated_constraints_hessian) + hess_min_mean_max[i, 1] = np.mean(evaluated_constraints_hessian) + hess_min_mean_max[i, 2] = np.max(evaluated_constraints_hessian) + + return jacobian_matrix, jacobian_rank, hess_min_mean_max + + +def hessian_objective(variables_vector, all_objectives): + """ + + Returns + ------- + A list with the hessian of objectives evaluate at initial time for each phase + A list with the condition numbers of each phases + A list that indicates if the objective is convexe or not + """ + + objectives_hess_func = Function( + "hessian", + [variables_vector], + [hessian(sum1(all_objectives), variables_vector)[0]], + ) + + return objectives_hess_func + + +def evaluate_hessian_objective(v, ocp): + """ + + Returns + ------- + A list with the hessian of objectives evaluate at initial time for each phase + A list with the condition numbers of each phases + A list that indicates if the objective is convexe or not + """ + + objectives_hess_func = ocp.conditioning_plots["objectives_hess_func"] + evaluated_objectives_hessian = objectives_hess_func(v) + hessian_matrix = np.array(evaluated_objectives_hessian) + + # Convexity checking (positive semi-definite hessian) + # On R (convex), the objective is convex if and only if the hessian is positive semi definite (psd) + # And, as the hessian is symmetric (Schwarz), the hessian is psd if and only if the eigenvalues are positive + eigen_values = np.linalg.eigvals(hessian_matrix) + ev_max = np.max(eigen_values) + ev_min = np.min(eigen_values) + if ev_min == 0: + condition_number = " /!\ min eigen value is 0" + if ev_min != 0: + condition_number = np.abs(ev_max) / np.abs(ev_min) + convexity = ( + "positive semi-definite" + if np.all(eigen_values > 0) + else f"not positive semi-definite (min: {ev_min}, max: {ev_max})" + ) + + return hessian_matrix, condition_number, convexity + + +def create_conditioning_plots(ocp): + + cmap = mcm.get_cmap("seismic") + cmap.set_bad(color="k") + interface = IpoptInterface(ocp) + variables_vector = ocp.variables_vector + all_g, _ = interface.dispatch_bounds(include_g=True, include_g_implicit=False, include_g_internal=False) + all_objectives = interface.dispatch_obj_func() + nb_variables = variables_vector.shape[0] + nb_constraints = all_g.shape[0] + + constraints_jac_func, constraints_hess_func = jacobian_hessian_constraints(variables_vector, all_g) + objectives_hess_func = hessian_objective(variables_vector, all_objectives) + + # PLOT CONSTRAINTS + fig_constraints, axis_constraints = plt.subplots(1, 2, num="Check conditioning for constraints") + + # Jacobian plot + fake_jacobian = np.zeros((nb_constraints, nb_variables)) + im_constraints_jacobian = axis_constraints[0].imshow(fake_jacobian, aspect="auto", cmap=cmap) + axis_constraints[0].set_title( + "Jacobian constraints \nMatrix rank = NA \n Number of constraints = NA", fontweight="bold", fontsize=12 + ) + # colorbar + cbar_ax = fig_constraints.add_axes([0.02, 0.4, 0.015, 0.3]) + fig_constraints.colorbar(im_constraints_jacobian, cax=cbar_ax) + + # Hessian constraints plot + fake_hessian = np.zeros((nb_constraints, 3)) + im_constraints_hessian = axis_constraints[1].imshow(fake_hessian, aspect="auto", cmap=cmap) + axis_constraints[1].set_title( + "Hessian constraint norms (Norms should be close to 0)", fontweight="bold", fontsize=12 + ) + axis_constraints[1].set_xticks([0, 1, 2]) + axis_constraints[1].set_xticklabels(["Min", "Mean", "Max"]) + # colobar + cbar_ax2 = fig_constraints.add_axes([0.95, 0.4, 0.015, 0.3]) + fig_constraints.colorbar(im_constraints_hessian, cax=cbar_ax2) + + fig_constraints.legend(["Black = 0"], loc="upper left") + plt.suptitle("The rank should be equal to the number of constraints", color="b", fontsize=15, fontweight="bold") + try: + figManager = plt.get_current_fig_manager() + figManager.window.showMaximized() + except: + pass + + # PLOT OBJECTIVES + fig_obj, axis_obj = plt.subplots(1, 1, num="Check conditioning for objectives") + + # Hessian objective plot + fake_hessian_obj = np.zeros((nb_variables, nb_variables)) + im_objectives_hessian = axis_obj.imshow(fake_hessian_obj, cmap=cmap) + axis_obj.set_title("Convexity = NA \n|λmax|/|λmin| = Condition number = NA", fontweight="bold", fontsize=12) + axis_obj.set_xlabel("Hessian objectives") + # colobar + cbar_ax3 = fig_obj.add_axes([0.02, 0.4, 0.015, 0.3]) + fig_obj.colorbar(im_objectives_hessian, cax=cbar_ax3) + + fig_obj.legend(["Black = 0"], loc="upper left") + plt.suptitle( + "Every hessian should be convexe (positive) and Condition numbers should be close to 0", + color="b", + fontsize=15, + fontweight="bold", + ) + try: + figManager = plt.get_current_fig_manager() + figManager.window.showMaximized() + except: + pass + + ocp.conditioning_plots = { + "axis_constraints": axis_constraints, + "axis_obj": axis_obj, + "im_constraints_jacobian": im_constraints_jacobian, + "im_constraints_hessian": im_constraints_hessian, + "im_objectives_hessian": im_objectives_hessian, + "constraints_jac_func": constraints_jac_func, + "constraints_hess_func": constraints_hess_func, + "objectives_hess_func": objectives_hess_func, + } + + +def update_constraints_plot(v, ocp): + + jacobian_matrix, jacobian_rank, hess_min_mean_max = evaluate_jacobian_hessian_constraints(v, ocp) + axis_constraints = ocp.conditioning_plots["axis_constraints"] + im_constraints_jacobian = ocp.conditioning_plots["im_constraints_jacobian"] + im_constraints_hessian = ocp.conditioning_plots["im_constraints_hessian"] + + # Jacobian plot + jacobian_matrix[jacobian_matrix == 0] = np.nan + jac_min = np.min(jacobian_matrix) if jacobian_matrix.shape[0] != 0 else 0 + jac_max = np.max(jacobian_matrix) if jacobian_matrix.shape[0] != 0 else 0 + norm = mcolors.TwoSlopeNorm(vmin=jac_min - 0.01, vmax=jac_max + 0.01, vcenter=0) + im_constraints_jacobian.set_data(jacobian_matrix) + im_constraints_jacobian.set_norm(norm) + axis_constraints[0].set_title( + f"Jacobian constraints \nMatrix rank = {str(jacobian_rank)}\n Number of constraints = {str(jacobian_matrix.shape[0])}", + fontweight="bold", + fontsize=12, + ) + + # Hessian constraints plot + hess_min = np.min(hess_min_mean_max) if hess_min_mean_max.shape[0] != 0 else 0 + hess_max = np.max(hess_min_mean_max) if hess_min_mean_max.shape[0] != 0 else 0 + norm = mcolors.TwoSlopeNorm(vmin=hess_min - 0.01, vmax=hess_max + 0.01, vcenter=0) + im_constraints_hessian.set_data(hess_min_mean_max) + im_constraints_hessian.set_norm(norm) + + +def update_objective_plot(v, ocp): + + hessian_matrix, condition_number, convexity = evaluate_hessian_objective(v, ocp) + axis_obj = ocp.conditioning_plots["axis_obj"] + im_objectives_hessian = ocp.conditioning_plots["im_objectives_hessian"] + cmap = mcm.get_cmap("seismic") + + # Hessian objective plot + hess_min = np.min(hessian_matrix) if hessian_matrix.shape[0] != 0 else 0 + hess_max = np.max(hessian_matrix) if hessian_matrix.shape[0] != 0 else 0 + norm = mcolors.TwoSlopeNorm(vmin=hess_min - 0.01, vmax=hess_max + 0.01, vcenter=0) + im_objectives_hessian.set_data(hessian_matrix) + im_objectives_hessian.set_norm(norm) + axis_obj.set_title( + f"Hessian objective \nConvexity = {convexity} \n|λmax|/|λmin| = Condition number = {condition_number}", + fontweight="bold", + fontsize=12, + ) + + +def update_conditioning_plots(v, ocp): + update_constraints_plot(v, ocp) + update_objective_plot(v, ocp) + plt.draw() + + +def check_conditioning(ocp): + """ + Visualisation of jacobian and hessian contraints and hessian objective for each phase at initial time + """ - check_constraints_plot() - check_objective_plot() + create_conditioning_plots(ocp) + v_init = ocp.init_vector + update_constraints_plot(v_init, ocp) + update_objective_plot(v_init, ocp) plt.show() diff --git a/bioptim/gui/ipopt_output_plot.py b/bioptim/gui/ipopt_output_plot.py new file mode 100644 index 000000000..83a3eed1b --- /dev/null +++ b/bioptim/gui/ipopt_output_plot.py @@ -0,0 +1,133 @@ +import numpy as np +from matplotlib import pyplot as plt +from matplotlib.cm import get_cmap +from casadi import jacobian, gradient, sum1, Function + + +def create_ipopt_output_plot(ocp, interface): + """ + This function creates the plots for the ipopt output: f, g, inf_pr, inf_du. + """ + ipopt_fig, axs = plt.subplots(3, 1, num="IPOPT output") + axs[0].set_ylabel("f", fontweight="bold") + axs[1].set_ylabel("inf_pr", fontweight="bold") + axs[2].set_ylabel("inf_du", fontweight="bold") + + plots = [] + colors = get_cmap("viridis") + for i in range(3): + plot = axs[i].plot([0], [1], linestyle="-", marker=".", color="k") + plots.append(plot[0]) + axs[i].grid(True) + axs[i].set_yscale("log") + + plot = axs[2].plot([0], [1], linestyle="-", marker=".", color=colors(0.1), label="grad_f") + plots.append(plot[0]) + plot = axs[2].plot([0], [1], linestyle="-", marker=".", color=colors(0.5), label="grad_g") + plots.append(plot[0]) + plot = axs[2].plot([0], [1], linestyle="-", marker=".", color=colors(0.9), label="lam_x") + plots.append(plot[0]) + axs[2].legend() + + ipopt_fig.tight_layout() + + try: + figManager = plt.get_current_fig_manager() + figManager.window.showMaximized() + except: + pass + + v_sym = interface.ocp.variables_vector + + all_objectives = interface.dispatch_obj_func() + all_g, all_g_bounds = interface.dispatch_bounds(include_g=True, include_g_internal=True, include_g_implicit=True) + + grad_f_func = Function("grad_f", [v_sym], [gradient(sum1(all_objectives), v_sym)]) + grad_g_func = Function("grad_g", [v_sym], [jacobian(all_g, v_sym).T]) + + ocp.ipopt_plots = { + "grad_f_func": grad_f_func, + "grad_g_func": grad_g_func, + "f": [], + "inf_pr": [], + "inf_du": [], + "grad_f": [], + "grad_g": [], + "lam_x": [], + "axs": axs, + "plots": plots, + "ipopt_fig": ipopt_fig, + } + + +def update_ipopt_output_plot(args, ocp): + """ + This function updated the plots for the ipopt output: x, f, g, inf_pr, inf_du. + We currently do not have access to the iteration number, weather we are currently in restoration, the lg(mu), the length of the current step, the alpha_du, or the alpha_pr. + inf_pr is obtained from the maximum absolute value of the constraints. + inf_du is obtained from the maximum absolute value of the equation 4a in the ipopt original paper. + """ + + x = args["x"] + f = args["f"] + lam_x = args["lam_x"] + lam_g = args["lam_g"] + lam_p = args["lam_p"] + + inf_pr = np.max(np.abs(args["g"])) + + grad_f_func = ocp.ipopt_plots["grad_f_func"] + grad_g_func = ocp.ipopt_plots["grad_g_func"] + + grad_f = grad_f_func(x) + grad_g_lam = grad_g_func(x) @ lam_g + eq_4a = np.max(np.abs(grad_f + grad_g_lam - lam_x)) + inf_du = np.max(np.abs(eq_4a)) + + ocp.ipopt_plots["f"].append(float(f)) + ocp.ipopt_plots["inf_pr"].append(float(inf_pr)) + ocp.ipopt_plots["inf_du"].append(float(inf_du)) + ocp.ipopt_plots["grad_f"].append(float(np.max(np.abs(grad_f)))) + ocp.ipopt_plots["grad_g"].append(float(np.max(np.abs(grad_g_lam)))) + ocp.ipopt_plots["lam_x"].append(float(np.max(np.abs(lam_x)))) + + ocp.ipopt_plots["plots"][0].set_ydata(ocp.ipopt_plots["f"]) + ocp.ipopt_plots["plots"][1].set_ydata(ocp.ipopt_plots["inf_pr"]) + ocp.ipopt_plots["plots"][2].set_ydata(ocp.ipopt_plots["inf_du"]) + ocp.ipopt_plots["plots"][3].set_ydata(ocp.ipopt_plots["grad_f"]) + ocp.ipopt_plots["plots"][4].set_ydata(ocp.ipopt_plots["grad_g"]) + ocp.ipopt_plots["plots"][5].set_ydata(ocp.ipopt_plots["lam_x"]) + + ocp.ipopt_plots["axs"][0].set_ylim(np.min(ocp.ipopt_plots["f"]), np.max(ocp.ipopt_plots["f"])) + ocp.ipopt_plots["axs"][1].set_ylim(np.min(ocp.ipopt_plots["inf_pr"]), np.max(ocp.ipopt_plots["inf_pr"])) + ocp.ipopt_plots["axs"][2].set_ylim( + np.min( + np.array( + [ + 1e8, + np.min(np.abs(ocp.ipopt_plots["inf_du"])), + np.min(np.abs(ocp.ipopt_plots["grad_f"])), + np.min(np.abs(ocp.ipopt_plots["grad_g"])), + np.min(np.abs(ocp.ipopt_plots["lam_x"])), + ] + ) + ), + np.max( + np.array( + [ + 1e-8, + np.max(np.abs(ocp.ipopt_plots["inf_du"])), + np.max(np.abs(ocp.ipopt_plots["grad_f"])), + np.max(np.abs(ocp.ipopt_plots["grad_g"])), + np.max(np.abs(ocp.ipopt_plots["lam_x"])), + ] + ) + ), + ) + + for i in range(6): + ocp.ipopt_plots["plots"][i].set_xdata(range(len(ocp.ipopt_plots["f"]))) + for i in range(3): + ocp.ipopt_plots["axs"][i].set_xlim(0, len(ocp.ipopt_plots["f"])) + + ocp.ipopt_plots["ipopt_fig"].canvas.draw() diff --git a/bioptim/gui/online_callback.py b/bioptim/gui/online_callback.py index 151cd9d62..d09b33494 100644 --- a/bioptim/gui/online_callback.py +++ b/bioptim/gui/online_callback.py @@ -58,7 +58,16 @@ def __init__(self, ocp, opts: dict = None, show_options: dict = None): Callback.__init__(self) self.ocp = ocp self.nx = self.ocp.variables_vector.shape[0] - self.ng = 0 + + # There must be an option to add an if here + from ..interfaces.ipopt_interface import IpoptInterface + + interface = IpoptInterface(ocp) + all_g, all_g_bounds = interface.dispatch_bounds() + self.ng = all_g.shape[0] + + v = interface.ocp.variables_vector + self.construct("AnimateCallback", opts) self.queue = mp.Queue() @@ -160,7 +169,10 @@ def eval(self, arg: list | tuple) -> list: A list of error index """ send = self.queue.put - send(arg[0]) + args_dict = {} + for i, s in enumerate(nlpsol_out()): + args_dict[s] = arg[i] + send(args_dict) return [0] class ProcessPlotter(object): @@ -221,8 +233,8 @@ def callback(self) -> bool: """ while not self.pipe.empty(): - v = self.pipe.get() - self.plot.update_data(v) + args = self.pipe.get() + self.plot.update_data(args) for i, fig in enumerate(self.plot.all_figures): fig.canvas.draw() diff --git a/bioptim/gui/plot.py b/bioptim/gui/plot.py index 784127717..c5c37787c 100644 --- a/bioptim/gui/plot.py +++ b/bioptim/gui/plot.py @@ -1,11 +1,11 @@ import tkinter -from itertools import accumulate from typing import Callable, Any import numpy as np -from casadi import DM +from casadi import DM, jacobian from matplotlib import pyplot as plt, lines from matplotlib.ticker import StrMethodFormatter +from matplotlib.cm import get_cmap from ..dynamics.ode_solver import OdeSolver from ..limits.path_conditions import Bounds @@ -71,6 +71,7 @@ def __init__( label: list = None, compute_derivative: bool = False, integration_rule: QuadratureRule = QuadratureRule.RECTANGLE_LEFT, + all_variables_in_one_subplot: bool = False, **parameters: Any, ): """ @@ -100,6 +101,8 @@ def __init__( Label of the curve to plot (to be added to the legend) compute_derivative: bool If the function should send the next node with x and u. Prevents from computing all at once (therefore a bit slower) + all_variables_in_one_subplot: bool + If all indices of the variables should be put on the same graph. This is not cute, but allows to display variables with a lot of entries. """ self.function = update_function @@ -125,6 +128,7 @@ def __init__( raise NotImplementedError(f"{integration_rule} has not been implemented yet.") self.integration_rule = integration_rule self.parameters = parameters + self.all_variables_in_one_subplot = all_variables_in_one_subplot class PlotOcp: @@ -285,6 +289,18 @@ def __init__( if self.plot_options["general_options"]["use_tight_layout"]: fig.tight_layout() + if self.ocp.plot_ipopt_outputs: + from ..gui.ipopt_output_plot import create_ipopt_output_plot + from ..interfaces.ipopt_interface import IpoptInterface + + interface = IpoptInterface(self.ocp) + create_ipopt_output_plot(ocp, interface) + + if self.ocp.plot_check_conditioning: + from ..gui.check_conditioning import create_conditioning_plots + + create_conditioning_plots(ocp) + def _update_time_vector(self, phase_times): """ Setup the time and time integrated vector, which is the x-axes of the graphs @@ -399,7 +415,12 @@ def legend_without_duplicate_labels(ax): for nlp in self.ocp.nlp ] ) - n_cols, n_rows = PlotOcp._generate_windows_size(nb_subplots) + # TODO: get rid of all_variables_in_one_subplot by fixing the mapping appropriately + if not nlp.plot[variable].all_variables_in_one_subplot: + n_cols, n_rows = PlotOcp._generate_windows_size(nb_subplots) + else: + n_cols = 1 + n_rows = 1 axes = self.__add_new_axis(variable, nb_subplots, n_rows, n_cols) self.axes[variable] = [nlp.plot[variable], axes] if not y_min_all[var_idx]: @@ -415,17 +436,24 @@ def legend_without_duplicate_labels(ax): mapping_to_first_index = nlp.plot[variable].phase_mappings.to_first.map_idx for ctr in mapping_to_first_index: - ax = axes[ctr] - if ctr in mapping_to_first_index: - index_legend = mapping_to_first_index.index(ctr) - if len(nlp.plot[variable].legend) > index_legend: - ax.set_title(nlp.plot[variable].legend[index_legend]) + if not nlp.plot[variable].all_variables_in_one_subplot: + ax = axes[ctr] + if ctr in mapping_to_first_index: + index_legend = mapping_to_first_index.index(ctr) + if len(nlp.plot[variable].legend) > index_legend: + ax.set_title(nlp.plot[variable].legend[index_legend]) + else: + ax = axes[0] ax.grid(**self.plot_options["grid"]) ax.set_xlim(self.t[-1][[0, -1]]) if nlp.plot[variable].ylim: ax.set_ylim(nlp.plot[variable].ylim) - elif self.show_bounds and nlp.plot[variable].bounds: + elif ( + self.show_bounds + and nlp.plot[variable].bounds + and not nlp.plot[variable].all_variables_in_one_subplot + ): if nlp.plot[variable].bounds.type != InterpolationType.CUSTOM: y_min = nlp.plot[variable].bounds.min[mapping_to_first_index.index(ctr), :].min() y_max = nlp.plot[variable].bounds.max[mapping_to_first_index.index(ctr), :].max() @@ -646,7 +674,10 @@ def show(): plt.show() - def update_data(self, v: np.ndarray): + def update_data( + self, + args: dict, + ): """ Update ydata from the variable a solution structure @@ -658,7 +689,7 @@ def update_data(self, v: np.ndarray): self.ydata = [] - sol = Solution.from_vector(self.ocp, v) + sol = Solution.from_vector(self.ocp, args["x"]) data_states_decision = sol.decision_states(scaled=True, to_merge=SolutionMerge.KEYS) data_states_stepwise = sol.stepwise_states(scaled=True, to_merge=SolutionMerge.KEYS) @@ -707,6 +738,16 @@ def update_data(self, v: np.ndarray): self.__update_axes() + if self.ocp.plot_ipopt_outputs: + from ..gui.ipopt_output_plot import update_ipopt_output_plot + + update_ipopt_output_plot(args, self.ocp) + + if self.ocp.plot_check_conditioning: + from ..gui.check_conditioning import update_conditioning_plots + + update_conditioning_plots(args["x"], self.ocp) + def _compute_y_from_plot_func( self, custom_plot: CustomPlot, phase_idx, time_stepwise, dt, x_decision, x_stepwise, u, p, a ): diff --git a/bioptim/interfaces/acados_interface.py b/bioptim/interfaces/acados_interface.py index e1fa5b55c..1460fcf11 100644 --- a/bioptim/interfaces/acados_interface.py +++ b/bioptim/interfaces/acados_interface.py @@ -180,7 +180,7 @@ def __acados_export_model(self, ocp): x_sym = vertcat(p_sym, x_sym) x_dot_sym = SX.sym("x_dot", x_sym.shape[0], x_sym.shape[1]) - f_expl = vertcat([0] * self.nparams, ocp.nlp[0].dynamics_func[0](t, x, u, p, a)) + f_expl = vertcat([0] * self.nparams, ocp.nlp[0].dynamics_func(t, x, u, p, a)) f_impl = x_dot_sym - f_expl self.acados_model.f_impl_expr = f_impl diff --git a/bioptim/interfaces/interface_utils.py b/bioptim/interfaces/interface_utils.py index 876eb74fc..9014516b8 100644 --- a/bioptim/interfaces/interface_utils.py +++ b/bioptim/interfaces/interface_utils.py @@ -45,6 +45,7 @@ def generic_solve(interface, expand_during_shake_tree=False) -> dict: ------- A reference to the solution """ + v = interface.ocp.variables_vector v_bounds = interface.ocp.bounds_vectors v_init = interface.ocp.init_vector @@ -152,42 +153,65 @@ def generic_set_lagrange_multiplier(interface, sol: Solution): return sol -def generic_dispatch_bounds(interface): +def generic_dispatch_bounds(interface, include_g: bool, include_g_internal: bool, include_g_implicit: bool): """ Parse the bounds of the full ocp to a SQP-friendly one + + Parameters + ---------- + interface: + A reference to the current interface + include_g: bool + If the g bounds should be included + include_g_internal: bool + If the g_internal bounds should be included + include_g_implicit: bool + If the g_implicit bounds should be included """ all_g = interface.ocp.cx() all_g_bounds = Bounds("all_g", interpolation=InterpolationType.CONSTANT) - all_g = vertcat(all_g, interface.get_all_penalties(interface.ocp, interface.ocp.g_internal)) - for g in interface.ocp.g_internal: - all_g_bounds.concatenate(g.bounds) - - all_g = vertcat(all_g, interface.get_all_penalties(interface.ocp, interface.ocp.g_implicit)) - for g in interface.ocp.g_implicit: - all_g_bounds.concatenate(g.bounds) - - all_g = vertcat(all_g, interface.get_all_penalties(interface.ocp, interface.ocp.g)) - for g in interface.ocp.g: - all_g_bounds.concatenate(g.bounds) - - for nlp in interface.ocp.nlp: - all_g = vertcat(all_g, interface.get_all_penalties(nlp, nlp.g_internal)) - for g in nlp.g_internal: - for _ in g.node_idx: + if include_g_internal: + all_g = vertcat(all_g, interface.get_all_penalties(interface.ocp, interface.ocp.g_internal)) + for g in interface.ocp.g_internal: + if g != []: all_g_bounds.concatenate(g.bounds) - all_g = vertcat(all_g, interface.get_all_penalties(nlp, nlp.g_implicit)) - for g in nlp.g_implicit: - for _ in g.node_idx: + if include_g_implicit: + all_g = vertcat(all_g, interface.get_all_penalties(interface.ocp, interface.ocp.g_implicit)) + for g in interface.ocp.g_implicit: + if g != []: all_g_bounds.concatenate(g.bounds) - all_g = vertcat(all_g, interface.get_all_penalties(nlp, nlp.g)) - for g in nlp.g: - for _ in g.node_idx: + if include_g: + all_g = vertcat(all_g, interface.get_all_penalties(interface.ocp, interface.ocp.g)) + for g in interface.ocp.g: + if g != []: all_g_bounds.concatenate(g.bounds) + for nlp in interface.ocp.nlp: + if include_g_internal: + all_g = vertcat(all_g, interface.get_all_penalties(nlp, nlp.g_internal)) + for g in nlp.g_internal: + if g != []: + for _ in g.node_idx: + all_g_bounds.concatenate(g.bounds) + + if include_g_implicit: + all_g = vertcat(all_g, interface.get_all_penalties(nlp, nlp.g_implicit)) + for g in nlp.g_implicit: + if g != []: + for _ in g.node_idx: + all_g_bounds.concatenate(g.bounds) + + if include_g: + all_g = vertcat(all_g, interface.get_all_penalties(nlp, nlp.g)) + for g in nlp.g: + if g != []: + for _ in g.node_idx: + all_g_bounds.concatenate(g.bounds) + if isinstance(all_g_bounds.min, (SX, MX)) or isinstance(all_g_bounds.max, (SX, MX)): raise RuntimeError(f"{interface.solver_name} doesn't support SX/MX types in constraints bounds") return all_g, all_g_bounds diff --git a/bioptim/interfaces/ipopt_interface.py b/bioptim/interfaces/ipopt_interface.py index b5fdaa63c..754d1be47 100644 --- a/bioptim/interfaces/ipopt_interface.py +++ b/bioptim/interfaces/ipopt_interface.py @@ -107,11 +107,13 @@ def set_lagrange_multiplier(self, sol: Solution): """ sol = generic_set_lagrange_multiplier(self, sol) - def dispatch_bounds(self): + def dispatch_bounds(self, include_g: bool = True, include_g_internal: bool = True, include_g_implicit: bool = True): """ Parse the bounds of the full ocp to a Ipopt-friendly one """ - return generic_dispatch_bounds(self) + return generic_dispatch_bounds( + self, include_g=include_g, include_g_internal=include_g_internal, include_g_implicit=include_g_implicit + ) def dispatch_obj_func(self): """ diff --git a/bioptim/interfaces/ipopt_options.py b/bioptim/interfaces/ipopt_options.py index 8c476a75b..4a4fcbef3 100644 --- a/bioptim/interfaces/ipopt_options.py +++ b/bioptim/interfaces/ipopt_options.py @@ -63,6 +63,8 @@ class IPOPT(GenericSolver): The valid range for this integer option is 0 ≤ print_level ≤ 12 and its default value is 5. _c_compile: bool True if you want to compile in C the code. + _check_derivatives_for_naninf: bool + If true, the Hessian will be checked for nan/inf values. If false this computational problem is silent. """ type: SolverType = SolverType.IPOPT @@ -92,6 +94,7 @@ class IPOPT(GenericSolver): _bound_frac: float = 0.01 _print_level: int = 5 _c_compile: bool = False + _check_derivatives_for_naninf: str = "no" # "yes" @property def tol(self): @@ -189,6 +192,10 @@ def print_level(self): def c_compile(self): return self._c_compile + @property + def check_derivatives_for_naninf(self): + return self._check_derivatives_for_naninf + def set_tol(self, val: float): self._tol = val @@ -261,6 +268,10 @@ def set_print_level(self, num: int): def set_c_compile(self, val: bool): self._c_compile = val + def set_check_derivatives_for_naninf(self, val: bool): + string_val = "yes" if val else "no" + self._check_derivatives_for_naninf = string_val + def set_convergence_tolerance(self, val: float): self._tol = val self._compl_inf_tol = val diff --git a/bioptim/interfaces/solve_ivp_interface.py b/bioptim/interfaces/solve_ivp_interface.py index 5c254fc44..b1dfab09e 100644 --- a/bioptim/interfaces/solve_ivp_interface.py +++ b/bioptim/interfaces/solve_ivp_interface.py @@ -5,10 +5,11 @@ from scipy.integrate import solve_ivp from scipy.interpolate import interp1d -from ..misc.enums import Shooting, ControlType, SolutionIntegrator +from ..misc.enums import Shooting, ControlType, SolutionIntegrator, PhaseDynamics def solve_ivp_interface( + list_of_dynamics: list[Callable], shooting_type: Shooting, nlp, t: list[np.ndarray], @@ -75,11 +76,10 @@ def solve_ivp_interface( if len(x0i.shape) > 1: x0i = x0i[:, 0] - func = nlp.dynamics_func[0] if len(nlp.dynamics_func) == 1 else nlp.dynamics_func[node] result = _solve_ivp_scipy_interface( - lambda t, x: np.array(func(t, x, _control_function(control_type, t, t_span, u[node]), p, a[node]))[ - :, 0 - ], + lambda t, x: np.array( + list_of_dynamics[node](t, x, _control_function(control_type, t, t_span, u[node]), p, a[node]) + )[:, 0], x0=x0i, t_span=np.array(t_span), t_eval=t_eval, diff --git a/bioptim/interfaces/sqp_interface.py b/bioptim/interfaces/sqp_interface.py index aaa67f9cb..f5baf9993 100644 --- a/bioptim/interfaces/sqp_interface.py +++ b/bioptim/interfaces/sqp_interface.py @@ -106,11 +106,13 @@ def set_lagrange_multiplier(self, sol: Solution): raise NotImplementedError("This is broken") # generic_set_lagrange_multiplier(self, sol) - def dispatch_bounds(self): + def dispatch_bounds(self, include_g: bool = True, include_g_internal: bool = True, include_g_implicit: bool = True): """ Parse the bounds of the full ocp to a SQP-friendly one """ - return generic_dispatch_bounds(self) + return generic_dispatch_bounds( + self, include_g=include_g, include_g_internal=include_g_internal, include_g_implicit=include_g_implicit + ) def dispatch_obj_func(self): """ diff --git a/bioptim/limits/constraints.py b/bioptim/limits/constraints.py index 60053e473..93ef3252f 100644 --- a/bioptim/limits/constraints.py +++ b/bioptim/limits/constraints.py @@ -122,6 +122,7 @@ def _add_penalty_to_pool(self, controller: list[PenaltyController, ...]): pool = controller.get_nlp.g if controller is not None and controller.get_nlp else controller.ocp.g else: raise ValueError(f"Invalid constraint type {self.penalty_type}.") + pool[self.list_index] = self def ensure_penalty_sanity(self, ocp, nlp): @@ -177,6 +178,7 @@ def add(self, constraint: Callable | Constraint | Any, **extra_arguments: Any): else: super(ConstraintList, self)._add(option_type=Constraint, constraint=constraint, **extra_arguments) + # TODO: add an InternalConstraint option type? Because now the list_index is wrong def print(self): """ @@ -1221,6 +1223,7 @@ def _add_penalty_to_pool(self, controller: list[PenaltyController, ...]): pool = controller.get_nlp.g if controller is not None and controller.get_nlp else controller.ocp.g else: raise ValueError(f"Invalid constraint type {self.penalty_type}.") + pool[self.list_index] = self def ensure_penalty_sanity(self, ocp, nlp): diff --git a/bioptim/limits/multinode_constraint.py b/bioptim/limits/multinode_constraint.py index a72d71812..030643dfb 100644 --- a/bioptim/limits/multinode_constraint.py +++ b/bioptim/limits/multinode_constraint.py @@ -3,7 +3,7 @@ import numpy as np from .path_conditions import Bounds -from ..misc.enums import InterpolationType +from ..misc.enums import InterpolationType, PenaltyType, ConstraintType from ..misc.fcn_enum import FcnEnum from .multinode_penalty import MultinodePenalty, MultinodePenaltyList, MultinodePenaltyFunctions @@ -49,7 +49,17 @@ def set_bounds(self): raise RuntimeError(f"bounds rows is {self.bounds.shape[0]} but should be {self.rows} or empty") def _get_pool_to_add_penalty(self, ocp, nlp): - return nlp.g_internal if nlp else ocp.g_internal + + if self.penalty_type == PenaltyType.INTERNAL: + pool = nlp.g_internal if nlp else ocp.g_internal + elif self.penalty_type == ConstraintType.IMPLICIT: + pool = nlp.g_implicit if nlp else ocp.g_implicit + elif self.penalty_type == PenaltyType.USER: + pool = nlp.g if nlp else ocp.g + else: + raise ValueError(f"Invalid constraint type {self.penalty_type}.") + + return pool class MultinodeConstraintList(MultinodePenaltyList): diff --git a/bioptim/limits/multinode_objective.py b/bioptim/limits/multinode_objective.py index 145ed2313..7ead4bf2f 100644 --- a/bioptim/limits/multinode_objective.py +++ b/bioptim/limits/multinode_objective.py @@ -1,5 +1,6 @@ from typing import Any +from ..misc.enums import PenaltyType from ..misc.fcn_enum import FcnEnum from .multinode_penalty import MultinodePenalty, MultinodePenaltyList, MultinodePenaltyFunctions from .objective_functions import ObjectiveFunction @@ -15,7 +16,14 @@ def __init__(self, *args, weight: float = 0, is_stochastic: bool = False, **kwar self.is_stochastic = is_stochastic def _get_pool_to_add_penalty(self, ocp, nlp): - return nlp.J_internal if nlp else ocp.J_internal + if self.penalty_type == PenaltyType.INTERNAL: + pool = nlp.J_internal if nlp else ocp.J_internal + elif self.penalty_type == PenaltyType.USER: + pool = nlp.J if nlp else ocp.J + else: + raise ValueError(f"Invalid objective type {self.penalty_type}.") + + return pool class MultinodeObjectiveList(MultinodePenaltyList): @@ -44,7 +52,7 @@ def add(self, multinode_objective: Any, **extra_arguments: Any): option_type=MultinodeObjective, multinode_penalty=multinode_objective, _multinode_penalty_fcn=MultinodeObjectiveFcn, - **extra_arguments + **extra_arguments, ) diff --git a/bioptim/limits/multinode_penalty.py b/bioptim/limits/multinode_penalty.py index cc19c79bd..2b4d959b2 100644 --- a/bioptim/limits/multinode_penalty.py +++ b/bioptim/limits/multinode_penalty.py @@ -83,8 +83,11 @@ def _get_pool_to_add_penalty(self, ocp, nlp): raise NotImplementedError("This is an abstract method and should be implemented by child") def _add_penalty_to_pool(self, controller: list[PenaltyController, ...]): - ocp = controller[0].ocp - nlp = controller[0].get_nlp + + controller = controller[0] # This is a special case of Node.TRANSITION + + ocp = controller.ocp + nlp = controller.get_nlp pool = self._get_pool_to_add_penalty(ocp, nlp) pool[self.list_index] = self diff --git a/bioptim/limits/objective_functions.py b/bioptim/limits/objective_functions.py index 0633faf3b..561251f2e 100644 --- a/bioptim/limits/objective_functions.py +++ b/bioptim/limits/objective_functions.py @@ -90,6 +90,7 @@ def _add_penalty_to_pool(self, controller: list[PenaltyController]): pool = controller.get_nlp.J if controller is not None and controller.get_nlp else controller.ocp.J else: raise ValueError(f"Invalid objective type {self.penalty_type}.") + pool[self.list_index] = self def ensure_penalty_sanity(self, ocp, nlp): @@ -499,6 +500,7 @@ def _add_penalty_to_pool(self, controller: list[PenaltyController]): pool = controller.ocp.J else: raise ValueError(f"Invalid objective type {self.penalty_type}.") + pool[self.list_index] = self def ensure_penalty_sanity(self, ocp, nlp): diff --git a/bioptim/limits/penalty.py b/bioptim/limits/penalty.py index e00e499ca..fbdde5c00 100644 --- a/bioptim/limits/penalty.py +++ b/bioptim/limits/penalty.py @@ -166,6 +166,10 @@ def stochastic_minimize_expected_feedback_efforts(penalty: PenaltyOption, contro controller : PenaltyController Controller to be used to compute the expected effort. """ + + if penalty.target is not None: + raise RuntimeError("It is not possible to use a target for the expected feedback effort.") + CX_eye = SX_eye if controller.ocp.cx == SX else MX_eye sensory_noise_matrix = controller.model.sensory_noise_magnitude * CX_eye( controller.model.sensory_noise_magnitude.shape[0] diff --git a/bioptim/limits/penalty_controller.py b/bioptim/limits/penalty_controller.py index d76ec37cd..0c1922de6 100644 --- a/bioptim/limits/penalty_controller.py +++ b/bioptim/limits/penalty_controller.py @@ -281,12 +281,10 @@ def integrate_extra_dynamics(self, dynamics_index): @property def dynamics(self): - return self._nlp.dynamics_func[0] + return self._nlp.dynamics_func def extra_dynamics(self, dynamics_index): - # +1 - index so "integrate_extra_dynamics" and "extra_dynamics" share the same index. - # This is a hack which should be dealt properly at some point - return self._nlp.dynamics_func[dynamics_index + 1] + return self._nlp.extra_dynamics_func[dynamics_index] @property def states_scaled(self) -> OptimizationVariableList: diff --git a/bioptim/limits/penalty_option.py b/bioptim/limits/penalty_option.py index b96e4dad2..e36542e31 100644 --- a/bioptim/limits/penalty_option.py +++ b/bioptim/limits/penalty_option.py @@ -420,6 +420,16 @@ def _set_penalty_function(self, controllers: list[PenaltyController], fcn: MX | if self.is_stochastic: sub_fcn = self.transform_penalty_to_stochastic(controller, sub_fcn, x) + from ..limits.constraints import ConstraintFcn + from ..limits.multinode_constraint import MultinodeConstraintFcn + + if ( + len(sub_fcn.shape) > 1 + and sub_fcn.shape[1] != 1 + and (isinstance(self.type, ConstraintFcn) or isinstance(self.type, MultinodeConstraintFcn)) + ): + raise RuntimeError("The constraint must return a vector not a matrix.") + is_trapezoidal = self.integration_rule in (QuadratureRule.APPROXIMATE_TRAPEZOIDAL, QuadratureRule.TRAPEZOIDAL) target_shape = tuple([len(self.rows), len(self.cols) + 1 if is_trapezoidal else len(self.cols)]) target_cx = controller.cx.sym("target", target_shape) @@ -844,7 +854,7 @@ def add_or_replace_to_penalty_pool(self, ocp, nlp): # The active controller is always the last one, and they all should be the same length anyway for node in range(len(controllers[-1])): # TODO WARNING THE NEXT IF STATEMENT IS A BUG DELIBERATELY INTRODUCED TO FIT THE PREVIOUS RESULTS. - # IT SHOULD BE REMOVED AS SOON AS THE MERGE IS DONE (AND VALUES OF THE TESTS ADJUSTED) + # IT SHOULD BE REMOVED AS SOON AS THE MERGE IS DONE (AND VALUES OF THE TESTS ADJUSTED) # @pariterre, can we remove? if self.integrate and self.target is not None: self.node_idx = controllers[0].t[:-1] if node not in self.node_idx: diff --git a/bioptim/optimization/non_linear_program.py b/bioptim/optimization/non_linear_program.py index ac0aeeb6a..5c43376fc 100644 --- a/bioptim/optimization/non_linear_program.py +++ b/bioptim/optimization/non_linear_program.py @@ -39,6 +39,8 @@ class NonLinearProgram: The dynamic MX or SX used during the current phase dynamics_func: Callable The dynamic function used during the current phase dxdt = f(x,u,p) + extra_dynamics_func: Callable + The extra dynamic function used during the current phase dxdt = f(x,u,p) implicit_dynamics_func: Callable The implicit dynamic function used during the current phase f(x,u,p,xdot) = 0 dynamics_type: Dynamics @@ -139,13 +141,12 @@ def __init__(self, phase_dynamics: PhaseDynamics): self.control_type = ControlType.CONSTANT self.cx = None self.dt = None - self.dynamics = ( - None # TODO Change this to a list to include extra_dynamics in a single vector (that matches dynamics_func) - ) + self.dynamics = [] self.extra_dynamics = [] self.dynamics_evaluation = DynamicsEvaluation() - self.dynamics_func: list = [] - self.implicit_dynamics_func: list = [] + self.dynamics_func = None + self.extra_dynamics_func: list = [] + self.implicit_dynamics_func = None self.dynamics_type = None self.external_forces: list[list[Any, ...], ...] | None = None # nodes x steps that are passed to the model self.g = [] diff --git a/bioptim/optimization/optimal_control_program.py b/bioptim/optimization/optimal_control_program.py index def862906..47766bfd0 100644 --- a/bioptim/optimization/optimal_control_program.py +++ b/bioptim/optimization/optimal_control_program.py @@ -6,6 +6,7 @@ import numpy as np from casadi import MX, SX, sum1, horzcat from matplotlib import pyplot as plt +import subprocess from .non_linear_program import NonLinearProgram as NLP from .optimization_vector import OptimizationVectorHelper @@ -584,6 +585,11 @@ def _check_arguments_and_build_nlp( NLP.add(self, "integrated_value_functions", integrated_value_functions, True) + # If we want to plot what is printed by IPOPT in the console + self.plot_ipopt_outputs = False + # If we want the conditioning of the problem to be plotted live + self.plot_check_conditioning = False + return ( constraints, objective_functions, @@ -1058,6 +1064,9 @@ def update_bounds( for key in nlp.controls.keys(): if f"{key}_controls" in nlp.plot and key in nlp.u_bounds.keys(): nlp.plot[f"{key}_controls"].bounds = nlp.u_bounds[key] + for key in nlp.algebraic_states.keys(): + if f"{key}_algebraic" in nlp.plot and key in nlp.a_bounds.keys(): + nlp.plot[f"{key}_algebraic"].bounds = nlp.a_bounds[key] def update_initial_guess( self, @@ -1305,6 +1314,12 @@ def add_penalty(_penalties): add_penalty(penalties_implicit) return + def add_plot_ipopt_outputs(self): + self.plot_ipopt_outputs = True + + def add_plot_check_conditioning(self): + self.plot_check_conditioning = True + def prepare_plots( self, automatically_organize: bool = True, diff --git a/bioptim/optimization/solution/solution.py b/bioptim/optimization/solution/solution.py index 188cbe281..786a6b330 100644 --- a/bioptim/optimization/solution/solution.py +++ b/bioptim/optimization/solution/solution.py @@ -4,7 +4,7 @@ import numpy as np from scipy import interpolate as sci_interp from scipy.interpolate import interp1d -from casadi import vertcat, DM +from casadi import vertcat, DM, Function from matplotlib import pyplot as plt from .solution_data import SolutionData, SolutionMerge, TimeAlignment, TimeResolution @@ -12,9 +12,19 @@ from ...limits.objective_functions import ObjectiveFcn from ...limits.path_conditions import InitialGuess, InitialGuessList from ...limits.penalty_helpers import PenaltyHelpers -from ...misc.enums import ControlType, CostType, Shooting, InterpolationType, SolverType, SolutionIntegrator, Node +from ...misc.enums import ( + ControlType, + CostType, + Shooting, + InterpolationType, + SolverType, + SolutionIntegrator, + Node, + PhaseDynamics, +) from ...dynamics.ode_solver import OdeSolver from ...interfaces.solve_ivp_interface import solve_ivp_interface +from ...models.protocols.stochastic_biomodel import StochasticBioModel class Solution: @@ -706,34 +716,14 @@ def copy(self, skip_data: bool = False) -> "Solution": new._parameters = deepcopy(self._parameters) return new - def integrate( - self, - shooting_type: Shooting = Shooting.SINGLE, - integrator: SolutionIntegrator = SolutionIntegrator.OCP, - to_merge: SolutionMerge | list[SolutionMerge] = None, - duplicated_times: bool = True, - return_time: bool = False, - ): + def _prepare_integrate(self, integrator: SolutionIntegrator): """ - Create a deepcopy of the Solution + Prepare the variables for the states integration and checks if the integrator is compatible with the ocp. Parameters ---------- - shooting_type: Shooting - The integration shooting type to use integrator: SolutionIntegrator - The type of integrator to use - to_merge: SolutionMerge | list[SolutionMerge, ...] - The type of merge to perform. If None, then no merge is performed. - duplicated_times: bool - If the times should be duplicated for each node. - If False, then the returned time vector will not have any duplicated times. - return_time: bool - If the time vector should be returned - - Returns - ------- - Return the integrated states + The integrator to use for the integration """ has_direct_collocation = sum([nlp.ode_solver.is_direct_collocation for nlp in self.ocp.nlp]) > 0 @@ -746,7 +736,7 @@ def integrate( ) has_trapezoidal = sum([isinstance(nlp.ode_solver, OdeSolver.TRAPEZOIDAL) for nlp in self.ocp.nlp]) > 0 - if has_trapezoidal: + if has_trapezoidal and integrator == SolutionIntegrator.OCP: raise ValueError( "When the ode_solver of the Optimal Control Problem is OdeSolver.TRAPEZOIDAL, " "we cannot use the SolutionIntegrator.OCP.\n" @@ -761,17 +751,52 @@ def integrate( x = self._decision_states.to_dict(to_merge=SolutionMerge.KEYS, scaled=False) u = self._stepwise_controls.to_dict(to_merge=SolutionMerge.KEYS, scaled=False) a = self._decision_algebraic_states.to_dict(to_merge=SolutionMerge.KEYS, scaled=False) + return t_spans, x, u, params, a + + def integrate( + self, + shooting_type: Shooting = Shooting.SINGLE, + integrator: SolutionIntegrator = SolutionIntegrator.OCP, + to_merge: SolutionMerge | list[SolutionMerge] = None, + duplicated_times: bool = True, + return_time: bool = False, + ): + """ + Create a deepcopy of the Solution + + Parameters + ---------- + shooting_type: Shooting + The integration shooting type to use + integrator: SolutionIntegrator + The type of integrator to use + to_merge: SolutionMerge | list[SolutionMerge, ...] + The type of merge to perform. If None, then no merge is performed. + duplicated_times: bool + If the times should be duplicated for each node. + If False, then the returned time vector will not have any duplicated times. + Default is True. + return_time: bool + If the time vector should be returned, default is False. + + Returns + ------- + Return the integrated states + """ + + t_spans, x, u, params, a = self._prepare_integrate(integrator=integrator) out: list = [None] * len(self.ocp.nlp) integrated_sol = None for p, nlp in enumerate(self.ocp.nlp): - next_x = self._states_for_phase_integration(shooting_type, p, integrated_sol, x, u, params, a) + first_x = self._states_for_phase_integration(shooting_type, p, integrated_sol, x, u, params, a) integrated_sol = solve_ivp_interface( + list_of_dynamics=[nlp.dynamics_func] * nlp.ns, shooting_type=shooting_type, nlp=nlp, t=t_spans[p], - x=next_x, + x=first_x, u=u[p], a=a[p], p=params, @@ -802,6 +827,110 @@ def integrate( else: return out if len(out) > 1 else out[0] + def noisy_integrate( + self, + integrator: SolutionIntegrator = SolutionIntegrator.OCP, + to_merge: SolutionMerge | list[SolutionMerge] = None, + size: int = 100, + ): + """ + TODO: Charbie! + """ + from ...optimization.stochastic_optimal_control_program import StochasticOptimalControlProgram + + if not isinstance(self.ocp, StochasticOptimalControlProgram): + raise ValueError("This method is only available for StochasticOptimalControlProgram.") + + t_spans, x, u, params, a = self._prepare_integrate(integrator=integrator) + + cov_index = self.ocp.nlp[0].algebraic_states["cov"].index + n_sub_nodes = x[0][0].shape[1] + motor_noise_index = self.ocp.nlp[0].parameters["motor_noise"].index + sensory_noise_index = ( + self.ocp.nlp[0].parameters["sensory_noise"].index + if len(list(self.ocp.nlp[0].parameters["sensory_noise"].index)) > 0 + else None + ) + + # initialize the out dictionary + out = [None] * len(self.ocp.nlp) + for p, nlp in enumerate(self.ocp.nlp): + out[p] = {} + for key in self.ocp.nlp[0].states.keys(): + out[p][key] = [None] * nlp.n_states_nodes + for i_node in range(nlp.ns): + out[p][key][i_node] = np.zeros((len(nlp.states[key].index), n_sub_nodes, size)) + out[p][key][nlp.ns] = np.zeros((len(nlp.states[key].index), 1, size)) + + cov_matrix = StochasticBioModel.reshape_to_matrix(a[0][0][cov_index, :], self.ocp.nlp[0].model.matrix_shape_cov) + first_x = np.random.multivariate_normal(x[0][0][:, 0], cov_matrix, size=size).T + for p, nlp in enumerate(self.ocp.nlp): + motor_noise = np.zeros((len(params[motor_noise_index]), nlp.ns, size)) + for i in range(len(params[motor_noise_index])): + motor_noise[i, :] = np.random.normal(0, params[motor_noise_index[i]], size=(nlp.ns, size)) + sensory_noise = ( + np.zeros((len(sensory_noise_index), nlp.ns, size)) if sensory_noise_index is not None else None + ) + if sensory_noise_index is not None: + for i in range(len(params[sensory_noise_index])): + sensory_noise[i, :] = np.random.normal(0, params[sensory_noise_index[i]], size=(nlp.ns, size)) + + without_noise_idx = [ + i for i in range(len(params)) if i not in motor_noise_index and i not in sensory_noise_index + ] + parameters_cx = nlp.parameters.cx[without_noise_idx] + parameters = params[without_noise_idx] + for i_random in range(size): + params_this_time = [] + list_of_dynamics = [] + for node in range(nlp.ns): + params_this_time += [nlp.parameters.cx] + params_this_time[node][motor_noise_index, :] = motor_noise[:, node, i_random] + if sensory_noise_index is not None: + params_this_time[node][sensory_noise_index, :] = sensory_noise[:, node, i_random] + + if len(nlp.extra_dynamics_func) > 1: + raise NotImplementedError("Noisy integration is not available for multiple extra dynamics.") + cas_func = Function( + "noised_extra_dynamics", + [nlp.time_cx, nlp.states.cx, nlp.controls.cx, parameters_cx, nlp.algebraic_states.cx], + [ + nlp.extra_dynamics_func[0]( + nlp.time_cx, + nlp.states.cx, + nlp.controls.cx, + params_this_time[node], + nlp.algebraic_states.cx, + ) + ], + ) + list_of_dynamics += [cas_func] + + integrated_sol = solve_ivp_interface( + list_of_dynamics=list_of_dynamics, + shooting_type=Shooting.SINGLE, + nlp=nlp, + t=t_spans[p], + x=[np.reshape(first_x[:, i_random], (-1, 1))], + u=u[p], # No need to add noise on the controls, the extra_dynamics should do it for us + a=a[p], + p=parameters, + method=integrator, + ) + for i_node in range(nlp.ns + 1): + for key in nlp.states.keys(): + states_integrated = ( + integrated_sol[i_node][nlp.states[key].index, :] + if n_sub_nodes > 1 + else integrated_sol[i_node][nlp.states[key].index, 0].reshape(-1, 1) + ) + out[p][key][i_node][:, :, i_random] = states_integrated + first_x[:, i_random] = np.reshape(integrated_sol[-1], (-1,)) + if to_merge: + out = SolutionData.from_unscaled(self.ocp, out, "x").to_dict(to_merge=to_merge, scaled=False) + + return out if len(out) > 1 else out[0] + def _states_for_phase_integration( self, shooting_type: Shooting, @@ -899,7 +1028,9 @@ def _integrate_stepwise(self) -> None: unscaled: list = [None] * len(self.ocp.nlp) for p, nlp in enumerate(self.ocp.nlp): + integrated_sol = solve_ivp_interface( + list_of_dynamics=[nlp.dynamics_func] * nlp.ns, shooting_type=Shooting.MULTIPLE, nlp=nlp, t=t_spans[p], @@ -1014,6 +1145,7 @@ def graphs( show_now: bool = True, shooting_type: Shooting = Shooting.MULTIPLE, integrator: SolutionIntegrator = SolutionIntegrator.OCP, + save_name: str = None, ): """ Show the graphs of the simulation @@ -1030,10 +1162,18 @@ def graphs( The type of interpolation integrator: SolutionIntegrator Use the scipy solve_ivp integrator for RungeKutta 45 instead of currently defined integrator + save_name: str + If a name is provided, the figures will be saved with this name """ plot_ocp = self.ocp.prepare_plots(automatically_organize, show_bounds, shooting_type, integrator) - plot_ocp.update_data(self.vector) + plot_ocp.update_data({"x": self.vector}) + if save_name: + if save_name.endswith(".png"): + save_name = save_name[:-4] + for i_fig, name_fig in enumerate(plt.get_figlabels()): + fig = plt.figure(i_fig + 1) + fig.savefig(f"{save_name}_{name_fig}.png", format="png") if show_now: plt.show() diff --git a/bioptim/optimization/solution/solution_data.py b/bioptim/optimization/solution/solution_data.py index df1451637..3db7a842d 100644 --- a/bioptim/optimization/solution/solution_data.py +++ b/bioptim/optimization/solution/solution_data.py @@ -237,8 +237,17 @@ def _to_scaled_values(unscaled: list, ocp, variable_type: str) -> list: if isinstance(unscaled[phase][key], list): # Nodes are not merged scaled[phase][key] = [] for node in range(len(unscaled[phase][key])): + scaled[phase][key].append( + np.zeros((unscaled[phase][key][node].shape[0], unscaled[phase][key][node].shape[1], 1)) + ) value = unscaled[phase][key][node] - scaled[phase][key].append(value / scale_factor.to_array(value.shape[1])) + if len(unscaled[phase][key][node].shape) == 3: + entry_size = unscaled[phase][key][node].shape[2] + scaling = np.repeat(scale_factor.to_array(value.shape[1])[:, :, np.newaxis], entry_size, axis=2) + else: + scaling = scale_factor.to_array(value.shape[1]) + scaled[phase][key][node] = value / scaling + elif isinstance(unscaled[phase][key], np.ndarray): # Nodes are merged value = unscaled[phase][key] scaled[phase][key] = value / scale_factor.to_array(value.shape[1]) diff --git a/bioptim/optimization/stochastic_optimal_control_program.py b/bioptim/optimization/stochastic_optimal_control_program.py index 351edcb04..908110021 100644 --- a/bioptim/optimization/stochastic_optimal_control_program.py +++ b/bioptim/optimization/stochastic_optimal_control_program.py @@ -25,9 +25,8 @@ from ..limits.path_conditions import InitialGuessList, InitialGuess from ..limits.penalty_controller import PenaltyController from ..limits.path_conditions import InitialGuessList -from ..misc.enums import PhaseDynamics, InterpolationType from ..misc.__version__ import __version__ -from ..misc.enums import Node, ControlType +from ..misc.enums import Node, ControlType, PenaltyType, PhaseDynamics, InterpolationType from ..misc.mapping import BiMappingList, Mapping, NodeMappingList, BiMapping from ..misc.utils import check_version from ..optimization.optimal_control_program import OptimalControlProgram @@ -209,7 +208,11 @@ def _prepare_stochastic_dynamics_explicit(self, constraints): Adds the internal constraint needed for the explicit formulation of the stochastic ocp. """ - constraints.add(ConstraintFcn.STOCHASTIC_MEAN_SENSORY_INPUT_EQUALS_REFERENCE, node=Node.ALL) + constraints.add( + ConstraintFcn.STOCHASTIC_MEAN_SENSORY_INPUT_EQUALS_REFERENCE, + node=Node.ALL, + # penalty_type=PenaltyType.INTERNAL, + ) penalty_m_dg_dz_list = MultinodeConstraintList() for i_phase, nlp in enumerate(self.nlp): @@ -218,12 +221,14 @@ def _prepare_stochastic_dynamics_explicit(self, constraints): MultinodeConstraintFcn.STOCHASTIC_HELPER_MATRIX_EXPLICIT, nodes_phase=(i_phase, i_phase), nodes=(i_node, i_node + 1), + # penalty_type=PenaltyType.INTERNAL, ) if i_phase > 0: penalty_m_dg_dz_list.add( MultinodeConstraintFcn.STOCHASTIC_HELPER_MATRIX_EXPLICIT, nodes_phase=(i_phase - 1, i_phase), nodes=(-1, 0), + # penalty_type=PenaltyType.INTERNAL, ) penalty_m_dg_dz_list.add_or_replace_to_penalty_pool(self) @@ -242,12 +247,14 @@ def _prepare_stochastic_dynamics_implicit(self, constraints): MultinodeConstraintFcn.STOCHASTIC_HELPER_MATRIX_IMPLICIT, nodes_phase=(i_phase, i_phase), nodes=(i_node, i_node + 1), + # penalty_type=PenaltyType.INTERNAL, ) if i_phase > 0 and i_phase < len(self.nlp) - 1: multi_node_penalties.add( MultinodeConstraintFcn.STOCHASTIC_HELPER_MATRIX_IMPLICIT, nodes_phase=(i_phase - 1, i_phase), nodes=(-1, 0), + # penalty_type=PenaltyType.INTERNAL, ) # Constraints for P @@ -256,6 +263,7 @@ def _prepare_stochastic_dynamics_implicit(self, constraints): ConstraintFcn.STOCHASTIC_COVARIANCE_MATRIX_CONTINUITY_IMPLICIT, node=Node.ALL, phase=i_phase, + # penalty_type=PenaltyType.INTERNAL, ) # Constraints for A @@ -264,6 +272,7 @@ def _prepare_stochastic_dynamics_implicit(self, constraints): ConstraintFcn.STOCHASTIC_DF_DX_IMPLICIT, node=Node.ALL, phase=i_phase, + # penalty_type=PenaltyType.INTERNAL, ) # Constraints for C @@ -273,12 +282,14 @@ def _prepare_stochastic_dynamics_implicit(self, constraints): MultinodeConstraintFcn.STOCHASTIC_DF_DW_IMPLICIT, nodes_phase=(i_phase, i_phase), nodes=(i_node, i_node + 1), + # penalty_type=PenaltyType.INTERNAL, ) if i_phase > 0 and i_phase < len(self.nlp) - 1: multi_node_penalties.add( MultinodeConstraintFcn.STOCHASTIC_DF_DW_IMPLICIT, nodes_phase=(i_phase, i_phase + 1), nodes=(-1, 0), + # penalty_type=PenaltyType.INTERNAL, ) multi_node_penalties.add_or_replace_to_penalty_pool(self) @@ -289,8 +300,15 @@ def _prepare_stochastic_dynamics_collocation(self, constraints, phase_transition integration. This is the real implementation suggested in Gillis 2013. """ - if "ref" in self.nlp[0].algebraic_states: - constraints.add(ConstraintFcn.STOCHASTIC_MEAN_SENSORY_INPUT_EQUALS_REFERENCE, node=Node.ALL) + # Constraints for ref + for i_phase, nlp in enumerate(self.nlp): + if "ref" in nlp.algebraic_states: + constraints.add( + ConstraintFcn.STOCHASTIC_MEAN_SENSORY_INPUT_EQUALS_REFERENCE, + node=Node.ALL, + phase=i_phase, + # penalty_type=PenaltyType.INTERNAL, + ) # Constraints for M for i_phase, nlp in enumerate(self.nlp): @@ -299,6 +317,7 @@ def _prepare_stochastic_dynamics_collocation(self, constraints, phase_transition node=Node.ALL_SHOOTING, phase=i_phase, expand=True, + # penalty_type=PenaltyType.INTERNAL, ) # Constraints for P inner-phase @@ -308,12 +327,17 @@ def _prepare_stochastic_dynamics_collocation(self, constraints, phase_transition node=Node.ALL_SHOOTING, phase=i_phase, expand=True, + # penalty_type=PenaltyType.INTERNAL, ) # Constraints for P inter-phase for i_phase, nlp in enumerate(self.nlp): if len(self.nlp) > 1 and i_phase < len(self.nlp) - 1: - phase_transition.add(PhaseTransitionFcn.COVARIANCE_CONTINUOUS, phase_pre_idx=i_phase) + phase_transition.add( + PhaseTransitionFcn.COVARIANCE_CONTINUOUS, + phase_pre_idx=i_phase, + # penalty_type=PenaltyType.INTERNAL + ) def _auto_initialize(self, x_init, u_init, parameter_init, a_init): def replace_initial_guess(key, n_var, var_init, a_init, i_phase): @@ -549,46 +573,6 @@ def _prepare_bounds_and_init( # Define the actual NLP problem OptimizationVectorHelper.declare_ocp_shooting_points(self) - @staticmethod - def load(file_path: str) -> list: - """ - Reload a previous optimization (*.bo) saved using save - - Parameters - ---------- - file_path: str - The path to the *.bo file - - Returns - ------- - The ocp and sol structure. If it was saved, the iterations are also loaded - """ - - with open(file_path, "rb") as file: - try: - data = pickle.load(file) - except BaseException as error_message: - raise ValueError( - f"The file '{file_path}' cannot be loaded, maybe the version of bioptim (version {__version__})\n" - f"is not the same as the one that created the file (version unknown). For more information\n" - "please refer to the original error message below\n\n" - f"{type(error_message).__name__}: {error_message}" - ) - ocp = StochasticOptimalControlProgram.from_loaded_data(data["ocp_initializer"]) - for key in data["versions"].keys(): - key_module = "biorbd_casadi" if key == "biorbd" else key - try: - check_version(sys.modules[key_module], data["versions"][key], ocp.version[key], exclude_max=False) - except ImportError: - raise ImportError( - f"Version of {key} from file ({data['versions'][key]}) is not the same as the " - f"installed version ({ocp.version[key]})" - ) - sol = data["sol"] - sol.ocp = Solution.SimplifiedOCP(ocp) - out = [ocp, sol] - return out - def _set_default_ode_solver(self): """It overrides the method in OptimalControlProgram that set a RK4 by default""" if isinstance(self.problem_type, SocpType.TRAPEZOIDAL_IMPLICIT) or isinstance( diff --git a/bioptim/optimization/variational_optimal_control_program.py b/bioptim/optimization/variational_optimal_control_program.py index 88ea8db85..f095ccfc1 100644 --- a/bioptim/optimization/variational_optimal_control_program.py +++ b/bioptim/optimization/variational_optimal_control_program.py @@ -243,20 +243,18 @@ def configure_dynamics_function( dynamics_dxdt = vertcat(*dynamics_dxdt) # Note: useless but needed to run bioptim as it need to test the size of xdot - nlp.dynamics_func = ( - Function( - "ForwardDyn", - [ - vertcat(nlp.time_mx, nlp.dt_mx), - nlp.states.scaled.mx_reduced, - nlp.controls.scaled.mx_reduced, - nlp.parameters.mx_reduced, - nlp.algebraic_states.scaled.mx_reduced, - ], - [dynamics_dxdt], - ["t_span", "x", "u", "p", "a"], - ["xdot"], - ), + nlp.dynamics_func = Function( + "ForwardDyn", + [ + vertcat(nlp.time_mx, nlp.dt_mx), + nlp.states.scaled.mx_reduced, + nlp.controls.scaled.mx_reduced, + nlp.parameters.mx_reduced, + nlp.algebraic_states.scaled.mx_reduced, + ], + [dynamics_dxdt], + ["t_span", "x", "u", "p", "a"], + ["xdot"], ) dt = MX.sym("time_step") @@ -289,66 +287,60 @@ def configure_dynamics_function( else: lambdas = None - nlp.implicit_dynamics_func = ( - Function( - "ThreeNodesIntegration", - three_nodes_input, - [ - self.bio_model.discrete_euler_lagrange_equations( - dt, - q_prev, - q_cur, - q_next, - control_prev, - control_cur, - control_next, - lambdas, - ) - ], - ), + nlp.implicit_dynamics_func = Function( + "ThreeNodesIntegration", + three_nodes_input, + [ + self.bio_model.discrete_euler_lagrange_equations( + dt, + q_prev, + q_cur, + q_next, + control_prev, + control_cur, + control_next, + lambdas, + ) + ], ) - nlp.implicit_dynamics_func_first_node = ( - Function( - "TwoFirstNodesIntegration", - two_first_nodes_input, - [ - self.bio_model.compute_initial_states( - dt, - q0, - qdot0, - q1, - control0, - control1, - lambdas, - ) - ], - ), + nlp.implicit_dynamics_func_first_node = Function( + "TwoFirstNodesIntegration", + two_first_nodes_input, + [ + self.bio_model.compute_initial_states( + dt, + q0, + qdot0, + q1, + control0, + control1, + lambdas, + ) + ], ) - nlp.implicit_dynamics_func_last_node = ( - Function( - "TwoLastNodesIntegration", - two_last_nodes_input, - [ - self.bio_model.compute_final_states( - dt, - q_penultimate, - q_ultimate, - qdot_ultimate, - controlN_minus_1, - controlN, - lambdas, - ) - ], - ), + nlp.implicit_dynamics_func_last_node = Function( + "TwoLastNodesIntegration", + two_last_nodes_input, + [ + self.bio_model.compute_final_states( + dt, + q_penultimate, + q_ultimate, + qdot_ultimate, + controlN_minus_1, + controlN, + lambdas, + ) + ], ) if expand: - nlp.dynamics_func = (nlp.dynamics_func[0].expand(),) - nlp.implicit_dynamics_func = (nlp.implicit_dynamics_func[0].expand(),) - nlp.implicit_dynamics_func_first_node = (nlp.implicit_dynamics_func_first_node[0].expand(),) - nlp.implicit_dynamics_func_last_node = (nlp.implicit_dynamics_func_last_node[0].expand(),) + nlp.dynamics_func = nlp.dynamics_func.expand() + nlp.implicit_dynamics_func = nlp.implicit_dynamics_func.expand() + nlp.implicit_dynamics_func_first_node = nlp.implicit_dynamics_func_first_node.expand() + nlp.implicit_dynamics_func_last_node = nlp.implicit_dynamics_func_last_node.expand() def configure_torque_driven(self, ocp: OptimalControlProgram, nlp: NonLinearProgram): """ @@ -401,7 +393,7 @@ def variational_integrator_three_nodes( """ if self.bio_model.has_holonomic_constraints: - return controllers[0].get_nlp.implicit_dynamics_func[0]( + return controllers[0].get_nlp.implicit_dynamics_func( controllers[0].dt.cx, controllers[0].states["q"].cx, controllers[1].states["q"].cx, @@ -412,7 +404,7 @@ def variational_integrator_three_nodes( controllers[1].states["lambdas"].cx, ) else: - return controllers[0].get_nlp.implicit_dynamics_func[0]( + return controllers[0].get_nlp.implicit_dynamics_func( controllers[0].dt.cx, controllers[0].states["q"].cx, controllers[1].states["q"].cx, @@ -443,7 +435,7 @@ def variational_integrator_initial( """ if self.bio_model.has_holonomic_constraints: - return controllers[0].get_nlp.implicit_dynamics_func_first_node[0]( + return controllers[0].get_nlp.implicit_dynamics_func_first_node( controllers[0].dt.cx, controllers[0].states["q"].cx, controllers[0].parameters.cx[:n_qdot], # hardcoded @@ -453,7 +445,7 @@ def variational_integrator_initial( controllers[0].states["lambdas"].cx, ) else: - return controllers[0].get_nlp.implicit_dynamics_func_first_node[0]( + return controllers[0].get_nlp.implicit_dynamics_func_first_node( controllers[0].dt.cx, controllers[0].states["q"].cx, controllers[0].parameters.cx[:n_qdot], # hardcoded @@ -484,7 +476,7 @@ def variational_integrator_final( """ if self.bio_model.has_holonomic_constraints: - return controllers[0].get_nlp.implicit_dynamics_func_last_node[0]( + return controllers[0].get_nlp.implicit_dynamics_func_last_node( controllers[0].dt.cx, controllers[0].states["q"].cx, controllers[1].states["q"].cx, @@ -494,7 +486,7 @@ def variational_integrator_final( controllers[1].states["lambdas"].cx, ) else: - return controllers[0].get_nlp.implicit_dynamics_func_last_node[0]( + return controllers[0].get_nlp.implicit_dynamics_func_last_node( controllers[0].dt.cx, controllers[0].states["q"].cx, controllers[1].states["q"].cx, diff --git a/tests/shard1/test__global_plots.py b/tests/shard1/test__global_plots.py index a5dd26c01..f99bc6707 100644 --- a/tests/shard1/test__global_plots.py +++ b/tests/shard1/test__global_plots.py @@ -6,6 +6,7 @@ import sys import os import pytest +import platform from casadi import Function, MX import numpy as np @@ -54,6 +55,37 @@ def test_plot_check_conditioning(phase_dynamics): sol.graphs(automatically_organize=False) +@pytest.mark.parametrize("phase_dynamics", [PhaseDynamics.SHARED_DURING_THE_PHASE, PhaseDynamics.ONE_PER_NODE]) +def test_plot_check_conditioning_live(phase_dynamics): + # Load graphs check conditioning + from bioptim.examples.getting_started import example_multiphase as ocp_module + + bioptim_folder = os.path.dirname(ocp_module.__file__) + + ocp = ocp_module.prepare_ocp( + biorbd_model_path=bioptim_folder + "/models/cube.bioMod", + long_optim=False, + phase_dynamics=phase_dynamics, + expand_dynamics=True, + ) + ocp.add_plot_check_conditioning() + + +@pytest.mark.parametrize("phase_dynamics", [PhaseDynamics.SHARED_DURING_THE_PHASE, PhaseDynamics.ONE_PER_NODE]) +def test_plot_ipopt_output_live(phase_dynamics): + from bioptim.examples.getting_started import example_multiphase as ocp_module + + bioptim_folder = os.path.dirname(ocp_module.__file__) + + ocp = ocp_module.prepare_ocp( + biorbd_model_path=bioptim_folder + "/models/cube.bioMod", + long_optim=False, + phase_dynamics=phase_dynamics, + expand_dynamics=True, + ) + ocp.add_plot_ipopt_outputs() + + @pytest.mark.parametrize("phase_dynamics", [PhaseDynamics.SHARED_DURING_THE_PHASE, PhaseDynamics.ONE_PER_NODE]) def test_plot_merged_graphs(phase_dynamics): # Load graphs_one_phase diff --git a/tests/shard1/test_dynamics.py b/tests/shard1/test_dynamics.py index 8a728d86e..514efa125 100644 --- a/tests/shard1/test_dynamics.py +++ b/tests/shard1/test_dynamics.py @@ -191,7 +191,7 @@ def test_torque_driven(with_contact, with_external_force, cx, rigidbody_dynamics params = np.random.rand(nlp.parameters.shape, nlp.ns) algebraic_states = np.random.rand(nlp.algebraic_states.shape, nlp.ns) time = np.random.rand(2) - x_out = np.array(nlp.dynamics_func[0](time, states, controls, params, algebraic_states)) + x_out = np.array(nlp.dynamics_func(time, states, controls, params, algebraic_states)) if rigidbody_dynamics == RigidBodyDynamics.ODE: if with_contact: contact_out = np.array(nlp.contact_forces_func(time, states, controls, params, algebraic_states)) @@ -339,7 +339,7 @@ def test_torque_driven_implicit(with_contact, cx, phase_dynamics): params = np.random.rand(nlp.parameters.shape, nlp.ns) algebraic_states = np.random.rand(nlp.algebraic_states.shape, nlp.ns) time = np.random.rand(2) - x_out = np.array(nlp.dynamics_func[0](time, states, controls, params, algebraic_states)) + x_out = np.array(nlp.dynamics_func(time, states, controls, params, algebraic_states)) if with_contact: contact_out = np.array(nlp.contact_forces_func(time, states, controls, params, algebraic_states)) @@ -414,7 +414,7 @@ def test_torque_driven_soft_contacts_dynamics(with_contact, cx, implicit_contact params = np.random.rand(nlp.parameters.shape, nlp.ns) algebraic_states = np.random.rand(nlp.algebraic_states.shape, nlp.ns) time = np.random.rand(2) - x_out = np.array(nlp.dynamics_func[0](time, states, controls, params, algebraic_states)) + x_out = np.array(nlp.dynamics_func(time, states, controls, params, algebraic_states)) if with_contact: contact_out = np.array(nlp.contact_forces_func(time, states, controls, params, algebraic_states)) @@ -581,7 +581,7 @@ def test_torque_derivative_driven(with_contact, with_external_force, cx, phase_d params = np.random.rand(nlp.parameters.shape, nlp.ns) algebraic_states = np.random.rand(nlp.algebraic_states.shape, nlp.ns) time = np.random.rand(2) - x_out = np.array(nlp.dynamics_func[0](time, states, controls, params, algebraic_states)) + x_out = np.array(nlp.dynamics_func(time, states, controls, params, algebraic_states)) if with_contact: contact_out = np.array(nlp.contact_forces_func(time, states, controls, params, algebraic_states)) @@ -718,7 +718,7 @@ def test_torque_derivative_driven_implicit(with_contact, cx, phase_dynamics): params = np.random.rand(nlp.parameters.shape, nlp.ns) algebraic_states = np.random.rand(nlp.algebraic_states.shape, nlp.ns) time = np.random.rand(2) - x_out = np.array(nlp.dynamics_func[0](time, states, controls, params, algebraic_states)) + x_out = np.array(nlp.dynamics_func(time, states, controls, params, algebraic_states)) if with_contact: contact_out = np.array(nlp.contact_forces_func(time, states, controls, params, algebraic_states)) @@ -825,7 +825,7 @@ def test_torque_derivative_driven_soft_contacts_dynamics(with_contact, cx, impli params = np.random.rand(nlp.parameters.shape, nlp.ns) algebraic_states = np.random.rand(nlp.algebraic_states.shape, nlp.ns) time = np.random.rand(2) - x_out = np.array(nlp.dynamics_func[0](time, states, controls, params, algebraic_states)) + x_out = np.array(nlp.dynamics_func(time, states, controls, params, algebraic_states)) if with_contact: contact_out = np.array(nlp.contact_forces_func(time, states, controls, params, algebraic_states)) @@ -1103,7 +1103,7 @@ def test_torque_activation_driven(with_contact, with_external_force, cx, phase_d params = np.random.rand(nlp.parameters.shape, nlp.ns) algebraic_states = np.random.rand(nlp.algebraic_states.shape, nlp.ns) time = np.random.rand(2) - x_out = np.array(nlp.dynamics_func[0](time, states, controls, params, algebraic_states)) + x_out = np.array(nlp.dynamics_func(time, states, controls, params, algebraic_states)) if with_contact: contact_out = np.array(nlp.contact_forces_func(time, states, controls, params, algebraic_states)) @@ -1306,7 +1306,7 @@ def test_torque_activation_driven_with_residual_torque( params = np.random.rand(nlp.parameters.shape, nlp.ns) algebraic_states = np.random.rand(nlp.algebraic_states.shape, nlp.ns) time = np.random.rand(2) - x_out = np.array(nlp.dynamics_func[0](time, states, controls, params, algebraic_states)) + x_out = np.array(nlp.dynamics_func(time, states, controls, params, algebraic_states)) if with_residual_torque: if with_external_force: @@ -1414,7 +1414,7 @@ def test_torque_driven_free_floating_base(cx, phase_dynamics): params = np.random.rand(nlp.parameters.shape, nlp.ns) algebraic_states = np.random.rand(nlp.algebraic_states.shape, nlp.ns) time = np.random.rand(2, 1) - x_out = np.array(nlp.dynamics_func[0](time, states, controls, params, algebraic_states)) + x_out = np.array(nlp.dynamics_func(time, states, controls, params, algebraic_states)) np.testing.assert_almost_equal( x_out[:, 0], @@ -1580,7 +1580,7 @@ def test_muscle_driven( params = np.random.rand(nlp.parameters.shape, nlp.ns) algebraic_states = np.random.rand(nlp.algebraic_states.shape, nlp.ns) time = np.random.rand(2) - x_out = np.array(nlp.dynamics_func[0](time, states, controls, params, algebraic_states)) + x_out = np.array(nlp.dynamics_func(time, states, controls, params, algebraic_states)) 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: @@ -2078,7 +2078,7 @@ def test_joints_acceleration_driven(cx, rigid_body_dynamics, phase_dynamics): params = np.random.rand(nlp.parameters.shape, nlp.ns) algebraic_states = np.random.rand(nlp.algebraic_states.shape, nlp.ns) time = np.random.rand(2) - x_out = np.array(nlp.dynamics_func[0](time, states, controls, params, algebraic_states)) + x_out = np.array(nlp.dynamics_func(time, states, controls, params, algebraic_states)) # 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]) @@ -2159,7 +2159,7 @@ def configure(ocp, nlp, with_contact=None): params = np.random.rand(nlp.parameters.shape, nlp.ns) algebraic_states = np.random.rand(nlp.algebraic_states.shape, nlp.ns) time = np.random.rand(2) - x_out = np.array(nlp.dynamics_func[0](time, states, controls, params, algebraic_states)) + x_out = np.array(nlp.dynamics_func(time, states, controls, params, algebraic_states)) if with_contact: contact_out = np.array(nlp.contact_forces_func(time, states, controls, params, algebraic_states)) diff --git a/tests/shard3/test_ligaments.py b/tests/shard3/test_ligaments.py index eb5098126..6414de347 100644 --- a/tests/shard3/test_ligaments.py +++ b/tests/shard3/test_ligaments.py @@ -83,7 +83,7 @@ def test_torque_driven_with_ligament(with_ligament, cx, phase_dynamics): params = np.random.rand(nlp.parameters.shape, nlp.ns) algebraic_states = np.random.rand(nlp.algebraic_states.shape, nlp.ns) time = np.random.rand(2) - x_out = np.array(nlp.dynamics_func[0](time, states, controls, params, algebraic_states)) + x_out = np.array(nlp.dynamics_func(time, states, controls, params, algebraic_states)) if with_ligament: np.testing.assert_almost_equal( x_out[:, 0], @@ -147,7 +147,7 @@ def test_torque_derivative_driven_with_ligament(with_ligament, cx, phase_dynamic params = np.random.rand(nlp.parameters.shape, nlp.ns) algebraic_states = np.random.rand(nlp.algebraic_states.shape, nlp.ns) time = np.random.rand(2) - x_out = np.array(nlp.dynamics_func[0](time, states, controls, params, algebraic_states)) + x_out = np.array(nlp.dynamics_func(time, states, controls, params, algebraic_states)) if with_ligament: np.testing.assert_almost_equal( x_out[:, 0], @@ -207,7 +207,7 @@ def test_torque_activation_driven_with_ligament(with_ligament, cx, phase_dynamic params = np.random.rand(nlp.parameters.shape, nlp.ns) algebraic_states = np.random.rand(nlp.algebraic_states.shape, nlp.ns) time = np.random.rand(2) - x_out = np.array(nlp.dynamics_func[0](time, states, controls, params, algebraic_states)) + x_out = np.array(nlp.dynamics_func(time, states, controls, params, algebraic_states)) if with_ligament: np.testing.assert_almost_equal( x_out[:, 0], @@ -275,7 +275,7 @@ def test_muscle_driven_with_ligament(with_ligament, cx, phase_dynamics): params = np.random.rand(nlp.parameters.shape, nlp.ns) algebraic_states = np.random.rand(nlp.algebraic_states.shape, nlp.ns) time = np.random.rand(2) - x_out = np.array(nlp.dynamics_func[0](time, states, controls, params, algebraic_states)) + x_out = np.array(nlp.dynamics_func(time, states, controls, params, algebraic_states)) if with_ligament: np.testing.assert_almost_equal( diff --git a/tests/shard3/test_passive_torque.py b/tests/shard3/test_passive_torque.py index c6794a254..846e31b0e 100644 --- a/tests/shard3/test_passive_torque.py +++ b/tests/shard3/test_passive_torque.py @@ -89,7 +89,7 @@ def test_torque_driven_with_passive_torque(with_passive_torque, cx, rigidbody_dy params = np.random.rand(nlp.parameters.shape, nlp.ns) algebraic_states = np.random.rand(nlp.algebraic_states.shape, nlp.ns) time = np.random.rand(2) - x_out = np.array(nlp.dynamics_func[0](time, states, controls, params, algebraic_states)) + x_out = np.array(nlp.dynamics_func(time, states, controls, params, algebraic_states)) if rigidbody_dynamics == RigidBodyDynamics.ODE: if with_passive_torque: np.testing.assert_almost_equal( @@ -175,7 +175,7 @@ def test_torque_derivative_driven_with_passive_torque(with_passive_torque, cx, p params = np.random.rand(nlp.parameters.shape, nlp.ns) algebraic_states = np.random.rand(nlp.algebraic_states.shape, nlp.ns) time = np.random.rand(2) - x_out = np.array(nlp.dynamics_func[0](time, states, controls, params, algebraic_states)) + x_out = np.array(nlp.dynamics_func(time, states, controls, params, algebraic_states)) if with_passive_torque: np.testing.assert_almost_equal( x_out[:, 0], @@ -267,7 +267,7 @@ def test_torque_activation_driven_with_passive_torque(with_passive_torque, with_ params = np.random.rand(nlp.parameters.shape, nlp.ns) algebraic_states = np.random.rand(nlp.algebraic_states.shape, nlp.ns) time = np.random.rand(2) - x_out = np.array(nlp.dynamics_func[0](time, states, controls, params, algebraic_states)) + x_out = np.array(nlp.dynamics_func(time, states, controls, params, algebraic_states)) if with_residual_torque: if with_passive_torque: np.testing.assert_almost_equal( @@ -388,7 +388,7 @@ def test_muscle_driven_with_passive_torque(with_passive_torque, rigidbody_dynami params = np.random.rand(nlp.parameters.shape, nlp.ns) algebraic_states = np.random.rand(nlp.algebraic_states.shape, nlp.ns) time = np.random.rand(2) - x_out = np.array(nlp.dynamics_func[0](time, states, controls, params, algebraic_states)) + x_out = np.array(nlp.dynamics_func(time, states, controls, params, algebraic_states)) if rigidbody_dynamics == RigidBodyDynamics.DAE_INVERSE_DYNAMICS: if with_passive_torque: diff --git a/tests/shard4/test_penalty.py b/tests/shard4/test_penalty.py index f68136c38..3d2e474de 100644 --- a/tests/shard4/test_penalty.py +++ b/tests/shard4/test_penalty.py @@ -1,5 +1,5 @@ import pytest -from casadi import DM, MX, vertcat +from casadi import DM, MX, vertcat, horzcat import numpy as np from bioptim import ( BiorbdModel, @@ -19,6 +19,7 @@ RigidBodyDynamics, ControlType, PhaseDynamics, + ConstraintList, ) from bioptim.limits.penalty_controller import PenaltyController from bioptim.limits.penalty import PenaltyOption @@ -1432,3 +1433,33 @@ def test_PenaltyFunctionAbstract_get_node(node, ns, phase_dynamics): np.testing.assert_almost_equal(controller.u_scaled, u_expected[2]) else: raise RuntimeError("Something went wrong") + + +def test_bad_shape_output_penalty(): + def bad_custom_function(controller: PenaltyController): + """ + This custom function returns a matrix, thus some terms will be ignored by bioptim! + """ + return horzcat(controller.states["q"].cx, controller.states["qdot"].cx, controller.controls["tau"].cx) + + def prepare_test_ocp_error(): + bioptim_folder = TestUtils.bioptim_folder() + + bio_model = BiorbdModel(bioptim_folder + "/examples/track/models/cube_and_line.bioMod") + dynamics = DynamicsList() + dynamics.add(DynamicsFcn.TORQUE_DRIVEN) + + constraints = ConstraintList() + constraints.add(bad_custom_function, node=Node.START) + + ocp = OptimalControlProgram( + bio_model, + dynamics, + 10, + 1.0, + constraints=constraints, + ) + return ocp + + with pytest.raises(RuntimeError, match="The constraint must return a vector not a matrix."): + ocp = prepare_test_ocp_error() diff --git a/tests/shard4/test_solver_options.py b/tests/shard4/test_solver_options.py index 3ba07ca3e..b477faffc 100644 --- a/tests/shard4/test_solver_options.py +++ b/tests/shard4/test_solver_options.py @@ -39,6 +39,7 @@ def test_ipopt_solver_options(): assert solver.bound_frac == 0.01 assert solver.print_level == 5 assert solver.c_compile is False + assert solver.check_derivatives_for_naninf == "no" solver.set_linear_solver("ma57") assert solver.linear_solver == "ma57" @@ -88,6 +89,8 @@ def test_ipopt_solver_options(): assert solver.print_level == 20 solver.set_c_compile(True) assert solver.c_compile is True + solver.set_check_derivatives_for_naninf(True) + assert solver.check_derivatives_for_naninf == "yes" solver.set_convergence_tolerance(21) assert solver.tol == 21 diff --git a/tests/shard5/test_global_stochastic_collocation.py b/tests/shard5/test_global_stochastic_collocation.py index 97d447ffa..e574048c9 100644 --- a/tests/shard5/test_global_stochastic_collocation.py +++ b/tests/shard5/test_global_stochastic_collocation.py @@ -3,7 +3,7 @@ import numpy as np from casadi import DM, vertcat -from bioptim import Solver, SocpType, SolutionMerge, PenaltyHelpers +from bioptim import Solver, SocpType, SolutionMerge, PenaltyHelpers, SolutionIntegrator @pytest.mark.parametrize("use_sx", [False, True]) @@ -332,17 +332,17 @@ def test_arm_reaching_torque_driven_collocations(use_sx: bool): penalty = socp.nlp[0].g_internal[0] for i_node in range(socp.n_shooting): x = PenaltyHelpers.states( - socp.nlp[0].g[9], + penalty, i_node, lambda p_idx, n_idx, sn_idx: states_sol[:, sn_idx, n_idx], ) u = PenaltyHelpers.controls( - socp.nlp[0].g[9], + penalty, i_node, lambda p_idx, n_idx, sn_idx: controls_sol[:, n_idx], ) a = PenaltyHelpers.states( - socp.nlp[0].g[9], + penalty, i_node, lambda p_idx, n_idx, sn_idx: algebraic_sol[:, n_idx], ) @@ -541,3 +541,13 @@ def test_obstacle_avoidance_direct_collocation(use_sx: bool): ), decimal=6, ) + + np.random.seed(42) + integrated_states = sol.noisy_integrate(integrator=SolutionIntegrator.SCIPY_RK45, to_merge=SolutionMerge.NODES) + integrated_stated_covariance = np.cov(integrated_states["q"][:, -1, :]) + np.testing.assert_almost_equal( + integrated_stated_covariance, np.array([[0.00404452, -0.00100082], [-0.00100082, 0.00382313]]), decimal=6 + ) + np.testing.assert_almost_equal( + cov[:, -1].reshape(4, 4)[:2, :2], np.array([[0.00266764, -0.0005587], [-0.0005587, 0.00134316]]), decimal=6 + ) diff --git a/tests/shard6/test_global_stochastic_except_collocation.py b/tests/shard6/test_global_stochastic_except_collocation.py index 167ddd239..76ed82f1b 100644 --- a/tests/shard6/test_global_stochastic_except_collocation.py +++ b/tests/shard6/test_global_stochastic_except_collocation.py @@ -5,7 +5,7 @@ import numpy as np from casadi import DM, vertcat -from bioptim import Solver, SolutionMerge +from bioptim import Solver, SolutionMerge, SolutionIntegrator from bioptim.examples.stochastic_optimal_control.arm_reaching_torque_driven_implicit import ExampleType @@ -607,6 +607,17 @@ def test_arm_reaching_torque_driven_implicit(with_cholesky, with_scaling, use_sx ] ), ) + + np.random.seed(42) + integrated_states = sol.noisy_integrate( + integrator=SolutionIntegrator.SCIPY_RK45, to_merge=SolutionMerge.NODES + ) + integrated_stated_covariance = np.cov(integrated_states["q"][:, -1, :]) + np.testing.assert_almost_equal( + integrated_stated_covariance, + np.array([[0.52918393, -1.36311773], [-1.36311773, 3.85744162]]), + decimal=6, + ) else: # Check some of the results k, ref, m, cov, a, c = (