KoopmanGraph

Topology-aware Koopman autoencoders for forecasting and analyzing networked dynamics

Documentation | Tutorials | API | Contributing | Citation

KoopmanGraph is an open-source PyTorch Geometric library for learning topology-aware Koopman autoencoders on graphs. GNN encoders lift node features into a latent space, a learned linear operator advances those states, and a matching decoder reconstructs physical node features for multi-step forecasting and spectral analysis.

It is aimed at researchers studying networked dynamical systems—smart grids, traffic networks, epidemic contact graphs, and similar domains—who want an inspectable linear latent operator instead of a flat-vector Koopman model or a purely nonlinear GNN forecaster.

Why KoopmanGraph?

Koopman theory maps nonlinear dynamics into a linear latent domain where multi-step forecasting and spectral analysis are natural. Existing deep Koopman packages typically ignore graph structure, while spatiotemporal GNN forecasters usually lack an explicit linear latent operator.

KoopmanGraph bridges that gap with GNN lifting/decoding, an inspectable Koopman matrix K, row-state latent advance $z \leftarrow z K^{\top}$, and a PyTorch Geometric-native fit / predict workflow.

The library sits in the consistent Koopman autoencoder lineage and is not claimed as a new theoretical contribution; it packages topology-aware lifting, linear latent evolution, and analysis tooling for networked dynamical systems.

Highlights

• Topology-aware learning — GCN/GAT/hypergraph encoders and decoders, delay embeddings, dynamic topology, optional self-adaptive edges, and edge weights
• Flexible dynamics — discrete, continuous-time, networked (koopman="graph"), hypergraph, multiplex / typed hetero (koopman="hetero_graph" with RelGraph), global/local, and continuous-graph operators, with soft or structural stability modes
• Forecasting stack — multi-step rollout, consistency losses, temporal evaluation metrics, and checkpointing
• Spectral analysis — eigendecomposition, mode shapes, dynamical similarity, anomaly helpers, SINDy, and spectral clustering
• Control and adaptation — additive/bilinear control, Koopman-MPC ([mpc]), online RLS adaptation, Kalman observation, and a Gymnasium RL wrapper
• Research tooling — classical DMD-family baselines, lightweight GNN teaching baselines, conformal UQ, and reproducible graph benchmarks
• Optional distributed trainers — native DDP / torchrun, Lightning Fabric, and Ray ensemble helpers under koopman_graph.distributed (power-user; compose with homo or hetero models; see installation extras lightning / ray / distributed)

Full inventory: Capabilities · Architecture

Scope. KoopmanGraph targets topology-aware Koopman autoencoders on graphs and hypergraphs, not traffic-forecasting leaderboards, simplicial/Hodge layers, or molecular-dynamics MSM toolchains. Measured limits (transfer, factorization cost, residual diagnostics, UQ assumptions) are consolidated in Scope and limitations.

Installation

Requires Python 3.10+, PyTorch, and PyTorch Geometric. Install those first, then:

pip install koopman-graph
# or: uv pip install koopman-graph

See the installation guide for editable installs, uv workflows, docs builds, and platform-specific wheels. Release notes: CHANGELOG.md.

Quickstart

import torch
from koopman_graph import GNNDecoder, GNNEncoder, GraphKoopmanModel
from koopman_graph.datasets import SyntheticDynamicGraphBenchmark

data_sequence = SyntheticDynamicGraphBenchmark.generate(
    num_nodes=20,
    num_timesteps=30,
    in_channels=3,
    seed=42,
    noise_std=0.01,
)

encoder = GNNEncoder(3, 64, 64)
decoder = GNNDecoder(64, 64, 3)
model = GraphKoopmanModel(
    encoder=encoder,
    decoder=decoder,
    latent_dim=64,
    time_step=0.1,
)

torch.manual_seed(0)
history = model.fit(data_sequence, epochs=20, lr=1e-3)
future_graphs = model.predict(data_sequence[0], steps=5)
spectrum = model.spectrum()

print(f"Final loss: {history.loss[-1]:.6f}")
print(f"Predicted {len(future_graphs)} snapshots, shape: {future_graphs[0].x.shape}")
print(f"K eigenvalues: {tuple(spectrum.eigenvalues.shape)}")
print(f"Top |λ|: {spectrum.magnitudes[:3].tolist()}")

Expected output:

Final loss: <float>
Predicted 5 snapshots, shape: torch.Size([20, 3])
K eigenvalues: (64,)
Top |λ|: [<float>, <float>, <float>]

More detail: Quickstart guide · API reference

See it in action

SIR epidemic on a ring: truth vs forecast from examples/06_epidemic_ring.ipynb.

METR-LA aggregate RMSE vs in-repo STGCN / DCRNN / Graph WaveNet teaching baselines (not dedicated-library SOTA) from examples/22_gnn_forecaster_comparison.ipynb.

Featured tutorials: 01 synthetic · 03 traffic · 06 epidemic · 22 GNN baselines · 39 hetero RelGraph · full gallery

Learn more

• Quickstart — train / predict walkthrough
• Capabilities — feature inventory and datasets
• Scope and limitations — when not to use; measured boundaries
• Architecture — public vs power-user API layers
• FAQ / troubleshooting — install, imports, checkpoints
• Installation — dependencies and install paths
• What’s new in 0.9.0: heterogeneous / multiplex RelGraph + koopman="hetero_graph", composed with 0.8 DDP / Fabric / Lightning / Ray trainers — see CHANGELOG.md.

Related software

• PyKoopman and DLKoopman target vector-valued Koopman / deep-Koopman workflows; they treat the state as a flat vector rather than propagating information along graph edges.
• PyTorch Geometric provides mature GNN infrastructure on irregular graphs; KoopmanGraph adds an explicit linear latent operator, consistency losses, and a documented fit / predict forecasting stack on that substrate.
• Spatiotemporal GNN forecasters such as STGCN, DCRNN, and Graph WaveNet typically learn nonlinear convolutional or recurrent maps on graphs; KoopmanGraph instead advances an inspectable linear Koopman matrix K (see in-repo teaching baselines in examples/22).

Community and citation

• Contribute, report issues, or seek support: CONTRIBUTING.md · Support · Code of Conduct
• Install / runtime troubleshooting: FAQ
• Security vulnerabilities (private): SECURITY.md
• Development checks and release process: CONTRIBUTING.md

If you use KoopmanGraph in research, please cite:

@software{koopmangraph2026,
  author       = {Travis Kessler},
  title        = {KoopmanGraph: Topology-Aware Koopman Autoencoders for Networked Dynamics},
  year         = {2026},
  publisher    = {Zenodo},
  doi          = {10.5281/zenodo.21711185},
  url          = {https://github.com/tjkessler/KoopmanGraph},
  version      = {0.9.0},
}

License

KoopmanGraph is released under the Apache License 2.0.