from __future__ import annotations import argparse from pathlib import Path import matplotlib.pyplot as plt import numpy as np import pandas as pd import torch try: from .config import ExperimentConfig, make_rms_forward_config from .dataset import get_dataloaders, report_to_text from .model import build_model except ImportError: from config import ExperimentConfig, make_rms_forward_config from dataset import get_dataloaders, report_to_text from model import build_model def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description="Evaluate direct RMS regression model on harmonic data.") parser.add_argument("--split", choices=("train", "val", "test"), default="val") parser.add_argument("--sample-index", type=int, default=0) parser.add_argument("--device", type=str, default=None) parser.add_argument("--checkpoint", type=str, default=None) parser.add_argument("--all-samples", action="store_true") return parser.parse_args() def resolve_device(device_name: str | None, config: ExperimentConfig) -> torch.device: requested = device_name or config.train.device if requested.startswith("cuda") and not torch.cuda.is_available(): return torch.device("cpu") return torch.device(requested) def resolve_checkpoint_path(config: ExperimentConfig, checkpoint_arg: str | None) -> Path: if checkpoint_arg: return Path(checkpoint_arg).resolve() return config.data.project_root / config.train.checkpoint_dir / "forward_rms" / config.train.best_model_name def relative_percent_error(true_value: float, pred_value: float) -> float: denominator = max(abs(true_value), 1e-12) return abs(pred_value - true_value) / denominator * 100.0 def inverse_waveform(array: torch.Tensor, mean: torch.Tensor, std: torch.Tensor) -> torch.Tensor: mean = mean.to(array.device, dtype=array.dtype) std = std.to(array.device, dtype=array.dtype) return array * std + mean def predict_sample( model: torch.nn.Module, sample: dict[str, object], device: torch.device, y_scaler, x_scaler, y_wave_mean: torch.Tensor, y_wave_std: torch.Tensor, ) -> dict[str, float | str | np.ndarray]: x = sample["x"].unsqueeze(0).to(device) aux = sample["aux"].unsqueeze(0).to(device) mask = sample["mask"].unsqueeze(0).to(device) x_rms_raw = sample["x_rms_raw"].unsqueeze(0).to(device) y_true = float(sample["y_raw"].detach().cpu().numpy().reshape(-1)[0]) frequency_hz = float(sample["frequency_hz"].item()) file_name = str(sample["file_name"]) valid_mask = sample["mask"].detach().cpu().numpy() > 0.5 with torch.no_grad(): outputs = model(x, mask, aux) pred_target = y_scaler.inverse_transform(outputs["rms"]) y_pred = float((torch.exp(pred_target) * x_rms_raw).detach().cpu().numpy().reshape(-1)[0]) pred_wave = inverse_waveform(outputs["waveform"], y_wave_mean, y_wave_std).detach().cpu().numpy().reshape(-1)[valid_mask] true_wave = inverse_waveform(sample["y_wave"], y_wave_mean.cpu(), y_wave_std.cpu()).detach().cpu().numpy().reshape(-1)[valid_mask] x_wave = x_scaler.inverse_transform(sample["x"]).reshape(-1)[valid_mask] return { "file_name": file_name, "frequency_hz": frequency_hz, "true_rms": y_true, "pred_rms": y_pred, "relative_error_percent": float(relative_percent_error(y_true, y_pred)), "x_wave": x_wave.numpy().reshape(-1), "true_wave": true_wave, "pred_wave": pred_wave, } def evaluate_sample( model: torch.nn.Module, sample: dict[str, object], device: torch.device, y_scaler, x_scaler, y_wave_mean: torch.Tensor, y_wave_std: torch.Tensor, ) -> dict[str, float | str]: result = predict_sample(model, sample, device, y_scaler, x_scaler, y_wave_mean, y_wave_std) return { "file_name": str(result["file_name"]), "frequency_hz": float(result["frequency_hz"]), "true_rms": float(result["true_rms"]), "pred_rms": float(result["pred_rms"]), "relative_error_percent": float(result["relative_error_percent"]), } def save_single_sample_figure(result: dict[str, float | str | np.ndarray], split: str, save_dir: Path, sample_index: int) -> Path: true_rms = float(result["true_rms"]) pred_rms = float(result["pred_rms"]) error_percent = relative_percent_error(true_rms, pred_rms) x_signal = np.asarray(result["x_wave"], dtype=np.float64).reshape(-1) true_wave = np.asarray(result["true_wave"], dtype=np.float64).reshape(-1) pred_wave = np.asarray(result["pred_wave"], dtype=np.float64).reshape(-1) time_steps = np.arange(x_signal.shape[0], dtype=np.float64) fig, axes = plt.subplots(3, 1, figsize=(12, 10)) fig.suptitle(f"Forward RMS regression | {split} | {result['file_name']}") axes[0].plot(time_steps, x_signal, linewidth=1.2) axes[0].set_title("Input Base Excitation (Steady-State Segment)") axes[0].set_xlabel("Time Step") axes[0].set_ylabel("Acceleration") axes[0].grid(True, alpha=0.3) axes[1].plot(time_steps, true_wave, label="True response", linewidth=1.2, color="tab:blue") axes[1].plot(time_steps, pred_wave, label="Pred response", linewidth=1.2, color="tab:orange", alpha=0.85) axes[1].set_title("Auxiliary Waveform Head: True vs Predicted Top Response") axes[1].set_xlabel("Time Step") axes[1].set_ylabel("Acceleration") axes[1].grid(True, alpha=0.3) axes[1].legend() axes[2].bar(["True RMS", "Pred RMS"], [true_rms, pred_rms], color=["tab:blue", "tab:orange"]) axes[2].set_title(f"True RMS: {true_rms:.4f}, Pred RMS: {pred_rms:.4f}, Error: {error_percent:.2f}%") axes[2].set_ylabel("RMS") axes[2].grid(True, axis="y", alpha=0.3) plt.tight_layout() save_dir.mkdir(parents=True, exist_ok=True) figure_path = save_dir / f"evaluation_{split}_s{sample_index}_waveform.png" plt.savefig(figure_path, dpi=180, bbox_inches="tight") plt.show() plt.close(fig) return figure_path def main() -> None: args = parse_args() config = make_rms_forward_config() device = resolve_device(args.device, config) checkpoint_path = resolve_checkpoint_path(config, args.checkpoint) loaders, datasets, reports, x_scaler, aux_scaler, y_scaler = get_dataloaders(config) print(report_to_text(reports)) if not checkpoint_path.exists(): raise FileNotFoundError(f"Checkpoint not found: {checkpoint_path}") dataset = datasets[args.split] sample = dataset[args.sample_index] normalization = datasets["train"].normalization model = build_model(config).to(device) checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=False) model.load_state_dict(checkpoint["model_state_dict"]) model.eval() save_dir = config.data.project_root / "evaluation_outputs" / "forward_rms" if args.all_samples: rows = [ evaluate_sample( model, dataset[index], device, y_scaler, x_scaler, normalization.y_wave_mean, normalization.y_wave_std, ) for index in range(len(dataset)) ] result_df = pd.DataFrame(rows) result_df = result_df.sort_values(["relative_error_percent", "file_name"]).reset_index(drop=True) csv_path = save_dir / f"evaluation_{args.split}_all_samples.csv" save_dir.mkdir(parents=True, exist_ok=True) result_df.to_csv(csv_path, index=False) summary = { "count": len(result_df), "mean_error_percent": float(result_df["relative_error_percent"].mean()), "median_error_percent": float(result_df["relative_error_percent"].median()), "max_error_percent": float(result_df["relative_error_percent"].max()), "min_error_percent": float(result_df["relative_error_percent"].min()), } print(f"Checkpoint: {checkpoint_path}") print(f"Summary: {summary}") print(f"CSV saved to: {csv_path}") print(result_df.to_string(index=False)) return result = predict_sample( model, sample, device, y_scaler, x_scaler, normalization.y_wave_mean, normalization.y_wave_std, ) figure_path = save_single_sample_figure( result=result, split=args.split, save_dir=save_dir, sample_index=args.sample_index, ) print(f"Checkpoint: {checkpoint_path}") print(f"Sample file: {result['file_name']}") print(f"True RMS: {result['true_rms']:.6f}") print(f"Pred RMS: {result['pred_rms']:.6f}") print(f"Relative RMS Error (%): {result['relative_error_percent']:.4f}") print(f"Figure saved to: {figure_path}") if __name__ == "__main__": main()