| |
| """Summarize manual P/S pick confidence and pick-level SNR.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import csv |
| import json |
| import math |
| import os |
| import sys |
| from collections import Counter, defaultdict |
| 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 h5py |
| import matplotlib.pyplot as plt |
| 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 |
|
|
|
|
| DEFAULT_H5 = ROOT / "data" / "credit-x1.h5" |
| DEFAULT_RECORDS = [ |
| ROOT / "outputs" / "snr_transfer_seed20260609" / "records_train_all.json", |
| ROOT / "outputs" / "snr_transfer_seed20260609" / "records_test_all.json", |
| ] |
| DEFAULT_CKPT = ROOT / "outputs" / "snr_transfer_seed20260609" / "pnsn.v3.transfer.full.pt" |
| DEFAULT_OUT_DIR = ROOT / "outputs" / "manual_phase_confidence_snr" |
| PHASE_TO_GROUP = {"Pg": "P", "Pn": "P", "Sg": "S", "Sn": "S"} |
| GROUP_TO_CHANNELS = {"P": [1, 3], "S": [2, 4]} |
| COMPONENT_ORDER = ("BHE", "BHN", "BHZ") |
| DTYPE = np.float32 |
|
|
|
|
| def component_keys(station) -> tuple[str, str, str] | None: |
| keys = set(station.keys()) |
| if all(k in keys for k in COMPONENT_ORDER): |
| return COMPONENT_ORDER |
| by_suffix: dict[str, str] = {} |
| for key in keys: |
| if key.endswith("HE"): |
| by_suffix["BHE"] = key |
| elif key.endswith("HN"): |
| by_suffix["BHN"] = key |
| elif key.endswith("HZ"): |
| by_suffix["BHZ"] = key |
| if all(k in by_suffix for k in COMPONENT_ORDER): |
| return tuple(by_suffix[k] for k in COMPONENT_ORDER) |
| return None |
|
|
|
|
| def window_std(wave: np.ndarray, start: int, end: int) -> float | None: |
| if start < 0 or end > len(wave) or end <= start: |
| return None |
| return float(np.std(wave[start:end])) |
|
|
|
|
| def pick_snr_db(waves: tuple[np.ndarray, np.ndarray, np.ndarray], phase: str, index: int) -> float | None: |
| """Use the same local SNR convention as scripts/snr_transfer_experiment.py.""" |
| east, north, vertical = waves |
| if PHASE_TO_GROUP[phase] == "P": |
| pre = window_std(vertical, index - 50, index) |
| aft = window_std(vertical, index, index + 50) |
| if pre is None or aft is None: |
| return None |
| return 10.0 * math.log10((aft + 1e-6) / (pre + 1e-6)) |
|
|
| pre_e = window_std(east, index - 150, index) |
| aft_e = window_std(east, index, index + 150) |
| pre_n = window_std(north, index - 150, index) |
| aft_n = window_std(north, index, index + 150) |
| if pre_e is None or aft_e is None or pre_n is None or aft_n is None: |
| return None |
| snr_e = 10.0 * math.log10((aft_e + 1e-6) / (pre_e + 1e-6)) |
| snr_n = 10.0 * math.log10((aft_n + 1e-6) / (pre_n + 1e-6)) |
| return 0.5 * (snr_e + snr_n) |
|
|
|
|
| def normalize_wave(wave: np.ndarray) -> np.ndarray: |
| wave = wave.astype(DTYPE, copy=False) |
| wave = wave - wave.mean(axis=0, keepdims=True) |
| denom = np.max(np.abs(wave), axis=0, keepdims=True) + 1e-6 |
| return wave / denom |
|
|
|
|
| def pick_window(station, comps: tuple[str, str, str], index: int, length: int) -> tuple[np.ndarray, int] | None: |
| station_len = min(int(station[c].shape[0]) for c in comps) |
| if station_len <= 0 or index < 0 or index >= station_len: |
| return None |
| start = int(np.clip(index - length // 2, 0, max(0, station_len - length))) |
| stop = min(start + length, station_len) |
| data = [station[c][start:stop] for c in comps] |
| wave = np.stack(data, axis=1) |
| if len(wave) < length: |
| padded = np.zeros((length, 3), dtype=DTYPE) |
| padded[: len(wave)] = wave |
| wave = padded |
| return normalize_wave(wave), index - start |
|
|
|
|
| def load_records(paths: list[Path], sample_size: int | None, seed: int) -> list[dict]: |
| rows: list[dict] = [] |
| seen: set[tuple[str, str, str, int]] = set() |
| for path in paths: |
| payload = json.loads(path.read_text()) |
| split = payload.get("split", path.stem) |
| for record in payload["records"]: |
| for pick in record["phases"]: |
| if str(pick.get("source", "")).startswith("MANUAL"): |
| key = (record["event"], record["station"], pick["phase"], int(pick["index"])) |
| if key in seen: |
| continue |
| seen.add(key) |
| rows.append( |
| { |
| "split": split, |
| "event": record["event"], |
| "station": record["station"], |
| "phase": pick["phase"], |
| "phase_group": PHASE_TO_GROUP[pick["phase"]], |
| "index": int(pick["index"]), |
| "distance_km": record.get("distance_km"), |
| } |
| ) |
| if sample_size is not None and sample_size < len(rows): |
| rng = np.random.default_rng(seed) |
| idx = np.sort(rng.choice(len(rows), size=sample_size, replace=False)) |
| rows = [rows[int(i)] for i in idx] |
| return rows |
|
|
|
|
| def scalar_attr(value): |
| if hasattr(value, "item"): |
| value = value.item() |
| if isinstance(value, bytes): |
| return value.decode("utf-8") |
| return value |
|
|
|
|
| def model_confidence( |
| ckpt_path: Path, |
| windows: np.ndarray, |
| rel_indices: list[int], |
| groups: list[str], |
| batch_size: int, |
| device: torch.device, |
| ) -> list[float]: |
| model = BRNN().to(device) |
| model.load_state_dict(torch.load(ckpt_path, map_location="cpu")) |
| model.eval() |
| values: list[float] = [] |
| with torch.no_grad(): |
| for start in range(0, len(windows), batch_size): |
| batch = torch.from_numpy(windows[start : start + batch_size]).to(device).permute(0, 2, 1) |
| out = model(batch).detach().cpu().numpy() |
| for j in range(out.shape[0]): |
| row_i = start + j |
| rel = rel_indices[row_i] |
| chans = GROUP_TO_CHANNELS[groups[row_i]] |
| values.append(float(out[j, chans, rel].max())) |
| return values |
|
|
|
|
| def collect_pick_rows( |
| h5_path: Path, |
| rows: list[dict], |
| ckpt_path: Path, |
| length: int, |
| batch_size: int, |
| device: torch.device, |
| ) -> list[dict]: |
| out: list[dict] = [] |
| windows: list[np.ndarray] = [] |
| rel_indices: list[int] = [] |
| groups: list[str] = [] |
| missing = Counter() |
| with h5py.File(h5_path, "r") as h5: |
| for i, row in enumerate(rows, start=1): |
| try: |
| station = h5[row["event"]][row["station"]] |
| except KeyError: |
| missing["station"] += 1 |
| continue |
| comps = component_keys(station) |
| if comps is None: |
| missing["components"] += 1 |
| continue |
| wt_key = f"MANUAL.TRAVTIME.{row['phase']}.WT" |
| if wt_key not in station.attrs: |
| missing["manual_wt"] += 1 |
| continue |
| try: |
| manual_wt = float(scalar_attr(station.attrs[wt_key])) |
| except (TypeError, ValueError): |
| missing["manual_wt_numeric"] += 1 |
| continue |
| waves = tuple(station[c][:] for c in comps) |
| snr = pick_snr_db(waves, row["phase"], row["index"]) |
| if snr is None or not math.isfinite(snr): |
| missing["snr"] += 1 |
| continue |
| window = pick_window(station, comps, row["index"], length) |
| if window is None: |
| missing["model_window"] += 1 |
| continue |
| wave_window, rel_index = window |
| out.append( |
| { |
| **row, |
| "manual_wt": manual_wt, |
| "snr_db": float(snr), |
| } |
| ) |
| windows.append(wave_window) |
| rel_indices.append(rel_index) |
| groups.append(row["phase_group"]) |
| if i % 25000 == 0: |
| print(f"processed {i}/{len(rows)} manual picks; usable={len(out)}", flush=True) |
| if missing: |
| print("skipped:", dict(missing), flush=True) |
| if out: |
| conf = model_confidence( |
| ckpt_path=ckpt_path, |
| windows=np.stack(windows, axis=0), |
| rel_indices=rel_indices, |
| groups=groups, |
| batch_size=batch_size, |
| device=device, |
| ) |
| for row, value in zip(out, conf): |
| row["model_confidence"] = value |
| return out |
|
|
|
|
| def describe(values: list[float]) -> dict: |
| arr = np.asarray(values, dtype=float) |
| if arr.size == 0: |
| return {"n": 0} |
| q = np.percentile(arr, [0, 5, 25, 50, 75, 95, 100]) |
| return { |
| "n": int(arr.size), |
| "mean": float(np.mean(arr)), |
| "std": float(np.std(arr, ddof=1)) if arr.size > 1 else 0.0, |
| "min": float(q[0]), |
| "p05": float(q[1]), |
| "p25": float(q[2]), |
| "median": float(q[3]), |
| "p75": float(q[4]), |
| "p95": float(q[5]), |
| "max": float(q[6]), |
| } |
|
|
|
|
| def make_summary(rows: list[dict]) -> dict: |
| summary = { |
| "notes": [ |
| "manual_wt is MANUAL.TRAVTIME.<phase>.WT from data/README.md.", |
| "model_confidence is the neural phase picker's P- or S-group probability at the manual pick sample.", |
| "snr_db uses the same pick-level local window definition as scripts/snr_transfer_experiment.py.", |
| ], |
| "overall": {}, |
| "by_phase_group": {}, |
| "by_phase": {}, |
| "manual_wt_counts_by_phase_group": {}, |
| } |
| for label, subset in [("all", rows)]: |
| summary["overall"][label] = { |
| "model_confidence": describe([r["model_confidence"] for r in subset]), |
| "manual_wt": describe([r["manual_wt"] for r in subset]), |
| "snr_db": describe([r["snr_db"] for r in subset]), |
| } |
| for key in ("phase_group", "phase"): |
| target = summary[f"by_{key}"] |
| grouped: dict[str, list[dict]] = defaultdict(list) |
| for row in rows: |
| grouped[row[key]].append(row) |
| for name, subset in sorted(grouped.items()): |
| target[name] = { |
| "model_confidence": describe([r["model_confidence"] for r in subset]), |
| "manual_wt": describe([r["manual_wt"] for r in subset]), |
| "snr_db": describe([r["snr_db"] for r in subset]), |
| } |
| counts: dict[str, Counter] = defaultdict(Counter) |
| for row in rows: |
| counts[row["phase_group"]][str(row["manual_wt"])] += 1 |
| summary["manual_wt_counts_by_phase_group"] = { |
| group: dict(sorted(counter.items(), key=lambda item: float(item[0]))) |
| for group, counter in sorted(counts.items()) |
| } |
| return summary |
|
|
|
|
| def write_csv(path: Path, rows: list[dict]) -> None: |
| fieldnames = [ |
| "split", |
| "event", |
| "station", |
| "phase", |
| "phase_group", |
| "index", |
| "distance_km", |
| "manual_wt", |
| "model_confidence", |
| "snr_db", |
| ] |
| with path.open("w", newline="") as f: |
| writer = csv.DictWriter(f, fieldnames=fieldnames) |
| writer.writeheader() |
| writer.writerows(rows) |
|
|
|
|
| def scatter_panel(ax, rows: list[dict], group: str, color: str) -> None: |
| subset = [r for r in rows if r["phase_group"] == group] |
| x = np.asarray([r["model_confidence"] for r in subset], dtype=float) |
| y = np.asarray([r["snr_db"] for r in subset], dtype=float) |
| ax.scatter( |
| x, |
| y, |
| s=9, |
| c=color, |
| alpha=0.28, |
| linewidths=0, |
| label=f"{group} picks (n={len(subset):,})", |
| rasterized=True, |
| ) |
| ax.axhline(0.0, color="#4d4d4d", lw=0.8, ls="--", alpha=0.75) |
| ax.set_xlabel("Neural-network pick confidence") |
| ax.set_ylabel("Pick-level SNR (dB)") |
| ax.set_title(f"{group} picks") |
| ax.grid(True, color="#d9d9d9", lw=0.6, alpha=0.7) |
| ax.legend(frameon=False, loc="best") |
|
|
|
|
| def plot_scatter(path: Path, rows: list[dict]) -> None: |
| colors = {"P": "#2868a8", "S": "#c44e52"} |
| fig, axes = plt.subplots(2, 1, figsize=(7.2, 8.0), dpi=220, sharex=True) |
| for ax, group in zip(axes, ("P", "S")): |
| scatter_panel(ax, rows, group, colors[group]) |
| axes[-1].set_xlim(-0.02, 1.02) |
| fig.suptitle("Manual Labels Scored by Neural Phase Picker: Confidence vs SNR", y=0.995) |
| fig.tight_layout() |
| fig.savefig(path) |
| plt.close(fig) |
|
|
|
|
| def plot_single_group(path: Path, rows: list[dict], group: str) -> None: |
| color = {"P": "#2868a8", "S": "#c44e52"}[group] |
| fig, ax = plt.subplots(figsize=(7.2, 5.0), dpi=220) |
| scatter_panel(ax, rows, group, color) |
| ax.set_xlim(-0.02, 1.02) |
| fig.tight_layout() |
| fig.savefig(path) |
| plt.close(fig) |
|
|
|
|
| def main() -> None: |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--h5", type=Path, default=DEFAULT_H5) |
| parser.add_argument("--records", type=Path, nargs="+", default=DEFAULT_RECORDS) |
| parser.add_argument("--ckpt", type=Path, default=DEFAULT_CKPT) |
| parser.add_argument("--out-dir", type=Path, default=DEFAULT_OUT_DIR) |
| parser.add_argument("--sample-size", type=int, default=2000) |
| parser.add_argument("--seed", type=int, default=20260626) |
| parser.add_argument("--length", type=int, default=5120) |
| parser.add_argument("--batch-size", type=int, default=64) |
| parser.add_argument("--device", default="cpu") |
| args = parser.parse_args() |
|
|
| args.out_dir.mkdir(parents=True, exist_ok=True) |
| device = torch.device(args.device) |
| candidates = load_records(args.records, args.sample_size, args.seed) |
| sample_note = "all" if args.sample_size is None else f"sampled {len(candidates)}" |
| print(f"loaded {sample_note} unique manual P/S picks from record caches", flush=True) |
| rows = collect_pick_rows(args.h5, candidates, args.ckpt, args.length, args.batch_size, device) |
| csv_path = args.out_dir / "manual_phase_confidence_snr.csv" |
| summary_path = args.out_dir / "manual_phase_confidence_snr_summary.json" |
| png_path = args.out_dir / "manual_phase_confidence_snr_scatter.png" |
| p_png_path = args.out_dir / "manual_phase_confidence_snr_scatter_P.png" |
| s_png_path = args.out_dir / "manual_phase_confidence_snr_scatter_S.png" |
| write_csv(csv_path, rows) |
| summary = make_summary(rows) |
| summary["inputs"] = { |
| "h5": str(args.h5), |
| "records": [str(p) for p in args.records], |
| "ckpt": str(args.ckpt), |
| "sample_size": args.sample_size, |
| "seed": args.seed, |
| "length": args.length, |
| "n_candidate_manual_picks_processed": len(candidates), |
| "n_usable_manual_picks": len(rows), |
| } |
| summary_path.write_text(json.dumps(summary, indent=2), encoding="utf-8") |
| plot_scatter(png_path, rows) |
| plot_single_group(p_png_path, rows, "P") |
| plot_single_group(s_png_path, rows, "S") |
| print(f"wrote {csv_path}") |
| print(f"wrote {summary_path}") |
| print(f"wrote {png_path}") |
| print(f"wrote {p_png_path}") |
| print(f"wrote {s_png_path}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
