diff --git a/bioptim/dynamics/configure_problem.py b/bioptim/dynamics/configure_problem.py index 333d1a376..9c0f4d622 100644 --- a/bioptim/dynamics/configure_problem.py +++ b/bioptim/dynamics/configure_problem.py @@ -1,7 +1,6 @@ -from typing import Callable, Any - import numpy as np from casadi import vertcat, Function, DM, horzcat +from typing import Callable, Any from .configure_new_variable import NewVariableConfiguration from .dynamics_functions import DynamicsFunctions @@ -1012,8 +1011,18 @@ def configure_lagrange_multipliers_function(ocp, nlp, dyn_func: Callable, **extr ["lagrange_multipliers"], ) - all_multipliers_names = [f"lagrange_multiplier_{i}" for i in range(nlp.model.nb_dependent_joints)] - all_multipliers_names_in_phase = [f"lagrange_multiplier_{i}" for i in range(nlp.model.nb_dependent_joints)] + all_multipliers_names = [] + for nlp_i in ocp.nlp: + if hasattr(nlp_i.model, "has_holonomic_constraints"): # making sure we have a HolonomicBiorbdModel + nlp_i_multipliers_names = [nlp_i.model.name_dof[i] for i in nlp_i.model.dependent_joint_index] + all_multipliers_names.extend( + [name for name in nlp_i_multipliers_names if name not in all_multipliers_names] + ) + + all_multipliers_names = [f"lagrange_multiplier_{name}" for name in all_multipliers_names] + all_multipliers_names_in_phase = [ + f"lagrange_multiplier_{nlp.model.name_dof[i]}" for i in nlp.model.dependent_joint_index + ] axes_idx = BiMapping( to_first=[i for i, c in enumerate(all_multipliers_names) if c in all_multipliers_names_in_phase], @@ -1032,7 +1041,7 @@ def configure_lagrange_multipliers_function(ocp, nlp, dyn_func: Callable, **extr @staticmethod def configure_qv(ocp, nlp, dyn_func: Callable, **extra_params): """ - Configure the contact points + Configure the qv, i.e. the dependent joint coordinates, to be plotted Parameters ---------- @@ -1066,9 +1075,15 @@ def configure_qv(ocp, nlp, dyn_func: Callable, **extra_params): ["q_v"], ) - all_multipliers_names = [nlp.model.name_dof[i] for i in nlp.model.independent_joint_index] - all_multipliers_names_in_phase = [nlp.model.name_dof[i] for i in nlp.model.independent_joint_index] + all_multipliers_names = [] + for nlp_i in ocp.nlp: + if hasattr(nlp_i.model, "has_holonomic_constraints"): # making sure we have a HolonomicBiorbdModel + nlp_i_multipliers_names = [nlp_i.model.name_dof[i] for i in nlp_i.model.dependent_joint_index] + all_multipliers_names.extend( + [name for name in nlp_i_multipliers_names if name not in all_multipliers_names] + ) + all_multipliers_names_in_phase = [nlp.model.name_dof[i] for i in nlp.model.dependent_joint_index] axes_idx = BiMapping( to_first=[i for i, c in enumerate(all_multipliers_names) if c in all_multipliers_names_in_phase], to_second=[i for i, c in enumerate(all_multipliers_names) if c in all_multipliers_names_in_phase], @@ -1086,7 +1101,7 @@ def configure_qv(ocp, nlp, dyn_func: Callable, **extra_params): @staticmethod def configure_qdotv(ocp, nlp, dyn_func: Callable, **extra_params): """ - Configure the contact points + Configure the qdot_v, i.e. the dependent joint velocities, to be plotted Parameters ---------- @@ -1123,9 +1138,15 @@ def configure_qdotv(ocp, nlp, dyn_func: Callable, **extra_params): ["qdot_v"], ) - all_multipliers_names = [nlp.model.name_dof[i] for i in nlp.model.independent_joint_index] - all_multipliers_names_in_phase = [nlp.model.name_dof[i] for i in nlp.model.independent_joint_index] + all_multipliers_names = [] + for nlp_i in ocp.nlp: + if hasattr(nlp_i.model, "has_holonomic_constraints"): # making sure we have a HolonomicBiorbdModel + nlp_i_multipliers_names = [nlp_i.model.name_dof[i] for i in nlp_i.model.dependent_joint_index] + all_multipliers_names.extend( + [name for name in nlp_i_multipliers_names if name not in all_multipliers_names] + ) + all_multipliers_names_in_phase = [nlp.model.name_dof[i] for i in nlp.model.dependent_joint_index] axes_idx = BiMapping( to_first=[i for i, c in enumerate(all_multipliers_names) if c in all_multipliers_names_in_phase], to_second=[i for i, c in enumerate(all_multipliers_names) if c in all_multipliers_names_in_phase], @@ -1321,7 +1342,7 @@ def configure_contact_function(ocp, nlp, dyn_func: Callable, **extra_params): nlp.plot["contact_forces"] = CustomPlot( lambda t0, phases_dt, node_idx, x, u, p, a, d: nlp.contact_forces_func( - [t0, t0 + phases_dt[nlp.phase_idx]], x, u, p, a, d + np.concatenate([t0, t0 + phases_dt[nlp.phase_idx]]), x, u, p, a, d ), plot_type=PlotType.INTEGRATED, axes_idx=axes_idx, @@ -1394,7 +1415,7 @@ def configure_soft_contact_function(ocp, nlp): ) nlp.plot[f"soft_contact_forces_{nlp.model.soft_contact_names[i_sc]}"] = CustomPlot( lambda t0, phases_dt, node_idx, x, u, p, a, d: nlp.soft_contact_forces_func( - [t0, t0 + phases_dt[nlp.phase_idx]], x, u, p, a, d + np.concatenate([t0, t0 + phases_dt[nlp.phase_idx]]), x, u, p, a, d )[(i_sc * 6) : ((i_sc + 1) * 6), :], plot_type=PlotType.INTEGRATED, axes_idx=phase_mappings, diff --git a/bioptim/gui/plot.py b/bioptim/gui/plot.py index 5ccbe9f17..d831bd2c5 100644 --- a/bioptim/gui/plot.py +++ b/bioptim/gui/plot.py @@ -393,11 +393,22 @@ def legend_without_duplicate_labels(ax): # No graph was setup in problem_type return - y_min_all = [None for _ in self.variable_sizes[0]] - y_max_all = [None for _ in self.variable_sizes[0]] + all_keys_across_phases = [] + for variable_sizes in self.variable_sizes: + keys_not_in_previous_phases = [ + key for key in list(variable_sizes.keys()) if key not in all_keys_across_phases + ] + all_keys_across_phases += keys_not_in_previous_phases + + y_min_all = [None for _ in all_keys_across_phases] + y_max_all = [None for _ in all_keys_across_phases] + self.custom_plots = {} for i, nlp in enumerate(self.ocp.nlp): + for var_idx, variable in enumerate(self.variable_sizes[i]): + y_range_var_idx = all_keys_across_phases.index(variable) + if nlp.plot[variable].combine_to: self.axes[variable] = self.axes[nlp.plot[variable].combine_to] axes = self.axes[variable][1] @@ -425,9 +436,10 @@ def legend_without_duplicate_labels(ax): 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]: - y_min_all[var_idx] = [np.inf] * nb_subplots - y_max_all[var_idx] = [-np.inf] * nb_subplots + + if not y_min_all[y_range_var_idx]: + y_min_all[y_range_var_idx] = [np.inf] * nb_subplots + y_max_all[y_range_var_idx] = [-np.inf] * nb_subplots if variable not in self.custom_plots: self.custom_plots[variable] = [ @@ -476,14 +488,16 @@ def legend_without_duplicate_labels(ax): for j in range(nlp.ns * repeat) ] ) - if y_min.__array__()[0] < y_min_all[var_idx][mapping_to_first_index.index(ctr)]: - y_min_all[var_idx][mapping_to_first_index.index(ctr)] = y_min - if y_max.__array__()[0] > y_max_all[var_idx][mapping_to_first_index.index(ctr)]: - y_max_all[var_idx][mapping_to_first_index.index(ctr)] = y_max + + if y_min.__array__()[0] < y_min_all[y_range_var_idx][mapping_to_first_index.index(ctr)]: + y_min_all[y_range_var_idx][mapping_to_first_index.index(ctr)] = y_min + + if y_max.__array__()[0] > y_max_all[y_range_var_idx][mapping_to_first_index.index(ctr)]: + y_max_all[y_range_var_idx][mapping_to_first_index.index(ctr)] = y_max y_range, _ = self.__compute_ylim( - y_min_all[var_idx][mapping_to_first_index.index(ctr)], - y_max_all[var_idx][mapping_to_first_index.index(ctr)], + y_min_all[y_range_var_idx][mapping_to_first_index.index(ctr)], + y_max_all[y_range_var_idx][mapping_to_first_index.index(ctr)], 1.25, ) ax.set_ylim(y_range) @@ -973,6 +987,7 @@ def __update_axes(self): def __compute_ylim(min_val: np.ndarray | DM, max_val: np.ndarray | DM, factor: float) -> tuple: """ Dynamically find the ylim + Parameters ---------- min_val: np.ndarray | DM diff --git a/bioptim/optimization/optimal_control_program.py b/bioptim/optimization/optimal_control_program.py index 78679e872..549a30cd9 100644 --- a/bioptim/optimization/optimal_control_program.py +++ b/bioptim/optimization/optimal_control_program.py @@ -13,8 +13,8 @@ from ..dynamics.ode_solver import OdeSolver, OdeSolverBase from ..gui.check_conditioning import check_conditioning from ..gui.graph import OcpToConsole, OcpToGraph -from ..gui.plot import CustomPlot, PlotOcp from ..gui.ipopt_output_plot import SaveIterationsInfo +from ..gui.plot import CustomPlot, PlotOcp from ..interfaces import Solver from ..interfaces.abstract_options import GenericSolver from ..limits.constraints import (