API Reference ¶

fit(data_sequence, *, epochs=100, lr=0.001, optimizer=<class 'torch.optim.adam.Adam'>, device=None, loss_weights=None, loss_weight_schedule=None, rollout_horizon=None, rollout_start_indices=None, rollout_starts_per_epoch=None, rollout_start_seed=None, lr_scheduler=None, window_length=None, batch_size=8, windows_per_epoch=None, window_seed=None, max_grad_norm=None, early_stopping_patience=None, early_stopping_min_delta=0.0, early_stopping_monitor='auto', validation_sequence=None, restore_best_weights=False, checkpoint_path=None, **optimizer_kwargs)[source]

Train encoder, Koopman operator, and decoder end-to-end.

Minimizes a weighted sum of one-step MSE and optional forward and backward consistency terms:

loss = w_r * recon_loss
     + w_f * ||K z_t - z_{t+1}||^2
     + w_b * ||z_t - z_{t+1} K^{\dagger}||^2

where z_t and z_{t+1} are encoder outputs for consecutive snapshots and weights (w_r, w_f, w_b) come from a LossWeights object or an optional per-epoch schedule.

When data_sequence is a list of GraphSnapshotSequence objects, losses are averaged across trajectories before each optimizer step.

Parameters:

data_sequence (GraphSnapshotSequence, sequence of Data, or sequence of GraphSnapshotSequence) – One training trajectory or multiple trajectories of the same system. A plain list of Data snapshots is treated as a single trajectory; a list whose first element is a GraphSnapshotSequence is treated as multiple trajectories.
epochs (int, optional) – Number of training epochs. Default is 100.
lr (float, optional) – Learning rate passed to the optimizer. Default is 1e-3.
optimizer (callable, optional) – Optimizer class. Default is torch.optim.Adam.
device (str, torch.device, or None, optional) – Device for training. Defaults to the model’s current device, or CPU if the model has no parameters.
loss_weights (LossWeights or None, optional) – Static loss weights for all epochs. When None and no schedule is provided, defaults to reconstruction-only training.
loss_weight_schedule (callable or None, optional) – Callable epoch -> LossWeights applied each epoch. Overrides loss_weights when set.
rollout_horizon (int or None, optional) – Number of autoregressive rollout steps used when loss_weights.rollout is non-zero. Defaults to num_timesteps - 1.
rollout_start_indices (sequence of int, "all", or None, optional) – Rollout-loss origin indices. None uses [0]; "all" uses every valid origin for the rollout horizon.
rollout_starts_per_epoch (int or None, optional) – When set, randomly sample this many rollout origins each epoch.
rollout_start_seed (int or None, optional) – Base seed for random rollout-origin sampling. The effective seed is rollout_start_seed + epoch.
lr_scheduler (LRScheduler or callable, optional) – Learning-rate scheduler instance or factory optimizer -> scheduler. Stepped once per epoch after the optimizer update.
window_length (int or None, optional) – Fixed number of snapshots per training window. When set, enables mini-batch training with multiple optimizer steps per epoch. None preserves full-sequence single-step training.
batch_size (int, optional) – Number of temporal windows averaged per optimizer step. Used only when window_length is set. Default is 8.
windows_per_epoch (int or None, optional) – Maximum sampled windows per epoch. None uses every valid window across all trajectories.
window_seed (int or None, optional) – Base seed for reproducible epoch-specific window shuffling.
max_grad_norm (float or None, optional) – When set, clip the global gradient norm before each optimizer step.
early_stopping_patience (int or None, optional) – Stop training when training loss fails to improve for this many consecutive epochs. Disabled when None.
early_stopping_min_delta (float, optional) – Minimum decrease in the monitored loss to count as improvement. Default is 0.0.
early_stopping_monitor ({"auto", "train", "val"}, optional) – Loss used for early stopping and best-epoch tracking. "auto" monitors validation loss when validation_sequence is provided, otherwise training loss. Default is "auto".
validation_sequence (GraphSnapshotSequence, sequence of Data, sequence of GraphSnapshotSequence, or None, optional) – Optional held-out snapshots for per-epoch validation loss. A single validation sequence is reused for all training trajectories; a list of validation sequences must match the training trajectory count.
restore_best_weights (bool, optional) – When True, reload in-memory weights from the lowest-loss epoch after training completes. Default is False.
checkpoint_path (str, Path, or None, optional) – When set, write a checkpoint at the lowest-loss epoch using save(). Default is None.
**optimizer_kwargs (Any) – Additional keyword arguments forwarded to the optimizer constructor.

Returns:

Per-epoch training and validation losses and early-stop metadata.

Return type:

FitHistory

Raises:

ValueError – If epochs < 1, early_stopping_patience < 1 when set, early_stopping_monitor="val" without validation_sequence, validation list length mismatches training trajectories, or fewer than two snapshots are provided for training or validation.

class koopman_graph.GraphSnapshotSequence(snapshots, *, allow_dynamic_topology=False, control_inputs=None)[source]

Bases: object

Container for a time-ordered sequence of PyG Data graph snapshots.

By default all snapshots must share the same edge_index, optional edge_weight, node count, and feature dimension. Set allow_dynamic_topology=True to permit per-snapshot edge_index while still requiring a fixed node count and feature dimension. Optional control_inputs store exogenous inputs u_t applied when advancing from snapshot t to t+1. Downstream training APIs should require at least two snapshots; construction here allows a single snapshot for inspection or prediction-only workflows.

Notes

Read-only views of sequence metadata are exposed as snapshots, edge_index, edge_weight, is_dynamic_topology, control_inputs, has_controls, control_dim, num_nodes, num_timesteps, and in_channels. The edge_index and edge_weight properties are only defined for static-topology sequences; use sequence[t].edge_index when is_dynamic_topology is True.

classmethod from_arrays(node_features, edge_index, *, edge_weight=None, control_inputs=None, dtype=torch.float32)[source]

Build a sequence from node feature arrays and a shared topology.

Parameters:

node_features (array-like) – Array with shape (num_timesteps, num_nodes, in_channels).
edge_index (array-like) – Shared edge index with shape (2, num_edges).
edge_weight (array-like, optional) – Shared scalar edge weights with shape (num_edges,). When provided, attached to every snapshot.
control_inputs (array-like, optional) – Per-timestep control inputs with shape (num_timesteps, control_dim) or (num_timesteps, num_nodes, control_dim).
dtype (torch.dtype, optional) – Floating dtype used when converting numpy inputs to torch tensors. Default is torch.float32.

Returns:

Validated snapshot sequence.

Return type:

Raises:

ValueError – If node_features, edge_index, or edge_weight have invalid shape.

classmethod from_dynamic_arrays(node_features, edge_indices, *, edge_weights=None, control_inputs=None, dtype=torch.float32)[source]

Build a dynamic-topology sequence from per-timestep edge indices.

Parameters:

node_features (array-like) – Array with shape (num_timesteps, num_nodes, in_channels).
edge_indices (sequence of array-like) – One edge index per timestep, each with shape (2, num_edges_t).
edge_weights (sequence of array-like or None, optional) – Optional per-timestep scalar edge weights aligned with edge_indices. When provided, must have the same length as edge_indices.
control_inputs (array-like, optional) – Per-timestep control inputs with shape (num_timesteps, control_dim) or (num_timesteps, num_nodes, control_dim).
dtype (torch.dtype, optional) – Floating dtype used when converting numpy inputs to torch tensors. Default is torch.float32.

Returns:

Validated snapshot sequence with allow_dynamic_topology=True.

Return type:

Raises:

ValueError – If shapes are inconsistent or edge_indices length mismatches num_timesteps.

property is_dynamic_topology: bool

Return whether snapshots use time-varying edge_index.

Returns:: True when the sequence was constructed with allow_dynamic_topology=True and at least one snapshot differs in edge_index from the first snapshot.
Return type:: bool

property allow_dynamic_topology: bool

Return whether dynamic topology mode was enabled at construction.

Returns:: True when per-snapshot edge_index values are permitted.
Return type:: bool

property control_inputs: Tensor | None

Return per-timestep control inputs when present.

Returns:: Control tensor with shape (num_timesteps, control_dim) or (num_timesteps, num_nodes, control_dim).
Return type:: Tensor or None

property has_controls: bool

Return whether the sequence carries control inputs.

Returns:: True when control_inputs is not None.
Return type:: bool

property control_dim: int

Return the control feature dimension.

Returns:: Control dimension when controls are present, otherwise 0.
Return type:: int

control_at(index)[source]

Return the control input driving transition from snapshot index.

Parameters:: index (int) – Timestep index in [0, num_timesteps - 1].
Returns:: Control vector for the transition index -> index + 1.
Return type:: Tensor
Raises:: ValueError – If controls are absent or index is out of range.

rollout_controls(start, steps)[source]

Return controls for an autoregressive rollout from a start snapshot.

Parameters:

start (int) – Index of the initial snapshot.
steps (int) – Number of rollout steps.

Returns:

Control inputs for each rollout step. Empty when the sequence has no controls.

Return type:

list of Tensor

Raises:

ValueError – If start or steps are invalid or controls are unavailable for the requested horizon.

property snapshots: list[Data]

Return the underlying list of graph snapshots.

Returns:: Time-ordered PyG graph snapshots.
Return type:: list of Data

property edge_index: Tensor

Return the shared edge index for static-topology sequences.

Returns:: Edge index with shape (2, num_edges).
Return type:: Tensor
Raises:: ValueError – If is_dynamic_topology is True.

property edge_weight: Tensor | None

Return the shared scalar edge weights for static-topology sequences.

Returns:: Edge weights with shape (num_edges,), or None when the sequence is unweighted.
Return type:: Tensor or None
Raises:: ValueError – If is_dynamic_topology is True.

property num_nodes: int

Return the number of nodes in the graph topology.

Returns:: Node count shared across all snapshots.
Return type:: int

property num_timesteps: int

Return the number of timesteps in the sequence.

Returns:: Length of the temporal sequence.
Return type:: int

property in_channels: int

Return the node feature dimension.

Returns:: Feature dimension shared across all snapshots.
Return type:: int

slice(start, stop)[source]

Return a contiguous temporal sub-sequence.

Parameters:

start (int) – Inclusive start index.
stop (int) – Exclusive stop index.

Returns:

Snapshots in [start, stop) with matching controls and topology policy.

Return type:

GraphSnapshotSequence

Raises:

ValueError – If the bounds are negative, empty, reversed, or exceed the sequence length.

class koopman_graph.HorizonMetrics(horizon, mae, rmse, mape)[source]

Bases: object

Forecast metrics at a single prediction horizon.

horizon

Forecast horizon in steps.

Type:: int

mae

Mean absolute error averaged over evaluation origins.

Type:: float

rmse

Root mean squared error averaged over evaluation origins.

Type:: float

mape

Mean absolute percentage error averaged over evaluation origins.

Type:: float

class koopman_graph.KoopmanOperator(latent_dim, *, init_mode='identity_noise', init_scale=0.01, parameterization='dense', max_spectral_radius=1.0, control_dim=0)[source]

Bases: Module

Learnable finite-dimensional Koopman operator matrix K.

Applies the same linear map to each node’s latent vector. For input z with trailing dimension latent_dim, the uncontrolled forward pass computes:

z_next = z @ K.T

When control_dim is positive, exogenous inputs drive the transition:

z_next = z @ K.T + u @ B

where K has shape (latent_dim, latent_dim) and B has shape (control_dim, latent_dim). Global controls u with shape (control_dim,) are broadcast to every node; per-node controls use shape (num_nodes, control_dim). Arbitrary leading dimensions are supported (e.g. (num_nodes, latent_dim) or (batch, num_nodes, latent_dim)).

latent_dim

Dimension of the latent space.

Type:: int

control_dim

Dimension of exogenous control inputs. Zero disables control.

Type:: int

init_mode

Weight initialization strategy for K.

Type:: str

init_scale

Noise scale used when init_mode="identity_noise".

Type:: float

parameterization

Parameterization used for K ("dense" or "odo").

Type:: str

max_spectral_radius

Upper bound on the spectral radius when parameterization="odo".

Type:: float

reset_control_parameters()[source]

Reinitialize the control input matrix B.

Return type:: None

control_term(u, *, num_nodes=None)[source]

Map control inputs to a latent-space offset u @ B.

Parameters:

u (Tensor) – Global control with shape (control_dim,) or per-node control with shape (num_nodes, control_dim).
num_nodes (int or None, optional) – Expected node count when u is per-node. Used for validation only.

Returns:

Latent offset with shape (latent_dim,) for global control or (num_nodes, latent_dim) for per-node control.

Return type:

Tensor

Raises:

ValueError – If control_dim is zero, u has invalid shape, or per-node u does not match num_nodes.

property K: Tensor

Assembled Koopman matrix with shape (latent_dim, latent_dim).

For parameterization="dense" this is the learnable parameter. For parameterization="odo" it is assembled from orthogonal factors and a bounded diagonal matrix.

Returns:: Current operator matrix K.
Return type:: Tensor

reset_parameters()[source]

Reinitialize operator parameters according to init_mode.

Return type:: None

spectral_radius()[source]

Return the spectral radius of the assembled operator matrix.

Returns:: Scalar tensor with the maximum eigenvalue magnitude.
Return type:: Tensor

forward(z, control=None)[source]

Advance latent states by one linear Koopman step.

When control_dim is positive, the controlled update is:

z_next = z @ K.T + control_effect

where control_effect is u @ B broadcast for global control u with shape (control_dim,) or applied per node when u has shape (num_nodes, control_dim).

Parameters:

z (Tensor) – Latent states with shape (..., latent_dim).
control (Tensor or None, optional) – Exogenous control input applied during this step. Required when control_dim is positive.

Returns:

Advanced latent states with the same shape as z.

Return type:

Tensor

Raises:

ValueError – If the trailing dimension of z does not match latent_dim, controls are missing for a controlled operator, or control has an invalid shape.

inverse_step(z, *, control=None, inverse_matrix=None)[source]

Apply one inverse Koopman step to recover the previous latent state.

For forward dynamics z_{t+1} = z_t @ K.T + u_t @ B, this returns an estimate of z_t from z_{t+1} and the control u_t that drove the transition.

Parameters:

z (Tensor) – Latent states at time t+1 with shape (..., latent_dim).
control (Tensor or None, optional) – Control input that drove the forward transition. Required when control_dim is positive.
inverse_matrix (Tensor or None, optional) – Precomputed K^{-1} for dense parameterization. When omitted, the inverse is computed on demand.

Returns:

Recovered latent states at time t, same shape as z.

Return type:

Tensor

dense_inverse_matrix()[source]

Return the inverse (or pseudo-inverse) of the assembled dense matrix.

Intended for reuse across multiple backward-consistency pair evaluations within one training step.

Returns:: Matrix K^{-1} (or K^{\dagger}) with shape (latent_dim, latent_dim).
Return type:: Tensor
Raises:: ValueError – If parameterization is not "dense".

class koopman_graph.KoopmanSpectrum(eigenvalues, eigenvectors, magnitudes, growth_rates, frequencies, time_step)[source]

Bases: object

Eigendecomposition and time scales of a discrete Koopman operator.

Eigenpairs are sorted by descending eigenvalue magnitude. Frequencies are reported in cycles per unit time; multiply by 2 * pi for angular frequency.

eigenvalues

Complex eigenvalues with shape (latent_dim,).

Type:: Tensor

eigenvectors

Complex right eigenvectors stored as columns, with shape (latent_dim, latent_dim).

Type:: Tensor

magnitudes

Eigenvalue magnitudes with shape (latent_dim,).

Type:: Tensor

growth_rates

Continuous-time exponential growth rates log(|lambda|) / time_step.

Type:: Tensor

frequencies

Signed continuous-time frequencies angle(lambda) / (2 * pi * time_step) in cycles per unit time.

Type:: Tensor

time_step

Physical duration represented by one discrete Koopman step.

Type:: float

mode_amplitudes(latent_states)[source]

Project latent states onto the Koopman eigenvector basis.

For a latent row vector z, the returned amplitudes a satisfy z.T = eigenvectors @ a. Any leading dimensions are preserved.

Parameters:

latent_states (Tensor) – Latent states with shape (..., latent_dim).

Returns:

Complex mode amplitudes with the same shape as latent_states.

Return type:

Tensor

Raises:

ValueError – If the trailing latent dimension does not match the spectrum.
RuntimeError – If the eigenvector matrix is singular.

class koopman_graph.LossWeights(reconstruction=1.0, forward=0.0, backward=0.0, rollout=0.0, eigenvalue=0.0)[source]

Bases: object

Weights for reconstruction and consistency loss terms.

reconstruction

Weight on the one-step reconstruction (MSE) loss.

Type:: float

forward

Weight on the forward consistency loss.

Type:: float

backward

Weight on the backward consistency loss.

Type:: float

rollout

Weight on the autoregressive rollout reconstruction loss.

Type:: float

eigenvalue

Weight on the unit-circle eigenvalue hinge penalty.

Type:: float

class koopman_graph.TemporalSplit(train, val, test)[source]

Bases: object

Train, validation, and test snapshot sequences from a temporal split.

train

Earliest contiguous snapshots used for training.

Type:: GraphSnapshotSequence

val

Middle contiguous snapshots used for validation.

Type:: GraphSnapshotSequence

test

Latest contiguous snapshots held out for evaluation.

Type:: GraphSnapshotSequence

class koopman_graph.WindowSampler(sequences, *, window_length, batch_size=8, windows_per_epoch=None, shuffle=True, seed=None)[source]

Bases: object

Sample fixed-length temporal windows from one or more trajectories.

Parameters:

sequences (GraphSnapshotSequence or sequence of GraphSnapshotSequence) – Source trajectories. Each must contain at least window_length snapshots.
window_length (int) – Number of snapshots per sampled window. Must be at least 2.
batch_size (int, optional) – Number of windows yielded together. Default is 8.
windows_per_epoch (int or None, optional) – Maximum number of windows sampled per epoch. None uses every valid window. Values larger than the available window count are capped.
shuffle (bool, optional) – Randomize window order each epoch. Default is True.
seed (int or None, optional) – Base seed for reproducible epoch-specific shuffling.

property num_windows: int

Return the total number of valid windows.

Returns:: Number of valid windows across every source trajectory.
Return type:: int

iter_epoch(epoch=0)[source]

Yield batches of windows for one epoch.

Parameters:: epoch (int, optional) – Zero-based epoch index mixed into seed. Default is 0.
Yields:: list of GraphSnapshotSequence – A batch containing at most batch_size temporal windows.
Return type:: Iterator[list[GraphSnapshotSequence]]

koopman_graph.compute_spectrum(operator, time_step)[source]

Compute the sorted spectrum and continuous-time mode characteristics.

Parameters:

operator (Tensor) – Square discrete-time Koopman matrix with shape (latent_dim, latent_dim).
time_step (float) – Positive physical duration represented by one operator step.

Returns:

Eigenpairs sorted by descending magnitude, plus growth rates and frequencies converted using time_step.

Return type:

KoopmanSpectrum

Raises:

ValueError – If operator is not a non-empty square matrix or time_step is not positive.
TypeError – If operator is not floating-point or complex.

koopman_graph.decode_mode_shapes(model, x_or_data, mode_indices=None, *, edge_index=None, perturbation=0.001)[source]

Decode latent Koopman directions into spatial node-feature mode shapes.

The decoder is generally nonlinear, so mode shapes are estimated with a centered finite-difference directional derivative around the encoded graph. Real and imaginary parts of complex eigenvectors are probed separately and combined into a complex-valued mode shape.

Parameters:

model (GraphKoopmanModel) – Model whose operator spectrum and decoder are analyzed.
x_or_data (Tensor or Data) – Reference graph used as the decoder linearization point.
mode_indices (sequence of int or None, optional) – Indices into the magnitude-sorted spectrum. Defaults to every mode.
edge_index (Tensor or None, optional) – Graph edges, required when x_or_data is a feature tensor.
perturbation (float, optional) – Positive centered finite-difference step. Default is 1e-3.

Returns:

Complex mode shapes with shape (num_modes, num_nodes, out_channels).

Return type:

Tensor

Raises:

ValueError – If perturbation is not positive or a mode index is out of range.

koopman_graph.evaluate_forecast(model, sequence, *, horizons=(3, 6, 12), start_indices=None)[source]

Evaluate autoregressive multi-horizon forecasts on a snapshot sequence.

For each forecast origin, the model predicts up to max(horizons) steps ahead and metrics are averaged across origins at each requested horizon.

Parameters:

model (nn.Module) – Model implementing predict().
sequence (GraphSnapshotSequence) – Evaluation snapshots with shared topology.
horizons (sequence of int, optional) – Forecast horizons to report. Default is (3, 6, 12).
start_indices (sequence of int or None, optional) – Forecast-origin indices. When None, uses every valid origin in sequence.

Returns:

Per-horizon and aggregate MAE, RMSE, and MAPE.

Return type:

Raises:

ValueError – If horizons is empty, any horizon is invalid, or the sequence is too short.

koopman_graph.mae(prediction, target)[source]

Compute mean absolute error.

Parameters:

prediction (Tensor) – Predicted values.
target (Tensor) – Ground-truth values with the same shape as prediction.

Returns:

Scalar mean absolute error.

Return type:

Tensor

koopman_graph.mape(prediction, target, *, eps=1e-08)[source]

Compute mean absolute percentage error.

Parameters:

prediction (Tensor) – Predicted values.
target (Tensor) – Ground-truth values with the same shape as prediction.
eps (float, optional) – Small constant added to the denominator for numerical stability. Default is 1e-8.

Returns:

Scalar mean absolute percentage error (not scaled to 0–100).

Return type:

Tensor

koopman_graph.rmse(prediction, target)[source]

Compute root mean squared error.

Parameters:

prediction (Tensor) – Predicted values.
target (Tensor) – Ground-truth values with the same shape as prediction.

Returns:

Scalar root mean squared error.

Return type:

Tensor

koopman_graph.temporal_split(sequence, *, train_ratio=0.7, val_ratio=0.1, test_ratio=0.2, min_train_timesteps=2, min_val_timesteps=2, min_test_timesteps=1)[source]

Split a snapshot sequence into contiguous train, validation, and test sets.

Earlier snapshots are assigned to training, later snapshots to validation and test. Ratios must sum to 1.0.

Parameters:

sequence (GraphSnapshotSequence) – Full time-ordered snapshot sequence to split.
train_ratio (float, optional) – Fraction of timesteps assigned to training. Default is 0.7.
val_ratio (float, optional) – Fraction assigned to validation. Default is 0.1.
test_ratio (float, optional) – Fraction assigned to test. Default is 0.2.
min_train_timesteps (int, optional) – Minimum training snapshots required. Default is 2.
min_val_timesteps (int, optional) – Minimum validation snapshots required. Default is 2.
min_test_timesteps (int, optional) – Minimum test snapshots required. Default is 1.

Returns:

Contiguous train, validation, and test sequences sharing topology.

Return type:

TemporalSplit

Raises:

ValueError – If ratios do not sum to 1.0, any minimum is violated, or the sequence is too short for the requested split.

Model¶

GraphKoopmanModel: encoder, Koopman operator, and decoder composition.

class koopman_graph.model.GraphKoopmanModel(encoder, decoder, latent_dim, time_step, *, koopman_init_mode='identity_noise', koopman_init_scale=0.01, koopman_parameterization='dense', koopman_max_spectral_radius=1.0, control_dim=0)[source]¶

Bases: Module

Topology-aware Koopman dynamics model for graph snapshots.

Composes a GNN encoder (lifting), a finite-dimensional Koopman operator (linear latent evolution), and a symmetric GNN decoder (reconstruction).

encoder¶

Topology-aware encoder for latent lifting.

Type:: GNNEncoder or GATEncoder

decoder¶

Symmetric GNN decoder for physical reconstruction.

Type:: GNNDecoder

latent_dim¶

Latent space dimension shared by encoder, operator, and decoder.

Type:: int

time_step¶

Physical time increment associated with one model step. Used by spectrum() to convert discrete eigenvalues into continuous-time growth rates and frequencies.

Type:: float

koopman¶

Learnable linear propagator in latent space.

Type:: KoopmanOperator

spectrum()[source]¶

Analyze the learned Koopman operator spectrum.

Uses time_step to convert discrete eigenvalues into continuous-time growth rates and frequencies.

Returns:: Magnitude-sorted eigenvalues, eigenvectors, and continuous-time mode characteristics.
Return type:: KoopmanSpectrum

save(path)[source]¶

Persist model weights and architecture configuration to disk.

Parameters:: path (str or Path) – Destination checkpoint file (.pt). Parent directories are created when missing.

classmethod load(path, *, map_location=None)[source]¶

Load a trained model from a checkpoint file.

Reconstructs encoder, decoder, and Koopman operator architecture from the saved configuration and restores learned weights.

Parameters:

path (str or Path) – Checkpoint file produced by save().
map_location (str, torch.device, or None, optional) – Device mapping forwarded to torch.load().

Returns:

Ready-to-use model in evaluation mode.

Return type:

GraphKoopmanModel

forward(x_or_data, edge_index=None, edge_weight=None, control=None)[source]¶

Predict the next graph snapshot from the current one.

Performs encode → linear Koopman advance → decode for a single step.

Parameters:

x_or_data (Tensor or Data) – Either a PyG Data object or node features x of shape (num_nodes, in_channels).
edge_index (Tensor, optional) – Edge index with shape (2, num_edges). Required when x_or_data is a tensor; ignored for Data input.
edge_weight (Tensor, optional) – Scalar edge weights with shape (num_edges,). Required when x_or_data is a tensor and weights are used; ignored for Data input.
control (Tensor or None, optional) – Exogenous control input for this step. Required when control_dim is positive.

Returns:

Predicted node features of shape (num_nodes, out_channels).

Return type:

Tensor

predict(initial_graph, steps, edge_index=None, edge_weight=None, controls=None, future_topologies=None)[source]¶

Autoregressively predict future graph snapshots.

Encodes the initial graph once, advances the latent state with the Koopman operator for steps iterations, and decodes after each step. Runs in evaluation mode without gradient tracking.

Parameters:

initial_graph (Tensor or Data) – Either a PyG Data object or node features x of shape (num_nodes, in_channels).
steps (int) – Number of future snapshots to predict (must be >= 1).
edge_index (Tensor, optional) – Edge index with shape (2, num_edges). Required when initial_graph is a tensor; ignored for Data input.
edge_weight (Tensor, optional) – Scalar edge weights with shape (num_edges,). Required when initial_graph is a tensor and weights are used; ignored for Data input.
controls (sequence of Tensor or None, optional) – Future control inputs for each rollout step. Required with length steps when control_dim is positive.
future_topologies (sequence of Data or None, optional) – Known topologies for rollout decode steps. Shorter sequences hold the last provided topology for remaining steps.

Returns:

steps predicted graph snapshots. Each Data.x has shape (num_nodes, out_channels) and carries the edge_index (and optional edge_weight) used for that step’s decode.

Return type:

list of Data

Raises:

ValueError – If steps < 1 or controls are missing/invalid for a controlled model.

evaluate(sequence, *, horizons=(3, 6, 12), start_indices=None)[source]¶

Evaluate multi-horizon forecast accuracy on a snapshot sequence.

Parameters:

sequence (GraphSnapshotSequence or sequence of Data) – Evaluation snapshots with shared topology.
horizons (sequence of int, optional) – Forecast horizons to report. Default is (3, 6, 12).
start_indices (sequence of int or None, optional) – Forecast-origin indices. When None, uses every valid origin in sequence.

Returns:

Per-horizon and aggregate MAE, RMSE, and MAPE.

Return type:

Train encoder, Koopman operator, and decoder end-to-end.

Minimizes a weighted sum of one-step MSE and optional forward and backward consistency terms:

loss = w_r * recon_loss
     + w_f * ||K z_t - z_{t+1}||^2
     + w_b * ||z_t - z_{t+1} K^{\dagger}||^2

where z_t and z_{t+1} are encoder outputs for consecutive snapshots and weights (w_r, w_f, w_b) come from a LossWeights object or an optional per-epoch schedule.

When data_sequence is a list of GraphSnapshotSequence objects, losses are averaged across trajectories before each optimizer step.

Parameters:

data_sequence (GraphSnapshotSequence, sequence of Data, or sequence of GraphSnapshotSequence) – One training trajectory or multiple trajectories of the same system. A plain list of Data snapshots is treated as a single trajectory; a list whose first element is a GraphSnapshotSequence is treated as multiple trajectories.
epochs (int, optional) – Number of training epochs. Default is 100.
lr (float, optional) – Learning rate passed to the optimizer. Default is 1e-3.
optimizer (callable, optional) – Optimizer class. Default is torch.optim.Adam.
device (str, torch.device, or None, optional) – Device for training. Defaults to the model’s current device, or CPU if the model has no parameters.
loss_weights (LossWeights or None, optional) – Static loss weights for all epochs. When None and no schedule is provided, defaults to reconstruction-only training.
loss_weight_schedule (callable or None, optional) – Callable epoch -> LossWeights applied each epoch. Overrides loss_weights when set.
rollout_horizon (int or None, optional) – Number of autoregressive rollout steps used when loss_weights.rollout is non-zero. Defaults to num_timesteps - 1.
rollout_start_indices (sequence of int, "all", or None, optional) – Rollout-loss origin indices. None uses [0]; "all" uses every valid origin for the rollout horizon.
rollout_starts_per_epoch (int or None, optional) – When set, randomly sample this many rollout origins each epoch.
rollout_start_seed (int or None, optional) – Base seed for random rollout-origin sampling. The effective seed is rollout_start_seed + epoch.
lr_scheduler (LRScheduler or callable, optional) – Learning-rate scheduler instance or factory optimizer -> scheduler. Stepped once per epoch after the optimizer update.
window_length (int or None, optional) – Fixed number of snapshots per training window. When set, enables mini-batch training with multiple optimizer steps per epoch. None preserves full-sequence single-step training.
batch_size (int, optional) – Number of temporal windows averaged per optimizer step. Used only when window_length is set. Default is 8.
windows_per_epoch (int or None, optional) – Maximum sampled windows per epoch. None uses every valid window across all trajectories.
window_seed (int or None, optional) – Base seed for reproducible epoch-specific window shuffling.
max_grad_norm (float or None, optional) – When set, clip the global gradient norm before each optimizer step.
early_stopping_patience (int or None, optional) – Stop training when training loss fails to improve for this many consecutive epochs. Disabled when None.
early_stopping_min_delta (float, optional) – Minimum decrease in the monitored loss to count as improvement. Default is 0.0.
early_stopping_monitor ({"auto", "train", "val"}, optional) – Loss used for early stopping and best-epoch tracking. "auto" monitors validation loss when validation_sequence is provided, otherwise training loss. Default is "auto".
validation_sequence (GraphSnapshotSequence, sequence of Data, sequence of GraphSnapshotSequence, or None, optional) – Optional held-out snapshots for per-epoch validation loss. A single validation sequence is reused for all training trajectories; a list of validation sequences must match the training trajectory count.
restore_best_weights (bool, optional) – When True, reload in-memory weights from the lowest-loss epoch after training completes. Default is False.
checkpoint_path (str, Path, or None, optional) – When set, write a checkpoint at the lowest-loss epoch using save(). Default is None.
**optimizer_kwargs (Any) – Additional keyword arguments forwarded to the optimizer constructor.

Returns:

Per-epoch training and validation losses and early-stop metadata.

Return type:

FitHistory

Raises:

Encoders¶

Graph Neural Network encoders for topology-aware latent lifting.

class koopman_graph.encoder.BaseGNNModule(*, input_channels, input_dim_name, num_layers, activation, convs)[source]¶

Bases: Module

Shared message-passing stack for GNN encoders and decoders.

input_channels¶

Expected input node feature dimension.

Type:: int

input_dim_name¶

Name of the input dimension used in validation errors.

Type:: str

num_layers¶

Number of message-passing layers.

Type:: int

activation_name¶

Identifier for the hidden-layer activation.

Type:: str

activation¶

Instantiated hidden-layer activation.

Type:: nn.Module

convs¶

Ordered graph convolution layers.

Type:: nn.ModuleList

forward(x_or_data, edge_index=None, edge_weight=None)[source]¶

Run the stacked message-passing layers on graph node features.

Parameters:

x_or_data (Tensor or Data) – Either a PyG Data object or node features x.
edge_index (Tensor or None, optional) – Edge index required when x_or_data is a tensor.
edge_weight (Tensor or None, optional) – Scalar edge weights with shape (num_edges,). Passed to GCNConv when present. Ignored by GAT layers.

Returns:

Transformed node features with shape (num_nodes, out_channels).

Return type:

Tensor

class koopman_graph.encoder.GNNEncoder(in_channels, hidden_channels, latent_dim, *, num_layers=2, activation='relu')[source]¶

Bases: BaseGNNModule

GCN encoder that lifts node features into a latent space.

in_channels¶

Input node feature dimension.

Type:: int

hidden_channels¶

Hidden GCN channel width.

Type:: int

latent_dim¶

Output latent dimension per node.

Type:: int

class koopman_graph.encoder.GATEncoder(in_channels, hidden_channels, latent_dim, *, num_layers=2, activation='relu', heads=1, dropout=0.0)[source]¶

Bases: BaseGNNModule

GAT encoder that lifts node features into a latent space.

Scalar edge_weight arguments are accepted for API symmetry with GNNEncoder but are ignored because GATConv does not consume scalar edge weights.

in_channels¶

Input node feature dimension.

Type:: int

hidden_channels¶

Hidden GAT channel width.

Type:: int

latent_dim¶

Output latent dimension per node.

Type:: int

heads¶

Number of attention heads per GAT layer.

Type:: int

dropout¶

Dropout probability inside GAT attention.

Type:: float

Decoder¶

Graph Neural Network decoders for latent-to-physical reconstruction.

class koopman_graph.decoder.GNNDecoder(latent_dim, hidden_channels, out_channels, *, num_layers=2, activation='relu')[source]¶

Bases: BaseGNNModule

GCN decoder that maps latent node features back to physical space.

latent_dim¶

Input latent dimension per node.

Type:: int

hidden_channels¶

Hidden GCN channel width.

Type:: int

out_channels¶

Output physical feature dimension per node.

Type:: int

Koopman Operator¶

Finite-dimensional Koopman operator for latent-state linear propagation.

class koopman_graph.operator.KoopmanOperator(latent_dim, *, init_mode='identity_noise', init_scale=0.01, parameterization='dense', max_spectral_radius=1.0, control_dim=0)[source]¶

Bases: Module

Learnable finite-dimensional Koopman operator matrix K.

Applies the same linear map to each node’s latent vector. For input z with trailing dimension latent_dim, the uncontrolled forward pass computes:

z_next = z @ K.T

When control_dim is positive, exogenous inputs drive the transition:

z_next = z @ K.T + u @ B

latent_dim¶

Dimension of the latent space.

Type:: int

control_dim¶

Dimension of exogenous control inputs. Zero disables control.

Type:: int

init_mode¶

Weight initialization strategy for K.

Type:: str

init_scale¶

Noise scale used when init_mode="identity_noise".

Type:: float

parameterization¶

Parameterization used for K ("dense" or "odo").

Type:: str

max_spectral_radius¶

Upper bound on the spectral radius when parameterization="odo".

Type:: float

reset_control_parameters()[source]¶

Reinitialize the control input matrix B.

Return type:: None

control_term(u, *, num_nodes=None)[source]¶

Map control inputs to a latent-space offset u @ B.

Parameters:

u (Tensor) – Global control with shape (control_dim,) or per-node control with shape (num_nodes, control_dim).
num_nodes (int or None, optional) – Expected node count when u is per-node. Used for validation only.

Returns:

Latent offset with shape (latent_dim,) for global control or (num_nodes, latent_dim) for per-node control.

Return type:

Tensor

Raises:

ValueError – If control_dim is zero, u has invalid shape, or per-node u does not match num_nodes.

property K: Tensor¶

Assembled Koopman matrix with shape (latent_dim, latent_dim).

For parameterization="dense" this is the learnable parameter. For parameterization="odo" it is assembled from orthogonal factors and a bounded diagonal matrix.

Returns:: Current operator matrix K.
Return type:: Tensor

reset_parameters()[source]¶

Reinitialize operator parameters according to init_mode.

Return type:: None

spectral_radius()[source]¶

Return the spectral radius of the assembled operator matrix.

Returns:: Scalar tensor with the maximum eigenvalue magnitude.
Return type:: Tensor

forward(z, control=None)[source]¶

Advance latent states by one linear Koopman step.

When control_dim is positive, the controlled update is:

z_next = z @ K.T + control_effect

where control_effect is u @ B broadcast for global control u with shape (control_dim,) or applied per node when u has shape (num_nodes, control_dim).

Parameters:

z (Tensor) – Latent states with shape (..., latent_dim).
control (Tensor or None, optional) – Exogenous control input applied during this step. Required when control_dim is positive.

Returns:

Advanced latent states with the same shape as z.

Return type:

Tensor

Raises:

ValueError – If the trailing dimension of z does not match latent_dim, controls are missing for a controlled operator, or control has an invalid shape.

inverse_step(z, *, control=None, inverse_matrix=None)[source]¶

Apply one inverse Koopman step to recover the previous latent state.

For forward dynamics z_{t+1} = z_t @ K.T + u_t @ B, this returns an estimate of z_t from z_{t+1} and the control u_t that drove the transition.

Parameters:

z (Tensor) – Latent states at time t+1 with shape (..., latent_dim).
control (Tensor or None, optional) – Control input that drove the forward transition. Required when control_dim is positive.
inverse_matrix (Tensor or None, optional) – Precomputed K^{-1} for dense parameterization. When omitted, the inverse is computed on demand.

Returns:

Recovered latent states at time t, same shape as z.

Return type:

Tensor

dense_inverse_matrix()[source]¶

Return the inverse (or pseudo-inverse) of the assembled dense matrix.

Intended for reuse across multiple backward-consistency pair evaluations within one training step.

Returns:: Matrix K^{-1} (or K^{\dagger}) with shape (latent_dim, latent_dim).
Return type:: Tensor
Raises:: ValueError – If parameterization is not "dense".

Spectral Analysis¶

Spectral analysis utilities for finite-dimensional Koopman operators.

class koopman_graph.analysis.KoopmanSpectrum(eigenvalues, eigenvectors, magnitudes, growth_rates, frequencies, time_step)[source]¶

Bases: object

Eigendecomposition and time scales of a discrete Koopman operator.

Eigenpairs are sorted by descending eigenvalue magnitude. Frequencies are reported in cycles per unit time; multiply by 2 * pi for angular frequency.

eigenvalues¶

Complex eigenvalues with shape (latent_dim,).

Type:: Tensor

eigenvectors¶

Complex right eigenvectors stored as columns, with shape (latent_dim, latent_dim).

Type:: Tensor

magnitudes¶

Eigenvalue magnitudes with shape (latent_dim,).

Type:: Tensor

growth_rates¶

Continuous-time exponential growth rates log(|lambda|) / time_step.

Type:: Tensor

frequencies¶

Signed continuous-time frequencies angle(lambda) / (2 * pi * time_step) in cycles per unit time.

Type:: Tensor

time_step¶

Physical duration represented by one discrete Koopman step.

Type:: float

mode_amplitudes(latent_states)[source]¶

Project latent states onto the Koopman eigenvector basis.

For a latent row vector z, the returned amplitudes a satisfy z.T = eigenvectors @ a. Any leading dimensions are preserved.

Parameters:

latent_states (Tensor) – Latent states with shape (..., latent_dim).

Returns:

Complex mode amplitudes with the same shape as latent_states.

Return type:

Tensor

Raises:

ValueError – If the trailing latent dimension does not match the spectrum.
RuntimeError – If the eigenvector matrix is singular.

koopman_graph.analysis.compute_spectrum(operator, time_step)[source]¶

Compute the sorted spectrum and continuous-time mode characteristics.

Parameters:

operator (Tensor) – Square discrete-time Koopman matrix with shape (latent_dim, latent_dim).
time_step (float) – Positive physical duration represented by one operator step.

Returns:

Eigenpairs sorted by descending magnitude, plus growth rates and frequencies converted using time_step.

Return type:

KoopmanSpectrum

Raises:

ValueError – If operator is not a non-empty square matrix or time_step is not positive.
TypeError – If operator is not floating-point or complex.

koopman_graph.analysis.decode_mode_shapes(model, x_or_data, mode_indices=None, *, edge_index=None, perturbation=0.001)[source]¶

Decode latent Koopman directions into spatial node-feature mode shapes.

Parameters:

model (GraphKoopmanModel) – Model whose operator spectrum and decoder are analyzed.
x_or_data (Tensor or Data) – Reference graph used as the decoder linearization point.
mode_indices (sequence of int or None, optional) – Indices into the magnitude-sorted spectrum. Defaults to every mode.
edge_index (Tensor or None, optional) – Graph edges, required when x_or_data is a feature tensor.
perturbation (float, optional) – Positive centered finite-difference step. Default is 1e-3.

Returns:

Complex mode shapes with shape (num_modes, num_nodes, out_channels).

Return type:

Tensor

Raises:

ValueError – If perturbation is not positive or a mode index is out of range.

Baselines¶

Classical topology-agnostic Koopman baselines.

class koopman_graph.baselines.DMDBaseline(*, time_step=1.0, rank=None)[source]¶

Bases: object

Dynamic Mode Decomposition baseline on flattened node states.

Parameters:

time_step (float, optional) – Physical duration represented by one snapshot transition. Used by spectrum(). Default is 1.0.
rank (int or None, optional) – Optional truncated-SVD rank for the data matrix. None uses the full least-squares solution. Default is None.

fit(sequence)[source]¶

Fit the DMD operator from consecutive graph snapshots.

Parameters:: sequence (GraphSnapshotSequence or sequence of Data) – Training snapshots with shared topology.
Returns:: The fitted baseline.
Return type:: DMDBaseline
Raises:: ValueError – If fewer than two snapshots are provided or rank is invalid.

predict(initial_graph, steps)[source]¶

Autoregressively predict future graph snapshots.

Parameters:

initial_graph (Data) – Initial graph snapshot. Its topology is copied to every prediction.
steps (int) – Number of future snapshots to predict.

Returns:

Predicted graph snapshots with the same node/feature shape as the fitted training data.

Return type:

list of Data

spectrum()[source]¶

Return the DMD operator spectrum.

Returns:: Eigendecomposition and continuous-time mode characteristics of the fitted DMD operator.
Return type:: KoopmanSpectrum

class koopman_graph.baselines.DMDcBaseline(*, time_step=1.0, rank=None)[source]¶

Bases: object

Dynamic Mode Decomposition with control on flattened node states.

Parameters:

time_step (float, optional) – Physical duration represented by one snapshot transition. Used by spectrum(). Default is 1.0.
rank (int or None, optional) – Optional truncated-SVD rank for the augmented regression. None uses the full least-squares solution. Default is None.

fit(sequence)[source]¶

Fit controlled DMD operators from consecutive graph snapshots.

Parameters:: sequence (GraphSnapshotSequence or sequence of Data) – Training snapshots with shared topology and control inputs.
Returns:: The fitted baseline.
Return type:: DMDcBaseline
Raises:: ValueError – If fewer than two snapshots are provided, controls are missing, or rank is invalid.

predict(initial_graph, steps, controls)[source]¶

Autoregressively predict future graph snapshots with future controls.

Parameters:

initial_graph (Data) – Initial graph snapshot. Its topology is copied to every prediction.
steps (int) – Number of future snapshots to predict.
controls (sequence of Tensor) – Future control inputs, one per rollout step. Global controls use shape (control_dim,); per-node controls use (num_nodes, control_dim).

Returns:

Predicted graph snapshots with the same node/feature shape as the fitted training data.

Return type:

list of Data

spectrum()[source]¶

Return the autonomous DMD operator spectrum.

Returns:: Eigendecomposition of the fitted state-transition operator K.
Return type:: KoopmanSpectrum

class koopman_graph.baselines.EDMDBaseline(*, time_step=1.0, rank=None, polynomial_degree=2)[source]¶

Bases: object

Extended DMD baseline with polynomial observables.

Parameters:

time_step (float, optional) – Physical duration represented by one snapshot transition. Used by spectrum(). Default is 1.0.
rank (int or None, optional) – Optional truncated-SVD rank for the observable data matrix. None uses the full least-squares solution. Default is None.
polynomial_degree ({1, 2}, optional) – Polynomial observable degree. Default is 2.

fit(sequence)[source]¶

Fit EDMD operator and linear reconstruction decoder.

Parameters:: sequence (GraphSnapshotSequence or sequence of Data) – Training snapshots with shared topology.
Returns:: The fitted baseline.
Return type:: EDMDBaseline
Raises:: ValueError – If fewer than two snapshots are provided or rank is invalid.

predict(initial_graph, steps)[source]¶

Autoregressively predict future graph snapshots.

Parameters:

initial_graph (Data) – Initial graph snapshot. Its topology is copied to every prediction.
steps (int) – Number of future snapshots to predict.

Returns:

Predicted graph snapshots with the same node/feature shape as the fitted training data.

Return type:

list of Data

Raises:

RuntimeError – If the baseline has not been fit.
ValueError – If steps < 1 or graph metadata does not match the fit data.

spectrum()[source]¶

Return the EDMD observable-space operator spectrum.

Returns:: Eigendecomposition and continuous-time mode characteristics of the fitted observable-space operator.
Return type:: KoopmanSpectrum

Data Utilities¶

Utilities for spatiotemporal graph snapshot sequences.

class koopman_graph.data.WindowSampler(sequences, *, window_length, batch_size=8, windows_per_epoch=None, shuffle=True, seed=None)[source]¶

Bases: object

Sample fixed-length temporal windows from one or more trajectories.

Parameters:

sequences (GraphSnapshotSequence or sequence of GraphSnapshotSequence) – Source trajectories. Each must contain at least window_length snapshots.
window_length (int) – Number of snapshots per sampled window. Must be at least 2.
batch_size (int, optional) – Number of windows yielded together. Default is 8.
windows_per_epoch (int or None, optional) – Maximum number of windows sampled per epoch. None uses every valid window. Values larger than the available window count are capped.
shuffle (bool, optional) – Randomize window order each epoch. Default is True.
seed (int or None, optional) – Base seed for reproducible epoch-specific shuffling.

property num_windows: int¶

Return the total number of valid windows.

Returns:: Number of valid windows across every source trajectory.
Return type:: int

iter_epoch(epoch=0)[source]¶

Yield batches of windows for one epoch.

Parameters:: epoch (int, optional) – Zero-based epoch index mixed into seed. Default is 0.
Yields:: list of GraphSnapshotSequence – A batch containing at most batch_size temporal windows.
Return type:: Iterator[list[GraphSnapshotSequence]]

class koopman_graph.data.TemporalSplit(train, val, test)[source]¶

Bases: object

Train, validation, and test snapshot sequences from a temporal split.

train¶

Earliest contiguous snapshots used for training.

Type:: GraphSnapshotSequence

val¶

Middle contiguous snapshots used for validation.

Type:: GraphSnapshotSequence

test¶

Latest contiguous snapshots held out for evaluation.

Type:: GraphSnapshotSequence

koopman_graph.data.temporal_split(sequence, *, train_ratio=0.7, val_ratio=0.1, test_ratio=0.2, min_train_timesteps=2, min_val_timesteps=2, min_test_timesteps=1)[source]¶

Split a snapshot sequence into contiguous train, validation, and test sets.

Earlier snapshots are assigned to training, later snapshots to validation and test. Ratios must sum to 1.0.

Parameters:

sequence (GraphSnapshotSequence) – Full time-ordered snapshot sequence to split.
train_ratio (float, optional) – Fraction of timesteps assigned to training. Default is 0.7.
val_ratio (float, optional) – Fraction assigned to validation. Default is 0.1.
test_ratio (float, optional) – Fraction assigned to test. Default is 0.2.
min_train_timesteps (int, optional) – Minimum training snapshots required. Default is 2.
min_val_timesteps (int, optional) – Minimum validation snapshots required. Default is 2.
min_test_timesteps (int, optional) – Minimum test snapshots required. Default is 1.

Returns:

Contiguous train, validation, and test sequences sharing topology.

Return type:

TemporalSplit

Raises:

ValueError – If ratios do not sum to 1.0, any minimum is violated, or the sequence is too short for the requested split.

class koopman_graph.data.GraphSnapshotSequence(snapshots, *, allow_dynamic_topology=False, control_inputs=None)[source]¶

Bases: object

Container for a time-ordered sequence of PyG Data graph snapshots.

Notes

classmethod from_arrays(node_features, edge_index, *, edge_weight=None, control_inputs=None, dtype=torch.float32)[source]¶

Build a sequence from node feature arrays and a shared topology.

Parameters:

node_features (array-like) – Array with shape (num_timesteps, num_nodes, in_channels).
edge_index (array-like) – Shared edge index with shape (2, num_edges).
edge_weight (array-like, optional) – Shared scalar edge weights with shape (num_edges,). When provided, attached to every snapshot.
control_inputs (array-like, optional) – Per-timestep control inputs with shape (num_timesteps, control_dim) or (num_timesteps, num_nodes, control_dim).
dtype (torch.dtype, optional) – Floating dtype used when converting numpy inputs to torch tensors. Default is torch.float32.

Returns:

Validated snapshot sequence.

Return type:

Raises:

ValueError – If node_features, edge_index, or edge_weight have invalid shape.

classmethod from_dynamic_arrays(node_features, edge_indices, *, edge_weights=None, control_inputs=None, dtype=torch.float32)[source]¶

Build a dynamic-topology sequence from per-timestep edge indices.

Parameters:

node_features (array-like) – Array with shape (num_timesteps, num_nodes, in_channels).
edge_indices (sequence of array-like) – One edge index per timestep, each with shape (2, num_edges_t).
edge_weights (sequence of array-like or None, optional) – Optional per-timestep scalar edge weights aligned with edge_indices. When provided, must have the same length as edge_indices.
control_inputs (array-like, optional) – Per-timestep control inputs with shape (num_timesteps, control_dim) or (num_timesteps, num_nodes, control_dim).
dtype (torch.dtype, optional) – Floating dtype used when converting numpy inputs to torch tensors. Default is torch.float32.

Returns:

Validated snapshot sequence with allow_dynamic_topology=True.

Return type:

Raises:

ValueError – If shapes are inconsistent or edge_indices length mismatches num_timesteps.

property is_dynamic_topology: bool¶

Return whether snapshots use time-varying edge_index.

Returns:: True when the sequence was constructed with allow_dynamic_topology=True and at least one snapshot differs in edge_index from the first snapshot.
Return type:: bool

property allow_dynamic_topology: bool¶

Return whether dynamic topology mode was enabled at construction.

Returns:: True when per-snapshot edge_index values are permitted.
Return type:: bool

property control_inputs: Tensor | None¶

Return per-timestep control inputs when present.

Returns:: Control tensor with shape (num_timesteps, control_dim) or (num_timesteps, num_nodes, control_dim).
Return type:: Tensor or None

property has_controls: bool¶

Return whether the sequence carries control inputs.

Returns:: True when control_inputs is not None.
Return type:: bool

property control_dim: int¶

Return the control feature dimension.

Returns:: Control dimension when controls are present, otherwise 0.
Return type:: int

control_at(index)[source]¶

Return the control input driving transition from snapshot index.

Parameters:: index (int) – Timestep index in [0, num_timesteps - 1].
Returns:: Control vector for the transition index -> index + 1.
Return type:: Tensor
Raises:: ValueError – If controls are absent or index is out of range.

rollout_controls(start, steps)[source]¶

Return controls for an autoregressive rollout from a start snapshot.

Parameters:

start (int) – Index of the initial snapshot.
steps (int) – Number of rollout steps.

Returns:

Control inputs for each rollout step. Empty when the sequence has no controls.

Return type:

list of Tensor

Raises:

ValueError – If start or steps are invalid or controls are unavailable for the requested horizon.

property snapshots: list[Data]¶

Return the underlying list of graph snapshots.

Returns:: Time-ordered PyG graph snapshots.
Return type:: list of Data

property edge_index: Tensor¶

Return the shared edge index for static-topology sequences.

Returns:: Edge index with shape (2, num_edges).
Return type:: Tensor
Raises:: ValueError – If is_dynamic_topology is True.

property edge_weight: Tensor | None¶

Return the shared scalar edge weights for static-topology sequences.

Returns:: Edge weights with shape (num_edges,), or None when the sequence is unweighted.
Return type:: Tensor or None
Raises:: ValueError – If is_dynamic_topology is True.

property num_nodes: int¶

Return the number of nodes in the graph topology.

Returns:: Node count shared across all snapshots.
Return type:: int

property num_timesteps: int¶

Return the number of timesteps in the sequence.

Returns:: Length of the temporal sequence.
Return type:: int

property in_channels: int¶

Return the node feature dimension.

Returns:: Feature dimension shared across all snapshots.
Return type:: int

slice(start, stop)[source]¶

Return a contiguous temporal sub-sequence.

Parameters:

start (int) – Inclusive start index.
stop (int) – Exclusive stop index.

Returns:

Snapshots in [start, stop) with matching controls and topology policy.

Return type:

GraphSnapshotSequence

Raises:

ValueError – If the bounds are negative, empty, reversed, or exceed the sequence length.

koopman_graph.data.resolve_sequence(sequence)[source]¶

Normalize input into a validated snapshot sequence.

Wraps a plain sequence of Data snapshots in GraphSnapshotSequence; existing sequences are returned unchanged.

Parameters:: sequence (GraphSnapshotSequence or sequence of Data) – Raw snapshot input from a training, baseline, or inference API.
Returns:: Validated sequence container.
Return type:: GraphSnapshotSequence

Losses¶

Loss functions for Koopman graph dynamics training.

class koopman_graph.losses.ForwardConsistencyLoss(*args, **kwargs)[source]¶

Bases: Module

Penalize deviation from linear latent evolution under the Koopman operator.

For latent encodings z_t and z_{t+1}, the loss is the mean squared error between K @ z_t (implemented as z_t @ K.T) and z_{t+1}:

\[\mathcal{L}_{\mathrm{fc}} = \| K z_t - z_{t+1} \|^2\]

Notes

This module is stateless. Call forward() with consecutive latent encodings and a KoopmanOperator.

forward(z_t, z_t1, koopman, *, control=None)[source]¶

Compute forward consistency loss between consecutive latent states.

Parameters:

z_t (Tensor) – Latent encoding at time t, shape (..., latent_dim).
z_t1 (Tensor) – Latent encoding at time t+1, same shape as z_t.
koopman (KoopmanOperator) – Learnable linear propagator applied to z_t.
control (Tensor or None, optional) – Control input driving the transition from t to t+1.

Returns:

Scalar mean-squared error between koopman(z_t, control) and z_t1.

Return type:

Tensor

class koopman_graph.losses.BackwardConsistencyLoss(*args, **kwargs)[source]¶

Bases: Module

Penalize deviation from inverse linear latent evolution under K.

\[\mathcal{L}_{\mathrm{bc}} = \| z_t - z_{t+1} K^{\dagger} \|^2\]

where K^{\dagger} denotes an inverse or pseudo-inverse of K.

Trade-offs¶

Notes

This module is stateless. Call forward() with consecutive latent encodings and a KoopmanOperator.

forward(z_t, z_t1, koopman, *, control=None, inverse_matrix=None)[source]¶

Compute backward consistency loss between consecutive latent states.

Parameters:

z_t (Tensor) – Latent encoding at time t, shape (..., latent_dim).
z_t1 (Tensor) – Latent encoding at time t+1, same shape as z_t.
koopman (KoopmanOperator) – Learnable linear propagator whose inverse step is applied to z_t1.
control (Tensor or None, optional) – Control input that drove the forward transition from t to t+1.
inverse_matrix (Tensor or None, optional) – Precomputed dense inverse matrix reused across pair evaluations.

Returns:

Scalar mean-squared error between z_t and the inverse propagation of z_t1.

Return type:

Tensor

class koopman_graph.losses.EigenvalueRegularizationLoss(*args, **kwargs)[source]¶

Bases: Module

Penalize Koopman eigenvalues outside the unit circle.

Implements a hinge-style eigenloss that activates only when eigenvalue magnitudes exceed one:

\[\mathcal{L}_{\mathrm{eig}} = \mathrm{mean}\big(\max(|\lambda_i| - 1, 0)^2\big)\]

For the ODO parameterization, eigenvalues are read directly from the bounded diagonal factor, avoiding an explicit eigendecomposition.

Trade-offs¶

Costs: Dense K requires torch.linalg.eigvals each evaluation. Use the ODO parameterization when a hard spectral-radius bound is preferred.

Notes

This module is stateless. Call forward() with a KoopmanOperator.

forward(koopman)[source]¶

Compute the unit-circle eigenvalue hinge penalty.

Parameters:: koopman (KoopmanOperator) – Operator whose eigenvalue magnitudes are penalized.
Returns:: Scalar hinge penalty (zero when all magnitudes are <= 1).
Return type:: Tensor

koopman_graph.losses.rollout_sequence_loss(model, sequence, *, horizon, start=0)[source]¶

Compute autoregressive rollout reconstruction loss from one start snapshot.

Encodes sequence[start] once, advances the latent state with the Koopman operator for horizon steps, and compares decoded predictions to the observed snapshots sequence[start + 1 : start + horizon + 1]. This term aligns training with predict().

Parameters:

model (nn.Module) – Model with encoder, koopman, and decoder attributes.
sequence (GraphSnapshotSequence) – Time-ordered snapshots. For dynamic-topology sequences, each decode step uses the target snapshot’s edge_index.
horizon (int) – Number of rollout steps (must be >= 1).
start (int, optional) – Index of the initial snapshot. Default is 0.

Returns:

Scalar mean rollout reconstruction loss over horizon steps.

Return type:

Tensor

Raises:

ValueError – If horizon < 1, start < 0, or the sequence is too short.

koopman_graph.losses.rollout_multi_start_loss(model, sequence, *, horizon, start_indices)[source]¶

Average rollout reconstruction loss over multiple start snapshots.

Parameters:

model (nn.Module) – Model with encoder, koopman, and decoder attributes.
sequence (GraphSnapshotSequence) – Time-ordered snapshots.
horizon (int) – Number of rollout steps (must be >= 1).
start_indices (sequence of int) – Zero-based origin indices for each rollout.

Returns:

Scalar mean rollout loss across origins.

Return type:

Tensor

Raises:

ValueError – If start_indices is empty or any origin is invalid.

Training¶

Training utilities for GraphKoopmanModel.

class koopman_graph.training.LossWeights(reconstruction=1.0, forward=0.0, backward=0.0, rollout=0.0, eigenvalue=0.0)[source]¶

Bases: object

Weights for reconstruction and consistency loss terms.

reconstruction¶

Weight on the one-step reconstruction (MSE) loss.

Type:: float

forward¶

Weight on the forward consistency loss.

Type:: float

backward¶

Weight on the backward consistency loss.

Type:: float

rollout¶

Weight on the autoregressive rollout reconstruction loss.

Type:: float

eigenvalue¶

Weight on the unit-circle eigenvalue hinge penalty.

Type:: float

class koopman_graph.training.TrainingLossBreakdown(reconstruction, forward, backward, rollout, eigenvalue, total)[source]¶

Bases: object

Unweighted and weighted training loss terms for one batch or epoch.

reconstruction¶

Mean one-step reconstruction loss.

Type:: Tensor

forward¶

Mean forward consistency loss.

Type:: Tensor

backward¶

Mean backward consistency loss.

Type:: Tensor

rollout¶

Mean rollout reconstruction loss.

Type:: Tensor

eigenvalue¶

Eigenvalue hinge regularization loss.

Type:: Tensor

total¶

Weighted sum of all active loss terms.

Type:: Tensor

classmethod zeros(device)[source]¶

Return a zero breakdown on device.

Parameters:: device (torch.device) – Device for the zero tensors.
Returns:: Breakdown with all scalar terms set to zero.
Return type:: TrainingLossBreakdown

to_floats()[source]¶

Convert tensor terms to detached Python floats.

Returns:: Mapping with keys reconstruction, forward, backward, rollout, eigenvalue, and total.
Return type:: dict of str to float

koopman_graph.training.mean_training_loss_breakdown(breakdowns)[source]¶

Average loss breakdowns across multiple trajectories.

Parameters:: breakdowns (sequence of TrainingLossBreakdown) – Per-trajectory breakdowns to average.
Returns:: Element-wise mean across breakdowns.
Return type:: TrainingLossBreakdown
Raises:: ValueError – If breakdowns is empty.

class koopman_graph.training.FitHistory(loss, epochs, reconstruction_loss=<factory>, forward_loss=<factory>, backward_loss=<factory>, rollout_loss=<factory>, eigenvalue_loss=<factory>, val_loss=None, val_reconstruction_loss=None, val_forward_loss=None, val_backward_loss=None, val_rollout_loss=None, val_eigenvalue_loss=None, stopped_early=False, best_epoch=None, best_loss=None)[source]¶

Bases: object

Training history returned by GraphKoopmanModel.fit().

loss¶

Per-epoch average combined training loss.

Type:: list of float

epochs¶

Number of training epochs completed (may be less than requested when early stopping triggers).

Type:: int

reconstruction_loss¶

Per-epoch unweighted reconstruction loss.

Type:: list of float

forward_loss¶

Per-epoch unweighted forward consistency loss.

Type:: list of float

backward_loss¶

Per-epoch unweighted backward consistency loss.

Type:: list of float

rollout_loss¶

Per-epoch unweighted rollout reconstruction loss.

Type:: list of float

eigenvalue_loss¶

Per-epoch unweighted eigenvalue regularization loss.

Type:: list of float

val_loss¶

Per-epoch validation loss when a validation sequence is provided.

Type:: list of float or None

val_reconstruction_loss¶

Per-epoch unweighted validation reconstruction loss.

Type:: list of float or None

val_forward_loss¶

Per-epoch unweighted validation forward consistency loss.

Type:: list of float or None

val_backward_loss¶

Per-epoch unweighted validation backward consistency loss.

Type:: list of float or None

val_rollout_loss¶

Per-epoch unweighted validation rollout loss.

Type:: list of float or None

val_eigenvalue_loss¶

Per-epoch unweighted validation eigenvalue loss.

Type:: list of float or None

stopped_early¶

Whether training stopped before the requested epoch count.

Type:: bool

best_epoch¶

Zero-based index of the lowest-loss epoch when best-epoch tracking is enabled in fit().

Type:: int or None

best_loss¶

Lowest monitored loss observed when best-epoch tracking is enabled.

Type:: float or None

koopman_graph.training.constant_loss_weights(*, reconstruction=1.0, forward=0.0, backward=0.0, rollout=0.0, eigenvalue=0.0)[source]¶

Build static loss weights for all training epochs.

Parameters:

reconstruction (float, optional) – Weight on the reconstruction loss. Default is 1.0.
forward (float, optional) – Weight on the forward consistency loss. Default is 0.0.
backward (float, optional) – Weight on the backward consistency loss. Default is 0.0.
rollout (float, optional) – Weight on the rollout reconstruction loss. Default is 0.0.
eigenvalue (float, optional) – Weight on the eigenvalue hinge penalty. Default is 0.0.

Returns:

Fixed weights applied uniformly across epochs.

Return type:

LossWeights

koopman_graph.training.linear_ramp_loss_weights(start, end, ramp_epochs)[source]¶

Linearly interpolate loss weights over the first ramp_epochs.

Parameters:

start (LossWeights) – Weights at epoch 0.
end (LossWeights) – Weights reached at epoch ramp_epochs - 1 and held thereafter.
ramp_epochs (int) – Number of epochs over which to interpolate. Must be >= 1.

Returns:

Schedule mapping epoch index to LossWeights.

Return type:

callable

Raises:

ValueError – If ramp_epochs < 1.

koopman_graph.training.one_step_loss(model, snapshot_t, snapshot_t1, *, control=None)[source]¶

Compute one-step MSE between model prediction and the next snapshot.

Parameters:

model (nn.Module) – Model implementing a single-step forward pass (e.g. GraphKoopmanModel).
snapshot_t (Data) – Graph snapshot at time t.
snapshot_t1 (Data) – Graph snapshot at time t+1 (prediction target).
control (Tensor or None, optional) – Control input driving the transition from t to t+1.

Returns:

Scalar mean-squared error loss.

Return type:

Tensor

koopman_graph.training.compute_sequence_loss(model, sequence)[source]¶

Average one-step prediction loss over consecutive snapshot pairs.

Parameters:

model (nn.Module) – Model implementing a single-step forward pass.
sequence (GraphSnapshotSequence) – Time-ordered snapshots with at least two timesteps.

Returns:

Scalar average loss over all consecutive pairs.

Return type:

Tensor

Raises:

ValueError – If sequence contains fewer than two snapshots.

koopman_graph.training.compute_forward_consistency_sequence_loss(model, sequence)[source]¶

Average forward consistency loss over consecutive snapshot pairs.

Parameters:

model (nn.Module) – Model with encoder and koopman attributes.
sequence (GraphSnapshotSequence) – Time-ordered snapshots with at least two timesteps.

Returns:

Scalar average forward consistency loss.

Return type:

Tensor

Raises:

ValueError – If sequence contains fewer than two snapshots.

koopman_graph.training.compute_backward_consistency_sequence_loss(model, sequence)[source]¶

Average backward consistency loss over consecutive snapshot pairs.

Parameters:

model (nn.Module) – Model with encoder and koopman attributes.
sequence (GraphSnapshotSequence) – Time-ordered snapshots with at least two timesteps.

Returns:

Scalar average backward consistency loss.

Return type:

Tensor

Raises:

ValueError – If sequence contains fewer than two snapshots.

koopman_graph.training.compute_eigenvalue_regularization_loss(model)[source]¶

Compute the eigenvalue hinge penalty for the model Koopman operator.

Parameters:: model (nn.Module) – Model with a koopman attribute.
Returns:: Scalar eigenvalue regularization loss.
Return type:: Tensor

koopman_graph.training.resolve_rollout_start_indices(sequence, *, horizon, rollout_start_indices=None, rollout_starts_per_epoch=None, rollout_start_seed=None, epoch=0)[source]¶

Resolve rollout-loss origin indices for one training epoch.

Parameters:

sequence (GraphSnapshotSequence) – Training snapshots.
horizon (int) – Rollout horizon in steps.
rollout_start_indices (sequence of int, "all", or None, optional) – Explicit origins, all valid origins, or None for the default [0].
rollout_starts_per_epoch (int or None, optional) – When set, randomly sample this many valid origins each epoch. Overrides rollout_start_indices except when that argument is an explicit non-empty index list.
rollout_start_seed (int or None, optional) – Base seed for random origin sampling. The effective seed is rollout_start_seed + epoch when provided.
epoch (int, optional) – Zero-based epoch index mixed into the random seed. Default is 0.

Returns:

Valid zero-based rollout origin indices.

Return type:

list of int

Raises:

ValueError – If horizon is invalid, no origins are available, or an origin is out of range.

koopman_graph.training.compute_rollout_loss(model, sequence, *, horizon, start_indices)[source]¶

Compute rollout reconstruction loss averaged over start origins.

Parameters:

model (nn.Module) – Model with encoder, Koopman operator, and decoder.
sequence (GraphSnapshotSequence) – Training snapshots.
horizon (int) – Number of rollout steps.
start_indices (sequence of int) – Rollout origin indices.

Returns:

Scalar mean rollout loss across origins.

Return type:

Tensor

koopman_graph.training.compute_training_loss(model, sequence, loss_weights, *, rollout_horizon=None, rollout_start_indices=None)[source]¶

Compute reconstruction, consistency, and rollout losses.

Parameters:

model (nn.Module) – Model implementing a single-step forward pass with encoder and Koopman operator.
sequence (GraphSnapshotSequence) – Time-ordered snapshots with at least two timesteps.
loss_weights (LossWeights) – Weights for reconstruction, forward, backward, and rollout terms.
rollout_horizon (int or None, optional) – Number of rollout steps when loss_weights.rollout is non-zero. Defaults to sequence.num_timesteps - 1.
rollout_start_indices (sequence of int or None, optional) – Rollout origin indices. Defaults to [0] when None.

Returns:

Unweighted per-term losses and the weighted total.

Return type:

koopman_graph.training.train_one_epoch(model, sequences, optimizer, loss_weights, *, max_grad_norm=None, rollout_horizon=None, rollout_start_indices=None)[source]¶

Run one training epoch and return the averaged loss breakdown.

Parameters:

model (nn.Module) – Model to train.
sequences (GraphSnapshotSequence or sequence of GraphSnapshotSequence) – One or more training trajectories.
optimizer (Optimizer) – PyTorch optimizer used for the parameter update.
loss_weights (LossWeights) – Weights for reconstruction and consistency terms this epoch.
max_grad_norm (float or None, optional) – When set, clip the global gradient norm to this value before optimizer.step().
rollout_horizon (int or None, optional) – Number of rollout steps when loss_weights.rollout is non-zero.
rollout_start_indices (sequence of int or None, optional) – Rollout origin indices for this epoch.

Returns:

Mean loss breakdown across trajectories.

Return type:

koopman_graph.training.train_windowed_epoch(model, sampler, optimizer, loss_weights, *, epoch=0, max_grad_norm=None, rollout_horizon=None, rollout_start_indices=None, rollout_starts_per_epoch=None, rollout_start_seed=None)[source]¶

Train on mini-batches of fixed-length temporal windows.

Each batch averages its window losses before one optimizer step. The returned breakdown is weighted by the number of windows in each batch, so a smaller final batch does not receive disproportionate weight.

Parameters:

model (nn.Module) – Model to train.
sampler (WindowSampler) – Window sampler defining trajectories, window size, and batch schedule.
optimizer (Optimizer) – Optimizer updated once per yielded batch.
loss_weights (LossWeights) – Active loss weights for the epoch.
epoch (int, optional) – Zero-based epoch index used for sampler shuffling. Default is 0.
max_grad_norm (float or None, optional) – Optional global gradient clipping threshold.
rollout_horizon (int or None, optional) – Rollout horizon. Defaults to window_length - 1.
rollout_start_indices (sequence of int, "all", or None, optional) – Rollout origins relative to each sampled window.
rollout_starts_per_epoch (int or None, optional) – Number of randomly sampled rollout origins.
rollout_start_seed (int or None, optional) – Base seed for rollout-origin sampling.

Returns:

Window-weighted mean loss breakdown for the epoch.

Return type:

koopman_graph.training.eval_one_epoch(model, sequences, loss_weights, *, rollout_horizon=None, rollout_start_indices=None)[source]¶

Compute validation loss for one epoch without parameter updates.

Parameters:

model (nn.Module) – Model to evaluate.
sequences (GraphSnapshotSequence or sequence of GraphSnapshotSequence) – One or more validation trajectories.
loss_weights (LossWeights) – Weights for reconstruction and consistency terms.
rollout_horizon (int or None, optional) – Number of rollout steps when loss_weights.rollout is non-zero.
rollout_start_indices (sequence of int or None, optional) – Rollout origin indices for this epoch.

Returns:

Mean loss breakdown across trajectories.

Return type:

koopman_graph.training.resolve_early_stopping_monitor(monitor, *, has_validation)[source]¶

Resolve which loss early stopping should monitor.

Parameters:

monitor ({"auto", "train", "val"}) – Requested monitor mode.
has_validation (bool) – Whether a validation sequence was provided to fit().

Returns:

Resolved monitor target.

Return type:

{“train”, “val”}

Raises:

ValueError – If monitor="val" but no validation sequence was provided.

koopman_graph.training.is_sequence_of_sequences(data)[source]¶

Return whether data is a list of snapshot sequences.

Parameters:: data (TrainingInput or ValidationInput) – Training or validation input passed to fit().
Returns:: True when the first element is a GraphSnapshotSequence.
Return type:: bool

koopman_graph.training.resolve_training_sequences(data_sequence)[source]¶

Normalize training input into one or more snapshot sequences.

Parameters:: data_sequence (TrainingInput) – Single sequence, list of Data snapshots, or list of sequences.
Returns:: One or more validated training trajectories.
Return type:: list of GraphSnapshotSequence
Raises:: ValueError – If a multi-sequence input is empty.

koopman_graph.training.resolve_validation_sequences(validation_sequence, *, num_training_sequences)[source]¶

Normalize validation input for fit().

A single validation sequence is reused for all training trajectories. A list of validation sequences must match the training trajectory count.

Parameters:

validation_sequence (ValidationInput) – Optional validation data.
num_training_sequences (int) – Number of training trajectories supplied to fit().

Returns:

Validation trajectories aligned with training input.

Return type:

list of GraphSnapshotSequence or None

Raises:

ValueError – If a validation list length does not match num_training_sequences.

koopman_graph.training.resolve_lr_scheduler(lr_scheduler, optimizer)[source]¶

Instantiate an optional learning-rate scheduler.

Parameters:

lr_scheduler (LRScheduler, callable, or None) – Scheduler instance or factory optimizer -> scheduler.
optimizer (Optimizer) – Optimizer passed to a scheduler factory.

Returns:

Resolved scheduler, if any.

Return type:

LRScheduler or None

koopman_graph.training.resolve_device(model, device)[source]¶

Resolve the training device from an explicit argument or model parameters.

Parameters:

model (nn.Module) – Model whose parameter device is used as fallback.
device (str, torch.device, or None) – Explicit device. When None, uses the device of the first model parameter, or CPU if the model has no parameters.

Returns:

Resolved device for training or inference.

Return type:

torch.device

koopman_graph.training.resolve_loss_weights_for_epoch(epoch, *, loss_weights, loss_weight_schedule)[source]¶

Resolve per-epoch loss weights from static weights or a schedule.

Parameters:

epoch (int) – Zero-based epoch index.
loss_weights (LossWeights or None) – Explicit static weights. When None and no schedule is provided, defaults to reconstruction-only weights.
loss_weight_schedule (callable or None) – Optional per-epoch schedule. Takes precedence over loss_weights.

Returns:

Weights to use for the given epoch.

Return type:

LossWeights

koopman_graph.training.should_stop_early(*, epoch_loss, best_loss, epochs_without_improvement, patience, min_delta)[source]¶

Update early-stopping state after an epoch.

Parameters:

epoch_loss (float) – Training loss for the completed epoch.
best_loss (float) – Best loss seen so far.
epochs_without_improvement (int) – Consecutive epochs without sufficient improvement.
patience (int) – Stop after this many non-improving epochs.
min_delta (float) – Minimum decrease in loss to count as improvement.

Returns:

Whether to stop, updated best loss, and updated non-improvement count.

Return type:

tuple of (bool, float, int)

Metrics¶

Forecast evaluation metrics for graph snapshot sequences.

koopman_graph.metrics.mae(prediction, target)[source]¶

Compute mean absolute error.

Parameters:

prediction (Tensor) – Predicted values.
target (Tensor) – Ground-truth values with the same shape as prediction.

Returns:

Scalar mean absolute error.

Return type:

Tensor

koopman_graph.metrics.rmse(prediction, target)[source]¶

Compute root mean squared error.

Parameters:

prediction (Tensor) – Predicted values.
target (Tensor) – Ground-truth values with the same shape as prediction.

Returns:

Scalar root mean squared error.

Return type:

Tensor

koopman_graph.metrics.mape(prediction, target, *, eps=1e-08)[source]¶

Compute mean absolute percentage error.

Parameters:

prediction (Tensor) – Predicted values.
target (Tensor) – Ground-truth values with the same shape as prediction.
eps (float, optional) – Small constant added to the denominator for numerical stability. Default is 1e-8.

Returns:

Scalar mean absolute percentage error (not scaled to 0–100).

Return type:

Tensor

class koopman_graph.metrics.HorizonMetrics(horizon, mae, rmse, mape)[source]¶

Bases: object

Forecast metrics at a single prediction horizon.

horizon¶

Forecast horizon in steps.

Type:: int

mae¶

Mean absolute error averaged over evaluation origins.

Type:: float

rmse¶

Root mean squared error averaged over evaluation origins.

Type:: float

mape¶

Mean absolute percentage error averaged over evaluation origins.

Type:: float

class koopman_graph.metrics.EvaluationResult(horizons, aggregate_mae, aggregate_rmse, aggregate_mape, num_origins)[source]¶

Bases: object

Multi-horizon forecast evaluation summary.

horizons¶

Per-horizon metrics in ascending horizon order.

Type:: tuple of HorizonMetrics

aggregate_mae¶

Mean of per-horizon MAE values.

Type:: float

aggregate_rmse¶

Mean of per-horizon RMSE values.

Type:: float

aggregate_mape¶

Mean of per-horizon MAPE values.

Type:: float

num_origins¶

Number of forecast origins averaged over.

Type:: int

koopman_graph.metrics.evaluate_forecast(model, sequence, *, horizons=(3, 6, 12), start_indices=None)[source]¶

Evaluate autoregressive multi-horizon forecasts on a snapshot sequence.

For each forecast origin, the model predicts up to max(horizons) steps ahead and metrics are averaged across origins at each requested horizon.

Parameters:

model (nn.Module) – Model implementing predict().
sequence (GraphSnapshotSequence) – Evaluation snapshots with shared topology.
horizons (sequence of int, optional) – Forecast horizons to report. Default is (3, 6, 12).
start_indices (sequence of int or None, optional) – Forecast-origin indices. When None, uses every valid origin in sequence.

Returns:

Per-horizon and aggregate MAE, RMSE, and MAPE.

Return type:

Raises:

ValueError – If horizons is empty, any horizon is invalid, or the sequence is too short.

Serialization¶

Checkpoint serialization for GraphKoopmanModel.

koopman_graph.serialization.build_model_config(model)[source]¶

Extract architecture configuration from a GraphKoopmanModel.

Parameters:: model (GraphKoopmanModel) – Model whose encoder, decoder, and Koopman settings will be serialized.
Returns:: JSON-serializable architecture configuration.
Return type:: dict

koopman_graph.serialization.reconstruct_model(config)[source]¶

Reconstruct a GraphKoopmanModel from a checkpoint configuration.

Parameters:: config (dict) – Architecture configuration produced by build_model_config().
Returns:: Uninitialized-weight model matching the saved architecture.
Return type:: GraphKoopmanModel

koopman_graph.serialization.build_checkpoint(model)[source]¶

Build a versioned checkpoint dictionary for a model.

Parameters:: model (GraphKoopmanModel) – Model whose weights and architecture will be serialized.
Returns:: Checkpoint payload suitable for torch.save().
Return type:: dict

koopman_graph.serialization.save_checkpoint(model, path)[source]¶

Persist a trained model checkpoint to disk.

Parameters:

model (GraphKoopmanModel) – Model to serialize.
path (str or Path) – Destination .pt file path.

koopman_graph.serialization.load_checkpoint(path, *, map_location=None)[source]¶

Load a trained model from a checkpoint file.

Parameters:

path (str or Path) – Checkpoint .pt file produced by save_checkpoint().
map_location (str, torch.device, or None, optional) – Device mapping forwarded to torch.load().

Returns:

Reconstructed model with restored weights in evaluation mode.

Return type:

GraphKoopmanModel

Raises:

ValueError – If the checkpoint format version is unsupported or the payload is invalid.
FileNotFoundError – If path does not exist.

koopman_graph.serialization.snapshot_state_dict(module)[source]¶

Return a detached copy of a module’s state_dict for checkpointing.

Parameters:: module (nn.Module) – Module whose parameters will be copied.
Returns:: Deep copy of nn.Module.state_dict() with detached tensors.
Return type:: dict

Datasets¶

Dataset utilities for KoopmanGraph benchmarks and examples.

Public benchmarks¶

SyntheticDynamicGraphBenchmark: Reproducible Laplacian-diffusion dynamics on path/ring topologies.
GridDynamicGraphBenchmark: Laplacian diffusion on a 4-connected 2D lattice graph.
AnisotropicAdvectionGridBenchmark: Directional advection on a grid with asymmetric neighbor weights.
IEEE118DynamicBenchmark: IEEE 118-bus MATPOWER topology with simulated voltage/load dynamics.
MetrLaTrafficBenchmark: METR-LA traffic-speed sensor graph with cached speed snapshots.

class koopman_graph.datasets.AnisotropicAdvectionGridBenchmark[source]¶

Bases: object

Directional advection on a 2D lattice with asymmetric neighbor weights.

Each node updates from a weighted mixture of neighbors where the west and north directions dominate. This breaks the symmetry that GCN layers assume when they aggregate neighbors uniformly, making graph attention a better fit than plain convolution for rollout forecasting.

The one-step update is

\[x_{t+1} = \text{decay\_rate} \cdot x_t + (1 - \text{decay\_rate}) \cdot \frac{\sum_{j \in \mathcal{N}(i)} w_{ij} x_{j,t}} {\sum_{j \in \mathcal{N}(i)} w_{ij}}\]

with w_{i,\text{west}} = west_weight, w_{i,\text{north}} = north_weight, and remaining neighbors sharing the leftover mass equally.

Parameters:

num_rows (int, optional) – Grid height. Default is 8.
num_cols (int, optional) – Grid width. Default is 8.
num_timesteps (int, optional) – Number of temporal snapshots. Default is 40.
in_channels (int, optional) – Node feature dimension. Default is 3.
decay_rate (float, optional) – Self-retention factor in (0, 1). Default is 0.85.
west_weight (float, optional) – Relative influence of the western neighbor. Default is 0.7.
north_weight (float, optional) – Relative influence of the northern neighbor. Default is 0.2.
noise_std (float, optional) – Standard deviation of additive Gaussian noise. Default is 0.005.
seed (int, optional) – Random seed for the initial state and noise.
initial_state ({"random", "ones"}, optional) – Initial node feature pattern. Default is "ones".
dtype (torch.dtype, optional) – Floating dtype for generated features. Default is torch.float32.

classmethod generate(*, num_rows=8, num_cols=8, num_timesteps=40, in_channels=3, decay_rate=0.85, west_weight=0.7, north_weight=0.2, noise_std=0.005, seed=42, initial_state='ones', dtype=torch.float32)[source]¶

Generate a directional advection snapshot sequence on a grid.

Parameters:

num_rows (int, optional) – Grid height. Default is 8.
num_cols (int, optional) – Grid width. Default is 8.
num_timesteps (int, optional) – Number of temporal snapshots. Default is 40.
in_channels (int, optional) – Node feature dimension. Default is 3.
decay_rate (float, optional) – Self-retention factor in (0, 1). Default is 0.85.
west_weight (float, optional) – Relative influence of the western neighbor. Default is 0.7.
north_weight (float, optional) – Relative influence of the northern neighbor. Default is 0.2.
noise_std (float, optional) – Standard deviation of additive Gaussian noise. Default is 0.005.
seed (int, optional) – Random seed for the initial state and noise.
initial_state ({"random", "ones"}, optional) – Initial node feature pattern. Default is "ones".
dtype (torch.dtype, optional) – Floating dtype for generated features. Default is torch.float32.

Returns:

Time-ordered snapshots on the grid graph.

Return type:

Raises:

ValueError – If any generation parameter is invalid.

class koopman_graph.datasets.GridDynamicGraphBenchmark[source]¶

Bases: object

Reproducible Laplacian diffusion on a 2D lattice graph.

Node features evolve via the same graph diffusion dynamics as SyntheticDynamicGraphBenchmark, but on a 4-connected grid. Corner, edge, and interior nodes have different degrees, which makes graph attention encoders a natural fit.

Parameters:

num_rows (int, optional) – Grid height. Default is 10.
num_cols (int, optional) – Grid width. Default is 10.
num_timesteps (int, optional) – Number of temporal snapshots. Default is 40.
in_channels (int, optional) – Node feature dimension. Default is 3.
diffusion_rate (float, optional) – Laplacian diffusion strength in [0, 1]. Default is 0.1.
decay_rate (float, optional) – Global amplitude decay applied each step. Default is 0.99.
noise_std (float, optional) – Standard deviation of additive Gaussian noise. Default is 0.01.
seed (int, optional) – Random seed for the initial state and noise.
initial_state ({"random", "ones"}, optional) – Initial node feature pattern. Default is "ones".
dtype (torch.dtype, optional) – Floating dtype for generated features. Default is torch.float32.

classmethod generate(*, num_rows=10, num_cols=10, num_timesteps=40, in_channels=3, diffusion_rate=0.1, decay_rate=0.99, noise_std=0.01, seed=42, initial_state='ones', dtype=torch.float32)[source]¶

Generate a dynamic 2D grid snapshot sequence.

Parameters:

num_rows (int, optional) – Grid height. Default is 10.
num_cols (int, optional) – Grid width. Default is 10.
num_timesteps (int, optional) – Number of temporal snapshots. Default is 40.
in_channels (int, optional) – Node feature dimension. Default is 3.
diffusion_rate (float, optional) – Laplacian diffusion strength in [0, 1]. Default is 0.1.
decay_rate (float, optional) – Global amplitude decay applied each step. Default is 0.99.
noise_std (float, optional) – Standard deviation of additive Gaussian noise. Default is 0.01.
seed (int, optional) – Random seed for the initial state and noise.
initial_state ({"random", "ones"}, optional) – Initial node feature pattern. Default is "ones".
dtype (torch.dtype, optional) – Floating dtype for generated features. Default is torch.float32.

Returns:

Time-ordered snapshots on the grid graph.

Return type:

Raises:

ValueError – If any generation parameter is invalid.

class koopman_graph.datasets.IEEE118DynamicBenchmark[source]¶

Bases: object

IEEE 118-bus spatiotemporal benchmark built from MATPOWER topology.

Node features are bus quantities [Vm, Va, Pd, Qd] (per-unit loads). Voltages and angles evolve via graph Laplacian diffusion on the real IEEE 118 transmission topology; loads follow a slow sinusoidal ramp to emulate changing grid conditions over time.

For large-scale optimal power flow snapshots, see the PowerGraph dataset (https://arxiv.org/abs/2402.02827).

NUM_BUSES¶

Number of buses in the IEEE 118 case.

Type:: int

IN_CHANNELS¶

Node feature dimension [Vm, Va, Pd, Qd].

Type:: int

classmethod load_topology(cache_dir=None, *, dtype=torch.float32)[source]¶

Load the cached IEEE 118 topology tables.

Parameters:

cache_dir (Path, optional) – Directory containing cached topology artifacts. Defaults to the package data/ieee118 directory.
dtype (torch.dtype, optional) – Floating dtype for node features. Default is torch.float32.

Returns:

Topology payload with keys base_mva, bus_ids, edge_index, initial_features, num_nodes, and source_url.

Return type:

classmethod generate(*, num_timesteps=40, diffusion_rate=0.35, decay_rate=0.98, noise_std=0.002, load_ramp_amplitude=0.15, load_ramp_period=20.0, expose_load_ramp_control=False, seed=42, cache_dir=None, dtype=torch.float32)[source]¶

Generate a dynamic IEEE 118-bus snapshot sequence.

Parameters:

num_timesteps (int, optional) – Number of temporal snapshots. Default is 40.
diffusion_rate (float, optional) – Laplacian diffusion strength for Vm and Va. Default is 0.35.
decay_rate (float, optional) – Global amplitude decay applied each step. Default is 0.98.
noise_std (float, optional) – Standard deviation of additive Gaussian noise. Default is 0.002.
load_ramp_amplitude (float, optional) – Peak fractional change applied to Pd and Qd. Default is 0.15.
load_ramp_period (float, optional) – Sinusoidal load ramp period in timesteps. Default is 20.0.
expose_load_ramp_control (bool, optional) – When True, attach the sinusoidal load-ramp multiplier as global control inputs with shape (num_timesteps, 1). Default is False.
seed (int, optional) – Random seed for noise. Default is 42.
cache_dir (Path, optional) – Directory containing cached topology artifacts.
dtype (torch.dtype, optional) – Floating dtype for generated features. Default is torch.float32.

Returns:

Time-ordered snapshots on the IEEE 118-bus graph.

Return type:

Raises:

ValueError – If any generation parameter is invalid.

class koopman_graph.datasets.MetrLaTrafficBenchmark[source]¶

Bases: object

METR-LA traffic-speed benchmark built from the DCRNN sensor graph.

Node features are per-sensor traffic speeds (mph), z-score normalized over the cached time window. The road-network adjacency follows the standard DCRNN Gaussian kernel on pairwise distances.

Full speed history is distributed with the DCRNN release (Google Drive / Baidu Yun). A public mirror is documented in scripts/download_metr_la.py.

NUM_SENSORS¶

Number of sensors in the METR-LA graph.

Type:: int

IN_CHANNELS¶

Node feature dimension (one normalized speed per sensor).

Type:: int

classmethod load_topology(cache_dir=None, *, dtype=torch.float32)[source]¶

Load cached METR-LA graph topology and metadata.

Parameters:

cache_dir (Path, optional) – Directory containing cached traffic artifacts. Defaults to the package data/metr_la directory.
dtype (torch.dtype, optional) – Floating dtype for returned tensors. Default is torch.float32.

Returns:

Metadata with keys sensor_ids, edge_index, edge_weight, num_nodes, source_h5_url, and normalized_k.

Return type:

classmethod load_sequence(cache_dir=None, *, dtype=torch.float32)[source]¶

Load the cached METR-LA speed snapshot sequence.

Parameters:

cache_dir (Path, optional) – Directory containing cached traffic artifacts. Defaults to the package data/metr_la directory.
dtype (torch.dtype, optional) – Floating dtype for speed features. Default is torch.float32.

Returns:

Time-ordered graph snapshots with one speed feature per sensor.

Return type:

class koopman_graph.datasets.SyntheticDynamicGraphBenchmark[source]¶

Bases: object

Reproducible synthetic graph dynamics for benchmarks and tutorials.

Node features evolve via graph Laplacian diffusion with optional global decay and additive Gaussian noise:

\[x_{t+1} = \text{decay\_rate} \cdot S x_t + \mathcal{N}(0, \sigma^2)\]

where S = (1 - diffusion_rate) * I + diffusion_rate * D^{-1/2} A D^{-1/2}.

Parameters:

num_nodes (int, optional) – Number of nodes in the graph. Default is 20.
num_timesteps (int, optional) – Number of temporal snapshots. Default is 50.
in_channels (int, optional) – Node feature dimension. Default is 3.
topology ({"path", "ring"}, optional) – Static graph topology shared across timesteps. Default is "path".
diffusion_rate (float, optional) – Laplacian diffusion strength in [0, 1]. Default is 0.5.
decay_rate (float, optional) – Global amplitude decay applied each step. Default is 0.95.
noise_std (float, optional) – Standard deviation of additive Gaussian noise. Default is 0.0.
seed (int, optional) – Random seed for the initial state and noise. None uses unseeded randomness.
initial_state ({"random", "ones"}, optional) – Initial node feature pattern. Default is "random".
dtype (torch.dtype, optional) – Floating dtype for generated features. Default is torch.float32.

classmethod generate(*, num_nodes=20, num_timesteps=50, in_channels=3, topology='path', diffusion_rate=0.5, decay_rate=0.95, noise_std=0.0, seed=None, initial_state='random', dtype=torch.float32)[source]¶

Generate a synthetic dynamic graph snapshot sequence.

Parameters:

num_nodes (int, optional) – Number of nodes in the graph. Default is 20.
num_timesteps (int, optional) – Number of temporal snapshots. Default is 50.
in_channels (int, optional) – Node feature dimension. Default is 3.
topology ({"path", "ring"}, optional) – Static graph topology shared across timesteps. Default is "path".
diffusion_rate (float, optional) – Laplacian diffusion strength in [0, 1]. Default is 0.5.
decay_rate (float, optional) – Global amplitude decay applied each step. Default is 0.95.
noise_std (float, optional) – Standard deviation of additive Gaussian noise. Default is 0.0.
seed (int, optional) – Random seed for the initial state and noise. None uses unseeded randomness.
initial_state ({"random", "ones"}, optional) – Initial node feature pattern. Default is "random".
dtype (torch.dtype, optional) – Floating dtype for generated features. Default is torch.float32.

Returns:

Time-ordered snapshots with shared topology and documented dynamics.

Return type:

Raises:

ValueError – If any generation parameter is invalid.

Synthetic spatiotemporal graph benchmarks for tests and tutorials.

class koopman_graph.datasets.synthetic.SyntheticDynamicGraphBenchmark[source]¶

Bases: object

Reproducible synthetic graph dynamics for benchmarks and tutorials.

Node features evolve via graph Laplacian diffusion with optional global decay and additive Gaussian noise:

\[x_{t+1} = \text{decay\_rate} \cdot S x_t + \mathcal{N}(0, \sigma^2)\]

where S = (1 - diffusion_rate) * I + diffusion_rate * D^{-1/2} A D^{-1/2}.

Parameters:

num_nodes (int, optional) – Number of nodes in the graph. Default is 20.
num_timesteps (int, optional) – Number of temporal snapshots. Default is 50.
in_channels (int, optional) – Node feature dimension. Default is 3.
topology ({"path", "ring"}, optional) – Static graph topology shared across timesteps. Default is "path".
diffusion_rate (float, optional) – Laplacian diffusion strength in [0, 1]. Default is 0.5.
decay_rate (float, optional) – Global amplitude decay applied each step. Default is 0.95.
noise_std (float, optional) – Standard deviation of additive Gaussian noise. Default is 0.0.
seed (int, optional) – Random seed for the initial state and noise. None uses unseeded randomness.
initial_state ({"random", "ones"}, optional) – Initial node feature pattern. Default is "random".
dtype (torch.dtype, optional) – Floating dtype for generated features. Default is torch.float32.

Generate a synthetic dynamic graph snapshot sequence.

Parameters:

num_nodes (int, optional) – Number of nodes in the graph. Default is 20.
num_timesteps (int, optional) – Number of temporal snapshots. Default is 50.
in_channels (int, optional) – Node feature dimension. Default is 3.
topology ({"path", "ring"}, optional) – Static graph topology shared across timesteps. Default is "path".
diffusion_rate (float, optional) – Laplacian diffusion strength in [0, 1]. Default is 0.5.
decay_rate (float, optional) – Global amplitude decay applied each step. Default is 0.95.
noise_std (float, optional) – Standard deviation of additive Gaussian noise. Default is 0.0.
seed (int, optional) – Random seed for the initial state and noise. None uses unseeded randomness.
initial_state ({"random", "ones"}, optional) – Initial node feature pattern. Default is "random".
dtype (torch.dtype, optional) – Floating dtype for generated features. Default is torch.float32.

Returns:

Time-ordered snapshots with shared topology and documented dynamics.

Return type:

Raises:

ValueError – If any generation parameter is invalid.

2D grid spatiotemporal graph benchmarks for tests and tutorials.

koopman_graph.datasets.grid.grid_node_index(row, col, *, num_cols)[source]¶

Return the flattened node index for a grid coordinate.

Parameters:

row (int) – Zero-based row index.
col (int) – Zero-based column index.
num_cols (int) – Number of columns in the grid.

Returns:

Flattened node index row * num_cols + col.

Return type:

int

class koopman_graph.datasets.grid.GridDynamicGraphBenchmark[source]¶

Bases: object

Reproducible Laplacian diffusion on a 2D lattice graph.

Parameters:

num_rows (int, optional) – Grid height. Default is 10.
num_cols (int, optional) – Grid width. Default is 10.
num_timesteps (int, optional) – Number of temporal snapshots. Default is 40.
in_channels (int, optional) – Node feature dimension. Default is 3.
diffusion_rate (float, optional) – Laplacian diffusion strength in [0, 1]. Default is 0.1.
decay_rate (float, optional) – Global amplitude decay applied each step. Default is 0.99.
noise_std (float, optional) – Standard deviation of additive Gaussian noise. Default is 0.01.
seed (int, optional) – Random seed for the initial state and noise.
initial_state ({"random", "ones"}, optional) – Initial node feature pattern. Default is "ones".
dtype (torch.dtype, optional) – Floating dtype for generated features. Default is torch.float32.

Generate a dynamic 2D grid snapshot sequence.

Parameters:

num_rows (int, optional) – Grid height. Default is 10.
num_cols (int, optional) – Grid width. Default is 10.
num_timesteps (int, optional) – Number of temporal snapshots. Default is 40.
in_channels (int, optional) – Node feature dimension. Default is 3.
diffusion_rate (float, optional) – Laplacian diffusion strength in [0, 1]. Default is 0.1.
decay_rate (float, optional) – Global amplitude decay applied each step. Default is 0.99.
noise_std (float, optional) – Standard deviation of additive Gaussian noise. Default is 0.01.
seed (int, optional) – Random seed for the initial state and noise.
initial_state ({"random", "ones"}, optional) – Initial node feature pattern. Default is "ones".
dtype (torch.dtype, optional) – Floating dtype for generated features. Default is torch.float32.

Returns:

Time-ordered snapshots on the grid graph.

Return type:

Raises:

ValueError – If any generation parameter is invalid.

class koopman_graph.datasets.grid.AnisotropicAdvectionGridBenchmark[source]¶

Bases: object

Directional advection on a 2D lattice with asymmetric neighbor weights.

The one-step update is

\[x_{t+1} = \text{decay\_rate} \cdot x_t + (1 - \text{decay\_rate}) \cdot \frac{\sum_{j \in \mathcal{N}(i)} w_{ij} x_{j,t}} {\sum_{j \in \mathcal{N}(i)} w_{ij}}\]

with w_{i,\text{west}} = west_weight, w_{i,\text{north}} = north_weight, and remaining neighbors sharing the leftover mass equally.

Parameters:

num_rows (int, optional) – Grid height. Default is 8.
num_cols (int, optional) – Grid width. Default is 8.
num_timesteps (int, optional) – Number of temporal snapshots. Default is 40.
in_channels (int, optional) – Node feature dimension. Default is 3.
decay_rate (float, optional) – Self-retention factor in (0, 1). Default is 0.85.
west_weight (float, optional) – Relative influence of the western neighbor. Default is 0.7.
north_weight (float, optional) – Relative influence of the northern neighbor. Default is 0.2.
noise_std (float, optional) – Standard deviation of additive Gaussian noise. Default is 0.005.
seed (int, optional) – Random seed for the initial state and noise.
initial_state ({"random", "ones"}, optional) – Initial node feature pattern. Default is "ones".
dtype (torch.dtype, optional) – Floating dtype for generated features. Default is torch.float32.

Generate a directional advection snapshot sequence on a grid.

Parameters:

num_rows (int, optional) – Grid height. Default is 8.
num_cols (int, optional) – Grid width. Default is 8.
num_timesteps (int, optional) – Number of temporal snapshots. Default is 40.
in_channels (int, optional) – Node feature dimension. Default is 3.
decay_rate (float, optional) – Self-retention factor in (0, 1). Default is 0.85.
west_weight (float, optional) – Relative influence of the western neighbor. Default is 0.7.
north_weight (float, optional) – Relative influence of the northern neighbor. Default is 0.2.
noise_std (float, optional) – Standard deviation of additive Gaussian noise. Default is 0.005.
seed (int, optional) – Random seed for the initial state and noise.
initial_state ({"random", "ones"}, optional) – Initial node feature pattern. Default is "ones".
dtype (torch.dtype, optional) – Floating dtype for generated features. Default is torch.float32.

Returns:

Time-ordered snapshots on the grid graph.

Return type:

Raises:

ValueError – If any generation parameter is invalid.

IEEE 118-bus MATPOWER benchmark for tutorials and tests.

koopman_graph.datasets.ieee118.parse_matpower_case(text)[source]¶

Parse bus and branch tables from a MATPOWER version-2 case file.

Parameters:: text (str) – Contents of a MATPOWER .m case file.
Returns:: Parsed fields including baseMVA, bus, and branch matrices.
Return type:: dict

koopman_graph.datasets.ieee118.topology_from_matpower_text(text, *, dtype=torch.float32)[source]¶

Convert MATPOWER case text into tensors used by the benchmark.

Parameters:

text (str) – Contents of a MATPOWER version-2 .m case file.
dtype (torch.dtype, optional) – Floating dtype for node features. Default is torch.float32.

Returns:

Topology payload with keys base_mva, bus_ids, edge_index, initial_features, num_nodes, and source_url.

Return type:

koopman_graph.datasets.ieee118.download_matpower_case118()[source]¶

Download the MATPOWER IEEE 118-bus case file text.

Returns:: Raw case118.m file contents.
Return type:: str

koopman_graph.datasets.ieee118.ensure_topology_cache(cache_dir=None, *, force=False, dtype=torch.float32)[source]¶

Download, parse, and cache IEEE 118 topology if needed.

Parameters:

cache_dir (Path, optional) – Directory used for cached topology artifacts. Defaults to data/ieee118 at the repository root.
force (bool, optional) – Rebuild the cache even when it already exists.
dtype (torch.dtype, optional) – Floating dtype stored in the cache. Default is torch.float32.

Returns:

Path to the cached topology.pt file.

Return type:

Path

koopman_graph.datasets.ieee118.load_topology(cache_dir=None, *, dtype=torch.float32)[source]¶

Load cached IEEE 118 topology, creating the cache on first use.

Parameters:

cache_dir (Path, optional) – Directory containing topology.pt.
dtype (torch.dtype, optional) – Floating dtype for returned tensors. Default is torch.float32.

Returns:

Topology payload with edge_index and initial_features.

Return type:

class koopman_graph.datasets.ieee118.IEEE118DynamicBenchmark[source]¶

Bases: object

IEEE 118-bus spatiotemporal benchmark built from MATPOWER topology.

For large-scale optimal power flow snapshots, see the PowerGraph dataset (https://arxiv.org/abs/2402.02827).

NUM_BUSES¶

Number of buses in the IEEE 118 case.

Type:: int

IN_CHANNELS¶

Node feature dimension [Vm, Va, Pd, Qd].

Type:: int

classmethod load_topology(cache_dir=None, *, dtype=torch.float32)[source]¶

Load the cached IEEE 118 topology tables.

Parameters:

cache_dir (Path, optional) – Directory containing cached topology artifacts. Defaults to the package data/ieee118 directory.
dtype (torch.dtype, optional) – Floating dtype for node features. Default is torch.float32.

Returns:

Topology payload with keys base_mva, bus_ids, edge_index, initial_features, num_nodes, and source_url.

Return type:

Generate a dynamic IEEE 118-bus snapshot sequence.

Parameters:

num_timesteps (int, optional) – Number of temporal snapshots. Default is 40.
diffusion_rate (float, optional) – Laplacian diffusion strength for Vm and Va. Default is 0.35.
decay_rate (float, optional) – Global amplitude decay applied each step. Default is 0.98.
noise_std (float, optional) – Standard deviation of additive Gaussian noise. Default is 0.002.
load_ramp_amplitude (float, optional) – Peak fractional change applied to Pd and Qd. Default is 0.15.
load_ramp_period (float, optional) – Sinusoidal load ramp period in timesteps. Default is 20.0.
expose_load_ramp_control (bool, optional) – When True, attach the sinusoidal load-ramp multiplier as global control inputs with shape (num_timesteps, 1). Default is False.
seed (int, optional) – Random seed for noise. Default is 42.
cache_dir (Path, optional) – Directory containing cached topology artifacts.
dtype (torch.dtype, optional) – Floating dtype for generated features. Default is torch.float32.

Returns:

Time-ordered snapshots on the IEEE 118-bus graph.

Return type:

Raises:

ValueError – If any generation parameter is invalid.

METR-LA traffic benchmark for tutorials and tests.

koopman_graph.datasets.metr_la.download_sensor_ids()[source]¶

Download the ordered METR-LA sensor ID list from the DCRNN repository.

Returns:: Sensor identifiers in graph node order.
Return type:: list of str

koopman_graph.datasets.metr_la.download_distances_csv()[source]¶

Download the METR-LA pairwise road-distance CSV from DCRNN.

Returns:: Raw CSV text with columns from, to, and cost.
Return type:: str

koopman_graph.datasets.metr_la.build_adjacency_matrix(distance_csv, sensor_ids, *, normalized_k=0.1)[source]¶

Build a Gaussian-kernel adjacency matrix from DCRNN distance CSV data.

Follows the preprocessing in the DCRNN gen_adj_mx script.

Parameters:

distance_csv (str) – CSV text with columns from, to, and cost.
sensor_ids (list of str) – Ordered sensor identifiers.
normalized_k (float, optional) – Entries below this threshold after kernel normalization are zeroed. Default is 0.1.

Returns:

Adjacency matrix with shape (len(sensor_ids), len(sensor_ids)).

Return type:

ndarray

koopman_graph.datasets.metr_la.adjacency_to_edge_index(adj_mx)[source]¶

Convert an adjacency matrix to a bidirectional PyG edge_index.

Parameters:: adj_mx (ndarray) – Weighted adjacency matrix.
Returns:: Long tensor with shape (2, num_edges).
Return type:: Tensor

koopman_graph.datasets.metr_la.adjacency_to_edge_weight(adj_mx)[source]¶

Extract scalar edge weights aligned with adjacency_to_edge_index().

Parameters:: adj_mx (ndarray) – Weighted adjacency matrix.
Returns:: Float tensor with shape (num_edges,).
Return type:: Tensor

koopman_graph.datasets.metr_la.read_h5_speed_window(h5_path, *, num_timesteps, offset=0)[source]¶

Read a window of METR-LA speed values from an HDF5 file.

Parameters:

h5_path (Path) – Path to metr-la.h5 in the DCRNN pandas HDF5 format.
num_timesteps (int) – Number of consecutive 5-minute readings to load.
offset (int, optional) – Starting row offset in the speed table. Default is 0.

Returns:

Speed array with shape (num_timesteps, num_sensors).

Return type:

ndarray

Raises:

ImportError – If h5py is not installed.
ValueError – If the requested window exceeds the available rows.

koopman_graph.datasets.metr_la.preprocess_speeds(speeds)[source]¶

Impute missing METR-LA readings before normalization.

The DCRNN HDF5 release encodes missing loop-detector samples as 0. Non-positive values are treated as missing and filled along time per sensor using forward fill followed by backward fill.

Parameters:: speeds (ndarray) – Raw speed array with shape (num_timesteps, num_sensors).
Returns:: Imputed speeds with the same shape.
Return type:: ndarray
Raises:: ValueError – If speeds is not two-dimensional.

koopman_graph.datasets.metr_la.normalize_speeds(speeds)[source]¶

Per-sensor z-score normalization along the time axis.

Parameters:: speeds (ndarray) – Speed array with shape (num_timesteps, num_sensors).
Returns:: Normalized speeds with the same shape.
Return type:: ndarray

koopman_graph.datasets.metr_la.build_traffic_cache_payload(speeds, sensor_ids, *, normalized_k=0.1, source_h5_url='https://huggingface.co/datasets/MintBruce/SkyTraffic/resolve/main/metr-la.h5', timestep_offset=0)[source]¶

Assemble a cache payload from speed readings and DCRNN graph metadata.

Parameters:

speeds (ndarray) – Raw speed readings with shape (num_timesteps, num_sensors).
sensor_ids (list of str) – Ordered sensor identifiers.
normalized_k (float, optional) – Adjacency sparsity threshold. Default is 0.1.
source_h5_url (str, optional) – Provenance URL for the speed source file.
timestep_offset (int, optional) – Row offset used when extracting speeds from the HDF5 table.

Returns:

Serializable cache payload for traffic.pt.

Return type:

koopman_graph.datasets.metr_la.ensure_traffic_cache(cache_dir=None, *, force=False, h5_path=None, num_timesteps=60, offset=0, normalized_k=0.1)[source]¶

Download graph metadata and build the METR-LA traffic cache if needed.

Parameters:

cache_dir (Path, optional) – Directory used for traffic.pt. Defaults to data/metr_la.
force (bool, optional) – Rebuild the cache even when it already exists.
h5_path (Path, optional) – Local metr-la.h5 file used to refresh speed readings. When omitted and force=True, an existing cache is reused without fetching HDF5.
num_timesteps (int, optional) – Number of timesteps stored in the cache. Default is 60.
offset (int, optional) – Starting row in the HDF5 speed table. Default is 0.
normalized_k (float, optional) – Adjacency sparsity threshold. Default is 0.1.

Returns:

Path to traffic.pt.

Return type:

Path

koopman_graph.datasets.metr_la.load_traffic_cache(cache_dir=None, *, dtype=torch.float32)[source]¶

Load cached METR-LA topology and speed readings.

Parameters:

cache_dir (Path, optional) – Directory containing traffic.pt.
dtype (torch.dtype, optional) – Floating dtype for returned speed tensors. Default is torch.float32.

Returns:

Cache payload with edge_index and speeds tensors.

Return type:

class koopman_graph.datasets.metr_la.MetrLaTrafficBenchmark[source]¶

Bases: object

METR-LA traffic-speed benchmark built from the DCRNN sensor graph.

Node features are per-sensor traffic speeds (mph), z-score normalized over the cached time window. The road-network adjacency follows the standard DCRNN Gaussian kernel on pairwise distances.

Full speed history is distributed with the DCRNN release (Google Drive / Baidu Yun). A public mirror is documented in scripts/download_metr_la.py.

NUM_SENSORS¶

Number of sensors in the METR-LA graph.

Type:: int

IN_CHANNELS¶

Node feature dimension (one normalized speed per sensor).

Type:: int

classmethod load_topology(cache_dir=None, *, dtype=torch.float32)[source]¶

Load cached METR-LA graph topology and metadata.

Parameters:

cache_dir (Path, optional) – Directory containing cached traffic artifacts. Defaults to the package data/metr_la directory.
dtype (torch.dtype, optional) – Floating dtype for returned tensors. Default is torch.float32.

Returns:

Metadata with keys sensor_ids, edge_index, edge_weight, num_nodes, source_h5_url, and normalized_k.

Return type:

classmethod load_sequence(cache_dir=None, *, dtype=torch.float32)[source]¶

Load the cached METR-LA speed snapshot sequence.

Parameters:

cache_dir (Path, optional) – Directory containing cached traffic artifacts. Defaults to the package data/metr_la directory.
dtype (torch.dtype, optional) – Floating dtype for speed features. Default is torch.float32.

Returns:

Time-ordered graph snapshots with one speed feature per sensor.

Return type: