code/scripts/grl_export_per_sample_outputs.py

#!/usr/bin/env python3
"""Export per-window and per-sample outputs for the GRL revision analyses.

This script does not retrain models. It reuses the saved checkpoints from the
matched-budget experiments and reconstructs the deterministic test windows or
test samples used by the original summaries. The derived files are intended as
inputs for paired bootstrap confidence intervals and SNR-stratified diagnostics.
"""

from __future__ import annotations

import argparse
import csv
import datetime as dt
import gzip
import json
import math
import os
import platform
import subprocess
import sys
from pathlib import Path
from typing import Iterable

os.environ.setdefault("KMP_DUPLICATE_LIB_OK", "TRUE")
os.environ.setdefault("OMP_NUM_THREADS", "1")
os.environ.setdefault("MKL_NUM_THREADS", "1")

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 scripts.reproduce_paper_stats import GROUP_TO_CHANNELS, find_peaks, make_specs, match_predictions, materialize_samples, run_model
from scripts.snr_transfer_experiment import compute_record_snr, record_from_dict
from scripts.disp_snr_transfer_experiment import compute_snr_cache, load_model as load_disp_model, load_v23_module, make_loader
from models.BRNNPNSN import BRNN


DEFAULT_REVISION_DIR = ROOT / "outputs" / "grl_revision_20260610"


def git_hash() -> str | None:
    try:
        return subprocess.check_output(["git", "rev-parse", "HEAD"], cwd=ROOT, text=True).strip()
    except Exception:
        return None


def command_line() -> list[str]:
    return [sys.executable, *sys.argv]


def write_metadata(path: Path, payload: dict) -> None:
    payload = {
        "generated_at_utc": dt.datetime.utcnow().isoformat(timespec="seconds") + "Z",
        "git_commit": git_hash(),
        "command": command_line(),
        "python": sys.version.replace("\n", " "),
        "platform": platform.platform(),
        "torch": torch.__version__,
        **payload,
    }
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(payload, indent=2, ensure_ascii=False), encoding="utf-8")


def load_cached_records(path: Path):
    if not path.exists():
        raise FileNotFoundError(f"Missing cached record file: {path}. Run scripts/snr_transfer_experiment.py first.")
    payload = json.loads(path.read_text(encoding="utf-8"))
    return [record_from_dict(row) for row in payload["records"]]


def phase_sample_metadata(records, specs, snr_by_record: dict[str, float]) -> list[dict]:
    rows = []
    for sample_id, spec in enumerate(specs):
        rec_ids, events, stations, starts, crop_snrs = [], [], [], [], []
        for crop in spec.crops:
            record = records[crop.rec_idx]
            rec_id = f"{record.event}/{record.station}"
            rec_ids.append(rec_id)
            events.append(record.event)
            stations.append(record.station)
            starts.append(crop.start)
            val = snr_by_record.get(rec_id)
            if val is not None and math.isfinite(float(val)):
                crop_snrs.append(float(val))
        rows.append(
            {
                "sample_id": sample_id,
                "window_kind": spec.kind,
                "record_ids": ";".join(rec_ids),
                "event_ids": ";".join(events),
                "station_ids": ";".join(stations),
                "crop_starts": ";".join(map(str, starts)),
                "test_snr_db": max(crop_snrs) if crop_snrs else float("nan"),
            }
        )
    return rows


def json_list(values: Iterable[int | float]) -> str:
    return json.dumps(list(values), separators=(",", ":"))


def export_phase(args: argparse.Namespace) -> None:
    out_dir = args.out_dir / "phase_picking"
    out_dir.mkdir(parents=True, exist_ok=True)
    phase_src = Path(args.phase_out_dir)
    records = load_cached_records(phase_src / "records_test_all.json")
    snr_by_record = compute_record_snr(Path(args.phase_h5), records, out_dir / "test_record_snr_db.json")
    specs = make_specs(
        records,
        n_samples=args.eval_samples,
        seed=args.seed + 17,
        length=args.phase_length,
        padlen=args.phase_padlen,
        double_prob=0.5,
    )
    waves, labels_all, _, kinds = materialize_samples(Path(args.phase_h5), records, specs, args.phase_length)
    sample_meta = phase_sample_metadata(records, specs, snr_by_record)
    device = torch.device(args.device)

    checkpoints = {
        "full": phase_src / "pnsn.v3.transfer.full.pt",
        "snr5": phase_src / "pnsn.v3.transfer.snr5.pt",
        "snr10": phase_src / "pnsn.v3.transfer.snr10.pt",
    }
    csv_path = out_dir / "phase_per_window_outputs.csv.gz"
    fieldnames = [
        "condition",
        "sample_id",
        "window_kind",
        "record_ids",
        "event_ids",
        "station_ids",
        "crop_starts",
        "test_snr_db",
        "phase",
        "true_indices",
        "pred_indices",
        "pred_probs",
        "tp",
        "fp",
        "fn",
    ]
    with gzip.open(csv_path, "wt", newline="", encoding="utf-8") as f:
        writer = csv.DictWriter(f, fieldnames=fieldnames)
        writer.writeheader()
        for condition, ckpt in checkpoints.items():
            if not ckpt.exists():
                raise FileNotFoundError(f"Missing phase checkpoint for {condition}: {ckpt}")
            model = BRNN().to(device)
            model.load_state_dict(torch.load(ckpt, map_location="cpu"))
            outputs = run_model(model, waves, device, args.phase_batch_size)
            for i, labels in enumerate(labels_all):
                for phase_group in ("P", "S"):
                    true_idx = [idx for _, group, idx in labels if group == phase_group]
                    prob = outputs[i, GROUP_TO_CHANNELS[phase_group], :].max(axis=0)
                    peaks = find_peaks(prob, args.phase_threshold, args.phase_min_sep)
                    pred_idx = [idx for idx, _ in peaks]
                    pred_prob = [prob for _, prob in peaks]
                    tp, fp, fn, _ = match_predictions(true_idx, pred_idx, args.phase_tolerance)
                    writer.writerow(
                        {
                            **sample_meta[i],
                            "condition": condition,
                            "phase": phase_group,
                            "true_indices": json_list(true_idx),
                            "pred_indices": json_list(pred_idx),
                            "pred_probs": json_list(round(float(x), 6) for x in pred_prob),
                            "tp": tp,
                            "fp": fp,
                            "fn": fn,
                        }
                    )
    write_metadata(
        out_dir / "metadata.json",
        {
            "task": "phase_picking",
            "input_h5": str(Path(args.phase_h5)),
            "source_output_dir": str(phase_src),
            "checkpoints": {k: str(v) for k, v in checkpoints.items()},
            "eval_samples": args.eval_samples,
            "threshold": args.phase_threshold,
            "min_sep_samples": args.phase_min_sep,
            "tolerance_samples": args.phase_tolerance,
            "output_csv": str(csv_path),
        },
    )
    print(f"[phase] wrote {csv_path}")


def export_dispersion(args: argparse.Namespace) -> None:
    out_dir = args.out_dir / "dispersion"
    out_dir.mkdir(parents=True, exist_ok=True)
    disp_src = Path(args.disp_out_dir)
    h5_path = Path(args.disp_h5)
    snr_rows = compute_snr_cache(h5_path, disp_src / "ncf_snr_cache.json")
    test_keys = sorted(
        key for key, row in snr_rows.items() if row["split"] == "test" and np.isfinite(float(row["snr_db"]))
    )
    v23 = load_v23_module()
    device = torch.device(args.device)
    ckpts = {
        "full": disp_src / "dispnet.v2.3.transfer.full.pt",
        "snr_q1": disp_src / "dispnet.v2.3.transfer.snr_q1.pt",
        "snr_q2": disp_src / "dispnet.v2.3.transfer.snr_q2.pt",
    }
    csv_path = out_dir / "dispersion_per_sample_metrics.csv.gz"
    fieldnames = [
        "condition",
        "sample_id",
        "key",
        "snr_db",
        "distance_km",
        "valid_period_count",
        "abs_error_sum",
        "squared_error_sum",
        "sample_mae",
        "sample_rmse",
    ]
    with gzip.open(csv_path, "wt", newline="", encoding="utf-8") as f:
        writer = csv.DictWriter(f, fieldnames=fieldnames)
        writer.writeheader()
        for condition, ckpt_path in ckpts.items():
            if not ckpt_path.exists():
                raise FileNotFoundError(f"Missing dispersion checkpoint for {condition}: {ckpt_path}")
            ckpt = torch.load(ckpt_path, map_location="cpu")
            model, cfg = load_disp_model(v23, ckpt, device)
            model.eval()
            loader = make_loader(
                h5_path,
                test_keys,
                args.disp_batch_size,
                args.num_workers,
                cfg.waveform_length,
                False,
                args.seed,
                False,
            )
            sample_offset = 0
            with torch.no_grad():
                for batch in loader:
                    waveform = batch["waveform"].float().to(device)
                    if waveform.ndim == 3 and waveform.size(1) == 1:
                        waveform = waveform.squeeze(1)
                    true = batch["disp"].float().to(device)
                    mask = batch["mask"].float().to(device)
                    pred = model(waveform)["disp_mu"]
                    abs_err = (pred - true).abs() * mask
                    sq_err = (pred - true).pow(2) * mask
                    valid = mask.sum(dim=1).clamp_min(1.0)
                    mae = abs_err.sum(dim=1) / valid
                    rmse = torch.sqrt(sq_err.sum(dim=1) / valid)
                    for j, key in enumerate(batch["key"]):
                        row = snr_rows[key]
                        writer.writerow(
                            {
                                "condition": condition,
                                "sample_id": sample_offset + j,
                                "key": key,
                                "snr_db": float(row["snr_db"]),
                                "distance_km": float(row.get("distance_km", float("nan"))),
                                "valid_period_count": int(valid[j].detach().cpu().item()),
                                "abs_error_sum": float(abs_err[j].sum().detach().cpu().item()),
                                "squared_error_sum": float(sq_err[j].sum().detach().cpu().item()),
                                "sample_mae": float(mae[j].detach().cpu().item()),
                                "sample_rmse": float(rmse[j].detach().cpu().item()),
                            }
                        )
                    sample_offset += len(batch["key"])
    write_metadata(
        out_dir / "metadata.json",
        {
            "task": "dispersion",
            "input_h5": str(h5_path),
            "source_output_dir": str(disp_src),
            "checkpoints": {k: str(v) for k, v in ckpts.items()},
            "test_samples": len(test_keys),
            "output_csv": str(csv_path),
        },
    )
    print(f"[dispersion] wrote {csv_path}")


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--task", choices=["phase", "dispersion", "both"], default="both")
    parser.add_argument("--out-dir", type=Path, default=DEFAULT_REVISION_DIR)
    parser.add_argument("--seed", type=int, default=20260609)
    parser.add_argument("--device", default="cpu")
    parser.add_argument("--num-workers", type=int, default=0)
    parser.add_argument("--phase-h5", default="data/credit-x1.h5")
    parser.add_argument("--phase-out-dir", default="outputs/snr_transfer_seed20260609")
    parser.add_argument("--eval-samples", type=int, default=10000)
    parser.add_argument("--phase-length", type=int, default=5120)
    parser.add_argument("--phase-padlen", type=int, default=512)
    parser.add_argument("--phase-batch-size", type=int, default=64)
    parser.add_argument("--phase-threshold", type=float, default=0.1)
    parser.add_argument("--phase-min-sep", type=int, default=50)
    parser.add_argument("--phase-tolerance", type=int, default=100)
    parser.add_argument("--disp-h5", default="data/ncf_data/ncf_disp_dataset_with_disp_image.h5")
    parser.add_argument("--disp-out-dir", default="outputs/disp_snr_transfer_seed20260609")
    parser.add_argument("--disp-batch-size", type=int, default=64)
    args = parser.parse_args()

    if args.task in ("phase", "both"):
        export_phase(args)
    if args.task in ("dispersion", "both"):
        export_dispersion(args)


if __name__ == "__main__":
    main()