53 lines
2.1 KiB
Python
53 lines
2.1 KiB
Python
import os
|
|
|
|
import numpy as np
|
|
import pandas as pd
|
|
|
|
"""
|
|
return ndarray (time_step, node_num, channel)
|
|
"""
|
|
|
|
|
|
def load_st_dataset(dataset):
|
|
# output B, N, D
|
|
if dataset == 'PEMSD3':
|
|
data_path = os.path.join('./data/trafficflow/PeMS03/PEMS03.npz')
|
|
data = np.load(data_path)['data'][:, :, 0] # only the first dimension, traffic flow data
|
|
data = data.astype(np.float32)
|
|
elif dataset == 'PEMSD4':
|
|
data_path = os.path.join('./data/trafficflow/PeMS04/PEMS04.npz')
|
|
data = np.load(data_path)['data'][:, :, 0] # only the first dimension, traffic flow data
|
|
data = data.astype(np.float32)
|
|
elif dataset == 'PEMSD7':
|
|
data_path = os.path.join('./data/trafficflow/PeMS07/PeMS07.npz')
|
|
data = np.load(data_path)['data'][:, :, 0] # only the first dimension, traffic flow data
|
|
data = data.astype(np.float32)
|
|
elif dataset == 'PEMSD8':
|
|
data_path = os.path.join('./data/trafficflow/PeMS08/PeMS08.npz')
|
|
data = np.load(data_path)['data'][:, :, 0] # only the first dimension, traffic flow data
|
|
data = data.astype(np.float32)
|
|
elif dataset == 'PEMSD7(L)':
|
|
data_path = os.path.join('./data/trafficflow/PeMS07(L)/PEMS07L.npz')
|
|
data = np.load(data_path)['data'][:, :, 0] # only the first dimension, traffic flow data
|
|
elif dataset == 'PEMSD7(M)':
|
|
data_path = os.path.join('./data/trafficflow/PeMS07(M)/V_228.csv')
|
|
data = np.array(pd.read_csv(data_path, header=None)) # only the first dimension, traffic flow data
|
|
elif dataset == 'METR-LA':
|
|
data_path = os.path.join('./data/trafficflow/METR-LA/METR.h5')
|
|
data = pd.read_hdf(data_path)
|
|
elif dataset == 'BJ':
|
|
data_path = os.path.join('./data/trafficflow/BJ/BJ500.csv')
|
|
data = np.array(pd.read_csv(data_path, header=0, index_col=0))
|
|
else:
|
|
raise ValueError
|
|
if len(data.shape) == 2:
|
|
data = np.expand_dims(data, axis=-1)
|
|
print('Load %s Dataset shaped: ' % dataset, data.shape, data.max(), data.min(), data.mean(), np.median(data))
|
|
return data
|
|
|
|
|
|
if __name__ == '__main__':
|
|
dataset = 'PEMSD8'
|
|
data = load_st_dataset(dataset)
|
|
print("Finished")
|