Robust Probabilistic Spatiotemporal Traffic Forecasting

Final course project repository for multi-horizon traffic speed forecasting on graph sensor networks (METR-LA and PEMS-BAY), with strict time-split evaluation, robustness stress tests, and uncertainty estimation.

Overview

Traffic forecasting is a coupled dynamical system problem: each sensor depends on its own history and nearby sensors.
This repo predicts speeds at 15/30/60 minutes ahead using:

classical baselines (persistence, historical average)
a neural no-graph baseline (GRU)
a graph-temporal model (GraphTemporalGRU) with fixed/correlation/adaptive adjacency

Proposal requirements implemented:

strict chronological splits with leakage prevention
mask-aware handling for METR-LA outages (0 treated as missing)
robustness experiments (random missing points, block outages, spatial outages)
split-conformal uncertainty intervals with coverage/width metrics

Datasets

Benchmarks used:

METR-LA: 207 sensors, 5-minute intervals
PEMS-BAY: 325 sensors, 5-minute intervals

Source (Zenodo):

https://zenodo.org/records/5146275/files

Expected raw files in data/raw/:

METR-LA.csv
PEMS-BAY.csv
PEMS-BAY-META.csv
adj_mx_bay.pkl

Download + integrity check

mkdir -p data/raw
cd data/raw

curl -L -o METR-LA.csv https://zenodo.org/records/5146275/files/METR-LA.csv?download=1
curl -L -o PEMS-BAY.csv https://zenodo.org/records/5146275/files/PEMS-BAY.csv?download=1
curl -L -o PEMS-BAY-META.csv https://zenodo.org/records/5146275/files/PEMS-BAY-META.csv?download=1
curl -L -o adj_mx_bay.pkl https://zenodo.org/records/5146275/files/adj_mx_bay.pkl?download=1
cd ../..

python - <<'PY'
import hashlib, pathlib
expected = {
    "METR-LA.csv": "d3adaa79856bf610f25558fc242c7190",
    "PEMS-BAY.csv": "c8dea58987a5882e946217c22fdb8256",
    "PEMS-BAY-META.csv": "6427ca7d8b37140535886dd66d6f0211",
    "adj_mx_bay.pkl": "55d25daceac847312b7748c59ded6f77",
}
root = pathlib.Path("data/raw")
for name, md5 in expected.items():
    p = root / name
    digest = hashlib.md5(p.read_bytes()).hexdigest()
    print(f"{name}: {'OK' if digest == md5 else 'MISMATCH'} ({digest})")
PY

Environment Setup

python3 -m venv .venv
source .venv/bin/activate
pip install --upgrade pip
pip install -r requirements.txt

Repository Layout

src/: training/evaluation scripts and core utilities
data/raw/: downloaded benchmark files (gitignored)
data/processed/: processed arrays + strict splits (gitignored)
reports/: all experiment outputs (JSON/CSV/MD/figures/checkpoints)
notebooks/99_demo_end_to_end.ipynb: polished 15-minute demo notebook

Main scripts:

src/make_strict_splits.py — builds strict no-leakage splits
src/eval_baselines.py — persistence + historical-average metrics
src/train_gru.py — no-graph GRU baseline
src/train_gnn.py — graph-temporal GRU (fixed/corr-kNN/adaptive adjacency)
src/eval_robustness.py — robustness degradation experiments
src/eval_uncertainty_conformal.py — split-conformal intervals
src/summarize_strict_results.py — markdown summary table from JSON artifacts

Reproducibility (Strict Splits)

All commands below assume repo root and active .venv.

If data/processed/traffic_datasets_W12_H3_6_12.npz is missing, run the data-prep notebooks once:

notebooks/00_data_audit.ipynb
notebooks/01_make_dataset.ipynb

1) Build strict splits

python -m src.make_strict_splits \
  --data-path data/processed/traffic_datasets_W12_H3_6_12.npz \
  --output-path data/processed/splits_strict_W12_H3_6_12.npz

2) Baselines (strict)

python -m src.eval_baselines \
  --data-path data/processed/traffic_datasets_W12_H3_6_12.npz \
  --splits-path data/processed/splits_strict_W12_H3_6_12.npz \
  --split-tag strict \
  --output-path reports/baselines_strict_metrics.json \
  --results-path reports/results.json

3) Train GRU baselines

python -m src.train_gru --dataset metr \
  --splits-path data/processed/splits_strict_W12_H3_6_12.npz \
  --time-features sincos --split-tag strict --epochs 30

python -m src.train_gru --dataset pems \
  --splits-path data/processed/splits_strict_W12_H3_6_12.npz \
  --time-features none --split-tag strict --epochs 20

4) Train GraphTemporalGRU

# PEMS: fixed adjacency from artifact
python -m src.train_gnn --dataset pems \
  --adjacency pems --split-tag strict --epochs 15 --patience 5

# METR: corr-kNN + adaptive adjacency + scheduler + long run
python -m src.train_gnn --dataset metr \
  --adjacency correlation_knn --knn-k 10 \
  --use-input-mask --time-features sincos \
  --adaptive-adj on --adaptive-adj-dim 16 --adaptive-adj-alpha learn \
  --scheduler plateau --scheduler-factor 0.5 --scheduler-patience 3 --min-lr 1e-5 \
  --horizon-weights uniform --epochs 200 --patience 20 \
  --split-tag strict

5) Robustness evaluation

python -m src.eval_robustness --dataset metr \
  --checkpoint reports/models/final_metr_graph_gru.pt \
  --splits-path data/processed/splits_strict_W12_H3_6_12.npz

python -m src.eval_robustness --dataset pems \
  --checkpoint reports/models/final_pems_graph_gru.pt \
  --splits-path data/processed/splits_strict_W12_H3_6_12.npz

6) Uncertainty evaluation (split conformal)

python -m src.eval_uncertainty_conformal --dataset metr \
  --checkpoint reports/models/final_metr_graph_gru.pt \
  --splits-path data/processed/splits_strict_W12_H3_6_12.npz

python -m src.eval_uncertainty_conformal --dataset pems \
  --checkpoint reports/models/final_pems_graph_gru.pt \
  --splits-path data/processed/splits_strict_W12_H3_6_12.npz

7) Regenerate strict summary table

python -m src.summarize_strict_results \
  --results-path reports/results.json \
  --baselines-path reports/baselines_strict_metrics.json \
  --output-path reports/summary_strict.md

Final Artifacts

Key outputs in reports/:

results.json: run-level metrics registry
summary_strict.md: model comparison table (strict split)
final_model_cards.md: frozen final checkpoint metadata
robustness_summary.md + figures/robustness_*.png
uncertainty_summary.md + uncertainty_conformal.json + reliability plots
final_results_onepager.md: concise final report page

Frozen checkpoints:

reports/models/final_metr_graph_gru.pt
reports/models/final_pems_graph_gru.pt

Run the Demo Notebook

source .venv/bin/activate
jupyter lab

Open:

notebooks/99_demo_end_to_end.ipynb

The notebook is designed for a ~15 minute walkthrough and reads saved artifacts instead of launching multi-hour training.

ZachariahAlzubi/traffic-forecasting-stgnn