code/scripts/evaluate_phase_direct_baseline.py

#!/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()