| |
| """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() |
|
|
