GEnformer#

Graph-Enformer (GEnformer) extends Enformer to the spatiotemporal setting, where each series is a node in an interconnected graph (sensors, regions, stations, …). It jointly models temporal dynamics, spatial interactions, and predictive uncertainty.

How it works#

Spatial encoding

A Graph Convolution (GCN) maps each node’s features over the graph structure to a spatially-aware latent embedding of size gcn_out_feat, before noise injection and the Transformer backbone.

Calibrated uncertainty

Alongside the Energy Score, an optional calibration loss nudges the predictive intervals toward a target coverage level (target_coverage), weighted by lambda_calib.

The noise injection, ensemble expansion, and Energy Score objective are inherited from the Enformer recipe.

Important

GEnformer operates and forecasts in the latent GCN space. Predictions are produced with genformer.utils.generate_forecasts() (not Darts’ model.predict) and have shape (M, T_out, N, gcn_out_feat). Average over the last (latent) dimension to recover one forecast per node.

Usage#

import torch
import numpy as np
import pandas as pd
from darts import TimeSeries
from genformer import GEnformer
from genformer.utils import generate_forecasts

# 1. Multivariate series: one column per spatial node
num_nodes = 3
df = pd.DataFrame(np.random.randn(150, num_nodes),
                  columns=[f"Node_{i}" for i in range(num_nodes)])
series = TimeSeries.from_dataframe(df)
train = series[:-30]

# 2. Adjacency matrix describing the spatial graph (num_nodes x num_nodes)
edges = torch.tensor([[0, 1, 1],
                      [1, 0, 1],
                      [1, 1, 0]], dtype=torch.float32)

# 3. Configure and train
model = GEnformer(
    input_chunk_length=20,
    output_chunk_length=10,
    edges=edges,
    num_nodes=num_nodes,
    gcn_out_feat=32,
    num_samples_engression=5,
    target_coverage=0.9,        # drives the calibration loss
    lambda_calib=2.0,
    n_epochs=30,
)
model.fit(train)

# 4. Forecast in latent space, then average over the latent dimension
samples = generate_forecasts(
    model=model,
    history=train[-20:],        # exactly input_chunk_length steps
    m_samples=30,
    device="cuda" if torch.cuda.is_available() else "cpu",
)                               # (M, T_out, N, gcn_out_feat)

node_forecasts = samples.mean(dim=-1).cpu().numpy()   # (M, T_out, N)

For a full, runnable walkthrough see Usage.

Key hyperparameters#

Argument

Symbol

Meaning

edges

Adjacency matrix (num_nodes × num_nodes) of the spatial graph.

num_nodes

\(N\)

Number of spatial locations.

gcn_out_feat

Dimensionality of the latent spatial embedding from the GCN.

num_samples_engression

\(M\)

Ensemble size used to estimate the Energy Score in training.

target_coverage

Target prediction-interval coverage for the calibration loss.

lambda_calib

Weight of the calibration term in the overall objective.

API reference#

class genformer.models.GEnformer(*args, **kwargs)[source]#

Bases: TransformerModel

Graph-Enformer (GEnformer): Spatiotemporal Probabilistic Forecasting.

An extension of the Enformer architecture for spatiotemporal contexts where spatial locations are represented as nodes in an interconnected graph. This model jointly captures temporal dynamics, complex spatial interactions, and predictive uncertainty.

Before noise injection and Transformer processing, GEnformer applies a Graph Convolutional Network (GCN) layer to map the target observations and the spatial topology (defined by an adjacency matrix) to spatially-aware latent embeddings. It then optimizes the Energy Score alongside an optional calibration objective.

Parameters:

  • input_chunk_length (int) – The length of the historical look-back window (( p )).

  • output_chunk_length (int) – The prediction horizon (( q )).

  • edges (torch.Tensor or List) – Graph edges defining the connectivity between nodes.

  • num_nodes (int) – The number of spatial locations (( D )).

  • output_chunk_shift (int, optional) – Shift for the output chunk. Defaults to 0.

  • gcn_out_feat (int, optional) – The dimensionality of the latent spatial embeddings outputted by the GCN. Defaults to 32.

  • graph_conv_params (dict, optional) – Dictionary of parameters for the GraphConv layer (e.g., aggregation and combination types). Defaults to None.

  • num_samples_engression (int, optional) – The number of in-sample forecast trajectories (( M )) for the ensemble. Defaults to 10.

  • noise_dist (str, optional) – The type of noise to inject (‘gaussian’ or ‘uniform’). Defaults to ‘gaussian’.

  • target_coverage (float, optional) – The target prediction interval coverage used for the calibration loss term. Defaults to 0.9.

  • noise_std (float, optional) – The standard deviation/scale (( sigma )) of the injected noise. Defaults to 1.

  • random_state (int, optional) – Seed for reproducibility. Defaults to 23.

  • save_checkpoints (bool, optional) – Whether to save PyTorch Lightning checkpoints. Defaults to False.

  • lambda_calib (float, optional) – The weight of the calibration loss term in the overall optimization objective. Defaults to 2.

  • **kwargs – Additional parameters passed to the underlying TransformerModel.

  • d_model (int)

  • nhead (int)

  • num_encoder_layers (int)

  • num_decoder_layers (int)

  • dim_feedforward (int)

  • dropout (float)

  • activation (str)

  • norm_type (str | torch.nn.Module | None)

  • custom_encoder (torch.nn.Module | None)

  • custom_decoder (torch.nn.Module | None)

genformer.utils.generate_forecasts(model, history, m_samples, past_covs=None, fut_covs=None, static_covs=None, unstandardize=None, device='cpu')[source]#