#!/usr/bin/env python3 """Compute test-SNR-stratified diagnostics for GRL revision outputs.""" from __future__ import annotations import argparse import csv import gzip import math from collections import defaultdict from pathlib import Path import numpy as np DEFAULT_REVISION_DIR = Path("outputs/grl_revision_20260610") def read_csv(path: Path) -> list[dict[str, str]]: opener = gzip.open if path.suffix == ".gz" else open with opener(path, "rt", newline="", encoding="utf-8") as f: return list(csv.DictReader(f)) def assign_quantile_bins(values: dict[int, float], n_bins: int) -> dict[int, str]: finite = np.array([v for v in values.values() if np.isfinite(v)], dtype=float) if finite.size == 0: raise RuntimeError("No finite SNR values available for stratification.") edges = np.unique(np.quantile(finite, np.linspace(0, 1, n_bins + 1))) if edges.size < 3: edges = np.unique(np.linspace(float(np.min(finite)), float(np.max(finite)), min(n_bins, finite.size) + 1)) out = {} for sid, value in values.items(): if not np.isfinite(value): out[sid] = "missing_snr" continue idx = int(np.searchsorted(edges[1:-1], value, side="right")) lo = edges[idx] hi = edges[idx + 1] out[sid] = f"Q{idx + 1}:{lo:.2f}to{hi:.2f}dB" return out def prf(tp: int, fp: int, fn: int) -> dict[str, float]: precision = tp / (tp + fp) if tp + fp else 0.0 recall = tp / (tp + fn) if tp + fn else 0.0 f1 = 2 * precision * recall / (precision + recall) if precision + recall else 0.0 return {"precision": precision, "recall": recall, "f1": f1} def phase_stratified(revision_dir: Path, n_bins: int) -> list[dict]: rows = read_csv(revision_dir / "phase_picking" / "phase_per_window_outputs.csv.gz") snr_by_sample = {} for row in rows: sid = int(row["sample_id"]) try: snr_by_sample[sid] = float(row["test_snr_db"]) except ValueError: snr_by_sample[sid] = float("nan") bins = assign_quantile_bins(snr_by_sample, n_bins) counts = defaultdict(lambda: {"tp": 0, "fp": 0, "fn": 0, "samples": set()}) for row in rows: sid = int(row["sample_id"]) key = (row["condition"], bins[sid], row["phase"]) counts[key]["tp"] += int(row["tp"]) counts[key]["fp"] += int(row["fp"]) counts[key]["fn"] += int(row["fn"]) counts[key]["samples"].add(sid) out = [] for (condition, bin_name, phase), vals in sorted(counts.items()): m = prf(vals["tp"], vals["fp"], vals["fn"]) out.append( { "task": "phase_picking", "condition": condition, "test_snr_bin": bin_name, "phase": phase, "n_test_samples": len(vals["samples"]), "tp": vals["tp"], "fp": vals["fp"], "fn": vals["fn"], **m, } ) by_cond_bin = defaultdict(dict) for row in out: by_cond_bin[(row["condition"], row["test_snr_bin"])][row["phase"]] = row for (condition, bin_name), phase_rows in sorted(by_cond_bin.items()): if "P" in phase_rows and "S" in phase_rows: out.append( { "task": "phase_picking", "condition": condition, "test_snr_bin": bin_name, "phase": "P_S_mean", "n_test_samples": max(phase_rows["P"]["n_test_samples"], phase_rows["S"]["n_test_samples"]), "tp": "", "fp": "", "fn": "", "precision": "", "recall": "", "f1": 0.5 * (phase_rows["P"]["f1"] + phase_rows["S"]["f1"]), } ) return out def dispersion_stratified(revision_dir: Path, n_bins: int) -> list[dict]: rows = read_csv(revision_dir / "dispersion" / "dispersion_per_sample_metrics.csv.gz") snr_by_sample = {} for row in rows: sid = int(row["sample_id"]) snr_by_sample[sid] = float(row["snr_db"]) bins = assign_quantile_bins(snr_by_sample, n_bins) accum = defaultdict(lambda: {"abs": 0.0, "sq": 0.0, "n": 0.0, "samples": set()}) for row in rows: sid = int(row["sample_id"]) key = (row["condition"], bins[sid]) accum[key]["abs"] += float(row["abs_error_sum"]) accum[key]["sq"] += float(row["squared_error_sum"]) accum[key]["n"] += float(row["valid_period_count"]) accum[key]["samples"].add(sid) out = [] for (condition, bin_name), vals in sorted(accum.items()): denom = vals["n"] out.append( { "task": "dispersion", "condition": condition, "test_snr_bin": bin_name, "n_test_samples": len(vals["samples"]), "mae": vals["abs"] / denom if denom else 0.0, "rmse": math.sqrt(vals["sq"] / denom) if denom else 0.0, } ) return out def write_rows(path: Path, rows: list[dict]) -> None: path.parent.mkdir(parents=True, exist_ok=True) if not rows: raise RuntimeError(f"No rows to write for {path}") keys = list(rows[0].keys()) for row in rows: for key in row: if key not in keys: keys.append(key) with path.open("w", newline="", encoding="utf-8") as f: writer = csv.DictWriter(f, fieldnames=keys) writer.writeheader() writer.writerows(rows) def write_summary(path: Path, phase_rows: list[dict], disp_rows: list[dict]) -> None: lines = [ "SNR-stratified diagnostics were computed by quantile-binning the unfiltered test set SNR.", "Phase metrics are recomputed from TP/FP/FN within each bin and condition.", "Dispersion MAE/RMSE are recomputed from summed per-period errors within each bin and condition.", f"Phase rows: {len(phase_rows)}; dispersion rows: {len(disp_rows)}.", ] path.parent.mkdir(parents=True, exist_ok=True) path.write_text("\n".join(lines) + "\n", encoding="utf-8") def main() -> None: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--revision-dir", type=Path, default=DEFAULT_REVISION_DIR) parser.add_argument("--n-bins", type=int, default=4) args = parser.parse_args() tables = args.revision_dir / "tables" phase = phase_stratified(args.revision_dir, args.n_bins) disp = dispersion_stratified(args.revision_dir, args.n_bins) write_rows(tables / "phase_snr_stratified_metrics.csv", phase) write_rows(tables / "dispersion_snr_stratified_metrics.csv", disp) write_summary(tables / "snr_stratified_summary_for_manuscript.txt", phase, disp) if __name__ == "__main__": main()