DualFlow/models/pressure-predictor/90天TMP预测模型源码/data_preprocessor.py

# data_preprocessor.py
import os
import torch
import joblib
import numpy as np
import pandas as pd
from tqdm import tqdm  # 进度条工具
from sklearn.preprocessing import MinMaxScaler  # 数据标准化
from torch.utils.data import DataLoader, TensorDataset  # PyTorch数据加载工具
from concurrent.futures import ThreadPoolExecutor  # 多线程并行处理

class DataPreprocessor:
    
    @staticmethod
    def load_and_process_data(args, data):
        """
        加载并处理数据，划分训练/验证/测试集，生成数据加载器
        :param args: 配置参数（包含日期范围、序列长度等）
        :param data: 原始数据（DataFrame格式）
        :return: 训练/验证/测试数据加载器、原始数据
        """
        # 处理日期列
        data['date'] = pd.to_datetime(data['date'])
        time_interval = pd.Timedelta(hours=(args.resolution / 60))
        window_time_span = time_interval * (args.seq_len + 314)

        # 划分训练/验证/测试集（调整起始日期以适应滑动窗口）
        val_start_date = pd.to_datetime(args.val_start_date)
        test_start_date = pd.to_datetime(args.test_start_date)
        
        # 调整验证集/测试集的起始日期（提前窗口跨度，确保能生成完整输入序列）
        adjusted_val_start = val_start_date - window_time_span
        adjusted_test_start = test_start_date - window_time_span

        # 生成训练/验证/测试集的掩码（布尔索引）
        train_mask = (data['date'] >= pd.to_datetime(args.train_start_date)) & \
                     (data['date'] <= pd.to_datetime(args.train_end_date))

        val_mask = (data['date'] >= adjusted_val_start) & \
                   (data['date'] <= pd.to_datetime(args.val_end_date))

        test_mask = (data['date'] >= adjusted_test_start) & \
                    (data['date'] <= pd.to_datetime(args.test_end_date))

        # 应用掩码并重置索引
        train_data = data[train_mask].reset_index(drop=True)
        val_data = data[val_mask].reset_index(drop=True)
        test_data = data[test_mask].reset_index(drop=True)
        
        # 移除日期列用于建模
        train_data = train_data.drop(columns=['date'])
        val_data = val_data.drop(columns=['date'])
        test_data = test_data.drop(columns=['date'])
    
        # 创建监督学习数据集（输入序列+目标序列）
        train_supervised = DataPreprocessor.create_supervised_dataset(
            args,
            train_data,
            1
        )
        
        val_supervised = DataPreprocessor.create_supervised_dataset(
            args,
            val_data,
            1
        )
        
        test_supervised = DataPreprocessor.create_supervised_dataset(
            args,
            test_data,
            args.step_size
        )
        
        # 转换为DataLoader
        train_loader = DataPreprocessor.load_data(
            args, 
            train_supervised,
            shuffle=True
        )
        
        val_loader = DataPreprocessor.load_data(
            args, 
            val_supervised,
            shuffle=False
        )
        
        test_loader = DataPreprocessor.load_data(
            args, 
            test_supervised,
            shuffle=False
        )
        
        return train_loader, val_loader, test_loader, data  # 返回原始数据用于后续处理
    
    @staticmethod
    def read_and_combine_csv_files(args):
        """
        读取并合并多个CSV文件（支持多线程加速）
        :param args: 配置参数（包含数据目录、文件命名模式等）
        :return: 合并并预处理后的DataFrame
        """
        current_dir = os.path.dirname(__file__)
        parent_dir = os.path.dirname(current_dir)
        args.data_dir = os.path.join(parent_dir, args.data_dir)
        
        def read_file(file_count):
            """内部函数：读取单个CSV文件"""
            file_name = args.file_pattern.format(file_count)
            file_path = os.path.join(args.data_dir, file_name)
            return pd.read_csv(file_path)
        
        # 生成文件索引列表（从start_files到end_files）
        file_indices = list(range(args.start_files, args.end_files + 1))
        
        # 多线程读取文件（加速大文件读取）
        max_workers = os.cpu_count()    # 按CPU核心数设置线程数
        with ThreadPoolExecutor(max_workers=max_workers) as executor:
            results = list(tqdm(executor.map(read_file, file_indices),
                                total=len(file_indices),
                                desc="正在读取文件"))
        
        # 合并所有数据并重置索引
        all_data = pd.concat(results, ignore_index=True)
        
        # 按分辨率下采样
        chunk = all_data.iloc[::args.resolution, :].reset_index(drop=True)
        
        # 处理日期和时间特征
        chunk = DataPreprocessor.process_date(chunk)
        # 归一化
        chunk = DataPreprocessor.scaler_data(chunk)
        
        return chunk
    
    @staticmethod
    def process_date(data):
        """
        处理日期列，生成周期性时间特征（年周期）
        :param data: 包含'index'列（日期字符串）的DataFrame
        :return: 增加时间特征后的DataFrame
        """
        data = data.rename(columns={'index': 'date'})
        data['date'] = pd.to_datetime(data['date'])

        # 生成周期性时间特征
        data['day_of_year'] = data['date'].dt.dayofyear

        data['day_year_sin'] = np.sin(2 * np.pi * data['day_of_year'] / 366)
        data['day_year_cos'] = np.cos(2 * np.pi * data['day_of_year'] / 366)

        # 移除原始时间列，保留特征列
        data.drop(columns=['day_of_year'], inplace=True)

        # 调整列顺序（日期+时间特征+其他特征）
        time_features = ['day_year_sin', 'day_year_cos']
        other_columns = [col for col in data.columns if col not in ['date'] and col not in time_features]
        data = data[['date'] + time_features + other_columns]

        return data
    
    @staticmethod
    def scaler_data(data):
        """
        对数据进行归一化（除日期列外），并保存标准化器
        :param data: 包含'date'列的DataFrame
        :return: 标准化后的DataFrame
        """
        date_col = data[['date']]
        data_to_scale = data.drop(columns=['date'])

        scaler = MinMaxScaler(feature_range=(0, 1))
        scaled_data = scaler.fit_transform(data_to_scale)
        joblib.dump(scaler, 'scaler.pkl')  # 保存归一化器

        scaled_data = pd.DataFrame(scaled_data, columns=data_to_scale.columns)
        
        # 拼接日期列和标准化后的数据
        scaled_data = pd.concat([date_col.reset_index(drop=True), scaled_data], axis=1)
        
        return scaled_data
    
    @staticmethod
    def create_supervised_dataset(args, data, step_size):
        """
        将时间序列数据转换为监督学习格式（输入序列+目标序列）
        :param args: 配置参数（序列长度、输出长度等）
        :param data: 输入数据（DataFrame，不含日期列）
        :param step_size: 滑动窗口步长
        :return: 监督学习数据集（DataFrame）
        """
        data = pd.DataFrame(data)
        cols = []
        col_names = []
        
        feature_columns = data.columns.tolist()

        # 输入序列（t-0到t-(seq_len-1)）
        for col in feature_columns:
            for i in range(args.seq_len - 1, -1, -1):
                cols.append(data[[col]].shift(i))
                col_names.append(f"{col}(t-{i})")
        
        # 目标序列（仅取最后labels_num列作为预测目标）
        target_columns = feature_columns[-args.labels_num:]
        for i in range(1, args.output_size + 1):
            for col in target_columns:
                cols.append(data[[col]].shift(-i))
                col_names.append(f"{col}(t+{i})")

        # 合并并清洗数据
        dataset = pd.concat(cols, axis=1)
        dataset.columns = col_names
        dataset = dataset.iloc[::step_size, :]  # 按步长采样
        dataset.dropna(inplace=True)  # 移除含缺失值的行
        
        return dataset

    @staticmethod
    def load_data(args, dataset, shuffle):
        """
        将监督学习数据集转换为PyTorch张量，并创建DataLoader
        :param args: 配置参数
        :param dataset: 监督学习数据集（DataFrame）
        :param shuffle: 是否打乱数据
        :return: DataLoader对象
        """
        input_length = args.seq_len
        n_features = args.feature_num
        labels_num = args.labels_num
    
        n_features_total = n_features * input_length  # 输入特征总维度
        n_labels_total = args.output_size * labels_num  # 目标总维度

        # 分割输入和目标
        X = dataset.values[:, :n_features_total]
        y = dataset.values[:, n_features_total:n_features_total + n_labels_total]
    
        # 重塑输入为[样本数, 序列长度, 特征数]
        X = X.reshape(X.shape[0], input_length, n_features)
        X = torch.tensor(X, dtype=torch.float32).to(args.device)
        y = torch.tensor(y, dtype=torch.float32).to(args.device)

        # 创建数据集和数据加载器
        dataset_tensor = TensorDataset(X, y)
        generator = torch.Generator()
        generator.manual_seed(args.random_seed)  # 固定随机种子确保可复现
        
        data_loader = DataLoader(
            dataset_tensor, 
            batch_size=args.batch_size, 
            shuffle=shuffle,
            generator=generator
        )
    
        return data_loader