from __future__ import annotations from dataclasses import dataclass, field from pathlib import Path import re from typing import Any import numpy as np import pandas as pd import torch import torch.nn.functional as F from torch.utils.data import DataLoader, Dataset, WeightedRandomSampler try: from .config import DataConfig, ExperimentConfig except ImportError: from config import DataConfig, ExperimentConfig def calculate_rms(signal: np.ndarray) -> float: signal = np.asarray(signal, dtype=np.float64).reshape(-1) return float(np.sqrt(np.mean(np.square(signal)))) def extract_frequency_hz(file_name: str) -> float | None: match = re.search(r"(\d+(?:\.\d+)?)Hz", file_name, flags=re.IGNORECASE) if match is None: return None return float(match.group(1)) @dataclass class NormalizationStats: x_mean: torch.Tensor x_std: torch.Tensor y_wave_mean: torch.Tensor y_wave_std: torch.Tensor aux_mean: torch.Tensor aux_std: torch.Tensor target_mean: torch.Tensor target_std: torch.Tensor @dataclass class TensorStandardScaler: mean: torch.Tensor std: torch.Tensor def transform(self, array: torch.Tensor | np.ndarray) -> torch.Tensor | np.ndarray: if isinstance(array, torch.Tensor): mean = self.mean.to(array.device, dtype=array.dtype) std = self.std.to(array.device, dtype=array.dtype) return (array - mean) / std np_array = np.asarray(array, dtype=np.float32) return (np_array - self.mean.cpu().numpy()) / self.std.cpu().numpy() def inverse_transform(self, array: torch.Tensor | np.ndarray) -> torch.Tensor | np.ndarray: if isinstance(array, torch.Tensor): mean = self.mean.to(array.device, dtype=array.dtype) std = self.std.to(array.device, dtype=array.dtype) return array * std + mean np_array = np.asarray(array, dtype=np.float32) return np_array * self.std.cpu().numpy() + self.mean.cpu().numpy() @dataclass class RMSRecord: file_path: Path split: str time: torch.Tensor x_full: torch.Tensor x_model_input: torch.Tensor y_model_target: torch.Tensor x_rms: torch.Tensor y_rms: torch.Tensor target_value: torch.Tensor aux_features: torch.Tensor sample_weight: float interpolation_count: int = 0 @dataclass class SplitLoadReport: split: str loaded_files: list[str] = field(default_factory=list) skipped_files: list[tuple[str, str]] = field(default_factory=list) interpolated_files: dict[str, int] = field(default_factory=dict) def to_lines(self) -> list[str]: lines = [f"[{self.split}] loaded={len(self.loaded_files)} skipped={len(self.skipped_files)}"] for file_name, reason in self.skipped_files: lines.append(f" - skipped {file_name}: {reason}") for file_name, count in self.interpolated_files.items(): if count > 0: lines.append(f" - interpolated {file_name}: missing_points={count}") return lines class RMSRegressionDataset(Dataset): def __init__( self, records: list[RMSRecord], normalization: NormalizationStats, max_sequence_length: int, ) -> None: self.records = records self.normalization = normalization self.max_sequence_length = max_sequence_length self.sample_weights = [record.sample_weight for record in records] def __len__(self) -> int: return len(self.records) def __getitem__(self, index: int) -> dict[str, Any]: record = self.records[index] x = record.x_model_input if x.shape[0] > self.max_sequence_length: x = x[-self.max_sequence_length :] pad_length = self.max_sequence_length - x.shape[0] if pad_length > 0: x = F.pad(x.transpose(0, 1), (pad_length, 0), value=0.0).transpose(0, 1) x_norm = (x - self.normalization.x_mean) / self.normalization.x_std y = record.y_model_target if y.shape[0] > self.max_sequence_length: y = y[-self.max_sequence_length :] aux_norm = (record.aux_features - self.normalization.aux_mean) / self.normalization.aux_std target_norm = (record.target_value - self.normalization.target_mean) / self.normalization.target_std if pad_length > 0: y = F.pad(y.transpose(0, 1), (pad_length, 0), value=0.0).transpose(0, 1) y_wave_norm = (y - self.normalization.y_wave_mean) / self.normalization.y_wave_std valid_length = min(record.x_model_input.shape[0], self.max_sequence_length) mask = torch.zeros(self.max_sequence_length, dtype=torch.float32) mask[-valid_length:] = 1.0 return { "x": x_norm, "aux": aux_norm, "y": target_norm, "y_wave": y_wave_norm, "y_raw": record.y_rms, "x_rms_raw": record.x_rms, "frequency_hz": torch.tensor(record.aux_features[0].item(), dtype=torch.float32), "mask": mask, "file_name": record.file_path.name, "file_path": str(record.file_path), "valid_length": torch.tensor(valid_length, dtype=torch.long), "sample_weight": torch.tensor(record.sample_weight, dtype=torch.float32), } def list_split_files(config: DataConfig, split: str) -> list[Path]: split_dir = getattr(config, f"{split}_dir") return sorted(split_dir.glob(config.csv_pattern)) def _value_frame(df: pd.DataFrame, config: DataConfig, sensor_code: str, value_column: str) -> pd.DataFrame: sensor_df = df.loc[df[config.code_column] == sensor_code, [config.time_column, value_column]].copy() sensor_df = sensor_df.sort_values(config.time_column) sensor_df = sensor_df.drop_duplicates(subset=config.time_column, keep="first") sensor_df[config.time_column] = sensor_df[config.time_column].astype("float64") sensor_df[value_column] = sensor_df[value_column].astype("float32") return sensor_df def select_model_segment(signal: np.ndarray, config: DataConfig) -> np.ndarray: if not config.use_steady_state_only: return signal start_index = int(len(signal) * config.steady_state_start_ratio) max_start = max(0, len(signal) - config.steady_state_min_samples) start_index = min(start_index, max_start) return signal[start_index:] def estimate_dominant_frequency(signal: np.ndarray, sampling_rate: float) -> float: signal = np.asarray(signal, dtype=np.float64).reshape(-1) if signal.size < 4: return 0.0 fft_values = np.fft.rfft(signal) freqs = np.fft.rfftfreq(signal.size, d=1.0 / sampling_rate) magnitudes = np.abs(fft_values) magnitudes[0] = 0.0 band_mask = (freqs >= 0.1) & (freqs <= 5.0) if not np.any(band_mask): return 0.0 masked_magnitudes = np.where(band_mask, magnitudes, 0.0) return float(freqs[int(np.argmax(masked_magnitudes))]) def estimate_sampling_rate(time_values: np.ndarray) -> float: if time_values.size < 2: return 100.0 dt = np.diff(time_values) dt = dt[np.isfinite(dt)] dt = dt[dt > 0.0] if dt.size == 0: return 100.0 return float(1.0 / np.median(dt)) def compute_sample_weight(y_rms: float, x_rms: float, freq_hz: float, config: DataConfig) -> float: gain = y_rms / max(x_rms, config.normalization_eps) low_freq_factor = (config.low_frequency_reference_hz / max(freq_hz, config.normalization_eps)) ** config.low_frequency_emphasis_power weight = (gain ** config.train_weight_power) * low_freq_factor return float(np.clip(weight, config.train_weight_min, config.train_weight_max)) def load_split_records(config: DataConfig, split: str) -> tuple[list[RMSRecord], SplitLoadReport]: records: list[RMSRecord] = [] report = SplitLoadReport(split=split) for file_path in list_split_files(config, split): df = pd.read_csv(file_path) base_df = _value_frame(df, config, config.base_sensor_code, config.base_axis) response_df = _value_frame(df, config, config.response_sensor_code, config.response_axis) if base_df.empty or response_df.empty: report.skipped_files.append((file_path.name, "missing required sensor")) continue aligned = base_df.rename(columns={config.base_axis: "base_signal"}) aligned = aligned.merge( response_df.rename(columns={config.response_axis: "response_signal"}), on=config.time_column, how="left", ) interpolation_count = int(aligned["response_signal"].isna().sum()) aligned["response_signal"] = aligned["response_signal"].interpolate( method=config.interpolation_method, limit_direction="both", ).ffill().bfill() if aligned["response_signal"].isna().any(): report.skipped_files.append((file_path.name, "remaining NaN after interpolation")) continue time_values = aligned[config.time_column].to_numpy(dtype=np.float64) x_values = aligned["base_signal"].to_numpy(dtype=np.float32) y_values = aligned["response_signal"].to_numpy(dtype=np.float32) if len(x_values) < config.min_sequence_length: report.skipped_files.append((file_path.name, f"sequence too short: {len(x_values)}")) continue x_model = select_model_segment(x_values, config) y_model = select_model_segment(y_values, config) time_model = select_model_segment(time_values, config) sampling_rate = estimate_sampling_rate(time_model) file_frequency = extract_frequency_hz(file_path.name) x_rms = calculate_rms(x_model) y_rms = calculate_rms(y_model) target_value = float(np.log(max(y_rms / max(x_rms, config.normalization_eps), config.normalization_eps))) dominant_freq = estimate_dominant_frequency(x_model, sampling_rate) feature_frequency = file_frequency if file_frequency is not None else dominant_freq sample_weight = compute_sample_weight(y_rms=y_rms, x_rms=x_rms, freq_hz=feature_frequency, config=config) if split == "train" else 1.0 aux_features = torch.tensor( [feature_frequency, x_rms, float(len(x_model)) / float(config.max_sequence_length)], dtype=torch.float32, ) records.append( RMSRecord( file_path=file_path, split=split, time=torch.tensor(time_model, dtype=torch.float64), x_full=torch.tensor(x_values[:, None], dtype=torch.float32), x_model_input=torch.tensor(x_model[:, None], dtype=torch.float32), y_model_target=torch.tensor(y_model[:, None], dtype=torch.float32), x_rms=torch.tensor([x_rms], dtype=torch.float32), y_rms=torch.tensor([y_rms], dtype=torch.float32), target_value=torch.tensor([target_value], dtype=torch.float32), aux_features=aux_features, sample_weight=sample_weight, interpolation_count=interpolation_count, ) ) report.loaded_files.append(file_path.name) if interpolation_count > 0: report.interpolated_files[file_path.name] = interpolation_count if not records: raise RuntimeError(f"No usable records found for split='{split}'.") return records, report def fit_normalization(records: list[RMSRecord], config: DataConfig) -> NormalizationStats: x_all = torch.cat([record.x_model_input for record in records], dim=0) y_all = torch.cat([record.y_model_target for record in records], dim=0) aux_all = torch.stack([record.aux_features for record in records], dim=0) target_all = torch.cat([record.target_value for record in records], dim=0) def safe_std(tensor: torch.Tensor, dim: int) -> torch.Tensor: std = tensor.std(dim=dim, unbiased=False) return torch.clamp(std, min=config.normalization_eps) return NormalizationStats( x_mean=x_all.mean(dim=0), x_std=safe_std(x_all, dim=0), y_wave_mean=y_all.mean(dim=0), y_wave_std=safe_std(y_all, dim=0), aux_mean=aux_all.mean(dim=0), aux_std=safe_std(aux_all, dim=0), target_mean=target_all.mean(dim=0, keepdim=True), target_std=safe_std(target_all, dim=0).view(1), ) def build_datasets( config: ExperimentConfig | DataConfig, ) -> tuple[dict[str, RMSRegressionDataset], NormalizationStats, dict[str, SplitLoadReport]]: data_config = config.data if isinstance(config, ExperimentConfig) else config train_records, train_report = load_split_records(data_config, "train") normalization = fit_normalization(train_records, data_config) val_records, val_report = load_split_records(data_config, "val") test_records, test_report = load_split_records(data_config, "test") datasets = { "train": RMSRegressionDataset(train_records, normalization, data_config.max_sequence_length), "val": RMSRegressionDataset(val_records, normalization, data_config.max_sequence_length), "test": RMSRegressionDataset(test_records, normalization, data_config.max_sequence_length), } return datasets, normalization, {"train": train_report, "val": val_report, "test": test_report} def build_dataloaders( config: ExperimentConfig | DataConfig, ) -> tuple[dict[str, DataLoader], dict[str, RMSRegressionDataset], dict[str, SplitLoadReport]]: data_config = config.data if isinstance(config, ExperimentConfig) else config datasets, _, reports = build_datasets(config) train_sampler = None train_shuffle = True if data_config.use_weighted_train_sampler: weights = torch.tensor(datasets["train"].sample_weights, dtype=torch.double) train_sampler = WeightedRandomSampler(weights, num_samples=len(weights), replacement=True) train_shuffle = False loaders = { "train": DataLoader( datasets["train"], batch_size=data_config.batch_size, shuffle=train_shuffle, sampler=train_sampler, num_workers=data_config.num_workers, pin_memory=data_config.pin_memory, ), "val": DataLoader( datasets["val"], batch_size=data_config.batch_size, shuffle=False, num_workers=data_config.num_workers, pin_memory=data_config.pin_memory, ), "test": DataLoader( datasets["test"], batch_size=data_config.batch_size, shuffle=False, num_workers=data_config.num_workers, pin_memory=data_config.pin_memory, ), } return loaders, datasets, reports def get_dataloaders( config: ExperimentConfig | DataConfig, ) -> tuple[ dict[str, DataLoader], dict[str, RMSRegressionDataset], dict[str, SplitLoadReport], TensorStandardScaler, TensorStandardScaler, TensorStandardScaler, ]: loaders, datasets, reports = build_dataloaders(config) normalization = datasets["train"].normalization x_scaler = TensorStandardScaler(normalization.x_mean, normalization.x_std) aux_scaler = TensorStandardScaler(normalization.aux_mean, normalization.aux_std) y_scaler = TensorStandardScaler(normalization.target_mean, normalization.target_std) return loaders, datasets, reports, x_scaler, aux_scaler, y_scaler def report_to_text(reports: dict[str, SplitLoadReport]) -> str: lines: list[str] = [] for split in ("train", "val", "test"): lines.extend(reports[split].to_lines()) return "\n".join(lines)