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
identitydictionary recovers(A, B, C)of a known linear controlled system to1e-9. - Koopman-invariant subspace — on a Brunton–Tu slow-manifold system whose
[x₁, x₂, x₁²]span is Koopman-invariant, thepolynomialdictionary predicts the original states over a long horizon to<1e-6, whileidentitycannot. - One-step lift gain — the
rbfdictionary 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.
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__ ¶
Validate and normalise immutable dictionary configuration.
Source code in src/scpn_phase_orchestrator/monitor/koopman_edmd.py
lift ¶
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
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¶
Methods:¶
__post_init__ ¶
Validate, copy, and freeze the fitted predictor evidence.
Source code in src/scpn_phase_orchestrator/monitor/koopman_edmd.py
lift ¶
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
predict ¶
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
Functions:¶
lift_states ¶
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
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
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
¶
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
¶
output_dim
property
¶
The lifted observable dimension N.
Returns¶
int
The lifted observable dimension N (constant + state + latent).
Methods:¶
__post_init__ ¶
Validate the frozen reducer and dictionary configuration.
Source code in src/scpn_phase_orchestrator/monitor/phase_koopman.py
lift ¶
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
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.