scikit-complexity 0.2.0__py3-none-any.whl

skcomplex/engine/solvers/_optimal_control.py

@@ -0,0 +1,368 @@
+"""Read optimal control problems from compiled systems and follow their solutions over the horizon of a run."""
+
+# Author: Georgios Douzas <gdouzas@icloud.com>
+# License: MIT
+
+import math
+from abc import abstractmethod
+from collections.abc import Mapping, Sequence
+from types import SimpleNamespace
+
+import numpy as np
+import numpy.typing as npt
+from sage.all__sagemath_symbolics import NN, SR, Expression, oo
+from sage.sets.real_set import RealSet
+
+from ...errors import SolverError, ValidationError
+from ...specs import ControlSolutionSpec
+from ...specs._control_spec import _ControlSpec
+from ...types import Array, FormKind
+from .._utils import symbols_of
+from ..compile import ObservableEvaluator
+from ..problems import ExpressionPrinter, ProblemSpec
+from ..problems._control_problem_spec import _ControlProblemSpec
+from ..types import Cost, GridFunction
+from ._solver import Solver
+
+TIME_TOLERANCE = 1e-9
+_SENSES = ('min', 'max')
+_NONNEGATIVE_DOMAINS = frozenset({RealSet.open(0, oo), RealSet.closed_open(0, oo), NN})
+
+
+class OptimalControlSolver(Solver):
+    """Base class of solvers that compute the controls minimising or maximising an objective over a run of a system.
+
+    The solver reads the bounded adaptive parameters of the system as its controls. Its objective is the running cost
+    `objective` integrated over the run plus the terminal cost `terminal` at the horizon, the time `until` of the run.
+    Each cost is a function of the namespace `s` of the states, parameters, and time of the system by dotted name, such
+    as `lambda s: s.plant.x**2 + s.plant.u**2`, or a number. The solver minimises the objective for the sense `'min'`
+    and maximises it for `'max'`. It solves the problem over the whole run from the initial state and the initial
+    parameters at the start of the run, and then follows the solution frame by frame. Subclasses take `objective`,
+    `terminal`, and `sense` as arguments, implement `_solve`, and may implement `_prepare`.
+
+    Attributes:
+        control_problem_:
+            The optimal control problem of the compiled system.
+
+        horizon_:
+            The end of the run, the horizon of the objective.
+
+        solution_:
+            The optimal solution of the run.
+    """
+
+    @classmethod
+    def _get_form(cls) -> FormKind:
+        """Return the form the solver solves."""
+        return FormKind.ODE
+
+    objective: Cost | None
+    terminal: Cost
+    sense: str
+
+    def _check(self) -> None:
+        """Raise an error for a missing objective or an unknown sense.
+
+        Raises:
+            SolverError:
+                The solver has no objective.
+
+            ValidationError:
+                The sense is neither `'min'` nor `'max'`.
+        """
+        if self.objective is None:
+            error_msg = (
+                f'{type(self).__name__} has no objective. Pass the running cost as `objective`, a function of the '
+                'namespace of the system such as `lambda s: s.plant.x**2`.'
+            )
+            raise SolverError(error_msg)
+        if self.sense not in _SENSES:
+            error_msg = f'Unknown sense {self.sense!r} of the objective. Use one of {", ".join(_SENSES)}.'
+            raise ValidationError(error_msg)
+
+    def _integrate(self, x0: Array, controls: Array, times: Array, p: Array) -> Array:
+        """Return the steered states on a grid from their initial values, by Runge-Kutta steps under given controls."""
+        control_problem = self.control_problem_
+        states = np.zeros((times.size, x0.size))
+        states[0] = x0
+        full = self.problem_.initial_state.copy()
+        for k in range(times.size - 1):
+            h = times[k + 1] - times[k]
+            start, end = controls[k], controls[k + 1]
+            middle = (start + end) / 2
+            x = states[k]
+            k1 = self.rates_(times[k], x, start, p)
+            k2 = self.rates_(times[k] + h / 2, x + h / 2 * k1, middle, p)
+            k3 = self.rates_(times[k] + h / 2, x + h / 2 * k2, middle, p)
+            k4 = self.rates_(times[k + 1], x + h * k3, end, p)
+            states[k + 1] = x + h / 6 * (k1 + 2 * k2 + 2 * k3 + k4)
+            if self.problem_.boundary is not None:
+                full[control_problem.indices] = states[k + 1]
+                full[control_problem.control_indices] = end
+                states[k + 1] = self.problem_.boundary(times[k + 1], full, p)[control_problem.indices]
+        return states
+
+    def _cost(self, times: Array, states: Array, controls: Array, p: Array) -> float:
+        """Return the value of the objective, in its own sense, for states and controls on a grid of times."""
+        running = self.running_(times, states, controls, p)[:, 0]
+        terminal = self.terminal_(times[-1], states[-1], controls[-1], p)[0]
+        return float(self.control_problem_.sign * (np.trapezoid(running, times) + terminal))
+
+    def _solution(
+        self,
+        x0: Array,
+        times: Array,
+        states: Array,
+        controls: Array,
+        p: Array,
+        *,
+        iterations: int,
+        costates: Mapping[str, Array] | None = None,
+    ) -> ControlSolutionSpec:
+        """Return the solution of steered states and controls on a grid, with its cost."""
+        control_problem = self.control_problem_
+        full = np.tile(x0, (times.size, 1))
+        full[:, control_problem.indices] = states
+        full[:, control_problem.control_indices] = controls
+        cost = self._cost(times, states, controls, p)
+        return ControlSolutionSpec(times, full, dict(costates or {}), cost, iterations)
+
+    def _prepare(self, context: ProblemSpec | None) -> None:
+        """Prepare the solver for a problem and read its optimal control problem.
+
+        Args:
+            context:
+                The compiled problem.
+
+        Raises:
+            SolverError:
+                The problem is not an ODE, or has no bounded adaptive parameter or laws the solver cannot
+                differentiate symbolically, or the run has no end.
+        """
+        super()._prepare(context)
+        problem = self.problem_
+        control_problem = self._control_problem(problem)
+        if not math.isfinite(self.until):
+            error_msg = 'An optimal control run needs a finite horizon. Run the simulation with `run(until=...)`.'
+            raise SolverError(error_msg)
+        self.control_problem_ = control_problem
+        self.horizon_ = float(self.until)
+        self.rates_ = self._function(control_problem, control_problem.rates)
+        self.running_ = self._function(control_problem, [control_problem.running])
+        self.terminal_ = self._function(control_problem, [control_problem.terminal])
+        self._set_up_control(control_problem)
+        self.solution_ = self._solve(self.initial_state_.copy(), problem.initial_parameters.copy(), self.horizon_)
+        self.initial_state_ = self.solution_.state_at(0.0)
+
+    def _advance(self, t: float, x: Array, p: Array, until: float) -> Array:
+        """Return the state of the optimal solution at a time of the run.
+
+        Args:
+            t:
+                The current time.
+
+            x:
+                The current state vector.
+
+            p:
+                The parameter vector.
+
+            until:
+                The time to advance to.
+
+        Returns:
+            state:
+                The optimal state vector at time `until`.
+
+        Raises:
+            SolverError:
+                The time lies after the horizon the problem was solved for.
+        """
+        del t, x, p
+        if until > self.horizon_ + TIME_TOLERANCE:
+            error_msg = (
+                f'The optimal control was solved until the horizon {self.horizon_}, so the run cannot go on to {until}.'
+                ' Run a new simulation with a later horizon.'
+            )
+            raise SolverError(error_msg)
+        return self.solution_.state_at(until)
+
+    def _set_up_control(self, control_problem: _ControlProblemSpec) -> None:
+        """Derive what the solver needs from the optimal control problem, once per compiled problem.
+
+        Args:
+            control_problem:
+                The optimal control problem.
+        """
+        del control_problem
+
+    @staticmethod
+    def _function(
+        control_problem: _ControlProblemSpec,
+        expressions: Sequence[Expression],
+        extra: Sequence[Expression] = (),
+    ) -> GridFunction:
+        """Return a function `f(t, x, u, p, *extra)` of expressions, whose last axes run over symbols and values."""
+        arguments = [*control_problem.states, *control_problem.control_symbols, *control_problem.parameters, *extra]
+        printer = ExpressionPrinter(time=control_problem.time, compiled=True).check().prepare(arguments)
+        compiled = printer.run([SR(item) for item in expressions]).code_
+        if not callable(compiled):
+            error_msg = f'The printer returned a {type(compiled).__name__}, not a compiled function.'
+            raise TypeError(error_msg)
+        count = len(expressions)
+
+        def evaluate(t: float | Array, x: Array, u: Array, p: Array, *values: Array) -> npt.NDArray[np.float64]:
+            if np.ndim(x) == 1 and np.ndim(t) == 0:
+                point = compiled(t, *x, *u, *p, *(item for value in values for item in value))
+                if not isinstance(point, list):
+                    error_msg = f'The compiled function returned a {type(point).__name__}, not a list of values.'
+                    raise TypeError(error_msg)
+                return np.array(point, dtype=np.float64)
+            times = np.asarray(t, dtype=float)
+            columns = [*np.moveaxis(x, -1, 0), *np.moveaxis(u, -1, 0), *p]
+            columns += [item for value in values for item in np.moveaxis(value, -1, 0)]
+            shape = np.broadcast_shapes(times.shape, *(np.shape(item) for item in columns))
+            if not count:
+                return np.zeros((*shape, 0))
+            results = compiled(times, *columns)
+            if not isinstance(results, list):
+                error_msg = f'The compiled function returned a {type(results).__name__}, not a list of values.'
+                raise TypeError(error_msg)
+            return np.stack([np.broadcast_to(np.asarray(item, dtype=float), shape) for item in results], axis=-1)
+
+        return evaluate
+
+    @staticmethod
+    def _check_symbols(control_problem: _ControlProblemSpec) -> None:
+        """Raise an error for a rate or a cost that reads a symbol other than states, controls, parameters, and time."""
+        expressions = [*control_problem.rates, control_problem.running, control_problem.terminal]
+        known = {
+            *control_problem.states,
+            *control_problem.control_symbols,
+            *control_problem.parameters,
+            control_problem.time,
+        }
+        unknown = sorted({str(item) for item in symbols_of([SR(value) for value in expressions]) - known})
+        if unknown:
+            error_msg = (
+                f'The dynamics or the objective read {", ".join(unknown)}, which are not variables of single '
+                'elements or interactions. Optimal control solvers take systems of single elements.'
+            )
+            raise SolverError(error_msg)
+
+    def _control_problem(self, problem: ProblemSpec) -> _ControlProblemSpec:
+        """Return the optimal control problem of a compiled problem, with its bounded controls and the objective."""
+        self._check_supported(problem)
+        objective = self.objective
+        if objective is None:
+            error_msg = f'{type(self).__name__} needs an objective. Pass `objective`, a function or a number.'
+            raise SolverError(error_msg)
+        symbols = self._single_symbols(problem)
+        controls = self._controls(problem, symbols)
+        rates = {problem.equation_symbols[name]: SR(value) for name, value in problem.equations.items()}
+        controlled = {control.index for control in controls}
+        indices = np.array([index for index in range(len(symbols)) if index not in controlled], dtype=np.intp)
+        names = [name for entry in problem.states.entries for name in entry.component_names]
+        sign = 1.0 if self.sense == 'min' else -1.0
+        namespace = ObservableEvaluator._namespace_of(problem)
+        control_problem = _ControlProblemSpec(
+            problem=problem,
+            names=tuple(names[index] for index in indices),
+            indices=indices,
+            states=tuple(symbols[index] for index in indices),
+            controls=controls,
+            parameters=tuple(problem.parameter_symbols),
+            rates=tuple(rates.get(symbols[index], SR(0)) for index in indices),
+            running=sign * self._cost_value(objective, namespace),
+            terminal=sign * self._cost_value(self.terminal, namespace),
+            sign=sign,
+            lower=self._lower_bounds(problem, indices),
+        )
+        self._check_symbols(control_problem)
+        return control_problem
+
+    @staticmethod
+    def _cost_value(cost: Cost, namespace: SimpleNamespace) -> Expression:
+        """Return the expression of a cost, a function of the namespace of the system or a number."""
+        return SR(cost(namespace) if callable(cost) else cost)
+
+    @staticmethod
+    def _single_symbols(problem: ProblemSpec) -> list[Expression]:
+        """Return the state symbols in state vector order, after checking that every state is a real element scalar."""
+        if any(element.members is not None for element in problem.elements):
+            error_msg = 'Optimal control solvers take systems of single elements, but the system has groups.'
+            raise SolverError(error_msg)
+        symbols = problem.state_symbols
+        if len(symbols) != problem.states.size:
+            error_msg = 'Optimal control solvers take real states, but the system has complex states.'
+            raise SolverError(error_msg)
+        return symbols
+
+    @staticmethod
+    def _check_supported(problem: ProblemSpec) -> None:
+        """Raise an error for a problem with laws that have no symbolic vector field or no finite state space."""
+        if problem.form != FormKind.ODE:
+            error_msg = (
+                f'Optimal control solvers take ordinary differential equations, but the system is a {problem.form}.'
+            )
+            raise SolverError(error_msg)
+        if problem.has_callables or problem.has_group_sums or problem.delays or problem.events:
+            error_msg = (
+                'Optimal control solvers derive the costates and the stationarity condition symbolically, so they take '
+                'symbolic rate laws without callables, sums over members, delays, or events.'
+            )
+            raise SolverError(error_msg)
+
+    @staticmethod
+    def _lower_bounds(problem: ProblemSpec, indices: Sequence[int] | npt.NDArray[np.intp]) -> Array:
+        """Return the lower bound of the states at indices: 0 for nonnegative domains and minus infinity for others."""
+        lower = np.full(problem.states.size, -np.inf)
+        for entry in problem.states.entries:
+            if entry.domain in _NONNEGATIVE_DOMAINS:
+                lower[entry.indices] = 0.0
+        return lower[list(indices)]
+
+    @staticmethod
+    def _controls(problem: ProblemSpec, symbols: Sequence[Expression]) -> tuple[_ControlSpec, ...]:
+        """Return the bounded adaptive parameters of a problem as its controls."""
+        controls = []
+        for name, declaration in problem.adaptive_parameters.items():
+            if declaration.bounds is None:
+                continue
+            entry = problem.states.entry(name)
+            if entry.size != 1:
+                error_msg = f'The control {name} is not a scalar. Declare every control as a scalar adaptive parameter.'
+                raise SolverError(error_msg)
+            lower, upper = declaration.bounds
+            controls.append(_ControlSpec(name, entry.start, symbols[entry.start], lower, upper))
+        if not controls:
+            error_msg = (
+                'The system has no bounded adaptive parameter to control. Declare a control input with '
+                '`Annotated[float, ParameterSpec(adaptive=True, bounds=(lower, upper))]`.'
+            )
+            raise SolverError(error_msg)
+        return tuple(controls)
+
+    @abstractmethod
+    def _solve(self, x0: Array, p: Array, horizon: float) -> ControlSolutionSpec:
+        """Return the optimal solution from an initial state over a horizon.
+
+        Args:
+            x0:
+                The initial state vector.
+
+            p:
+                The parameter vector.
+
+            horizon:
+                The end of the run.
+
+        Returns:
+            solution:
+                The optimal solution.
+
+        Raises:
+            NotImplementedError:
+                The subclass does not implement the method.
+        """
+        raise NotImplementedError

skcomplex/engine/solvers/_pontryagin.py

@@ -0,0 +1,298 @@
+"""Solve optimal control problems by Pontryagin's principle, with symbolic costates and a forward-backward sweep."""
+
+# Author: Georgios Douzas <gdouzas@icloud.com>
+# License: MIT
+
+
+import math
+
+import numpy as np
+import numpy.typing as npt
+from sage.all__sagemath_symbolics import SR, Expression, solve
+
+from ...core.laws import BoundaryLaw, EquationLaw, FluxLaw, RateLaw
+from ...core.spaces import ContinuousTime
+from ...errors import SolverError, ValidationError
+from ...specs import ControlSolutionSpec, SupportSpec
+from ...types import Array, OrderKind
+from ..problems._control_problem_spec import _ControlProblemSpec
+from ..types import Cost, GridFunction
+from ._optimal_control import OptimalControlSolver
+
+_MIN_RELAXATION = 1e-3
+MIN_NODES = 2
+
+_CONTROL_SUPPORT = SupportSpec(
+    time={ContinuousTime},
+    space={None},
+    laws={RateLaw, FluxLaw, EquationLaw, BoundaryLaw},
+    orders=frozenset({OrderKind.INTEGER}),
+)
+
+
+class PontryaginSolver(OptimalControlSolver):
+    """Optimal control by Pontryagin's principle, with costates and the stationarity condition derived symbolically.
+
+    The solver builds the Hamiltonian `H = L + λ · f` from the running cost `L` and the rates `f` of the states, and
+    derives the costate equations `λ' = -∂H/∂x` with `λ(T) = ∂φ/∂x` at the horizon `T` for the terminal cost `φ`. The
+    controls solve `∂H/∂u = 0`, projected onto their bounds, or take the bound that minimises `H` where `H` is linear
+    in them. A forward-backward sweep integrates the states forward and the costates backward with fourth-order
+    Runge-Kutta steps on a grid, and relaxes the controls towards their update until they converge. The weight of the
+    update halves whenever the distance of the controls to their update stops falling, which breaks the cycles of the
+    sweep. An objective to maximise is minimised with its sign changed. The solution holds the costates.
+
+    Args:
+        objective:
+            The running cost, a function of the namespace `s` of the states, parameters, and time of the system by
+            dotted name, such as `lambda s: s.plant.x**2`, or a number.
+
+        terminal:
+            The terminal cost at the horizon, a function of the same namespace or a number.
+
+        sense:
+            `'min'` to minimise the objective or `'max'` to maximise it.
+
+        nodes:
+            The number of times of the grid, from the start to the horizon.
+
+        relaxation:
+            The initial weight of the update of the controls at every sweep, between 0 and 1.
+
+        tolerance:
+            The largest relative change of the states, costates, and controls between sweeps at convergence.
+
+        max_iterations:
+            The largest number of sweeps.
+
+        random_state:
+            The seed or generator of stochastic solvers.
+
+        until:
+            The end time of the run the solver prepares for, infinite for a run that goes on until it is stopped.
+
+    Attributes:
+        solution_:
+            The optimal solution of the run.
+
+    Raises:
+        SolverError:
+            The solver has no objective, the sweep diverges or does not converge, or the stationarity condition has
+            no unique explicit solution.
+
+        ValidationError:
+            The sense is unknown, the grid has fewer than 2 nodes, or the relaxation is not in `(0, 1]`.
+
+    Examples:
+        >>> from typing import Annotated
+        >>> from skcomplex.core.elements import Element
+        >>> from skcomplex.core.laws import RateLaw
+        >>> from skcomplex.specs import ParameterSpec, StateSpec
+        >>> from skcomplex.core.system import System
+        >>> from skcomplex.engine.runtime import Simulation
+        >>> from skcomplex.engine.solvers import PontryaginSolver
+        >>> class PlantElement(Element):
+        ...     def __init__(
+        ...         self,
+        ...         label,
+        ...         x: Annotated[float, StateSpec()] = 1.0,
+        ...         u: Annotated[float, ParameterSpec(adaptive=True, bounds=(-10, 10))] = 0.0,
+        ...     ):
+        ...         super().__init__(label)
+        ...         self.x = x
+        ...         self.u = u
+        ...
+        ...     def _get_laws(self, s):
+        ...         return [RateLaw(s.x, s.u)]
+        >>> solver = PontryaginSolver(objective=lambda s: (s.plant.x**2 + s.plant.u**2) / 2)
+        >>> simulation = Simulation(System(elements=[PlantElement('plant')]), solver=solver)
+        >>> simulation = simulation.check().prepare().run(until=1.0)
+        >>> float(round(simulation.trajectory_.values('plant.u')[0], 3))
+        -0.762
+    """
+
+    @classmethod
+    def _get_support(cls) -> SupportSpec:
+        """Return the times, spaces, laws, noise, and orders the solver solves."""
+        return _CONTROL_SUPPORT
+
+    def __init__(
+        self,
+        objective: Cost | None = None,
+        *,
+        terminal: Cost = 0.0,
+        sense: str = 'min',
+        nodes: int = 201,
+        relaxation: float = 0.5,
+        tolerance: float = 1e-6,
+        max_iterations: int = 1000,
+        random_state: int | np.random.Generator | None = None,
+        until: float = math.inf,
+    ) -> None:
+        self.objective = objective
+        self.terminal = terminal
+        self.sense = sense
+        self.nodes = nodes
+        self.relaxation = relaxation
+        self.tolerance = tolerance
+        self.max_iterations = max_iterations
+        self.random_state = random_state
+        self.until = until
+
+    def _check(self) -> None:
+        """Raise an error for a missing objective, an unknown sense, too few nodes, or a bad relaxation.
+
+        Raises:
+            SolverError:
+                The solver has no objective.
+
+            ValidationError:
+                The sense is unknown, the grid has fewer than 2 nodes, or the relaxation is not in `(0, 1]`.
+        """
+        super()._check()
+        if not (isinstance(self.nodes, int) and self.nodes >= MIN_NODES):
+            error_msg = f'nodes = {self.nodes!r} must be an integer of at least {MIN_NODES}.'
+            raise ValidationError(error_msg)
+        if not 0 < self.relaxation <= 1:
+            error_msg = f'relaxation = {self.relaxation!r} must lie in (0, 1].'
+            raise ValidationError(error_msg)
+
+    def _costates(self, times: Array, states: Array, controls: Array, p: Array) -> Array:
+        """Return the costates on a grid, integrated backward from the horizon by Runge-Kutta steps."""
+        costates = np.zeros_like(states)
+        costates[-1] = self._final(times[-1], states[-1], controls[-1], p)
+        for k in range(times.size - 1, 0, -1):
+            h = times[k] - times[k - 1]
+            middle_time = times[k] - h / 2
+            x_middle, u_middle = (states[k] + states[k - 1]) / 2, (controls[k] + controls[k - 1]) / 2
+            lam = costates[k]
+            k1 = self._adjoint(times[k], states[k], controls[k], p, lam)
+            k2 = self._adjoint(middle_time, x_middle, u_middle, p, lam - h / 2 * k1)
+            k3 = self._adjoint(middle_time, x_middle, u_middle, p, lam - h / 2 * k2)
+            k4 = self._adjoint(times[k - 1], states[k - 1], controls[k - 1], p, lam - h * k3)
+            costates[k - 1] = lam - h / 6 * (k1 + 2 * k2 + 2 * k3 + k4)
+        return costates
+
+    def _update(self, times: Array, states: Array, controls: Array, costates: Array, p: Array) -> Array:
+        """Return the controls that minimise the Hamiltonian on a grid, within their bounds."""
+        lower, upper = self.control_problem_.bounds
+        values = self._law(times, states, controls, p, costates)
+        if self._bang_bang:
+            return np.where(values > 0, lower, upper)
+        return np.clip(values, lower, upper)
+
+    def _set_up_control(self, control_problem: _ControlProblemSpec) -> None:
+        """Derive the costate equations, the terminal costates, and the control law symbolically.
+
+        Args:
+            control_problem:
+                The optimal control problem.
+
+        Raises:
+            SolverError:
+                The stationarity condition has no unique explicit solution in the controls.
+        """
+        costates = [SR.var(f'_costate_{index}') for index in range(len(control_problem.states))]
+        hamiltonian = control_problem.running + sum(
+            (lam * rate for lam, rate in zip(costates, control_problem.rates, strict=True)),
+            SR(0),
+        )
+        controls = control_problem.control_symbols
+        gradients = [hamiltonian.diff(control) for control in controls]
+        self._bang_bang = self._is_linear(gradients, controls)
+        law = gradients if self._bang_bang else self._stationary_controls(gradients, controls)
+        self._adjoint: GridFunction = self._function(
+            control_problem,
+            [-hamiltonian.diff(item) for item in control_problem.states],
+            costates,
+        )
+        self._final: GridFunction = self._function(
+            control_problem,
+            [control_problem.terminal.diff(item) for item in control_problem.states],
+        )
+        self._law: GridFunction = self._function(control_problem, law, costates)
+
+    def _solve(self, x0: Array, p: Array, horizon: float) -> ControlSolutionSpec:
+        """Return the optimal solution by a forward-backward sweep with relaxed control updates.
+
+        Args:
+            x0:
+                The initial state vector.
+
+            p:
+                The parameter vector.
+
+            horizon:
+                The end of the run.
+
+        Returns:
+            solution:
+                The optimal solution with its costates.
+
+        Raises:
+            SolverError:
+                The sweep diverges or does not converge within the largest number of sweeps.
+        """
+        control_problem = self.control_problem_
+        times = np.linspace(0.0, horizon, self.nodes)
+        lower, upper = control_problem.bounds
+        start = np.clip(x0[control_problem.control_indices], lower, upper)
+        controls = np.tile(start, (times.size, 1))
+        initial = x0[control_problem.indices]
+        states = self._integrate(initial, controls, times, p)
+        costates = self._costates(times, states, controls, p)
+        weight, distance = self.relaxation, np.inf
+        for iteration in range(1, self.max_iterations + 1):
+            update = self._update(times, states, controls, costates, p)
+            gap = float(np.abs(update - controls).sum())
+            if gap >= distance:
+                weight = max(weight / 2, _MIN_RELAXATION)
+            distance = gap
+            new_controls = (1 - weight) * controls + weight * update
+            new_states = self._integrate(initial, new_controls, times, p)
+            new_costates = self._costates(times, new_states, new_controls, p)
+            if not (np.all(np.isfinite(new_states)) and np.all(np.isfinite(new_costates))):
+                error_msg = (
+                    f'The forward-backward sweep diverged at sweep {iteration}, as it does near unstable saddle paths. '
+                    'Lower the relaxation, shorten the horizon, or use the CollocationSolver solver.'
+                )
+                raise SolverError(error_msg)
+            pairs = [(new_controls, controls), (new_states, states), (new_costates, costates)]
+            controls, states, costates = new_controls, new_states, new_costates
+            if self._converged(self.tolerance, pairs):
+                named = {name: costates[:, index] for index, name in enumerate(control_problem.names)}
+                return self._solution(x0, times, states, controls, p, iterations=iteration, costates=named)
+        error_msg = (
+            f'The forward-backward sweep did not converge in {self.max_iterations} sweeps. Lower the relaxation, raise '
+            'max_iterations, or use the CollocationSolver solver.'
+        )
+        raise SolverError(error_msg)
+
+    @staticmethod
+    def _is_linear(gradients: list[Expression], controls: list[Expression]) -> bool:
+        """Check whether no gradient of the Hamiltonian by a control depends on any control."""
+        return all(gradient.diff(control).is_trivial_zero() for gradient in gradients for control in controls)
+
+    @staticmethod
+    def _stationary_controls(gradients: list[Expression], controls: list[Expression]) -> list[Expression]:
+        """Return the controls that make every gradient of the Hamiltonian vanish, solved symbolically."""
+        solutions = solve([gradient == 0 for gradient in gradients], *controls, solution_dict=True)
+        explicit = [
+            solution
+            for solution in solutions
+            if set(solution) == set(controls)
+            and all(not set(SR(value).variables()) & set(controls) for value in solution.values())
+        ]
+        if len(explicit) != 1:
+            error_msg = (
+                'The condition that the Hamiltonian is stationary in the controls has no unique explicit solution. Use '
+                'the CollocationSolver solver, which needs no such solution.'
+            )
+            raise SolverError(error_msg)
+        return [SR(explicit[0][control]) for control in controls]
+
+    @staticmethod
+    def _converged(tolerance: float, pairs: list[tuple[npt.NDArray[np.float64], npt.NDArray[np.float64]]]) -> bool:
+        """Check whether every new array differs from the old one by at most a relative tolerance in the sum norm."""
+        return all(
+            np.abs(new - old).sum() <= tolerance * max(np.abs(new).sum(), np.finfo(float).tiny) for new, old in pairs
+        )