Skip to content

Koopman EDMD — Data-Driven Linear Predictor with Control

monitor.koopman_edmd fits a linear predictive model of a nonlinear controlled system directly from data. It is the model layer of the grid-forming (dVOC) oscillation pack: the matrices it produces are consumed unchanged by the convex Koopman-MPC controller (actuation.koopman_mpc).

1. Mathematical formalism

Koopman operator theory lifts a nonlinear controlled system x_{k+1} = f(x_k, u_k) into a higher-dimensional space of observables ψ : ℝⁿ → ℝᴺ where the evolution is approximately linear. Extended Dynamic Mode Decomposition with control (Korda & Mezić 2018) fits that lifted linear system from snapshot triples (x_i, u_i, y_i) with y_i = f(x_i, u_i):

z_{k+1} = A z_k + B u_k,        x̂_k = C z_k,        z_0 = ψ(x_0),

with A ∈ ℝᴺˣᴺ, B ∈ ℝᴺˣᵐ, C ∈ ℝⁿˣᴺ.

1.1 The two least-squares problems

The matrices minimise the lifted one-step residual and the reconstruction residual (Korda eq. 17 and eq. 20):

[A, B] = argmin Σ_i ‖ψ(y_i) − A ψ(x_i) − B u_i‖²
C      = argmin Σ_i ‖x_i − C ψ(x_i)‖²

Both are solved in closed form through Tikhonov-regularised normal equations (row-major snapshot convention, X_lift is K×N):

(ΦᵀΦ + ρI) [Aᵀ; Bᵀ] = Φᵀ Y_lift,     Φ = [X_lift | U]
(X_liftᵀ X_lift + ρI) Cᵀ = X_liftᵀ X

The ridge ρ ≥ 0 keeps the Gram matrices well-posed. When the dictionary contains the state coordinates, C reduces to the selection [I, 0] and the recovery is exact on the training data.

1.2 Observable dictionaries

kind ψ(x) (before the optional constant) Use
identity x linear systems; exact closure
polynomial monomials up to degree polynomial vector fields, Koopman-invariant subspaces
rbf x plus Gaussian radial bases at centres smooth general nonlinearities
phase x, cos θ_i, sin θ_i, and the Kuramoto order-parameter components R cos Ψ, R sin Ψ phase oscillators / Sakaguchi–Kuramoto

The phase dictionary is the SPO-specific value-add: the first-harmonic Fourier features render the phase vector field close to linear in the lifted space.

2. Python API

from scpn_phase_orchestrator.monitor.koopman_edmd import (
    KoopmanDictionary, fit_koopman_predictor,
)

dictionary = KoopmanDictionary(kind="polynomial", state_dim=2, degree=2)
predictor = fit_koopman_predictor(states, next_states, inputs, dictionary=dictionary)
trajectory = predictor.predict(initial_state, input_sequence)   # (T+1, n)

fit_koopman_predictor(states, next_states, inputs, *, dictionary, regularisation=1e-8) returns a frozen KoopmanPredictor carrying (A, B, C), the dictionary, and the RMS one-step lift residual. KoopmanPredictor.predict rolls the linear model forward; an empty input sequence yields the single-row reconstruction.

All public state, successor-state, control, centre, lift, and rollout arrays must contain finite real numeric evidence before conversion. Boolean, complex, numeric-string, and broken array-protocol payloads fail closed. A directly constructed KoopmanPredictor independently replays the (A, B, C) shape and dictionary-dimension contracts, requires a finite non-negative fit residual, and owns read-only copies of all three matrices so caller mutation cannot alter published prediction evidence.

3. Multi-backend fallback chain

The heavy step is the least-squares solve; it runs on the standard chain Rust → Mojo → Julia → Go → Python, selected fastest-first. The dictionary lift and predictor roll-out are control flow over that kernel and stay Python-side.

3.1 Parity budget

Every backend reproduces the NumPy reference solve to within 1e-9 (absolute and relative) on the same lifted snapshots; the parity gate (tests/test_koopman_edmd_backends.py) fails the build otherwise. Rust, Julia and Go carry the matrices numerically; Mojo round-trips them through a text subprocess protocol, which is why the shared budget is 1e-9 rather than the 1e-12 of the in-process backends.

The direct Go, Julia, and Mojo bridges also share a fail-closed matrix contract before and after backend execution. Snapshot inputs and returned (A, B, C) matrices must be two-dimensional, finite, real numeric payloads with no boolean, complex, or numeric-string aliases; output matrices must keep the exact (N, N), (N, m), and (n, N) shapes before publication.

3.2 Building the backends

Backend Build
Rust maturin develop --release -m spo-kernel/crates/spo-ffi/Cargo.toml
Go cd go && go build -buildmode=c-shared -o libkoopman_edmd.so koopman_edmd.go
Mojo mojo build mojo/koopman_edmd.mojo -o mojo/koopman_edmd_mojo -Xlinker -lm
Julia interpreted; loaded through juliacall at first use

4. Tested invariants

  • Exact linear recovery — the identity dictionary recovers (A, B, C) of a known linear controlled system to 1e-9.
  • Koopman-invariant subspace — on a Brunton–Tu slow-manifold system whose [x₁, x₂, x₁²] span is Koopman-invariant, the polynomial dictionary predicts the original states over a long horizon to <1e-6, while identity cannot.
  • One-step lift gain — the rbf dictionary lowers the one-step state-prediction error below the linear dictionary.
  • Backend parity — all backends match the reference to 1e-9.

5. Pipeline position

oscillation_modes / modal_participation (monitor) → koopman_edmd (model)koopman_mpc (actuation, review-only) → prc_oscillation (assurance). The predictor never actuates; it supplies the linear model the controller's convex programme is built from.

6. References

  • Korda & Mezić 2018, Automatica 93, 149–160 (arXiv:1611.03537) — linear predictors for nonlinear dynamical systems: Koopman operator meets MPC.
  • Williams, Kevrekidis & Rowley 2015, J. Nonlinear Sci. 25, 1307 — a data-driven approximation of the Koopman operator (EDMD).

7. API reference

koopman_edmd

Extended Dynamic Mode Decomposition with control — a data-driven linear predictor.

Koopman operator theory lifts a nonlinear controlled system x_{k+1} = f(x_k, u_k) into a higher-dimensional space of observables ψ(x) where the evolution is approximately linear. Extended Dynamic Mode Decomposition (EDMD) fits that lifted linear model from data; with the control extension of Korda & Mezić (2018) the fitted model is a linear controlled system

z_{k+1} = A z_k + B u_k,        x̂_k = C z_k,        z_0 = ψ(x_0),

with A ∈ ℝ^{N×N}, B ∈ ℝ^{N×m}, C ∈ ℝ^{n×N} and N the dictionary size. Given snapshot triples (x_i, u_i, y_i) with y_i = f(x_i, u_i) the matrices are the (regularised) least-squares solutions

[A, B] = argmin Σ_i ‖ψ(y_i) − A ψ(x_i) − B u_i‖²              (Korda eq. 17)
C      = argmin Σ_i ‖x_i − C ψ(x_i)‖²                          (Korda eq. 20)

solved in closed form through the lifted data matrices (Korda eq. 22). When the dictionary contains the state coordinates themselves, C reduces to the selection [I, 0] and the recovery is exact on the training data; we still fit C so the predictor degrades gracefully for partial dictionaries.

The lifted predictor is the model layer of the grid-forming (dVOC) oscillation pack: it feeds the convex Koopman-MPC controller (actuation.koopman_mpc) whose quadratic programme is built directly from (A, B, C) and whose online cost is independent of the lift dimension N.

The heavy step is the least-squares solve over the lifted matrices; it runs on the standard five-backend chain (Rust → Mojo → Julia → Go → Python). The dictionary lift and the predictor roll-out are control flow over that kernel and stay Python-side.

References

  • Korda & Mezić 2018, Automatica 93, 149-160 (arXiv:1611.03537) — linear predictors for nonlinear dynamical systems: Koopman operator meets MPC.
  • Williams, Kevrekidis & Rowley 2015, J. Nonlinear Sci. 25, 1307 — a data-driven approximation of the Koopman operator (EDMD).

Classes

KoopmanObservables

Bases: Protocol

The observable-map interface consumed by the EDMD fit and predictor.

A Koopman observable map declares its original state dimension n and lifts a batch of states (K, n) to observables (K, N). Both the analytic :class:KoopmanDictionary and the learned monitor.phase_koopman.LearnedKoopmanDictionary satisfy it, so the EDMD fit and the rolled-out predictor are agnostic to how the lift is produced.

Attributes
state_dim property
state_dim: int

The original state dimension n.

Returns

int The original state dimension n.

Methods:
lift
lift(states: FloatArray) -> FloatArray

Lift a batch of states (K, n) to observables (K, N).

Parameters

states : numpy.ndarray The state batch of shape (K, n).

Returns

numpy.ndarray The lifted batch of shape (K, N).

Source code in src/scpn_phase_orchestrator/monitor/koopman_edmd.py
def lift(self, states: FloatArray) -> FloatArray:
    """Lift a batch of states ``(K, n)`` to observables ``(K, N)``.

    Parameters
    ----------
    states : numpy.ndarray
        The state batch of shape ``(K, n)``.

    Returns
    -------
    numpy.ndarray
        The lifted batch of shape ``(K, N)``.
    """
    ...

KoopmanDictionary dataclass

KoopmanDictionary(
    kind: str,
    state_dim: int,
    degree: int = 2,
    centres: FloatArray | None = None,
    width: float = 1.0,
    include_constant: bool = True,
)

A dictionary of observables ψ : ℝ^n → ℝ^N for the lift.

Parameters

kind : str One of "identity", "polynomial", "rbf" or "phase". state_dim : int The state dimension n. degree : int Polynomial degree ("polynomial" only); monomials up to and including this total degree are appended to the state. centres : numpy.ndarray | None RBF centres of shape (n_centres, n) ("rbf" only). width : float Gaussian RBF width σ ("rbf" only). include_constant : bool Whether to prepend the constant observable 1.

Notes

The "phase" dictionary is tuned for phase oscillators: it appends the first-harmonic Fourier features cos θ_i, sin θ_i and the global Kuramoto order-parameter components R cos Ψ, R sin Ψ, which render the Sakaguchi–Kuramoto vector field close to linear in the lifted space.

Methods:
__post_init__
__post_init__() -> None

Validate and normalise immutable dictionary configuration.

Source code in src/scpn_phase_orchestrator/monitor/koopman_edmd.py
def __post_init__(self) -> None:
    """Validate and normalise immutable dictionary configuration."""
    if not isinstance(self.kind, str):
        raise TypeError("kind must be a string")
    if self.kind not in _DICTIONARY_KINDS:
        raise ValueError("kind must be one of: " + ", ".join(_DICTIONARY_KINDS))
    if not isinstance(self.include_constant, bool):
        raise TypeError("include_constant must be a boolean")
    state_dim = _validate_int_at_least(self.state_dim, name="state_dim", minimum=1)
    object.__setattr__(self, "state_dim", state_dim)
    degree = _validate_int_at_least(self.degree, name="degree", minimum=1)
    object.__setattr__(self, "degree", degree)
    object.__setattr__(
        self, "width", _validate_non_negative_real(self.width, name="width")
    )
    if self.kind == "rbf":
        if self.centres is None:
            raise ValueError("rbf dictionary requires centres")
        centres = _validate_matrix(self.centres, name="centres")
        if centres.shape[1] != state_dim:
            raise ValueError(
                f"centres must have {state_dim} columns, got {centres.shape[1]}"
            )
        if self.width <= 0.0:
            raise ValueError("rbf dictionary requires a positive width")
        centres.setflags(write=False)
        object.__setattr__(self, "centres", centres)
    elif self.centres is not None:
        centres = _validate_matrix(self.centres, name="centres")
        centres.setflags(write=False)
        object.__setattr__(self, "centres", centres)
    object.__setattr__(self, "output_dim", self._compute_output_dim())
lift
lift(states: FloatArray) -> FloatArray

Lift a batch of states (K, n) to observables (K, N).

Parameters

states : numpy.ndarray State batch of shape (K, n).

Returns

numpy.ndarray The lifted batch of shape (K, output_dim).

Raises

ValueError If states is not a finite (K, state_dim) array.

Source code in src/scpn_phase_orchestrator/monitor/koopman_edmd.py
def lift(self, states: FloatArray) -> FloatArray:
    """Lift a batch of states ``(K, n)`` to observables ``(K, N)``.

    Parameters
    ----------
    states : numpy.ndarray
        State batch of shape ``(K, n)``.

    Returns
    -------
    numpy.ndarray
        The lifted batch of shape ``(K, output_dim)``.

    Raises
    ------
    ValueError
        If ``states`` is not a finite ``(K, state_dim)`` array.
    """
    matrix = _validate_matrix(states, name="states")
    if matrix.shape[1] != self.state_dim:
        raise ValueError(
            f"states must have {self.state_dim} columns, got {matrix.shape[1]}"
        )
    if self.kind == "identity":
        features = matrix
    elif self.kind == "polynomial":
        features = self._lift_polynomial(matrix)
    elif self.kind == "rbf":
        features = self._lift_rbf(matrix)
    else:
        features = self._lift_phase(matrix)
    if self.include_constant:
        constant = np.ones((matrix.shape[0], 1), dtype=np.float64)
        features = np.hstack((constant, features))
    return np.ascontiguousarray(features, dtype=np.float64)

KoopmanPredictor dataclass

KoopmanPredictor(
    state_matrix: FloatArray,
    input_matrix: FloatArray,
    output_matrix: FloatArray,
    dictionary: KoopmanObservables,
    fit_residual: float,
)

A fitted Koopman linear predictor z_{k+1}=Az_k+Bu_k, x̂=Cz_k.

Parameters

state_matrix : numpy.ndarray A of shape (N, N). input_matrix : numpy.ndarray B of shape (N, m). output_matrix : numpy.ndarray C of shape (n, N). dictionary : KoopmanDictionary The observable dictionary used for the lift. fit_residual : float Root-mean-square one-step lift residual on the training snapshots.

Attributes
lift_dim property
lift_dim: int

The lifted-state dimension N.

input_dim property
input_dim: int

The control dimension m.

state_dim property
state_dim: int

The original state dimension n.

Methods:
__post_init__
__post_init__() -> None

Validate, copy, and freeze the fitted predictor evidence.

Source code in src/scpn_phase_orchestrator/monitor/koopman_edmd.py
def __post_init__(self) -> None:
    """Validate, copy, and freeze the fitted predictor evidence."""
    state_matrix = _validate_matrix(self.state_matrix, name="state_matrix")
    input_matrix = _validate_matrix(self.input_matrix, name="input_matrix")
    output_matrix = _validate_matrix(self.output_matrix, name="output_matrix")
    if state_matrix.shape[0] != state_matrix.shape[1]:
        raise ValueError("state_matrix must be square")
    lift_dim = int(state_matrix.shape[0])
    if input_matrix.shape[0] != lift_dim:
        raise ValueError(f"input_matrix must have {lift_dim} rows")
    if output_matrix.shape[1] != lift_dim:
        raise ValueError(f"output_matrix must have {lift_dim} columns")

    try:
        dictionary_state_dim = self.dictionary.state_dim
        dictionary_lift = self.dictionary.lift
    except AttributeError as exc:
        raise TypeError("dictionary must implement KoopmanObservables") from exc
    if not callable(dictionary_lift):
        raise TypeError("dictionary must implement KoopmanObservables")
    state_dim = _validate_int_at_least(
        dictionary_state_dim, name="dictionary.state_dim", minimum=1
    )
    if output_matrix.shape[0] != state_dim:
        raise ValueError(
            f"output_matrix must have {state_dim} rows to match the dictionary"
        )
    try:
        probe = _validate_matrix(
            dictionary_lift(np.zeros((1, state_dim), dtype=np.float64)),
            name="dictionary lift output",
        )
    except (TypeError, ValueError) as exc:
        raise ValueError(
            "dictionary must produce a finite real lift matrix"
        ) from exc
    if probe.shape != (1, lift_dim):
        raise ValueError(
            "dictionary lift output must have shape "
            f"(1, {lift_dim}), got {probe.shape}"
        )

    residual = _validate_non_negative_real(self.fit_residual, name="fit_residual")
    for matrix in (state_matrix, input_matrix, output_matrix):
        matrix.setflags(write=False)
    object.__setattr__(self, "state_matrix", state_matrix)
    object.__setattr__(self, "input_matrix", input_matrix)
    object.__setattr__(self, "output_matrix", output_matrix)
    object.__setattr__(self, "fit_residual", residual)
lift
lift(state: FloatArray) -> FloatArray

Lift a single state (n,) to its observable z = ψ(x) (N,).

Parameters

state : numpy.ndarray The state vector of shape (n,).

Returns

numpy.ndarray The lifted observable z = ψ(x) of shape (N,).

Raises

ValueError If state is not a finite vector of length state_dim.

Source code in src/scpn_phase_orchestrator/monitor/koopman_edmd.py
def lift(self, state: FloatArray) -> FloatArray:
    """Lift a single state ``(n,)`` to its observable ``z = ψ(x)`` ``(N,)``.

    Parameters
    ----------
    state : numpy.ndarray
        The state vector of shape ``(n,)``.

    Returns
    -------
    numpy.ndarray
        The lifted observable ``z = ψ(x)`` of shape ``(N,)``.

    Raises
    ------
    ValueError
        If ``state`` is not a finite vector of length ``state_dim``.
    """
    vector = _validate_vector(state, name="state")
    if vector.shape[0] != self.state_dim:
        raise ValueError(
            f"state must have {self.state_dim} entries, got {vector.shape[0]}"
        )
    return cast("FloatArray", self.dictionary.lift(vector[None, :])[0])
predict
predict(
    initial_state: FloatArray, input_sequence: FloatArray
) -> FloatArray

Roll the linear predictor forward over an input sequence.

Parameters

initial_state : numpy.ndarray The initial state x_0 of shape (n,). input_sequence : numpy.ndarray The control sequence of shape (T, m).

Returns

numpy.ndarray Predicted states x̂_0 … x̂_T of shape (T + 1, n); the first row is the reconstruction C ψ(x_0). An empty input sequence (0, m) yields the single-row reconstruction.

Raises

ValueError If input_sequence is not a finite (T, m) array.

Source code in src/scpn_phase_orchestrator/monitor/koopman_edmd.py
def predict(
    self, initial_state: FloatArray, input_sequence: FloatArray
) -> FloatArray:
    """Roll the linear predictor forward over an input sequence.

    Parameters
    ----------
    initial_state : numpy.ndarray
        The initial state ``x_0`` of shape ``(n,)``.
    input_sequence : numpy.ndarray
        The control sequence of shape ``(T, m)``.

    Returns
    -------
    numpy.ndarray
        Predicted states ``x̂_0 … x̂_T`` of shape ``(T + 1, n)``; the first
        row is the reconstruction ``C ψ(x_0)``. An empty input sequence
        ``(0, m)`` yields the single-row reconstruction.

    Raises
    ------
    ValueError
        If ``input_sequence`` is not a finite ``(T, m)`` array.
    """
    raw = np.asarray(input_sequence)
    if raw.ndim == 2 and raw.shape == (0, self.input_dim):
        reconstruction = self.output_matrix @ self.lift(initial_state)
        return np.ascontiguousarray(reconstruction[None, :], dtype=np.float64)
    inputs = _validate_matrix(input_sequence, name="input_sequence")
    if inputs.shape[1] != self.input_dim:
        raise ValueError(
            f"input_sequence must have {self.input_dim} columns, "
            f"got {inputs.shape[1]}"
        )
    z = self.lift(initial_state)
    horizon = inputs.shape[0]
    states = np.empty((horizon + 1, self.state_dim), dtype=np.float64)
    states[0] = self.output_matrix @ z
    for step in range(horizon):
        z = self.state_matrix @ z + self.input_matrix @ inputs[step]
        states[step + 1] = self.output_matrix @ z
    return states

Functions:

lift_states

lift_states(
    dictionary: KoopmanDictionary, states: FloatArray
) -> FloatArray

Lift states through dictionary — a free-function alias of lift.

Parameters

dictionary : KoopmanDictionary The observable dictionary. states : numpy.ndarray State batch of shape (K, n).

Returns

numpy.ndarray The lifted batch of shape (K, output_dim).

Source code in src/scpn_phase_orchestrator/monitor/koopman_edmd.py
def lift_states(dictionary: KoopmanDictionary, states: FloatArray) -> FloatArray:
    """Lift ``states`` through ``dictionary`` — a free-function alias of ``lift``.

    Parameters
    ----------
    dictionary : KoopmanDictionary
        The observable dictionary.
    states : numpy.ndarray
        State batch of shape ``(K, n)``.

    Returns
    -------
    numpy.ndarray
        The lifted batch of shape ``(K, output_dim)``.
    """
    return dictionary.lift(states)

fit_koopman_predictor

fit_koopman_predictor(
    states: FloatArray,
    next_states: FloatArray,
    inputs: FloatArray,
    *,
    dictionary: KoopmanObservables,
    regularisation: float = 1e-08,
) -> KoopmanPredictor

Fit an EDMD-with-control linear predictor from snapshot triples.

Parameters

states : numpy.ndarray Snapshot states x_i of shape (K, n). next_states : numpy.ndarray Successor states y_i = f(x_i, u_i) of shape (K, n). inputs : numpy.ndarray Applied controls u_i of shape (K, m). dictionary : KoopmanDictionary The observable dictionary defining the lift ψ. regularisation : float Tikhonov ridge ρ ≥ 0 added to the normal-equation Gram matrices for a well-posed solve.

Returns

KoopmanPredictor The fitted predictor with matrices (A, B, C) and the RMS one-step lift residual.

Raises

ValueError If the snapshot shapes are inconsistent or the dictionary state dimension does not match the data.

Source code in src/scpn_phase_orchestrator/monitor/koopman_edmd.py
def fit_koopman_predictor(
    states: FloatArray,
    next_states: FloatArray,
    inputs: FloatArray,
    *,
    dictionary: KoopmanObservables,
    regularisation: float = 1.0e-8,
) -> KoopmanPredictor:
    """Fit an EDMD-with-control linear predictor from snapshot triples.

    Parameters
    ----------
    states : numpy.ndarray
        Snapshot states ``x_i`` of shape ``(K, n)``.
    next_states : numpy.ndarray
        Successor states ``y_i = f(x_i, u_i)`` of shape ``(K, n)``.
    inputs : numpy.ndarray
        Applied controls ``u_i`` of shape ``(K, m)``.
    dictionary : KoopmanDictionary
        The observable dictionary defining the lift ``ψ``.
    regularisation : float
        Tikhonov ridge ``ρ ≥ 0`` added to the normal-equation Gram matrices for
        a well-posed solve.

    Returns
    -------
    KoopmanPredictor
        The fitted predictor with matrices ``(A, B, C)`` and the RMS one-step
        lift residual.

    Raises
    ------
    ValueError
        If the snapshot shapes are inconsistent or the dictionary state
        dimension does not match the data.
    """
    state_matrix = _validate_matrix(states, name="states")
    next_matrix = _validate_matrix(next_states, name="next_states")
    input_matrix = _validate_matrix(inputs, name="inputs")
    regulariser = _validate_non_negative_real(regularisation, name="regularisation")
    if state_matrix.shape != next_matrix.shape:
        raise ValueError(
            f"states {state_matrix.shape} and next_states {next_matrix.shape} "
            "must have the same shape"
        )
    if input_matrix.shape[0] != state_matrix.shape[0]:
        raise ValueError(
            f"inputs must have {state_matrix.shape[0]} rows, "
            f"got {input_matrix.shape[0]}"
        )
    if state_matrix.shape[1] != dictionary.state_dim:
        raise ValueError(
            f"states must have {dictionary.state_dim} columns to match the "
            f"dictionary, got {state_matrix.shape[1]}"
        )

    x_lift = dictionary.lift(state_matrix)
    y_lift = dictionary.lift(next_matrix)
    a, b, c = _edmd_solve(x_lift, input_matrix, y_lift, state_matrix, regulariser)
    predicted_lift = x_lift @ a.T + input_matrix @ b.T
    residual = float(
        np.sqrt(np.mean((predicted_lift - y_lift) ** 2)) if y_lift.size else 0.0
    )
    return KoopmanPredictor(
        state_matrix=a,
        input_matrix=b,
        output_matrix=c,
        dictionary=dictionary,
        fit_residual=residual,
    )

8. Learned phase-autoencoder observables

The analytic dictionaries above are fixed feature maps. monitor.phase_koopman instead uses a trained phase autoencoder (nn.phase_autoencoder, frozen to the pure-NumPy oscillators.phase_reduction evaluator) as the observable map: the learned latent is the coordinate in which a nonlinear oscillator's dynamics are close to linear, so a predictor fitted in it captures dynamics the analytic dictionaries miss. LearnedKoopmanDictionary satisfies the same KoopmanObservables protocol the fit and predictor consume, and the lift is state-inclusive (ψ(x) = [x, g(x)]) so the output map reconstructs the state exactly while the learned block sharpens the linear evolution. There is no JAX on the control path.

The learned dictionary is also an explicit evidence boundary. Construction requires a real PhaseReducer, a positive integer reducer state dimension, and a canonical boolean constant flag. lift() rejects boolean, complex, numeric-text, coercive-object, overflowing, and broken array-protocol state batches before the reducer runs while preserving legitimate Python and NumPy real numeric objects. The reducer's returned latent is independently required to be a finite, non-coercive (K, 3) array before it can be joined to the state- inclusive lift. This keeps malformed or contradictory learned-observable output from entering EDMD fitting or a downstream review path.

phase_koopman

Use a trained phase autoencoder as the Koopman observable dictionary.

The analytic dictionaries (identity, polynomial, rbf, phase) are fixed feature maps. A phase autoencoder (nn.phase_autoencoder), trained so its latent evolves by an exactly-linear normal-form flow, learns observables in which a nonlinear oscillator's dynamics are close to linear — which is exactly what the Koopman operator wants. :class:LearnedKoopmanDictionary wraps the trained encoder (frozen to the pure-NumPy oscillators.phase_reduction evaluator) as a :class:~scpn_phase_orchestrator.monitor.koopman_edmd.KoopmanObservables map, so the EDMD fit and the condensed Koopman MPC consume learned observables with no change to their machinery and no JAX on the control path.

The lift is state-inclusive — ψ(x) = [x, g(x)] with the encoder latent g(x) — so the output map C reconstructs the state exactly from the identity block while the learned block sharpens the linear evolution A.

Classes

LearnedKoopmanDictionary dataclass

LearnedKoopmanDictionary(
    reducer: PhaseReducer, include_constant: bool = False
)

A Koopman observable map backed by a trained phase-autoencoder encoder.

Parameters

reducer : PhaseReducer The frozen-weights evaluator of the trained phase autoencoder. include_constant : bool Whether to prepend a constant observable for the affine term.

Attributes
state_dim property
state_dim: int

The original state dimension n.

Returns

int The original state dimension n.

output_dim property
output_dim: int

The lifted observable dimension N.

Returns

int The lifted observable dimension N (constant + state + latent).

Methods:
__post_init__
__post_init__() -> None

Validate the frozen reducer and dictionary configuration.

Source code in src/scpn_phase_orchestrator/monitor/phase_koopman.py
def __post_init__(self) -> None:
    """Validate the frozen reducer and dictionary configuration."""
    if not isinstance(self.reducer, PhaseReducer):
        raise ValueError("reducer must be a PhaseReducer")
    state_dim = self.reducer.weights.state_dim
    if isinstance(state_dim, (bool, np.bool_)) or not isinstance(
        state_dim, (Integral, np.integer)
    ):
        raise ValueError("reducer state_dim must be a positive integer")
    if int(state_dim) < 1:
        raise ValueError("reducer state_dim must be a positive integer")
    if type(self.include_constant) is not bool:
        raise ValueError("include_constant must be a boolean")
lift
lift(states: FloatArray) -> FloatArray

Lift a batch of states (K, n) to [x, g(x)] observables.

Parameters

states : numpy.ndarray The state batch of shape (K, n).

Returns

numpy.ndarray The lifted batch of shape (K, output_dim).

Raises

ValueError If states is not a finite (K, state_dim) array.

Source code in src/scpn_phase_orchestrator/monitor/phase_koopman.py
def lift(self, states: FloatArray) -> FloatArray:
    """Lift a batch of states ``(K, n)`` to ``[x, g(x)]`` observables.

    Parameters
    ----------
    states : numpy.ndarray
        The state batch of shape ``(K, n)``.

    Returns
    -------
    numpy.ndarray
        The lifted batch of shape ``(K, output_dim)``.

    Raises
    ------
    ValueError
        If ``states`` is not a finite ``(K, state_dim)`` array.
    """
    matrix = _validate_state_batch(states, self.state_dim)
    latent = _validate_latent(
        self.reducer.encode_observables(matrix), matrix.shape[0]
    )
    features = np.hstack((matrix, latent))
    if self.include_constant:
        constant = np.ones((features.shape[0], 1), dtype=np.float64)
        features = np.hstack((constant, features))
    return np.ascontiguousarray(features, dtype=np.float64)

Functions:

fit_phase_koopman_predictor

fit_phase_koopman_predictor(
    reducer: PhaseReducer,
    states: FloatArray,
    next_states: FloatArray,
    inputs: FloatArray,
    *,
    include_constant: bool = False,
    regularisation: float = 1e-08,
) -> KoopmanPredictor

Fit an EDMD-with-control predictor in learned phase-autoencoder observables.

Parameters

reducer : PhaseReducer The trained phase-autoencoder evaluator providing the observables. states, next_states : numpy.ndarray Snapshot states x_i and successors y_i of shape (K, n). inputs : numpy.ndarray Applied controls u_i of shape (K, m). include_constant : bool Whether the lift prepends a constant observable. regularisation : float Tikhonov regularisation of the least-squares solve.

Returns

KoopmanPredictor The fitted predictor over learned observables.

Source code in src/scpn_phase_orchestrator/monitor/phase_koopman.py
def fit_phase_koopman_predictor(
    reducer: PhaseReducer,
    states: FloatArray,
    next_states: FloatArray,
    inputs: FloatArray,
    *,
    include_constant: bool = False,
    regularisation: float = 1.0e-8,
) -> KoopmanPredictor:
    """Fit an EDMD-with-control predictor in learned phase-autoencoder observables.

    Parameters
    ----------
    reducer : PhaseReducer
        The trained phase-autoencoder evaluator providing the observables.
    states, next_states : numpy.ndarray
        Snapshot states ``x_i`` and successors ``y_i`` of shape ``(K, n)``.
    inputs : numpy.ndarray
        Applied controls ``u_i`` of shape ``(K, m)``.
    include_constant : bool
        Whether the lift prepends a constant observable.
    regularisation : float
        Tikhonov regularisation of the least-squares solve.

    Returns
    -------
    KoopmanPredictor
        The fitted predictor over learned observables.
    """
    dictionary = LearnedKoopmanDictionary(reducer, include_constant=include_constant)
    return fit_koopman_predictor(
        states,
        next_states,
        inputs,
        dictionary=dictionary,
        regularisation=regularisation,
    )