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