| |
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| from pathlib import Path |
| from typing import Any |
|
|
| import h5py |
| import numpy as np |
| import pandas as pd |
| import torch |
| from scipy.signal import find_peaks |
| from seisbench.models import EQTransformer |
| from torch.utils.data import DataLoader, Dataset |
| from tqdm import tqdm |
|
|
|
|
| def normalize_std(waveform: np.ndarray) -> np.ndarray: |
| if waveform.ndim != 2: |
| raise ValueError(f"Expected 2D waveform, got {waveform.shape}") |
| if waveform.shape[-1] != 3 and waveform.shape[0] == 3: |
| waveform = waveform.T |
| if waveform.shape[-1] != 3: |
| raise ValueError(f"Expected waveform as (T, 3) or (3, T), got {waveform.shape}") |
| mean = waveform.mean(axis=0, keepdims=True) |
| std = waveform.std(axis=0, keepdims=True) |
| std[std < 1e-12] = 1.0 |
| return ((waveform - mean) / std).T.astype(np.float32, copy=False) |
|
|
|
|
| class H5TestDataset(Dataset): |
| def __init__(self, h5_path: Path, data_key: str = "data", max_samples: int = 0, cache_waveforms: bool = False) -> None: |
| self.h5_path = h5_path |
| self.data_key = data_key |
| self._h5: h5py.File | None = None |
| self._cached_waveforms: np.ndarray | None = None |
| with h5py.File(h5_path, "r") as f: |
| self.n = int(f[data_key].shape[0]) |
| self.npts = int(f[data_key].shape[1] if f[data_key].shape[-1] == 3 else f[data_key].shape[-1]) |
| if max_samples > 0: |
| self.n = min(self.n, int(max_samples)) |
| if cache_waveforms: |
| waves = np.asarray(f[data_key][: self.n], dtype=np.float32) |
| if waves.ndim != 3: |
| raise ValueError(f"Expected waveform data with 3 dims, got {waves.shape}") |
| if waves.shape[1] == 3: |
| mean = waves.mean(axis=2, keepdims=True) |
| std = waves.std(axis=2, keepdims=True) |
| std[std < 1e-12] = 1.0 |
| waves = (waves - mean) / std |
| elif waves.shape[-1] == 3: |
| mean = waves.mean(axis=1, keepdims=True) |
| std = waves.std(axis=1, keepdims=True) |
| std[std < 1e-12] = 1.0 |
| waves = ((waves - mean) / std).transpose(0, 2, 1) |
| else: |
| raise ValueError(f"Expected waveform shape (N, 3, T) or (N, T, 3), got {waves.shape}") |
| self._cached_waveforms = np.ascontiguousarray(waves, dtype=np.float32) |
| if max_samples > 0: |
| self.n = min(self.n, int(max_samples)) |
|
|
| def __len__(self) -> int: |
| return self.n |
|
|
| def _ensure_h5(self) -> h5py.File: |
| if self._h5 is None: |
| self._h5 = h5py.File(self.h5_path, "r") |
| return self._h5 |
|
|
| def __getitem__(self, idx: int) -> dict[str, torch.Tensor | int]: |
| if self._cached_waveforms is not None: |
| waveform = self._cached_waveforms[idx] |
| else: |
| h5 = self._ensure_h5() |
| waveform = np.asarray(h5[self.data_key][idx], dtype=np.float32) |
| waveform = normalize_std(waveform) |
| return {"waveform": torch.from_numpy(waveform), "index": int(idx)} |
|
|
|
|
| def read_h5_phase_labels(h5_path: Path, n: int) -> tuple[np.ndarray, np.ndarray] | None: |
| with h5py.File(h5_path, "r") as f: |
| if "labels" not in f: |
| return None |
| labels = np.asarray(f["labels"][:n], dtype=np.float64) |
| if labels.ndim != 2 or labels.shape[1] < 3: |
| return None |
| return labels[:, 1], labels[:, 2] |
|
|
|
|
| def validate_meta_against_h5_labels(h5_path: Path, meta: pd.DataFrame, n: int) -> None: |
| h5_labels = read_h5_phase_labels(h5_path, n) |
| if h5_labels is None or not {"p_index", "s_index"}.issubset(meta.columns): |
| return |
| h5_p, h5_s = h5_labels |
| meta_p = pd.to_numeric(meta["p_index"], errors="coerce").to_numpy(dtype=float) |
| meta_s = pd.to_numeric(meta["s_index"], errors="coerce").to_numpy(dtype=float) |
| p_ok = np.isclose(h5_p, meta_p, atol=0.5, equal_nan=False) |
| s_ok = np.isclose(h5_s, meta_s, atol=0.5, equal_nan=False) |
| mismatch = int(np.sum(~(p_ok & s_ok))) |
| if mismatch: |
| raise ValueError( |
| f"test_meta is not aligned with HDF5 labels: {mismatch}/{n} rows differ. " |
| "Use the metadata generated with this HDF5 file, or add an explicit source_index mapping before evaluation." |
| ) |
|
|
|
|
| def phase_peak_summary( |
| prob: np.ndarray, |
| true_sample: int, |
| threshold: float, |
| distance: int, |
| tolerance_samples: int, |
| ) -> dict[str, Any]: |
| peaks, props = find_peaks(prob, height=threshold, distance=distance) |
| if peaks.size == 0: |
| peak_idx = int(np.argmax(prob)) if prob.size else -1 |
| return { |
| "pred": peak_idx, |
| "score": float(prob[peak_idx]) if peak_idx >= 0 else float("nan"), |
| "detected": False, |
| "correct": False, |
| "delta": float(peak_idx - true_sample) if peak_idx >= 0 and true_sample >= 0 else float("nan"), |
| "delta_recalled": float("nan"), |
| "peak_count": 0, |
| "correct_peak_count": 0, |
| "peak_samples": "", |
| "peak_scores": "", |
| "matched_peak_samples": "", |
| } |
| heights = props["peak_heights"] |
| deltas = peaks.astype(np.int64) - int(true_sample) |
| correct_mask = np.abs(deltas) <= int(tolerance_samples) if true_sample >= 0 else np.zeros(peaks.shape, dtype=bool) |
| matched_count = 1 if np.any(correct_mask) else 0 |
| if np.any(correct_mask): |
| correct_indices = np.where(correct_mask)[0] |
| best = int(correct_indices[np.argmin(np.abs(deltas[correct_indices]))]) |
| else: |
| best = int(np.argmax(heights)) |
| return { |
| "pred": int(peaks[best]), |
| "score": float(heights[best]), |
| "detected": True, |
| "correct": bool(np.any(correct_mask)), |
| "delta": float(deltas[best]) if true_sample >= 0 else float("nan"), |
| "delta_recalled": float(deltas[best]) if np.any(correct_mask) else float("nan"), |
| "peak_count": int(peaks.size), |
| "correct_peak_count": int(matched_count), |
| "peak_samples": ";".join(str(int(value)) for value in peaks), |
| "peak_scores": ";".join(f"{float(value):.4f}" for value in heights), |
| "matched_peak_samples": str(int(peaks[best])) if matched_count else "", |
| } |
|
|
|
|
| def summarize_predictions(df: pd.DataFrame, tolerance_samples: int) -> dict[str, Any]: |
| summary: dict[str, Any] = {"n": int(len(df)), "tolerance_samples": int(tolerance_samples)} |
| for phase in ["p", "s"]: |
| correct_col = f"{phase}_correct" |
| peak_col = f"{phase}_num_peaks" |
| correct_peak_col = f"{phase}_correct_peak_count" |
| delta_col = f"{phase}_delta_recalled_samples" |
| if correct_col not in df.columns: |
| continue |
| correct = df[correct_col].fillna(False).astype(bool).to_numpy() |
| recall = float(correct.mean()) if len(correct) else np.nan |
| correct_peak_sum = pd.to_numeric(df[correct_peak_col], errors="coerce").fillna(0).sum() |
| peak_sum = pd.to_numeric(df[peak_col], errors="coerce").fillna(0).sum() |
| precision = float(correct_peak_sum / max(float(peak_sum), 1.0)) |
| f1 = 2.0 * recall * precision / (recall + precision) if np.isfinite(recall) and np.isfinite(precision) and (recall + precision) > 0 else np.nan |
| err = pd.to_numeric(df[delta_col], errors="coerce").to_numpy(dtype=float) |
| summary[f"{phase}_recall"] = float(recall) |
| summary[f"{phase}_precision"] = float(precision) |
| summary[f"{phase}_f1"] = float(f1) |
| summary[f"{phase}_mae_samples"] = float(np.nanmean(np.abs(err))) |
| summary[f"{phase}_bias_samples"] = float(np.nanmean(err)) |
| summary[f"{phase}_bias_sd_samples"] = float(np.nanstd(err, ddof=1)) if np.sum(np.isfinite(err)) > 1 else float("nan") |
| return summary |
|
|
|
|
| def print_summary(summary: dict[str, Any], out_dir: Path) -> None: |
| parts = [] |
| for phase in ["p", "s"]: |
| if f"{phase}_recall" not in summary: |
| continue |
| parts.append( |
| f"{phase.upper()} R={summary.get(f'{phase}_recall', float('nan')):.3f} " |
| f"Prec={summary.get(f'{phase}_precision', float('nan')):.3f} " |
| f"F1={summary.get(f'{phase}_f1', float('nan')):.3f} " |
| f"bias={summary.get(f'{phase}_bias_samples', float('nan')):.1f}+/-{summary.get(f'{phase}_bias_sd_samples', float('nan')):.1f}" |
| ) |
| print(f"n={summary.get('n', 0)} | " + " | ".join(parts) + f" | out={out_dir}") |
|
|
|
|
| def parse_args() -> argparse.Namespace: |
| parser = argparse.ArgumentParser(description="Predict Iquique test set with a trained EQTransformer checkpoint.") |
| parser.add_argument("--checkpoint", default="work/EQT/eqtransformer_original_state_dict.pt", help="Path to best_val_state_dict.pt or another EQT state dict.") |
| parser.add_argument("--test-h5", default="data/processed/iquique_decay_arias90/test/stead_100hz_60s_test.h5") |
| parser.add_argument("--test-meta", default="data/processed/iquique_decay_arias90/test/stead_100hz_60s_test_meta.csv") |
| parser.add_argument("--output-dir", required=True) |
| parser.add_argument("--batch-size", type=int, default=128) |
| parser.add_argument("--workers", type=int, default=8) |
| parser.add_argument("--device", default="cuda") |
| parser.add_argument("--threshold", type=float, default=0.1) |
| parser.add_argument("--peak-distance-samples", type=int, default=100) |
| parser.add_argument("--tolerance-samples", type=int, default=50) |
| parser.add_argument("--max-samples", type=int, default=0) |
| parser.add_argument("--phase", choices=["both", "p", "s"], default="both", help="Which phase picks to decode and report.") |
| parser.add_argument("--save-probs", action="store_true", help="Save detection/P/S probability curves as float16 npz.") |
| parser.add_argument("--cache-waveforms", action="store_true", help="Load and standardize the whole test HDF5 into memory before prediction.") |
| return parser.parse_args() |
|
|
|
|
| def build_eqtransformer_for_state(npts: int, state: dict[str, torch.Tensor]) -> EQTransformer: |
| uses_pretrained_lstm = any(".lstm.cell_f." in key or ".lstm.cell_b." in key for key in state) |
| if uses_pretrained_lstm: |
| model = EQTransformer.from_pretrained("original") |
| if getattr(model, "in_samples", npts) != npts: |
| print(f"WARNING: pretrained original in_samples={getattr(model, 'in_samples', None)}, data npts={npts}.") |
| return model |
| return EQTransformer(in_samples=npts) |
|
|
|
|
| def main() -> None: |
| args = parse_args() |
| checkpoint = Path(args.checkpoint).expanduser().resolve() |
| test_h5 = Path(args.test_h5).expanduser().resolve() |
| test_meta = Path(args.test_meta).expanduser().resolve() |
| out_dir = Path(args.output_dir).expanduser().resolve() |
| out_dir.mkdir(parents=True, exist_ok=True) |
|
|
| ds = H5TestDataset(test_h5, max_samples=int(args.max_samples), cache_waveforms=bool(args.cache_waveforms)) |
| if args.cache_waveforms: |
| print(f"Cached test waveforms in memory: n={len(ds)}, npts={ds.npts}") |
| loader = DataLoader( |
| ds, |
| batch_size=int(args.batch_size), |
| shuffle=False, |
| num_workers=int(args.workers), |
| pin_memory=torch.cuda.is_available() and str(args.device).startswith("cuda"), |
| persistent_workers=int(args.workers) > 0, |
| ) |
| device = torch.device(str(args.device) if torch.cuda.is_available() else "cpu") |
| state = torch.load(checkpoint, map_location=device) |
| model = build_eqtransformer_for_state(ds.npts, state).to(device) |
| model.load_state_dict(state) |
| model.eval() |
|
|
| meta = pd.read_csv(test_meta).iloc[: len(ds)].reset_index(drop=True) if test_meta.exists() else pd.DataFrame(index=np.arange(len(ds))) |
| validate_meta_against_h5_labels(test_h5, meta, len(ds)) |
| rows: list[dict[str, Any]] = [] |
| prob_chunks: list[np.ndarray] = [] |
| with torch.no_grad(): |
| for batch in tqdm(loader, desc="predict", unit="batch"): |
| x = batch["waveform"].to(device, non_blocking=True) |
| det, p, s = model(x, logits=True) |
| probs = torch.stack([torch.sigmoid(det), torch.sigmoid(p), torch.sigmoid(s)], dim=1).cpu().numpy() |
| indices = batch["index"].cpu().numpy() |
| if args.save_probs: |
| prob_chunks.append(probs.astype(np.float16)) |
| for j, idx in enumerate(indices): |
| det_i, p_i, s_i = probs[j] |
| p_true = int(round(float(meta.loc[int(idx), "p_index"]))) if "p_index" in meta.columns else -1 |
| s_true = int(round(float(meta.loc[int(idx), "s_index"]))) if "s_index" in meta.columns else -1 |
| row: dict[str, Any] = { |
| "index": int(idx), |
| "det_max": float(np.max(det_i)), |
| } |
| if args.phase in {"both", "p"}: |
| p_pick = phase_peak_summary(p_i, p_true, float(args.threshold), int(args.peak_distance_samples), int(args.tolerance_samples)) |
| row.update( |
| { |
| "p_pred": int(p_pick["pred"]), |
| "p_score": float(p_pick["score"]), |
| "p_num_peaks": int(p_pick["peak_count"]), |
| "p_correct_peak_count": int(p_pick["correct_peak_count"]), |
| "p_peak_samples": str(p_pick["peak_samples"]), |
| "p_peak_scores": str(p_pick["peak_scores"]), |
| "p_matched_peak_samples": str(p_pick["matched_peak_samples"]), |
| "p_correct": bool(p_pick["correct"]), |
| "p_delta_samples": float(p_pick["delta"]), |
| "p_delta_recalled_samples": float(p_pick["delta_recalled"]), |
| } |
| ) |
| if args.phase in {"both", "s"}: |
| s_pick = phase_peak_summary(s_i, s_true, float(args.threshold), int(args.peak_distance_samples), int(args.tolerance_samples)) |
| row.update( |
| { |
| "s_pred": int(s_pick["pred"]), |
| "s_score": float(s_pick["score"]), |
| "s_num_peaks": int(s_pick["peak_count"]), |
| "s_correct_peak_count": int(s_pick["correct_peak_count"]), |
| "s_peak_samples": str(s_pick["peak_samples"]), |
| "s_peak_scores": str(s_pick["peak_scores"]), |
| "s_matched_peak_samples": str(s_pick["matched_peak_samples"]), |
| "s_correct": bool(s_pick["correct"]), |
| "s_delta_samples": float(s_pick["delta"]), |
| "s_delta_recalled_samples": float(s_pick["delta_recalled"]), |
| } |
| ) |
| if "sample_id" in meta.columns: |
| row["sample_id"] = meta.loc[int(idx), "sample_id"] |
| if "trace_name" in meta.columns: |
| row["trace_name"] = meta.loc[int(idx), "trace_name"] |
| if "p_index" in meta.columns and args.phase in {"both", "p"}: |
| row["p_true"] = p_true |
| row["p_error_samples"] = row["p_delta_samples"] |
| if "s_index" in meta.columns and args.phase in {"both", "s"}: |
| row["s_true"] = s_true |
| row["s_error_samples"] = row["s_delta_samples"] |
| rows.append(row) |
|
|
| pred_df = pd.DataFrame(rows) |
| pred_df.to_csv(out_dir / "eqt_test_predictions.csv", index=False) |
| summary = summarize_predictions(pred_df, int(args.tolerance_samples)) |
| summary.update( |
| { |
| "checkpoint": str(checkpoint), |
| "test_h5": str(test_h5), |
| "test_meta": str(test_meta), |
| "threshold": float(args.threshold), |
| "peak_distance_samples": int(args.peak_distance_samples), |
| "phase": args.phase, |
| } |
| ) |
| with (out_dir / "eqt_test_summary.json").open("w", encoding="utf-8") as f: |
| json.dump(summary, f, indent=2) |
| if args.save_probs and prob_chunks: |
| np.savez_compressed(out_dir / "eqt_test_probabilities_float16.npz", probs=np.concatenate(prob_chunks, axis=0)) |
| print_summary(summary, out_dir) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
