329 lines
13 KiB
Python
329 lines
13 KiB
Python
import pandas as pd
|
||
import numpy as np
|
||
from sklearn.preprocessing import StandardScaler
|
||
import joblib
|
||
from utils.timefeatures import time_features
|
||
import os
|
||
|
||
def get_ett_dataset_borders(dataset_name, data_len, input_len):
|
||
"""
|
||
ETT系列数据集的特定边界处理函数
|
||
|
||
Args:
|
||
dataset_name (str): 数据集名称（如 'ETTm1', 'ETTh1'）
|
||
data_len (int): 数据总长度
|
||
input_len (int): 输入序列长度
|
||
|
||
Returns:
|
||
tuple: (border1s, border2s) 边界点列表
|
||
"""
|
||
if dataset_name.startswith('ETTm'):
|
||
# ETTm1, ETTm2: 15分钟间隔，每天96个点
|
||
border1s = [0, 12 * 30 * 96 - input_len, 12 * 30 * 96 + 4 * 30 * 96 - input_len]
|
||
border2s = [12 * 30 * 96, 12 * 30 * 96 + 4 * 30 * 96, 12 * 30 * 96 + 8 * 30 * 96]
|
||
elif dataset_name.startswith('ETTh'):
|
||
# ETTh1, ETTh2: 小时间隔，每天24个点
|
||
border1s = [0, 12 * 30 * 24 - input_len, 12 * 30 * 24 + 4 * 30 * 24 - input_len]
|
||
border2s = [12 * 30 * 24, 12 * 30 * 24 + 4 * 30 * 24, 12 * 30 * 24 + 8 * 30 * 24]
|
||
else:
|
||
raise ValueError(f"Unknown ETT dataset: {dataset_name}")
|
||
|
||
return border1s, border2s
|
||
|
||
# 示例：可以添加其他特定数据集的处理函数
|
||
# def get_weather_dataset_borders(dataset_name, data_len, input_len):
|
||
# """
|
||
# Weather数据集的特定边界处理函数
|
||
# """
|
||
# # 假设weather数据集使用不同的分割策略
|
||
# # 比如：前80%训练，中间10%验证，后10%测试
|
||
# train_end = int(data_len * 0.8)
|
||
# val_end = int(data_len * 0.9)
|
||
#
|
||
# border1s = [0, train_end - input_len, val_end - input_len]
|
||
# border2s = [train_end, val_end, data_len]
|
||
#
|
||
# return border1s, border2s
|
||
|
||
# 数据集处理函数映射表
|
||
DATASET_HANDLERS = {
|
||
'ETTm1': get_ett_dataset_borders,
|
||
'ETTm2': get_ett_dataset_borders,
|
||
'ETTh1': get_ett_dataset_borders,
|
||
'ETTh2': get_ett_dataset_borders,
|
||
# 可以在这里添加更多数据集的处理函数
|
||
# 'weather': get_weather_dataset_borders,
|
||
}
|
||
|
||
def preprocess_time_series(
|
||
csv_data,
|
||
input_len,
|
||
pred_len,
|
||
slide_step,
|
||
dataset_name=None, # 新增：数据集名称参数
|
||
data_path_name='ETTm1.csv', # 保留向后兼容，但优先使用dataset_name
|
||
selected_columns=None,
|
||
date_column='date',
|
||
freq='h', # 按照分析，原文 ETTm1/ETTh1 实验均使用 'h'
|
||
split_method='auto', # 'auto', 'specific', 'ratio'
|
||
train_ratio=0.7,
|
||
val_ratio=0.1,
|
||
test_ratio=0.2,
|
||
has_time_column=True, # 新增：是否包含时间列
|
||
):
|
||
"""
|
||
修改版：根据 TimesNet 原文逻辑预处理时序数据。
|
||
1. 支持三种分割方法：auto（自动选择）、specific（特定数据集）、ratio（比例分割）
|
||
2. 支持基于数据集名称的特定处理函数调用
|
||
3. 滑动窗口的目标 y 长度为 pred_len (按用户要求)。
|
||
4. 支持无时间列的数据集处理
|
||
|
||
Args:
|
||
csv_data (pd.DataFrame or str): CSV data as DataFrame or path to CSV file
|
||
input_len (int): Length of input sequence (seq_len in original paper)
|
||
pred_len (int): Length of prediction sequence
|
||
slide_step (int): Step size for sliding window
|
||
dataset_name (str): 数据集名称（如 'ETTm1', 'weather'），优先使用此参数
|
||
data_path_name (str): 数据文件名（如 'ETTm1.csv'），向后兼容用
|
||
selected_columns (list): List of column names to use (default: None, uses all).
|
||
date_column (str): Name of the date column (default: 'date').
|
||
freq (str): Frequency for time features ('h' for hourly, 't' for minutely).
|
||
split_method (str): Data split method - 'auto', 'specific', or 'ratio'
|
||
- 'auto': automatically choose based on dataset_name
|
||
- 'specific': use dataset-specific split function
|
||
- 'ratio': use ratio-based split
|
||
train_ratio (float): Training set ratio (only used when split_method='ratio')
|
||
val_ratio (float): Validation set ratio (only used when split_method='ratio')
|
||
test_ratio (float): Test set ratio (only used when split_method='ratio')
|
||
has_time_column (bool): Whether the dataset has a time column
|
||
|
||
Returns:
|
||
dict: Dictionary containing processed data.
|
||
"""
|
||
# 1. 加载数据
|
||
if isinstance(csv_data, str):
|
||
try:
|
||
data = pd.read_csv(csv_data)
|
||
except FileNotFoundError:
|
||
raise FileNotFoundError(f"CSV file not found: {csv_data}")
|
||
except Exception as e:
|
||
raise Exception(f"Error loading CSV file: {e}")
|
||
else:
|
||
data = csv_data.copy()
|
||
|
||
# 2. 提取时间特征（仅在有时间列时）
|
||
if has_time_column and date_column in data.columns:
|
||
date_index = pd.to_datetime(data[date_column])
|
||
if isinstance(date_index, pd.Series):
|
||
date_index = pd.DatetimeIndex(date_index)
|
||
time_stamp = time_features(date_index, freq=freq)
|
||
time_stamp = time_stamp.transpose(1, 0) # Shape: (n_samples, n_time_features)
|
||
elif has_time_column:
|
||
raise ValueError(f"Date column '{date_column}' not found in data")
|
||
else:
|
||
# 没有时间列，创建空的时间戳数组
|
||
time_stamp = None
|
||
|
||
# 3. 选择数据列
|
||
if selected_columns is not None:
|
||
data = data[selected_columns]
|
||
else:
|
||
if has_time_column:
|
||
feature_columns = [col for col in data.columns if col != date_column]
|
||
else:
|
||
feature_columns = list(data.columns)
|
||
data = data[feature_columns]
|
||
|
||
# 4. 【核心修改】根据split_method选择数据集分割方式
|
||
# 确定使用的数据集名称
|
||
if dataset_name is None:
|
||
# 向后兼容：从文件路径提取数据集名称
|
||
dataset_name = os.path.splitext(data_path_name)[0]
|
||
|
||
if split_method == 'auto':
|
||
# 自动选择：特定数据集用specific，其他用ratio
|
||
if dataset_name in DATASET_HANDLERS:
|
||
split_method = 'specific'
|
||
else:
|
||
split_method = 'ratio'
|
||
|
||
if split_method == 'specific':
|
||
# 使用特定数据集的处理函数
|
||
if dataset_name in DATASET_HANDLERS:
|
||
handler_func = DATASET_HANDLERS[dataset_name]
|
||
border1s, border2s = handler_func(dataset_name, len(data), input_len)
|
||
print(f"Using specific split for dataset '{dataset_name}'")
|
||
else:
|
||
print(f"Warning: No specific handler for dataset '{dataset_name}'. Falling back to ratio split.")
|
||
split_method = 'ratio'
|
||
|
||
if split_method == 'ratio':
|
||
# 使用比例分割数据集
|
||
# 验证比例和为1
|
||
if abs(train_ratio + val_ratio + test_ratio - 1.0) > 1e-6:
|
||
raise ValueError(f"Ratios must sum to 1.0, got {train_ratio + val_ratio + test_ratio}")
|
||
|
||
total_len = len(data)
|
||
num_train = int(total_len * train_ratio)
|
||
num_val = int(total_len * val_ratio)
|
||
num_test = total_len - num_train - num_val # 确保所有数据都被使用
|
||
|
||
border1s = [0, num_train - input_len, num_train + num_val - input_len]
|
||
border2s = [num_train, num_train + num_val, total_len]
|
||
|
||
print(f"Using ratio split for dataset '{dataset_name}': train={train_ratio:.1%}, val={val_ratio:.1%}, test={test_ratio:.1%}")
|
||
print(f"Data points: train={num_train}, val={num_val}, test={num_test}")
|
||
|
||
train_data = data.iloc[border1s[0]:border2s[0]].values
|
||
val_data = data.iloc[border1s[1]:border2s[1]].values
|
||
test_data = data.iloc[border1s[2]:border2s[2]].values
|
||
|
||
# 处理时间戳（仅在有时间列时）
|
||
if time_stamp is not None:
|
||
train_time_stamp = time_stamp[border1s[0]:border2s[0]]
|
||
val_time_stamp = time_stamp[border1s[1]:border2s[1]]
|
||
test_time_stamp = time_stamp[border1s[2]:border2s[2]]
|
||
else:
|
||
train_time_stamp = None
|
||
val_time_stamp = None
|
||
test_time_stamp = None
|
||
|
||
# 5. 归一化 (Fit on training data only)
|
||
scaler = StandardScaler()
|
||
scaler.fit(train_data)
|
||
|
||
train_data_scaled = scaler.transform(train_data)
|
||
val_data_scaled = scaler.transform(val_data)
|
||
test_data_scaled = scaler.transform(test_data)
|
||
|
||
# 6. 【核心修改】使用您的滑窗逻辑创建样本
|
||
train_x, train_y = create_sliding_windows(train_data_scaled, input_len, pred_len, slide_step)
|
||
val_x, val_y = create_sliding_windows(val_data_scaled, input_len, pred_len, slide_step)
|
||
test_x, test_y = create_sliding_windows(test_data_scaled, input_len, pred_len, slide_step)
|
||
|
||
# 处理时间标记（仅在有时间列时）
|
||
if train_time_stamp is not None:
|
||
train_x_mark, train_y_mark = create_sliding_windows(train_time_stamp, input_len, pred_len, slide_step)
|
||
val_x_mark, val_y_mark = create_sliding_windows(val_time_stamp, input_len, pred_len, slide_step)
|
||
test_x_mark, test_y_mark = create_sliding_windows(test_time_stamp, input_len, pred_len, slide_step)
|
||
else:
|
||
train_x_mark, train_y_mark = None, None
|
||
val_x_mark, val_y_mark = None, None
|
||
test_x_mark, test_y_mark = None, None
|
||
|
||
return {
|
||
'train_x': train_x, 'train_y': train_y,
|
||
'train_x_mark': train_x_mark, 'train_y_mark': train_y_mark,
|
||
'val_x': val_x, 'val_y': val_y,
|
||
'val_x_mark': val_x_mark, 'val_y_mark': val_y_mark,
|
||
'test_x': test_x, 'test_y': test_y,
|
||
'test_x_mark': test_x_mark, 'test_y_mark': test_y_mark,
|
||
'scaler': scaler
|
||
}
|
||
|
||
|
||
def create_sliding_windows(data, input_len, pred_len, slide_step):
|
||
"""
|
||
Create sliding windows from time series data.
|
||
Target `y` has length `pred_len`.
|
||
|
||
Args:
|
||
data (np.ndarray): Time series data (features or time marks)
|
||
input_len (int): Length of input sequence
|
||
pred_len (int): Length of prediction sequence
|
||
slide_step (int): Step size for sliding window
|
||
|
||
Returns:
|
||
tuple: (X, y) where X is input sequences and y is target sequences
|
||
"""
|
||
total_window_len = input_len + pred_len
|
||
X, y = [], []
|
||
|
||
n_samples = len(data)
|
||
|
||
for start_idx in range(0, n_samples, slide_step):
|
||
end_idx = start_idx + total_window_len
|
||
|
||
# Skip if there's not enough data for a full window
|
||
if end_idx > n_samples:
|
||
break
|
||
|
||
# Split window into input and target
|
||
input_window = data[start_idx : start_idx + input_len]
|
||
target_window = data[start_idx + input_len : end_idx]
|
||
|
||
X.append(input_window)
|
||
y.append(target_window)
|
||
|
||
return np.array(X), np.array(y)
|
||
|
||
|
||
def process_and_save_time_series(
|
||
csv_path,
|
||
output_file,
|
||
input_len,
|
||
pred_len,
|
||
slide_step,
|
||
dataset_name=None, # 新增：数据集名称参数
|
||
selected_columns=None,
|
||
date_column='date',
|
||
freq='h',
|
||
split_method='auto',
|
||
train_ratio=0.7,
|
||
val_ratio=0.1,
|
||
test_ratio=0.2,
|
||
has_time_column=True, # 新增：是否包含时间列
|
||
):
|
||
"""
|
||
Process time series data and save it as an NPZ file along with the fitted scaler.
|
||
This function now calls the modified preprocess_time_series with flexible split methods.
|
||
|
||
Args:
|
||
dataset_name (str): 数据集名称（如 'ETTm1', 'weather'），优先使用此参数
|
||
split_method (str): Data split method - 'auto', 'specific', or 'ratio'
|
||
train_ratio (float): Training set ratio (only used when split_method='ratio')
|
||
val_ratio (float): Validation set ratio (only used when split_method='ratio')
|
||
test_ratio (float): Test set ratio (only used when split_method='ratio')
|
||
has_time_column (bool): Whether the dataset has a time column
|
||
"""
|
||
# Create output directory if it doesn't exist
|
||
output_dir = os.path.dirname(os.path.abspath(output_file))
|
||
os.makedirs(output_dir, exist_ok=True)
|
||
|
||
# Extract data file name from path
|
||
data_path_name = os.path.basename(csv_path)
|
||
|
||
# Load and preprocess the time series data using the new logic
|
||
result = preprocess_time_series(
|
||
csv_data=csv_path,
|
||
input_len=input_len,
|
||
pred_len=pred_len,
|
||
slide_step=slide_step,
|
||
dataset_name=dataset_name,
|
||
data_path_name=data_path_name,
|
||
selected_columns=selected_columns,
|
||
date_column=date_column,
|
||
freq=freq,
|
||
split_method=split_method,
|
||
train_ratio=train_ratio,
|
||
val_ratio=val_ratio,
|
||
test_ratio=test_ratio,
|
||
has_time_column=has_time_column
|
||
)
|
||
|
||
# Extract the processed data
|
||
scaler = result.pop('scaler') # Pop scaler to not save it in the npz
|
||
|
||
# Save the scaler object separately
|
||
scaler_file = output_file.replace('.npz', '_scaler.gz')
|
||
joblib.dump(scaler, scaler_file)
|
||
print(f"Saved scaler to {scaler_file}")
|
||
|
||
# Save the processed data arrays as .npz file
|
||
np.savez(output_file, **result)
|
||
print(f"Saved processed data to {output_file}")
|
||
|
||
return result
|
||
|
||
|