#!/usr/bin/env python3 """Compare near-distance and all-distance transfer training at matched pick counts.""" from __future__ import annotations import argparse import csv import json import math import os import sys from pathlib import Path os.environ.setdefault("KMP_DUPLICATE_LIB_OK", "TRUE") os.environ.setdefault("OMP_NUM_THREADS", "1") os.environ.setdefault("MKL_NUM_THREADS", "1") import torch ROOT = Path(__file__).resolve().parents[1] if str(ROOT) not in sys.path: sys.path.insert(0, str(ROOT)) from scripts import magnitude_pick_transfer_experiment as exp def finite_distance(sample: exp.ManualPickSample) -> bool: return sample.distance_km is not None and sample.distance_km == sample.distance_km def all_pick_samples(record_cache: Path, split: str) -> list[exp.ManualPickSample]: records = exp.load_record_cache(record_cache) samples: list[exp.ManualPickSample] = [] for record in records: distance = record.get("distance_km") for pick in record["phases"]: phase = pick["phase"] samples.append( exp.ManualPickSample( split=split, event=record["event"], station=record["station"], phase=phase, phase_group=exp.PHASE_TO_GROUP[phase], index=int(pick["index"]), length=int(record["length"]), magnitude=float("nan"), distance_km=None if distance is None or not math.isfinite(float(distance)) else float(distance), ) ) return samples def main() -> None: parser = argparse.ArgumentParser() parser.add_argument("--h5", type=Path, default=ROOT / "data" / "credit-x1.h5") parser.add_argument("--base-ckpt", type=Path, default=ROOT / "ckpt" / "pnsn.v3.pt") parser.add_argument("--train-record-cache", type=Path, default=ROOT / "outputs" / "snr_transfer_seed20260609" / "records_train_all.json") parser.add_argument("--test-record-cache", type=Path, default=ROOT / "outputs" / "snr_transfer_seed20260609" / "records_test_all.json") parser.add_argument("--out-dir", type=Path, default=ROOT / "outputs" / "distance150_pick_transfer_seed20260628") parser.add_argument("--seed", type=int, default=20260628) parser.add_argument("--distance-threshold-km", type=float, default=150.0) parser.add_argument("--label-source", choices=["manual", "all"], default="all") parser.add_argument("--length", type=int, default=5120) parser.add_argument("--padlen", type=int, default=512) parser.add_argument("--train-steps", type=int, default=2000) parser.add_argument("--train-batch", type=int, default=16) parser.add_argument("--eval-batch", type=int, default=64) parser.add_argument("--init-mode", choices=["transfer", "scratch"], default="transfer") parser.add_argument("--max-eval-picks", type=int, default=0) parser.add_argument("--lr", type=float, default=1e-4) parser.add_argument("--thresholds", type=float, nargs="+", default=[0.1, 0.2, 0.3, 0.5]) parser.add_argument("--tolerance-samples", type=int, default=100) parser.add_argument("--min-sep", type=int, default=50) parser.add_argument("--resume", action="store_true") parser.add_argument("--device", default=None) args = parser.parse_args() if args.device is None: device = torch.device("mps") if torch.backends.mps.is_available() else torch.device("cpu") else: device = torch.device(args.device) print(f"device={device}", flush=True) args.out_dir.mkdir(parents=True, exist_ok=True) if args.label_source == "manual": train_samples = exp.manual_pick_samples(args.h5, args.train_record_cache, "train", args.out_dir / "manual_train_picks.json") test_samples = exp.manual_pick_samples(args.h5, args.test_record_cache, "test", args.out_dir / "manual_test_picks.json") source_note = "MANUAL.TRAVTIME Pg/Sg/Pn/Sn picks" else: train_samples = all_pick_samples(args.train_record_cache, "train") test_samples = all_pick_samples(args.test_record_cache, "test") source_note = "all cached Pg/Sg/Pn/Sn picks (MANUAL preferred, RNN.tagY used where manual is absent)" train_all_valid = [s for s in train_samples if s.length >= args.length and finite_distance(s)] train_near = [s for s in train_all_valid if float(s.distance_km) <= args.distance_threshold_km] test_near = [ s for s in test_samples if s.length >= args.length and finite_distance(s) and float(s.distance_km) <= args.distance_threshold_km ] test_all_valid = [s for s in test_samples if s.length >= args.length and finite_distance(s)] if args.max_eval_picks > 0: test_near = exp.matched_pick_subset(test_near, min(args.max_eval_picks, len(test_near)), args.seed + 17) test_all_valid = exp.matched_pick_subset( test_all_valid, min(args.max_eval_picks, len(test_all_valid)), args.seed + 18 ) train_budget = len(train_near) train_all_matched = exp.matched_pick_subset(train_all_valid, train_budget, args.seed) distance_slug = f"distance_le_{str(args.distance_threshold_km).replace('.', 'p').rstrip('0').rstrip('p')}" subsets = { "all_distance_matched": train_all_matched, distance_slug: sorted(train_near, key=exp.stable_pick_id), } ckpt_prefix = "pnsn.v3.transfer" if args.init_mode == "transfer" else "pnsn.v3.scratch" print( f"train budget picks={train_budget}; eval <= {args.distance_threshold_km:g} km picks={len(test_near)}; " f"eval all-distance picks={len(test_all_valid)}; init={args.init_mode}", flush=True, ) models: dict[str, exp.BRNN] = {} for slug, subset in subsets.items(): print(f"training {slug}: {exp.summarize_counts(subset)}", flush=True) models[slug] = exp.train_one( args.h5, subset, args.base_ckpt, args.out_dir / f"{ckpt_prefix}.{slug}.pt", args.out_dir / f"train_log_{slug}.csv", args.seed, args.train_steps, args.train_batch, args.length, args.padlen, args.lr, device, args.resume, args.init_mode, ) eval_subsets = { distance_slug: sorted(test_near, key=exp.stable_pick_id), "all_distance": sorted(test_all_valid, key=exp.stable_pick_id), } rows = [] metrics_by_eval = {} for eval_scope, eval_samples in eval_subsets.items(): print(f"evaluating {eval_scope}: {exp.summarize_counts(eval_samples)}", flush=True) waves, rel_indices, groups, phases = exp.materialize_eval(args.h5, eval_samples, args.length, args.padlen) metrics_by_eval[eval_scope] = {} for slug, model in models.items(): outputs = exp.run_model(model, waves, device, args.eval_batch) metrics = exp.evaluate_pick_recall( outputs, rel_indices, groups, phases, args.thresholds, args.tolerance_samples, args.min_sep ) metrics_by_eval[eval_scope][slug] = metrics for scope, labels in [("by_group", metrics["by_group"]), ("by_phase", metrics["by_phase"])]: for label, vals in labels.items(): for row in vals: rows.append({"eval_scope": eval_scope, "model": slug, "scope": scope, "label": label, **row}) for row in metrics["combined"]: rows.append({"eval_scope": eval_scope, "model": slug, "scope": "combined", "label": "all", **row}) safe_distance = str(args.distance_threshold_km).replace(".", "p").rstrip("0").rstrip("p") metrics_path = args.out_dir / f"distance{safe_distance}_pick_metrics.csv" with metrics_path.open("w", newline="") as f: fieldnames = [ "eval_scope", "model", "scope", "label", "threshold", "n", "tp", "fp", "fn", "precision", "recall", "f1", ] writer = csv.DictWriter(f, fieldnames=fieldnames) writer.writeheader() writer.writerows(rows) delta = [] for eval_scope in eval_subsets: for threshold in args.thresholds: for scope, label in [ ("combined", "all"), ("by_group", "P"), ("by_group", "S"), ("by_phase", "Pg"), ("by_phase", "Pn"), ("by_phase", "Sg"), ("by_phase", "Sn"), ]: vals = { r["model"]: r for r in rows if r["eval_scope"] == eval_scope and r["scope"] == scope and r["label"] == label and abs(float(r["threshold"]) - threshold) < 1e-9 } if set(vals) != {"all_distance_matched", distance_slug}: continue all_row = vals["all_distance_matched"] near_row = vals[distance_slug] delta.append( { "eval_scope": eval_scope, "threshold": threshold, "scope": scope, "label": label, "n_test_picks": int(all_row["n"]), "all_distance_precision": float(all_row["precision"]), f"{distance_slug}_precision": float(near_row["precision"]), f"delta_precision_{distance_slug}_minus_all": float(near_row["precision"]) - float(all_row["precision"]), "all_distance_recall": float(all_row["recall"]), f"{distance_slug}_recall": float(near_row["recall"]), f"delta_recall_{distance_slug}_minus_all": float(near_row["recall"]) - float(all_row["recall"]), "all_distance_f1": float(all_row["f1"]), f"{distance_slug}_f1": float(near_row["f1"]), f"delta_f1_{distance_slug}_minus_all": float(near_row["f1"]) - float(all_row["f1"]), } ) delta_path = args.out_dir / f"distance{safe_distance}_pick_delta.csv" with delta_path.open("w", newline="") as f: writer = csv.DictWriter(f, fieldnames=list(delta[0])) writer.writeheader() writer.writerows(delta) summary = { "experiment": f"distance_pick_transfer_{distance_slug}", "notes": [ f"Training/evaluation samples are individual {source_note}.", f"The all-distance and <={args.distance_threshold_km:g} km training conditions are matched to the same pick count.", f"Evaluation is reported on both <={args.distance_threshold_km:g} km and all-distance test-set picks, one centered window per pick.", ], "config": { "seed": args.seed, "distance_threshold_km": args.distance_threshold_km, "label_source": args.label_source, "init_mode": args.init_mode, "base_ckpt": str(args.base_ckpt.resolve()) if args.init_mode == "transfer" else None, "train_steps": args.train_steps, "train_batch": args.train_batch, "length": args.length, "padlen": args.padlen, "thresholds": args.thresholds, "tolerance_samples": args.tolerance_samples, "min_sep": args.min_sep, }, "counts": { "train_all_distance_valid": exp.summarize_counts(train_all_valid), f"train_{distance_slug}_valid": exp.summarize_counts(train_near), "train_all_distance_matched": exp.summarize_counts(train_all_matched), f"test_{distance_slug}_valid": exp.summarize_counts(test_near), "test_all_distance_valid": exp.summarize_counts(test_all_valid), }, "metrics": metrics_by_eval, } summary_path = args.out_dir / "summary.json" summary_path.write_text(json.dumps(summary, indent=2), encoding="utf-8") print(f"wrote {summary_path}") print(f"wrote {metrics_path}") print(f"wrote {delta_path}") if __name__ == "__main__": main()