code/scripts/distance_pick_transfer_experiment.py

#!/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()