267 lines
9.2 KiB
Python
Executable File
267 lines
9.2 KiB
Python
Executable File
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
|