#!/usr/bin/env python3 """Evaluate the pretrained phase picker directly on CREDIT-X1local test crops. This script recomputes the "Direct" phase-picking baseline used in the manuscript learning-summary figure. It uses the public CREDIT-X1local HDF5 file and split-key file, the archived base PNSN checkpoint, and the same crop materialization/evaluation functions as the matched SNR training experiment. """ from __future__ import annotations import argparse import csv import json import math import random import sys from pathlib import Path from statistics import mean, stdev import numpy as np import torch ROOT = Path(__file__).resolve().parents[1] if str(ROOT) not in sys.path: sys.path.insert(0, str(ROOT)) from models.BRNNPNSN import BRNN from scripts.reproduce_paper_stats import evaluate_outputs, run_model, strip_errors from scripts.snr_transfer_experiment import metric_row from scripts.snr_transfer_phase_balanced_experiment import ( build_h5_access_index, build_records_sequential, load_or_materialize_eval_samples, ) METRICS = [ "P_precision", "P_recall", "P_f1", "S_precision", "S_recall", "S_f1", "mean_f1", ] def choose_device(name: str) -> torch.device: if name == "auto": if torch.backends.mps.is_available(): return torch.device("mps") if torch.cuda.is_available(): return torch.device("cuda") return torch.device("cpu") return torch.device(name) def summarize(values: list[float]) -> dict[str, float | int | list[float]]: return { "mean": float(mean(values)) if values else math.nan, "std": float(stdev(values)) if len(values) > 1 else 0.0, "n": len(values), "values": values, } def write_csvs(rows_long: list[dict], out_dir: Path) -> None: out_dir.mkdir(parents=True, exist_ok=True) with (out_dir / "phase_direct_baseline_long.csv").open("w", newline="", encoding="utf-8") as handle: writer = csv.DictWriter( handle, fieldnames=["task", "seed", "condition_slug", "condition_label", "metric", "value"], ) writer.writeheader() writer.writerows(rows_long) by_metric: dict[str, list[float]] = {metric: [] for metric in METRICS} for row in rows_long: by_metric[row["metric"]].append(float(row["value"])) summary_rows = [] for metric in METRICS: stat = summarize(by_metric[metric]) summary_rows.append( { "task": "phase", "condition_slug": "pretrained_direct", "condition_label": "Pretrained direct use", "metric": metric, "mean": stat["mean"], "std": stat["std"], "n": stat["n"], "values": json.dumps(stat["values"]), } ) with (out_dir / "phase_direct_baseline_summary.csv").open("w", newline="", encoding="utf-8") as handle: writer = csv.DictWriter( handle, fieldnames=["task", "condition_slug", "condition_label", "metric", "mean", "std", "n", "values"], ) writer.writeheader() writer.writerows(summary_rows) def main() -> None: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--h5", type=Path, required=True) parser.add_argument("--keys", type=Path, required=True) parser.add_argument("--base-ckpt", type=Path, default=ROOT.parent / "checkpoints" / "base" / "pnsn.v3.pt") parser.add_argument("--out-dir", type=Path, default=ROOT.parent / "outputs" / "phase_direct_baseline") parser.add_argument("--cache-dir", type=Path, default=None) parser.add_argument("--seeds", nargs="+", type=int, default=[20260609, 20260610, 20260611]) parser.add_argument("--length", type=int, default=5120) parser.add_argument("--padlen", type=int, default=512) parser.add_argument("--eval-samples", type=int, default=10000) parser.add_argument("--eval-batch", type=int, default=64) parser.add_argument("--max-test-events", type=int, default=0) parser.add_argument("--device", default="auto") args = parser.parse_args() random.seed(min(args.seeds)) np.random.seed(min(args.seeds)) torch.manual_seed(min(args.seeds)) device = choose_device(args.device) out_dir = args.out_dir out_dir.mkdir(parents=True, exist_ok=True) cache_dir = args.cache_dir or (out_dir / "cache") cache_dir.mkdir(parents=True, exist_ok=True) max_test = None if args.max_test_events == 0 else args.max_test_events test_tag = "all" if max_test is None else str(max_test) test_records = build_records_sequential( args.h5, args.keys, "test", max_test, cache_dir / f"records_test_{test_tag}.json", ) test_access_index = build_h5_access_index( args.h5, test_records, cache_dir / f"h5_access_test_{test_tag}.json", ) model = BRNN().to(device) model.load_state_dict(torch.load(args.base_ckpt, map_location="cpu")) model.eval() thresholds = [round(x, 1) for x in np.arange(0.1, 1.0, 0.1)] rows_long: list[dict] = [] per_seed_rows = [] for seed in args.seeds: eval_cache = cache_dir / ( f"eval_direct_seed{seed}_test{test_tag}_n{args.eval_samples}_" f"len{args.length}_pad{args.padlen}.npz" ) waves, labels, kinds = load_or_materialize_eval_samples( h5_path=args.h5, test_records=test_records, access_index=test_access_index, cache_path=eval_cache, seed=seed + 17, n_samples=args.eval_samples, length=args.length, padlen=args.padlen, ) outputs = run_model(model, waves, device, args.eval_batch) metrics = evaluate_outputs(outputs, labels, kinds, thresholds, min_sep=50, tolerance=100) row = { "slug": "pretrained_direct", "label": "Pretrained direct use", "seed": seed, **metric_row(metrics), } per_seed_rows.append(row) for metric in METRICS: rows_long.append( { "task": "phase", "seed": seed, "condition_slug": "pretrained_direct", "condition_label": "Pretrained direct use", "metric": metric, "value": row[metric], } ) print("direct baseline metrics " + json.dumps(row, ensure_ascii=False), flush=True) write_csvs(rows_long, out_dir) summary = { "experiment": "phase_direct_baseline", "h5": str(args.h5), "keys": str(args.keys), "base_ckpt": str(args.base_ckpt), "seeds": args.seeds, "eval_samples": args.eval_samples, "device": str(device), "rows": per_seed_rows, "metrics_by_model": {"pretrained_direct": strip_errors(metrics)}, "long_csv": str((out_dir / "phase_direct_baseline_long.csv").resolve()), "summary_csv": str((out_dir / "phase_direct_baseline_summary.csv").resolve()), } (out_dir / "summary.json").write_text(json.dumps(summary, indent=2, ensure_ascii=False), encoding="utf-8") print(json.dumps(summary, indent=2, ensure_ascii=False), flush=True) if __name__ == "__main__": main()