Merge pull request #885 from Ipuch/copyholo

holonomic graph multiphase multi all combinations

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,

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

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 (