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