diff --git a/bioptim/__init__.py b/bioptim/__init__.py index 7ccf9b7c2..af9ff4a75 100644 --- a/bioptim/__init__.py +++ b/bioptim/__init__.py @@ -249,6 +249,7 @@ ) from .optimization.receding_horizon_optimization import MovingHorizonEstimator, NonlinearModelPredictiveControl from .optimization.solution.solution import Solution +from .optimization.solution.utils import adapt_solution_to_initial_guesses from .optimization.solution.solution_data import SolutionMerge, TimeAlignment from .optimization.stochastic_optimal_control_program import StochasticOptimalControlProgram from .optimization.variable_scaling import VariableScalingList, VariableScaling diff --git a/bioptim/optimization/solution/solution.py b/bioptim/optimization/solution/solution.py index cbeea6273..718f0c37f 100644 --- a/bioptim/optimization/solution/solution.py +++ b/bioptim/optimization/solution/solution.py @@ -736,6 +736,25 @@ def copy(self, skip_data: Bool = False) -> "Solution": new._parameters = deepcopy(self._parameters) return new + def to_initial_guesses( + self, + state_initial_guesses: InitialGuessList, + control_initial_guesses: InitialGuessList, + parameter_initial_guesses: InitialGuessList | None = None, + algebraic_state_initial_guesses: InitialGuessList | None = None, + ) -> tuple[InitialGuessList, InitialGuessList, InitialGuessList, InitialGuessList]: + """Adapt this solution's primal variables to new initial-guess grids.""" + + from .utils import adapt_solution_to_initial_guesses + + return adapt_solution_to_initial_guesses( + self, + state_initial_guesses, + control_initial_guesses, + parameter_initial_guesses, + algebraic_state_initial_guesses, + ) + def _prepare_integrate(self, integrator: SolutionIntegrator) -> AnyTuple: """ Prepare the variables for the states integration and checks if the integrator is compatible with the ocp. diff --git a/bioptim/optimization/solution/utils.py b/bioptim/optimization/solution/utils.py index a529dc53d..c2a21a362 100644 --- a/bioptim/optimization/solution/utils.py +++ b/bioptim/optimization/solution/utils.py @@ -6,6 +6,112 @@ NpArrayDict, ) +from ...limits.path_conditions import InitialGuessList +from ...misc.enums import InterpolationType +from .solution_data import SolutionMerge + + +def _resample_initial_guess(values: np.ndarray, n_columns: int) -> np.ndarray: + """Linearly resample node values while preserving both endpoints.""" + + values = np.asarray(values, dtype=float) + if values.ndim == 1: + values = values[:, np.newaxis] + finite_columns = np.all(np.isfinite(values), axis=0) + values = values[:, finite_columns] + if values.shape[1] == 0: + raise ValueError("A solution variable contains no finite node that can be transferred") + if n_columns == 1: + return values[:, :1].copy() + if values.shape[1] == 1: + return np.repeat(values, n_columns, axis=1) + + source_grid = np.linspace(0.0, 1.0, values.shape[1]) + target_grid = np.linspace(0.0, 1.0, n_columns) + return np.vstack([np.interp(target_grid, source_grid, row) for row in values]) + + +def _target_column_count(initial_guess) -> int: + interpolation = initial_guess.type + if interpolation == InterpolationType.CONSTANT: + return 1 + elif interpolation == InterpolationType.CONSTANT_WITH_FIRST_AND_LAST_DIFFERENT: + return 3 + elif interpolation == InterpolationType.LINEAR: + return 2 + elif interpolation in (InterpolationType.EACH_FRAME, InterpolationType.ALL_POINTS): + return initial_guess.init.shape[1] + else: + raise NotImplementedError(f"Adapting a solution to {interpolation} is not implemented") + + +def _as_phase_list(data, n_phases: int) -> list[dict]: + if n_phases == 1 and isinstance(data, dict): + return [data] + if not isinstance(data, list) or len(data) != n_phases: + raise ValueError(f"The solution has {len(data) if isinstance(data, list) else 1} phases, expected {n_phases}") + return data + + +def _adapt_variable_group(source, target: InitialGuessList, group_name: str) -> InitialGuessList: + adapted = InitialGuessList() + n_phases = len(target.options) + source_phases = _as_phase_list(source, n_phases) + for phase, target_phase in enumerate(target.options): + for key, target_guess in target_phase.items(): + if key not in source_phases[phase]: + raise KeyError(f"{group_name} '{key}' is absent from the solution phase {phase}") + if target_guess.type in (InterpolationType.SPLINE, InterpolationType.CUSTOM): + raise NotImplementedError( + f"{group_name} '{key}' cannot be linearly resampled to {target_guess.type}. " + "Spline time vectors and custom interpolation callbacks cannot be inferred from solution samples" + ) + values = _resample_initial_guess(source_phases[phase][key], _target_column_count(target_guess)) + adapted.add(key, values, interpolation=target_guess.type, phase=phase) + return adapted + + +def adapt_solution_to_initial_guesses( + solution, + state_initial_guesses: InitialGuessList, + control_initial_guesses: InitialGuessList, + parameter_initial_guesses: InitialGuessList | None = None, + algebraic_state_initial_guesses: InitialGuessList | None = None, +) -> tuple[InitialGuessList, InitialGuessList, InitialGuessList, InitialGuessList]: + """Adapt a solution's primal variables to the grids described by new initial-guess lists. + + Resampling is performed on a normalized phase grid, so it supports different numbers of shooting and control + nodes. ``SPLINE`` and ``CUSTOM`` targets are rejected because their time vectors or callbacks cannot be inferred + from solution samples. Solver multipliers are deliberately not transferred by this function. + """ + + states = _adapt_variable_group( + solution.decision_states(to_merge=SolutionMerge.NODES), state_initial_guesses, "State" + ) + controls = _adapt_variable_group( + solution.decision_controls(to_merge=SolutionMerge.NODES), control_initial_guesses, "Control" + ) + + parameters = InitialGuessList() + if parameter_initial_guesses is not None: + source_parameters = solution.decision_parameters() + for phase, target_phase in enumerate(parameter_initial_guesses.options): + for key, target_guess in target_phase.items(): + if key not in source_parameters: + raise KeyError(f"Parameter '{key}' is absent from the solution") + values = np.asarray(source_parameters[key], dtype=float).reshape((-1, 1)) + parameters.add(key, values, interpolation=target_guess.type, phase=phase) + + algebraic_states = InitialGuessList() + if algebraic_state_initial_guesses is not None: + algebraic_states = _adapt_variable_group( + solution.decision_algebraic_states(to_merge=SolutionMerge.NODES), + algebraic_state_initial_guesses, + "Algebraic state", + ) + + return states, controls, parameters, algebraic_states + def concatenate_optimization_variables_dict(variable: list[NpArrayDict], continuous: Bool = True) -> list[NpArrayDict]: """ diff --git a/tests/shard4/test_solution_grid_adaptation.py b/tests/shard4/test_solution_grid_adaptation.py new file mode 100644 index 000000000..0ab031d41 --- /dev/null +++ b/tests/shard4/test_solution_grid_adaptation.py @@ -0,0 +1,71 @@ +import numpy as np +import pytest + +from bioptim import InitialGuessList, InterpolationType, adapt_solution_to_initial_guesses + + +class _FakeSolution: + def decision_states(self, **_): + return {"q": np.array([[0.0, 1.0, 2.0]])} + + def decision_controls(self, **_): + return {"tau": np.array([[0.0, 2.0, np.nan]])} + + def decision_parameters(self): + return {"mass": np.array([3.0])} + + def decision_algebraic_states(self, **_): + return {"contact": np.array([[1.0, 3.0, 5.0]])} + + +def test_adapt_solution_between_different_grids(): + x_init = InitialGuessList() + x_init.add("q", np.zeros((1, 5)), interpolation=InterpolationType.EACH_FRAME) + u_init = InitialGuessList() + u_init.add("tau", np.zeros((1, 4)), interpolation=InterpolationType.EACH_FRAME) + p_init = InitialGuessList() + p_init.add("mass", [0.0], interpolation=InterpolationType.CONSTANT) + a_init = InitialGuessList() + a_init.add("contact", np.zeros((1, 7)), interpolation=InterpolationType.ALL_POINTS) + + states, controls, parameters, algebraic_states = adapt_solution_to_initial_guesses( + _FakeSolution(), x_init, u_init, p_init, a_init + ) + + np.testing.assert_allclose(states[0]["q"].init, [[0.0, 0.5, 1.0, 1.5, 2.0]]) + np.testing.assert_allclose(controls[0]["tau"].init, [[0.0, 2 / 3, 4 / 3, 2.0]]) + np.testing.assert_allclose(parameters[0]["mass"].init, [[3.0]]) + np.testing.assert_allclose(algebraic_states[0]["contact"].init, [[1, 5 / 3, 7 / 3, 3, 11 / 3, 13 / 3, 5]]) + + +def test_adapt_solution_preserves_control_grid_semantics(): + x_init = InitialGuessList() + x_init.add("q", np.zeros((1, 3)), interpolation=InterpolationType.EACH_FRAME) + u_init = InitialGuessList() + u_init.add("tau", np.zeros((1, 2)), interpolation=InterpolationType.LINEAR) + + _, controls, _, _ = _FakeSolutionAdapter().to_initial_guesses(x_init, u_init) + + assert controls[0]["tau"].type == InterpolationType.LINEAR + np.testing.assert_allclose(controls[0]["tau"].init, [[0.0, 2.0]]) + + +@pytest.mark.parametrize( + ("interpolation", "initial_guess", "extra_arguments"), + ( + (InterpolationType.SPLINE, [[0.0, 0.0]], {"t": [0.0, 1.0]}), + (InterpolationType.CUSTOM, lambda _index: np.array([0.0]), {}), + ), +) +def test_adapt_solution_rejects_interpolations_that_cannot_be_inferred(interpolation, initial_guess, extra_arguments): + x_init = InitialGuessList() + x_init.add("q", initial_guess, interpolation=interpolation, **extra_arguments) + u_init = InitialGuessList() + u_init.add("tau", np.zeros((1, 2)), interpolation=InterpolationType.EACH_FRAME) + + with pytest.raises(NotImplementedError, match="cannot be linearly resampled"): + adapt_solution_to_initial_guesses(_FakeSolution(), x_init, u_init) + + +class _FakeSolutionAdapter(_FakeSolution): + to_initial_guesses = __import__("bioptim").Solution.to_initial_guesses