### Install Index-Batching Support

Source: https://github.com/benedekrozemberczki/pytorch_geometric_temporal/blob/master/docs/source/notes/installation.md

Install additional dependencies required for memory-efficient index-batching.

```bash
$ pip install torch-geometric-temporal[index]
```

--------------------------------

### Install with Optional Features

Source: https://github.com/benedekrozemberczki/pytorch_geometric_temporal/blob/master/README.md

Installation commands for specific library features like index-batching and DDP support.

```sh
pip install torch-geometric-temporal[index]
```

```sh
pip install torch-geometric-temporal[ddp]
```

--------------------------------

### Install PyTorch Geometric Temporal

Source: https://github.com/benedekrozemberczki/pytorch_geometric_temporal/blob/master/README.md

Standard installation command for the library.

```sh
pip install torch-geometric-temporal
```

--------------------------------

### Import Libraries and Setup

Source: https://github.com/benedekrozemberczki/pytorch_geometric_temporal/blob/master/notebooks/a3tgcn_for_traffic_forecasting.ipynb

Imports necessary libraries for data manipulation, visualization, and PyTorch Geometric Temporal. Sets up the device for GPU acceleration if available.

```python
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
import networkx as nx

import torch
import torch.nn.functional as F
from torch_geometric.nn import GCNConv
from torch_geometric_temporal.nn.recurrent import A3TGCN2
from torch_geometric_temporal.signal import temporal_signal_split

# GPU support
DEVICE = torch.device('cuda') # cuda
shuffle=True
batch_size = 32
```

--------------------------------

### Dataset Loading

Source: https://github.com/benedekrozemberczki/pytorch_geometric_temporal/blob/master/docs/source/notes/introduction.md

Example of how to load a benchmark dataset using the provided loaders.

```APIDOC
## GET /dataset/chickenpox

### Description
Loads the Hungarian Chickenpox Dataset using the ChickenpoxDatasetLoader.

### Request Example
```python
from torch_geometric_temporal.dataset import ChickenpoxDatasetLoader
loader = ChickenpoxDatasetLoader()
dataset = loader.get_dataset()
```

### Response
- **dataset** (StaticGraphTemporalSignal) - The returned dataset object.
```

--------------------------------

### Install PyTorch Geometric Temporal

Source: https://github.com/benedekrozemberczki/pytorch_geometric_temporal/blob/master/docs/source/notes/installation.md

Standard installation command for the library via pip.

```bash
$ pip install torch-geometric-temporal
```

--------------------------------

### Install Distributed Data Parallel Support

Source: https://github.com/benedekrozemberczki/pytorch_geometric_temporal/blob/master/docs/source/notes/installation.md

Install additional dependencies required for distributed data parallel training.

```bash
$ pip install torch-geometric-temporal[ddp]
```

--------------------------------

### Check Package Version

Source: https://github.com/benedekrozemberczki/pytorch_geometric_temporal/blob/master/docs/source/notes/installation.md

Verify the currently installed version of the library.

```bash
$ pip freeze | grep torch-geometric-temporal
```

--------------------------------

### Install PyTorch and Dependencies

Source: https://github.com/benedekrozemberczki/pytorch_geometric_temporal/blob/master/notebooks/a3tgcn_for_traffic_forecasting.ipynb

Installs PyTorch, specific PyTorch Geometric temporal dependencies, and the PyTorch Geometric Temporal library from GitHub. Ensure CUDA compatibility for GPU acceleration.

```python
!python -c "import torch; print(torch.__version__)"
```

```python
!python -c "import torch; print(torch.version.cuda)"
```

```python
!pip -q install torch-spline-conv==latest+cu102 torch-scatter==latest+cu102 torch-cluster==latest+cu102 torch-sparse==latest+cu102  -f https://pytorch-geometric.com/whl/torch-1.6.0.html
```

```python
!pip install -q git+https://github.com/elmahyai/pytorch_geometric_temporal   # remove after merging
```

--------------------------------

### Instantiate Data Loaders

Source: https://github.com/benedekrozemberczki/pytorch_geometric_temporal/blob/master/notebooks/astgcn_for_traffic_flow_forecasting.ipynb

Example of how to call the `load_graphdata_channel1` function with specific parameters to obtain data loaders for a traffic dataset. Adjust `graph_signal_matrix_filename`, `num_of_hours`, `num_of_days`, `num_of_weeks`, and `batch_size` as needed.

```python
graph_signal_matrix_filename = '../input/pems-dataset/data/PEMS04/PEMS04.npz'
batch_size = 32
num_of_weeks = 0
num_of_days = 0
num_of_hours = 1

train_loader, train_target_tensor, val_loader, val_target_tensor, test_loader, test_target_tensor, _mean, _std = load_graphdata_channel1(
    graph_signal_matrix_filename, num_of_hours, num_of_days, num_of_weeks, batch_size)
```

--------------------------------

### Implement Recurrent Graph Convolutional Network in Python

Source: https://github.com/benedekrozemberczki/pytorch_geometric_temporal/blob/master/README.md

This example demonstrates how to implement a recurrent graph convolutional network using two GConvGRU cells and a linear layer. Ensure PyTorch and PyTorch Geometric Temporal are installed.

```python
import torch
import torch.nn.functional as F
from torch_geometric_temporal.nn.recurrent import GConvGRU

class RecurrentGCN(torch.nn.Module):

    def __init__(self, node_features, num_classes):
        super(RecurrentGCN, self).__init__()
        self.recurrent_1 = GConvGRU(node_features, 32, 5)
        self.recurrent_2 = GConvGRU(32, 16, 5)
        self.linear = torch.nn.Linear(16, num_classes)

    def forward(self, x, edge_index, edge_weight):
        x = self.recurrent_1(x, edge_index, edge_weight)
        x = F.relu(x)
        x = F.dropout(x, training=self.training)
        x = self.recurrent_2(x, edge_index, edge_weight)
        x = F.relu(x)
        x = F.dropout(x, training=self.training)
        x = self.linear(x)
        return F.log_softmax(x, dim=1)

```

--------------------------------

### Train ST-GNNs with Index-Batching and DCRNN

Source: https://github.com/benedekrozemberczki/pytorch_geometric_temporal/blob/master/docs/source/notes/introduction.md

Example training loop demonstrating the use of index-batching with the DCRNN model on the PeMS-Bay dataset. Initialize the DatasetLoader with index=True and use loader.get_index_dataset().

```python
model = BatchedDCRNN(2, 2, K=3)
model = DDP(model)

loader = PemsBayDatasetLoader(index=True)
train_dataloader, _, _, edges, edge_weights, means, stds = loader.get_index_dataset(world_size=world_size, ddp_rank=worker_rank, batch_size=batch_size)
for batch in train_dataloader:
    X_batch, y_batch = batch

    # Forward pass
    outputs = model(X_batch, edges, edge_weights)

    # Calculate loss
    loss = masked_mae_loss((outputs * std) + mean, (y_batch * std) + mean)

    # Backward pass
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
```

--------------------------------

### Import Core PyTorch Geometric Temporal and Utility Libraries

Source: https://github.com/benedekrozemberczki/pytorch_geometric_temporal/blob/master/notebooks/astgcn_for_traffic_flow_forecasting.ipynb

Imports essential libraries for deep learning, graph manipulation, and temporal data processing, including PyTorch, PyTorch Geometric, and various utility modules. Ensure these libraries are installed before running.

```python
import os
from time import time
import numpy as np
import matplotlib.pyplot as plt
import networkx as nx



import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
```

```python
import math
from typing import Optional, List, Union

import torch
import torch.nn as nn
from torch.nn import Parameter
import torch.nn.functional as F

from torch_geometric.data import Data
from torch_geometric.typing import OptTensor
from torch_geometric.nn.conv import MessagePassing
from torch_geometric.transforms import LaplacianLambdaMax
from torch_geometric.utils import remove_self_loops, add_self_loops, get_laplacian
from torch_geometric.utils import to_dense_adj
```

--------------------------------

### Upgrade PyTorch Geometric Temporal

Source: https://github.com/benedekrozemberczki/pytorch_geometric_temporal/blob/master/docs/source/notes/installation.md

Command to update an existing installation to the latest version.

```bash
$ pip install torch-geometric-temporal --upgrade
```

--------------------------------

### search_data

Source: https://github.com/benedekrozemberczki/pytorch_geometric_temporal/blob/master/notebooks/processing_traffic_data_for_deep_learning_projects.ipynb

Calculates the start and end indices for historical data dependencies based on temporal units.

```APIDOC
## search_data

### Description
Calculates a list of (start_idx, end_idx) tuples representing historical data segments based on specified dependencies.

### Parameters
- **sequence_length** (int) - Total length of history data.
- **num_of_depend** (int) - Number of dependencies.
- **label_start_idx** (int) - First index of the predicting target.
- **num_for_predict** (int) - Number of points to predict per sample.
- **units** (int) - Temporal unit (e.g., 168 for week, 24 for day, 1 for hour).
- **points_per_hour** (int) - Number of data points per hour.

### Response
- **list[(start_idx, end_idx)]** - List of index ranges or None if invalid.
```

--------------------------------

### Evaluate Model and Calculate Test MSE

Source: https://github.com/benedekrozemberczki/pytorch_geometric_temporal/blob/master/notebooks/a3tgcn_for_traffic_forecasting.ipynb

Iterates through the test data loader, gets model predictions, calculates Mean Squared Error (MSE) loss, and prints the average test MSE. Ensure the model is in evaluation mode.

```python
model.eval()
step = 0
# Store for analysis
total_loss = []
for encoder_inputs, labels in test_loader:
    # Get model predictions
    y_hat = model(encoder_inputs, static_edge_index)
    # Mean squared error
    loss = loss_fn(y_hat, labels)
    total_loss.append(loss.item())
    # Store for analysis below
    #test_labels.append(labels)
    #predictions.append(y_hat)
    

print("Test MSE: {:.4f}".format(sum(total_loss)/len(total_loss)))
```

--------------------------------

### Initialize Model and Inspect Parameters

Source: https://github.com/benedekrozemberczki/pytorch_geometric_temporal/blob/master/notebooks/a3tgcn_for_traffic_forecasting.ipynb

Sets up the model and optimizer, then prints the state dictionary and total parameter count.

```python
# Create model and optimizers
model = TemporalGNN(node_features=2, periods=12, batch_size=batch_size).to(DEVICE)
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
loss_fn = torch.nn.MSELoss()


print('Net\'s state_dict:')
total_param = 0
for param_tensor in model.state_dict():
    print(param_tensor, '\t', model.state_dict()[param_tensor].size())
    total_param += np.prod(model.state_dict()[param_tensor].size())
print('Net\'s total params:', total_param)
#--------------------------------------------------
print('Optimizer\'s state_dict:')
for var_name in optimizer.state_dict():
    print(var_name, '\t', optimizer.state_dict()[var_name])
```

--------------------------------

### Initialize Optimizer and Print Model/Optimizer States

Source: https://github.com/benedekrozemberczki/pytorch_geometric_temporal/blob/master/notebooks/astgcn_for_traffic_flow_forecasting.ipynb

Sets up the Adam optimizer for the ASTGCN model and prints the model's state dictionary, including parameter names, sizes, and devices, along with the total number of parameters. It also prints the optimizer's state dictionary.

```python
#------------------------------------------------------
learning_rate = 0.001
optimizer = optim.Adam(net.parameters(), lr=learning_rate)

print('Net\'s state_dict:')
total_param = 0
for param_tensor in net.state_dict():
    print(param_tensor, '\t', net.state_dict()[param_tensor].size(), '\t', net.state_dict()[param_tensor].device)
    total_param += np.prod(net.state_dict()[param_tensor].size())
print('Net\'s total params:', total_param)
#--------------------------------------------------
print('Optimizer\'s state_dict:')
for var_name in optimizer.state_dict():
    print(var_name, '\t', optimizer.state_dict()[var_name])
```

--------------------------------

### Initialize TensorBoard SummaryWriter

Source: https://github.com/benedekrozemberczki/pytorch_geometric_temporal/blob/master/notebooks/astgcn_for_traffic_flow_forecasting.ipynb

Initializes a SummaryWriter for TensorBoard logging. This allows visualization of training metrics and graphs. The log directory is set to the current directory.

```python
from tensorboardX import SummaryWriter
sw = SummaryWriter(logdir='.', flush_secs=5)
```

--------------------------------

### Run the test suite

Source: https://github.com/benedekrozemberczki/pytorch_geometric_temporal/blob/master/CONTRIBUTING.md

Execute the entire test suite located in the test directory.

```bash
python setup.py test
```

--------------------------------

### Initialize training state variables

Source: https://github.com/benedekrozemberczki/pytorch_geometric_temporal/blob/master/notebooks/astgcn_for_traffic_flow_forecasting.ipynb

Sets initial values for tracking training progress and best model performance.

```python
global_step = 0
best_epoch = 0
best_val_loss = np.inf
start_time= time()
```

--------------------------------

### Initialize ASTGCN Model

Source: https://github.com/benedekrozemberczki/pytorch_geometric_temporal/blob/master/notebooks/astgcn_for_traffic_flow_forecasting.ipynb

Initializes the ASTGCN model with specified parameters. Ensure all required parameters like `nb_block`, `in_channels`, `K`, `nb_chev_filter`, `nb_time_filter`, `time_strides`, `num_for_predict`, `len_input`, `num_of_vertices`, and `DEVICE` are defined before this step.

```python
nb_block = 2
in_channels = 1
K = 3
nb_chev_filter = 64
nb_time_filter = 64
time_strides = num_of_hours
num_for_predict = 12
len_input = 12
#L_tilde = scaled_Laplacian(adj_mx)
#cheb_polynomials = [torch.from_numpy(i).type(torch.FloatTensor).to(DEVICE) for i in cheb_polynomial(L_tilde, K)]
net = ASTGCN( nb_block, in_channels, K, nb_chev_filter, nb_time_filter, time_strides, num_for_predict, len_input, num_of_vertices).to(DEVICE)

print(net)
```

--------------------------------

### Calculate Historical Data Indices

Source: https://github.com/benedekrozemberczki/pytorch_geometric_temporal/blob/master/notebooks/processing_traffic_data_for_deep_learning_projects.ipynb

Calculates start and end indices for historical data segments based on specified lookback periods and prediction lengths. Raises ValueError if points_per_hour is not positive and returns None if indices fall outside the sequence length.

```python
def search_data(sequence_length, num_of_depend, label_start_idx,num_for_predict, units, points_per_hour):
    '''
    Parameters
    ----------
    sequence_length: int, length of all history data
    num_of_depend: int,
    label_start_idx: int, the first index of predicting target
    num_for_predict: int, the number of points will be predicted for each sample
    units: int, week: 7 * 24, day: 24, recent(hour): 1
    points_per_hour: int, number of points per hour, depends on data
    Returns
    ----------
    list[(start_idx, end_idx)]
    '''

    if points_per_hour < 0:
        raise ValueError("points_per_hour should be greater than 0!")

    if label_start_idx + num_for_predict > sequence_length:
        return None

    x_idx = []
    for i in range(1, num_of_depend + 1):
        start_idx = label_start_idx - points_per_hour * units * i
        end_idx = start_idx + num_for_predict
        if start_idx >= 0:
            x_idx.append((start_idx, end_idx))
        else:
            return None

    if len(x_idx) != num_of_depend:
        return None

    return x_idx[::-1]
```

--------------------------------

### Training Loop with Index-Batching

Source: https://github.com/benedekrozemberczki/pytorch_geometric_temporal/blob/master/examples/indexBatching/README.md

This code snippet demonstrates a typical training loop utilizing index-batching with a static graph dataset and temporal signals in PyTorch Geometric Temporal. It shows data loading, forward pass, loss calculation, and backpropagation.

```python
train_dataloader, _, _, edges, edge_weights, means, stds = loader.get_index_dataset(batch_size=batch_size)

for batch in train_dataloader:
            X_batch, y_batch = batch

            # Forward pass
            outputs = model(X_batch, edges, edge_weights) 

            # Calculate loss 
            loss = masked_mae_loss((outputs * std) + mean, (y_batch * std) + mean)

            # Backward pass
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()


```

--------------------------------

### Training Console Output

Source: https://github.com/benedekrozemberczki/pytorch_geometric_temporal/blob/master/notebooks/astgcn_for_traffic_flow_forecasting.ipynb

Sample output generated during the training process, showing validation losses and global step progress.

```text
validation batch 1 / 107, loss: 277.87
validation batch 101 / 107, loss: 334.15
save parameters to file: ./epoch_0.params
global step: 200, training loss: 103.08, time: 38.94s
validation batch 1 / 107, loss: 54.61
validation batch 101 / 107, loss: 83.22
save parameters to file: ./epoch_1.params
global step: 400, training loss: 39.34, time: 77.39s
global step: 600, training loss: 26.75, time: 108.20s
validation batch 1 / 107, loss: 39.82
validation batch 101 / 107, loss: 43.75
save parameters to file: ./epoch_2.params
global step: 800, training loss: 27.49, time: 146.22s
validation batch 1 / 107, loss: 33.06
validation batch 101 / 107, loss: 36.45
save parameters to file: ./epoch_3.params
global step: 1000, training loss: 26.49, time: 184.33s
global step: 1200, training loss: 26.64, time: 215.03s
validation batch 1 / 107, loss: 28.73
validation batch 101 / 107, loss: 34.21
save parameters to file: ./epoch_4.params
global step: 1400, training loss: 23.15, time: 253.47s
validation batch 1 / 107, loss: 28.88
validation batch 101 / 107, loss: 33.24
save parameters to file: ./epoch_5.params
global step: 1600, training loss: 26.00, time: 291.49s
global step: 1800, training loss: 20.62, time: 322.22s
validation batch 1 / 107, loss: 27.75
validation batch 101 / 107, loss: 32.38
save parameters to file: ./epoch_6.params
global step: 2000, training loss: 21.06, time: 360.26s
global step: 2200, training loss: 21.80, time: 390.96s
validation batch 1 / 107, loss: 27.43
validation batch 101 / 107, loss: 32.48
global step: 2400, training loss: 19.55, time: 428.93s
validation batch 1 / 107, loss: 27.28
validation batch 101 / 107, loss: 32.15
global step: 2600, training loss: 21.54, time: 467.05s
global step: 2800, training loss: 19.25, time: 497.67s
validation batch 1 / 107, loss: 28.14
validation batch 101 / 107, loss: 31.68
save parameters to file: ./epoch_9.params
global step: 3000, training loss: 23.09, time: 535.62s
validation batch 1 / 107, loss: 28.35
validation batch 101 / 107, loss: 31.51
save parameters to file: ./epoch_10.params
global step: 3200, training loss: 19.72, time: 573.80s
global step: 3400, training loss: 21.37, time: 604.43s
validation batch 1 / 107, loss: 28.92
validation batch 101 / 107, loss: 31.78
global step: 3600, training loss: 20.88, time: 642.30s
global step: 3800, training loss: 20.98, time: 672.92s
validation batch 1 / 107, loss: 28.47
validation batch 101 / 107, loss: 31.58
global step: 4000, training loss: 17.25, time: 710.76s
validation batch 1 / 107, loss: 27.29
validation batch 101 / 107, loss: 31.51
save parameters to file: ./epoch_13.params
global step: 4200, training loss: 19.48, time: 748.75s
global step: 4400, training loss: 19.07, time: 779.38s
validation batch 1 / 107, loss: 26.81
validation batch 101 / 107, loss: 31.57
save parameters to file: ./epoch_14.params
global step: 4600, training loss: 19.82, time: 817.15s
validation batch 1 / 107, loss: 27.52
validation batch 101 / 107, loss: 31.10
global step: 4800, training loss: 19.51, time: 854.94s
global step: 5000, training loss: 18.88, time: 885.64s
validation batch 1 / 107, loss: 26.25
validation batch 101 / 107, loss: 31.66
save parameters to file: ./epoch_16.params
global step: 5200, training loss: 17.55, time: 923.73s
global step: 5400, training loss: 19.11, time: 954.32s
validation batch 1 / 107, loss: 26.75
validation batch 101 / 107, loss: 31.25
save parameters to file: ./epoch_17.params
global step: 5600, training loss: 20.55, time: 992.44s
```

--------------------------------

### Prepare Testing Data Tensors

Source: https://github.com/benedekrozemberczki/pytorch_geometric_temporal/blob/master/notebooks/a3tgcn_for_traffic_forecasting.ipynb

Converts the testing data features and targets into NumPy arrays and then into PyTorch tensors. Moves the tensors to the specified device (GPU if available) and creates a DataLoader for batching.

```python
test_input = np.array(test_dataset.features) # (, 207, 2, 12)
test_target = np.array(test_dataset.targets) # (, 207, 12)
test_x_tensor = torch.from_numpy(test_input).type(torch.FloatTensor).to(DEVICE)  # (B, N, F, T)
test_target_tensor = torch.from_numpy(test_target).type(torch.FloatTensor).to(DEVICE)  # (B, N, T)
test_dataset_new = torch.utils.data.TensorDataset(test_x_tensor, test_target_tensor)
test_loader = torch.utils.data.DataLoader(test_dataset_new, batch_size=batch_size, shuffle=shuffle,drop_last=True)
```

--------------------------------

### Visualize Predictions vs. Ground Truth

Source: https://github.com/benedekrozemberczki/pytorch_geometric_temporal/blob/master/notebooks/astgcn_for_traffic_flow_forecasting.ipynb

Visualizes the predictions of the first 50 detectors against the ground truth over a 12-time step window. Requires matplotlib and torch.

```python
from matplotlib.pyplot import figure

figure(figsize=(30,4), dpi=80)
for i in range(50):
    new_i = i * 12
    plt.plot(range(0+new_i,12+new_i),sample_output[i].detach().cpu().numpy(), color = 'red')
    plt.plot(range(0+new_i,12+new_i),sample_labels[i].cpu().numpy(), color='blue')
plt.show()
```

--------------------------------

### Prepare Training Data Tensors

Source: https://github.com/benedekrozemberczki/pytorch_geometric_temporal/blob/master/notebooks/a3tgcn_for_traffic_forecasting.ipynb

Converts the training data features and targets into NumPy arrays and then into PyTorch tensors. Moves the tensors to the specified device (GPU if available) and creates a DataLoader for batching.

```python
train_input = np.array(train_dataset.features) # (27399, 207, 2, 12)
train_target = np.array(train_dataset.targets) # (27399, 207, 12)
train_x_tensor = torch.from_numpy(train_input).type(torch.FloatTensor).to(DEVICE)  # (B, N, F, T)
train_target_tensor = torch.from_numpy(train_target).type(torch.FloatTensor).to(DEVICE)  # (B, N, T)
train_dataset_new = torch.utils.data.TensorDataset(train_x_tensor, train_target_tensor)
train_loader = torch.utils.data.DataLoader(train_dataset_new, batch_size=batch_size, shuffle=shuffle,drop_last=True)
```

--------------------------------

### Load Hungarian Chickenpox Dataset

Source: https://github.com/benedekrozemberczki/pytorch_geometric_temporal/blob/master/docs/source/notes/introduction.md

Initializes the ChickenpoxDatasetLoader and retrieves the dataset as a StaticGraphTemporalSignal object.

```python
from torch_geometric_temporal.dataset import ChickenpoxDatasetLoader

loader = ChickenpoxDatasetLoader()

dataset = loader.get_dataset()
```

--------------------------------

### Load and Split Chickenpox Dataset

Source: https://github.com/benedekrozemberczki/pytorch_geometric_temporal/blob/master/docs/source/notes/introduction.md

Initializes the Chickenpox dataset loader and splits the temporal signal into training and testing sets.

```python
from torch_geometric_temporal.dataset import ChickenpoxDatasetLoader
from torch_geometric_temporal.signal import temporal_signal_split

loader = ChickenpoxDatasetLoader()

dataset = loader.get_dataset()

train_dataset, test_dataset = temporal_signal_split(dataset, train_ratio=0.2)
```

--------------------------------

### Split Dataset for Training and Testing

Source: https://github.com/benedekrozemberczki/pytorch_geometric_temporal/blob/master/notebooks/a3tgcn_for_traffic_forecasting.ipynb

Splits the loaded dataset into training and testing sets using a specified training ratio. This is crucial for evaluating the model's performance on unseen data.

```python
train_dataset, test_dataset = temporal_signal_split(dataset, train_ratio=0.8)

print("Number of train buckets: ", len(set(train_dataset)))
print("Number of test buckets: ", len(set(test_dataset)))
```

--------------------------------

### Execute Training Loop

Source: https://github.com/benedekrozemberczki/pytorch_geometric_temporal/blob/master/notebooks/a3tgcn_for_traffic_forecasting.ipynb

Iterates through the training loader to perform forward passes, loss calculation, and backpropagation.

```python
model.train()

for epoch in range(1):
    step = 0
    loss_list = []
    for encoder_inputs, labels in train_loader:
        y_hat = model(encoder_inputs, static_edge_index)         # Get model predictions
        loss = loss_fn(y_hat, labels) # Mean squared error #loss = torch.mean((y_hat-labels)**2)  sqrt to change it to rmse
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()
        step= step+ 1
        loss_list.append(loss.item())
        if step % 100 == 0 :
            print(sum(loss_list)/len(loss_list))
    print("Epoch {} train RMSE: {:.4f}".format(epoch, sum(loss_list)/len(loss_list)))
```

--------------------------------

### Configure loss function selection

Source: https://github.com/benedekrozemberczki/pytorch_geometric_temporal/blob/master/notebooks/astgcn_for_traffic_flow_forecasting.ipynb

Sets up the criterion and masking flag based on the chosen loss function string.

```python
masked_flag=0
criterion = nn.L1Loss().to(DEVICE)
criterion_masked = masked_mae
loss_function = 'mse'

metric_method = 'unmask'
missing_value=0.0


if loss_function=='masked_mse':
    criterion_masked = masked_mse         #nn.MSELoss().to(DEVICE)
    masked_flag=1
elif loss_function=='masked_mae':
    criterion_masked = masked_mae
    masked_flag = 1
elif loss_function == 'mae':
    criterion = nn.L1Loss().to(DEVICE)
    masked_flag = 0
elif loss_function == 'rmse':
    criterion = nn.MSELoss().to(DEVICE)
    masked_flag= 0
```

--------------------------------

### BibTeX Citation for PGT-I Scaling

Source: https://github.com/benedekrozemberczki/pytorch_geometric_temporal/blob/master/README.md

Use this citation for the PGT-I extension, focusing on scaling spatiotemporal GNNs with memory-efficient distributed training.

```bibtex
@misc{ockerman2025pgtiscalingspatiotemporalgnns,
      title={PGT-I: Scaling Spatiotemporal GNNs with Memory-Efficient Distributed Training},
      author={Seth Ockerman and Amal Gueroudji and Tanwi Mallick and Yixuan He and Line Pouchard and Robert Ross and Shivaram Venkataraman},
      year={2025},
      eprint={2507.11683},
      archivePrefix={arXiv},
      primaryClass={cs.DC},
      url={https://arxiv.org/abs/2507.11683},
}

```

--------------------------------

### Generate Dataset Splits

Source: https://github.com/benedekrozemberczki/pytorch_geometric_temporal/blob/master/notebooks/processing_traffic_data_for_deep_learning_projects.ipynb

Configures parameters and executes the dataset generation process for training, validation, and testing.

```python
data = np.load(graph_signal_matrix_filename)
data['data'].shape

num_of_vertices = 307
points_per_hour = 12
num_for_predict = 12
num_of_weeks = 0
num_of_days = 0
num_of_hours = 1

training_set, validation_set, testing_set = read_and_generate_dataset(graph_signal_matrix_filename, 0, 0, num_of_hours, 
                                                                      num_for_predict, points_per_hour=points_per_hour)
```

--------------------------------

### Run Project Tests

Source: https://github.com/benedekrozemberczki/pytorch_geometric_temporal/blob/master/README.md

Command to execute the test suite using pytest.

```sh
$ python -m pytest test
```

--------------------------------

### Visualize and Prepare Graph Data

Source: https://github.com/benedekrozemberczki/pytorch_geometric_temporal/blob/master/notebooks/astgcn_for_traffic_flow_forecasting.ipynb

Uses the generated adjacency matrix to visualize the graph with NetworkX and prepare edge indices for PyTorch Geometric.

```python
id_filename = None
adj_filename = '../input/pems-dataset/data/PEMS04/PEMS04.csv'
num_of_vertices = 307
adj_mx, distance_mx = get_adjacency_matrix(adj_filename, num_of_vertices, id_filename) #  adj_mx and distance_mx (307, 307)

rows, cols = np.where(adj_mx == 1)
edges = zip(rows.tolist(), cols.tolist())
gr = nx.Graph()
gr.add_edges_from(edges)
nx.draw(gr, node_size=3)
plt.show()
rows, cols = np.where(adj_mx == 1)
edges = zip(rows.tolist(), cols.tolist())
edge_index_data = torch.LongTensor(np.array([rows, cols])).to(DEVICE)
```

--------------------------------

### Training Loop Output

Source: https://github.com/benedekrozemberczki/pytorch_geometric_temporal/blob/master/notebooks/a3tgcn_for_traffic_forecasting.ipynb

The console output showing training progress and final RMSE.

```text
Output:
0.797569936811924
0.6679292801022529
0.6100249779224396
0.5751089830696583
0.5526966672539712
0.5415776373445987
0.5327171741213117
0.5293058890849351
Epoch 0 train RMSE: 0.5256
```

--------------------------------

### Print ASTGCN Model Structure

Source: https://github.com/benedekrozemberczki/pytorch_geometric_temporal/blob/master/notebooks/astgcn_for_traffic_flow_forecasting.ipynb

Prints the detailed structure of the initialized ASTGCN model, showing all layers and their configurations. This output is useful for verifying the model architecture.

```text
ASTGCN(
  (_blocklist): ModuleList(
    (0): ASTGCNBlock(
      (_temporal_attention): TemporalAttention()
      (_spatial_attention): SpatialAttention()
      (_chebconv_attention): ChebConvAttention(1, 64, K=3, normalization=None)
      (_time_convolution): Conv2d(64, 64, kernel_size=(1, 3), stride=(1, 1), padding=(0, 1))
      (_residual_convolution): Conv2d(1, 64, kernel_size=(1, 1), stride=(1, 1))
      (_layer_norm): LayerNorm((64,), eps=1e-05, elementwise_affine=True)
    )
    (1): ASTGCNBlock(
      (_temporal_attention): TemporalAttention()
      (_spatial_attention): SpatialAttention()
      (_chebconv_attention): ChebConvAttention(64, 64, K=3, normalization=None)
      (_time_convolution): Conv2d(64, 64, kernel_size=(1, 3), stride=(1, 1), padding=(0, 1))
      (_residual_convolution): Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1))
      (_layer_norm): LayerNorm((64,), eps=1e-05, elementwise_affine=True)
    )
  )
  (_final_conv): Conv2d(12, 12, kernel_size=(1, 64), stride=(1, 1))
)
```

--------------------------------

### Load Static Graph

Source: https://github.com/benedekrozemberczki/pytorch_geometric_temporal/blob/master/notebooks/a3tgcn_for_traffic_forecasting.ipynb

Extracts the edge index from the first snapshot of the training dataset.

```python
for snapshot in train_dataset:
    static_edge_index = snapshot.edge_index.to(DEVICE)
    break;
```

--------------------------------

### read_and_generate_dataset

Source: https://github.com/benedekrozemberczki/pytorch_geometric_temporal/blob/master/notebooks/processing_traffic_data_for_deep_learning_projects.ipynb

Reads a graph signal matrix file and generates a full dataset for training.

```APIDOC
## read_and_generate_dataset

### Description
Loads a graph signal matrix from a file and iterates through the sequence to generate feature and target arrays.

### Parameters
- **graph_signal_matrix_filename** (str) - Path to the .npz file.
- **num_of_weeks** (int) - Number of weeks.
- **num_of_days** (int) - Number of days.
- **num_of_hours** (int) - Number of hours.
- **num_for_predict** (int) - Number of points to predict.
- **points_per_hour** (int) - Points per hour (default 12).

### Response
- **feature** (np.ndarray) - Generated feature matrix.
- **target** (np.ndarray) - Generated target matrix.
```

--------------------------------

### Process and Split PEMS Dataset

Source: https://github.com/benedekrozemberczki/pytorch_geometric_temporal/blob/master/notebooks/processing_traffic_data_for_deep_learning_projects.ipynb

Prepares temporal samples by expanding dimensions and transposing arrays, then splits the collection into training, validation, and testing sets.

```python
        week_sample, day_sample, hour_sample, target = sample #  week_sample, day_sample are None because we are predicting per hour
        #print(target.shape) # hour_sample and target (12, 307, 3)
        sample = []  # [(week_sample),(day_sample),(hour_sample),target,time_sample]
#-------------------------------- Ignore
        if num_of_weeks > 0:
            week_sample = np.expand_dims(week_sample, axis=0).transpose((0, 2, 3, 1))  # (1,N,F,T)
            sample.append(week_sample)

        if num_of_days > 0:
            day_sample = np.expand_dims(day_sample, axis=0).transpose((0, 2, 3, 1))  # (1,N,F,T)
            sample.append(day_sample)
#----------------------------------Continue
        if num_of_hours > 0:
            hour_sample = np.expand_dims(hour_sample, axis=0).transpose((0, 2, 3, 1))  # (1,N,F,T)
            sample.append(hour_sample)

        target = np.expand_dims(target, axis=0).transpose((0, 2, 3, 1))[:, :, 0, :]  # (1,N,T)
        sample.append(target)
        time_sample = np.expand_dims(np.array([idx]), axis=0)  # (1,1)
        sample.append(time_sample)
        all_samples.append(sample)#sampe：[(week_sample),(day_sample),(hour_sample),target,time_sample] = [(1,N,F,Tw),(1,N,F,Td),(1,N,F,Th),(1,N,Tpre),(1,1)]

    split_line1 = int(len(all_samples) * 0.6)
    split_line2 = int(len(all_samples) * 0.8)

    training_set = [np.concatenate(i, axis=0)  for i in zip(*all_samples[:split_line1])] #[(B,N,F,Tw),(B,N,F,Td),(B,N,F,Th),(B,N,Tpre),(B,1)]
    validation_set = [np.concatenate(i, axis=0) for i in zip(*all_samples[split_line1: split_line2])]
    testing_set = [np.concatenate(i, axis=0) for i in zip(*all_samples[split_line2:])]

    return training_set, validation_set, testing_set
```

--------------------------------

### Model Architecture Output

Source: https://github.com/benedekrozemberczki/pytorch_geometric_temporal/blob/master/notebooks/a3tgcn_for_traffic_forecasting.ipynb

The string representation of the initialized TemporalGNN model structure.

```text
Result:
TemporalGNN(
  (tgnn): A3TGCN2(
    (_base_tgcn): TGCN2(
      (conv_z): GCNConv(2, 32)
      (linear_z): Linear(in_features=64, out_features=32, bias=True)
      (conv_r): GCNConv(2, 32)
      (linear_r): Linear(in_features=64, out_features=32, bias=True)
      (conv_h): GCNConv(2, 32)
      (linear_h): Linear(in_features=64, out_features=32, bias=True)
    )
  )
  (linear): Linear(in_features=32, out_features=12, bias=True)
)
```

--------------------------------

### ASTGCN Model State Dictionary Output

Source: https://github.com/benedekrozemberczki/pytorch_geometric_temporal/blob/master/notebooks/astgcn_for_traffic_flow_forecasting.ipynb

Displays the state dictionary of the ASTGCN model, listing each parameter's name, its tensor size, and the device it resides on (e.g., CUDA). This is crucial for understanding model parameters and their distribution across devices.

```text
Net's state_dict:
_blocklist.0._temporal_attention._U1 	 torch.Size([307]) 	 cuda:0
_blocklist.0._temporal_attention._U2 	 torch.Size([1, 307]) 	 cuda:0
_blocklist.0._temporal_attention._U3 	 torch.Size([1]) 	 cuda:0
_blocklist.0._temporal_attention._be 	 torch.Size([1, 12, 12]) 	 cuda:0
_blocklist.0._temporal_attention._Ve 	 torch.Size([12, 12]) 	 cuda:0
_blocklist.0._spatial_attention._W1 	 torch.Size([12]) 	 cuda:0
_blocklist.0._spatial_attention._W2 	 torch.Size([1, 12]) 	 cuda:0
_blocklist.0._spatial_attention._W3 	 torch.Size([1]) 	 cuda:0
_blocklist.0._spatial_attention._bs 	 torch.Size([1, 307, 307]) 	 cuda:0
_blocklist.0._spatial_attention._Vs 	 torch.Size([307, 307]) 	 cuda:0
_blocklist.0._chebconv_attention._weight 	 torch.Size([3, 1, 64]) 	 cuda:0
_blocklist.0._chebconv_attention._bias 	 torch.Size([64]) 	 cuda:0
_blocklist.0._time_convolution.weight 	 torch.Size([64, 64, 1, 3]) 	 cuda:0
_blocklist.0._time_convolution.bias 	 torch.Size([64]) 	 cuda:0
_blocklist.0._residual_convolution.weight 	 torch.Size([64, 1, 1, 1]) 	 cuda:0
_blocklist.0._residual_convolution.bias 	 torch.Size([64]) 	 cuda:0
_blocklist.0._layer_norm.weight 	 torch.Size([64]) 	 cuda:0
_blocklist.0._layer_norm.bias 	 torch.Size([64]) 	 cuda:0
_blocklist.1._temporal_attention._U1 	 torch.Size([307]) 	 cuda:0
_blocklist.1._temporal_attention._U2 	 torch.Size([64, 307]) 	 cuda:0
_blocklist.1._temporal_attention._U3 	 torch.Size([64]) 	 cuda:0
_blocklist.1._temporal_attention._be 	 torch.Size([1, 12, 12]) 	 cuda:0
_blocklist.1._temporal_attention._Ve 	 torch.Size([12, 12]) 	 cuda:0
_blocklist.1._spatial_attention._W1 	 torch.Size([12]) 	 cuda:0
_blocklist.1._spatial_attention._W2 	 torch.Size([64, 12]) 	 cuda:0
_blocklist.1._spatial_attention._W3 	 torch.Size([64]) 	 cuda:0
_blocklist.1._spatial_attention._bs 	 torch.Size([1, 307, 307]) 	 cuda:0
_blocklist.1._spatial_attention._Vs 	 torch.Size([307, 307]) 	 cuda:0
_blocklist.1._chebconv_attention._weight 	 torch.Size([3, 64, 64]) 	 cuda:0
_blocklist.1._chebconv_attention._bias 	 torch.Size([64]) 	 cuda:0
_blocklist.1._time_convolution.weight 	 torch.Size([64, 64, 1, 3]) 	 cuda:0
_blocklist.1._time_convolution.bias 	 torch.Size([64]) 	 cuda:0
_blocklist.1._residual_convolution.weight 	 torch.Size([64, 64, 1, 1]) 	 cuda:0
_blocklist.1._residual_convolution.bias 	 torch.Size([64]) 	 cuda:0
_blocklist.1._layer_norm.weight 	 torch.Size([64]) 	 cuda:0
_blocklist.1._layer_norm.bias 	 torch.Size([64]) 	 cuda:0
_final_conv.weight 	 torch.Size([12, 12, 1, 64]) 	 cuda:0
_final_conv.bias 	 torch.Size([12]) 	 cuda:0
Net's total params: 450159
```

--------------------------------

### Prepare and Normalize Data

Source: https://github.com/benedekrozemberczki/pytorch_geometric_temporal/blob/master/notebooks/processing_traffic_data_for_deep_learning_projects.ipynb

Concatenates input sets and applies the normalization function to prepare the final data dictionary.

```python
train_x = np.concatenate(training_set[:-2], axis=-1)  # (B,N,F,T')
val_x = np.concatenate(validation_set[:-2], axis=-1)
test_x = np.concatenate(testing_set[:-2], axis=-1)

train_target = training_set[-2]  # (B,N,T)
val_target = validation_set[-2]
test_target = testing_set[-2]

train_timestamp = training_set[-1]  # (B,1)
val_timestamp = validation_set[-1]
test_timestamp = testing_set[-1]

(stats, train_x_norm, val_x_norm, test_x_norm) = normalization(train_x, val_x, test_x)

all_data = {'train': { 'x': train_x_norm, 'target': train_target,'timestamp': train_timestamp},
            'val': {'x': val_x_norm, 'target': val_target, 'timestamp': val_timestamp},
            'test': {'x': test_x_norm, 'target': test_target, 'timestamp': test_timestamp},
            'stats': {'_mean': stats['_mean'], '_std': stats['_std']} }
```

--------------------------------

### Load Temporal Graph Data Function

Source: https://github.com/benedekrozemberczki/pytorch_geometric_temporal/blob/master/notebooks/astgcn_for_traffic_flow_forecasting.ipynb

Use this function to load and preprocess temporal graph data. It splits the data into training, validation, and testing sets and returns PyTorch DataLoaders. Ensure the input filename points to a .npz file containing 'train_x', 'train_target', 'val_x', 'val_target', 'test_x', 'test_target', 'mean', and 'std'.

```python
def load_graphdata_channel1(graph_signal_matrix_filename, num_of_hours, num_of_days, num_of_weeks,
                            batch_size,
                            shuffle=True, DEVICE = torch.device('cuda:0')):
    '''
    :param graph_signal_matrix_filename: str
    :param num_of_hours: int
    :param num_of_days: int
    :param num_of_weeks: int
    :param DEVICE:
    :param batch_size: int
    :return:
    three DataLoaders, each dataloader contains:
    test_x_tensor: (B, N_nodes, in_feature, T_input)
    test_decoder_input_tensor: (B, N_nodes, T_output)
    test_target_tensor: (B, N_nodes, T_output)
    '''

    file = os.path.basename(graph_signal_matrix_filename).split('.')[0]
    filename = os.path.join('../input/processing-traffic-data-for-deep-learning-projects/', file + '_r' + str(num_of_hours) + '_d' + str(num_of_days) + '_w' + str(num_of_weeks)) +'_astcgn'
    print('load file:', filename)

    file_data = np.load(filename + '.npz')
    train_x = file_data['train_x']  # (10181, 307, 3, 12)
    train_x = train_x[:, :, 0:1, :]
    train_target = file_data['train_target']  # (10181, 307, 12)

    val_x = file_data['val_x']
    val_x = val_x[:, :, 0:1, :]
    val_target = file_data['val_target']

    test_x = file_data['test_x']
    test_x = test_x[:, :, 0:1, :]
    test_target = file_data['test_target']

    mean = file_data['mean'][:, :, 0:1, :]  # (1, 1, 3, 1)
    std = file_data['std'][:, :, 0:1, :]  # (1, 1, 3, 1)

    # ------- train_loader ------- 
    train_x_tensor = torch.from_numpy(train_x).type(torch.FloatTensor).to(DEVICE)  # (B, N, F, T)
    train_target_tensor = torch.from_numpy(train_target).type(torch.FloatTensor).to(DEVICE)  # (B, N, T)
    train_dataset = torch.utils.data.TensorDataset(train_x_tensor, train_target_tensor)
    train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=shuffle)

    # ------- val_loader ------- 
    val_x_tensor = torch.from_numpy(val_x).type(torch.FloatTensor).to(DEVICE)  # (B, N, F, T)
    val_target_tensor = torch.from_numpy(val_target).type(torch.FloatTensor).to(DEVICE)  # (B, N, T)
    val_dataset = torch.utils.data.TensorDataset(val_x_tensor, val_target_tensor)
    val_loader = torch.utils.data.DataLoader(val_dataset, batch_size=batch_size, shuffle=False)

    # ------- test_loader ------- 
    test_x_tensor = torch.from_numpy(test_x).type(torch.FloatTensor).to(DEVICE)  # (B, N, F, T)
    test_target_tensor = torch.from_numpy(test_target).type(torch.FloatTensor).to(DEVICE)  # (B, N, T)
    test_dataset = torch.utils.data.TensorDataset(test_x_tensor, test_target_tensor)
    test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=batch_size, shuffle=False)

    # print
    print('train:', train_x_tensor.size(), train_target_tensor.size())
    print('val:', val_x_tensor.size(), val_target_tensor.size())
    print('test:', test_x_tensor.size(), test_target_tensor.size())

    return train_loader, train_target_tensor, val_loader, val_target_tensor, test_loader, test_target_tensor, mean, std
```