| |
| """Create the matched SNR filtering design figure for the GRL manuscript.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import math |
| from pathlib import Path |
|
|
| import matplotlib.pyplot as plt |
|
|
|
|
| ROOT = Path(__file__).resolve().parents[1] |
| PHASE_SUMMARY = ROOT / "outputs" / "snr_transfer_seed20260609" / "summary.json" |
| DISP_SUMMARY = ROOT / "outputs" / "disp_snr_transfer_seed20260609" / "summary.json" |
| DISP_SNR_CACHE = ROOT / "outputs" / "disp_snr_transfer_seed20260609" / "ncf_snr_cache.json" |
| DEFAULT_OUTPUT = ROOT / "grl_overleaf" / "figures" / "snr_matched_design.png" |
|
|
|
|
| def read_json(path: Path) -> dict: |
| with path.open("r", encoding="utf-8") as f: |
| return json.load(f) |
|
|
|
|
| def dispersion_candidate_counts(summary: dict) -> list[int]: |
| cache = read_json(DISP_SNR_CACHE) |
| q1 = summary["snr_thresholds"]["q1"] |
| q2 = summary["snr_thresholds"]["q2"] |
| train_snr = [ |
| item["snr_db"] |
| for item in cache.values() |
| if item.get("split") == "train" and math.isfinite(float(item.get("snr_db", float("nan")))) |
| ] |
| return [ |
| len(train_snr), |
| sum(float(value) > q1 for value in train_snr), |
| sum(float(value) > q2 for value in train_snr), |
| ] |
|
|
|
|
| def main() -> None: |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--output", type=Path, default=DEFAULT_OUTPUT) |
| args = parser.parse_args() |
|
|
| phase = read_json(PHASE_SUMMARY) |
| disp = read_json(DISP_SUMMARY) |
|
|
| phase_candidate = [ |
| phase["candidate_train_records"]["full"], |
| phase["candidate_train_records"]["snr5"], |
| phase["candidate_train_records"]["snr10"], |
| ] |
| phase_matched = [row["train_records"] for row in phase["rows"]] |
| disp_matched = [row["train_records"] for row in disp["rows"]] |
| disp_candidate = dispersion_candidate_counts(disp) |
|
|
| labels = ["Full", "SNR > 5", "SNR > 10"] |
| disp_labels = ["Full", "SNR > 3.04", "SNR > 6.77"] |
| colors = ["#2f6f73", "#c47f2b", "#8f4d6b"] |
|
|
| fig, axes = plt.subplots(1, 2, figsize=(10.8, 4.6), dpi=220) |
| panels = [ |
| (axes[0], "A", "Phase picking", labels, phase_candidate, phase_matched, "records"), |
| (axes[1], "B", "Dispersion", disp_labels, disp_candidate, disp_matched, "samples"), |
| ] |
| for ax, panel_label, title, xlabels, candidate, matched, unit in panels: |
| x = range(len(xlabels)) |
| ax.bar(x, candidate, width=0.62, color="#d7d7d7", edgecolor="#555555", label="candidate pool") |
| ax.bar(x, matched, width=0.38, color=colors, edgecolor="#222222", label="matched training") |
| for idx, value in enumerate(candidate): |
| ax.text(idx, value * 1.02, f"{value:,}", ha="center", va="bottom", fontsize=7) |
| for idx, value in enumerate(matched): |
| ax.text(idx, value * 0.50, f"{value:,}", ha="center", va="center", fontsize=8, color="white") |
| ax.set_title(title, fontsize=11) |
| ax.set_xticks(list(x)) |
| ax.set_xticklabels(xlabels, fontsize=8) |
| ax.set_ylabel(f"training {unit}") |
| ax.spines[["top", "right"]].set_visible(False) |
| ax.grid(axis="y", color="#eeeeee", linewidth=0.8) |
| ax.set_axisbelow(True) |
| ax.text( |
| -0.12, |
| 1.08, |
| panel_label, |
| transform=ax.transAxes, |
| fontsize=13, |
| fontweight="bold", |
| va="top", |
| ) |
| axes[0].legend(frameon=False, loc="upper right", fontsize=8) |
| fig.suptitle("Matched training budgets isolate the effect of SNR filtering", fontsize=13, y=0.98) |
| fig.tight_layout(rect=(0, 0, 1, 0.94)) |
|
|
| args.output.parent.mkdir(parents=True, exist_ok=True) |
| fig.savefig(args.output, bbox_inches="tight") |
| print(f"Wrote {args.output}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
