#!/usr/bin/env python3 """Matched-SNR transfer experiment for DispNet v2.3. This script trains DispNet v2.3 from scratch on equal-size NCF-SNR subsets and evaluates every model on the same unfiltered test split. SNR is estimated from the NCF waveform by comparing RMS energy in the surface-wave arrival window with RMS energy outside that window. """ from __future__ import annotations import argparse import csv import importlib.util import json import math import os import random import sys from dataclasses import dataclass from pathlib import Path from typing import Dict, Iterable, List, Sequence os.environ.setdefault("KMP_DUPLICATE_LIB_OK", "TRUE") os.environ.setdefault("OMP_NUM_THREADS", "1") os.environ.setdefault("MKL_NUM_THREADS", "1") import h5py import matplotlib.pyplot as plt import numpy as np import torch from torch.utils.data import DataLoader, Dataset ROOT = Path(__file__).resolve().parents[1] if str(ROOT) not in sys.path: sys.path.insert(0, str(ROOT)) def load_v23_module(): script_path = ROOT / "dispnet.v2.3.py" spec = importlib.util.spec_from_file_location("dispnet_v23_local", script_path) if spec is None or spec.loader is None: raise RuntimeError(f"Cannot import {script_path}") module = importlib.util.module_from_spec(spec) spec.loader.exec_module(module) return module def decode_key(x) -> str: return x.decode("utf-8") if isinstance(x, bytes) else str(x) def haversine_km(coord1: np.ndarray, coord2: np.ndarray) -> float: lon1, lat1 = map(float, coord1[:2]) lon2, lat2 = map(float, coord2[:2]) r = 6371.0 phi1 = math.radians(lat1) phi2 = math.radians(lat2) dphi = math.radians(lat2 - lat1) dlambda = math.radians(lon2 - lon1) a = math.sin(dphi / 2.0) ** 2 + math.cos(phi1) * math.cos(phi2) * math.sin(dlambda / 2.0) ** 2 return 2.0 * r * math.asin(min(1.0, math.sqrt(a))) def rms(x: np.ndarray) -> float: if x.size == 0: return float("nan") x = np.asarray(x, dtype=np.float64) return float(np.sqrt(np.mean(np.square(x)))) def ncf_snr_db(time: np.ndarray, wave: np.ndarray, distance_km: float, periods: np.ndarray, velocity: np.ndarray, mask: np.ndarray) -> float: valid = (mask > 0) & np.isfinite(velocity) & (velocity > 0) & np.isfinite(periods) & (periods > 0) if valid.sum() < 3 or distance_km <= 0: return float("nan") v = velocity[valid].astype(float) p = periods[valid].astype(float) vmin = float(np.nanpercentile(v, 5)) vmax = float(np.nanpercentile(v, 95)) if not (np.isfinite(vmin) and np.isfinite(vmax)) or vmin <= 0 or vmax <= 0: return float("nan") pad = max(10.0, float(np.nanmedian(p))) t0 = max(float(time[0]), distance_km / vmax - pad) t1 = min(float(time[-1]), distance_km / vmin + pad) if t1 <= t0: return float("nan") signal = (time >= t0) & (time <= t1) noise = ~signal if signal.sum() < 10 or noise.sum() < 10: n = len(time) edge = max(10, n // 5) noise = np.zeros(n, dtype=bool) noise[:edge] = True noise[-edge:] = True signal = ~noise srms = rms(wave[signal]) nrms = rms(wave[noise]) if not (np.isfinite(srms) and np.isfinite(nrms)) or nrms <= 0: return float("nan") return 20.0 * math.log10((srms + 1e-12) / (nrms + 1e-12)) def valid_key(f: h5py.File, key: str) -> bool: if key not in f["paths"]: return False grp = f["paths"][key] for name in ("disp_periods", "disp_velocity", "disp_mask", "ncf"): if name not in grp: return False try: mask = np.asarray(grp["disp_mask"][()]).astype(bool) velocity = np.asarray(grp["disp_velocity"][()], dtype=float) if mask.sum() <= 0: return False if not np.all(np.isfinite(velocity[mask])): return False if not np.all(velocity[mask] > 0): return False except Exception: return False for item_name in grp["ncf"].keys(): item = grp["ncf"][item_name] if "time" in item and "waveform" in item and "sta1_coord" in item and "sta2_coord" in item: return True return False def compute_snr_cache(h5_path: Path, out_path: Path) -> Dict[str, Dict[str, float | str]]: if out_path.exists(): with out_path.open() as f: return json.load(f) out_path.parent.mkdir(parents=True, exist_ok=True) rows: Dict[str, Dict[str, float | str]] = {} with h5py.File(h5_path, "r") as f: for split_name in ("train_keys", "test_keys"): keys = [decode_key(x) for x in f[split_name][()]] for i, key in enumerate(keys): if not valid_key(f, key): continue grp = f["paths"][key] periods = np.asarray(grp["disp_periods"][()], dtype=np.float32) velocity = np.asarray(grp["disp_velocity"][()], dtype=np.float32) mask = np.asarray(grp["disp_mask"][()], dtype=np.float32) values = [] distances = [] for item_name in grp["ncf"].keys(): item = grp["ncf"][item_name] if "time" not in item or "waveform" not in item: continue time = np.asarray(item["time"][()], dtype=np.float32) wave = np.asarray(item["waveform"][()], dtype=np.float32) distance = haversine_km(np.asarray(item["sta1_coord"][()]), np.asarray(item["sta2_coord"][()])) snr = ncf_snr_db(time, wave, distance, periods, velocity, mask) if np.isfinite(snr): values.append(float(snr)) distances.append(float(distance)) if values: rows[key] = { "split": "train" if split_name == "train_keys" else "test", "snr_db": float(np.median(values)), "distance_km": float(np.median(distances)) if distances else float("nan"), "valid_period_count": int(mask.sum()), } if i % 5000 == 0: print(f"{split_name}: scanned {i}/{len(keys)} keys; cached={len(rows)}", flush=True) tmp = out_path.with_suffix(".tmp") with tmp.open("w") as f: json.dump(rows, f, allow_nan=False) tmp.replace(out_path) return rows def matched_threshold_subsets(rows: Dict[str, Dict], seed: int) -> tuple[Dict[str, List[str]], Dict[str, float]]: train = [(key, float(row["snr_db"])) for key, row in rows.items() if row["split"] == "train" and np.isfinite(row["snr_db"])] if len(train) < 3: raise RuntimeError("Not enough finite-SNR train records.") values = np.array([snr for _, snr in train], dtype=float) q1, q2 = np.quantile(values, [1 / 3, 2 / 3]) all_keys = sorted([key for key, _ in train]) snr_gt_q1 = sorted([key for key, snr in train if snr > q1]) snr_gt_q2 = sorted([key for key, snr in train if snr > q2]) n = len(snr_gt_q2) if n == 0: raise RuntimeError("Highest SNR threshold produced no records.") rng = np.random.default_rng(seed) full_idx = np.sort(rng.choice(len(all_keys), size=n, replace=False)) q1_idx = np.sort(rng.choice(len(snr_gt_q1), size=n, replace=False)) subsets = { "full": [all_keys[int(i)] for i in full_idx], "snr_q1": [snr_gt_q1[int(i)] for i in q1_idx], "snr_q2": snr_gt_q2, } thresholds = {"q1": float(q1), "q2": float(q2)} return subsets, thresholds class NCFKeyDataset(Dataset): def __init__(self, h5_path: Path, keys: Sequence[str], waveform_length: int = 1536, random_ncf: bool = True, seed: int = 42): self.h5_path = str(h5_path) self.keys = list(keys) self.waveform_length = waveform_length self.random_ncf = random_ncf self.rng = random.Random(seed) self._h5_file = None def __len__(self): return len(self.keys) def _h5(self): if self._h5_file is None: self._h5_file = h5py.File(self.h5_path, "r") return self._h5_file @staticmethod def _pad_or_truncate(arr: np.ndarray, n: int, pad_value: float = 0.0): arr = np.asarray(arr, dtype=np.float32).reshape(-1) if len(arr) == n: return arr if len(arr) > n: return arr[:n] out = np.full(n, pad_value, dtype=np.float32) out[: len(arr)] = arr return out def __getitem__(self, index: int): f = self._h5() key = self.keys[index] grp = f["paths"][key] ncf = grp["ncf"] names = [name for name in sorted(ncf.keys()) if "time" in ncf[name] and "waveform" in ncf[name]] item = ncf[self.rng.choice(names) if self.random_ncf and len(names) > 1 else names[0]] waveform = self._pad_or_truncate(np.asarray(item["waveform"][()], dtype=np.float32), self.waveform_length) return { "key": key, "waveform": torch.from_numpy(waveform), "disp": torch.from_numpy(np.asarray(grp["disp_velocity"][()], dtype=np.float32)), "mask": torch.from_numpy(np.asarray(grp["disp_mask"][()], dtype=np.float32)), "periods": torch.from_numpy(np.asarray(grp["disp_periods"][()], dtype=np.float32)), } def collate(batch): return { "key": [x["key"] for x in batch], "waveform": torch.stack([x["waveform"] for x in batch]), "disp": torch.stack([x["disp"] for x in batch]), "mask": torch.stack([x["mask"] for x in batch]), "periods": torch.stack([x["periods"] for x in batch]), } def make_loader(h5_path: Path, keys: Sequence[str], batch_size: int, num_workers: int, waveform_length: int, random_ncf: bool, seed: int, shuffle: bool): ds = NCFKeyDataset(h5_path, keys, waveform_length=waveform_length, random_ncf=random_ncf, seed=seed) return DataLoader(ds, batch_size=batch_size, shuffle=shuffle, num_workers=num_workers, collate_fn=collate, drop_last=False) def load_model(v23, ckpt: Dict, device: torch.device): cfg = v23.TrainConfig() for key, value in ckpt.get("config", {}).items(): if hasattr(cfg, key): setattr(cfg, key, value) reference = ckpt.get("reference_disp", None) if reference is None and "model" in ckpt: reference = ckpt["model"].get("reference_disp") model = v23.DispNetCNNV23( input_length=cfg.waveform_length, base_channels=cfg.base_channels, output_dim=cfg.output_dim, dropout=cfg.dropout, reference_disp=reference, ).to(device) model.load_state_dict(ckpt["model"]) return model, cfg def new_config(v23, args) -> object: cfg = v23.TrainConfig() cfg.h5_path = args.h5 cfg.epochs = args.epochs cfg.lr = args.lr cfg.batch_size = args.batch_size cfg.num_workers = args.num_workers cfg.use_amp = False cfg.device = args.device return cfg def scratch_model(v23, cfg, reference_disp: torch.Tensor, device: torch.device): return v23.DispNetCNNV23( input_length=cfg.waveform_length, base_channels=cfg.base_channels, output_dim=cfg.output_dim, dropout=cfg.dropout, reference_disp=reference_disp, ).to(device) def run_epoch(v23, model, loader, device, optimizer=None, cfg=None): return v23.run_one_epoch( model=model, loader=loader, device=device, optimizer=optimizer, scaler=None, use_amp=False, grad_clip=(cfg.grad_clip if optimizer is not None and cfg is not None else None), lambda_certainty=cfg.lambda_certainty if cfg is not None else 0.2, lambda_slope=cfg.lambda_slope if cfg is not None else 0.25, lambda_curvature=cfg.lambda_curvature if cfg is not None else 0.05, lambda_pairwise=cfg.lambda_pairwise if cfg is not None else 0.35, lambda_std=cfg.lambda_std if cfg is not None else 0.75, huber_delta=cfg.huber_delta if cfg is not None else 0.05, certainty_pos_weight=None, ) def summarize_bin(rows: Dict[str, Dict], keys: Sequence[str]) -> Dict[str, float]: snr = np.array([rows[k]["snr_db"] for k in keys], dtype=float) return { "count": int(len(keys)), "snr_min": float(np.min(snr)), "snr_median": float(np.median(snr)), "snr_max": float(np.max(snr)), } def plot_dispersion(rows: Sequence[Dict], out: Path) -> None: labels = [row["label"] for row in rows] x = np.arange(len(rows)) mae = [row["val_mae"] for row in rows] rmse = [row["val_rmse"] for row in rows] fig, ax = plt.subplots(figsize=(7.2, 4.2), dpi=220) width = 0.34 ax.bar(x - width / 2, mae, width, label="MAE", color="#0072B2") ax.bar(x + width / 2, rmse, width, label="RMSE", color="#D55E00") ax.set_xticks(x) ax.set_xticklabels(labels) ax.set_ylabel("Phase velocity error (km/s)") ax.grid(axis="y", alpha=0.25) ax.legend(frameon=False) fig.tight_layout() fig.savefig(out, bbox_inches="tight") plt.close(fig) def write_table(rows: Sequence[Dict], out: Path) -> None: lines = [ "\\begin{tabular}{@{}lrrrrr@{}}", "\\toprule", "Subset & Records & Median SNR & MAE & RMSE & Cert. F1 \\\\", "\\midrule", ] for row in rows: lines.append( f"{row['label']} & {row['train_records']} & {row['snr_median']:.2f} & " f"{row['val_mae']:.4f} & {row['val_rmse']:.4f} & {row['val_certainty_f1']:.3f} \\\\" ) lines.extend(["\\bottomrule", "\\end{tabular}", ""]) out.write_text("\n".join(lines)) def main() -> None: parser = argparse.ArgumentParser() parser.add_argument("--h5", default="data/ncf_data/ncf_disp_dataset_with_disp_image.h5") parser.add_argument("--base-ckpt", default="/Users/yuziye/machinelearning/disp/ckpt_large/checkpoints_dispnet_v2.3_residual_cnn/best.pt") parser.add_argument("--out-dir", default="outputs/disp_snr_transfer_seed20260609") parser.add_argument("--seed", type=int, default=20260609) parser.add_argument("--epochs", type=int, default=8) parser.add_argument("--batch-size", type=int, default=64) parser.add_argument("--num-workers", type=int, default=0) parser.add_argument("--lr", type=float, default=2e-4) parser.add_argument("--max-train-per-bin", type=int, default=0) parser.add_argument("--device", default="cpu") parser.add_argument("--mode", choices=["scratch", "finetune"], default="scratch") parser.add_argument("--eval-every", type=int, default=0, help="Evaluate during training every N epochs; 0 means final only.") args = parser.parse_args() random.seed(args.seed) np.random.seed(args.seed) torch.manual_seed(args.seed) out_dir = Path(args.out_dir) out_dir.mkdir(parents=True, exist_ok=True) h5_path = Path(args.h5) rows = compute_snr_cache(h5_path, out_dir / "ncf_snr_cache.json") bins, snr_thresholds = matched_threshold_subsets(rows, args.seed) if args.max_train_per_bin > 0: rng = np.random.default_rng(args.seed) for name, keys in bins.items(): idx = np.sort(rng.choice(len(keys), size=min(args.max_train_per_bin, len(keys)), replace=False)) bins[name] = [keys[int(i)] for i in idx] test_keys = [key for key, row in rows.items() if row["split"] == "test" and np.isfinite(row["snr_db"])] test_keys = sorted(test_keys) v23 = load_v23_module() device = torch.device(args.device) ckpt = torch.load(args.base_ckpt, map_location="cpu") if args.mode == "finetune" else None if args.mode == "finetune": _, base_cfg = load_model(v23, ckpt, device) base_cfg.lr = args.lr base_cfg.epochs = args.epochs base_cfg.batch_size = args.batch_size base_cfg.num_workers = args.num_workers base_cfg.use_amp = False else: base_cfg = new_config(v23, args) test_loader = make_loader(h5_path, test_keys, args.batch_size, args.num_workers, base_cfg.waveform_length, False, args.seed, False) subset_defs = [ ("full", "Full matched"), ("snr_q1", f"SNR>{snr_thresholds['q1']:.2f} dB matched"), ("snr_q2", f"SNR>{snr_thresholds['q2']:.2f} dB matched"), ] result_rows = [] bin_summaries = {name: summarize_bin(rows, keys) for name, keys in bins.items()} print("bin summaries=" + json.dumps(bin_summaries, ensure_ascii=False), flush=True) print(f"test records={len(test_keys)} device={device} mode={args.mode}", flush=True) for i, (name, label) in enumerate(subset_defs): keys = bins[name] cfg = new_config(v23, args) train_loader = make_loader(h5_path, keys, args.batch_size, args.num_workers, cfg.waveform_length, True, args.seed + i * 1000, True) if args.mode == "finetune": model, cfg = load_model(v23, ckpt, device) else: print(f"{name}: estimating reference dispersion from scratch-training subset...", flush=True) reference_disp, reference_periods, reference_counts = v23.estimate_reference_disp( train_loader, output_dim=cfg.output_dim, device=device, ) model = scratch_model(v23, cfg, reference_disp, device) cfg.lr = args.lr cfg.batch_size = args.batch_size cfg.num_workers = args.num_workers cfg.use_amp = False optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=cfg.weight_decay) scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=max(args.epochs, 1), eta_min=cfg.min_lr) log_path = out_dir / f"train_log_{name}.csv" with log_path.open("w", newline="") as f: writer = csv.writer(f) writer.writerow(["epoch", "train_loss", "train_mae", "val_loss", "val_mae", "val_rmse", "val_certainty_f1"]) best = None best_state = None for epoch in range(1, args.epochs + 1): train_stats = run_epoch(v23, model, train_loader, device, optimizer=optimizer, cfg=cfg) scheduler.step() should_eval = args.eval_every > 0 and (epoch % args.eval_every == 0 or epoch == args.epochs) val_stats = run_epoch(v23, model, test_loader, device, optimizer=None, cfg=cfg) if should_eval else None writer.writerow([ epoch, train_stats["loss"], train_stats["mae"], val_stats["loss"] if val_stats else "", val_stats["mae"] if val_stats else "", val_stats["rmse"] if val_stats else "", val_stats["certainty_f1"] if val_stats else "", ]) msg = f"{name} epoch {epoch:03d}/{args.epochs}: train_mae={train_stats['mae']:.4f}" if val_stats: msg += f" val_mae={val_stats['mae']:.4f} val_rmse={val_stats['rmse']:.4f}" print(msg, flush=True) if val_stats and (best is None or val_stats["mae"] < best): best = val_stats["mae"] best_state = {k: v.detach().cpu() for k, v in model.state_dict().items()} if best_state is not None: model.load_state_dict(best_state) final_stats = run_epoch(v23, model, test_loader, device, optimizer=None, cfg=cfg) for stat_name in ("loss", "mae", "rmse", "certainty_f1"): if not np.isfinite(float(final_stats[stat_name])): raise RuntimeError(f"Non-finite final {stat_name} for subset {name}: {final_stats[stat_name]}") torch.save({"model": model.state_dict(), "config": cfg.__dict__, "stats": final_stats}, out_dir / f"dispnet.v2.3.transfer.{name}.pt") row = { "slug": name, "label": label, "train_records": len(keys), "snr_min": bin_summaries[name]["snr_min"], "snr_median": bin_summaries[name]["snr_median"], "snr_max": bin_summaries[name]["snr_max"], "val_loss": float(final_stats["loss"]), "val_mae": float(final_stats["mae"]), "val_rmse": float(final_stats["rmse"]), "val_certainty_f1": float(final_stats["certainty_f1"]), } result_rows.append(row) print("metrics " + json.dumps(row, ensure_ascii=False), flush=True) plot_dispersion(result_rows, out_dir / "disp_snr_transfer_error_summary.png") write_table(result_rows, out_dir / "disp_metrics_table.tex") summary = { "seed": args.seed, "epochs": args.epochs, "batch_size": args.batch_size, "lr": args.lr, "mode": args.mode, "h5": str(h5_path), "base_ckpt": args.base_ckpt, "snr_definition": "20*log10(signal-window RMS / off-window RMS); signal window predicted from interstation distance and valid dispersion velocity range.", "snr_thresholds": snr_thresholds, "bin_summaries": bin_summaries, "test_records": len(test_keys), "rows": result_rows, "figure": str((out_dir / "disp_snr_transfer_error_summary.png").resolve()), "table": str((out_dir / "disp_metrics_table.tex").resolve()), } with (out_dir / "summary.json").open("w") as f: json.dump(summary, f, ensure_ascii=False, indent=2, allow_nan=False) print(json.dumps(summary, ensure_ascii=False, indent=2), flush=True) if __name__ == "__main__": main()