| |
| """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() |
|
|
