Skip to content

Commit

Permalink
Merge pull request #814 from pariterre/master
Browse files Browse the repository at this point in the history 
Major refactor towards 3.2.0
  • Loading branch information
@pariterre
pariterre committed Dec 22, 2023
2 parents aca4e6d + 5715e1e commit 31b8f4a
Show file tree
Hide file tree
Showing 143 changed files with 7,982 additions and 10,346 deletions.
7 changes: 7 additions & 0 deletions Notes
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,7 @@
currently broken:
- ACADOS
- Linear_continous control
- No control
- Collocation and objective NODE.ALL
- Time vector in collocation (and integration)
- CVODES
1 change: 1 addition & 0 deletions bioptim/__init__.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -218,6 +218,7 @@
)
from .optimization.parameters import ParameterList
from .optimization.solution.solution import Solution
from .optimization.solution.solution_data import SolutionMerge, TimeAlignment
from .optimization.optimization_variable import OptimizationVariableList
from .optimization.variable_scaling import VariableScalingList, VariableScaling
from .optimization.variational_optimal_control_program import VariationalOptimalControlProgram
Expand Down
129 changes: 72 additions & 57 deletions bioptim/dynamics/configure_new_variable.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,15 +2,14 @@
import numpy as np

from .fatigue.fatigue_dynamics import FatigueList, MultiFatigueInterface
from .ode_solver import OdeSolver
from ..gui.plot import CustomPlot
from ..limits.path_conditions import Bounds
from ..misc.enums import PlotType, ControlType, VariableType, PhaseDynamics
from ..misc.mapping import BiMapping


def variable_type_from_booleans_to_enums(
as_states: bool, as_controls: bool, as_states_dot: bool, as_stochastic: bool
as_states: bool, as_controls: bool, as_states_dot: bool, as_algebraic_states: bool
) -> list[VariableType]:
"""
Convert the booleans to enums
Expand All @@ -23,8 +22,8 @@ def variable_type_from_booleans_to_enums(
If the new variable should be added to the state_dot variable set
as_controls: bool
If the new variable should be added to the control variable set
as_stochastic: bool
If the new variable should be added to the stochastic variable set
as_algebraic_states: bool
If the new variable should be added to the algebraic states variable set
Returns
-------
Expand All @@ -38,21 +37,21 @@ def variable_type_from_booleans_to_enums(
variable_type.append(VariableType.STATES_DOT)
if as_controls:
variable_type.append(VariableType.CONTROLS)
if as_stochastic:
variable_type.append(VariableType.STOCHASTIC)
if as_algebraic_states:
variable_type.append(VariableType.ALGEBRAIC_STATES)
return variable_type


class NewVariableConfiguration:
# todo: add a way to remove the if as_states, as_controls, as_states_dot, as_stochastic, etc...
# todo: add a way to remove the if as_states, as_controls, as_states_dot, as_algebraic_states, etc...
# if we want to remove ocp, nlp, it
# should be a method of ocp, and specify the phase_idx where the variable is added
# ocp.configure_new_variable(
# phase_idx,
# name,
# name_elements,
# variable_type=variable_types,
# # VariableType.CONTROL, VariableType.STATE_DOT, VariableType.STOCHASTIC, VariableType.ALGEBRAIC_STATE
# # VariableType.CONTROL, VariableType.STATE_DOT, VariableType.ALGEBRAIC_STATE
# )
def __init__(
self,
Expand All @@ -63,7 +62,7 @@ def __init__(
as_states: bool,
as_controls: bool,
as_states_dot: bool = False,
as_stochastic: bool = False,
as_algebraic_states: bool = False,
fatigue: FatigueList = None,
combine_name: str = None,
combine_state_control_plot: bool = False,
Expand All @@ -89,8 +88,8 @@ def __init__(
If the new variable should be added to the state_dot variable set
as_controls: bool
If the new variable should be added to the control variable set
as_stochastic: bool
If the new variable should be added to the stochastic variable set
as_algebraic_states: bool
If the new variable should be added to the algebraic states variable set
fatigue: FatigueList
The list of fatigable item
combine_name: str
Expand All @@ -110,7 +109,7 @@ def __init__(
self.as_states = as_states
self.as_controls = as_controls
self.as_states_dot = as_states_dot
self.as_stochastic = as_stochastic
self.as_algebraic_states = as_algebraic_states
self.fatigue = fatigue
self.combine_name = combine_name
self.combine_state_control_plot = combine_state_control_plot
Expand All @@ -124,7 +123,7 @@ def __init__(
self.mx_states = None
self.mx_states_dot = None
self.mx_controls = None
self.mx_stochastic = None
self.mx_algebraic_states = None

self._check_combine_state_control_plot()

Expand Down Expand Up @@ -189,6 +188,10 @@ def _declare_phase_copy_booleans(self):
nlp, phase_idx, self.nlp.use_states_dot_from_phase_idx, self.name, "states_dot"
)

self.copy_algebraic_states = self.check_variable_copy_condition(
nlp, phase_idx, self.nlp.use_states_from_phase_idx, self.name, "algebraic_states"
)

@staticmethod
def check_variable_copy_condition(
nlp, phase_idx: int, use_from_phase_idx: int, name: str, decision_variable_attribute: str
Expand Down Expand Up @@ -302,8 +305,8 @@ def _declare_initial_guess(self):
self.name, initial_guess=np.zeros(len(self.nlp.variable_mappings[self.name].to_first.map_idx))
)

if self.as_stochastic and self.name not in self.nlp.s_init:
self.nlp.s_init.add(
if self.as_algebraic_states and self.name not in self.nlp.a_init:
self.nlp.a_init.add(
self.name, initial_guess=np.zeros(len(self.nlp.variable_mappings[self.name].to_first.map_idx))
)

Expand All @@ -320,8 +323,8 @@ def _declare_variable_scaling(self):
self.nlp.u_scaling.add(
self.name, scaling=np.ones(len(self.nlp.variable_mappings[self.name].to_first.map_idx))
)
if self.as_stochastic and self.name not in self.nlp.s_scaling:
self.nlp.s_scaling.add(
if self.as_algebraic_states and self.name not in self.nlp.a_scaling:
self.nlp.a_scaling.add(
self.name, scaling=np.ones(len(self.nlp.variable_mappings[self.name].to_first.map_idx))
)

Expand All @@ -341,7 +344,11 @@ def _use_copy(self):
if not self.copy_controls
else [self.ocp.nlp[self.nlp.use_controls_from_phase_idx].controls[0][self.name].mx]
)
self.mx_stochastic = []
self.mx_algebraic_states = (
[]
if not self.copy_algebraic_states
else [self.ocp.nlp[self.nlp.use_states_from_phase_idx].algebraic_states[0][self.name].mx]
)

# todo: if mapping on variables, what do we do with mapping on the nodes
for i in self.nlp.variable_mappings[self.name].to_second.map_idx:
Expand All @@ -360,12 +367,12 @@ def _use_copy(self):
if not self.copy_controls:
self.mx_controls.append(MX.sym(var_name, 1, 1))

self.mx_stochastic.append(MX.sym(var_name, 1, 1))
self.mx_algebraic_states.append(MX.sym(var_name, 1, 1))

self.mx_states = vertcat(*self.mx_states)
self.mx_states_dot = vertcat(*self.mx_states_dot)
self.mx_controls = vertcat(*self.mx_controls)
self.mx_stochastic = vertcat(*self.mx_stochastic)
self.mx_algebraic_states = vertcat(*self.mx_algebraic_states)

def _declare_auto_axes_idx(self):
"""Declare the axes index if not already declared"""
Expand All @@ -388,9 +395,9 @@ def _declare_legend(self):
def _declare_cx_and_plot(self):
if self.as_states:
for node_index in range(
(0 if self.nlp.phase_dynamics == PhaseDynamics.SHARED_DURING_THE_PHASE else self.nlp.ns) + 1
self.nlp.n_states_nodes if self.nlp.phase_dynamics == PhaseDynamics.ONE_PER_NODE else 1
):
n_cx = self.nlp.ode_solver.n_cx if isinstance(self.nlp.ode_solver, OdeSolver.COLLOCATION) else 3
n_cx = self.nlp.ode_solver.n_required_cx + 2
cx_scaled = (
self.ocp.nlp[self.nlp.use_states_from_phase_idx].states[node_index][self.name].original_cx
if self.copy_states
Expand All @@ -411,7 +418,9 @@ def _declare_cx_and_plot(self):
)
if not self.skip_plot:
self.nlp.plot[f"{self.name}_states"] = CustomPlot(
lambda t, x, u, p, s: x[self.nlp.states[self.name].index, :],
lambda t0, phases_dt, node_idx, x, u, p, a: x[self.nlp.states.key_index(self.name), :]
if x.any()
else np.ndarray((cx[0][0].shape[0], 1)) * np.nan,
plot_type=PlotType.INTEGRATED,
axes_idx=self.axes_idx,
legend=self.legend,
Expand All @@ -420,19 +429,7 @@ def _declare_cx_and_plot(self):

if self.as_controls:
for node_index in range(
(
1
if self.nlp.phase_dynamics == PhaseDynamics.SHARED_DURING_THE_PHASE
else (
self.nlp.ns
+ +(
1
if self.nlp.control_type
in (ControlType.LINEAR_CONTINUOUS, ControlType.CONSTANT_WITH_LAST_NODE)
else 0
)
)
)
self.nlp.n_controls_nodes if self.nlp.phase_dynamics == PhaseDynamics.ONE_PER_NODE else 1
):
cx_scaled = (
self.ocp.nlp[self.nlp.use_controls_from_phase_idx].controls[node_index][self.name].original_cx
Expand All @@ -456,7 +453,9 @@ def _declare_cx_and_plot(self):
plot_type = PlotType.PLOT if self.nlp.control_type == ControlType.LINEAR_CONTINUOUS else PlotType.STEP
if not self.skip_plot:
self.nlp.plot[f"{self.name}_controls"] = CustomPlot(
lambda t, x, u, p, s: u[self.nlp.controls[self.name].index, :],
lambda t0, phases_dt, node_idx, x, u, p, a: u[self.nlp.controls.key_index(self.name), :]
if u.any()
else np.ndarray((cx[0][0].shape[0], 1)) * np.nan,
plot_type=plot_type,
axes_idx=self.axes_idx,
legend=self.legend,
Expand All @@ -467,15 +466,9 @@ def _declare_cx_and_plot(self):

if self.as_states_dot:
for node_index in range(
(0 if self.nlp.phase_dynamics == PhaseDynamics.SHARED_DURING_THE_PHASE else self.nlp.ns) + 1
self.nlp.n_states_nodes if self.nlp.phase_dynamics == PhaseDynamics.ONE_PER_NODE else 1
):
n_cx = (
self.nlp.ode_solver.polynomial_degree + 1
if isinstance(self.nlp.ode_solver, OdeSolver.COLLOCATION)
else 3
)
if n_cx < 3:
n_cx = 3
n_cx = self.nlp.ode_solver.n_required_cx + 2
cx_scaled = (
self.ocp.nlp[self.nlp.use_states_dot_from_phase_idx].states_dot[node_index][self.name].original_cx
if self.copy_states_dot
Expand All @@ -495,19 +488,27 @@ def _declare_cx_and_plot(self):
node_index,
)

if self.as_stochastic:
if self.as_algebraic_states:
for node_index in range(
(0 if self.nlp.phase_dynamics == PhaseDynamics.SHARED_DURING_THE_PHASE else self.nlp.ns) + 1
self.nlp.n_states_nodes if self.nlp.phase_dynamics == PhaseDynamics.ONE_PER_NODE else 1
):
n_cx = 3
cx_scaled = self.define_cx_scaled(n_col=n_cx, n_shooting=1, initial_node=node_index)
cx = self.define_cx_unscaled(cx_scaled, self.nlp.s_scaling[self.name].scaling)
n_cx = 2
cx_scaled = (
self.ocp.nlp[self.nlp.use_states_from_phase_idx].algebraic_states[node_index][self.name].original_cx
if self.copy_algebraic_states
else self.define_cx_scaled(n_col=n_cx, n_shooting=0, initial_node=node_index)
)
cx = (
self.ocp.nlp[self.nlp.use_states_from_phase_idx].algebraic_states[node_index][self.name].original_cx
if self.copy_algebraic_states
else self.define_cx_unscaled(cx_scaled, self.nlp.a_scaling[self.name].scaling)
)

self.nlp.stochastic_variables.append(
self.nlp.algebraic_states.append(
self.name,
cx[0],
cx_scaled[0],
self.mx_stochastic,
self.mx_states,
self.nlp.variable_mappings[self.name],
node_index,
)
Expand Down Expand Up @@ -570,14 +571,22 @@ def _manage_fatigue_to_new_variable(
legend = [f"{name}_{i}" for i in name_elements]
fatigue_plot_name = f"fatigue_{name}"
nlp.plot[fatigue_plot_name] = CustomPlot(
lambda t, x, u, p, s: x[:n_elements, :] * np.nan,
lambda t0, phases_dt, node_idx, x, u, p, a: (
x[:n_elements, :] if x.any() else np.ndarray((len(name_elements), 1))
)
* np.nan,
plot_type=PlotType.INTEGRATED,
legend=legend,
bounds=Bounds(None, -1, 1),
)
control_plot_name = f"{name}_controls" if not multi_interface and split_controls else f"{name}"
nlp.plot[control_plot_name] = CustomPlot(
lambda t, x, u, p, s: u[:n_elements, :] * np.nan, plot_type=PlotType.STEP, legend=legend
lambda t0, phases_dt, node_idx, x, u, p, a: (
u[:n_elements, :] if u.any() else np.ndarray((len(name_elements), 1))
)
* np.nan,
plot_type=PlotType.STEP,
legend=legend,
)

var_names_with_suffix = []
Expand All @@ -592,7 +601,9 @@ def _manage_fatigue_to_new_variable(
var_names_with_suffix[-1], name_elements, ocp, nlp, as_states, as_controls, skip_plot=True
)
nlp.plot[f"{var_names_with_suffix[-1]}_controls"] = CustomPlot(
lambda t, x, u, p, s, key: u[nlp.controls[key].index, :],
lambda t0, phases_dt, node_idx, x, u, p, a, key: u[nlp.controls.key_index(key), :]
if u.any()
else np.ndarray((len(name_elements), 1)) * np.nan,
plot_type=PlotType.STEP,
combine_to=control_plot_name,
key=var_names_with_suffix[-1],
Expand All @@ -601,7 +612,9 @@ def _manage_fatigue_to_new_variable(
elif i == 0:
NewVariableConfiguration(f"{name}", name_elements, ocp, nlp, as_states, as_controls, skip_plot=True)
nlp.plot[f"{name}_controls"] = CustomPlot(
lambda t, x, u, p, s, key: u[nlp.controls[key].index, :],
lambda t0, phases_dt, node_idx, x, u, p, a, key: u[nlp.controls.key_index(key), :]
if u.any()
else np.ndarray((len(name_elements), 1)) * np.nan,
plot_type=PlotType.STEP,
combine_to=control_plot_name,
key=f"{name}",
Expand All @@ -612,7 +625,9 @@ def _manage_fatigue_to_new_variable(
name_tp = f"{var_names_with_suffix[-1]}_{params}"
NewVariableConfiguration(name_tp, name_elements, ocp, nlp, True, False, skip_plot=True)
nlp.plot[name_tp] = CustomPlot(
lambda t, x, u, p, s, key, mod: mod * x[nlp.states[key].index, :],
lambda t0, phases_dt, node_idx, x, u, p, a, key, mod: mod * x[nlp.states.key_index(key), :]
if x.any()
else np.ndarray((len(name_elements), 1)) * np.nan,
plot_type=PlotType.INTEGRATED,
combine_to=fatigue_plot_name,
key=name_tp,
Expand Down
Loading

0 comments on commit 31b8f4a

Please sign in to comment.