from __future__ import annotations from dataclasses import dataclass, field from pathlib import Path import re import numpy as np import pandas as pd try: from .config import CORE_FEATURE_NAMES, DataConfig, ExperimentConfig except ImportError: from config import CORE_FEATURE_NAMES, 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)) 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 get_middle_segment( time_values: np.ndarray, x_values: np.ndarray, y_values: np.ndarray, config: DataConfig, ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: length = len(x_values) search_start = int(length * config.middle_segment_start_ratio) search_end = int(length * config.middle_segment_end_ratio) search_start = max(0, min(search_start, length - 1)) search_end = max(search_start + 1, min(search_end, length)) if search_end - search_start < config.min_segment_length: center = length // 2 half = config.min_segment_length // 2 search_start = max(0, center - half) search_end = min(length, search_start + config.min_segment_length) search_start = max(0, search_end - config.min_segment_length) search_length = search_end - search_start window_length = max(config.min_segment_length, int(length * config.steady_window_ratio)) window_length = min(window_length, search_length) if search_length <= window_length: return ( time_values[search_start:search_end], x_values[search_start:search_end], y_values[search_start:search_end], ) stride = max(1, int(length * config.steady_window_stride_ratio)) candidate_y = y_values[search_start:search_end] candidate_rms = calculate_rms(candidate_y) best_score = float("inf") best_slice = slice(search_start, search_start + window_length) for window_start in range(search_start, search_end - window_length + 1, stride): window_end = window_start + window_length window_y = y_values[window_start:window_end] split_windows = np.array_split(window_y, config.stability_subwindow_count) split_rms = np.asarray([calculate_rms(chunk) for chunk in split_windows], dtype=np.float64) split_peaks = np.asarray([float(np.max(np.abs(chunk))) for chunk in split_windows], dtype=np.float64) rms_cv = float(split_rms.std() / max(split_rms.mean(), config.normalization_eps)) peak_cv = float(split_peaks.std() / max(split_peaks.mean(), config.normalization_eps)) window_rms = calculate_rms(window_y) score = rms_cv + 0.35 * peak_cv - 0.05 * (window_rms / max(candidate_rms, config.normalization_eps)) if score < best_score: best_score = score best_slice = slice(window_start, window_end) return time_values[best_slice], x_values[best_slice], y_values[best_slice] def _compute_windowed_spectrum(signal: np.ndarray, sampling_rate: float) -> tuple[np.ndarray, np.ndarray]: signal = np.asarray(signal, dtype=np.float64).reshape(-1) if signal.size < 4: return np.asarray([], dtype=np.float64), np.asarray([], dtype=np.float64) centered = signal - np.mean(signal) window = np.hanning(signal.size) scale = max(np.sum(window), 1e-12) fft_values = np.fft.rfft(centered * window) freqs = np.fft.rfftfreq(signal.size, d=1.0 / sampling_rate) magnitudes = (2.0 / scale) * np.abs(fft_values) return freqs, magnitudes def _parabolic_peak_frequency(freqs: np.ndarray, magnitudes: np.ndarray, peak_index: int) -> float: if peak_index <= 0 or peak_index >= magnitudes.size - 1: return float(freqs[peak_index]) alpha = magnitudes[peak_index - 1] beta = magnitudes[peak_index] gamma = magnitudes[peak_index + 1] denominator = alpha - 2.0 * beta + gamma if abs(denominator) < 1e-12: return float(freqs[peak_index]) offset = 0.5 * (alpha - gamma) / denominator bin_width = float(freqs[1] - freqs[0]) return float(freqs[peak_index] + offset * bin_width) def get_dominant_frequency(signal: np.ndarray, sampling_rate: float) -> float: freqs, magnitudes = _compute_windowed_spectrum(signal, sampling_rate) if magnitudes.size == 0: return 0.0 magnitudes[0] = 0.0 band_mask = (freqs >= 0.1) & (freqs <= 5.0) if not np.any(band_mask): return 0.0 band_magnitudes = np.where(band_mask, magnitudes, 0.0) peak_index = int(np.argmax(band_magnitudes)) return _parabolic_peak_frequency(freqs, magnitudes, peak_index) def theoretical_acceleration(frequency_hz: float, amplitude_m: float) -> float: return float((2.0 * np.pi * frequency_hz) ** 2 * amplitude_m) def extract_core_features( signal: np.ndarray, sampling_rate: float, config: DataConfig, known_frequency_hz: float | None = None, ) -> np.ndarray: frequency_hz = known_frequency_hz if known_frequency_hz is not None else get_dominant_frequency(signal, sampling_rate) x_rms = calculate_rms(signal) return np.asarray( [ frequency_hz, frequency_hz * frequency_hz, x_rms, theoretical_acceleration(frequency_hz, config.harmonic_amplitude_m), ], dtype=np.float32, ) 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 load_aligned_signals(file_path: Path, config: DataConfig) -> tuple[np.ndarray, np.ndarray, np.ndarray, int]: 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: raise ValueError(f"Missing required sensor in {file_path.name}") aligned = base_df.rename(columns={config.base_axis: "base_signal"}).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(): raise ValueError(f"Remaining NaN after interpolation in {file_path.name}") 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) return time_values, x_values, y_values, interpolation_count @dataclass class SampleRecord: file_path: Path frequency_hz: float features: np.ndarray x_rms: float y_rms: float target_tr: float time_middle: np.ndarray x_middle: np.ndarray y_middle: np.ndarray sampling_rate: float interpolation_count: int = 0 @dataclass class LoadReport: 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"[all_data] 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 def list_harmonic_files(config: DataConfig) -> list[Path]: files = sorted(config.data_root.rglob(config.harmonic_pattern)) return [file_path for file_path in files if extract_frequency_hz(file_path.name) is not None] def build_record_from_file(file_path: Path, config: DataConfig) -> SampleRecord: time_values, x_values, y_values, interpolation_count = load_aligned_signals(file_path, config) time_middle, x_middle, y_middle = get_middle_segment(time_values, x_values, y_values, config) sampling_rate = estimate_sampling_rate(time_middle) frequency_hz = extract_frequency_hz(file_path.name) if frequency_hz is None: frequency_hz = get_dominant_frequency(x_middle, sampling_rate) x_rms = calculate_rms(x_middle) y_rms = calculate_rms(y_middle) target_tr = float(y_rms / max(x_rms, config.normalization_eps)) features = extract_core_features(x_middle, sampling_rate, config, known_frequency_hz=frequency_hz) return SampleRecord( file_path=file_path, frequency_hz=float(frequency_hz), features=features, x_rms=float(x_rms), y_rms=float(y_rms), target_tr=target_tr, time_middle=np.asarray(time_middle, dtype=np.float64), x_middle=np.asarray(x_middle, dtype=np.float32).reshape(-1), y_middle=np.asarray(y_middle, dtype=np.float32).reshape(-1), sampling_rate=float(sampling_rate), interpolation_count=interpolation_count, ) def load_all_records(config: ExperimentConfig | DataConfig) -> tuple[list[SampleRecord], LoadReport]: data_config = config.data if isinstance(config, ExperimentConfig) else config records: list[SampleRecord] = [] report = LoadReport() for file_path in list_harmonic_files(data_config): try: record = build_record_from_file(file_path, data_config) except ValueError as error: report.skipped_files.append((file_path.name, str(error))) continue records.append(record) report.loaded_files.append(file_path.name) if record.interpolation_count > 0: report.interpolated_files[file_path.name] = record.interpolation_count if not records: raise RuntimeError("No usable harmonic records found in data_root.") return records, report def records_to_frame(records: list[SampleRecord]) -> pd.DataFrame: rows = [] for record in records: row = { "file_name": record.file_path.name, "frequency_hz": record.frequency_hz, "x_rms": record.x_rms, "y_rms": record.y_rms, "target_tr": record.target_tr, "sampling_rate": record.sampling_rate, } for feature_name, feature_value in zip(CORE_FEATURE_NAMES, record.features.tolist()): row[feature_name] = float(feature_value) rows.append(row) return pd.DataFrame(rows) def report_to_text(report: LoadReport) -> str: return "\n".join(report.to_lines())