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 CORE_FEATURE_NAMES, ExperimentConfig, make_experiment_config from .dataset import build_dataloaders, normalization_from_payload, report_to_text from .model import build_model except ImportError: from config import CORE_FEATURE_NAMES, ExperimentConfig, make_experiment_config from dataset import build_dataloaders, normalization_from_payload, report_to_text from model import build_model def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description="Evaluate feature MLP on train/val/test splits.") parser.add_argument("--split", choices=("train", "val", "test"), default="val") parser.add_argument("--sample-index", type=int, default=0) parser.add_argument("--all-samples", action="store_true") parser.add_argument("--device", type=str, default=None) parser.add_argument("--checkpoint", type=str, default=None) 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 / "task1_feature_mlp" / 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 predict_rms( model: torch.nn.Module, feature_norm: torch.Tensor, x_rms_raw: torch.Tensor, target_mean: torch.Tensor, target_std: torch.Tensor, ) -> torch.Tensor: pred_norm = model(feature_norm) pred_tr = torch.clamp(pred_norm * target_std + target_mean, min=1e-6) return pred_tr * x_rms_raw def evaluate_sample( model: torch.nn.Module, dataset, raw_record, sample_index: int, device: torch.device, ) -> dict[str, float | str]: sample = dataset[sample_index] normalization = dataset.normalization feature_norm = sample["x"].unsqueeze(0).to(device) x_rms_raw = sample["x_rms_raw"].unsqueeze(0).to(device) target_mean = normalization.target_mean.to(device).view(1, 1) target_std = normalization.target_std.to(device).view(1, 1) with torch.no_grad(): pred_rms = predict_rms(model, feature_norm, x_rms_raw, target_mean, target_std) true_rms = float(sample["y_rms_raw"].item()) pred_rms_value = float(pred_rms.detach().cpu().numpy().reshape(-1)[0]) row = { "file_name": str(sample["file_name"]), "frequency_hz": float(sample["frequency_hz"].item()), "true_rms": true_rms, "pred_rms": pred_rms_value, "relative_error_percent": relative_percent_error(true_rms, pred_rms_value), } features_raw = sample["features_raw"].numpy().reshape(-1) for name, value in zip(CORE_FEATURE_NAMES, features_raw): row[name] = float(value) row["sampling_rate"] = float(raw_record.sampling_rate) return row def save_all_samples_plot(result_df: pd.DataFrame, split: str, save_dir: Path) -> Path | None: if result_df["frequency_hz"].isna().any(): return None plot_df = result_df.sort_values("frequency_hz").reset_index(drop=True) fig, axes = plt.subplots(2, 1, figsize=(12, 8)) fig.suptitle(f"Task1 Feature-MLP Evaluation | {split}") axes[0].plot(plot_df["frequency_hz"], plot_df["true_rms"], marker="o", label="True RMS") axes[0].plot(plot_df["frequency_hz"], plot_df["pred_rms"], marker="o", label="Pred RMS") axes[0].set_xlabel("Frequency (Hz)") axes[0].set_ylabel("RMS") axes[0].grid(True, alpha=0.3) axes[0].legend() axes[1].bar(plot_df["frequency_hz"].astype(str), plot_df["relative_error_percent"], color="tab:orange") axes[1].set_xlabel("Frequency (Hz)") axes[1].set_ylabel("Relative Error (%)") axes[1].grid(True, axis="y", alpha=0.3) axes[1].tick_params(axis="x", labelrotation=45) plt.tight_layout() save_dir.mkdir(parents=True, exist_ok=True) figure_path = save_dir / f"evaluation_{split}_curve.png" plt.savefig(figure_path, dpi=180, bbox_inches="tight") plt.close(fig) return figure_path def save_single_sample_plot(raw_record, result: dict[str, float | str], split: str, save_dir: Path, sample_index: int) -> Path: time_middle = raw_record.time_middle.detach().cpu().numpy().reshape(-1) x_middle = raw_record.x_middle.detach().cpu().numpy().reshape(-1) y_middle = raw_record.y_middle.detach().cpu().numpy().reshape(-1) frequency_hz = float(result["frequency_hz"]) pred_rms = float(result["pred_rms"]) true_rms = float(result["true_rms"]) fig, axes = plt.subplots(3, 1, figsize=(12, 10)) fig.suptitle(f"Task1 Feature-MLP | {split} | {result['file_name']}") axes[0].plot(time_middle, x_middle, color="tab:blue") axes[0].set_title("Input Base Excitation (Middle Segment)") axes[0].set_xlabel("Time") axes[0].set_ylabel("Acceleration") axes[0].grid(True, alpha=0.3) axes[1].plot(time_middle, y_middle, color="tab:green") axes[1].set_title("True Top Response (Middle Segment)") axes[1].set_xlabel("Time") axes[1].set_ylabel("Acceleration") axes[1].grid(True, alpha=0.3) axes[2].bar(["True RMS", "Pred RMS"], [true_rms, pred_rms], color=["tab:green", "tab:orange"]) axes[2].set_title( f"Freq: {frequency_hz:.4f} Hz | True RMS: {true_rms:.6f} | Pred RMS: {pred_rms:.6f} | " f"Error: {float(result['relative_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}.png" plt.savefig(figure_path, dpi=180, bbox_inches="tight") plt.close(fig) return figure_path def main() -> None: args = parse_args() config = make_experiment_config() device = resolve_device(args.device, config) checkpoint_path = resolve_checkpoint_path(config, args.checkpoint) if not checkpoint_path.exists(): raise FileNotFoundError(f"Checkpoint not found: {checkpoint_path}") checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=False) if "normalization" in checkpoint: normalization = normalization_from_payload(checkpoint["normalization"]) else: normalization = None loaders, datasets, raw_records, reports = build_dataloaders(config) del loaders print(report_to_text(reports)) dataset = datasets[args.split] if normalization is not None: dataset.normalization = normalization model = build_model(config).to(device) model.load_state_dict(checkpoint["model_state_dict"]) model.eval() save_dir = config.data.project_root / "evaluation_outputs" / "task1_feature_mlp" if args.all_samples: rows = [ evaluate_sample(model, dataset, raw_records[args.split][index], index, device) for index in range(len(dataset)) ] result_df = pd.DataFrame(rows).sort_values(["relative_error_percent", "file_name"]).reset_index(drop=True) save_dir.mkdir(parents=True, exist_ok=True) csv_path = save_dir / f"evaluation_{args.split}_all_samples.csv" result_df.to_csv(csv_path, index=False) figure_path = save_all_samples_plot(result_df, args.split, save_dir) 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}") if figure_path is not None: print(f"Figure saved to: {figure_path}") print(result_df.to_string(index=False)) return result = evaluate_sample(model, dataset, raw_records[args.split][args.sample_index], args.sample_index, device) figure_path = save_single_sample_plot(raw_records[args.split][args.sample_index], result, args.split, save_dir, args.sample_index) print(f"Checkpoint: {checkpoint_path}") print(f"Sample file: {result['file_name']}") print(f"True RMS: {float(result['true_rms']):.6f}") print(f"Pred RMS: {float(result['pred_rms']):.6f}") print(f"Relative RMS Error (%): {float(result['relative_error_percent']):.4f}") print(f"Figure saved to: {figure_path}") if __name__ == "__main__": main()