Build you STGNN

Spatiotemporal graph neural networks (STGNNs) forecast signals observed on the nodes of a graph. They combine an encoder, a stack of spatiotemporal message-passing (STMP) layers, and a decoder.

The composable classes in this notebook follow the template architecture presented by Cini et al., Taming Local Effects in Graph-based Spatiotemporal Forecasting (2023).

The encoder–STMP–decoder template

Let \(x_\tau^i\) denote the observation at node \(i\in\{1,\ldots,N\}\) and time \(\tau\). At forecast origin \(t\), the input window is \(\tau=t-T+1,\ldots,t\) and the target window is \(t+1,\ldots,t+H\). Let \(u_\tau^i\) be observed covariates, \(v^i\) static attributes, \(c_\tau^i\) known future covariates, and \(\omega^i\) a learnable node embedding.

The encoder produces the initial hidden sequence:

\[\boldsymbol{h}_\tau^{i,(0)} = \mathsf{ENCODER}\left(\boldsymbol{x}_\tau^i, \boldsymbol{u}_\tau^i, \boldsymbol{v}^i, \boldsymbol{\omega}^i\right),\qquad \tau=t-T+1,\ldots,t.\]

Writing \(\mathbf{H}^{(\ell)}_{t-T:t}\) for the hidden states of all nodes in the input window and \(\mathcal{E}\) for the set of edges encoding pairwise relationships, a stack of \(L\) STMP layers is

\[\mathbf{H}^{(\ell)}_{t-T:t}=\mathsf{STMP}^{(\ell)}\left(\mathbf{H}^{(\ell-1)}_{t-T:t}\,,\ \mathcal{E}\right),\qquad \ell=1,\ldots,L.\]

Finally, the decoder maps the final hidden sequence, future covariates, and embeddings to the forecast:

\[\hat{\boldsymbol{x}}^{i}_{t:t+H}=\mathsf{DEC}\left(\boldsymbol{h}^{i,(L)}_{t},\boldsymbol{u}^{i}_{t:t+H},\boldsymbol{\omega}^{i}\right).\]

This indexing keeps the observed input window and the strictly future forecast window separate. A causal STMP layer may, internally, restrict the state used at time \(\tau\) to earlier states \(\mathbf{H}_{t-T:\tau}^{(\ell)}\).

The matching base classes in tsl

STGNN mirrors the three-stage template. A subclass constructs its encoder, each STMP layer, and its decoder; the inherited forward always executes the same workflow. DisjointSTGNN splits every STMP layer into temporal (TMP) and spatial (SMP) stacks, while TimeThenSpace fixes the outer stack to one such temporal-then-spatial block.

The argument-mapping hooks decouple that workflow from the native signature of a chosen module. Each returns (args, kwargs): it can select a subset of inputs, rename them as keyword arguments, or preprocess them. Consequently, subclasses do not implement forward.

class STGNN(BaseModel):
    def build_encoder(self) -> nn.Module: ...
    def build_stmp_layer(self, layer: int) -> nn.Module: ...
    def build_decoder(self) -> nn.Module: ...

    def map_encoder_args(self, x, u=None, v=None, emb=None, **kwargs): ...
    def map_stmp_args(self, h, edge_index, edge_weight=None, *, layer, **kwargs): ...
    def map_decoder_args(self, h, u_h=None, emb=None, **kwargs): ...

class DisjointSTGNN(STGNN):
    def build_tmp_layer(self, layer: int, inner_layer: int) -> nn.Module: ...
    def build_smp_layer(self, layer: int, inner_layer: int) -> nn.Module: ...
    def map_tmp_args(self, h, *, layer, inner_layer, **kwargs): ...
    def map_smp_args(self, h, edge_index, edge_weight=None, *, layer, inner_layer, **kwargs): ...

Hint

Use the summary() method to print the architecture tree.

Now we import all the modules that we use to build custom STGNNs following this template.

from tsl.nn.blocks.decoders import MLPDecoder
from tsl.nn.blocks.encoders import DCRNN, RNN, ConditionalEncoder
from tsl.nn.layers.graph_convs import DiffusionConv
from tsl.nn.models.stgn import STGNN, TimeThenSpace

A time-and-space STGNN: DCRNN

The diffusion convolutional recurrent neural network (DCRNN) of Li et al. (2018) is a time-and-space (T&S) model: temporal and spatial processing are coupled inside the same recurrent operation.

ConditionalEncoder is a non-concatenating endpoint: it independently projects \(x\), \(u\), \(v\), and \(\omega\), broadcasts the projected terms to \([B,T,N,F]\), and sums them. MLPDecoder only accepts hidden states, so map_decoder_args selects \(h\); embeddings are used only by the encoder.

class DCRNNModel(STGNN):
    def build_encoder(self):
        return ConditionalEncoder(
            input_size=self.input_size,
            output_size=self.hidden_size,
            exog_size=self.exog_size,
            static_size=self.static_size,
            emb_size=self.emb_size,
        )

    def build_stmp_layer(self, layer):
        return DCRNN(
            input_size=self.hidden_size,
            hidden_size=self.hidden_size,
            n_layers=self.dcrnn_layers,
            k=self.diffusion_order,
            return_only_last_state=True,
        )

    def build_decoder(self):
        return MLPDecoder(
            self.hidden_size,
            self.decoder_size,
            self.output_size,
            horizon=self.horizon,
            n_layers=2,
        )

    def map_decoder_args(self, h, u_h=None, emb=None, **kwargs):
        return (h,), {}


dcrnn_model = DCRNNModel(
    input_size=1,
    output_size=1,
    horizon=12,
    n_nodes=207,
    emb_size=8,
    n_layers=1,
    hidden_size=64,
    exog_size=2,
    static_size=3,
    decoder_size=64,
    dcrnn_layers=2,
    diffusion_order=2,
)
print(dcrnn_model.summary())
# print(dcrnn_model.summary(show_workflow=True))

DCRNNModel architecture
  ├─ encoder: ConditionalEncoder
  ├─ STMP stack (1 layer)
  │   [0] DCRNN
  └─ decoder: MLPDecoder

A time-then-space STGNN

A time-then-space (TTS) layer separates the two operations. For outer layer \(\ell\), temporal processing first evolves each node independently, then spatial processing mixes the resulting hidden states over the graph:

\[\mathbf{Z}^{(\ell)}=\operatorname{TMP}^{(\ell)}_{L_T}\circ\cdots\circ\operatorname{TMP}^{(\ell)}_1\left(\mathbf{H}^{(\ell-1)}\right),\qquad\mathbf{H}^{(\ell)}=\operatorname{SMP}^{(\ell)}_{L_S}\circ\cdots\circ\operatorname{SMP}^{(\ell)}_1\left(\mathbf{Z}^{(\ell)},\mathcal{E}\right).\]

Here RNN(..., cell="gru") processes every node trajectory independently: its recurrent state is not graph-aware. DiffusionConv then applies the same static edge_index and optional edge_weight to every time step. Unlike DCRNN, graph propagation occurs only in the SMP stage, after node-wise temporal processing.

class GRUDiffusionTTS(TimeThenSpace):
    def build_encoder(self):
        return ConditionalEncoder(
            input_size=self.input_size,
            output_size=self.hidden_size,
            exog_size=self.exog_size,
            static_size=self.static_size,
            emb_size=self.emb_size,
        )

    def build_tmp_layer(self, layer, inner_layer):
        return RNN(self.hidden_size, self.hidden_size, cell="gru")

    def build_smp_layer(self, layer, inner_layer):
        return DiffusionConv(self.hidden_size, self.hidden_size, k=self.diffusion_order)

    def build_decoder(self):
        return MLPDecoder(
            self.hidden_size,
            self.decoder_size,
            self.output_size,
            horizon=self.horizon,
            n_layers=2,
        )

    def map_decoder_args(self, h, u_h=None, emb=None, **kwargs):
        return (h,), {}


tts_model = GRUDiffusionTTS(
    input_size=1,
    output_size=1,
    horizon=12,
    n_nodes=207,
    emb_size=8,
    n_temporal_layers=1,
    n_spatial_layers=1,
    hidden_size=64,
    exog_size=2,
    static_size=3,
    decoder_size=64,
    diffusion_order=2,
)
print(tts_model.summary())

GRUDiffusionTTS architecture
  ├─ encoder: ConditionalEncoder
  ├─ STMP stack (1 layer)
  │   [0]
  │   ├─ temporal: RNN
  │   ├─ spatial: DiffusionConv
  └─ decoder: MLPDecoder