Skip to content

Model predictive control of Duffing, Van der Pol, and a bilinear DC motor through Koopman surrogates

A plant that a discrete-time controller steers with model predictive control on a fitted surrogate.

  • Name: koopman_model_predictive_control
  • Field: engineering/control
  • Formalism: co-simulation
  • References: Korda and Mezić 2018; Brunton et al. 2022

Select this system

from skcomplex.systems import create_simulation, create_system, describe

system = create_system('koopman_model_predictive_control')
entry = describe('koopman_model_predictive_control')
trajectory = create_simulation('koopman_model_predictive_control').prepare().run().trajectory_

Factory

@register_system(
    field='engineering/control',
    tags=('control', 'model predictive control', 'koopman', 'co-simulation', 'surrogate'),
    references=('Korda and Mezić 2018', 'Brunton et al. 2022'),
    defaults={'until': 100.0, 'dt': 1.0},
    observables=_observables,
)
def koopman_model_predictive_control(
    plant: str = 'duffing',
    *,
    controlled: bool = True,
    dictionary: TransformerMixin | None = None,
    x: float = 0.5,
    y: float = 0.0,
    n_runs: int = 100,
    duration: int = 1000,
    random_state: int | np.random.Generator | None = 0,
) -> System:
    """Return a plant that a discrete-time controller steers with model predictive control on a fitted surrogate.

    The factory runs the plant under a random binary input at the limits of its range, drawn every sample, and fits
    `ModeDecompositionAnalysis` with the input to the runs. The controller lifts the states with the dictionary,
    condenses the quadratic tracking cost over the horizon with the matrices of the surrogate, solves the bounded least
    squares problem with `scipy.optimize.lsq_linear`, and sets the input to the first value of the solution. The
    plant reads the input through a role, and `CoSimulationSolver` holds it between samples. The plants are those of
    figure 6.6 of the review of Brunton et al.: the forced Duffing oscillator
    `x'' = a x + b x^3 + d x' + f cos(w t) + (3 + cos(x)) u`, the Van der Pol oscillator
    `x'' = -x - mu (x^2 - 1) x' + u`, and the bilinear DC motor, whose second variable is the output that tracks the
    reference. The motor reads four times the input, as Korda and Mezić scale it. One time unit of the system is one
    sample of the controller.

    Args:
        plant:
            `'duffing'`, `'van_der_pol'`, or `'motor'`.

        controlled:
            Whether the controller solves the tracking problem, or holds the input at 0 with no fit.

        dictionary:
            A scikit-learn transformer that computes the observables and keeps the variables among them, or `None` for
            the monomials up to degree 3.

        x:
            The initial value of the first variable.

        y:
            The initial value of the second variable.

        n_runs:
            The number of training runs.

        duration:
            The number of samples of every training run.

        random_state:
            The seed or generator of the training runs.

    Returns:
        system:
            The system, labelled by the plant, with the variables `plant/state.x` and `plant/state.y`, the input
            `controller/input.u`, and the observable `reference`.

    Raises:
        ValidationError:
            The plant is unknown.

    Examples:
        >>> from skcomplex.systems.engineering.control import koopman_model_predictive_control
        >>> system = koopman_model_predictive_control('motor', controlled=False)
        >>> [subsystem.label for subsystem in system.systems]
        ['controller', 'plant']
    """
    if plant not in PLANTS:
        error_msg = f'The plant {plant!r} is unknown. Use one of {", ".join(PLANTS)}.'
        raise ValidationError(error_msg)
    case = PLANTS[plant]
    problem = _fit_problem(case, dictionary, n_runs, duration, random_state) if controlled else None
    controller = MapInteraction(
        'policy',
        target='controller/input',
        state='u',
        sources=['plant/state'],
        function=functools.partial(_hold, problem=problem),
    )
    return System(
        plant,
        systems=[
            System('controller', elements=[ControlInputElement('input')], time=DiscreteTime()),
            System('plant', elements=[PlantStateElement('state', x=x, y=y)], time=ContinuousTime()),
        ],
        interactions=[controller, *_plant_interactions(case, 'plant/state', 'controller/input')],
    )

Components

Elements:

  • ControlInputElement, from Element
  • PlantStateElement, from Element

Interactions:

  • MapInteraction
  • RateInteraction
class MapInteraction(Interaction):
    """Map of a state of discrete time, whose next value is a function, or the state plus it.

    In the mode `'replace'` the next value of the state is `f`, and in the mode `'increment'` it is the state plus `f`.
    Under the `'synchronous'` schedule every member of a bound group updates at once, from the states at the start of
    the step. Under the `'sequential'` schedule the members update one after the other, in their order, and under the
    `'random-sequential'` schedule in a new random order every step, each member reading the values the members before
    it set. Under the `'member'` schedule of the agent form, every member updates in its own activation, together with
    the rules of its element and the other member interactions, in the order the system lists them.

    In the `'symbolic'` form, `function(namespace)` returns a symbolic expression, or a vector of them for a vector
    state, where `namespace.target` holds the states and parameters of the target, `namespace.sources` the namespaces
    of the sources, and `namespace.t` the time, and may hold draws such as `normal('noise')`. In the `'agent'` form,
    the target is an agent group and `function(agent, system)` returns the value as a number, or as an array of the
    shape of a vector state, so it can read the agent, its neighbours, and the history and storage of the system view.

    Args:
        label:
            The label of the interaction, unique among the parts of its system.

        target:
            The path of the element or group whose state the map sets.

        state:
            The name of the state of the target.

        function:
            The function `f(namespace)` in the symbolic form, or `f(agent, system)` in the agent form.

        sources:
            The paths of the elements, or members of groups such as `cells[0]`, that the function reads.

        mode:
            `'replace'` or `'increment'`.

        schedule:
            `'synchronous'`, `'sequential'`, `'random-sequential'`, or `'member'` for the agent form.

        form:
            `'symbolic'` for update laws of the states of the target, or `'agent'` for a rule over an agent group.

        phase:
            The phase of the step the agent form runs in, or `None` for the default phase.

    Raises:
        ValidationError:
            The mode, the schedule, or the form is unknown, the symbolic form has the schedule `'member'`, the
            function is not callable, or the agent form returns a value that does not match the state, raised by
            `run` of the simulation.

    Examples:
        >>> from skcomplex.core.components.elements import Element
        >>> from skcomplex.core.components.interactions import MapInteraction
        >>> from skcomplex.core.components.spaces import DiscreteTime
        >>> from skcomplex.core.components.system import System
        >>> from skcomplex.core.engine.runtime import Simulation
        >>> from skcomplex.core.specs import StateSpec
        >>> from typing import Annotated, SupportsFloat
        >>> class ValueElement(Element):
        ...     def __init__(self, label, x: Annotated[float, StateSpec()] = 0.5):
        ...         super().__init__(label)
        ...         self.x = x
        >>> logistic = MapInteraction(
        ...     'map', target='value', state='x', function=lambda s: 3 * s.target.x * (1 - s.target.x)
        ... )
        >>> system = System(elements=[ValueElement('value')], interactions=[logistic], time=DiscreteTime())
        >>> Simulation(system).prepare().run(until=2).trajectory_.values('value.x').tolist()
        [0.5, 0.75, 0.5625]
        >>> agents = MapInteraction('map', target='walkers', state='x', function=lambda a, m: a.x + 1, form='agent')
        >>> type(agents.prepare().run(SimpleNamespace()).laws_[0]).__name__
        'RuleLaw'
        >>> from skcomplex.core.components.elements import GroupElement
        >>> class WeightsElement(Element):
        ...     def __init__(self, label, visual=None, *, w: Annotated[object, StateSpec(vector=True)] = 0.0):
        ...         super().__init__(label, visual=visual)
        ...         self.w = w
        >>> group = GroupElement('weights', element=WeightsElement, size=1, values={'w': [1.0, 2.0]})
        >>> doubling = MapInteraction(
        ...     'double', target='weights', state='w', function=lambda a, m: [2 * x for x in a.w], form='agent'
        ... )
        >>> trajectory = Simulation(System(elements=[group], interactions=[doubling])).prepare().run(until=2)
        >>> trajectory.trajectory_.agents('weights', 2).values['w'].tolist()
        [[4.0, 8.0]]
    """

    law_classes: ClassVar[tuple[type[BaseLaw], ...]] = (RuleLaw, UpdateLaw)

    def __init__(
        self,
        label: str,
        *,
        target: Annotated[str, RoleSpec(elements=(ElementKind.ANY,), states=('state',))],
        state: str,
        function: Callable[..., object],
        sources: Annotated[Sequence[str], RoleSpec(variadic=True, elements=(ElementKind.ANY,))] = (),
        mode: str = 'replace',
        schedule: str = 'synchronous',
        form: str = 'symbolic',
        phase: str | None = None,
    ) -> None:
        super().__init__(label)
        self.target = target
        self.state = state
        self.function = function
        self.sources = sources
        self.mode = mode
        self.schedule = schedule
        self.form = form
        self.phase = phase

    def _check(self) -> None:
        """Raise an error for an unknown mode, schedule, or form, or a function that cannot be called.

        Raises:
            ValidationError:
                The mode, the schedule, or the form is unknown, the symbolic form has the schedule `'member'`, or the
                function is not callable.
        """
        super()._check()
        self._check_variant('mode', self.mode, MODES, (BaseSpec,))
        for name, value, variants in (
            ('schedule', self.schedule, SCHEDULES),
            ('form', self.form, FORMS),
        ):
            if value not in variants:
                error_msg = f'The {name} of map {self.label!r} is {value!r}. Use one of {", ".join(variants)}.'
                raise ValidationError(error_msg)
        if self.schedule == 'member' and self.form != 'agent':
            error_msg = (
                f'The map {self.label!r} updates every member in its activation, which only the agent form does.'
            )
            raise ValidationError(error_msg)
        if not callable(self.function):
            error_msg = f'The function of map {self.label!r} must be callable, but it is {self.function!r}.'
            raise ValidationError(error_msg)

    def _apply_mode(self, current: Value, value: Value) -> Value:
        """Return the next value of a state from its current value and the value of the function, by the mode."""
        if self.mode == 'increment':
            return current + value
        if self.mode == 'replace':
            return value
        raise self._get_variant_error(self.mode, 'mode', '_apply_mode')

    def _get_vector(self, agent: AgentViewProtocol, system: SystemViewProtocol) -> Value:
        """Return the next value of a vector state of an agent."""
        current = np.asarray(getattr(agent, self.state), dtype=np.float64)
        value = np.asarray(self.function(agent, system), dtype=np.float64)
        if value.shape != current.shape:
            error_msg = (
                f'The function of map {self.label!r} returns the shape {value.shape} for the state {self.state!r} of '
                f'the shape {current.shape}. Return an array of the shape of the state.'
            )
            raise ValidationError(error_msg)
        return self._apply_mode(current, value)

    def _get_value(self, agent: AgentViewProtocol, system: SystemViewProtocol) -> Value:
        """Return the next value of the state of an agent, a number or the array of a vector state."""
        if np.ndim(getattr(agent, self.state)) > 0:
            return self._get_vector(agent, system)
        found = self.function(agent, system)
        if not isinstance(found, SupportsFloat):
            error_msg = f'The function of map {self.label!r} returns {found!r} for an agent. Return a number.'
            raise ValidationError(error_msg)
        value = float(found)
        return self._apply_mode(float(getattr(agent, self.state)), value)

    def _set_value(self, agent: AgentViewProtocol, value: Value) -> None:
        """Set the state of an agent to its next value."""
        setattr(agent, self.state, value)

    def _update_member(self, agent: AgentViewProtocol, system: SystemViewProtocol) -> None:
        """Set the state of one member to its next value."""
        if agent.alive:
            self._set_value(agent, self._get_value(agent, system))

    def update(self, system: SystemViewProtocol) -> None:
        """Set the state of every agent of the target to its next value, by the schedule.

        Args:
            system:
                The system view of the interaction.
        """
        self._update_agents(system, self.schedule, self._get_value, self._set_value)

    def _get_order(self) -> object:
        """Return the order of the members of the symbolic form: none, their own order, or a uniform draw."""
        if self.schedule == 'synchronous':
            return None
        if self.schedule == 'sequential':
            return SR(0)
        return uniform(re.sub(r'\W', '_', f'order_{self.label}'))

    def _get_laws(self, namespace: SimpleNamespace) -> list[BasePart]:
        """Return the update of the state to its next value, or the rule of the agent form.

        Args:
            namespace:
                The namespace of the symbols the laws can use, with `namespace.target` and `namespace.sources`.

        Returns:
            laws:
                The update law of the state, or the rule that runs the map of the agents.
        """
        if self.form == 'agent' and self.schedule == 'member':
            return [RuleLaw(self._update_member, phase=self.phase, members=self._get_role('target'))]
        if self.form == 'agent':
            return [RuleLaw(self.update, phase=self.phase)]
        value = getattr(self._get_member(namespace, 'target'), self.state)
        mapped = cast(Value, self.function(namespace))
        next_value = self._apply_mode(value, mapped)
        return [UpdateLaw(value, next_value, increment=self.mode == 'increment', order=self._get_order())]

class RateInteraction(Interaction):
    """Rate `rate(ns)` and noise `noise(ns) dW` added to the highest time derivative of a state of a target.

    On a state of order 1 the rate is its first derivative, and on a state of order `k` it adds to the `k`-th
    derivative, so rates on one state superpose. With a divisor `d` the rate and the noise are divided by `d`, the name
    of a state or parameter of the target, such as an inertia, a friction, or a capacitance, which backends can read.

    A rate, a noise, or a divisor function takes the namespace of the interaction, with the namespaces `target` and
    `sources` of the bound elements, the time `t`, and the symbols of the parameters a subclass declares. A number
    stands for a constant, and a number on a vector state applies to every component. The noise is a Wiener process per
    component and member by default, a process of a `WienerProcessSpec` to share or correlate it across states and
    members, or a fractional Brownian motion of a Hurst index.

    Args:
        label:
            The label of the interaction, unique among the parts of its system.

        target:
            The path of the element or group whose state the rate changes.

        state:
            The name of the state of the target that the rate changes.

        rate:
            A function of the namespace that returns the rate, or a number.

        sources:
            The paths of the elements or groups the rate and the noise read but do not change.

        noise:
            A function of the namespace that returns the amplitude of the noise, a number, or `None` for no noise.

        divisor:
            The name of the state or parameter of the target that divides the rate and the noise, a function of the
            namespace of the target that returns the divisor, such as a moment of inertia `m l^2`, or `None`.

        process:
            The Wiener processes that drive the noise, or `None` for one independent process per component and member.

        index:
            The process of `process` that drives the noise, or `None` for every process, one per component of a vector
            state.

        convention:
            The stochastic calculus of the noise, `'ito'`, `'stratonovich'`, or `'anti-ito'`.

        hurst:
            The Hurst index in `(0, 1)` of a fractional Brownian motion that drives the noise, a number or a function
            of the namespace that returns a parameter, or `None` for a Wiener process.

        conserved:
            Whether the noise is the divergence of a random flux on the links of the lattice of a field, so it keeps the
            total of the state.

    Raises:
        ValidationError:
            The state is not a name, a function is neither callable nor a number, the convention is unknown, the Hurst
            index is outside `(0, 1)`, or the noise options cannot combine.

    Examples:
        >>> from skcomplex.core.components.elements import CompartmentElement
        >>> from skcomplex.core.components.interactions import RateInteraction
        >>> from skcomplex.core.components.system import System
        >>> from skcomplex.core.engine.runtime import Simulation
        >>> decay = RateInteraction('decay', target='stock', state='n', rate=lambda s: -0.5 * s.target.n)
        >>> system = System(elements=[CompartmentElement('stock', n=8.0)], interactions=[decay])
        >>> trajectory = Simulation(system, dt=0.01).prepare().run(until=2.0).trajectory_
        >>> round(float(trajectory.values('stock.n')[-1]), 3)
        2.943
    """

    law_classes: ClassVar[tuple[type[BaseLaw], ...]] = (RateLaw, NoiseLaw, FractionalNoiseLaw)

    def __init__(
        self,
        label: str,
        *,
        target: Annotated[str, RoleSpec(elements=(ElementKind.ANY,), states=('state', 'divisor'))],
        state: str = 'x',
        rate: Callable[[SimpleNamespace], object] | float = 0.0,
        sources: Annotated[Sequence[str], RoleSpec(variadic=True, elements=(ElementKind.ANY,))] = (),
        noise: Callable[[SimpleNamespace], object] | float | None = None,
        divisor: str | Callable[[SimpleNamespace], object] | None = None,
        process: Annotated[WienerProcessSpec | None, RoleSpec(noise=True)] = None,
        index: int | None = None,
        convention: str = 'ito',
        hurst: Callable[[SimpleNamespace], object] | float | None = None,
        conserved: bool = False,
    ) -> None:
        super().__init__(label)
        self.target = target
        self.state = state
        self.rate = rate
        self.sources = sources
        self.noise = noise
        self.divisor = divisor
        self.process = process
        self.index = index
        self.convention = convention
        self.hurst = hurst
        self.conserved = conserved

    def _check_functions(self) -> None:
        """Raise an error for a state that is not a name, or a function that is neither callable nor a number."""
        if not isinstance(self.state, str) or not self.state.isidentifier():
            error_msg = (
                f'{type(self).__name__} {self.label!r} changes the state {self.state!r}, which is not a name of a '
                'state.'
            )
            raise ValidationError(error_msg)
        for name in ('rate', 'noise'):
            value = getattr(self, name)
            if not (value is None or callable(value) or isinstance(value, int | float)):
                error_msg = f'The {name} of {self.label!r} is {value!r}. Give a function of the namespace or a number.'
                raise ValidationError(error_msg)
        if not (
            self.divisor is None
            or callable(self.divisor)
            or (isinstance(self.divisor, str) and self.divisor.isidentifier())
        ):
            error_msg = (
                f'The divisor of {self.label!r} is {self.divisor!r}. Give the name of a state or parameter, or a '
                'function of the namespace of the target.'
            )
            raise ValidationError(error_msg)

    def _check_noise(self) -> None:
        """Raise an error for an unknown convention, a Hurst index outside `(0, 1)`, or noise options that conflict."""
        if self.convention not in CONVENTIONS:
            error_msg = f'Unknown noise convention {self.convention!r}. Use one of {", ".join(CONVENTIONS)}.'
            raise ValidationError(error_msg)
        if not (
            self.hurst is None or callable(self.hurst) or (isinstance(self.hurst, int | float) and 0 < self.hurst < 1)
        ):
            error_msg = f'The Hurst index of {self.label!r} is {self.hurst!r}. Give a number in (0, 1) or a function.'
            raise ValidationError(error_msg)
        if self.index is not None and self.process is None:
            error_msg = f'{self.label!r} picks the process {self.index!r} but binds no processes. Give `process`.'
            raise ValidationError(error_msg)
        options = [
            name
            for name, given in (('hurst', self.hurst), ('process', self.process), ('conserved', self.conserved))
            if given
        ]
        if len(options) > 1:
            error_msg = f'The noise of {self.label!r} takes {" and ".join(options)}. Give at most one of them.'
            raise ValidationError(error_msg)

    def _check(self) -> None:
        """Raise an error for a state, functions, or noise options the interaction cannot take.

        Raises:
            ValidationError:
                The type declares states, the state is not a name, a function is neither callable nor a number, the
                convention is unknown, the Hurst index is outside `(0, 1)`, or the noise options cannot combine.
        """
        super()._check()
        self._check_functions()
        self._check_noise()

    def _evaluate(self, function: Callable[[SimpleNamespace], object] | float, namespace: SimpleNamespace) -> object:
        """Return the value of a function of the namespace, or a number as it is."""
        return function(namespace) if callable(function) else function

    def _broadcast(self, value: object, state: object) -> object:
        """Return a value with one component per component of a vector state, repeating a number."""
        if isinstance(state, FreeModuleElement) and not isinstance(value, FreeModuleElement | list | tuple):
            return vector(SR, [value] * len(state))
        if isinstance(value, list | tuple):
            return vector(SR, value)
        return value

    def _divide(self, value: object, namespace: SimpleNamespace) -> object:
        """Return a value divided by the divisor of the target, or the value without a divisor."""
        if self.divisor is None:
            return value
        divisor = (
            self.divisor(self._get_member(namespace, 'target'))
            if callable(self.divisor)
            else getattr(self._get_member(namespace, 'target'), self.divisor)
        )
        return operator.truediv(value, divisor)

    def _get_rate(self, namespace: SimpleNamespace) -> object:
        """Return the rate before the divisor, the hook a subclass overrides to derive it from its parameters."""
        return self._evaluate(self.rate, namespace)

    def _get_noise(self, namespace: SimpleNamespace) -> object:
        """Return the noise amplitude before the divisor, the hook a subclass overrides, or `None` for no noise."""
        return None if self.noise is None else self._evaluate(self.noise, namespace)

    def _get_noise_law(self, state: object, amplitude: object, namespace: SimpleNamespace) -> BasePart:
        """Return the noise law of an amplitude, by the process, the convention, and the Hurst index."""
        if self.hurst is not None:
            hurst = self._evaluate(self.hurst, namespace)
            return FractionalNoiseLaw(state, amplitude, hurst=hurst, convention=self.convention)
        process = None
        if self.process is not None:
            handle = self._get_member(namespace, 'process').dW
            process = handle if self.index is None else handle[self.index]
        return NoiseLaw(state, amplitude, convention=self.convention, process=process, conserved=self.conserved)

    def _get_laws(self, namespace: SimpleNamespace) -> list[BasePart]:
        """Return the rate law of the state and its noise law, both divided by the divisor.

        Args:
            namespace:
                The namespace of the symbols the laws can use, with the namespaces of the target and the sources.

        Returns:
            laws:
                The rate law, and the noise law when the interaction has noise.
        """
        state = getattr(self._get_member(namespace, 'target'), self.state)
        rate = self._divide(self._broadcast(self._get_rate(namespace), state), namespace)
        laws: list[BasePart] = [RateLaw(state, rate)]
        noise = self._get_noise(namespace)
        if noise is not None:
            laws.append(self._get_noise_law(state, self._divide(noise, namespace), namespace))
        return laws

class ControlInputElement(Element):
    """Control input that a plant reads, a bounded adaptive parameter scaled to the interval from -1 to 1.

    Args:
        label:
            The label of the input.

        visual:
            How the element is drawn, or `None` for `default_visual_`.

        u:
            The value of the input.

    Examples:
        >>> from skcomplex.systems.engineering.control import ControlInputElement
        >>> ControlInputElement('input', u=0.5).u
        0.5
    """

    def __init__(
        self,
        label: str,
        visual: VisualSpec | None = None,
        *,
        u: Annotated[float, ParameterSpec(adaptive=True, bounds=(-1.0, 1.0))] = 0.0,
    ) -> None:
        super().__init__(label, visual=visual)
        self.u = u

class PlantStateElement(Element):
    """State of a plant with two variables, whose rates a pair of interactions gives.

    Args:
        label:
            The label of the plant.

        visual:
            How the element is drawn, or `None` for `default_visual_`.

        x:
            The value of the first variable.

        y:
            The value of the second variable, the output that tracks the reference.

    Examples:
        >>> from skcomplex.systems.engineering.control import PlantStateElement
        >>> PlantStateElement('plant', x=0.5).x
        0.5
    """

    def __init__(
        self,
        label: str,
        visual: VisualSpec | None = None,
        *,
        x: Annotated[float, StateSpec()] = 0.0,
        y: Annotated[float, StateSpec()] = 0.0,
    ) -> None:
        super().__init__(label, visual=visual)
        self.x = x
        self.y = y