STONE: A Spatio-temporal OOD Learning Framework Kills Both Structural and Temporal Shifts (KDD 2024)

This is the official repository of our KDD 2024 paper. This paper represents the first comprehensive investigation into the traffic prediction considering both temporal and spatial shifts. We propose a novel spatio-temporal OOD-network framework with Fréchet Embedding called STONE. STONE effectively models spatial heterogeneity and generates temporal and spatial semantic graphs. Additionally, we introduce a graph perturbation mechanism to enhance the model’s environmental modeling capability for better generalization. We implement extensive experiments on both datasets with spatio-temporal shifts and datasets only with temporal shift, and results demonstrate that STONE achieves competitive performance in terms of both generalization and scalability.

We show the pseudocode for the main algorithms, including the Fréchet embedding process, STONE forward process, and the training process of STONE. We will add it to the new version for a clear presentation.

1. Introduction about the datasets

1.1 Generating the SD and GBA sub-datasets from CA dataset

In the experiments of our paper, we used SD and GBA datasets with years from 2019 to 2021, which were generated from CA dataset, followed by LargeST. For example, you can download CA dataset from the provided link and please place the downloaded archive.zip file in the data/ca folder and unzip the file.

First of all, you should go through a jupyter notebook process_ca_his.ipynb in the folder data/ca to process and generate a cleaned version of the flow data. Then, please go through all the cells in the provided jupyter notebooks generate_sd_dataset.ipynb in the folder data/sd and generate_gla_dataset.ipynb in the folder data/gla respectively. Finally use the commands below to generate traffic flow data for our experiments.

python data/generate_data_for_training.py --dataset sd_gba --years 2019_2020_2021

Moreover, you can also generate the other years of data, as well as the two additional remaining subdatasets.

1.2 Generating the additional PM2.5 Knowair dataset

We implement extra experiments on Knowair. For example, you can download Knowair dataset from the provided link and please place the downloaded Knowair.npy file in the Knowair folder and complete the files in the Knowair/data folder.

2. Environmental Requirments

The experiment requires the same environment as LargeST, and need to add the libraries mentioned in the requirements in Knowair.

3. Model Running

To run STONE on LargeST, for example, you may execute this command in the terminal:

bash experiments/stone/run.sh

or directly execute the Python file in the terminal:

python experiments/stone/main.py --device cuda:0 --dataset SD --years 2019 --model_name stone --seed 0 --bs 64

To run STONE on Knowair, you may directly execute the Pyhon file in the terminal:

python Knowair/train.py