TrafficWheel/dataloader/TSloader.py

from dataloader.data_selector import load_st_dataset
from utils.normalization import normalize_dataset

import numpy as np
import torch


def get_dataloader(args, normalizer="std", single=True):
    data = load_st_dataset(args)
    # data = data[..., 0:1]

    args = args["data"]
    L, N, F = data.shape
    # data = data.reshape(L, N*F)  # [L, N*F]

    # Generate sliding windows for main data and add time features
    x, y = _prepare_data_with_windows(data, args, single)
    
    # Split data [b,t,n,c]
    split_fn = split_data_by_days if args["test_ratio"] > 1 else split_data_by_ratio
    x_train, x_val, x_test = split_fn(x, args["val_ratio"], args["test_ratio"])
    y_train, y_val, y_test = split_fn(y, args["val_ratio"], args["test_ratio"])

    # Normalize x and y using the same scaler
    scaler = _normalize_data(x_train, x_val, x_test, args, normalizer)
    _apply_existing_scaler(y_train, y_val, y_test, scaler, args)

    # reshape [b,t,n,c] -> [b*n, t, c]
    x_train, x_val, x_test, y_train, y_val, y_test = \
        _reshape_tensor(x_train, x_val, x_test, y_train, y_val, y_test)

    # Create dataloaders
    return (
        _create_dataloader(x_train, y_train, args["batch_size"], True, False),
        _create_dataloader(x_val, y_val, args["batch_size"], False, False),
        _create_dataloader(x_test, y_test, args["batch_size"], False, False),
        scaler
    )

def _reshape_tensor(*tensors):
    """Reshape tensors from [b, t, n, c] -> [b*n, t, c]."""
    reshaped = []
    for x in tensors:
        # x 是 ndarray：shape (b, t, n, c)
        b, t, n, c = x.shape
        x_new = x.transpose(0, 2, 1, 3).reshape(b * n, t, c)
        reshaped.append(x_new)
    return reshaped

def _prepare_data_with_windows(data, args, single):
    # Generate sliding windows for main data
    x = add_window_x(data, args["lag"], args["horizon"], single)
    y = add_window_y(data, args["lag"], args["horizon"], single)
    return x, y

def _normalize_data(train_data, val_data, test_data, args, normalizer):
    scaler = normalize_dataset(train_data[..., : args["num_nodes"]], normalizer, args["column_wise"])
    
    for data in [train_data, val_data, test_data]:
        data[..., : args["num_nodes"]] = scaler.transform(data[..., : args["num_nodes"]])
    
    return scaler


def _apply_existing_scaler(train_data, val_data, test_data, scaler, args):
    for data in [train_data, val_data, test_data]:
        data[..., : args["num_nodes"]] = scaler.transform(data[..., : args["num_nodes"]])


def _create_dataloader(X_data, Y_data, batch_size, shuffle, drop_last):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    X_tensor = torch.tensor(X_data, dtype=torch.float32, device=device)
    Y_tensor = torch.tensor(Y_data, dtype=torch.float32, device=device)
    dataset = torch.utils.data.TensorDataset(X_tensor, Y_tensor)
    return torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=shuffle, drop_last=drop_last)


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 _generate_windows(data, window=3, horizon=1, offset=0):
    """
    Internal helper function to generate sliding windows.
    
    :param data: Input data, shape [L, T, C]
    :param window: Window size
    :param horizon: Horizon size
    :param offset: Offset from window start
    :return: Windowed data, shape [num_windows, window, T, C]
    """
    length = len(data)
    end_index = length - horizon - window + 1
    windows = []
    index = 0

    if end_index <= 0:
        raise ValueError(f"end_index is non-positive: {end_index}, length={length}, horizon={horizon}, window={window}")
    
    while index < end_index:
        window_data = data[index + offset : index + offset + window]
        windows.append(window_data)
        index += 1

    if not windows:
        raise ValueError("No windows generated")
    
    # Check window shapes
    first_shape = windows[0].shape
    for i, w in enumerate(windows):
        if w.shape != first_shape:
            raise ValueError(f"Window {i} has shape {w.shape}, expected {first_shape}")
    
    return np.array(windows)

def add_window_x(data, window=3, horizon=1, single=False):
    """
    Generate windowed X values from the input data.

    :param data: Input data, shape [L, T, C]
    :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 [num_windows, window, T, C]
    """
    return _generate_windows(data, window, horizon, offset=0)

def add_window_y(data, window=3, horizon=1, single=False):
    """
    Generate windowed Y values from the input data.

    :param data: Input data, shape [L, T, C]
    :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 [num_windows, horizon, T, C]
    """
    offset = window if not single else window + horizon - 1
    return _generate_windows(data, window=1 if single else horizon, horizon=horizon, offset=offset)

if __name__ == "__main__":
    
    # Test with a dummy config using METR-LA dataset
    dummy_args = {
        "basic": {
            "dataset": "METR-LA"
        },
        "data": {
            "lag": 3,
            "horizon": 1,
            "val_ratio": 0.1,
            "test_ratio": 0.2,
            "steps_per_day": 288,
            "days_per_week": 7,
            "input_dim": 1,
            "column_wise": False,
            "batch_size": 32,
            "time_dim": 1  # Add time dimension parameter
        }
    }
    
    try:
        # Load data
        data = load_st_dataset(dummy_args)
        print(f"Original data shape: {data.shape}")
        
        # Get dataloader
        train_loader, val_loader, test_loader, scaler = get_dataloader(dummy_args)
        
        # Test data loader
        for batch_x, batch_y in train_loader:
            print(f"Batch X shape: {batch_x.shape}")
            print(f"Batch Y shape: {batch_y.shape}")
            break
        
        print("Test passed successfully!")
        
    except Exception as e:
        print(f"Test failed with error: {e}")
        import traceback
        traceback.print_exc()