diff --git a/bioptim/__init__.py b/bioptim/__init__.py index 7ccf9b7c2..7bea1e999 100644 --- a/bioptim/__init__.py +++ b/bioptim/__init__.py @@ -246,6 +246,9 @@ CyclicNonlinearModelPredictiveControl, CyclicMovingHorizonEstimator, MultiCyclicNonlinearModelPredictiveControl, + RecedingHorizonFailurePolicy, + RecedingHorizonWindowResult, + CyclicVariableShift, ) from .optimization.receding_horizon_optimization import MovingHorizonEstimator, NonlinearModelPredictiveControl from .optimization.solution.solution import Solution diff --git a/bioptim/interfaces/acados_interface.py b/bioptim/interfaces/acados_interface.py index a7517996b..a09cec047 100644 --- a/bioptim/interfaces/acados_interface.py +++ b/bioptim/interfaces/acados_interface.py @@ -1,5 +1,6 @@ from time import perf_counter from datetime import datetime +from typing import TYPE_CHECKING import numpy as np from scipy import linalg @@ -13,6 +14,9 @@ from ..limits.path_conditions import Bounds from ..misc.enums import InterpolationType +if TYPE_CHECKING: + from ..optimization.solution.solution import Solution + from ..misc.parameters_types import ( Str, @@ -20,6 +24,14 @@ AnyListorDict, ) +ACADOS_STATUS_LABELS = { + 0: "success", + 1: "maximum_iterations_reached", + 2: "minimum_step_reached", + 3: "qp_solver_failure", + 4: "ready", +} + class AcadosInterface(SolverInterface): """ @@ -268,9 +280,19 @@ def __set_constraints(self, ocp) -> None: ) for key in ocp.nlp[0].controls.keys(): - if not np.all(np.all(ocp.nlp[0].u_bounds[key].min.T == ocp.nlp[0].u_bounds[key].min.T[0, :], axis=0)): + if not np.all( + np.all( + ocp.nlp[0].u_bounds[key].min.T == ocp.nlp[0].u_bounds[key].min.T[0, :], + axis=0, + ) + ): raise NotImplementedError("u_bounds min must be the same at each shooting point with ACADOS") - if not np.all(np.all(ocp.nlp[0].u_bounds[key].max.T == ocp.nlp[0].u_bounds[key].max.T[0, :], axis=0)): + if not np.all( + np.all( + ocp.nlp[0].u_bounds[key].max.T == ocp.nlp[0].u_bounds[key].max.T[0, :], + axis=0, + ) + ): raise NotImplementedError("u_bounds max must be the same at each shooting point with ACADOS") if ( @@ -478,7 +500,11 @@ def _adjust_dim(): v_var[rows] = 1.0 return v_var, rows - if objectives.node[0] not in [Node.INTERMEDIATES, Node.PENULTIMATE, Node.END]: + if objectives.node[0] not in [ + Node.INTERMEDIATES, + Node.PENULTIMATE, + Node.END, + ]: v_var, rows = _adjust_dim() if is_state: acados.Vx0 = np.vstack((acados.Vx0, np.diag(v_var))) @@ -517,7 +543,8 @@ def add_nonlinear_ls_lagrange(acados, objectives, t, dt, x, u, p, a, d): u = vertcat(u, u) acados.lagrange_costs = vertcat( - acados.lagrange_costs, objectives.function[0](t, dt, x, u, p, a, d).reshape((-1, 1)) + acados.lagrange_costs, + objectives.function[0](t, dt, x, u, p, a, d).reshape((-1, 1)), ) acados.W = linalg.block_diag( acados.W, np.diag(objectives.weight.evaluate_at(0, objectives.function[0].numel_out())) @@ -530,7 +557,11 @@ def add_nonlinear_ls_lagrange(acados, objectives, t, dt, x, u, p, a, d): acados.y_ref.append([np.zeros((objectives.function[0].numel_out(), 1)) for _ in node_idx]) def add_nonlinear_ls_mayer(acados, objectives, t, dt, x, u, p, a, d, node=None): - if objectives.node[0] not in [Node.INTERMEDIATES, Node.PENULTIMATE, Node.END]: + if objectives.node[0] not in [ + Node.INTERMEDIATES, + Node.PENULTIMATE, + Node.END, + ]: acados.W_0 = linalg.block_diag( acados.W_0, np.diag(objectives.weight.evaluate_at(0, objectives.function[0].numel_out())) ) @@ -546,7 +577,8 @@ def add_nonlinear_ls_mayer(acados, objectives, t, dt, x, u, p, a, d, node=None): u_tp = u_tp if objectives.function[0].size_in("u") != (0, 0) else [] acados.mayer_costs = vertcat( - acados.mayer_costs, objectives.function[0](t, dt, x_tp, u_tp, p, a, d).reshape((-1, 1)) + acados.mayer_costs, + objectives.function[0](t, dt, x_tp, u_tp, p, a, d).reshape((-1, 1)), ) if objectives.target is not None: @@ -569,7 +601,8 @@ def add_nonlinear_ls_mayer(acados, objectives, t, dt, x, u, p, a, d, node=None): u_tp = u_tp if objectives.function[-1].size_in("u") != (0, 0) else [] acados.mayer_costs_e = vertcat( - acados.mayer_costs_e, objectives.function[-1](t, dt, x_tp, u_tp, p, a, d).reshape((-1, 1)) + acados.mayer_costs_e, + objectives.function[-1](t, dt, x_tp, u_tp, p, a, d).reshape((-1, 1)), ) if objectives.target is not None: @@ -590,7 +623,10 @@ def add_nonlinear_ls_mayer(acados, objectives, t, dt, x, u, p, a, d, node=None): self.W_e = np.zeros((0, 0)) self.W_0 = np.zeros((0, 0)) allowed_control_objectives = [ObjectiveFcn.Lagrange.MINIMIZE_CONTROL] - allowed_state_objectives = [ObjectiveFcn.Lagrange.MINIMIZE_STATE, ObjectiveFcn.Mayer.TRACK_STATE] + allowed_state_objectives = [ + ObjectiveFcn.Lagrange.MINIMIZE_STATE, + ObjectiveFcn.Mayer.TRACK_STATE, + ] if self.acados_ocp.cost.cost_type == "LINEAR_LS": n_states = ocp.nlp[0].states.shape @@ -875,6 +911,7 @@ def get_optimized_value(self) -> AnyListorDict: "iter": self.ocp_solver.get_stats("sqp_iter"), "status": self.status, "solver": SolverType.ACADOS, + "solver_diagnostics": self.get_solver_diagnostics(), } out["x"] = vertcat(out["x"], acados_x.reshape(-1, 1, order="F")) @@ -893,6 +930,91 @@ def get_optimized_value(self) -> AnyListorDict: return out[0] if len(out) == 1 else out + def get_solver_diagnostics(self) -> dict: + """Return stable, best-effort diagnostics for the last Acados solve.""" + + diagnostics = { + "status": self.status, + "status_label": ACADOS_STATUS_LABELS.get(self.status, "unknown"), + "wall_time": self.real_time_to_optimize, + } + if self.ocp_solver is None: + return diagnostics + + stat_names = ( + "res_stat", + "res_eq", + "res_ineq", + "res_comp", + "sqp_iter", + "qp_iter", + "alpha", + "time_tot", + "statistics", + ) + for name in stat_names: + try: + diagnostics[name] = self.ocp_solver.get_stats(name) + except Exception: + # The available statistics depend on the acados_template version and solver configuration. + continue + diagnostics["iterates"] = self.get_iterates() + return diagnostics + + def get_iterates(self) -> list[dict]: + """Return the current Acados iterate without exposing acados_template internals.""" + + if self.ocp_solver is None: + raise RuntimeError("Acados must be solved before its iterates can be retrieved") + iterates = [] + for node in range(self.acados_ocp.dims.N + 1): + iterate = {"x": np.asarray(self.ocp_solver.get(node, "x")).copy()} + for field in ("u", "z", "pi", "lam", "sl", "su"): + if field in ("u", "pi") and node == self.acados_ocp.dims.N: + iterate[field] = None + continue + try: + iterate[field] = np.asarray(self.ocp_solver.get(node, field)).copy() + except Exception: + iterate[field] = None + iterates.append(iterate) + return iterates + + def set_iterates(self, iterates: list[dict], include_multipliers: bool = False) -> None: + """Initialize Acados from compatible iterates. + + Primal fields (``x``, ``u`` and ``z``) are always transferred. Multipliers (``pi``, ``lam``, ``sl`` and + ``su``) are transferred only when explicitly requested; their dimensions must already match the Acados OCP. + """ + + if self.ocp_solver is None: + raise RuntimeError("The Acados solver must be created before its iterates can be initialized") + if len(iterates) != self.acados_ocp.dims.N + 1: + raise ValueError(f"Expected {self.acados_ocp.dims.N + 1} Acados nodes, got {len(iterates)}") + + primal_fields = ("x", "u", "z") + multiplier_fields = ("pi", "lam", "sl", "su") if include_multipliers else () + for node, iterate in enumerate(iterates): + for field in primal_fields + multiplier_fields: + value = iterate.get(field) + if value is None: + continue + if field in ("u", "pi") and node == self.acados_ocp.dims.N: + raise ValueError(f"Acados field '{field}' is not defined at the terminal node") + self.ocp_solver.set(node, field, np.asarray(value)) + + def set_lagrange_multiplier(self, sol: "Solution") -> None: + """Transfer Acados multipliers only from a compatible Acados iterate.""" + + diagnostics = getattr(sol, "solver_diagnostics", None) or {} + iterates = diagnostics.get("iterates") + if iterates is None: + raise ValueError( + "IPOPT/NLP multipliers cannot be converted generically to Acados QP multipliers. " + "Use a primal-only warm start, or provide compatible Acados iterates." + ) + self.set_iterates(iterates, include_multipliers=True) + def solve(self, expand_during_shake_tree: Bool = False) -> AnyListorDict: """ Solve the prepared ocp diff --git a/bioptim/optimization/optimal_control_program.py b/bioptim/optimization/optimal_control_program.py index d9c03f1ab..d531ad926 100644 --- a/bioptim/optimization/optimal_control_program.py +++ b/bioptim/optimization/optimal_control_program.py @@ -1475,7 +1475,7 @@ def set_ocp_solver(self, solver: Solver) -> None: "You cannot change the solver once it has been set. Please create a new OptimalControlProgram." ) - def set_warm_start(self, sol: Solution) -> None: + def set_warm_start(self, sol: Solution, transfer_multipliers: Bool | None = None) -> None: """ Modify x and u initial guess based on a solution. @@ -1483,6 +1483,9 @@ def set_warm_start(self, sol: Solution) -> None: ---------- sol: Solution The solution to initiate the OCP from + transfer_multipliers: bool | None + Whether to transfer solver multipliers. By default they are reused only for IPOPT, whose NLP multiplier + layout is compatible. Acados QP multipliers require explicit opt-in and Acados-origin iterates. """ state = sol.decision_states(to_merge=SolutionMerge.NODES) @@ -1516,7 +1519,9 @@ def set_warm_start(self, sol: Solution) -> None: self.update_initial_guess(x_init=x_init_guess, u_init=u_init_guess, parameter_init=param_init_guess) - if self.ocp_solver: + if transfer_multipliers is None: + transfer_multipliers = self.ocp_solver is not None and self.ocp_solver.opts.type == SolverType.IPOPT + if self.ocp_solver and transfer_multipliers: self.ocp_solver.set_lagrange_multiplier(sol) self._is_warm_starting = True diff --git a/bioptim/optimization/receding_horizon_optimization.py b/bioptim/optimization/receding_horizon_optimization.py index a0dd3e4af..329b400fa 100644 --- a/bioptim/optimization/receding_horizon_optimization.py +++ b/bioptim/optimization/receding_horizon_optimization.py @@ -1,4 +1,6 @@ from copy import deepcopy +from dataclasses import dataclass +from enum import Enum from math import inf from time import perf_counter @@ -12,7 +14,13 @@ from ..limits.constraints import ConstraintFcn, ConstraintList from ..limits.objective_functions import ObjectiveFcn, ObjectiveList from ..limits.path_conditions import InitialGuessList -from ..misc.enums import SolverType, InterpolationType, MultiCyclicCycleSolutions, ControlType, OnlineOptim +from ..misc.enums import ( + SolverType, + InterpolationType, + MultiCyclicCycleSolutions, + ControlType, + OnlineOptim, +) from ..interfaces import Solver from ..interfaces.abstract_options import GenericSolver from ..models.protocols.biomodel import BioModel @@ -32,6 +40,38 @@ ) +class RecedingHorizonFailurePolicy(Enum): + """Behavior to adopt when a receding-horizon window does not converge.""" + + STOP = "stop" + CONTINUE_DIAGNOSTIC = "continue_diagnostic" + + +@dataclass(frozen=True) +class RecedingHorizonWindowResult: + """Outcome of one receding-horizon solve, independently of trajectory export.""" + + solution: Solution + solver_succeeded: bool + trajectory_available: bool + physically_acceptable: bool + exported: bool + + @property + def accepted(self) -> bool: + return self.solver_succeeded or self.physically_acceptable + + +@dataclass(frozen=True) +class CyclicVariableShift: + """Add ``turns * period`` to selected state rows when a cyclic window advances.""" + + key: str + indices: int | tuple[int, ...] | list[int] + period: float + turns: int = 1 + + class RecedingHorizonOptimization(OptimalControlProgram): """ The main class to define an MHE. This class prepares the full program and gives all @@ -106,6 +146,10 @@ def solve( max_consecutive_failing: Int = inf, update_function_extra_params: AnyDictOptional = None, get_all_iterations: Bool = False, + failure_policy: RecedingHorizonFailurePolicy | str = RecedingHorizonFailurePolicy.CONTINUE_DIAGNOSTIC, + window_evaluation_function: Callable | None = None, + before_window_solve: Callable | None = None, + after_window_solve: Callable | None = None, **advance_options, ) -> Solution | AnyTuple: """ @@ -138,6 +182,18 @@ def solve( Any parameters to pass to the update function get_all_iterations: bool If an extra output value that includes all the individual solution should be returned + failure_policy: RecedingHorizonFailurePolicy | str + Stop immediately on a solver failure or continue collecting diagnostics. Failed windows are never + exported or used to advance the horizon unless ``window_evaluation_function`` accepts them. + window_evaluation_function: Callable | None + Optional callback ``(rhe, window_index, solution) -> bool`` declaring a numerically failed window + physically acceptable. This does not change ``solver_succeeded`` or the solver status. + before_window_solve: Callable | None + Optional callback ``(rhe, window_index, previous_solution)`` called immediately before each solve. + Returning ``False`` stops before solving the window. + after_window_solve: Callable | None + Optional callback ``(rhe, window_index, window_result)`` called after each solve and export decision. + Returning ``False`` prevents the next window from being solved. advance_options: Any The extra options to pass to the advancing methods @@ -149,6 +205,19 @@ def solve( if len(self.nlp) != 1: raise NotImplementedError("MHE is only available for 1 phase program") + if isinstance(failure_policy, str): + failure_policy = RecedingHorizonFailurePolicy(failure_policy) + + for key in self.nlp[0].x_bounds.keys(): + if self.nlp[0].x_bounds[key].type not in ( + InterpolationType.CONSTANT, + InterpolationType.CONSTANT_WITH_FIRST_AND_LAST_DIFFERENT, + ): + raise NotImplementedError( + "The MHE is not implemented yet for x_bounds not being " + "CONSTANT or CONSTANT_WITH_FIRST_AND_LAST_DIFFERENT" + ) + sol = None states = [] controls = [] @@ -167,6 +236,8 @@ def solve( real_time = perf_counter() all_solutions = [] split_solutions = [] + window_results = [] + last_exported_solution = None consecutive_failing = 0 update_function_extra_params = {} if update_function_extra_params is None else update_function_extra_params @@ -175,12 +246,32 @@ def solve( update_function(self, self.total_optimization_run, sol, **update_function_extra_params) and consecutive_failing < max_consecutive_failing ): + if before_window_solve is not None and before_window_solve(self, self.total_optimization_run, sol) is False: + break sol = super(RecedingHorizonOptimization, self).solve( solver=solver_current, warm_start=warm_start, ) consecutive_failing = 0 if sol.status == 0 else consecutive_failing + 1 + solver_succeeded = sol.status == 0 + trajectory_available = sol.vector is not None + physically_acceptable = bool( + not solver_succeeded + and window_evaluation_function is not None + and window_evaluation_function(self, self.total_optimization_run, sol) + ) + export_window = trajectory_available and (solver_succeeded or physically_acceptable) + window_results.append( + RecedingHorizonWindowResult( + solution=sol, + solver_succeeded=solver_succeeded, + trajectory_available=trajectory_available, + physically_acceptable=physically_acceptable, + exported=export_window, + ) + ) + # Set the option for the next iteration if self.total_optimization_run == 0: # Update the solver if first and the rest are different @@ -200,26 +291,57 @@ def solve( real_time = perf_counter() # Reset timer to skip the compiling time (so skip the first call to solve) # Solve and save the current window of interest - _states, _controls, _parameters = self.export_data(sol) - states.append(_states) - controls.append(_controls) - parameters.append(_parameters) + if export_window: + _states, _controls, _parameters = self.export_data(sol) + states.append(_states) + controls.append(_controls) + parameters.append(_parameters) + last_exported_solution = sol # Solve and save the full window of the OCP if get_all_iterations: all_solutions.append(sol) # Update the initial frame bounds and initial guess - self.advance_window(sol, **advance_options) + if export_window: + self.advance_window(sol, **advance_options) self.total_optimization_run += 1 - states.append({key: sol.decision_states()[key][-1] for key in sol.decision_states().keys()}) + if ( + after_window_solve is not None + and after_window_solve(self, self.total_optimization_run - 1, window_results[-1]) is False + ): + break + if not solver_succeeded and failure_policy == RecedingHorizonFailurePolicy.STOP: + break + + if sol is None: + raise RuntimeError("No receding-horizon window was solved") + + if not states: + sol.window_results = window_results + return (sol, all_solutions, split_solutions) if get_all_iterations else sol + + states.append( + { + key: last_exported_solution.decision_states()[key][-1] + for key in last_exported_solution.decision_states().keys() + } + ) real_time = perf_counter() - real_time # Prepare the modified ocp that fits the solution dimension - dt = sol.t_span()[0][-1] + dt = last_exported_solution.t_span()[0][-1] + attempted_windows = self.total_optimization_run + self.total_optimization_run = len(controls) final_sol = self._initialize_solution(float(dt), states, controls, parameters) + self.total_optimization_run = attempted_windows final_sol.solver_time_to_optimize = total_time final_sol.real_time_to_optimize = real_time + final_sol.status = next( + (result.solution.status for result in window_results if not result.solver_succeeded), + 0, + ) + final_sol.window_results = window_results return (final_sol, all_solutions, split_solutions) if get_all_iterations else final_sol @@ -336,7 +458,10 @@ def advance_window_initial_guess_states(self, sol: Solution, **advance_options) if self.nlp[0].x_init[key].type != InterpolationType.EACH_FRAME: # Override the previous x_init self.nlp[0].x_init.add( - key, np.ndarray(states[key].shape), interpolation=InterpolationType.EACH_FRAME, phase=0 + key, + np.ndarray(states[key].shape), + interpolation=InterpolationType.EACH_FRAME, + phase=0, ) self.nlp[0].x_init[key].check_and_adjust_dimensions(len(self.nlp[0].states[key]), self.nlp[0].ns) @@ -370,7 +495,12 @@ def advance_window_initial_guess_parameters(self, sol: Solution, **advance_optio parameters = sol.parameters for key in parameters.keys(): # Override the previous param_init - self.parameter_init.add(key, parameters[key][:, None], interpolation=InterpolationType.CONSTANT, phase=0) + self.parameter_init.add( + key, + parameters[key][:, None], + interpolation=InterpolationType.CONSTANT, + phase=0, + ) return True def export_data(self, sol: Solution) -> AnyTuple: @@ -394,7 +524,10 @@ def export_data(self, sol: Solution) -> AnyTuple: frames = self.frame_to_export if frames.stop is not None and frames.stop == self.nlp[0].n_controls_nodes: - if self.nlp[0].control_type in (ControlType.CONSTANT, ControlType.CONSTANT_WITH_LAST_NODE): + if self.nlp[0].control_type in ( + ControlType.CONSTANT, + ControlType.CONSTANT_WITH_LAST_NODE, + ): frames = slice(frames.start, frames.stop - 1) for key in self.nlp[0].controls.keys(): controls[key] = merged_controls[key][:, frames] @@ -479,6 +612,7 @@ def solve( if not cyclic_options: cyclic_options = {} self._initialize_state_idx_to_cycle(cyclic_options) + self._initialize_cyclic_variable_shifts(cyclic_options) self._set_cyclic_bound() if solver.type == SolverType.IPOPT: @@ -502,7 +636,10 @@ def export_data(self, sol: Solution) -> AnyTuple: if frames.stop is not None and frames.stop != self.nlp[0].n_controls_nodes: # The "not" conditions are there because if they are true, super() already avec done it. # Otherwise since it is cyclic it should always be done anyway - if self.nlp[0].control_type in (ControlType.CONSTANT, ControlType.CONSTANT_WITH_LAST_NODE): + if self.nlp[0].control_type in ( + ControlType.CONSTANT, + ControlType.CONSTANT_WITH_LAST_NODE, + ): frames = slice(self.frame_to_export.start, self.frame_to_export.stop - 1) for key in self.nlp[0].controls.keys(): @@ -515,7 +652,10 @@ def _initialize_solution(self, dt: Float, states: AnyList, controls: AnyList, pa for key in self.nlp[0].states.keys(): x_init.add( key, - np.concatenate([state[key][:, :-1] for state in states] + [states[-1][key][:, -1:]], axis=1), + np.concatenate( + [state[key][:, :-1] for state in states] + [states[-1][key][:, -1:]], + axis=1, + ), interpolation=InterpolationType.EACH_FRAME, phase=0, ) @@ -565,6 +705,29 @@ def _initialize_state_idx_to_cycle(self, options: AnyDict) -> None: states = self.nlp[0].states self.state_idx_to_cycle = {key: range(len(states[key])) for key in options["states"]} + def _initialize_cyclic_variable_shifts(self, options: AnyDict) -> None: + self.cyclic_variable_shifts = [] + for shift in options.get("variable_shifts", []): + if isinstance(shift, dict): + shift = CyclicVariableShift(**shift) + if not isinstance(shift, CyclicVariableShift): + raise TypeError("Each cyclic variable shift must be a CyclicVariableShift or a compatible dictionary") + if shift.key not in self.nlp[0].states: + raise KeyError(f"Cyclic state '{shift.key}' is not defined") + indices = (shift.indices,) if isinstance(shift.indices, int) else tuple(shift.indices) + if any(index < 0 or index >= len(self.nlp[0].states[shift.key]) for index in indices): + raise IndexError(f"A cyclic index for state '{shift.key}' is out of range") + self.cyclic_variable_shifts.append( + CyclicVariableShift(shift.key, indices, float(shift.period), int(shift.turns)) + ) + + def _apply_cyclic_variable_shifts(self, key: str, values: np.ndarray) -> np.ndarray: + shifted = np.asarray(values).copy() + for shift in self.cyclic_variable_shifts: + if shift.key == key: + shifted[list(shift.indices), ...] += shift.turns * shift.period + return shifted + def _set_cyclic_bound(self, sol: Solution | None = None) -> None: if self.nlp[0].x_bounds.type != InterpolationType.CONSTANT_WITH_FIRST_AND_LAST_DIFFERENT: raise ValueError( @@ -586,8 +749,9 @@ def _set_cyclic_bound(self, sol: Solution | None = None) -> None: else: t = self.time_idx_to_cycle * self.nb_intermediate_frames states = sol.decision_states(to_merge=SolutionMerge.NODES) - self.nlp[0].x_bounds[key].min[s, 2] = states[key][s, t] - range_of_motion * 0.01 - self.nlp[0].x_bounds[key].max[s, 2] = states[key][s, t] + range_of_motion * 0.01 + terminal_state = self._apply_cyclic_variable_shifts(key, states[key][:, t]) + self.nlp[0].x_bounds[key].min[s, 2] = terminal_state[s] - range_of_motion * 0.01 + self.nlp[0].x_bounds[key].max[s, 2] = terminal_state[s] + range_of_motion * 0.01 def advance_window(self, sol: Solution, steps: Int = 0, **advance_options) -> None: super(CyclicRecedingHorizonOptimization, self).advance_window(sol, steps, **advance_options) @@ -599,7 +763,8 @@ def advance_window_bounds_states(self, sol: Solution, **advance_options) -> Bool # Update the initial frame bounds for key in states.keys(): - self.nlp[0].x_bounds[key][:, 0] = states[key][:, self.time_idx_to_cycle * self.nb_intermediate_frames] + initial_state = states[key][:, self.time_idx_to_cycle * self.nb_intermediate_frames] + self.nlp[0].x_bounds[key][:, 0] = self._apply_cyclic_variable_shifts(key, initial_state) self._set_cyclic_bound(sol) return True @@ -609,11 +774,14 @@ def advance_window_initial_guess_states(self, sol: Solution, **advance_options) for key in states.keys(): if self.nlp[0].x_init[key].type != InterpolationType.EACH_FRAME: self.nlp[0].x_init.add( - key, np.ndarray(states[key].shape), interpolation=InterpolationType.EACH_FRAME, phase=0 + key, + np.ndarray(states[key].shape), + interpolation=InterpolationType.EACH_FRAME, + phase=0, ) self.nlp[0].x_init[key].check_and_adjust_dimensions(len(self.nlp[0].states[key]), self.nlp[0].ns) - self.nlp[0].x_init[key].init[:, :] = states[key] + self.nlp[0].x_init[key].init[:, :] = self._apply_cyclic_variable_shifts(key, states[key]) return True def advance_window_initial_guess_controls(self, sol: Solution, **advance_options) -> Bool: @@ -670,7 +838,12 @@ def __init__( self.initial_guess_frames = [] for _ in range(self.n_cycles): self.initial_guess_frames.extend( - list(range(self.n_cycles_to_advance * self.cycle_len, (self.n_cycles_to_advance + 1) * self.cycle_len)) + list( + range( + self.n_cycles_to_advance * self.cycle_len, + (self.n_cycles_to_advance + 1) * self.cycle_len, + ) + ) ) self.initial_guess_frames.append((self.n_cycles_to_advance + 1) * self.cycle_len) @@ -692,12 +865,18 @@ def advance_window_initial_guess_states(self, sol: Solution, **advance_options) if self.nlp[0].x_init[key].type != InterpolationType.ALL_POINTS: self.nlp[0].x_init.add( key, - np.ndarray((states[key].shape[0], self.nlp[0].ns * self.nb_intermediate_frames + 1)), + np.ndarray( + ( + states[key].shape[0], + self.nlp[0].ns * self.nb_intermediate_frames + 1, + ) + ), interpolation=InterpolationType.ALL_POINTS, phase=0, ) self.nlp[0].x_init[key].check_and_adjust_dimensions( - self.nlp[0].states[key].shape, self.nlp[0].ns * self.nb_intermediate_frames + self.nlp[0].states[key].shape, + self.nlp[0].ns * self.nb_intermediate_frames, ) else: initial_guess_frames = [] @@ -743,7 +922,10 @@ def advance_window_initial_guess_controls(self, sol: Solution, **advance_options self.nlp[0].controls[key].shape, self.nlp[0].n_controls_nodes - 1 ) - if self.nlp[0].control_type in (ControlType.CONSTANT, ControlType.CONSTANT_WITH_LAST_NODE): + if self.nlp[0].control_type in ( + ControlType.CONSTANT, + ControlType.CONSTANT_WITH_LAST_NODE, + ): frames = self.initial_guess_frames[:-1] elif self.nlp[0].control_type == ControlType.LINEAR_CONTINUOUS: frames = self.initial_guess_frames @@ -790,7 +972,10 @@ def solve( final_solution.append(solution[1]) cycle_solutions_output = [] - if cycle_solutions in (MultiCyclicCycleSolutions.FIRST_CYCLES, MultiCyclicCycleSolutions.ALL_CYCLES): + if cycle_solutions in ( + MultiCyclicCycleSolutions.FIRST_CYCLES, + MultiCyclicCycleSolutions.ALL_CYCLES, + ): for sol in solution[1]: _states, _controls, _parameters = self.export_cycles(sol) dt = float(sol.t_span()[0][-1]) @@ -802,7 +987,10 @@ def solve( dt = float(sol.t_span()[0][-1]) cycle_solutions_output.append(self._initialize_one_cycle(dt, _states, _controls, _parameters)) - if cycle_solutions in (MultiCyclicCycleSolutions.FIRST_CYCLES, MultiCyclicCycleSolutions.ALL_CYCLES): + if cycle_solutions in ( + MultiCyclicCycleSolutions.FIRST_CYCLES, + MultiCyclicCycleSolutions.ALL_CYCLES, + ): final_solution.append(cycle_solutions_output) return tuple(final_solution) if len(final_solution) > 1 else final_solution[0] @@ -824,7 +1012,10 @@ def export_cycles(self, sol: Solution, cycle_number: Int = 0) -> AnyTuple: for key in self.nlp[0].states.keys(): states[key] = decision_states[key][:, window_slice] - if self.nlp[0].control_type in (ControlType.CONSTANT, ControlType.CONSTANT_WITH_LAST_NODE): + if self.nlp[0].control_type in ( + ControlType.CONSTANT, + ControlType.CONSTANT_WITH_LAST_NODE, + ): window_slice = slice(cycle_number * self.cycle_len, (cycle_number + 1) * self.cycle_len) for key in self.nlp[0].controls.keys(): controls[key] = decision_controls[key][:, window_slice] @@ -839,7 +1030,10 @@ def _initialize_solution(self, dt: Float, states: AnyList, controls: AnyList, pa for key in self.nlp[0].states.keys(): x_init.add( key, - np.concatenate([state[key][:, :-1] for state in states] + [states[-1][key][:, -1:]], axis=1), + np.concatenate( + [state[key][:, :-1] for state in states] + [states[-1][key][:, -1:]], + axis=1, + ), interpolation=self.nlp[0].x_init.type, phase=0, ) @@ -888,7 +1082,7 @@ def _initialize_one_cycle(self, dt: Float, states: AnyDict, controls: AnyDict, p x_init.add( key, states[key], - interpolation=self.nlp[0].x_init.type, + interpolation=InterpolationType.EACH_FRAME, phase=0, ) diff --git a/bioptim/optimization/solution/solution.py b/bioptim/optimization/solution/solution.py index cbeea6273..565a64e34 100644 --- a/bioptim/optimization/solution/solution.py +++ b/bioptim/optimization/solution/solution.py @@ -127,6 +127,7 @@ def __init__( real_time_to_optimize: FloatOptional = None, iterations: IntOptional = None, status: IntOptional = None, + solver_diagnostics: AnyDict | None = None, ): """ Parameters @@ -168,8 +169,18 @@ def __init__( # Solver options self.status, self.iterations = status, iterations - self.lam_g, self.lam_p, self.lam_x, self.inf_pr, self.inf_du = lam_g, lam_p, lam_x, inf_pr, inf_du - self.solver_time_to_optimize, self.real_time_to_optimize = solver_time_to_optimize, real_time_to_optimize + self.solver_diagnostics = solver_diagnostics + self.lam_g, self.lam_p, self.lam_x, self.inf_pr, self.inf_du = ( + lam_g, + lam_p, + lam_x, + inf_pr, + inf_du, + ) + self.solver_time_to_optimize, self.real_time_to_optimize = ( + solver_time_to_optimize, + real_time_to_optimize, + ) # Extract the data now for further use self._decision_states = None @@ -223,6 +234,7 @@ def from_dict(cls, ocp: "OptimalControlProgram", sol: AnyDict): real_time_to_optimize=sol["real_time_to_optimize"], iterations=sol["iter"], status=sol["status"], + solver_diagnostics=sol.get("solver_diagnostics"), ) @classmethod @@ -292,7 +304,7 @@ def from_initial_guess(cls, ocp: "OptimalControlProgram", sol: AnyList): ns = ( ocp.nlp[p].ns * nb_intermediate_frames if ss[key].init.type == InterpolationType.ALL_POINTS - else ocp.nlp[p].ns + 1 if ss[key].init.type != InterpolationType.EACH_FRAME else ocp.nlp[p].ns + else (ocp.nlp[p].ns + 1 if ss[key].init.type != InterpolationType.EACH_FRAME else ocp.nlp[p].ns) ) ss[key].init.check_and_adjust_dimensions(len(ocp.nlp[p].states[key]), ns, "states") @@ -722,6 +734,8 @@ def copy(self, skip_data: Bool = False) -> "Solution": new.solver_time_to_optimize = deepcopy(self.solver_time_to_optimize) new.real_time_to_optimize = deepcopy(self.real_time_to_optimize) new.iterations = deepcopy(self.iterations) + new.status = deepcopy(self.status) + new.solver_diagnostics = deepcopy(self.solver_diagnostics) new.phases_dt = deepcopy(self.phases_dt) new._stepwise_times = deepcopy(self._stepwise_times) @@ -861,7 +875,7 @@ def integrate( if return_time: time_vector = self._return_time_vector(to_merge=to_merge, duplicated_times=duplicated_times) - return out if len(out) > 1 else out[0], time_vector if len(time_vector) > 1 else time_vector[0] + return out if len(out) > 1 else out[0], (time_vector if len(time_vector) > 1 else time_vector[0]) else: return out if len(out) > 1 else out[0] @@ -874,7 +888,9 @@ def noisy_integrate( """ Integrated the states with different noise values sampled from the covariance matrix. """ - from ...optimization.stochastic_optimal_control_program import StochasticOptimalControlProgram + from ...optimization.stochastic_optimal_control_program import ( + StochasticOptimalControlProgram, + ) from ...interfaces.interface_utils import get_numerical_timeseries if not isinstance(self.ocp, StochasticOptimalControlProgram): @@ -1326,7 +1342,9 @@ def _get_penalty_cost(self, penalty: PenaltyOption) -> FloatTuple: phases_dt = PenaltyHelpers.phases_dt(penalty, self.ocp, lambda p: np.array([self.phases_dt[idx] for idx in p])) params = PenaltyHelpers.parameters( - penalty, 0, lambda p_idx, n_idx, sn_idx: self._dispatch_params(self._parameters.scaled[0]) + penalty, + 0, + lambda p_idx, n_idx, sn_idx: self._dispatch_params(self._parameters.scaled[0]), ) merged_x = self._decision_states.to_dict(to_merge=SolutionMerge.KEYS, scaled=True) @@ -1422,12 +1440,20 @@ def _compute_detailed_cost(self) -> None: continue val, val_weighted = self._get_penalty_cost(penalty) self._detailed_cost += [ - {"name": penalty.type.__str__(), "cost_value_weighted": val_weighted, "cost_value": val} + { + "name": penalty.type.__str__(), + "cost_value_weighted": val_weighted, + "cost_value": val, + } ] for penalty in self.ocp.J: val, val_weighted = self._get_penalty_cost(penalty) self._detailed_cost += [ - {"name": penalty.type.__str__(), "cost_value_weighted": val_weighted, "cost_value": val} + { + "name": penalty.type.__str__(), + "cost_value_weighted": val_weighted, + "cost_value": val, + } ] return diff --git a/tests/shard2/test_receding_horizon_hooks_and_shifts.py b/tests/shard2/test_receding_horizon_hooks_and_shifts.py new file mode 100644 index 000000000..b82507c8b --- /dev/null +++ b/tests/shard2/test_receding_horizon_hooks_and_shifts.py @@ -0,0 +1,45 @@ +from types import SimpleNamespace + +import numpy as np + +from bioptim import CyclicVariableShift, RecedingHorizonFailurePolicy +from bioptim.misc.enums import SolverType +from bioptim.optimization.optimal_control_program import OptimalControlProgram +from bioptim.optimization.receding_horizon_optimization import ( + CyclicRecedingHorizonOptimization, + RecedingHorizonOptimization, +) + + +def test_window_hooks_can_stop_after_collecting_diagnostics(monkeypatch): + solved = SimpleNamespace(status=1, vector=np.ones((1, 1)), real_time_to_optimize=0.1) + monkeypatch.setattr(OptimalControlProgram, "solve", lambda *args, **kwargs: solved) + events = [] + + rhe = object.__new__(RecedingHorizonOptimization) + rhe.nlp = [SimpleNamespace(x_bounds={})] + solver = SimpleNamespace(type=SolverType.IPOPT, online_optim=False) + returned = rhe.solve( + update_function=lambda *args: True, + solver=solver, + failure_policy=RecedingHorizonFailurePolicy.CONTINUE_DIAGNOSTIC, + before_window_solve=lambda _, index, __: events.append(("before", index)), + after_window_solve=lambda _, index, result: events.append(("after", index, result.solver_succeeded)) or False, + ) + + assert returned is solved + assert events == [("before", 0), ("after", 0, False)] + + +def test_cyclic_shift_uses_configured_period_and_state_index(): + cyclic = object.__new__(CyclicRecedingHorizonOptimization) + cyclic.nlp = [SimpleNamespace(states={"q": [object(), object()], "qdot": [object(), object()]})] + cyclic._initialize_cyclic_variable_shifts( + {"variable_shifts": [CyclicVariableShift(key="q", indices=1, period=3.5, turns=-2)]} + ) + + q = cyclic._apply_cyclic_variable_shifts("q", np.array([[1.0, 2.0], [10.0, 20.0]])) + qdot = cyclic._apply_cyclic_variable_shifts("qdot", np.array([[1.0], [2.0]])) + + np.testing.assert_allclose(q, [[1.0, 2.0], [3.0, 13.0]]) + np.testing.assert_allclose(qdot, [[1.0], [2.0]]) diff --git a/tests/shard2/test_receding_horizon_status.py b/tests/shard2/test_receding_horizon_status.py new file mode 100644 index 000000000..0facc6a8d --- /dev/null +++ b/tests/shard2/test_receding_horizon_status.py @@ -0,0 +1,105 @@ +from types import SimpleNamespace + +import numpy as np +import pytest + +from bioptim import RecedingHorizonFailurePolicy +from bioptim.misc.enums import SolverType +from bioptim.optimization.optimal_control_program import OptimalControlProgram +from bioptim.optimization.receding_horizon_optimization import ( + RecedingHorizonOptimization, +) + + +def test_failed_window_is_reported_and_not_exported(monkeypatch): + failed_solution = SimpleNamespace( + status=3, + vector=np.ones((1, 1)), + real_time_to_optimize=0.1, + ) + monkeypatch.setattr(OptimalControlProgram, "solve", lambda *args, **kwargs: failed_solution) + + rhe = object.__new__(RecedingHorizonOptimization) + rhe.nlp = [SimpleNamespace(x_bounds={})] + solver = SimpleNamespace(type=SolverType.IPOPT, online_optim=False) + solution = rhe.solve( + update_function=lambda *args: True, + solver=solver, + failure_policy=RecedingHorizonFailurePolicy.STOP, + ) + + assert solution is failed_solution + assert solution.status == 3 + assert len(solution.window_results) == 1 + assert not solution.window_results[0].solver_succeeded + assert solution.window_results[0].trajectory_available + assert not solution.window_results[0].physically_acceptable + assert not solution.window_results[0].exported + + +class _FakeAcadosSolver: + def __init__(self): + self.set_calls = [] + + def get_stats(self, name): + if name == "unavailable": + raise RuntimeError + return { + "res_stat": 1e-8, + "sqp_iter": 4, + "time_tot": 0.02, + "statistics": np.ones((2, 2)), + }[name] + + def get(self, node, field): + return np.array([node]) if field == "x" else np.array([-node]) + + def set(self, node, field, value): + self.set_calls.append((node, field, np.asarray(value))) + + +def test_acados_diagnostics_and_iterates_are_public(): + pytest.importorskip("acados_template") + from bioptim.interfaces.acados_interface import AcadosInterface + + interface = object.__new__(AcadosInterface) + interface.status = 3 + interface.real_time_to_optimize = 0.03 + interface.ocp_solver = _FakeAcadosSolver() + interface.acados_ocp = SimpleNamespace(dims=SimpleNamespace(N=2)) + + diagnostics = interface.get_solver_diagnostics() + iterates = interface.get_iterates() + + assert diagnostics["status"] == 3 + assert diagnostics["status_label"] == "qp_solver_failure" + assert diagnostics["res_stat"] == 1e-8 + assert diagnostics["sqp_iter"] == 4 + assert len(iterates) == 3 + np.testing.assert_array_equal(iterates[1]["u"], [-1]) + assert iterates[-1]["u"] is None + + interface.set_iterates(iterates, include_multipliers=False) + assert not any(field in ("pi", "lam", "sl", "su") for _, field, _ in interface.ocp_solver.set_calls) + + +def test_acados_multiplier_transfer_requires_explicit_opt_in(): + pytest.importorskip("acados_template") + from bioptim.interfaces.acados_interface import AcadosInterface + + interface = object.__new__(AcadosInterface) + interface.status = 0 + interface.real_time_to_optimize = 0.03 + interface.ocp_solver = _FakeAcadosSolver() + interface.acados_ocp = SimpleNamespace(dims=SimpleNamespace(N=1)) + iterates = interface.get_iterates() + + interface.set_iterates(iterates) + primal_fields = [field for _, field, _ in interface.ocp_solver.set_calls] + assert "x" in primal_fields and "u" in primal_fields + assert "pi" not in primal_fields and "lam" not in primal_fields + + interface.ocp_solver.set_calls.clear() + interface.set_iterates(iterates, include_multipliers=True) + dual_fields = [field for _, field, _ in interface.ocp_solver.set_calls] + assert "pi" in dual_fields and "lam" in dual_fields