解决合并冲突，整合dev和main分支的更改

config/EXP/PEMSD3.yaml

@@ -0,0 +1,52 @@
+data:
+  num_nodes: 358
+  lag: 12
+  horizon: 12
+  val_ratio: 0.2
+  test_ratio: 0.2
+  tod: False
+  normalizer: std
+  column_wise: False
+  default_graph: True
+  add_time_in_day: True
+  add_day_in_week: True
+  steps_per_day: 288
+  days_per_week: 7
+
+model:
+  input_dim: 1
+  output_dim: 1
+  embed_dim: 10
+  in_len: 12
+  rnn_units: 64
+  num_layers: 1
+  cheb_order: 2
+  use_day: True
+  use_week: True
+  graph_size: 30
+  expert_nums: 8
+  top_k: 2
+
+train:
+  loss_func: mae
+  seed: 10
+  batch_size: 64
+  epochs: 300
+  lr_init: 0.003
+  weight_decay: 0
+  lr_decay: False
+  lr_decay_rate: 0.3
+  lr_decay_step: "5,20,40,70"
+  early_stop: True
+  early_stop_patience: 15
+  grad_norm: False
+  max_grad_norm: 5
+  real_value: True
+
+test:
+  mae_thresh: null
+  mape_thresh: 0.0
+
+log:
+  log_step: 2000
+  plot: False

config/EXP/PEMSD4.yaml

@@ -0,0 +1,48 @@
+data:
+  num_nodes: 307
+  lag: 12
+  horizon: 12
+  val_ratio: 0.2
+  test_ratio: 0.2
+  tod: False
+  normalizer: std
+  column_wise: False
+  default_graph: True
+  add_time_in_day: True
+  add_day_in_week: True
+  steps_per_day: 288
+  days_per_week: 7
+  cycle: 288
+
+model:
+  batch_size: 64
+  input_dim: 1
+  output_dim: 1
+  in_len: 12
+  cycle_len: 288
+
+
+train:
+  loss_func: mae
+  seed: 10
+  batch_size: 64
+  epochs: 300
+  lr_init: 0.003
+  weight_decay: 0
+  lr_decay: False
+  lr_decay_rate: 0.5
+  lr_decay_step: "5,20,40,65"
+  early_stop: True
+  early_stop_patience: 15
+  grad_norm: False
+  max_grad_norm: 5
+  real_value: True
+
+
+test:
+  mae_thresh: null
+  mape_thresh: 0.0
+
+log:
+  log_step: 200
+  plot: False

config/EXP/PEMSD7.yaml

@@ -0,0 +1,45 @@
+data:
+  num_nodes: 883
+  lag: 12
+  horizon: 12
+  val_ratio: 0.2
+  test_ratio: 0.2
+  tod: False
+  normalizer: std
+  column_wise: False
+  default_graph: True
+  add_time_in_day: True
+  add_day_in_week: True
+  steps_per_day: 288
+  days_per_week: 7
+
+model:
+  batch_size: 64
+  input_dim: 1
+  output_dim: 1
+  in_len: 12
+
+
+train:
+  loss_func: mae
+  seed: 10
+  batch_size: 64
+  epochs: 300
+  lr_init: 0.003
+  weight_decay: 0
+  lr_decay: False
+  lr_decay_rate: 0.3
+  lr_decay_step: "5,20,40,70"
+  early_stop: True
+  early_stop_patience: 15
+  grad_norm: False
+  max_grad_norm: 5
+  real_value: True
+
+test:
+  mae_thresh: null
+  mape_thresh: 0.0
+
+log:
+  log_step: 200
+  plot: False

config/EXP/PEMSD8.yaml

@@ -0,0 +1,45 @@
+data:
+  num_nodes: 170
+  lag: 12
+  horizon: 12
+  val_ratio: 0.2
+  test_ratio: 0.2
+  tod: False
+  normalizer: std
+  column_wise: False
+  default_graph: True
+  add_time_in_day: True
+  add_day_in_week: True
+  steps_per_day: 288
+  days_per_week: 7
+
+model:
+  batch_size: 64
+  input_dim: 1
+  output_dim: 1
+  in_len: 12
+
+
+train:
+  loss_func: mae
+  seed: 10
+  batch_size: 64
+  epochs: 300
+  lr_init: 0.003
+  weight_decay: 0
+  lr_decay: False
+  lr_decay_rate: 0.3
+  lr_decay_step: "5,20,40,70"
+  early_stop: True
+  early_stop_patience: 15
+  grad_norm: False
+  max_grad_norm: 5
+  real_value: True
+
+test:
+  mae_thresh: null
+  mape_thresh: 0.0
+
+log:
+  log_step: 200
+  plot: False

config/EXP/SD.yaml

@@ -0,0 +1,44 @@
+data:
+  num_nodes: 716
+  lag: 12
+  horizon: 12
+  val_ratio: 0.2
+  test_ratio: 0.2
+  tod: False
+  normalizer: std
+  column_wise: False
+  default_graph: True
+  add_time_in_day: True
+  add_day_in_week: True
+  steps_per_day: 288
+  days_per_week: 7
+
+model:
+  batch_size: 64
+  input_dim: 1
+  output_dim: 1
+  in_len: 12
+
+train:
+  loss_func: mae
+  seed: 10
+  batch_size: 64
+  epochs: 300
+  lr_init: 0.003
+  weight_decay: 0
+  lr_decay: False
+  lr_decay_rate: 0.3
+  lr_decay_step: "5,20,40,70"
+  early_stop: True
+  early_stop_patience: 15
+  grad_norm: False
+  max_grad_norm: 5
+  real_value: True
+
+test:
+  mae_thresh: null
+  mape_thresh: 0.0
+
+log:
+  log_step: 2000
+  plot: False

config/EXPB/PEMSD4.yaml

@@ -0,0 +1,49 @@
+data:
+  num_nodes: 307
+  lag: 12
+  horizon: 12
+  val_ratio: 0.2
+  test_ratio: 0.2
+  tod: False
+  normalizer: std
+  column_wise: False
+  default_graph: True
+  add_time_in_day: True
+  add_day_in_week: True
+  steps_per_day: 288
+  days_per_week: 7
+
+model:
+  input_dim: 1
+  output_dim: 1
+  embed_dim: 10
+  rnn_units: 64
+  num_layers: 1
+  cheb_order: 2
+  patch_size: 3
+  use_day: True
+  use_week: True
+
+train:
+  loss_func: mae
+  seed: 10
+  batch_size: 64
+  epochs: 300
+  lr_init: 0.003
+  weight_decay: 0
+  lr_decay: False
+  lr_decay_rate: 0.3
+  lr_decay_step: "5,20,40,70"
+  early_stop: True
+  early_stop_patience: 15
+  grad_norm: False
+  max_grad_norm: 5
+  real_value: True
+
+test:
+  mae_thresh: null
+  mape_thresh: 0.0
+
+log:
+  log_step: 200
+  plot: False

config/STAEFormer/PEMSD3.yaml

@@ -0,0 +1,56 @@
+data:
+  num_nodes: 358
+  lag: 12
+  horizon: 12
+  val_ratio: 0.2
+  test_ratio: 0.2
+  tod: False
+  normalizer: std
+  column_wise: False
+  default_graph: True
+  add_time_in_day: True
+  add_day_in_week: True
+  steps_per_day: 288
+  days_per_week: 7
+
+model:
+  num_nodes: 358
+  in_steps: 12
+  out_steps: 12
+  steps_per_day: 288
+  input_dim: 1
+  output_dim: 1
+  input_embedding_dim: 24
+  tod_embedding_dim: 24
+  dow_embedding_dim: 24
+  spatial_embedding_dim: 0
+  adaptive_embedding_dim: 80
+  feed_forward_dim: 256
+  num_heads: 4
+  num_layers: 3
+  dropout: 0.1
+  use_mixed_proj: true
+
+train:
+  loss_func: mae
+  seed: 10
+  batch_size: 64
+  epochs: 300
+  lr_init: 0.003
+  weight_decay: 0
+  lr_decay: False
+  lr_decay_rate: 0.3
+  lr_decay_step: "5,20,40,70"
+  early_stop: True
+  early_stop_patience: 15
+  grad_norm: False
+  max_grad_norm: 5
+  real_value: True
+
+test:
+  mae_thresh: null
+  mape_thresh: 0.0
+
+log:
+  log_step: 200
+  plot: False

config/STAEFormer/PEMSD4.yaml

@@ -0,0 +1,55 @@
+data:
+  num_nodes: 307
+  lag: 12
+  horizon: 12
+  val_ratio: 0.1
+  test_ratio: 0.2
+  tod: False
+  normalizer: std
+  column_wise: False
+  default_graph: True
+  add_time_in_day: True
+  add_day_in_week: True
+  steps_per_day: 288
+  days_per_week: 7
+
+model:
+  num_nodes: 307
+  in_steps: 12
+  out_steps: 12
+  steps_per_day: 288
+  input_dim: 1
+  output_dim: 1
+  input_embedding_dim: 24
+  tod_embedding_dim: 24
+  dow_embedding_dim: 24
+  spatial_embedding_dim: 0
+  adaptive_embedding_dim: 80
+  feed_forward_dim: 256
+  num_heads: 4
+  num_layers: 3
+  dropout: 0.1
+  use_mixed_proj: true
+
+train:
+  loss_func: Huber
+  seed: 10
+  batch_size: 16
+  epochs: 200
+  lr_init: 0.001
+  weight_decay: 0.0003
+  lr_decay: True
+  lr_decay_rate: 0.1
+  lr_decay_step: "5,20,40,70"
+  early_stop: True
+  early_stop_patience: 30
+  grad_norm: False
+  real_value: True
+
+test:
+  mae_thresh: null
+  mape_thresh: 0.0
+
+log:
+  log_step: 2000
+  plot: False

config/STAEFormer/PEMSD7.yaml

@@ -0,0 +1,56 @@
+data:
+  num_nodes: 883
+  lag: 12
+  horizon: 12
+  val_ratio: 0.2
+  test_ratio: 0.2
+  tod: False
+  normalizer: std
+  column_wise: False
+  default_graph: True
+  add_time_in_day: True
+  add_day_in_week: True
+  steps_per_day: 288
+  days_per_week: 7
+
+model:
+  num_nodes: 883
+  in_steps: 12
+  out_steps: 12
+  steps_per_day: 288
+  input_dim: 1
+  output_dim: 1
+  input_embedding_dim: 24
+  tod_embedding_dim: 24
+  dow_embedding_dim: 24
+  spatial_embedding_dim: 0
+  adaptive_embedding_dim: 80
+  feed_forward_dim: 256
+  num_heads: 4
+  num_layers: 3
+  dropout: 0.1
+  use_mixed_proj: true
+
+train:
+  loss_func: mae
+  seed: 10
+  batch_size: 64
+  epochs: 300
+  lr_init: 0.003
+  weight_decay: 0
+  lr_decay: False
+  lr_decay_rate: 0.3
+  lr_decay_step: "5,20,40,70"
+  early_stop: True
+  early_stop_patience: 15
+  grad_norm: False
+  max_grad_norm: 5
+  real_value: True
+
+test:
+  mae_thresh: null
+  mape_thresh: 0.0
+
+log:
+  log_step: 200
+  plot: False

config/STAEFormer/PEMSD8.yaml

@@ -0,0 +1,56 @@
+data:
+  num_nodes: 170
+  lag: 12
+  horizon: 12
+  val_ratio: 0.2
+  test_ratio: 0.2
+  tod: False
+  normalizer: std
+  column_wise: False
+  default_graph: True
+  add_time_in_day: True
+  add_day_in_week: True
+  steps_per_day: 288
+  days_per_week: 7
+
+model:
+  num_nodes: 170
+  in_steps: 12
+  out_steps: 12
+  steps_per_day: 288
+  input_dim: 1
+  output_dim: 1
+  input_embedding_dim: 24
+  tod_embedding_dim: 24
+  dow_embedding_dim: 24
+  spatial_embedding_dim: 0
+  adaptive_embedding_dim: 80
+  feed_forward_dim: 256
+  num_heads: 4
+  num_layers: 3
+  dropout: 0.1
+  use_mixed_proj: true
+
+train:
+  loss_func: mae
+  seed: 10
+  batch_size: 64
+  epochs: 300
+  lr_init: 0.003
+  weight_decay: 0
+  lr_decay: False
+  lr_decay_rate: 0.3
+  lr_decay_step: "5,20,40,70"
+  early_stop: True
+  early_stop_patience: 15
+  grad_norm: False
+  max_grad_norm: 5
+  real_value: True
+
+test:
+  mae_thresh: null
+  mape_thresh: 0.0
+
+log:
+  log_step: 200
+  plot: False

config/STID/PEMSD4.yaml

@@ -0,0 +1,58 @@
+data:
+  num_nodes: 307
+  lag: 12
+  horizon: 12
+  val_ratio: 0.2
+  test_ratio: 0.2
+  tod: False
+  normalizer: std
+  column_wise: False
+  default_graph: True
+  add_time_in_day: True
+  add_day_in_week: True
+  steps_per_day: 288
+  days_per_week: 7
+
+model:
+  input_dim: 3
+  output_dim: 1
+  history: 12
+  horizon: 12
+  num_nodes: 307
+  input_len: 12
+  embed_dim": 32
+  output_len: 12
+  num_layer: 3
+  if_node: True
+  node_dim: 32
+  if_T_i_D: True
+  if_D_i_W: True
+  temp_dim_tid: 32
+  temp_dim_diw: 32
+  time_of_day_size: 288
+  day_of_week_size: 7
+
+
+train:
+  loss_func: mae
+  seed: 1
+  batch_size: 64
+  epochs: 300
+  lr_init: 0.002
+  weight_decay: 0.0001
+  lr_decay: False
+  lr_decay_rate: 0.3
+  lr_decay_step: "1,50,80"
+  early_stop: True
+  early_stop_patience: 15
+  grad_norm: False
+  max_grad_norm: 5
+  real_value: True
+
+test:
+  mae_thresh: null
+  mape_thresh: 0.0
+
+log:
+  log_step: 200
+  plot: False

config/TWDGCN/PEMSD8.yaml

@@ -18,7 +18,7 @@ log:
   plot: false
 model:
   cheb_order: 2
-  embed_dim: 5
+  embed_dim: 12
   input_dim: 1
   num_layers: 1
   output_dim: 1
@@ -29,10 +29,10 @@ test:
   mae_thresh: None
   mape_thresh: 0.001
 train:
-  batch_size: 64
+  batch_size: 12
   early_stop: true
-  early_stop_patience: 15
-  epochs: 100
+  early_stop_patience: 30
+  epochs: 200
   grad_norm: false
   loss_func: mae
   lr_decay: true
@@ -41,5 +41,5 @@ train:
   lr_init: 0.003
   max_grad_norm: 5
   real_value: true
-  seed: 12
+  seed: 3407
   weight_decay: 0

dataloader/EXPdataloader.py

@@ -0,0 +1,213 @@
+import numpy as np
+import gc
+import os
+import torch
+import h5py
+from lib.normalization import normalize_dataset
+
+
+def get_dataloader(args, normalizer='std', single=True):
+    # args should now include 'cycle'
+    data = load_st_dataset(args['type'], args['sample'])  # [T, N, F]
+    L, N, F = data.shape
+
+    # compute cycle index
+    cycle_arr = np.arange(L) % args['cycle']  # length-L array
+
+    # Step 1: sliding windows for X and Y
+    x = add_window_x(data, args['lag'], args['horizon'], single)
+    y = add_window_y(data, args['lag'], args['horizon'], single)
+    # window count = M = L - lag - horizon + 1
+    M = x.shape[0]
+
+    # Step 2: time features
+    time_in_day = np.tile(
+        np.array([i % args['steps_per_day'] / args['steps_per_day'] for i in range(L)]),
+        (N, 1)
+    ).T.reshape(L, N, 1)
+    day_in_week = np.tile(
+        np.array([(i // args['steps_per_day']) % args['days_per_week'] for i in range(L)]),
+        (N, 1)
+    ).T.reshape(L, N, 1)
+
+    x_day = add_window_x(time_in_day, args['lag'], args['horizon'], single)
+    x_week = add_window_x(day_in_week, args['lag'], args['horizon'], single)
+    x = np.concatenate([x, x_day, x_week], axis=-1)
+    # del x_day, x_week
+    # gc.collect()
+
+    # Step 3: extract cycle index per window: take value at end of sequence
+    cycle_win = np.array([cycle_arr[i + args['lag']] for i in range(M)])  # shape [M]
+
+    # Step 4: split into train/val/test
+    if args['test_ratio'] > 1:
+        x_train, x_val, x_test = split_data_by_days(x, args['val_ratio'], args['test_ratio'])
+        y_train, y_val, y_test = split_data_by_days(y, args['val_ratio'], args['test_ratio'])
+        c_train, c_val, c_test = split_data_by_days(cycle_win, args['val_ratio'], args['test_ratio'])
+    else:
+        x_train, x_val, x_test = split_data_by_ratio(x, args['val_ratio'], args['test_ratio'])
+        y_train, y_val, y_test = split_data_by_ratio(y, args['val_ratio'], args['test_ratio'])
+        c_train, c_val, c_test = split_data_by_ratio(cycle_win, args['val_ratio'], args['test_ratio'])
+    # del x, y, cycle_win
+    # gc.collect()
+
+    # Step 5: normalization on X only
+    scaler = normalize_dataset(x_train[..., :args['input_dim']], normalizer, args['column_wise'])
+    x_train[..., :args['input_dim']] = scaler.transform(x_train[..., :args['input_dim']])
+    x_val[..., :args['input_dim']] = scaler.transform(x_val[..., :args['input_dim']])
+    x_test[..., :args['input_dim']] = scaler.transform(x_test[..., :args['input_dim']])
+
+    # add time features to Y
+    y_day = add_window_y(time_in_day, args['lag'], args['horizon'], single)
+    y_week = add_window_y(day_in_week, args['lag'], args['horizon'], single)
+    y = np.concatenate([y, y_day, y_week], axis=-1)
+    # del y_day, y_week, time_in_day, day_in_week
+    # gc.collect()
+
+    # split Y time-augmented
+    if args['test_ratio'] > 1:
+        y_train, y_val, y_test = split_data_by_days(y, args['val_ratio'], args['test_ratio'])
+    else:
+        y_train, y_val, y_test = split_data_by_ratio(y, args['val_ratio'], args['test_ratio'])
+    # del y
+
+    # Step 6: create dataloaders including cycle index
+    train_loader = data_loader_with_cycle(x_train, y_train, c_train, args['batch_size'], shuffle=True, drop_last=True)
+    val_loader = data_loader_with_cycle(x_val, y_val, c_val, args['batch_size'], shuffle=False, drop_last=True)
+    test_loader = data_loader_with_cycle(x_test, y_test, c_test, args['batch_size'], shuffle=False, drop_last=False)
+
+    return train_loader, val_loader, test_loader, scaler
+
+
+def data_loader_with_cycle(X, Y, C, batch_size, shuffle=True, drop_last=True):
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    X_t = torch.tensor(X, dtype=torch.float32, device=device)
+    Y_t = torch.tensor(Y, dtype=torch.float32, device=device)
+    C_t = torch.tensor(C, dtype=torch.long, device=device).unsqueeze(-1)  # [B,1]
+    dataset = torch.utils.data.TensorDataset(X_t, Y_t, C_t)
+    loader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=shuffle, drop_last=drop_last)
+    return loader
+
+# Rest of the helper functions (load_st_dataset, split_data..., add_window_x/y) unchanged
+
+
+def load_st_dataset(dataset, sample):
+    # output B, N, D
+    match dataset:
+        case 'PEMSD3':
+            data_path = os.path.join('./data/PEMS03/PEMS03.npz')
+            data = np.load(data_path)['data'][:, :, 0]  # only the first dimension, traffic flow data
+        case 'PEMSD4':
+            data_path = os.path.join('./data/PEMS04/PEMS04.npz')
+            data = np.load(data_path)['data'][:, :, 0]  # only the first dimension, traffic flow data
+        case 'PEMSD7':
+            data_path = os.path.join('./data/PEMS07/PEMS07.npz')
+            data = np.load(data_path)['data'][:, :, 0]  # only the first dimension, traffic flow data
+        case 'PEMSD8':
+            data_path = os.path.join('./data/PEMS08/PEMS08.npz')
+            data = np.load(data_path)['data'][:, :, 0]  # only the first dimension, traffic flow data
+        case 'PEMSD7(L)':
+            data_path = os.path.join('./data/PEMS07(L)/PEMS07L.npz')
+            data = np.load(data_path)['data'][:, :, 0]  # only the first dimension, traffic flow data
+        case 'PEMSD7(M)':
+            data_path = os.path.join('./data/PEMS07(M)/V_228.csv')
+            data = np.genfromtxt(data_path, delimiter=',')  # Read CSV directly with numpy
+        case 'METR-LA':
+            data_path = os.path.join('./data/METR-LA/METR.h5')
+            with h5py.File(data_path, 'r') as f:  # Use h5py to handle HDF5 files without pandas
+                data = np.array(f['data'])
+        case 'BJ':
+            data_path = os.path.join('./data/BJ/BJ500.csv')
+            data = np.genfromtxt(data_path, delimiter=',', skip_header=1)  # Skip header if present
+        case 'Hainan':
+            data_path = os.path.join('./data/Hainan/Hainan.npz')
+            data = np.load(data_path)['data'][:, :, 0]
+        case 'SD':
+            data_path = os.path.join('./data/SD/data.npz')
+            data = np.load(data_path)["data"][:, :, 0].astype(np.float32)
+        case _:
+            raise ValueError(f"Unsupported dataset: {dataset}")
+
+    # Ensure data shape compatibility
+    if len(data.shape) == 2:
+        data = np.expand_dims(data, axis=-1)
+
+    print('加载 %s 数据集中... ' % dataset)
+    return data[::sample]
+
+def split_data_by_days(data, val_days, test_days, interval=30):
+    t = int((24 * 60) / interval)
+    test_data = data[-t * int(test_days):]
+    val_data = data[-t * int(test_days + val_days):-t * int(test_days)]
+    train_data = data[:-t * int(test_days + val_days)]
+    return train_data, val_data, test_data
+
+
+def split_data_by_ratio(data, val_ratio, test_ratio):
+    data_len = data.shape[0]
+    test_data = data[-int(data_len * test_ratio):]
+    val_data = data[-int(data_len * (test_ratio + val_ratio)):-int(data_len * test_ratio)]
+    train_data = data[:-int(data_len * (test_ratio + val_ratio))]
+    return train_data, val_data, test_data
+
+
+def data_loader(X, Y, batch_size, shuffle=True, drop_last=True):
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    X = torch.tensor(X, dtype=torch.float32, device=device)
+    Y = torch.tensor(Y, dtype=torch.float32, device=device)
+    data = torch.utils.data.TensorDataset(X, Y)
+    dataloader = torch.utils.data.DataLoader(data, batch_size=batch_size,
+                                             shuffle=shuffle, drop_last=drop_last)
+    return dataloader
+
+
+def add_window_x(data, window=3, horizon=1, single=False):
+    """
+    Generate windowed X values from the input data.
+
+    :param data: Input data, shape [B, ...]
+    :param window: Size of the sliding window
+    :param horizon: Horizon size
+    :param single: If True, generate single-step windows, else multi-step
+    :return: X with shape [B, W, ...]
+    """
+    length = len(data)
+    end_index = length - horizon - window + 1
+    x = []  # Sliding windows
+    index = 0
+
+    while index < end_index:
+        x.append(data[index:index + window])
+        index += 1
+
+    return np.array(x)
+
+
+def add_window_y(data, window=3, horizon=1, single=False):
+    """
+    Generate windowed Y values from the input data.
+
+    :param data: Input data, shape [B, ...]
+    :param window: Size of the sliding window
+    :param horizon: Horizon size
+    :param single: If True, generate single-step windows, else multi-step
+    :return: Y with shape [B, H, ...]
+    """
+    length = len(data)
+    end_index = length - horizon - window + 1
+    y = []  # Horizon values
+    index = 0
+
+    while index < end_index:
+        if single:
+            y.append(data[index + window + horizon - 1:index + window + horizon])
+        else:
+            y.append(data[index + window:index + window + horizon])
+        index += 1
+
+    return np.array(y)
+
+
+if __name__ == '__main__':
+    res = load_st_dataset('SD', 1)
+    k = 1

dataloader/PeMSDdataloader.py

@@ -121,6 +121,9 @@ def load_st_dataset(dataset, sample):
         case 'Hainan':
             data_path = os.path.join('./data/Hainan/Hainan.npz')
             data = np.load(data_path)['data'][:, :, 0]
+        case 'SD':
+            data_path = os.path.join('./data/SD/data.npz')
+            data = np.load(data_path)["data"][:, :, 0].astype(np.float32)
         case _:
             raise ValueError(f"Unsupported dataset: {dataset}")
 
@@ -204,3 +207,6 @@ def add_window_y(data, window=3, horizon=1, single=False):
     return np.array(y)
 
 
+if __name__ == '__main__':
+    res = load_st_dataset('SD', 1)
+    k = 1

dataloader/loader_selector.py

@@ -1,11 +1,13 @@
 from dataloader.cde_loader.cdeDataloader import get_dataloader as cde_loader
 from dataloader.PeMSDdataloader import get_dataloader as normal_loader
 from dataloader.DCRNNdataloader import get_dataloader as DCRNN_loader
+from dataloader.EXPdataloader import get_dataloader as EXP_loader
 
 def get_dataloader(config, normalizer, single):
     match config['model']['type']:
         case 'STGNCDE': return cde_loader(config['data'], normalizer, single)
         case 'DCRNN': return DCRNN_loader(config['data'], normalizer, single)
+        case 'EXP': return EXP_loader(config['data'], normalizer, single)
         case _: return normal_loader(config['data'], normalizer, single)
 
 

lib/LargeST.py

@@ -0,0 +1,267 @@
+import pickle
+import torch
+import numpy as np
+import os
+import gc
+# ! X shape: (B, T, N, C)
+
+def load_pkl(pickle_file: str) -> object:
+    """
+    Load data from a pickle file.
+
+    Args:
+        pickle_file (str): Path to the pickle file.
+
+    Returns:
+        object: Loaded object from the pickle file.
+    """
+
+    try:
+        with open(pickle_file, "rb") as f:
+            pickle_data = pickle.load(f)
+    except UnicodeDecodeError:
+        with open(pickle_file, "rb") as f:
+            pickle_data = pickle.load(f, encoding="latin1")
+    except Exception as e:
+        print(f"Unable to load data from {pickle_file}: {e}")
+        raise
+    return pickle_data
+
+def get_dataloaders_from_index_data(
+    data_dir, tod=False, dow=False, batch_size=64, log=None, train_size=0.6
+):
+    data = np.load(os.path.join(data_dir, "data.npz"))["data"].astype(np.float32)
+
+    features = [0]
+    if tod:
+        features.append(1)
+    if dow:
+        features.append(2)
+    # if dom:
+    #     features.append(3)
+    data = data[..., features]
+
+    index = np.load(os.path.join(data_dir, "index.npz"))
+
+    train_index = index["train"]  # (num_samples, 3)
+    val_index = index["val"]
+    test_index = index["test"]
+
+    x_train_index = vrange(train_index[:, 0], train_index[:, 1])
+    y_train_index = vrange(train_index[:, 1], train_index[:, 2])
+    x_val_index = vrange(val_index[:, 0], val_index[:, 1])
+    y_val_index = vrange(val_index[:, 1], val_index[:, 2])
+    x_test_index = vrange(test_index[:, 0], test_index[:, 1])
+    y_test_index = vrange(test_index[:, 1], test_index[:, 2])
+
+    x_train = data[x_train_index]
+    y_train = data[y_train_index][..., :1]
+    x_val = data[x_val_index]
+    y_val = data[y_val_index][..., :1]
+    x_test = data[x_test_index]
+    y_test = data[y_test_index][..., :1]
+
+    scaler = StandardScaler(mean=x_train[..., 0].mean(), std=x_train[..., 0].std())
+
+    x_train[..., 0] = scaler.transform(x_train[..., 0])
+    x_val[..., 0] = scaler.transform(x_val[..., 0])
+    x_test[..., 0] = scaler.transform(x_test[..., 0])
+
+    print_log(f"Trainset:\tx-{x_train.shape}\ty-{y_train.shape}", log=log)
+    print_log(f"Valset:  \tx-{x_val.shape}  \ty-{y_val.shape}", log=log)
+    print_log(f"Testset:\tx-{x_test.shape}\ty-{y_test.shape}", log=log)
+
+    trainset = torch.utils.data.TensorDataset(
+        torch.FloatTensor(x_train), torch.FloatTensor(y_train)
+    )
+    valset = torch.utils.data.TensorDataset(
+        torch.FloatTensor(x_val), torch.FloatTensor(y_val)
+    )
+    testset = torch.utils.data.TensorDataset(
+        torch.FloatTensor(x_test), torch.FloatTensor(y_test)
+    )
+    if train_size != 0.6:
+        drop_last=True
+    else:
+        drop_last=False
+    trainset_loader = torch.utils.data.DataLoader(
+        trainset, batch_size=batch_size, shuffle=True, drop_last=drop_last
+    )
+    valset_loader = torch.utils.data.DataLoader(
+        valset, batch_size=batch_size, shuffle=False, drop_last=drop_last
+    )
+    testset_loader = torch.utils.data.DataLoader(
+        testset, batch_size=batch_size, shuffle=False, drop_last=drop_last
+    )
+
+    return trainset_loader, valset_loader, testset_loader, scaler
+
+def get_dataloaders_from_index_data_MTS(
+    data_dir,
+    in_steps=12,
+    out_steps=12,
+    tod=False,
+    dow=False,
+    y_tod=False,
+    y_dow=False,
+    batch_size=64,
+    log=None,
+):
+    data = np.load(os.path.join(data_dir, f"data.npz"))["data"].astype(np.float32)
+    index = np.load(os.path.join(data_dir, f"index_{in_steps}_{out_steps}.npz"))
+
+    x_features = [0]
+    if tod:
+        x_features.append(1)
+    if dow:
+        x_features.append(2)
+
+    y_features = [0]
+    if y_tod:
+        y_features.append(1)
+    if y_dow:
+        y_features.append(2)
+
+    train_index = index["train"]  # (num_samples, 3)
+    val_index = index["val"]
+    test_index = index["test"]
+
+    # Parallel
+    # x_train_index = vrange(train_index[:, 0], train_index[:, 1])
+    # y_train_index = vrange(train_index[:, 1], train_index[:, 2])
+    # x_val_index = vrange(val_index[:, 0], val_index[:, 1])
+    # y_val_index = vrange(val_index[:, 1], val_index[:, 2])
+    # x_test_index = vrange(test_index[:, 0], test_index[:, 1])
+    # y_test_index = vrange(test_index[:, 1], test_index[:, 2])
+
+    # x_train = data[x_train_index][..., x_features]
+    # y_train = data[y_train_index][..., y_features]
+    # x_val = data[x_val_index][..., x_features]
+    # y_val = data[y_val_index][..., y_features]
+    # x_test = data[x_test_index][..., x_features]
+    # y_test = data[y_test_index][..., y_features]
+
+    # Iterative
+    x_train = np.stack([data[idx[0] : idx[1]] for idx in train_index])[..., x_features]
+    y_train = np.stack([data[idx[1] : idx[2]] for idx in train_index])[..., y_features]
+    x_val = np.stack([data[idx[0] : idx[1]] for idx in val_index])[..., x_features]
+    y_val = np.stack([data[idx[1] : idx[2]] for idx in val_index])[..., y_features]
+    x_test = np.stack([data[idx[0] : idx[1]] for idx in test_index])[..., x_features]
+    y_test = np.stack([data[idx[1] : idx[2]] for idx in test_index])[..., y_features]
+
+    scaler = StandardScaler(mean=x_train[..., 0].mean(), std=x_train[..., 0].std())
+
+    x_train[..., 0] = scaler.transform(x_train[..., 0])
+    x_val[..., 0] = scaler.transform(x_val[..., 0])
+    x_test[..., 0] = scaler.transform(x_test[..., 0])
+
+    print_log(f"Trainset:\tx-{x_train.shape}\ty-{y_train.shape}", log=log)
+    print_log(f"Valset:  \tx-{x_val.shape}  \ty-{y_val.shape}", log=log)
+    print_log(f"Testset:\tx-{x_test.shape}\ty-{y_test.shape}", log=log)
+
+    trainset = torch.utils.data.TensorDataset(
+        torch.FloatTensor(x_train), torch.FloatTensor(y_train)
+    )
+    valset = torch.utils.data.TensorDataset(
+        torch.FloatTensor(x_val), torch.FloatTensor(y_val)
+    )
+    testset = torch.utils.data.TensorDataset(
+        torch.FloatTensor(x_test), torch.FloatTensor(y_test)
+    )
+
+    trainset_loader = torch.utils.data.DataLoader(
+        trainset, batch_size=batch_size, shuffle=True
+    )
+    valset_loader = torch.utils.data.DataLoader(
+        valset, batch_size=batch_size, shuffle=False
+    )
+    testset_loader = torch.utils.data.DataLoader(
+        testset, batch_size=batch_size, shuffle=False
+    )
+
+    return trainset_loader, valset_loader, testset_loader, scaler
+
+def get_dataloaders_from_index_data_Test(
+    data_dir,
+    in_steps=12,
+    out_steps=12,
+    tod=False,
+    dow=False,
+    y_tod=False,
+    y_dow=False,
+    batch_size=64,
+    log=None,
+):
+    data = np.load(os.path.join(data_dir, f"data.npz"))["data"].astype(np.float32)
+    index = np.load(os.path.join(data_dir, f"index_{in_steps}_{out_steps}.npz"))
+
+    x_features = [0]
+    if tod:
+        x_features.append(1)
+    if dow:
+        x_features.append(2)
+
+    y_features = [0]
+    if y_tod:
+        y_features.append(1)
+    if y_dow:
+        y_features.append(2)
+
+    train_index = index["train"]  # (num_samples, 3)
+    # val_index = index["val"]
+    test_index = index["test"]
+
+    # Parallel
+    # x_train_index = vrange(train_index[:, 0], train_index[:, 1])
+    # y_train_index = vrange(train_index[:, 1], train_index[:, 2])
+    # x_val_index = vrange(val_index[:, 0], val_index[:, 1])
+    # y_val_index = vrange(val_index[:, 1], val_index[:, 2])
+    # x_test_index = vrange(test_index[:, 0], test_index[:, 1])
+    # y_test_index = vrange(test_index[:, 1], test_index[:, 2])
+
+    # x_train = data[x_train_index][..., x_features]
+    # y_train = data[y_train_index][..., y_features]
+    # x_val = data[x_val_index][..., x_features]
+    # y_val = data[y_val_index][..., y_features]
+    # x_test = data[x_test_index][..., x_features]
+    # y_test = data[y_test_index][..., y_features]
+
+    # Iterative
+    x_train = np.stack([data[idx[0] : idx[1]] for idx in train_index])[..., x_features]
+    # y_train = np.stack([data[idx[1] : idx[2]] for idx in train_index])[..., y_features]
+    # x_val = np.stack([data[idx[0] : idx[1]] for idx in val_index])[..., x_features]
+    # y_val = np.stack([data[idx[1] : idx[2]] for idx in val_index])[..., y_features]
+    x_test = np.stack([data[idx[0] : idx[1]] for idx in test_index])[..., x_features]
+    y_test = np.stack([data[idx[1] : idx[2]] for idx in test_index])[..., y_features]
+
+    scaler = StandardScaler(mean=x_train[..., 0].mean(), std=x_train[..., 0].std())
+
+    # x_train[..., 0] = scaler.transform(x_train[..., 0])
+    # x_val[..., 0] = scaler.transform(x_val[..., 0])
+    x_test[..., 0] = scaler.transform(x_test[..., 0])
+
+    # print_log(f"Trainset:\tx-{x_train.shape}\ty-{y_train.shape}", log=log)
+    # print_log(f"Valset:  \tx-{x_val.shape}  \ty-{y_val.shape}", log=log)
+    print_log(f"Testset:\tx-{x_test.shape}\ty-{y_test.shape}", log=log)
+
+    # trainset = torch.utils.data.TensorDataset(
+    #     torch.FloatTensor(x_train), torch.FloatTensor(y_train)
+    # )
+    # valset = torch.utils.data.TensorDataset(
+    #     torch.FloatTensor(x_val), torch.FloatTensor(y_val)
+    # )
+    testset = torch.utils.data.TensorDataset(
+        torch.FloatTensor(x_test), torch.FloatTensor(y_test)
+    )
+
+    # trainset_loader = torch.utils.data.DataLoader(
+    #     trainset, batch_size=batch_size, shuffle=True
+    # )
+    # valset_loader = torch.utils.data.DataLoader(
+    #     valset, batch_size=batch_size, shuffle=False
+    # )
+    testset_loader = torch.utils.data.DataLoader(
+        testset, batch_size=batch_size, shuffle=False
+    )
+
+    return testset_loader, scaler

lib/initializer.py

@@ -12,6 +12,8 @@ def init_model(args, device):
             nn.init.xavier_uniform_(p)
         else:
             nn.init.uniform_(p)
+    total_params = sum(p.numel() for p in model.parameters())
+    print(f"Model has {total_params} parameters")
     return model
 
 def init_optimizer(model, args):