scikit-complexity 0.2.0__py3-none-any.whl

skcomplex/engine/solvers/_collocation.py

@@ -0,0 +1,334 @@
+"""Solve optimal control problems by direct collocation, a transcription solved by nonlinear programming."""
+
+# Author: Georgios Douzas <gdouzas@icloud.com>
+# License: MIT
+
+import math
+
+import numpy as np
+from scipy.optimize import Bounds, minimize
+
+from ...core.laws import 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
+
+SCHEMES = ('trapezoidal', 'hermite-simpson')
+METHODS = ('SLSQP', 'trust-constr')
+MIN_NODES = 2
+
+
+class CollocationSolver(OptimalControlSolver):
+    """Optimal control by direct collocation, a transcription of the problem solved by nonlinear programming.
+
+    The states and the controls at the points of a grid are the variables of a nonlinear programme. The trapezoidal
+    scheme requires the states of neighbouring nodes to differ by the trapezoidal integral of their rates, and the
+    Hermite-Simpson scheme adds a midpoint to every interval, with a cubic interpolation of the states and Simpson's
+    integral. The cost is the quadrature of the running cost over the points plus the terminal cost. The controls lie
+    within their bounds and the states of nonnegative domains stay nonnegative. SciPy's `minimize` solves the programme
+    with the gradient of the cost and the Jacobian of the constraints derived symbolically. An objective to maximise is
+    minimised with its sign changed.
+
+    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.
+
+        scheme:
+            `'trapezoidal'` or `'hermite-simpson'`.
+
+        method:
+            The method of `scipy.optimize.minimize`, `'SLSQP'` or `'trust-constr'`.
+
+        tolerance:
+            The tolerance of the optimiser.
+
+        max_iterations:
+            The largest number of iterations of the optimiser.
+
+        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:
+        ValidationError:
+            The sense, the scheme, or the method is unknown, or the grid has fewer than 2 nodes.
+
+        SolverError:
+            The solver has no objective, or the optimiser does not converge.
+
+    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 CollocationSolver
+        >>> 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)]
+        >>> cost = lambda s: (s.plant.x**2 + s.plant.u**2) / 2
+        >>> solver = CollocationSolver(objective=cost, nodes=51, scheme='hermite-simpson')
+        >>> 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 SupportSpec(
+            time={ContinuousTime},
+            space={None},
+            laws={RateLaw, FluxLaw, EquationLaw},
+            orders=frozenset({OrderKind.INTEGER}),
+        )
+
+    def __init__(
+        self,
+        objective: Cost | None = None,
+        *,
+        terminal: Cost = 0.0,
+        sense: str = 'min',
+        nodes: int = 101,
+        scheme: str = 'hermite-simpson',
+        method: str = 'SLSQP',
+        tolerance: float = 1e-10,
+        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.scheme = scheme
+        self.method = method
+        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, scheme, or method, or too few nodes.
+
+        Raises:
+            SolverError:
+                The solver has no objective.
+
+            ValidationError:
+                The sense, the scheme, or the method is unknown, or the grid has fewer than 2 nodes.
+        """
+        super()._check()
+        if self.scheme not in SCHEMES:
+            error_msg = f'Unknown collocation scheme {self.scheme!r}. Use one of {", ".join(SCHEMES)}.'
+            raise ValidationError(error_msg)
+        if self.method not in METHODS:
+            error_msg = f'Unknown optimiser method {self.method!r}. Use one of {", ".join(METHODS)}.'
+            raise ValidationError(error_msg)
+        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)
+
+    def _unpack(self, z: Array, points: int) -> tuple[Array, Array]:
+        """Return the states and the controls of the points from the vector of variables."""
+        n, m = len(self.control_problem_.states), len(self.control_problem_.controls)
+        return z[: points * n].reshape(points, n), z[points * n :].reshape(points, m)
+
+    def _cost_and_gradient(self, z: Array, points: Array, weights: Array, p: Array) -> tuple[float, Array]:
+        """Return the transcribed cost and its gradient by the variables."""
+        x, u = self._unpack(z, points.size)
+        values = self.running_(points, x, u, p)[:, 0]
+        cost = float(weights @ values + self.terminal_(points[-1], x[-1], u[-1], p)[0])
+        gradient_x = weights[:, None] * self._running_states(points, x, u, p)
+        gradient_x[-1] += self._terminal_states(points[-1], x[-1], u[-1], p)
+        gradient_u = weights[:, None] * self._running_controls(points, x, u, p)
+        return cost, np.concatenate([gradient_x.ravel(), gradient_u.ravel()])
+
+    def _defects(self, z: Array, points: Array, coefficients: tuple[Array, Array], x0: Array, p: Array) -> Array:
+        """Return the defects of the collocation conditions and of the initial state."""
+        x, u = self._unpack(z, points.size)
+        states, rates = coefficients
+        defects = states @ x + rates @ self.rates_(points, x, u, p)
+        return np.concatenate([defects.T.ravel(), x[0] - x0])
+
+    def _jacobian(self, z: Array, points: Array, coefficients: tuple[Array, Array], p: Array) -> Array:
+        """Return the Jacobian of the defects by the variables, from the symbolic Jacobians of the rates."""
+        x, u = self._unpack(z, points.size)
+        states, rates = coefficients
+        n, m, size = x.shape[1], u.shape[1], points.size
+        by_states = self._rates_states(points, x, u, p).reshape(size, n, n)
+        by_controls = self._rates_controls(points, x, u, p).reshape(size, n, m)
+        rows = states.shape[0]
+        jacobian = np.zeros((rows * n + n, size * (n + m)))
+        for i in range(n):
+            block = slice(i * rows, (i + 1) * rows)
+            for j in range(n):
+                jacobian[block, j : size * n : n] = rates * by_states[:, i, j] + (states if i == j else 0.0)
+            for j in range(m):
+                jacobian[block, size * n + j :: m] = rates * by_controls[:, i, j]
+            jacobian[rows * n + i, i] = 1.0
+        return jacobian
+
+    def _bounds(self, points: int) -> Bounds:
+        """Return the bounds of the variables: the lower bounds of the states and the bounds of the controls."""
+        lower, upper = self.control_problem_.bounds
+        state_lower = np.tile(self.control_problem_.lower, points)
+        state_upper = np.full(state_lower.size, np.inf)
+        return Bounds(
+            np.concatenate([state_lower, np.tile(lower, points)]),
+            np.concatenate([state_upper, np.tile(upper, points)]),
+        )
+
+    def _set_up_control(self, control_problem: _ControlProblemSpec) -> None:
+        """Derive the Jacobians of the rates and the gradients of the costs symbolically.
+
+        Args:
+            control_problem:
+                The optimal control problem.
+        """
+        states, controls = control_problem.states, control_problem.control_symbols
+        rates, running = control_problem.rates, control_problem.running
+        self._rates_states: GridFunction = self._function(
+            control_problem,
+            [rate.diff(item) for rate in rates for item in states],
+        )
+        self._rates_controls: GridFunction = self._function(
+            control_problem,
+            [rate.diff(item) for rate in rates for item in controls],
+        )
+        self._running_states: GridFunction = self._function(control_problem, [running.diff(item) for item in states])
+        self._running_controls: GridFunction = self._function(
+            control_problem,
+            [running.diff(item) for item in controls],
+        )
+        self._terminal_states: GridFunction = self._function(
+            control_problem,
+            [control_problem.terminal.diff(item) for item in states],
+        )
+
+    def _solve(self, x0: Array, p: Array, horizon: float) -> ControlSolutionSpec:
+        """Return the optimal solution of the transcribed nonlinear programme.
+
+        Args:
+            x0:
+                The initial state vector.
+
+            p:
+                The parameter vector.
+
+            horizon:
+                The end of the run.
+
+        Returns:
+            solution:
+                The optimal solution at the points of the grid.
+
+        Raises:
+            SolverError:
+                The optimiser does not converge.
+        """
+        control_problem = self.control_problem_
+        transcribe = self._trapezoidal if self.scheme == 'trapezoidal' else self._hermite_simpson
+        points, states, rates, weights = transcribe(np.linspace(0.0, horizon, self.nodes))
+        lower, upper = control_problem.bounds
+        initial = x0[control_problem.indices]
+        guess_controls = np.tile(np.clip(x0[control_problem.control_indices], lower, upper), (points.size, 1))
+        with np.errstate(all='ignore'):
+            guess_states = np.maximum(self._integrate(initial, guess_controls, points, p), control_problem.lower)
+        if not np.all(np.isfinite(guess_states)):
+            guess_states = np.tile(initial, (points.size, 1))
+        coefficients = (states, rates)
+        constraint: dict[str, object] = {
+            'type': 'eq',
+            'fun': lambda z: self._defects(z, points, coefficients, initial, p),
+            'jac': lambda z: self._jacobian(z, points, coefficients, p),
+        }
+        result = minimize(
+            lambda z: self._cost_and_gradient(z, points, weights, p),
+            np.concatenate([guess_states.ravel(), guess_controls.ravel()]),
+            jac=True,
+            bounds=self._bounds(points.size),
+            constraints=[constraint],
+            method=self.method,
+            tol=self.tolerance,
+            options={'maxiter': self.max_iterations},
+        )
+        if not result.success:
+            error_msg = (
+                f'The collocation programme did not converge: {result.message} Raise max_iterations, change the '
+                'number of nodes, or try the other scheme or method.'
+            )
+            raise SolverError(error_msg)
+        x, u = self._unpack(result.x, points.size)
+        return self._solution(x0, points, x, u, p, iterations=int(result.nit))
+
+    @staticmethod
+    def _trapezoidal(times: Array) -> tuple[Array, Array, Array, Array]:
+        """Return the points, the defect coefficients of states and rates, and the quadrature weights of trapezoids."""
+        steps = np.diff(times)
+        rows = np.arange(steps.size)
+        states = np.zeros((steps.size, times.size))
+        rates = np.zeros((steps.size, times.size))
+        states[rows, rows], states[rows, rows + 1] = -1.0, 1.0
+        rates[rows, rows], rates[rows, rows + 1] = -steps / 2, -steps / 2
+        weights = np.zeros(times.size)
+        np.add.at(weights, rows, steps / 2)
+        np.add.at(weights, rows + 1, steps / 2)
+        return times, states, rates, weights
+
+    @staticmethod
+    def _hermite_simpson(times: Array) -> tuple[Array, Array, Array, Array]:
+        """Return the points with midpoints, the defect coefficients of states and rates, and Simpson's weights."""
+        steps = np.diff(times)
+        points = np.empty(2 * times.size - 1)
+        points[0::2], points[1::2] = times, times[:-1] + steps / 2
+        intervals = np.arange(steps.size)
+        first, middle, last = 2 * intervals, 2 * intervals + 1, 2 * intervals + 2
+        interpolation, quadrature = 2 * intervals, 2 * intervals + 1
+        states = np.zeros((2 * steps.size, points.size))
+        rates = np.zeros((2 * steps.size, points.size))
+        states[interpolation, middle], states[interpolation, first], states[interpolation, last] = 1.0, -0.5, -0.5
+        rates[interpolation, first], rates[interpolation, last] = -steps / 8, steps / 8
+        states[quadrature, last], states[quadrature, first] = 1.0, -1.0
+        rates[quadrature, first], rates[quadrature, middle], rates[quadrature, last] = (
+            -steps / 6,
+            -2 * steps / 3,
+            -steps / 6,
+        )
+        weights = np.zeros(points.size)
+        np.add.at(weights, first, steps / 6)
+        np.add.at(weights, middle, 2 * steps / 3)
+        np.add.at(weights, last, steps / 6)
+        return points, states, rates, weights

skcomplex/engine/solvers/_convolution.py

@@ -0,0 +1,155 @@
+"""Base class of explicit schemes that weigh the displacements of the states from the start of a run."""
+
+# Author: Georgios Douzas <gdouzas@icloud.com>
+# License: MIT
+
+import math
+from abc import abstractmethod
+
+import numpy as np
+
+from ...errors import SolverError
+from ...types import Array
+from ._caputo import CaputoSolver
+from ._rows import _INITIAL_ROWS, _Rows
+
+_SLOWEST_MODE = 1e-7
+
+
+_FASTEST_MODE = 60.0
+
+
+_MODE_SPACING = 0.35
+
+
+class ConvolutionSolver(CaputoSolver):
+    """Base class of explicit schemes that weigh the displacements of the states from the start of a run.
+
+    A step sets `g dt^alpha f(t_(k-1), x_(k-1)) = sum_(j=0)^k c_j (x_(k-j) - x_0)` with `c_0 = 1`. With a memory
+    window, the displacements older than the window enter through exponential modes fitted to the weights.
+    """
+
+    def __init__(
+        self,
+        dt: float = 0.01,
+        memory: float | None = None,
+        *,
+        random_state: int | np.random.Generator | None = None,
+        until: float = math.inf,
+    ) -> None:
+        self.dt = dt
+        self.memory = memory
+        self.random_state = random_state
+        self.until = until
+
+    @abstractmethod
+    def _weights(self, alpha: Array, count: int) -> Array:
+        """Return the weights `c_0` to `c_count`, one column per order.
+
+        Args:
+            alpha:
+                The distinct orders.
+
+            count:
+                The largest lag.
+
+        Returns:
+            weights:
+                The weights, one row per lag.
+
+        Raises:
+            NotImplementedError:
+                The subclass does not implement the method.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def _tail(self, alpha: Array, rates: Array, window: int) -> Array:
+        """Return the coefficients of the exponential modes that give the weights past a window, one column per order.
+
+        Args:
+            alpha:
+                The distinct orders.
+
+            rates:
+                The decay rates of the modes per step.
+
+            window:
+                The number of steps of the window.
+
+        Returns:
+            coefficients:
+                The coefficients, one row per mode.
+
+        Raises:
+            NotImplementedError:
+                The subclass does not implement the method.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def _gain(self, alpha: Array) -> Array:
+        """Return the factor `g` of the drift of every order.
+
+        Args:
+            alpha:
+                The distinct orders.
+
+        Returns:
+            gain:
+                The factor of every order.
+
+        Raises:
+            NotImplementedError:
+                The subclass does not implement the method.
+        """
+        raise NotImplementedError
+
+    def _restart(self, t: float, x: Array) -> None:
+        """Start the history of displacements, the weights, and the exponential modes of the tail."""
+        del t
+        if self.memory is not None and not self.memory > 0:
+            error_msg = f'memory = {self.memory!r} must be positive or None.'
+            raise SolverError(error_msg)
+        self._window = None if self.memory is None else max(1, round(self.memory / self.dt))
+        self._rows = _Rows(np.zeros(x.size), self._window)
+        self._table = self._weights(self._unique, _INITIAL_ROWS if self._window is None else self._window)
+        self._scale = self._per_state(self._gain(self._unique) * self.dt**self._unique)
+        if self._window is not None:
+            rates = self._tail_rates(self._window)
+            self._decay = np.exp(-rates)[:, None]
+            self._modes = self._tail(self._unique, rates, self._window)[:, self._groups]
+            self._amplitudes = np.zeros((rates.size, x.size))
+
+    def _carry(self, x: Array) -> None:
+        """Replace the last displacement of the states of order 1 by the given state."""
+        ordinary = ~self._fractional
+        self._state[ordinary] = x[ordinary]
+        self._rows.row(self._step_index)[ordinary] = x[ordinary] - self._initial[ordinary]
+
+    def _history_sum(self, step: int) -> Array:
+        """Return the weighted sum of the displacements before a step, with the tail of a memory window."""
+        window = self._window
+        lags = step - 1 if window is None else min(step - 1, window)
+        if lags >= self._table.shape[0]:
+            self._table = self._weights(self._unique, 2 * lags)
+        total = self._convolve(self._table[lags:0:-1], self._rows.rows(step - lags, step)) if lags else 0.0
+        if window is None or step <= window + 1:
+            return np.asarray(total + np.zeros(self._state.size))
+        self._amplitudes = self._decay * self._amplitudes + self._rows.row(step - window - 1)
+        return np.asarray(total + np.sum(self._modes * self._amplitudes, axis=0))
+
+    def _step(self, p: Array) -> None:
+        """Advance the displacements by one explicit step."""
+        step = self._step_index + 1
+        time = self._origin + self._step_index * self.dt
+        drift = self._drift()(time, self._state, p)
+        displacement = self._scale * drift - self._history_sum(step)
+        self._state = self._bounded(time + self.dt, self._initial + displacement, p)
+        self._rows.append(self._state - self._initial)
+
+    @staticmethod
+    def _tail_rates(window: int) -> Array:
+        """Return the decay rates per step of the exponential modes that hold the history older than a window."""
+        exponents = np.arange(math.log(_SLOWEST_MODE), math.log(_FASTEST_MODE / (window + 1)), _MODE_SPACING)
+        return np.exp(exponents)