#!/usr/bin/env python3 """Fine-tune the Pn/Sn picker with phase-aware SNR-filtered labels. This variant keeps the original P thresholds (5 and 10 dB) but chooses lower S thresholds so that the retained P- and S-label counts are approximately balanced in the medium- and high-SNR training pools. Unlike the original record-level SNR experiment, filtering is applied at the phase-label level: labels below their phase-specific threshold are removed from the training target, and records with no retained labels are excluded. For stricter P/S joint training, use ``--filter-mode record-complete`` with ``--s-threshold-mode same-as-p``. That mode keeps a training record only when every original P/S label in the record has finite SNR and passes its threshold, and it keeps the full original label set in the target. Add ``--match-mode phase-composition`` when P-only, P+S, and S-only waveform counts must be identical across the full, medium-SNR, and high-SNR training pools. Use ``--filter-mode record-any`` when a less restrictive record-level filter is desired: if any original P/S label in the record passes its phase-specific SNR threshold, the whole waveform record is kept with all original labels. """ from __future__ import annotations import argparse import csv import json import math import random import sys from collections import Counter from pathlib import Path from typing import Any import h5py import numpy as np import torch ROOT = Path(__file__).resolve().parents[1] if str(ROOT) not in sys.path: sys.path.insert(0, str(ROOT)) from models.BRNNPNSN import BRNN, Loss from scripts.reproduce_paper_stats import ( COMPONENTS, DTYPE, PHASE_TO_GROUP, PhasePick, Record, evaluate_outputs, labels_to_target, load_crop, make_specs, materialize_samples, normalize_wave, run_model, strip_errors, ) from scripts.snr_transfer_experiment import ( build_records_sequential, component_keys, matched_records, metric_row, pick_snr_db, plot_summary, sample_batch, write_metrics_table, ) def record_key(record: Record) -> str: return f"{record.event}/{record.station}" H5AccessIndex = dict[str, tuple[str, str, str]] def build_h5_access_index(h5_path: Path, records: list[Record], cache_path: Path) -> H5AccessIndex: if cache_path.exists(): raw = json.loads(cache_path.read_text(encoding="utf-8")) paths = raw.get("paths", raw) print(f"loaded HDF5 access index {cache_path} entries={len(paths):,}", flush=True) return {str(k): tuple(v) for k, v in paths.items()} print(f"building HDF5 access index {cache_path} from local record keys", flush=True) paths: dict[str, list[str]] = { record_key(record): [f"/{record.event}/{record.station}/{component}" for component in COMPONENTS] for record in records } cache_path.parent.mkdir(parents=True, exist_ok=True) tmp = cache_path.with_suffix(".tmp") tmp.write_text( json.dumps({"h5": str(h5_path), "record_count": len(records), "paths": paths}, ensure_ascii=False), encoding="utf-8", ) tmp.replace(cache_path) return {k: tuple(v) for k, v in paths.items()} def phase_snr_key(record: Record, pick: PhasePick, pick_index: int) -> str: return f"{record.event}/{record.station}/{pick_index}:{pick.phase}:{pick.index}:{pick.source}" def phase_counts(records: list[Record]) -> dict[str, int]: counts = Counter() for record in records: for pick in record.phases: counts[PHASE_TO_GROUP[pick.phase]] += 1 return {"P": int(counts["P"]), "S": int(counts["S"]), "total": int(counts["P"] + counts["S"])} PHASE_COMPOSITION_ORDER = ("P_only", "PS", "S_only", "none") def record_phase_composition(record: Record) -> str: groups = {PHASE_TO_GROUP[pick.phase] for pick in record.phases} if groups == {"P", "S"}: return "PS" if groups == {"P"}: return "P_only" if groups == {"S"}: return "S_only" return "none" def phase_composition_counts(records: list[Record]) -> dict[str, int]: counts = Counter(record_phase_composition(record) for record in records) return {key: int(counts[key]) for key in PHASE_COMPOSITION_ORDER} def matched_records_by_phase_composition( records: list[Record], targets: dict[str, int], seed: int, ) -> list[Record]: by_composition: dict[str, list[Record]] = {key: [] for key in PHASE_COMPOSITION_ORDER} for record in records: by_composition[record_phase_composition(record)].append(record) rng = np.random.default_rng(seed) selected: list[Record] = [] for key in PHASE_COMPOSITION_ORDER: target = int(targets.get(key, 0)) if target == 0: continue pool = sorted(by_composition[key], key=record_key) if len(pool) < target: raise RuntimeError(f"Cannot draw {target} {key} records from a pool of {len(pool)} records.") idx = np.sort(rng.choice(len(pool), size=target, replace=False)) selected.extend(pool[int(i)] for i in idx) return sorted(selected, key=record_key) def compute_phase_snr_cache(h5_path: Path, records: list[Record], cache_path: Path) -> dict[str, float]: if cache_path.exists(): return {k: float(v) for k, v in json.loads(cache_path.read_text(encoding="utf-8")).items()} cache_path.parent.mkdir(parents=True, exist_ok=True) values: dict[str, float] = {} with h5py.File(h5_path, "r") as h5: for i, record in enumerate(records): station = h5[record.event][record.station] comps = component_keys(station) if comps is None: continue waves = [station[c][:] for c in comps] for pick_index, pick in enumerate(record.phases): snr = pick_snr_db(waves, pick.phase, pick.index) if snr is not None and math.isfinite(snr): values[phase_snr_key(record, pick, pick_index)] = float(snr) if i % 5000 == 0: print(f"computed phase SNR for {i:,}/{len(records):,} records", flush=True) tmp = cache_path.with_suffix(".tmp") tmp.write_text(json.dumps(values, ensure_ascii=False), encoding="utf-8") tmp.replace(cache_path) return values def collect_phase_snr(records: list[Record], phase_snr: dict[str, float]) -> dict[str, list[float]]: by_group = {"P": [], "S": []} for record in records: for pick_index, pick in enumerate(record.phases): value = phase_snr.get(phase_snr_key(record, pick, pick_index)) if value is None: continue by_group[PHASE_TO_GROUP[pick.phase]].append(float(value)) return by_group def choose_s_threshold(s_values: list[float], target_count: int) -> tuple[float, int]: if target_count <= 0 or not s_values: raise RuntimeError("Cannot choose S threshold without P target and S values.") ordered = sorted(float(v) for v in s_values if math.isfinite(v)) if target_count >= len(ordered): return ordered[0] - 1e-6, len(ordered) # Count kept is values >= threshold. Use the boundary value giving the # nearest count; subtract a tiny epsilon so equality is included. threshold = ordered[-target_count] - 1e-6 kept = sum(v >= threshold for v in ordered) return float(threshold), int(kept) def filter_records_by_phase_thresholds( records: list[Record], phase_snr: dict[str, float], p_threshold: float | None, s_threshold: float | None, ) -> list[Record]: if p_threshold is None and s_threshold is None: return list(records) filtered: list[Record] = [] for record in records: kept: list[PhasePick] = [] for pick_index, pick in enumerate(record.phases): value = phase_snr.get(phase_snr_key(record, pick, pick_index)) if value is None: continue group = PHASE_TO_GROUP[pick.phase] threshold = p_threshold if group == "P" else s_threshold if threshold is not None and value >= threshold: kept.append(pick) if kept: filtered.append( Record( event=record.event, station=record.station, length=record.length, delta=record.delta, distance_km=record.distance_km, phases=tuple(kept), ) ) return filtered def filter_records_by_complete_thresholds( records: list[Record], phase_snr: dict[str, float], p_threshold: float | None, s_threshold: float | None, ) -> list[Record]: if p_threshold is None and s_threshold is None: return list(records) filtered: list[Record] = [] for record in records: keep_record = True for pick_index, pick in enumerate(record.phases): value = phase_snr.get(phase_snr_key(record, pick, pick_index)) if value is None or not math.isfinite(value): keep_record = False break group = PHASE_TO_GROUP[pick.phase] threshold = p_threshold if group == "P" else s_threshold if threshold is not None and value < threshold: keep_record = False break if keep_record: filtered.append(record) return filtered def filter_records_by_any_thresholds( records: list[Record], phase_snr: dict[str, float], p_threshold: float | None, s_threshold: float | None, ) -> list[Record]: if p_threshold is None and s_threshold is None: return list(records) filtered: list[Record] = [] for record in records: keep_record = False for pick_index, pick in enumerate(record.phases): value = phase_snr.get(phase_snr_key(record, pick, pick_index)) if value is None or not math.isfinite(value): continue group = PHASE_TO_GROUP[pick.phase] threshold = p_threshold if group == "P" else s_threshold if threshold is not None and value >= threshold: keep_record = True break if keep_record: filtered.append(record) return filtered def summarize_snr(values: list[float]) -> dict[str, float | int]: arr = np.asarray(values, dtype=float) arr = arr[np.isfinite(arr)] if arr.size == 0: return {"count": 0} return { "count": int(arr.size), "min": float(np.min(arr)), "p25": float(np.percentile(arr, 25)), "median": float(np.median(arr)), "p75": float(np.percentile(arr, 75)), "p95": float(np.percentile(arr, 95)), "max": float(np.max(arr)), } def write_phase_balance_table(rows: list[dict[str, Any]], out: Path) -> None: with out.open("w", newline="", encoding="utf-8") as f: writer = csv.DictWriter( f, fieldnames=[ "slug", "label", "p_threshold_db", "s_threshold_db", "candidate_records", "candidate_P_only_records", "candidate_PS_records", "candidate_S_only_records", "candidate_none_records", "candidate_P_labels", "candidate_S_labels", "matched_records", "matched_P_only_records", "matched_PS_records", "matched_S_only_records", "matched_none_records", "matched_P_labels", "matched_S_labels", ], ) writer.writeheader() writer.writerows(rows) def load_or_materialize_eval_samples( h5_path: Path, test_records: list[Record], access_index: H5AccessIndex, cache_path: Path, seed: int, n_samples: int, length: int, padlen: int, ) -> tuple[np.ndarray, list[list[tuple[str, str, int]]], list[str]]: if cache_path.exists(): print(f"loading cached eval samples {cache_path}", flush=True) data = np.load(cache_path, allow_pickle=True) return data["waves"], data["labels"].tolist(), data["kinds"].tolist() print(f"materializing eval samples {cache_path}", flush=True) eval_specs = make_specs( test_records, n_samples=n_samples, seed=seed, length=length, padlen=padlen, double_prob=0.5, ) waves = np.zeros((len(eval_specs), length, 3), dtype=DTYPE) labels: list[list[tuple[str, str, int]]] = [] kinds: list[str] = [] with h5py.File(h5_path, "r") as h5: for i, spec in enumerate(eval_specs): if i % 1000 == 0 or i == len(eval_specs) - 1: print(f"materialized eval sample {i:,}/{len(eval_specs):,}", flush=True) mixed = np.zeros((length, 3), dtype=DTYPE) sample_labels: list[tuple[str, str, int]] = [] for crop in spec.crops: wave, crop_labels = load_crop_indexed(h5, test_records[crop.rec_idx], crop, length, access_index) mixed += wave sample_labels.extend(crop_labels) waves[i] = normalize_wave(mixed) labels.append(sample_labels) kinds.append(spec.kind) cache_path.parent.mkdir(parents=True, exist_ok=True) np.savez( cache_path, waves=waves, labels=np.asarray(labels, dtype=object), kinds=np.asarray(kinds, dtype=object), ) print(f"wrote cached eval samples {cache_path}", flush=True) return waves, labels, kinds def load_crop_indexed( h5: h5py.File, record: Record, crop: Any, length: int, access_index: H5AccessIndex, ) -> tuple[np.ndarray, list[tuple[str, str, int]]]: paths = access_index.get(record_key(record)) if paths is None: return load_crop(h5, record, crop, length) try: data = [h5[path][crop.start : crop.start + length] for path in paths] except KeyError: return load_crop(h5, record, crop, length) wave = np.stack(data, axis=1) wave = normalize_wave(wave) * crop.amp labels: list[tuple[str, str, int]] = [] for pick in record.phases: rel = pick.index - crop.start if 0 <= rel < length: labels.append((pick.phase, PHASE_TO_GROUP[pick.phase], rel)) return wave, labels def materialize_specs_from_open_h5( h5: h5py.File, records: list[Record], specs: list, length: int, access_index: H5AccessIndex, ) -> tuple[np.ndarray, np.ndarray]: waves = np.zeros((len(specs), length, 3), dtype=DTYPE) targets = np.zeros((len(specs), 5, length), dtype=DTYPE) for i, spec in enumerate(specs): mixed = np.zeros((length, 3), dtype=DTYPE) labels: list[tuple[str, str, int]] = [] for crop in spec.crops: wave, crop_labels = load_crop_indexed(h5, records[crop.rec_idx], crop, length, access_index) mixed += wave labels.extend(crop_labels) mixed = normalize_wave(mixed) waves[i] = mixed targets[i] = labels_to_target(labels, length) return waves, targets def sample_batch_open_h5( h5: h5py.File, records: list[Record], access_index: H5AccessIndex, valid_indices: list[int], seed: int, batch_size: int, length: int, padlen: int, step: int, ) -> tuple[np.ndarray, np.ndarray]: specs = make_specs( records, n_samples=batch_size, seed=seed + step * 104729, length=length, padlen=padlen, double_prob=0.5, valid_indices=valid_indices, ) return materialize_specs_from_open_h5(h5, records, specs, length, access_index) def train_model( h5_path: Path, records: list[Record], access_index: H5AccessIndex, base_ckpt: Path, out_ckpt: Path, log_csv: Path, seed: int, steps: int, batch_size: int, length: int, padlen: int, lr: float, device: torch.device, resume: bool, init_mode: str, checkpoint_every_steps: int, ) -> BRNN: torch.manual_seed(seed) np.random.seed(seed) random.seed(seed) model = BRNN().to(device) progress_ckpt = out_ckpt.with_name(f"{out_ckpt.stem}.progress.pt") train_state_path = out_ckpt.with_name(f"{out_ckpt.stem}.train_state.json") if resume and out_ckpt.exists(): state = {} if train_state_path.exists(): state = json.loads(train_state_path.read_text(encoding="utf-8")) if state.get("complete", True): model.load_state_dict(torch.load(out_ckpt, map_location="cpu")) model.eval() print(f"loaded existing checkpoint {out_ckpt}", flush=True) return model if init_mode == "finetune": model.load_state_dict(torch.load(base_ckpt, map_location="cpu")) elif init_mode == "scratch": pass else: raise ValueError(f"Unsupported init mode: {init_mode}") model.train() loss_fn = Loss().to(device) opt = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=1e-4) start_step = 0 if resume and progress_ckpt.exists(): progress = torch.load(progress_ckpt, map_location="cpu") if ( progress.get("init_mode") == init_mode and progress.get("target_steps") == steps and progress.get("seed") == seed ): model.load_state_dict(progress["model_state"]) opt.load_state_dict(progress["optimizer_state"]) start_step = int(progress.get("completed_steps", 0)) print(f"resuming partial checkpoint {progress_ckpt} at step {start_step}/{steps}", flush=True) else: print(f"ignoring incompatible partial checkpoint {progress_ckpt}", flush=True) out_ckpt.parent.mkdir(parents=True, exist_ok=True) log_csv.parent.mkdir(parents=True, exist_ok=True) log_mode = "a" if start_step > 0 and log_csv.exists() else "w" valid_indices = [i for i, record in enumerate(records) if record.length >= length] if not valid_indices: raise RuntimeError("No training records are long enough for the requested window length.") with log_csv.open(log_mode, newline="", encoding="utf-8") as f: writer = csv.writer(f) if log_mode == "w": writer.writerow(["step", "loss"]) with h5py.File(h5_path, "r") as train_h5: for step in range(start_step, steps): waves, targets = sample_batch_open_h5( train_h5, records, access_index, valid_indices, seed, batch_size, length, padlen, step, ) xb = torch.from_numpy(waves).to(device).permute(0, 2, 1) yb = torch.from_numpy(targets).to(device) out = model(xb) loss = loss_fn(out, yb) opt.zero_grad(set_to_none=True) loss.backward() opt.step() loss_value = float(loss.detach().cpu()) writer.writerow([step, loss_value]) if step % 50 == 0 or step == steps - 1: print( f"{out_ckpt.stem}: step {step:05d}/{steps} loss={loss_value:.3f}", flush=True, ) f.flush() completed_steps = step + 1 if checkpoint_every_steps > 0 and ( completed_steps % checkpoint_every_steps == 0 or completed_steps == steps ): torch.save( { "model_state": model.state_dict(), "optimizer_state": opt.state_dict(), "completed_steps": completed_steps, "target_steps": steps, "init_mode": init_mode, "seed": seed, }, progress_ckpt, ) train_state_path.write_text( json.dumps( { "complete": False, "completed_steps": completed_steps, "target_steps": steps, "init_mode": init_mode, "seed": seed, "progress_checkpoint": str(progress_ckpt), }, ensure_ascii=False, indent=2, ), encoding="utf-8", ) torch.save(model.state_dict(), out_ckpt) train_state_path.write_text( json.dumps( { "complete": True, "completed_steps": steps, "target_steps": steps, "init_mode": init_mode, "seed": seed, "checkpoint": str(out_ckpt), }, ensure_ascii=False, indent=2, ), encoding="utf-8", ) if progress_ckpt.exists(): progress_ckpt.unlink() model.eval() return model def main() -> None: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--h5", default="data/credit-x1.h5") parser.add_argument("--keys", default="data/creditkeys.npz") parser.add_argument("--base-ckpt", default="ckpt/pnsn.v3.pt") parser.add_argument("--out-dir", default="outputs/snr_transfer_phase_balanced_seed20260609") parser.add_argument( "--cache-dir", default=None, help="Directory for reusable record and phase-SNR caches. Defaults to --out-dir.", ) parser.add_argument("--seed", type=int, default=20260609) 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( "--scratch-train-steps", type=int, default=0, help="Override --train-steps for scratch initialization. Defaults to --train-steps.", ) parser.add_argument("--train-batch", type=int, default=16) parser.add_argument("--eval-samples", type=int, default=10000) parser.add_argument("--eval-batch", type=int, default=64) parser.add_argument("--max-train-events", type=int, default=0) parser.add_argument("--max-test-events", type=int, default=0) parser.add_argument("--lr", type=float, default=1e-4) parser.add_argument("--resume", action="store_true") parser.add_argument( "--checkpoint-every-steps", type=int, default=0, help="Save resumable model and optimizer state during training every N steps.", ) parser.add_argument( "--init-modes", nargs="+", choices=["finetune", "scratch"], default=["finetune"], help="Initialization modes to run on the same matched training pools.", ) parser.add_argument( "--force-retrain-init-modes", nargs="+", choices=["finetune", "scratch"], default=[], help="Ignore existing checkpoints for these initialization modes even when --resume is set.", ) parser.add_argument( "--subset-slugs", nargs="+", choices=["full", "p5_s_bal", "p10_s_bal"], default=["full", "p5_s_bal", "p10_s_bal"], help="Training subsets to run. Candidate matching is still computed across all subsets.", ) parser.add_argument( "--merge-existing-summary", action="store_true", help="Merge newly evaluated rows into an existing summary.json by slug.", ) parser.add_argument( "--train-only", action="store_true", help="Train or resume selected models and skip evaluation/summary plotting.", ) parser.add_argument( "--no-match-train-size", action="store_true", help="Use all phase-filtered candidate records; training steps remain equal.", ) parser.add_argument( "--match-mode", choices=["record-count", "phase-composition"], default="record-count", help=( "record-count matches only the total record count; phase-composition also " "matches the number of P-only, P+S, and S-only waveform records." ), ) parser.add_argument( "--filter-mode", choices=["phase-label", "record-complete", "record-any"], default="phase-label", help=( "phase-label removes only labels below threshold; record-complete keeps " "only records whose original labels all pass and then trains on the full label set; " "record-any keeps the full label set when any original label passes." ), ) parser.add_argument( "--s-threshold-mode", choices=["balanced", "same-as-p"], default="balanced", help="Choose S thresholds by P/S label-count balancing or set S thresholds equal to the P thresholds.", ) args = parser.parse_args() h5_path = Path(args.h5) key_path = Path(args.keys) base_ckpt = Path(args.base_ckpt) out_dir = Path(args.out_dir) out_dir.mkdir(parents=True, exist_ok=True) cache_dir = Path(args.cache_dir) if args.cache_dir else out_dir cache_dir.mkdir(parents=True, exist_ok=True) device = torch.device("mps") if torch.backends.mps.is_available() else torch.device("cpu") print(f"device={device}", flush=True) max_train = None if args.max_train_events == 0 else args.max_train_events max_test = None if args.max_test_events == 0 else args.max_test_events train_tag = "all" if max_train is None else str(max_train) test_tag = "all" if max_test is None else str(max_test) train_records = build_records_sequential( h5_path, key_path, "train", max_train, cache_dir / f"records_train_{train_tag}.json" ) test_records = build_records_sequential( h5_path, key_path, "test", max_test, cache_dir / f"records_test_{test_tag}.json" ) print(f"train records={len(train_records):,} test records={len(test_records):,}", flush=True) train_access_index = build_h5_access_index( h5_path, train_records, cache_dir / f"h5_access_index_train_{train_tag}.json", ) test_access_index = build_h5_access_index( h5_path, test_records, cache_dir / f"h5_access_index_test_{test_tag}.json", ) phase_snr = compute_phase_snr_cache(h5_path, train_records, cache_dir / "train_phase_snr_db.json") by_group = collect_phase_snr(train_records, phase_snr) p5_count = sum(v >= 5.0 for v in by_group["P"]) p10_count = sum(v >= 10.0 for v in by_group["P"]) if args.s_threshold_mode == "balanced": s5, s5_count = choose_s_threshold(by_group["S"], p5_count) s10, s10_count = choose_s_threshold(by_group["S"], p10_count) else: s5 = 5.0 s10 = 10.0 s5_count = sum(v >= s5 for v in by_group["S"]) s10_count = sum(v >= s10 for v in by_group["S"]) print( f"phase-aware thresholds ({args.s_threshold_mode}, {args.filter_mode}): " f"P>=5/S>={s5:.3f} gives P={p5_count:,} S={s5_count:,}; " f"P>=10/S>={s10:.3f} gives P={p10_count:,} S={s10_count:,}", flush=True, ) p5_label = f"P>=5 dB, S>={s5:.2f} dB" p10_label = f"P>=10 dB, S>={s10:.2f} dB" if args.filter_mode == "record-complete": p5_label += " record-complete" p10_label += " record-complete" elif args.filter_mode == "record-any": p5_label += " record-any" p10_label += " record-any" subset_defs = [ ("full", "Full", None, None), ("p5_s_bal", p5_label, 5.0, s5), ("p10_s_bal", p10_label, 10.0, s10), ] if args.filter_mode == "record-complete": filter_fn = filter_records_by_complete_thresholds elif args.filter_mode == "record-any": filter_fn = filter_records_by_any_thresholds else: filter_fn = filter_records_by_phase_thresholds candidate_records = { slug: filter_fn(train_records, phase_snr, p_thr, s_thr) for slug, _label, p_thr, s_thr in subset_defs } candidate_counts = {slug: len(records) for slug, records in candidate_records.items()} candidate_phase_counts = {slug: phase_counts(records) for slug, records in candidate_records.items()} candidate_phase_composition_counts = { slug: phase_composition_counts(records) for slug, records in candidate_records.items() } if args.no_match_train_size: train_pools = candidate_records matched_train_records = None matched_train_composition = None elif args.match_mode == "phase-composition": matched_train_composition = { key: min(counts[key] for counts in candidate_phase_composition_counts.values()) for key in PHASE_COMPOSITION_ORDER } matched_train_records = sum(matched_train_composition.values()) train_pools = { slug: matched_records_by_phase_composition(records, matched_train_composition, args.seed + i * 8191) for i, (slug, _label, _p_thr, _s_thr) in enumerate(subset_defs) for records in [candidate_records[slug]] } else: matched_train_records = min(candidate_counts.values()) matched_train_composition = None train_pools = { slug: matched_records(records, matched_train_records, args.seed + i * 8191) for i, (slug, _label, _p_thr, _s_thr) in enumerate(subset_defs) for records in [candidate_records[slug]] } matched_phase_counts = {slug: phase_counts(records) for slug, records in train_pools.items()} matched_phase_composition_counts = { slug: phase_composition_counts(records) for slug, records in train_pools.items() } balance_rows = [] for slug, label, p_thr, s_thr in subset_defs: candidate_composition = candidate_phase_composition_counts[slug] matched_composition = matched_phase_composition_counts[slug] balance_rows.append( { "slug": slug, "label": label, "p_threshold_db": "" if p_thr is None else p_thr, "s_threshold_db": "" if s_thr is None else s_thr, "candidate_records": candidate_counts[slug], "candidate_P_only_records": candidate_composition["P_only"], "candidate_PS_records": candidate_composition["PS"], "candidate_S_only_records": candidate_composition["S_only"], "candidate_none_records": candidate_composition["none"], "candidate_P_labels": candidate_phase_counts[slug]["P"], "candidate_S_labels": candidate_phase_counts[slug]["S"], "matched_records": len(train_pools[slug]), "matched_P_only_records": matched_composition["P_only"], "matched_PS_records": matched_composition["PS"], "matched_S_only_records": matched_composition["S_only"], "matched_none_records": matched_composition["none"], "matched_P_labels": matched_phase_counts[slug]["P"], "matched_S_labels": matched_phase_counts[slug]["S"], } ) write_phase_balance_table(balance_rows, out_dir / "phase_balance_table.csv") print("candidate counts=" + json.dumps(candidate_counts, ensure_ascii=False), flush=True) print("phase balance=" + json.dumps(balance_rows, ensure_ascii=False), flush=True) eval_cache_path = ( cache_dir / f"eval_seed{args.seed}_test{test_tag}_n{args.eval_samples}_len{args.length}_pad{args.padlen}.npz" ) eval_data: tuple[np.ndarray, list[list[tuple[str, str, int]]], list[str]] | None = None def get_eval_data() -> tuple[np.ndarray, list[list[tuple[str, str, int]]], list[str]]: nonlocal eval_data if eval_data is None: eval_data = load_or_materialize_eval_samples( h5_path=h5_path, test_records=test_records, access_index=test_access_index, cache_path=eval_cache_path, seed=args.seed + 17, n_samples=args.eval_samples, length=args.length, padlen=args.padlen, ) return eval_data thresholds = [round(x, 1) for x in np.arange(0.1, 1.0, 0.1)] train_steps_by_init_mode = { "finetune": args.train_steps, "scratch": args.scratch_train_steps if args.scratch_train_steps > 0 else args.train_steps, } force_retrain_init_modes = set(args.force_retrain_init_modes) rows = [] model_metrics = {} run_subset_slugs = set(args.subset_slugs) run_subset_defs = [(idx, item) for idx, item in enumerate(subset_defs) if item[0] in run_subset_slugs] init_seed_offsets = {"finetune": 0, "scratch": 10000} for init_mode in args.init_modes: mode_steps = train_steps_by_init_mode[init_mode] for idx, (subset_slug, label, p_thr, s_thr) in run_subset_defs: subset_records = train_pools[subset_slug] model_slug = f"{init_mode}_{subset_slug}" model_label = f"{init_mode} {label if args.no_match_train_size else f'{label} matched'}" model = train_model( h5_path=h5_path, records=subset_records, access_index=train_access_index, base_ckpt=base_ckpt, out_ckpt=out_dir / f"pnsn.v3.{init_mode}.{subset_slug}.pt", log_csv=out_dir / f"loss_{init_mode}_{subset_slug}.csv", seed=args.seed + init_seed_offsets[init_mode] + idx * 1000, steps=mode_steps, batch_size=args.train_batch, length=args.length, padlen=args.padlen, lr=args.lr, device=device, resume=args.resume and init_mode not in force_retrain_init_modes, init_mode=init_mode, checkpoint_every_steps=args.checkpoint_every_steps, ) if args.train_only: continue eval_waves, eval_labels, eval_kinds = get_eval_data() outputs = run_model(model, eval_waves, device, args.eval_batch) metrics = evaluate_outputs(outputs, eval_labels, eval_kinds, thresholds, min_sep=50, tolerance=100) model_metrics[model_slug] = strip_errors(metrics) row = { "slug": model_slug, "subset_slug": subset_slug, "init_mode": init_mode, "label": model_label, "p_threshold_db": p_thr, "s_threshold_db": s_thr, "train_records": len(subset_records), "train_steps": mode_steps, "force_retrained": init_mode in force_retrain_init_modes, "candidate_train_records": candidate_counts[subset_slug], "candidate_phase_counts": candidate_phase_counts[subset_slug], "candidate_phase_composition_counts": candidate_phase_composition_counts[subset_slug], "matched_phase_counts": matched_phase_counts[subset_slug], "matched_phase_composition_counts": matched_phase_composition_counts[subset_slug], **metric_row(metrics), } rows.append(row) print(f"metrics {model_label}: {json.dumps(row, ensure_ascii=False)}", flush=True) if args.train_only: print("train-only run complete; skipped evaluation and summary plotting", flush=True) return if args.merge_existing_summary and (out_dir / "summary.json").exists(): existing_summary = json.loads((out_dir / "summary.json").read_text(encoding="utf-8")) merged_rows_by_slug = {row["slug"]: row for row in existing_summary.get("rows", [])} for row in rows: merged_rows_by_slug[row["slug"]] = row rows = list(merged_rows_by_slug.values()) model_metrics = { **existing_summary.get("metrics_by_model", {}), **model_metrics, } plot_summary(rows, out_dir / "snr_transfer_phase_balanced_f1_summary.png") write_metrics_table(rows, out_dir / "metrics_table.tex") summary = { "experiment": "phase_balanced_snr_transfer", "seed": args.seed, "device": str(device), "train_steps": args.train_steps, "scratch_train_steps": args.scratch_train_steps if args.scratch_train_steps > 0 else args.train_steps, "train_steps_by_init_mode": train_steps_by_init_mode, "train_batch": args.train_batch, "init_modes": args.init_modes, "force_retrain_init_modes": sorted(force_retrain_init_modes), "checkpoint_every_steps": args.checkpoint_every_steps, "subset_slugs": args.subset_slugs, "merge_existing_summary": args.merge_existing_summary, "filter_mode": args.filter_mode, "s_threshold_mode": args.s_threshold_mode, "match_mode": args.match_mode, "base_ckpt": str(base_ckpt.resolve()), "cache_dir": str(cache_dir.resolve()), "matched_train_records": matched_train_records, "matched_train_composition": matched_train_composition, "thresholds": { "medium": {"P_db": 5.0, "S_db": s5, "target_P_labels": p5_count, "retained_S_labels": s5_count}, "high": {"P_db": 10.0, "S_db": s10, "target_P_labels": p10_count, "retained_S_labels": s10_count}, }, "snr_summary_db": {"P": summarize_snr(by_group["P"]), "S": summarize_snr(by_group["S"])}, "candidate_train_records": candidate_counts, "candidate_phase_counts": candidate_phase_counts, "candidate_phase_composition_counts": candidate_phase_composition_counts, "matched_phase_counts": matched_phase_counts, "matched_phase_composition_counts": matched_phase_composition_counts, "eval_samples": args.eval_samples, "eval_samples_single": int(sum(k == "single" for k in get_eval_data()[2])), "eval_samples_double": int(sum(k == "double" for k in get_eval_data()[2])), "rows": rows, "metrics_by_model": model_metrics, "phase_balance_table": str((out_dir / "phase_balance_table.csv").resolve()), "figure": str((out_dir / "snr_transfer_phase_balanced_f1_summary.png").resolve()), "table": str((out_dir / "metrics_table.tex").resolve()), } (out_dir / "summary.json").write_text(json.dumps(summary, ensure_ascii=False, indent=2), encoding="utf-8") print(json.dumps(summary, ensure_ascii=False, indent=2)[:5000], flush=True) if __name__ == "__main__": main()