Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

holonomic graph multiphase multi all combinations #885

Merged
merged 7 commits into from
Jul 19, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
45 changes: 33 additions & 12 deletions bioptim/dynamics/configure_problem.py
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -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],
Expand All @@ -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
----------
Expand Down Expand Up @@ -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],
Expand All @@ -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
----------
Expand Down Expand Up @@ -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],
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down
37 changes: 26 additions & 11 deletions bioptim/gui/plot.py
Original file line number Diff line number Diff line change
Expand Up @@ -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]
Expand Down Expand Up @@ -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] = [
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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
Expand Down
2 changes: 1 addition & 1 deletion bioptim/optimization/optimal_control_program.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 (
Expand Down
Loading