feat: adapt solutions to new initial guess grids

Author: mickaelbegon

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

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.

bioptim/optimization/solution/utils.py

@@ -6,6 +6,105 @@
     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
+    if interpolation == InterpolationType.CONSTANT_WITH_FIRST_AND_LAST_DIFFERENT:
+        return 3
+    if interpolation == InterpolationType.LINEAR:
+        return 2
+    if interpolation in (InterpolationType.EACH_FRAME, InterpolationType.ALL_POINTS):
+        return initial_guess.init.shape[1]
+    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}")
+            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 nodes,
+    collocation points and control nodes. 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]:
     """

tests/shard4/test_solution_grid_adaptation.py

@@ -0,0 +1,53 @@
+import numpy as np
+
+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]])
+
+
+class _FakeSolutionAdapter(_FakeSolution):
+    to_initial_guesses = __import__("bioptim").Solution.to_initial_guesses