Skip to content

Commit

Permalink
Merge remote-tracking branch 'pyomeca/master' into Cholesky_for_impli…
Browse files Browse the repository at this point in the history 
…cit_stochastic
  • Loading branch information
@EveCharbie
EveCharbie committed Jul 24, 2023
2 parents 691a73a + b41eaa9 commit 9cb1519
Show file tree
Hide file tree
Showing 3 changed files with 74 additions and 25 deletions.
6 changes: 5 additions & 1 deletion bioptim/examples/getting_started/example_multinode_objective.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@
BiorbdModel,
PenaltyController,
MultinodeObjectiveList,
CostType,
)


Expand Down Expand Up @@ -116,12 +117,15 @@ def main():
# --- Prepare the ocp --- #
n_shooting = 30
ocp = prepare_ocp(biorbd_model_path="models/pendulum.bioMod", final_time=1, n_shooting=n_shooting)
ocp.add_plot_penalty(CostType.ALL)

# --- Solve the ocp --- #
sol = ocp.solve(Solver.IPOPT(show_online_optim=False)) # show_online_optim=platform.system() == "Linux"
sol = ocp.solve(Solver.IPOPT(show_online_optim=platform.system() == "Linux"))

# --- Show the results in a bioviz animation --- #
sol.animate(n_frames=100)
# sol.graphs()
# sol.print_cost()


if __name__ == "__main__":
Expand Down
63 changes: 46 additions & 17 deletions bioptim/gui/plot.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -345,14 +345,38 @@ def legend_without_duplicate_labels(ax):
nlp.controls.node_index = node_index
nlp.stochastic_variables.node_index = node_index

# If multi-node penalties = None, stays zero
size_x = 0
size_u = 0
size_p = 0
size_s = 0
if "penalty" in nlp.plot[key].parameters:
casadi_function = nlp.plot[key].parameters["penalty"].weighted_function_non_threaded[0]
if nlp.plot[key].parameters["penalty"].multinode_penalty:
if casadi_function is not None:
size_x = len(casadi_function.nominal_in(0))
size_u = len(casadi_function.nominal_in(1))
size_p = len(casadi_function.nominal_in(2))
size_s = len(casadi_function.nominal_in(3))
else:
size_x = nlp.states.shape
size_u = nlp.controls.shape
size_p = nlp.parameters.shape
size_s = nlp.stochastic_variables.shape
else:
size_x = nlp.states.shape
size_u = nlp.controls.shape
size_p = nlp.parameters.shape
size_s = nlp.stochastic_variables.shape

size = (
nlp.plot[key]
.function(
node_index,
np.zeros((nlp.states.shape, 2)),
np.zeros((nlp.controls.shape, 2)),
np.zeros((nlp.parameters.shape, 2)),
np.zeros((nlp.stochastic_variables.shape, 2)),
np.zeros((size_x, 1)),
np.zeros((size_u, 1)),
np.zeros((size_p, 1)),
np.zeros((size_s, 1)),
**nlp.plot[key].parameters,
)
.shape[0]
Expand Down Expand Up @@ -786,6 +810,8 @@ def update_data(self, v: dict):
penalty: MultinodeConstraint = self.plot_func[key][i].parameters["penalty"]

x_phase = np.ndarray((0, len(penalty.nodes_phase)))
if sol.ocp.n_phases == 1 and isinstance(data_states, dict):
data_states = [data_states]
for state_key in data_states[i].keys():
x_phase_tp = np.ndarray((data_states[i][state_key].shape[0], 0))
for tp in range(len(penalty.nodes_phase)):
Expand All @@ -799,23 +825,28 @@ def update_data(self, v: dict):
)
x_phase = np.vstack((x_phase, x_phase_tp))

u_phase = np.ndarray((0, len(penalty.nodes_phase)))
u_phase = None
if sol.ocp.n_phases == 1 and isinstance(data_controls, dict):
data_controls = [data_controls]
for control_key in data_controls[i].keys():
u_phase_tp = np.ndarray((data_controls[i][control_key].shape[0], 0))
u_phase_tp = None
for tp in range(len(penalty.nodes_phase)):
phase_tp = penalty.nodes_phase[tp]
node_idx_tp = penalty.all_nodes_index[tp]
u_phase_tp = np.hstack(
(
u_phase_tp,
data_controls[phase_tp][control_key][
:, node_idx_tp - (1 if node_idx_tp == nlp.ns else 0)
][:, np.newaxis],
if node_idx_tp != nlp.ns or (
node_idx_tp == nlp.ns and not np.isnan(data_controls[i][control_key][0, -1])
):
new_value = data_controls[phase_tp][control_key][:, node_idx_tp][:, np.newaxis]
u_phase_tp = (
np.vstack((u_phase_tp, new_value))
if isinstance(u_phase_tp, np.ndarray)
else new_value
)
)
u_phase = np.vstack((u_phase, u_phase_tp))
u_phase = u_phase_tp if u_phase is None else np.hstack((u_phase, u_phase_tp))

s_phase = np.ndarray((0, len(penalty.nodes_phase)))
if sol.ocp.n_phases == 1 and isinstance(data_stochastic, dict):
data_stochastic = [data_stochastic]
for stochastic_key in data_stochastic[i].keys():
s_phase_tp = np.ndarray((data_stochastic[i][stochastic_key].shape[0], 0))
for tp in range(len(penalty.nodes_phase)):
Expand All @@ -824,9 +855,7 @@ def update_data(self, v: dict):
s_phase_tp = np.hstack(
(
s_phase_tp,
data_stochastic[phase_tp][stochastic_key][
:, node_idx_tp - (1 if node_idx_tp == nlp.ns else 0)
][:, np.newaxis],
data_stochastic[phase_tp][stochastic_key][:, node_idx_tp][:, np.newaxis],
)
)
s_phase = np.vstack((s_phase, s_phase_tp))
Expand Down
30 changes: 23 additions & 7 deletions bioptim/optimization/optimal_control_program.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1327,14 +1327,20 @@ def compute_penalty_values(t, x, u, p, s, penalty, dt: int | Callable):
if len(s.shape) < 2:
s = s.reshape((-1, 1))

if penalty.multinode_penalty:
# TODO: Fix the scaling of multi_node_penalty (This is a hack, it should be computed at each phase)
penalty_phase = penalty.nodes_phase[0]
else:
penalty_phase = penalty.phase

# if time is parameter of the ocp, we need to evaluate with current parameters
if isinstance(dt, Function):
dt = dt(p)
# The division is to account for the steps in the integration. The else is for Mayer term
dt = dt / (x.shape[1] - 1) if x.shape[1] > 1 else dt
if not isinstance(penalty.dt, (float, int)):
if dt.shape[0] > 1:
dt = dt[penalty.phase]
dt = dt[penalty_phase]

_target = (
np.hstack([p[..., penalty.node_idx.index(t)] for p in penalty.target])
Expand All @@ -1344,28 +1350,38 @@ def compute_penalty_values(t, x, u, p, s, penalty, dt: int | Callable):

x_scaling = np.concatenate(
[
np.repeat(self.nlp[penalty.phase].x_scaling[key].scaling[:, np.newaxis], x.shape[1], axis=1)
for key in self.nlp[penalty.phase].states
np.repeat(self.nlp[penalty_phase].x_scaling[key].scaling[:, np.newaxis], x.shape[1], axis=1)
for key in self.nlp[penalty_phase].states
]
)
if penalty.multinode_penalty:
len_x = sum(self.nlp[penalty_phase].states[key].shape for key in self.nlp[penalty_phase].states)
complete_scaling = np.array(x_scaling)
number_of_repeat = x.shape[0] // len_x
x_scaling = np.repeat(complete_scaling, number_of_repeat, axis=0)
x /= x_scaling

if u.size != 0:
u_scaling = np.concatenate(
[
np.repeat(self.nlp[penalty.phase].u_scaling[key].scaling[:, np.newaxis], u.shape[1], axis=1)
for key in self.nlp[penalty.phase].controls
np.repeat(self.nlp[penalty_phase].u_scaling[key].scaling[:, np.newaxis], u.shape[1], axis=1)
for key in self.nlp[penalty_phase].controls
]
)
if penalty.multinode_penalty:
len_u = sum(self.nlp[penalty_phase].controls[key].shape for key in self.nlp[penalty_phase].controls)
complete_scaling = np.array(u_scaling)
number_of_repeat = u.shape[0] // len_u
u_scaling = np.repeat(complete_scaling, number_of_repeat, axis=0)
u /= u_scaling

out = []
if penalty.transition or penalty.multinode_penalty:
out.append(
penalty.weighted_function_non_threaded[t](
x.reshape((-1, 1)), u.reshape((-1, 1)), p, s, penalty.weight, _target, dt
x.reshape((-1, 1)), u.reshape((-1, 1)), p, s.reshape((-1, 1)), penalty.weight, _target, 1
)
)
) # dt=1 because multinode penalties behave like Mayer functions

elif penalty.derivative or penalty.explicit_derivative:
if not np.all(
Expand Down

0 comments on commit 9cb1519

Please sign in to comment.