Spatiotemporal Forecasting with Graph-Enformer (GEnformer)#
This notebook demonstrates the usage of the GEnformer for spatiotemporal data (data with both temporal sequences and spatial topology).
[1]:
import torch
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from darts import TimeSeries
from genformer.models import GEnformer
plt.style.use('seaborn-v0_8-darkgrid')
/home/docs/checkouts/readthedocs.org/user_builds/genformer/envs/latest/lib/python3.12/site-packages/gluonts/json.py:102: UserWarning: Using `json`-module for json-handling. Consider installing one of `orjson`, `ujson` to speed up serialization and deserialization.
warnings.warn(
[2]:
# Dummy spatiotemporal data generation: 4 spatial nodes with distinct periodic patterns
time_steps = 150
num_nodes = 4
x = np.linspace(0, 40, time_steps)
data = np.zeros((time_steps, num_nodes), dtype=np.float32)
# Node 0: Sine wave
data[:, 0] = np.sin(x)
# Node 1: Cosine wave (dependent on Node 0)
data[:, 1] = np.cos(x) + 0.3 * data[:, 0]
# Node 2: Faster sine wave (dependent on Node 1)
data[:, 2] = np.sin(1.5 * x) - 0.2 * data[:, 1]
# Node 3: Modulated wave (dependent on Node 0 and Node 2)
data[:, 3] = np.sin(x) * np.cos(2 * x) + 0.15 * data[:, 0] + 0.15 * data[:, 2]
# Add some noise
data += np.random.normal(0, 0.1, (time_steps, num_nodes)).astype(np.float32)
df = pd.DataFrame(data, columns=['Node_0', 'Node_1', 'Node_2', 'Node_3'])
series = TimeSeries.from_dataframe(df)
train, val = series[:-10], series[-10:]
fig, axes = plt.subplots(2, 2, figsize=(12, 8), sharex=True)
axes = axes.flatten()
for i in range(num_nodes):
train[f'Node_{i}'].plot(ax=axes[i], label=f'Node {i} (Train)')
axes[i].legend(loc='upper right')
axes[i].set_ylabel(f'Node {i}')
fig.suptitle('Spatiotemporal Dummy Data (4 Nodes)', fontsize=14)
plt.tight_layout()
plt.show()
[3]:
# Create dummy adjacency matrix (edges) for 4 nodes
edges = torch.tensor([
[0, 1, 1, 1],
[1, 0, 1, 0],
[1, 1, 0, 1],
[1, 0, 1, 0]
], dtype=torch.float32)
model = GEnformer(
input_chunk_length=20,
output_chunk_length=10,
edges=edges,
num_nodes=num_nodes,
num_samples_engression=5,
n_epochs=10, # Demo
batch_size=8,
d_model=64,
nhead=4,
num_encoder_layers=2,
num_decoder_layers=2,
dim_feedforward=128,
dropout=0.1
)
model.fit(train)
/home/docs/checkouts/readthedocs.org/user_builds/genformer/envs/latest/lib/python3.12/site-packages/torch/nn/modules/transformer.py:143: UserWarning: enable_nested_tensor is True, but self.use_nested_tensor is False because encoder_layer.self_attn.batch_first was not True(use batch_first for better inference performance)
self.encoder = TransformerEncoder(
/home/docs/checkouts/readthedocs.org/user_builds/genformer/envs/latest/lib/python3.12/site-packages/genformer/utils.py:81: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.detach().clone() or sourceTensor.detach().clone().requires_grad_(True), rather than torch.tensor(sourceTensor).
self.register_buffer("src_nodes", torch.tensor(edges[0], dtype=torch.long))
/home/docs/checkouts/readthedocs.org/user_builds/genformer/envs/latest/lib/python3.12/site-packages/genformer/utils.py:82: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.detach().clone() or sourceTensor.detach().clone().requires_grad_(True), rather than torch.tensor(sourceTensor).
self.register_buffer("dst_nodes", torch.tensor(edges[1], dtype=torch.long))
GPU available: False, used: False
TPU available: False, using: 0 TPU cores
💡 Tip: For seamless cloud logging and experiment tracking, try installing [litlogger](https://pypi.org/project/litlogger/) to enable LitLogger, which logs metrics and artifacts automatically to the Lightning Experiments platform.
| Name | Type | Params | Mode | FLOPs
----------------------------------------------------------------------------
0 | criterion | MSELoss | 0 | train | 0
1 | train_criterion | MSELoss | 0 | train | 0
2 | val_criterion | MSELoss | 0 | train | 0
3 | train_metrics | MetricCollection | 0 | train | 0
4 | val_metrics | MetricCollection | 0 | train | 0
5 | encoder | Sequential | 8.3 K | train | 0
6 | positional_encoding | _PositionalEncoding | 0 | train | 0
7 | transformer | Transformer | 167 K | train | 0
8 | decoder | Linear | 83.2 K | train | 0
9 | gcn | GraphConv | 128 | train | 0
----------------------------------------------------------------------------
259 K Trainable params
0 Non-trainable params
259 K Total params
1.037 Total estimated model params size (MB)
68 Modules in train mode
0 Modules in eval mode
0 Total Flops
/home/docs/checkouts/readthedocs.org/user_builds/genformer/envs/latest/lib/python3.12/site-packages/pytorch_lightning/utilities/_pytree.py:21: `isinstance(treespec, LeafSpec)` is deprecated, use `isinstance(treespec, TreeSpec) and treespec.is_leaf()` instead.
/home/docs/checkouts/readthedocs.org/user_builds/genformer/envs/latest/lib/python3.12/site-packages/torch/utils/data/dataloader.py:752: UserWarning: 'pin_memory' argument is set as true but no accelerator is found, then device pinned memory won't be used.
super().__init__(loader)
Epoch 9: 100%|██████████| 14/14 [00:00<00:00, 26.30it/s, train_loss_step=8.700, train_coverage_step=0.573, train_loss_epoch=7.800, train_coverage_epoch=0.625]
`Trainer.fit` stopped: `max_epochs=10` reached.
Epoch 9: 100%|██████████| 14/14 [00:00<00:00, 26.17it/s, train_loss_step=8.700, train_coverage_step=0.573, train_loss_epoch=7.800, train_coverage_epoch=0.625]
[3]:
GEnformer(output_chunk_shift=0, d_model=64, nhead=4, num_encoder_layers=2, num_decoder_layers=2, dim_feedforward=128, dropout=0.1, activation=relu, norm_type=None, custom_encoder=None, custom_decoder=None, lambda_calib=2, input_chunk_length=20, output_chunk_length=10)
[4]:
from genformer.utils import generate_forecasts
# Generate forecasts using the custom spatial method
# Output shape will be (M, T_out, N, D_gcn) where D_gcn is the latent spatial dimension
# The model expects exactly input_chunk_length as the history window
predictions_tensor = generate_forecasts(
model=model,
history=train[-20:], # 20 is the input_chunk_length
m_samples=30,
device="cuda" if torch.cuda.is_available() else "cpu"
)
# We average over the latent spatial dimension to get the 3 node forecasts
predictions_nodes = predictions_tensor.mean(dim=-1).cpu().numpy() # (M, T_out, N)