Koopman Methods

Benchmark harness for comparing Koopman/DMD-based methods on three nonlinear dynamical systems:

• Duffing(): 2D bistable oscillator
• VanDerPol(): 2D limit-cycle oscillator
• Lorenz(): 3D chaotic system

All methods share the same data generation, evaluation, and plotting pipeline so results are directly comparable.

---

Paper

Implementing and Comparing Koopman Operator and DMD-Based Methods for Nonlinear Dynamics
Cameron Gordon, Tamanna Iyyani, Roberto Julian Campos, Jeffrey Shin, Evan Jones — August 7, 2026

📄 paper/AM170B_Final_Paper.pdf

The paper writes up the experiments in this repo: four Koopman/DMD approximations (DMD, EDMD, Time-Delay DMD, and two neural Koopman autoencoders) evaluated on Duffing, Van der Pol, and Lorenz under one harness. Headline results:

• EDMD beats DMD on Duffing and Van der Pol when the polynomial dictionary matches the system's nonlinearity, but the gap collapses on Lorenz.
• Time-Delay DMD with truncated SVD and modal reconstruction is the strongest purely linear method on Duffing (horizon 2.20 s in the single-run configuration of Table 5).
• The neural Koopman autoencoder only trains reliably with three specific choices: zero-initialized auxiliary output layers, system-dependent eigenvalue clamping, and a scale-invariant linearization loss.
• Lorenz exposes a ceiling shared by every method: chaotic mixing plus a continuous Koopman spectrum.

Paper → code map

Paper section	Implementation
§3.1 DMD	examples/method_dmd.py
§3.2 EDMD	examples/method_roberto_edmd.py
§3.3 Time-Delay DMD	examples/method_timedelay.py
§3.4.1 Lusch et al. model (+ our fixes)	src/koopman_methods/nn/ — reported as Neural Koopman
§3.4.2 Our parameterization (K = BlockDiag{A_i(z_t)})	koopman-rewrite/ (training) + src/koopman_methods/aux_koopman.py (inference) — reported as Aux Koopman
§3.4.5 Per-system hyperparameters (Table 1)	configs/duffing.toml, configs/vdp.toml, configs/lorenz.toml
§4 Test systems	src/koopman_methods/systems/
§4 Metrics (RMSE, horizon at ε = 0.1)	src/koopman_methods/metrics.py
§5 Per-method results and figures	outputs/<method>/
§6 / Figure 15 cross-method summary	outputs/comparison/ (table.md, rmse_bars.png, horizon_heatmap.png)

Our parameterization differs from Lusch et al. in one place: each auxiliary MLP sees the full latent state rather than only the squared radius of its own block. In the code this is the omega_input = "full_y" setting (the Lusch-style scalar input is "radius", the default).

Results at a glance

Per-method results (single tuned run, paper Tables 3–5):

Method	Duffing RMSE / horizon (s)	Van der Pol RMSE / horizon (s)	Lorenz RMSE / horizon (s)
DMD	0.695 / 0.55	1.092 / 0.10	13.801 / 0.01
EDMD (degree 3)	0.260 / 0.75	0.407 / 1.10	17.35 / 0.09
Time-Delay DMD (d = 10)	0.052 / 2.20	0.483 / 0.75	14.62 / 0.00

Cross-method summary (mean ± std over 5 seeds at σ = 0, paper Figure 15) lives in outputs/comparison/table.md and is regenerated by examples/compare_all.py. Because it re-tunes and averages over seeds, its numbers differ from the single-run tables above — most visibly for Time-Delay DMD on Duffing (0.08 s mean horizon vs the 2.20 s of the best single run).

Citing

@misc{gordon2026koopman,
  title  = {Implementing and Comparing Koopman Operator and DMD-Based Methods for Nonlinear Dynamics},
  author = {Gordon, Cameron and Iyyani, Tamanna and Campos, Roberto Julian and Shin, Jeffrey and Jones, Evan},
  year   = {2026},
  note   = {AM170B final paper. \url{https://github.com/am170b-group2/koopman-methods}}
}

References cited by the paper:

1. Brunton, S. L., Budišić, M., Kaiser, E., and Kutz, J. N. (2021). Modern Koopman theory for dynamical systems. SIAM Review, 64(2), 229–340.
2. Lusch, B., Kutz, J. N., and Brunton, S. L. (2018). Deep learning for universal linear embeddings of nonlinear dynamics. Nature Communications, 9(1), 4950.

---

Implemented Methods

Linear methods (examples/)

File	Method
method_dmd.py	Standard DMD with tunable rank
method_roberto_edmd.py	EDMD with polynomial dictionary, tunable degree and regularisation
method_timedelay.py	Time-delay DMD with tunable number of delays

Neural Koopman (src/koopman_methods/nn/)

The Lusch et al. architecture with the training fixes from §3.4.6 of the paper. Autoencoder with a block-diagonal Koopman operator (complex conjugate pairs + real blocks), trained with reconstruction, multi-step prediction, linearisation, and invertibility losses. See the system-specific training scripts:

Script	System
train_duffing.py	Duffing (latent dim 6, 15 epochs)
train_van_der_pol.py	Van der Pol (latent dim 6, 10 epochs)
train_lorenz.py	Lorenz (latent dim 20, 15 epochs)

Aux Koopman (koopman-rewrite/ + src/koopman_methods/aux_koopman.py)

The paper's own parameterization (§3.4.2): same block-diagonal operator, but every auxiliary MLP consumes the full latent state (omega_input = "full_y") instead of just its own block's squared radius, so the eigenvalue groups are coupled. This is the best-performing method in the paper on all three systems.

Path	Role
koopman-rewrite/train.py	training entry point, driven by a TOML config
configs/{duffing,vdp,lorenz}.toml	per-system architecture and schedule (paper Table 1)
koopman-rewrite/eval_*.py	trajectory, eigenvalue, horizon, and latent-space evaluations
src/koopman_methods/aux_koopman.py	inference-only loader for trained checkpoints
examples/method_aux_koopman.py	fit/predict adapter so it joins the cross-method comparison

python koopman-rewrite/train.py --config configs/duffing.toml

Each run writes its own directory under runs/<name>/ (config.json, best.pt, latest.pt, scaling.npz, train.csv, figures/). To include a trained model in the cross-method comparison, copy that run's config.json, best.pt, and scaling.npz into weights/aux_koopman/<system>/, which is where examples/method_aux_koopman.py looks for them.

---

Comparison Pipeline

Run these in order:

1. Generate training data

python -m koopman_methods.nn.generate_data --system duffing --n 50
python -m koopman_methods.nn.generate_data --system van_der_pol --n 50
python -m koopman_methods.nn.generate_data --system lorenz --n 50

2. Train neural Koopman models

python -m koopman_methods.nn.train_duffing
python -m koopman_methods.nn.train_van_der_pol
python -m koopman_methods.nn.train_lorenz

Checkpoints saved to weights/<system>/.

3. Tune linear method hyperparameters

python examples/tune_linear.py

Runs a grid search over DMD rank, EDMD polynomial degree and regularisation, and time-delay count. Outputs:

• outputs/tune_linear/all_results.json
• outputs/tune_linear/best_configs.json

4. Robustness study

python examples/robustness.py

Evaluates all methods across 5 random seeds and 3 noise levels (σ = 0, 0.01, 0.05). Requires best_configs.json from step 3 and trained neural Koopman weights from step 2. Output: outputs/robustness/results.json.

5. Final comparison figures

python examples/compare_all.py

Produces from results.json:

• outputs/comparison/table.md — mean ± std RMSE and horizon across seeds at σ=0
• outputs/comparison/rmse_bars.png — grouped bar chart (log scale)
• outputs/comparison/horizon_heatmap.png — prediction horizon heatmap

6. Evaluate a neural Koopman checkpoint

python -m koopman_methods.nn.eval_checkpoint weights/duffing/model_epoch_14.pt

Writes trajectory comparison plots, eigenvalue plot, error-over-time, latent trajectory, and phase portrait to plots/<system>/<checkpoint>/.

---

Google Colab

colab_train.ipynb runs the full pipeline on a GPU. Before running:

1. Set runtime to GPU: Runtime → Change runtime type → T4 GPU
2. Add a Colab secret named GITHUB_TOKEN with a personal access token that has repo read access
3. Run cells top to bottom — Drive is mounted and all data, weights, outputs, and plots persist across sessions

---

Adding a New Method

Copy the template:

cp examples/method_template.py examples/method_yourname.py

Implement two functions:

def fit(train_trajs):
    # train_trajs: (n_traj, n_steps, dim)
    return model

def predict(model, x0, n_steps):
    # returns: (n_steps, dim)
    return predicted_trajectory

Everything else — data loading, evaluation, plotting — is already wired up.

---

Adding a New Neural Koopman Version

examples/generate_all_figures.py automatically picks up any NN checkpoint placed in the right location. There are two cases depending on whether your architecture uses the existing Autoencoder class or a custom one.

Case 1 — same Autoencoder architecture

Just drop your weights into the standard location:

weights/
└── <system>/          # duffing | van_der_pol | lorenz
    ├── meta.json
    └── model_epoch_<N>.pt

meta.json must contain at minimum:

{
  "dt": 0.05,
  "window_dt": 2.4
}

dt is the integration timestep used during training; window_dt is window_size * dt (the length of the multi-step prediction window). Both are written automatically by the training scripts. generate_all_figures.py picks the highest-epoch checkpoint automatically.

Case 2 — custom architecture

If your model is not an Autoencoder, add a loader function to generate_all_figures.py. Your model must expose this interface for the existing eval code to work:

model.state_dim          # int — dimension of the observed state
model.latent_dim         # int — dimension of the latent space
model.history_len        # int — number of warmup steps needed before simulation
model.num_delays         # int — number of delay embeddings (1 if none)
model.delay_stride       # int — stride between delay frames (1 if no delays)
model.encode(x)          # (batch, latent_dim) tensor → latent states
model.operator(z)        # one latent step forward
model.operator.as_matrix(z)  # returns the latent operator as a Tensor
model.simulate(history, max_t)  # (history_len, state_dim) → (n_steps, state_dim) numpy array

Then in generate_all_figures.py, replace the Autoencoder.from_checkpoint(ckpt) line in run_nn() with your own loader.

---

Data Shapes

Shape	Meaning
(dim,)	Single state
(n_steps, dim)	Single trajectory
(n_traj, n_steps, dim)	Batch of trajectories
(dim, n_samples)	Snapshot matrix for DMD/EDMD

Useful helpers:

X, X_next = data.make_snapshot_pair(train_trajs)   # for DMD/EDMD
H = data.make_hankel(train_trajs[0], n_delays=20)   # for time-delay DMD

---

Install & Tests

pip install -e .[dev]
pytest -q

---

File Tree

koopman-methods/
├── README.md
├── paper/
│   └── AM170B_Final_Paper.pdf      # final write-up of the results in this repo
├── colab_train.ipynb               # Colab notebook: data generation, training, eval, plots
├── pyproject.toml
├── configs/                        # aux-Koopman training configs (paper Table 1)
│   ├── duffing.toml
│   ├── vdp.toml
│   └── lorenz.toml
├── weights/                        # trained NN checkpoints (not committed)
│   ├── duffing/
│   ├── van_der_pol/
│   ├── lorenz/
│   └── aux_koopman/<system>/
├── outputs/                        # committed figures, summaries, and metrics
│   ├── dmd/  edmd/  timedelay/
│   ├── neural_koopman/  aux_koopman/
│   ├── robustness/results.json
│   └── comparison/                 # paper Figure 15 table + plots
├── examples/
│   ├── generate_all_figures.py     # one-shot: all linear + NN figures
│   ├── quickstart.py               # baseline persistence predictor demo
│   ├── method_template.py          # copy this to add a new linear method
│   ├── method_dmd.py               # DMD
│   ├── method_roberto_edmd.py      # EDMD with polynomial dictionary
│   ├── method_timedelay.py         # time-delay DMD
│   ├── method_aux_koopman.py       # aux-Koopman network (paper §3.4.2)
│   ├── tune_linear.py              # hyperparameter grid search for linear methods
│   ├── robustness.py               # multi-seed + noise robustness study
│   └── compare_all.py              # final comparison figures and table
├── koopman-rewrite/                # aux-Koopman training + evaluation
│   ├── train.py                    # training entry point
│   ├── model.py                    # encoder/decoder + state-dependent operator
│   ├── losses.py                   # reconstruction / prediction / linearity losses
│   ├── config.py  data.py  runs.py  generate_data.py
│   ├── eval_*.py                   # trajectories, eigenvalues, horizon, latent space
│   └── systems/                    # duffing, vdp, lorenz, pendulum
├── src/
│   └── koopman_methods/
│       ├── data.py                 # train/test splits, snapshots, Hankel, noise
│       ├── metrics.py              # RMSE, error curves, prediction horizon
│       ├── plotting.py             # trajectory, eigenvalue, phase portrait plots
│       ├── aux_koopman.py          # inference for aux-Koopman checkpoints
│       ├── systems/
│       │   ├── base.py             # dynamical system base class
│       │   ├── duffing.py
│       │   ├── van_der_pol.py
│       │   └── lorenz.py
│       └── nn/
│           ├── autoencoder.py      # encoder, decoder, block-diagonal Koopman operator
│           ├── training.py         # training loop with multi-loss objective
│           ├── data.py             # WindowDataset for multi-step rollout training
│           ├── generate_data.py    # trajectory generation for NN training
│           ├── plotting.py
│           ├── train_duffing.py
│           ├── train_van_der_pol.py
│           ├── train_lorenz.py
│           ├── eval_checkpoint.py  # per-checkpoint evaluation and plotting
│           ├── eval_run.py
│           └── eval_heatmap.py
└── tests/
    ├── test_systems.py
    └── test_plotting.py