| |
| """Reproduce paper statistics with a fixed seed. |
| |
| The script uses CREDIT-X1Local records to build deterministic single-event and |
| two-event mixture windows, fine-tunes the pretrained Pn/Sn picker, and exports |
| metrics plus publication-ready composite figures. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import csv |
| import datetime as dt |
| 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, Tuple |
|
|
| 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, Loss |
|
|
|
|
| PHASE_TO_CHANNEL = {"Pg": 1, "Sg": 2, "Pn": 3, "Sn": 4} |
| PHASE_TO_GROUP = {"Pg": "P", "Pn": "P", "Sg": "S", "Sn": "S"} |
| GROUP_TO_CHANNELS = {"P": [1, 3], "S": [2, 4]} |
| COMPONENTS = ("BHE", "BHN", "BHZ") |
| DTYPE = np.float32 |
|
|
|
|
| @dataclass(frozen=True) |
| class PhasePick: |
| phase: str |
| index: int |
| source: str |
|
|
|
|
| @dataclass(frozen=True) |
| class Record: |
| event: str |
| station: str |
| length: int |
| delta: float |
| distance_km: float |
| phases: Tuple[PhasePick, ...] |
|
|
|
|
| @dataclass(frozen=True) |
| class CropSpec: |
| rec_idx: int |
| start: int |
| amp: float |
|
|
|
|
| @dataclass(frozen=True) |
| class SampleSpec: |
| crops: Tuple[CropSpec, ...] |
| kind: str |
|
|
|
|
| def parse_time(value: str) -> dt.datetime: |
| return dt.datetime.strptime(value, "%Y/%m/%d %H:%M:%S.%f") |
|
|
|
|
| def choose_phase(station, phase: str, prefer_manual: bool = True) -> Tuple[str, str] | None: |
| manual_key = f"MANUAL.TRAVTIME.{phase}" |
| rnn_key = f"RNN.TRAVTIME.{phase}" |
| if prefer_manual and manual_key in station.attrs: |
| value = station.attrs[manual_key] |
| if isinstance(value, str) and "/" in value: |
| return value, "MANUAL" |
| if rnn_key in station.attrs: |
| value = station.attrs[rnn_key] |
| tag = station.attrs.get(f"{rnn_key}.tag", "") |
| if isinstance(value, str) and "/" in value and tag == "Y": |
| return value, "RNN.tagY" |
| if not prefer_manual and manual_key in station.attrs: |
| value = station.attrs[manual_key] |
| if isinstance(value, str) and "/" in value: |
| return value, "MANUAL" |
| return None |
|
|
|
|
| def component_keys(station) -> Tuple[str, str, str] | None: |
| keys = set(station.keys()) |
| if all(k in keys for k in COMPONENTS): |
| return COMPONENTS |
| by_suffix = {} |
| 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 COMPONENTS): |
| return tuple(by_suffix[k] for k in COMPONENTS) |
| return None |
|
|
|
|
| def build_records( |
| h5_path: Path, |
| key_path: Path, |
| split: str, |
| max_events: int | None, |
| prefer_manual: bool = True, |
| ) -> List[Record]: |
| keys = np.load(key_path)[split] |
| if max_events is not None: |
| keys = keys[:max_events] |
| records: List[Record] = [] |
| with h5py.File(h5_path, "r") as h5: |
| for event_key in keys: |
| event_key = str(event_key) |
| event = h5[event_key] |
| for station_key in event.keys(): |
| station = event[station_key] |
| comps = component_keys(station) |
| if comps is None: |
| continue |
| first = station[comps[0]] |
| delta = float(first.attrs.get("delta_sec", 0.01)) |
| if abs(delta - 0.01) > 1e-6: |
| continue |
| start_time = first.attrs.get("start_time") |
| if not isinstance(start_time, str): |
| continue |
| btime = parse_time(start_time) |
| lengths = [int(station[c].shape[0]) for c in comps] |
| length = min(lengths) |
| picks: List[PhasePick] = [] |
| for phase in ("Pg", "Sg", "Pn", "Sn"): |
| chosen = choose_phase(station, phase, prefer_manual=prefer_manual) |
| if chosen is None: |
| continue |
| ptime, source = chosen |
| idx = int(round((parse_time(ptime) - btime).total_seconds() / delta)) |
| if 0 <= idx < length: |
| picks.append(PhasePick(phase, idx, source)) |
| if not picks: |
| continue |
| distances = [] |
| for phase in ("Pg", "Sg", "Pn", "Sn"): |
| for prefix in ("MANUAL.TRAVTIME", "RNN.TRAVTIME"): |
| dk = f"{prefix}.{phase}.dist_km" |
| if dk in station.attrs: |
| try: |
| distances.append(float(station.attrs[dk])) |
| except (TypeError, ValueError): |
| pass |
| dist = float(np.median(distances)) if distances else float("nan") |
| records.append( |
| Record( |
| event=event_key, |
| station=str(station_key), |
| length=length, |
| delta=delta, |
| distance_km=dist, |
| phases=tuple(picks), |
| ) |
| ) |
| return records |
|
|
|
|
| def crop_start_for_record(record: Record, rng: np.random.Generator, length: int, padlen: int) -> int | None: |
| if record.length < length: |
| return None |
| phase_indices = np.array([p.index for p in record.phases], dtype=np.int64) |
| for _ in range(24): |
| anchor = int(rng.choice(phase_indices)) |
| offset = int(rng.integers(padlen, max(padlen + 1, length - padlen))) |
| start = anchor - offset |
| if 0 <= start <= record.length - length: |
| return start |
| anchor = int(rng.choice(phase_indices)) |
| return int(np.clip(anchor - length // 2, 0, record.length - length)) |
|
|
|
|
| def make_specs( |
| records: Sequence[Record], |
| n_samples: int, |
| seed: int, |
| length: int, |
| padlen: int, |
| double_prob: float = 0.5, |
| valid_indices: Sequence[int] | None = None, |
| ) -> List[SampleSpec]: |
| valid = list(valid_indices) if valid_indices is not None else [i for i, r in enumerate(records) if r.length >= length] |
| if not valid: |
| raise RuntimeError("No records are long enough for the requested window length.") |
| rng = np.random.default_rng(seed) |
| specs: List[SampleSpec] = [] |
| while len(specs) < n_samples: |
| is_double = rng.random() < double_prob |
| crop_count = 2 if is_double else 1 |
| crops: List[CropSpec] = [] |
| for j in range(crop_count): |
| rec_idx = int(rng.choice(valid)) |
| start = crop_start_for_record(records[rec_idx], rng, length, padlen) |
| if start is None: |
| crops = [] |
| break |
| amp = 1.0 if j == 0 else float(rng.uniform(0.2, 5.0)) |
| crops.append(CropSpec(rec_idx, start, amp)) |
| if crops: |
| specs.append(SampleSpec(tuple(crops), "double" if is_double else "single")) |
| return specs |
|
|
|
|
| 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 load_crop(h5, record: Record, crop: CropSpec, length: int) -> Tuple[np.ndarray, List[Tuple[str, str, int]]]: |
| station = h5[record.event][record.station] |
| comps = component_keys(station) |
| if comps is None: |
| raise RuntimeError(f"Missing components for {record.event}/{record.station}") |
| data = [station[c][crop.start : crop.start + length] for c in comps] |
| 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 labels_to_target(labels: Sequence[Tuple[str, str, int]], length: int, sigma_samples: float = 10.0) -> np.ndarray: |
| target = np.zeros((5, length), dtype=DTYPE) |
| x = np.arange(length, dtype=DTYPE) |
| for phase, _, idx in labels: |
| ch = PHASE_TO_CHANNEL[phase] |
| pulse = np.exp(-0.5 * ((x - idx) / sigma_samples) ** 2) |
| target[ch] = np.maximum(target[ch], pulse.astype(DTYPE)) |
| target[0] = np.clip(1.0 - target[1:].sum(axis=0), 0.0, 1.0) |
| return target |
|
|
|
|
| def materialize_samples( |
| h5_path: Path, |
| records: Sequence[Record], |
| specs: Sequence[SampleSpec], |
| length: int, |
| progress_interval: int = 0, |
| ) -> Tuple[np.ndarray, List[List[Tuple[str, str, int]]], np.ndarray, List[str]]: |
| waves = np.zeros((len(specs), length, 3), dtype=DTYPE) |
| labels_all: List[List[Tuple[str, str, int]]] = [] |
| targets = np.zeros((len(specs), 5, length), dtype=DTYPE) |
| kinds: List[str] = [] |
| with h5py.File(h5_path, "r") as h5: |
| for i, spec in enumerate(specs): |
| if progress_interval > 0 and (i % progress_interval == 0 or i == len(specs) - 1): |
| print(f"materialized eval sample {i:,}/{len(specs):,}", flush=True) |
| mixed = np.zeros((length, 3), dtype=DTYPE) |
| labels: List[Tuple[str, str, int]] = [] |
| for crop in spec.crops: |
| wave, crop_labels = load_crop(h5, records[crop.rec_idx], crop, length) |
| mixed += wave |
| labels.extend(crop_labels) |
| mixed = normalize_wave(mixed) |
| waves[i] = mixed |
| labels_all.append(labels) |
| targets[i] = labels_to_target(labels, length) |
| kinds.append(spec.kind) |
| return waves, labels_all, targets, kinds |
|
|
|
|
| def sample_batch( |
| h5_path: Path, |
| records: Sequence[Record], |
| 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, |
| ) |
| waves, _, targets, _ = materialize_samples(h5_path, records, specs, length) |
| return waves, targets |
|
|
|
|
| def find_peaks(prob: np.ndarray, threshold: float, min_sep: int) -> List[Tuple[int, float]]: |
| if len(prob) < 3: |
| return [] |
| mask = (prob[1:-1] >= threshold) & (prob[1:-1] >= prob[:-2]) & (prob[1:-1] > prob[2:]) |
| candidates = np.where(mask)[0] + 1 |
| if candidates.size == 0: |
| return [] |
| order = candidates[np.argsort(prob[candidates])[::-1]] |
| selected: List[int] = [] |
| for idx in order: |
| if all(abs(int(idx) - old) >= min_sep for old in selected): |
| selected.append(int(idx)) |
| selected.sort() |
| return [(idx, float(prob[idx])) for idx in selected] |
|
|
|
|
| def match_predictions( |
| true_indices: Sequence[int], |
| pred_indices: Sequence[int], |
| tolerance: int, |
| ) -> Tuple[int, int, int, List[float]]: |
| unmatched = set(range(len(true_indices))) |
| tp = 0 |
| fp = 0 |
| errors: List[float] = [] |
| for pred in pred_indices: |
| best = None |
| best_dist = tolerance + 1 |
| for ti in unmatched: |
| dist = abs(pred - true_indices[ti]) |
| if dist <= tolerance and dist < best_dist: |
| best = ti |
| best_dist = dist |
| if best is None: |
| fp += 1 |
| else: |
| unmatched.remove(best) |
| tp += 1 |
| errors.append((pred - true_indices[best]) * 0.01) |
| fn = len(unmatched) |
| return tp, fp, fn, errors |
|
|
|
|
| def run_model( |
| model: BRNN, |
| waves: np.ndarray, |
| device: torch.device, |
| batch_size: int, |
| ) -> np.ndarray: |
| model.eval() |
| outputs: List[np.ndarray] = [] |
| with torch.no_grad(): |
| for start in range(0, len(waves), batch_size): |
| batch = torch.from_numpy(waves[start : start + batch_size]).to(device) |
| batch = batch.permute(0, 2, 1) |
| out = model(batch).detach().cpu().numpy() |
| outputs.append(out) |
| return np.concatenate(outputs, axis=0) |
|
|
|
|
| def evaluate_outputs( |
| outputs: np.ndarray, |
| labels_all: Sequence[Sequence[Tuple[str, str, int]]], |
| kinds: Sequence[str], |
| thresholds: Sequence[float], |
| min_sep: int, |
| tolerance: int, |
| ) -> Dict: |
| result = {"thresholds": list(map(float, thresholds)), "all": {}, "double": {}} |
| for subset_name, subset_mask in { |
| "all": np.ones(len(labels_all), dtype=bool), |
| "double": np.array([k == "double" for k in kinds], dtype=bool), |
| }.items(): |
| subset_result = {} |
| for group in ("P", "S"): |
| rows = [] |
| for thr in thresholds: |
| tp = fp = fn = 0 |
| errors: List[float] = [] |
| for i, labels in enumerate(labels_all): |
| if not subset_mask[i]: |
| continue |
| true_idx = [idx for _, g, idx in labels if g == group] |
| prob = outputs[i, GROUP_TO_CHANNELS[group], :].max(axis=0) |
| pred_idx = [idx for idx, _ in find_peaks(prob, thr, min_sep)] |
| mtp, mfp, mfn, merr = match_predictions(true_idx, pred_idx, tolerance) |
| tp += mtp |
| fp += mfp |
| fn += mfn |
| errors.extend(merr) |
| precision = tp / (tp + fp) if tp + fp else 0.0 |
| recall = tp / (tp + fn) if tp + fn else 0.0 |
| f1 = 2 * precision * recall / (precision + recall) if precision + recall else 0.0 |
| row = { |
| "threshold": float(thr), |
| "tp": int(tp), |
| "fp": int(fp), |
| "fn": int(fn), |
| "precision": float(precision), |
| "recall": float(recall), |
| "f1": float(f1), |
| "mean_error_s": float(np.mean(errors)) if errors else None, |
| "std_error_s": float(np.std(errors)) if errors else None, |
| "errors_s": errors if abs(thr - 0.1) < 1e-9 else [], |
| } |
| rows.append(row) |
| subset_result[group] = rows |
| result[subset_name] = subset_result |
| return result |
|
|
|
|
| def metric_at(metrics: Dict, subset: str, group: str, threshold: float) -> Dict: |
| rows = metrics[subset][group] |
| return min(rows, key=lambda r: abs(r["threshold"] - threshold)) |
|
|
|
|
| def train_transfer( |
| h5_path: Path, |
| records: Sequence[Record], |
| 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, |
| ) -> BRNN: |
| torch.manual_seed(seed) |
| np.random.seed(seed) |
| random.seed(seed) |
| model = BRNN().to(device) |
| model.load_state_dict(torch.load(base_ckpt, map_location="cpu")) |
| model.train() |
| loss_fn = Loss().to(device) |
| opt = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=1e-4) |
| log_csv.parent.mkdir(parents=True, exist_ok=True) |
| with log_csv.open("w", newline="") as f: |
| writer = csv.writer(f) |
| writer.writerow(["step", "loss"]) |
| for step in range(steps): |
| waves, targets = sample_batch(h5_path, records, 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 % 25 == 0 or step == steps - 1: |
| print(f"train step {step:05d}/{steps} loss={loss_value:.3f}", flush=True) |
| torch.save(model.state_dict(), out_ckpt) |
| return model |
|
|
|
|
| def load_model(ckpt: Path, device: torch.device) -> BRNN: |
| model = BRNN().to(device) |
| model.load_state_dict(torch.load(ckpt, map_location="cpu")) |
| model.eval() |
| return model |
|
|
|
|
| def plot_training_samples(waves: np.ndarray, labels_all, out: Path) -> None: |
| single_idx = 0 |
| double_idx = next((i for i, labels in enumerate(labels_all) if len(labels) >= 3), min(1, len(labels_all) - 1)) |
| fig, axes = plt.subplots(2, 2, figsize=(11, 6.8), dpi=220) |
| time = np.arange(waves.shape[1]) * 0.01 |
| for ax, idx, title in [ |
| (axes[0, 0], single_idx, "(a) Single-event waveform"), |
| (axes[0, 1], double_idx, "(b) Two-event mixed waveform"), |
| ]: |
| for ci, lab in enumerate(["E", "N", "Z"]): |
| ax.plot(time, waves[idx, :, ci] + (2 - ci) * 2.2, lw=0.55, color="black") |
| ax.text(time[0] - 1.0, (2 - ci) * 2.2, lab, va="center", ha="right", fontsize=8) |
| ax.set_title(title, loc="left", fontsize=10) |
| ax.set_xlabel("Time (s)") |
| ax.set_yticks([]) |
| ax.set_xlim(time[0], time[-1]) |
| for ax, idx, title in [ |
| (axes[1, 0], single_idx, "(c) Single-event labels"), |
| (axes[1, 1], double_idx, "(d) Mixed-event labels"), |
| ]: |
| target = labels_to_target(labels_all[idx], waves.shape[1]) |
| ax.plot(time, target[1] + target[3], color="#d62728", label="P", lw=0.9) |
| ax.plot(time, target[2] + target[4], color="#1f77b4", label="S", lw=0.9) |
| ax.set_title(title, loc="left", fontsize=10) |
| ax.set_xlabel("Time (s)") |
| ax.set_ylabel("Probability") |
| ax.set_ylim(-0.05, 1.05) |
| ax.set_xlim(time[0], time[-1]) |
| ax.legend(frameon=False, fontsize=8, loc="upper right") |
| fig.tight_layout() |
| fig.savefig(out, bbox_inches="tight") |
| plt.close(fig) |
|
|
|
|
| def plot_pr(base_metrics: Dict, multi_metrics: Dict, out: Path) -> None: |
| fig, axes = plt.subplots(2, 2, figsize=(9.2, 6.8), dpi=220, sharex=True, sharey=True) |
| panels = [ |
| (axes[0, 0], base_metrics, "P", "(a) RNN, P"), |
| (axes[0, 1], base_metrics, "S", "(b) RNN, S"), |
| (axes[1, 0], multi_metrics, "P", "(c) Multi-RNN, P"), |
| (axes[1, 1], multi_metrics, "S", "(d) Multi-RNN, S"), |
| ] |
| for ax, metrics, group, title in panels: |
| rows = metrics["all"][group] |
| x = [r["threshold"] for r in rows] |
| ax.plot(x, [r["precision"] for r in rows], marker="o", ms=3, label="Precision") |
| ax.plot(x, [r["recall"] for r in rows], marker="s", ms=3, label="Recall") |
| ax.plot(x, [r["f1"] for r in rows], marker="^", ms=3, label="F1") |
| ax.set_title(title, loc="left", fontsize=10) |
| ax.grid(True, alpha=0.25) |
| ax.set_ylim(0, 1.02) |
| ax.set_xlabel("Minimum confidence threshold") |
| ax.set_ylabel("Value") |
| axes[0, 0].legend(frameon=False, fontsize=8, loc="best") |
| fig.tight_layout() |
| fig.savefig(out, bbox_inches="tight") |
| plt.close(fig) |
|
|
|
|
| def plot_errors(base_metrics: Dict, multi_metrics: Dict, subset: str, out: Path) -> None: |
| fig, axes = plt.subplots(2, 2, figsize=(9.2, 6.8), dpi=220, sharex=True) |
| panels = [ |
| (axes[0, 0], base_metrics, "P", "(a) RNN, P", "#ff6b6b"), |
| (axes[0, 1], base_metrics, "S", "(b) RNN, S", "#6b6bff"), |
| (axes[1, 0], multi_metrics, "P", "(c) Multi-RNN, P", "#ff6b6b"), |
| (axes[1, 1], multi_metrics, "S", "(d) Multi-RNN, S", "#6b6bff"), |
| ] |
| for ax, metrics, group, title, color in panels: |
| row = metric_at(metrics, subset, group, 0.1) |
| errors = np.array(row["errors_s"], dtype=float) |
| ax.hist(errors, bins=np.linspace(-2, 2, 81), color=color, alpha=0.82) |
| ax.axvline(0, color="#2ca6df", ls="--", lw=0.9) |
| text = ( |
| f"P={row['precision']:.3f}\n" |
| f"R={row['recall']:.3f}\n" |
| f"F1={row['f1']:.3f}\n" |
| f"mean={row['mean_error_s'] * 1000:.1f} ms\n" |
| f"std={row['std_error_s'] * 1000:.1f} ms" |
| ) |
| ax.text(0.03, 0.95, text, transform=ax.transAxes, va="top", ha="left", fontsize=8) |
| ax.set_title(title, loc="left", fontsize=10) |
| ax.set_xlabel("Error (s)") |
| ax.set_ylabel("Count") |
| ax.set_xlim(-2, 2) |
| ax.grid(True, alpha=0.18) |
| fig.tight_layout() |
| fig.savefig(out, bbox_inches="tight") |
| plt.close(fig) |
|
|
|
|
| def select_continuous_event(records: Sequence[Record], min_records: int = 24) -> str: |
| by_event: Dict[str, List[Record]] = {} |
| for record in records: |
| if math.isnan(record.distance_km): |
| continue |
| by_event.setdefault(record.event, []).append(record) |
| candidates = sorted( |
| ((event, recs) for event, recs in by_event.items() if len(recs) >= min_records), |
| key=lambda item: len(item[1]), |
| reverse=True, |
| ) |
| if not candidates: |
| return max(by_event.items(), key=lambda item: len(item[1]))[0] |
| return candidates[0][0] |
|
|
|
|
| def continuous_specs_for_event(records: Sequence[Record], event: str, length: int) -> List[Tuple[int, int]]: |
| out = [] |
| for idx, record in enumerate(records): |
| if record.event != event or record.length < length: |
| continue |
| phase_indices = [p.index for p in record.phases] |
| if not phase_indices: |
| continue |
| start = int(np.clip(min(phase_indices) - 800, 0, record.length - length)) |
| out.append((idx, start)) |
| out.sort(key=lambda item: records[item[0]].distance_km) |
| return out[:60] |
|
|
|
|
| def evaluate_continuous( |
| model: BRNN, |
| h5_path: Path, |
| records: Sequence[Record], |
| specs: Sequence[Tuple[int, int]], |
| length: int, |
| device: torch.device, |
| threshold: float, |
| ) -> Tuple[np.ndarray, List[List[Tuple[str, str, int]]], np.ndarray, Dict]: |
| sample_specs = [SampleSpec((CropSpec(rec_idx, start, 1.0),), "continuous") for rec_idx, start in specs] |
| waves, labels_all, _, _ = materialize_samples(h5_path, records, sample_specs, length) |
| outputs = run_model(model, waves, device, batch_size=32) |
| metrics = evaluate_outputs( |
| outputs, |
| labels_all, |
| ["continuous"] * len(labels_all), |
| [threshold], |
| min_sep=50, |
| tolerance=100, |
| ) |
| return waves, labels_all, outputs, metrics |
|
|
|
|
| def plot_continuous( |
| records: Sequence[Record], |
| specs: Sequence[Tuple[int, int]], |
| waves: np.ndarray, |
| base_outputs: np.ndarray, |
| multi_outputs: np.ndarray, |
| threshold: float, |
| out: Path, |
| ) -> None: |
| fig, axes = plt.subplots(1, 2, figsize=(9.2, 7.2), dpi=220, sharex=True, sharey=True) |
| time = np.arange(waves.shape[1]) * 0.01 |
| for ax, outputs, title in [ |
| (axes[0], base_outputs, "(a) RNN"), |
| (axes[1], multi_outputs, "(b) Multi-RNN"), |
| ]: |
| for i, (rec_idx, _) in enumerate(specs): |
| rec = records[rec_idx] |
| dist = rec.distance_km if not math.isnan(rec.distance_km) else i |
| trace = waves[i, :, 2] |
| trace = trace / (np.max(np.abs(trace)) + 1e-6) * 2.0 + dist |
| ax.plot(time, trace, color="black", lw=0.35, alpha=0.75) |
| for group, color in [("P", "#d62728"), ("S", "#1f77b4")]: |
| prob = outputs[i, GROUP_TO_CHANNELS[group], :].max(axis=0) |
| peaks = find_peaks(prob, threshold, min_sep=50) |
| if peaks: |
| ax.scatter([p[0] * 0.01 for p in peaks], [dist] * len(peaks), s=6, color=color, alpha=0.9) |
| ax.set_title(title, loc="left", fontsize=10) |
| ax.set_xlabel("Time (s)") |
| ax.grid(True, alpha=0.15) |
| axes[0].set_ylabel("Epicentral distance (km)") |
| axes[1].scatter([], [], s=12, color="#d62728", label="P") |
| axes[1].scatter([], [], s=12, color="#1f77b4", label="S") |
| axes[1].legend(frameon=False, fontsize=8, loc="upper right") |
| fig.tight_layout() |
| fig.savefig(out, bbox_inches="tight") |
| plt.close(fig) |
|
|
|
|
| def strip_errors(metrics: Dict) -> Dict: |
| clean = json.loads(json.dumps(metrics)) |
| for subset in ("all", "double"): |
| for group in ("P", "S"): |
| for row in clean[subset][group]: |
| row["error_count"] = len(row.get("errors_s", [])) |
| row.pop("errors_s", None) |
| return clean |
|
|
|
|
| def summarize_records(records: Sequence[Record]) -> Dict: |
| source_counts: Dict[str, int] = {} |
| phase_counts: Dict[str, int] = {} |
| for record in records: |
| for pick in record.phases: |
| source_counts[pick.source] = source_counts.get(pick.source, 0) + 1 |
| phase_counts[pick.phase] = phase_counts.get(pick.phase, 0) + 1 |
| return { |
| "records": len(records), |
| "source_counts": source_counts, |
| "phase_counts": phase_counts, |
| } |
|
|
|
|
| def main() -> None: |
| parser = argparse.ArgumentParser() |
| 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/repro_seed20260607") |
| parser.add_argument("--seed", type=int, default=20260607) |
| parser.add_argument("--length", type=int, default=5120) |
| parser.add_argument("--padlen", type=int, default=512) |
| parser.add_argument("--train-steps", type=int, default=400) |
| parser.add_argument("--train-batch", type=int, default=16) |
| parser.add_argument("--eval-samples", type=int, default=20000) |
| parser.add_argument("--eval-batch", type=int, default=64) |
| parser.add_argument("--max-train-events", type=int, default=8000) |
| parser.add_argument("--max-test-events", type=int, default=0) |
| parser.add_argument("--lr", type=float, default=1e-4) |
| parser.add_argument("--skip-train", action="store_true") |
| args = parser.parse_args() |
|
|
| out_dir = Path(args.out_dir) |
| out_dir.mkdir(parents=True, exist_ok=True) |
| h5_path = Path(args.h5) |
| key_path = Path(args.keys) |
| base_ckpt = Path(args.base_ckpt) |
| multi_ckpt = out_dir / f"pnsn.v3.multirnn.seed{args.seed}.pt" |
|
|
| torch.manual_seed(args.seed) |
| np.random.seed(args.seed) |
| random.seed(args.seed) |
| device = torch.device("mps") if torch.backends.mps.is_available() else torch.device("cpu") |
| print(f"device={device}", flush=True) |
|
|
| max_test = None if args.max_test_events == 0 else args.max_test_events |
| train_records = build_records(h5_path, key_path, "train", args.max_train_events) |
| test_records = build_records(h5_path, key_path, "test", max_test) |
| print(f"train records={len(train_records)} test records={len(test_records)}", flush=True) |
|
|
| with (out_dir / "record_summary.json").open("w") as f: |
| json.dump( |
| { |
| "train": summarize_records(train_records), |
| "test": summarize_records(test_records), |
| "seed": args.seed, |
| "length_samples": args.length, |
| "length_seconds": args.length * 0.01, |
| "label_rule": "MANUAL.TRAVTIME preferred; RNN.TRAVTIME used only when tag=Y.", |
| }, |
| f, |
| ensure_ascii=False, |
| indent=2, |
| ) |
|
|
| if args.skip_train and multi_ckpt.exists(): |
| multi_model = load_model(multi_ckpt, device) |
| else: |
| multi_model = train_transfer( |
| h5_path=h5_path, |
| records=train_records, |
| base_ckpt=base_ckpt, |
| out_ckpt=multi_ckpt, |
| log_csv=out_dir / "transfer_loss.csv", |
| seed=args.seed, |
| steps=args.train_steps, |
| batch_size=args.train_batch, |
| length=args.length, |
| padlen=args.padlen, |
| lr=args.lr, |
| device=device, |
| ) |
|
|
| eval_specs = make_specs( |
| test_records, |
| n_samples=args.eval_samples, |
| seed=args.seed + 17, |
| length=args.length, |
| padlen=args.padlen, |
| double_prob=0.5, |
| ) |
| waves, labels_all, _, kinds = materialize_samples(h5_path, test_records, eval_specs, args.length) |
| np.savez_compressed( |
| out_dir / "eval_sample_summary.npz", |
| kinds=np.array(kinds), |
| label_counts=np.array([len(x) for x in labels_all], dtype=np.int16), |
| ) |
| plot_training_samples(waves[:100], labels_all[:100], out_dir / "fig5_training_samples_composite.png") |
|
|
| thresholds = [round(x, 1) for x in np.arange(0.1, 1.0, 0.1)] |
| base_model = load_model(base_ckpt, device) |
| base_outputs = run_model(base_model, waves, device, args.eval_batch) |
| multi_outputs = run_model(multi_model, waves, device, args.eval_batch) |
| base_metrics = evaluate_outputs(base_outputs, labels_all, kinds, thresholds, min_sep=50, tolerance=100) |
| multi_metrics = evaluate_outputs(multi_outputs, labels_all, kinds, thresholds, min_sep=50, tolerance=100) |
|
|
| plot_pr(base_metrics, multi_metrics, out_dir / "fig6_pr_composite.png") |
| plot_errors(base_metrics, multi_metrics, "all", out_dir / "fig7_error_all_composite.png") |
| plot_errors(base_metrics, multi_metrics, "double", out_dir / "fig8_error_double_composite.png") |
|
|
| event = select_continuous_event(test_records) |
| cont_specs = continuous_specs_for_event(test_records, event, min(args.length * 2, 10240)) |
| cont_len = min(args.length * 2, 10240) |
| if len(cont_specs) >= 4: |
| cont_base_waves, cont_labels, cont_base_outputs, cont_base_metrics = evaluate_continuous( |
| base_model, h5_path, test_records, cont_specs, cont_len, device, threshold=0.1 |
| ) |
| _, _, cont_multi_outputs, cont_multi_metrics = evaluate_continuous( |
| multi_model, h5_path, test_records, cont_specs, cont_len, device, threshold=0.1 |
| ) |
| plot_continuous( |
| test_records, |
| cont_specs, |
| cont_base_waves, |
| cont_base_outputs, |
| cont_multi_outputs, |
| 0.1, |
| out_dir / "fig9_continuous_composite.png", |
| ) |
| else: |
| cont_base_metrics = {} |
| cont_multi_metrics = {} |
|
|
| report = { |
| "seed": args.seed, |
| "device": str(device), |
| "length_samples": args.length, |
| "length_seconds": args.length * 0.01, |
| "eval_samples": args.eval_samples, |
| "eval_single_samples": int(sum(k == "single" for k in kinds)), |
| "eval_double_samples": int(sum(k == "double" for k in kinds)), |
| "threshold_for_error_stats": 0.1, |
| "match_tolerance_seconds": 1.0, |
| "min_peak_separation_seconds": 0.5, |
| "base_metrics": strip_errors(base_metrics), |
| "multi_metrics": strip_errors(multi_metrics), |
| "continuous_event": event, |
| "continuous_station_count": len(cont_specs), |
| "continuous_base_metrics": strip_errors(cont_base_metrics) if cont_base_metrics else {}, |
| "continuous_multi_metrics": strip_errors(cont_multi_metrics) if cont_multi_metrics else {}, |
| "figures": { |
| "fig5": str((out_dir / "fig5_training_samples_composite.png").resolve()), |
| "fig6": str((out_dir / "fig6_pr_composite.png").resolve()), |
| "fig7": str((out_dir / "fig7_error_all_composite.png").resolve()), |
| "fig8": str((out_dir / "fig8_error_double_composite.png").resolve()), |
| "fig9": str((out_dir / "fig9_continuous_composite.png").resolve()), |
| }, |
| } |
| with (out_dir / "metrics_report.json").open("w") as f: |
| json.dump(report, f, ensure_ascii=False, indent=2) |
|
|
| print(json.dumps(report, ensure_ascii=False, indent=2)[:4000], flush=True) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
