Building/src_old/dataset.py

import os
import glob
import pandas as pd
import numpy as np
import torch
from torch.utils.data import Dataset, DataLoader
from sklearn.preprocessing import StandardScaler
from config import *

class MultiOutputStandardizer:
    """Per-channel standardization that ignores missing labels via masks."""
    def __init__(self, n_outputs):
        self.n_outputs = n_outputs
        self.mean_ = np.zeros(n_outputs, dtype=np.float32)
        self.scale_ = np.ones(n_outputs, dtype=np.float32)
        self.fitted = False

    def fit(self, y_sequences, mask_sequences):
        means = []
        scales = []
        for c in range(self.n_outputs):
            valid_values = []
            for y_seq, m_seq in zip(y_sequences, mask_sequences):
                valid = m_seq[:, c] > 0.5
                if np.any(valid):
                    valid_values.append(y_seq[valid, c])
            if len(valid_values) == 0:
                means.append(0.0)
                scales.append(1.0)
                continue
            vals = np.concatenate(valid_values, axis=0)
            mean = float(np.mean(vals))
            std = float(np.std(vals))
            if std < 1e-6:
                std = 1.0
            means.append(mean)
            scales.append(std)

        self.mean_ = np.asarray(means, dtype=np.float32)
        self.scale_ = np.asarray(scales, dtype=np.float32)
        self.fitted = True

    def transform(self, y):
        if not self.fitted:
            raise RuntimeError("MultiOutputStandardizer must be fitted before transform.")
        return (y - self.mean_) / self.scale_

    def inverse_transform(self, y):
        if not self.fitted:
            raise RuntimeError("MultiOutputStandardizer must be fitted before inverse_transform.")
        return y * self.scale_ + self.mean_


class BuildingDataset(Dataset):
    def __init__(self, file_paths, seq_len, step_size, scaler_X=None, scaler_Y=None, fit_scaler=False):
        self.seq_len = seq_len
        self.X_data = []
        self.Y_data = []
        self.M_data = []
        self.scaler_X = scaler_X if scaler_X is not None else StandardScaler()
        self.scaler_Y = scaler_Y if scaler_Y is not None else MultiOutputStandardizer(len(OUTPUT_SENSORS))

        raw_X = []
        raw_Y = []
        raw_M = []

        for f in file_paths:
            # 读取数据
            df = pd.read_csv(f)
            # 使用长表格式: code, type, time, value1, value2, value3
            
            # 提取 012 的输入轴作为基准
            df_in = df[df['code'] == INPUT_SENSOR][['time', INPUT_AXIS]].rename(columns={INPUT_AXIS: 'input_signal'})
            if len(df_in) == 0:
                # 若无输入传感器(如自由衰减数据)，则补零
                df_in = pd.DataFrame({'time': df['time'].unique()})
                df_in['input_signal'] = 0.0
            
            # 提取 007~011 的 Z 轴并按时间戳逐步合并 (使用 left join 确保以 df_in 的时间为基准)
            df_merged = df_in
            for sens in OUTPUT_SENSORS:
                df_out_sens = df[df['code'] == sens][['time', OUTPUT_AXIS]].rename(columns={OUTPUT_AXIS: f'out_{sens}'})
                df_merged = pd.merge(df_merged, df_out_sens, on='time', how='left')
                
            df_merged = df_merged.sort_values('time').reset_index(drop=True)
            
            if len(df_merged) == 0:
                print(f"Warning: Skipping file {f} due to no overlapping timestamps across required sensors.")
                continue

            x_seq = df_merged['input_signal'].values.reshape(-1, 1).astype(np.float32)

            # 提取所有 target 传感器列与可用性掩码
            out_cols = [f'out_{sens}' for sens in OUTPUT_SENSORS]
            y_seq = np.zeros((len(df_merged), len(OUTPUT_SENSORS)), dtype=np.float32)
            m_seq = np.zeros((len(df_merged), len(OUTPUT_SENSORS)), dtype=np.float32)
            for c, col in enumerate(out_cols):
                series = df_merged[col]
                observed = ~series.isna()
                m_seq[:, c] = observed.astype(np.float32)
                if observed.any():
                    filled = series.interpolate(method='linear').bfill().ffill()
                    y_seq[:, c] = filled.fillna(0.0).values.astype(np.float32)
                else:
                    y_seq[:, c] = 0.0
            
            raw_X.append(x_seq)
            raw_Y.append(y_seq)
            raw_M.append(m_seq)

        if len(raw_X) == 0:
            raise ValueError("未能从文件中构造出有效序列，请检查数据路径与传感器编码配置。")
            
        # 拼接所有文件数据进行 fit
        X_all = np.vstack(raw_X)
        
        if fit_scaler:
            self.scaler_X.fit(X_all)
            self.scaler_Y.fit(raw_Y, raw_M)
            
        # 切分窗口
        for x_seq, y_seq, m_seq in zip(raw_X, raw_Y, raw_M):
            x_seq_scaled = self.scaler_X.transform(x_seq)
            y_seq_scaled = self.scaler_Y.transform(y_seq)
            y_seq_scaled = np.where(m_seq > 0.5, y_seq_scaled, 0.0).astype(np.float32)
            
            for i in range(0, len(x_seq_scaled) - seq_len + 1, step_size):
                x_win = x_seq_scaled[i:i+seq_len]
                y_win = y_seq_scaled[i:i+seq_len]
                m_win = m_seq[i:i+seq_len]
                if np.sum(m_win) <= 0:
                    continue
                self.X_data.append(x_win)
                self.Y_data.append(y_win)
                self.M_data.append(m_win)
                
        self.X_data = np.array(self.X_data)
        self.Y_data = np.array(self.Y_data)
        self.M_data = np.array(self.M_data)

    def __len__(self):
        return len(self.X_data)

    def __getitem__(self, idx):
        return (
            torch.tensor(self.X_data[idx], dtype=torch.float32),
            torch.tensor(self.Y_data[idx], dtype=torch.float32),
            torch.tensor(self.M_data[idx], dtype=torch.float32),
        )

def get_dataloaders(condition='Non_TMD', include_free_vib=False):
    """
    condition: 'Non_TMD' 或者是 'TMD'
    include_free_vib: 是否在训练集中加入自由振动与自由衰减数据
    """
    base_dir = os.path.join(DATA_DIR, condition)
    
    # 手动区分的子目录
    train_files = glob.glob(os.path.join(base_dir, 'train', '*.csv'))
    if include_free_vib:
        train_files += glob.glob(os.path.join(base_dir, 'free_vib', '*.csv'))
        
    val_files = glob.glob(os.path.join(base_dir, 'val', '*.csv'))
    test_files = glob.glob(os.path.join(base_dir, 'test', '*.csv'))
    
    print(f"[{condition}] Train files: {len(train_files)}, Val files: {len(val_files)}, Test files: {len(test_files)}")
    
    if len(train_files) == 0:
        raise ValueError(f"错误: 在 {base_dir}/train 目录下未找到训练文件！请检查路径是否正确。")
    if len(val_files) == 0:
        raise ValueError(f"错误: 在 {base_dir}/val 目录下未找到验证文件！请检查路径是否正确。")
    
    train_dataset = BuildingDataset(train_files, SEQ_LEN, STEP_SIZE, fit_scaler=True)
    val_dataset = BuildingDataset(val_files, SEQ_LEN, STEP_SIZE, 
                                  scaler_X=train_dataset.scaler_X, 
                                  scaler_Y=train_dataset.scaler_Y, fit_scaler=False)
    # 若某条件（如 TMD）下没有 test 数据，可以处理一下防止报错
    test_loader = None
    if len(test_files) > 0:
        test_dataset = BuildingDataset(test_files, SEQ_LEN, STEP_SIZE, 
                                       scaler_X=train_dataset.scaler_X, 
                                       scaler_Y=train_dataset.scaler_Y, fit_scaler=False)
        test_loader = DataLoader(test_dataset, batch_size=BATCH_SIZE, shuffle=False)
                                   
    train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)
    val_loader = DataLoader(val_dataset, batch_size=BATCH_SIZE, shuffle=False)
    
    return train_loader, val_loader, test_loader, train_dataset.scaler_X, train_dataset.scaler_Y