#!/usr/bin/env python3 """Recompute manuscript statistics from public inputs and regenerate figures. This is the primary GRL reproducibility entrypoint for the SNR-bias package. It does not read the archived plotted-data files in ``results/manuscript_figures``. Instead, it orchestrates the public-data workflow: 1. CREDIT-X1local HDF5 + split keys -> phase-picking SNR caches, matched training/evaluation summaries, and the pretrained direct baseline. 2. SeisDispFusion-NCF HDF5 -> dispersion SNR cache and matched training summaries. 3. SeismicX-Cont picker streams + labels/index -> continuous phase-recall sensitivity. 4. Optional phase-balanced association outputs -> event-geometry Figure 3. CSV/JSON files written under ``--work-dir`` are intermediate products generated by this run, not source plotted data. For a lightweight visual check from the archived plotted-data exports, use ``plot_all_paper_figures.py`` instead. """ from __future__ import annotations import argparse import csv import importlib.util import json import os import subprocess import sys from pathlib import Path from statistics import mean, stdev from typing import Any ROOT = Path(__file__).resolve().parents[2] CODE_ROOT = ROOT / "code" SCRIPTS = CODE_ROOT / "scripts" ODATA = CODE_ROOT / "odata" DEFAULT_SEEDS = [20260609, 20260610, 20260611] PHASE_METRICS = [ "P_precision", "P_recall", "P_f1", "S_precision", "S_recall", "S_f1", "mean_f1", ] DISP_METRICS = ["val_mae", "val_rmse", "val_certainty_f1"] def load_module(name: str, path: Path): spec = importlib.util.spec_from_file_location(name, path) if spec is None or spec.loader is None: raise RuntimeError(f"Cannot load module from {path}") module = importlib.util.module_from_spec(spec) spec.loader.exec_module(module) return module def run(cmd: list[str], *, cwd: Path, env: dict[str, str], dry_run: bool) -> None: printable = " ".join(str(part) for part in cmd) print(f"\n$ {printable}", flush=True) if dry_run: return subprocess.run(cmd, cwd=cwd, env=env, check=True) def require_path(path: Path | None, label: str) -> Path: if path is None: raise SystemExit(f"Missing required input: {label}") if not path.exists(): raise SystemExit(f"Input path does not exist for {label}: {path}") return path def summary_stats(values: list[float]) -> tuple[float, float, int]: if not values: return float("nan"), float("nan"), 0 return float(mean(values)), float(stdev(values)) if len(values) > 1 else 0.0, len(values) def aggregate_seed_summaries( *, task: str, metrics: list[str], source_dirs: list[Path], out_dir: Path, ) -> Path: by_condition: dict[str, dict[str, list[float]]] = {} labels: dict[str, str] = {} long_rows: list[dict[str, Any]] = [] for source_dir in source_dirs: summary_path = source_dir / "summary.json" summary = json.loads(summary_path.read_text(encoding="utf-8")) seed = summary.get("seed", source_dir.name) for row in summary["rows"]: slug = str(row["slug"]) labels.setdefault(slug, str(row.get("label", slug))) by_condition.setdefault(slug, {metric: [] for metric in metrics}) for metric in metrics: value = float(row[metric]) by_condition[slug][metric].append(value) long_rows.append( { "task": task, "seed": seed, "condition_slug": slug, "condition_label": labels[slug], "metric": metric, "value": value, } ) out_dir.mkdir(parents=True, exist_ok=True) long_path = out_dir / f"{task}_multiseed_long.csv" with long_path.open("w", newline="", encoding="utf-8") as handle: writer = csv.DictWriter( handle, fieldnames=["task", "seed", "condition_slug", "condition_label", "metric", "value"], ) writer.writeheader() writer.writerows(long_rows) summary_path = out_dir / f"{task}_multiseed_summary.csv" summary_rows: list[dict[str, Any]] = [] for slug, metric_values in by_condition.items(): for metric, values in metric_values.items(): m, s, n = summary_stats(values) summary_rows.append( { "task": task, "condition_slug": slug, "condition_label": labels.get(slug, slug), "metric": metric, "mean": m, "std": s, "n": n, "values": json.dumps(values), } ) with summary_path.open("w", newline="", encoding="utf-8") as handle: writer = csv.DictWriter( handle, fieldnames=["task", "condition_slug", "condition_label", "metric", "mean", "std", "n", "values"], ) writer.writeheader() writer.writerows(summary_rows) return summary_path def write_manifest(args: argparse.Namespace, work_dir: Path, figure_dir: Path) -> None: manifest = { "workflow": "open-data manuscript figure reproduction", "script": str(Path(__file__).resolve()), "work_dir": str(work_dir.resolve()), "figure_dir": str(figure_dir.resolve()), "seeds": args.seeds, "data_sources": { "phase": { "name": "CREDIT-X1local", "citation": "Li et al. (2024), doi:10.1016/j.eqs.2024.01.018", "h5": str(args.credit_h5) if args.credit_h5 else None, "keys": str(args.credit_keys) if args.credit_keys else None, }, "continuous": { "name": "SeismicX-Cont", "revision": "96367f8", "doi": "10.57967/hf/9006", "pick_dir": str(args.continuous_pick_dir) if args.continuous_pick_dir else None, "label_json": str(args.continuous_label_json) if args.continuous_label_json else None, "waveform_db": str(args.continuous_waveform_db) if args.continuous_waveform_db else None, "phase_balanced_root": str(args.phase_balanced_root) if args.phase_balanced_root else None, }, "dispersion": { "name": "SeisDispFusion-NCF", "revision": "afcd805", "doi": "10.57967/hf/9114", "h5": str(args.ncf_h5) if args.ncf_h5 else None, }, }, "notes": [ "Figures are rendered from intermediate outputs generated in this work directory.", "The archived results/manuscript_figures plotted-data files are not used by this workflow.", ], } (work_dir / "open_data_reproduction_manifest.json").write_text( json.dumps(manifest, indent=2, ensure_ascii=False), encoding="utf-8", ) def run_phase_workflow(args: argparse.Namespace, outputs: Path, env: dict[str, str], dry_run: bool) -> list[Path]: credit_h5 = require_path(args.credit_h5, "CREDIT-X1local HDF5") credit_keys = require_path(args.credit_keys, "CREDIT-X1local split keys") base_ckpt = require_path(args.base_ckpt, "base PNSN checkpoint") cache_dir = outputs / "snr_transfer_phase_balanced_cache" source_dirs = [] for seed in args.seeds: out_dir = outputs / f"snr_transfer_phase_any_seed{seed}" source_dirs.append(out_dir) cmd = [ sys.executable, str(SCRIPTS / "snr_transfer_phase_balanced_experiment.py"), "--h5", str(credit_h5), "--keys", str(credit_keys), "--base-ckpt", str(base_ckpt), "--out-dir", str(out_dir), "--cache-dir", str(cache_dir), "--seed", str(seed), "--train-steps", str(args.phase_train_steps), "--scratch-train-steps", str(args.phase_scratch_train_steps), "--train-batch", str(args.phase_train_batch), "--eval-samples", str(args.phase_eval_samples), "--eval-batch", str(args.phase_eval_batch), "--init-modes", "finetune", "scratch", "--filter-mode", "record-any", "--s-threshold-mode", "same-as-p", "--match-mode", "phase-composition", "--checkpoint-every-steps", str(args.phase_checkpoint_every), ] if args.resume: cmd.append("--resume") cmd.append("--merge-existing-summary") run(cmd, cwd=ROOT, env=env, dry_run=dry_run) direct_dir = outputs / "phase_direct_baseline" cmd = [ sys.executable, str(SCRIPTS / "evaluate_phase_direct_baseline.py"), "--h5", str(credit_h5), "--keys", str(credit_keys), "--base-ckpt", str(base_ckpt), "--out-dir", str(direct_dir), "--cache-dir", str(cache_dir), "--seeds", *[str(seed) for seed in args.seeds], "--eval-samples", str(args.phase_eval_samples), "--eval-batch", str(args.phase_eval_batch), "--device", args.phase_direct_device, ] run(cmd, cwd=ROOT, env=env, dry_run=dry_run) if not dry_run: aggregate_seed_summaries( task="phase", metrics=PHASE_METRICS, source_dirs=source_dirs, out_dir=outputs / "grl_phase_any_multiseed_seed20260609_20260611", ) return source_dirs def run_dispersion_workflow(args: argparse.Namespace, outputs: Path, env: dict[str, str], dry_run: bool) -> list[Path]: ncf_h5 = require_path(args.ncf_h5, "SeisDispFusion-NCF HDF5") source_dirs = [] for seed in args.seeds: out_dir = outputs / f"disp_snr_transfer_seed{seed}" source_dirs.append(out_dir) cmd = [ sys.executable, str(SCRIPTS / "disp_snr_transfer_experiment.py"), "--h5", str(ncf_h5), "--out-dir", str(out_dir), "--seed", str(seed), "--mode", "scratch", "--epochs", str(args.disp_epochs), "--batch-size", str(args.disp_batch_size), "--num-workers", str(args.disp_num_workers), "--device", args.disp_device, ] run(cmd, cwd=ROOT, env=env, dry_run=dry_run) if not dry_run: aggregate_seed_summaries( task="dispersion", metrics=DISP_METRICS, source_dirs=source_dirs, out_dir=outputs / "grl_multiseed_seed20260609_20260611", ) return source_dirs def run_continuous_recall(args: argparse.Namespace, outputs: Path, env: dict[str, str], dry_run: bool) -> Path: pick_dir = require_path(args.continuous_pick_dir, "SeismicX-Cont continuous picker JSONL directory") label_json = require_path(args.continuous_label_json, "SeismicX-Cont label JSON") waveform_db = require_path(args.continuous_waveform_db, "SeismicX-Cont waveform index SQLite") out_dir = outputs / "continuous_phase_recall_snr_conf_sweep" cmd = [ sys.executable, str(SCRIPTS / "continuous_phase_recall_snr_conf_sweep.py"), "--root", str(args.continuous_root or pick_dir.parent), "--pick-dir", str(pick_dir), "--label-json", str(label_json), "--waveform-db", str(waveform_db), "--eval-script", str(ODATA / "evaluate_pick_recall_no_nms.py"), "--outdir", str(out_dir), "--days", *args.continuous_days, "--thresholds", args.continuous_thresholds, ] if args.include_validation_thresholds: cmd.append("--include-validation-thresholds") if args.force_rebuild_continuous_index: cmd.append("--force-rebuild") run(cmd, cwd=ROOT, env=env, dry_run=dry_run) return out_dir / "continuous_phase_recall_snr_conf_sweep.csv" def run_precision_table(args: argparse.Namespace, outputs: Path, env: dict[str, str], dry_run: bool) -> None: credit_h5 = require_path(args.credit_h5, "CREDIT-X1local HDF5") credit_keys = require_path(args.credit_keys, "CREDIT-X1local split keys") source_dirs = [outputs / f"snr_transfer_phase_any_seed{seed}" for seed in args.seeds] cmd = [ sys.executable, str(SCRIPTS / "snr_filtered_test_precision_table.py"), "--h5", str(credit_h5), "--keys", str(credit_keys), "--out-dir", str(outputs / "snr_filtered_test_precision_table"), "--source-dirs", *[str(path) for path in source_dirs], "--cache-dir", str(outputs / "snr_transfer_phase_balanced_cache"), "--eval-samples", str(args.phase_eval_samples), ] run(cmd, cwd=ROOT, env=env, dry_run=dry_run) def export_training_manifests(args: argparse.Namespace, outputs: Path, env: dict[str, str], dry_run: bool) -> None: cmd = [ sys.executable, str(SCRIPTS / "export_training_manifests.py"), "--out-dir", str(args.work_dir / "training_manifests"), "--seeds", *[str(seed) for seed in args.seeds], ] has_sources = False phase_records = outputs / "snr_transfer_phase_balanced_cache" / "records_train_all.json" phase_snr = outputs / "snr_transfer_phase_balanced_cache" / "train_phase_snr_db.json" if phase_records.exists() and phase_snr.exists(): cmd.extend(["--phase-records-json", str(phase_records), "--phase-snr-json", str(phase_snr)]) has_sources = True elif not args.skip_phase_training: raise FileNotFoundError(f"Cannot export phase manifests; missing {phase_records} or {phase_snr}") dispersion_snr = outputs / f"disp_snr_transfer_seed{args.seeds[0]}" / "ncf_snr_cache.json" if dispersion_snr.exists(): cmd.extend(["--dispersion-snr-json", str(dispersion_snr)]) has_sources = True elif not args.skip_dispersion_training: raise FileNotFoundError(f"Cannot export dispersion manifests; missing {dispersion_snr}") if not has_sources: print("Skipping training-manifest export because no source caches are available.", flush=True) return run(cmd, cwd=ROOT, env=env, dry_run=dry_run) def render_figures(args: argparse.Namespace, outputs: Path, figure_dir: Path) -> None: figure_dir.mkdir(parents=True, exist_ok=True) ncf_h5 = require_path(args.ncf_h5, "SeisDispFusion-NCF HDF5 for Figure 1 period bins") observability = load_module("grl_observability_open_data", SCRIPTS / "grl_make_observability_real_data_figure.py") observability.PHASE_RECORDS = outputs / "snr_transfer_phase_balanced_cache" / "records_train_all.json" observability.PHASE_SNR = outputs / "snr_transfer_phase_balanced_cache" / "train_phase_snr_db.json" observability.DISP_CACHE = outputs / f"disp_snr_transfer_seed{args.seeds[0]}" / "ncf_snr_cache.json" observability.DISP_H5 = ncf_h5 observability.FIG_DIR = figure_dir observability.OUT_PDF = figure_dir / "fig_observability_real_data_v1.pdf" observability.OUT_PNG = figure_dir / "fig_observability_real_data_v1.png" observability.OUT_CSV = figure_dir / "fig_observability_real_data_v1_data.csv" observability.OUT_JSON = figure_dir / "fig_observability_real_data_v1_summary.json" observability.make_figure() learning = load_module("grl_learning_open_data", SCRIPTS / "grl_make_learning_summary_figure.py") learning.PHASE_RECORDS = outputs / "snr_transfer_phase_balanced_cache" / "records_train_all.json" learning.PHASE_SNR = outputs / "snr_transfer_phase_balanced_cache" / "train_phase_snr_db.json" learning.DISP_SNR = outputs / f"disp_snr_transfer_seed{args.seeds[0]}" / "ncf_snr_cache.json" learning.PHASE_SUMMARY = outputs / "grl_phase_any_multiseed_seed20260609_20260611" / "phase_multiseed_summary.csv" learning.PHASE_DIRECT_BASELINE = outputs / "phase_direct_baseline" / "phase_direct_baseline_summary.csv" learning.DISP_SUMMARY = outputs / "grl_multiseed_seed20260609_20260611" / "dispersion_multiseed_summary.csv" learning.CONTINUOUS_SWEEP = outputs / "continuous_phase_recall_snr_conf_sweep" / "continuous_phase_recall_snr_conf_sweep.csv" learning.FIG_DIR = figure_dir learning.OUT_PDF = figure_dir / "fig_learning_selection_generalization_summary_v2.pdf" learning.OUT_PNG = figure_dir / "fig_learning_selection_generalization_summary_v2.png" learning.OUT_DATA = figure_dir / "fig_learning_selection_generalization_summary_v2_data.csv" learning.main() if args.skip_phase_balanced_geometry: print("Skipping phase-balanced event geometry figure by request.", flush=True) return phase_balanced_root = require_path(args.phase_balanced_root, "phase-balanced association output root") label_json = require_path(args.continuous_label_json, "SeismicX-Cont label JSON for Figure 3") waveform_db = require_path(args.continuous_waveform_db, "SeismicX-Cont waveform index SQLite for Figure 3") geometry = load_module("grl_geometry_open_data", SCRIPTS / "grl_make_phase_balanced_event_geometry_figure.py") geometry.PHASE_BALANCED_ROOT = phase_balanced_root geometry.PICK_EVAL_SCRIPT = ODATA / "evaluate_pick_recall_no_nms.py" geometry.LABEL_JSON = label_json geometry.WAVEFORM_DB = waveform_db geometry.FIG_DIR = figure_dir geometry.FIG_PDF = figure_dir / "fig_event_geometry_distribution_polished.pdf" geometry.FIG_PNG = figure_dir / "fig_event_geometry_distribution_polished.png" geometry.DATA_CSV = figure_dir / "fig_event_geometry_distribution_polished_data.csv" geometry.SUMMARY_JSON = figure_dir / "fig_event_geometry_distribution_polished_summary.json" geometry.main() def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--credit-h5", type=Path) parser.add_argument("--credit-keys", type=Path) parser.add_argument("--ncf-h5", type=Path) parser.add_argument("--continuous-root", type=Path, default=None) parser.add_argument("--continuous-pick-dir", type=Path, default=None) parser.add_argument("--continuous-label-json", type=Path, default=None) parser.add_argument("--continuous-waveform-db", type=Path, default=None) parser.add_argument("--phase-balanced-root", type=Path, default=None) parser.add_argument("--base-ckpt", type=Path, default=ROOT / "checkpoints" / "base" / "pnsn.v3.pt") parser.add_argument("--work-dir", type=Path, default=ROOT / "open_data_work") parser.add_argument("--figure-dir", type=Path, default=None) parser.add_argument("--seeds", nargs="+", type=int, default=DEFAULT_SEEDS) parser.add_argument("--resume", action="store_true") parser.add_argument("--dry-run", action="store_true") parser.add_argument("--skip-phase-training", action="store_true") parser.add_argument("--skip-dispersion-training", action="store_true") parser.add_argument("--skip-continuous-sweep", action="store_true") parser.add_argument("--skip-precision-table", action="store_true") parser.add_argument("--skip-training-manifest-export", action="store_true") parser.add_argument("--skip-figures", action="store_true") parser.add_argument("--skip-phase-balanced-geometry", action="store_true") parser.add_argument("--phase-train-steps", type=int, default=2000) parser.add_argument("--phase-scratch-train-steps", type=int, default=10000) parser.add_argument("--phase-train-batch", type=int, default=16) parser.add_argument("--phase-eval-samples", type=int, default=10000) parser.add_argument("--phase-eval-batch", type=int, default=64) parser.add_argument("--phase-checkpoint-every", type=int, default=100) parser.add_argument("--phase-direct-device", default="auto") parser.add_argument("--disp-epochs", type=int, default=5) parser.add_argument("--disp-batch-size", type=int, default=256) parser.add_argument("--disp-num-workers", type=int, default=0) parser.add_argument("--disp-device", default="cpu") parser.add_argument("--continuous-days", nargs="+", default=["20190706", "20211113"]) parser.add_argument("--continuous-thresholds", default="0:10:1") parser.add_argument("--include-validation-thresholds", action="store_true") parser.add_argument("--force-rebuild-continuous-index", action="store_true") return parser.parse_args() def main() -> None: args = parse_args() work_dir = args.work_dir outputs = work_dir / "outputs" figure_dir = args.figure_dir or (work_dir / "figures") work_dir.mkdir(parents=True, exist_ok=True) outputs.mkdir(parents=True, exist_ok=True) write_manifest(args, work_dir, figure_dir) env = os.environ.copy() env["PYTHONPATH"] = str(CODE_ROOT) + os.pathsep + env.get("PYTHONPATH", "") env.setdefault("OMP_NUM_THREADS", "1") env.setdefault("MKL_NUM_THREADS", "1") env.setdefault("KMP_DUPLICATE_LIB_OK", "TRUE") if not args.skip_phase_training: run_phase_workflow(args, outputs, env, args.dry_run) if not args.skip_dispersion_training: run_dispersion_workflow(args, outputs, env, args.dry_run) if not args.skip_continuous_sweep: run_continuous_recall(args, outputs, env, args.dry_run) if not args.skip_training_manifest_export: export_training_manifests(args, outputs, env, args.dry_run) if not args.skip_precision_table: run_precision_table(args, outputs, env, args.dry_run) if not args.skip_figures and not args.dry_run: render_figures(args, outputs, figure_dir) elif args.dry_run: print("Dry run complete; no commands were executed and no figures were rendered.", flush=True) if not args.dry_run: print(f"\nOpen-data reproduction outputs: {work_dir.resolve()}", flush=True) print(f"Regenerated figures: {figure_dir.resolve()}", flush=True) if __name__ == "__main__": main()