18 lines
559 B
Python
18 lines
559 B
Python
import numpy as np
|
|
import os
|
|
|
|
def load_dataset(config):
|
|
|
|
dataset_name = config['basic']['dataset']
|
|
node_num = config['data']['num_nodes']
|
|
input_dim = config['data']['input_dim']
|
|
data = None
|
|
match dataset_name:
|
|
case 'EcoSolar':
|
|
data_path = os.path.join('./data/EcoSolar.npy')
|
|
data = np.load(data_path)[:, :node_num, :input_dim]
|
|
case 'PEMS08':
|
|
data_path = os.path.join('./data/PEMS08/PEMS08.npz')
|
|
data = np.load(data_path)['data'][:, :node_num, :input_dim]
|
|
|
|
return data |