bepad 0.1.0__py3-none-any.whl

bepad/stats.py

@@ -0,0 +1,140 @@
+"""Dataset statistics: tables that describe what is in a BEPAD split.
+
+Every function takes a loaded split (``bepad.load(split)``, with labels) and
+returns a pandas DataFrame; ``bepad.plots`` draws them.
+
+    network_scales(split)     network -> "local" | "regional" | "global"
+    source_events(split)      one row per source earthquake (catalog values + how often it is used)
+    event_pairs(split)        one row per pair of earthquakes in the same window: distance, time apart
+    window_geometry(split)    one row per window: stations, spacing, aperture, distances
+    pick_distances(split)     one row per pick: type, phase, distance to its earthquake, travel time
+
+Large splits are sampled with ``max_windows`` (per network) where the table is per pick or
+per station, so the functions stay fast on the full release.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+import pandas as pd
+
+from .build.geo import haversine_km
+
+#: typical pick distance [km] separating the three scales
+SCALE_EDGES = (150.0, 1500.0)
+SCALES = ("local", "regional", "global")
+
+
+def _sample_ids(split, max_windows, seed=0):
+    w = split.windows
+    if max_windows is None:
+        return set(w.window_id)
+    return set(w.groupby("network", group_keys=False)
+               .apply(lambda g: g.sample(min(len(g), max_windows), random_state=seed)).window_id)
+
+
+def _coords(split):
+    st = split.window_stations.drop_duplicates(["network", "station"]).set_index(["network", "station"])
+    return st[["latitude", "longitude"]]
+
+
+def pick_distances(split, max_windows=200):
+    """Picks of (a sample of) windows with the hypocentral distance to their own earthquake
+    (event picks) or to the nearest earthquake (noise), and the travel time of event picks."""
+    p = _pick_distances(split, max_windows)
+    p["scale"] = p.network.map(network_scales(split))
+    return p[["window_id", "network", "scale", "station", "phase", "pick_type", "prob", "event_idx",
+              "distance_km", "travel_time_s"]]
+
+
+def _pick_distances(split, max_windows):
+    ids = _sample_ids(split, max_windows)
+    net = split.windows.set_index("window_id").network
+    p = split.picks[split.picks.window_id.isin(ids)].copy()
+    p["network"] = p.window_id.map(net)
+    xy = _coords(split).reindex(pd.MultiIndex.from_arrays([p.network, p.station])).to_numpy()
+    ev = split.events[split.events.window_id.isin(ids)]
+    out = np.full(len(p), np.nan)
+    tt = np.full(len(p), np.nan)
+    for wid, rows in p.groupby("window_id").indices.items():
+        e = ev[ev.window_id == wid]
+        d = np.hypot(haversine_km(e.latitude.to_numpy()[:, None], e.longitude.to_numpy()[:, None],
+                                  xy[rows, 0][None, :], xy[rows, 1][None, :]),
+                     (e.depth.to_numpy() / 1e3)[:, None])                        # (E, n)
+        own = p.event_idx.to_numpy()[rows]
+        pos = {k: i for i, k in enumerate(e.event_idx)}
+        idx = np.array([pos.get(k, -1) for k in own])
+        out[rows] = np.where(idx >= 0, d[np.maximum(idx, 0), np.arange(len(rows))], d.min(axis=0))
+        t0 = e.set_index("event_idx").t0_rel_s.reindex(own).to_numpy()      # NaN for noise (-1)
+        tt[rows] = p.t_rel_s.to_numpy()[rows] - t0
+    p["distance_km"], p["travel_time_s"] = out, tt
+    return p
+
+
+def network_scales(split, max_windows=50):
+    """Scale of each network from the 90th percentile of the hypocentral distance of its real picks:
+    local (< 150 km), regional (150-1500 km) or global (> 1500 km)."""
+    p = _pick_distances(split, max_windows)
+    p90 = p[p.pick_type == "real"].groupby("network").distance_km.quantile(0.9)
+    return p90.map(lambda v: SCALES[int(np.searchsorted(SCALE_EDGES, v))])
+
+
+def source_events(split):
+    """One row per source earthquake: catalog location, depth, magnitude, farthest station,
+    stations and gap in its first window, and the number of windows it appears in."""
+    ev = split.events.merge(split.windows[["window_id", "network"]], on="window_id")
+    uses = ev.groupby("src_event_id").size().rename("n_windows")
+    first = ev.sort_values(["src_event_id", "is_edge"]).drop_duplicates("src_event_id")
+    out = first.set_index("src_event_id").join(uses).reset_index()
+    out["depth_km"] = out.depth / 1e3
+    out["scale"] = out.network.map(network_scales(split))
+    return out[["src_event_id", "network", "scale", "latitude", "longitude", "depth_km", "magnitude",
+                "magnitude_type", "last_station_km", "n_stations", "gap_deg", "n_windows"]]
+
+
+def event_pairs(split):
+    """Every pair of earthquakes in the same window: epicentral distance and origin-time difference.
+
+    Pairs close in both are the hardest to associate: their picks overlap in time at the same
+    stations."""
+    ev = split.events.merge(split.windows[["window_id", "network", "scenario"]], on="window_id")
+    rows = []
+    for wid, g in ev.groupby("window_id"):
+        if len(g) < 2:
+            continue
+        la, lo, t = g.latitude.to_numpy(), g.longitude.to_numpy(), g.t0_rel_s.to_numpy()
+        i, j = np.triu_indices(len(g), 1)
+        rows.append(pd.DataFrame({"window_id": wid, "network": g.network.iloc[0], "scenario": g.scenario.iloc[0],
+                                  "distance_km": haversine_km(la[i], lo[i], la[j], lo[j]),
+                                  "dt_s": np.abs(t[i] - t[j])}))
+    out = pd.concat(rows, ignore_index=True) if rows else pd.DataFrame(
+        columns=["window_id", "network", "scenario", "distance_km", "dt_s"])
+    out["scale"] = out.network.map(network_scales(split))
+    return out
+
+
+def window_geometry(split, max_windows=300):
+    """Per window: number of stations, median nearest-neighbour station spacing, aperture
+    (largest inter-station distance), and median distance from the earthquakes to the stations
+    that picked them."""
+    ids = _sample_ids(split, max_windows)
+    w = split.windows[split.windows.window_id.isin(ids)]
+    ws = split.window_stations[split.window_stations.window_id.isin(ids)]
+    real = split.picks[split.picks.window_id.isin(ids) & (split.picks.real == 1)]
+    from scipy.spatial import cKDTree
+
+    n_pick = real.groupby("window_id").station.nunique()
+    rows = []
+    for wid, g in ws.groupby("window_id"):
+        la, lo = np.radians(g.latitude.to_numpy()), np.radians(g.longitude.to_numpy())
+        xyz = np.c_[np.cos(la) * np.cos(lo), np.cos(la) * np.sin(lo), np.sin(la)]      # unit sphere
+        chord, _ = cKDTree(xyz).query(xyz, k=2)
+        nn_km = 2 * 6371.0 * np.arcsin(np.clip(chord[:, 1] / 2, 0, 1))
+        # aperture: farthest station from the farthest station of an arbitrary one (two passes)
+        a = np.argmax(haversine_km(g.latitude.iloc[0], g.longitude.iloc[0], g.latitude, g.longitude))
+        aperture = float(np.max(haversine_km(g.latitude.iloc[a], g.longitude.iloc[a], g.latitude, g.longitude)))
+        rows.append({"window_id": wid, "n_stations": len(g), "spacing_km": float(np.median(nn_km)),
+                     "aperture_km": aperture, "n_picking_stations": int(n_pick.get(wid, 0))})
+    out = pd.DataFrame(rows).merge(w[["window_id", "network", "scenario", "max_hyp_dist_km"]], on="window_id")
+    out["scale"] = out.network.map(network_scales(split))
+    return out

bepad/torch.py

@@ -0,0 +1,159 @@
+"""PyTorch datasets for training associators on BEPAD.
+
+- ``ReleaseWindowDataset``: the fixed windows of a release split (train/val).
+- ``GeneratedWindowDataset``: fresh windows from the pool on every epoch
+  (same generator as the release, different seeds), for unlimited training data.
+
+Each item is a dict of tensors for ONE window (variable number of picks and
+stations). ``collate_windows`` concatenates a batch PyG-style, with
+``pick_batch``/``station_batch`` vectors and per-window offsets applied to
+``pick_station``.
+
+Item keys::
+
+    window_id        str
+    station_coords   (S, 3) float32  latitude, longitude [deg], elevation [m]
+    pick_station     (N,)   int64    index into station_coords
+    pick_phase       (N,)   int64    0 = P, 1 = S
+    pick_time        (N,)   float32  seconds from window start
+    pick_label       (N,)   int64    event index in the window, -1 = noise
+    pick_prob        (N,)   float32  1.0-0.5 real, >= 0.5 generated, 0 noise (soft target)
+    pick_real        (N,)   bool     pick comes from UTDQuake (not generated, not noise)
+    pick_type        (N,)   int64    0 = real, 1 = generated, 2 = noise
+    pick_sigma       (N,)   float32  timing uncertainty [s] at that distance (NaN for noise)
+    pick_scorable    (N,)   bool     pick belongs to a scorable event
+    event_origin     (E, 4) float32  t0 [s from window start], latitude, longitude, depth [km]
+    event_scorable   (E,)   bool
+
+Normalisation is left to the model.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+import torch
+from torch.utils.data import Dataset
+
+from .build.geo import stable_seed
+
+PHASE_CODE = {"P": 0, "S": 1}
+TYPE_CODE = {"real": 0, "generated": 1, "noise": 2}
+
+
+def window_to_tensors(window_id, picks, stations, events=None):
+    stations = stations.reset_index(drop=True)
+    pos = {s: i for i, s in enumerate(stations["station"])}
+    item = {
+        "window_id": window_id,
+        "station_coords": torch.tensor(stations[["latitude", "longitude", "elevation"]].to_numpy(np.float32)),
+        "pick_station": torch.tensor(picks["station"].map(pos).to_numpy(np.int64)),
+        "pick_phase": torch.tensor(picks["phase"].map(PHASE_CODE).to_numpy(np.int64)),
+        "pick_time": torch.tensor(picks["t_rel_s"].to_numpy(np.float32)),
+    }
+    if "event_idx" in picks:
+        item["pick_label"] = torch.tensor(picks["event_idx"].to_numpy(np.int64))
+        item["pick_prob"] = torch.tensor(picks["prob"].to_numpy(np.float32))
+        item["pick_real"] = torch.tensor(picks["real"].to_numpy(bool))
+        item["pick_type"] = torch.tensor(picks["pick_type"].map(TYPE_CODE).to_numpy(np.int64))
+        item["pick_sigma"] = torch.tensor(picks["sigma_s"].to_numpy(np.float32))
+        item["pick_scorable"] = torch.tensor(picks["in_scorable_event"].to_numpy(bool))
+    if events is not None:
+        e = events.sort_values("event_idx")
+        item["event_origin"] = torch.tensor(np.stack([
+            e["t0_rel_s"], e["latitude"], e["longitude"], e["depth"] / 1e3], axis=1).astype(np.float32))
+        item["event_scorable"] = torch.tensor(e["scorable"].to_numpy(bool))
+    return item
+
+
+class ReleaseWindowDataset(Dataset):
+    """Fixed windows of one split (use train/val; never tune on test).
+
+    ``release``: a :class:`bepad.data.Release`, a local release folder, or None
+    to download ``split`` from Hugging Face.
+    """
+
+    def __init__(self, release=None, split="train", networks=None):
+        from .data import Release, load
+
+        if release is None:
+            self.split = load(split, networks=networks, labels=True)
+        else:
+            rel = release if isinstance(release, Release) else Release(release)
+            self.split = rel.load(split, networks=networks, labels=True)
+
+    def __len__(self):
+        return len(self.split)
+
+    def __getitem__(self, i):
+        w = self.split[i]
+        return window_to_tensors(w.window_id, w.picks, w.stations, w.events)
+
+
+class GeneratedWindowDataset(Dataset):
+    """On-the-fly windows from the pool.
+
+    Item ``i`` of epoch ``e`` uses seed sha256(base_seed, e, i), so a run is
+    reproducible and every epoch sees new windows. Scenario configs (dicts of
+    window parameters, e.g. from ``bepad.materialize.expand_scenarios``) are
+    drawn uniformly per item, as is the network.
+    """
+
+    def __init__(self, data_root, networks, scenarios, split="train", length=10_000, base_seed=0):
+        from .build.config import WINDOW_DEFAULTS
+        from .build.labels import LabelModel
+        from .build.pool import load_network, pool_dir
+        from .build.windows import NetworkIndex
+
+        model = LabelModel.load(pool_dir(data_root) / "label_model.json")
+        self.index = {n: NetworkIndex(*load_network(data_root, n), network=n, label_model=model) for n in networks}
+        self.networks = list(networks)
+        self.scenarios = [(k, {**WINDOW_DEFAULTS, **v}) for k, v in
+                          (scenarios.items() if isinstance(scenarios, dict) else enumerate(scenarios))]
+        self.split = split
+        self.length = length
+        self.base_seed = base_seed
+        self.epoch = 0
+
+    def set_epoch(self, epoch):
+        self.epoch = epoch
+
+    def __len__(self):
+        return self.length
+
+    def __getitem__(self, i):
+        from .build.windows import generate_window
+
+        seed = stable_seed(self.base_seed, self.epoch, i)
+        rng = np.random.default_rng(seed)
+        net = self.networks[rng.integers(len(self.networks))]
+        name, cfg = self.scenarios[rng.integers(len(self.scenarios))]
+        for attempt in range(10):
+            w = generate_window(self.index[net], self.split, cfg, seed + attempt, f"gen-{self.epoch}-{i}", str(name))
+            if w is not None:
+                break
+        else:
+            raise RuntimeError(f"could not generate a window for network {net}, scenario {name}")
+        st = self.index[net].stations.loc[w.stations["station"]].reset_index()
+        return window_to_tensors(w.window["window_id"], w.picks, st, w.events)
+
+
+def collate_windows(batch):
+    """Concatenate variable-size windows; adds pick_batch / station_batch / event_batch."""
+    out = {"window_id": [b["window_id"] for b in batch]}
+    offset = 0
+    parts = {k: [] for k in batch[0] if k != "window_id"}
+    parts.update(pick_batch=[], station_batch=[], event_batch=[])
+    for k, b in enumerate(batch):
+        n_sta = len(b["station_coords"])
+        for key in parts:
+            if key in b:
+                parts[key].append(b[key] + offset if key == "pick_station" else b[key])
+        parts["pick_batch"].append(torch.full((len(b["pick_time"]),), k, dtype=torch.long))
+        parts["station_batch"].append(torch.full((n_sta,), k, dtype=torch.long))
+        if "event_origin" in b:
+            parts["event_batch"].append(torch.full((len(b["event_origin"]),), k, dtype=torch.long))
+        offset += n_sta
+    for key, vals in parts.items():
+        if vals:
+            out[key] = torch.cat(vals)
+    return out

bepad-0.1.0.dist-info/METADATA

@@ -0,0 +1,85 @@
+Metadata-Version: 2.4
+Name: bepad
+Version: 0.1.0
+Summary: BEPAD: a benchmark dataset for seismic phase association built from UTDQuake
+License-Expression: GPL-3.0-or-later
+Project-URL: Repository, https://github.com/ecastillot/BEPAD
+Project-URL: Dataset, https://huggingface.co/datasets/ecastillot/BEPAD
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: numpy
+Requires-Dist: pandas>=2.0
+Requires-Dist: pyarrow
+Requires-Dist: scipy
+Requires-Dist: pyyaml
+Requires-Dist: matplotlib
+Requires-Dist: huggingface_hub
+Provides-Extra: dl
+Requires-Dist: torch>=2.1; extra == "dl"
+Provides-Extra: baselines
+Requires-Dist: pyocto; extra == "baselines"
+Provides-Extra: build
+Requires-Dist: utdquake>=1.0.1; extra == "build"
+Provides-Extra: docs
+Requires-Dist: sphinx>=7; extra == "docs"
+Requires-Dist: sphinx-rtd-theme; extra == "docs"
+Requires-Dist: myst-parser; extra == "docs"
+Requires-Dist: sphinx-copybutton; extra == "docs"
+Dynamic: license-file
+
+# BEPAD — a benchmark dataset for seismic phase association
+
+A phase associator groups the P and S picks of a seismic network into earthquakes. BEPAD gives
+thousands of association problems built from **real earthquakes and real picks**
+([UTDQuake](https://huggingface.co/datasets/ecastillot/UTDQuake)), with the answer for every pick.
+
+![a BEPAD window](https://raw.githubusercontent.com/ecastillot/BEPAD/main/docs/figures/window_anatomy_tx.png)
+
+*One window (Texas network): (a) earthquakes and stations, (b) every pick at its distance from
+its earthquake, (c) picks by type.*
+
+```python
+import bepad
+
+test = bepad.load("test", labels=False)          # download the test split from Hugging Face
+for w in test:                                   # one window = one association problem
+    w.picks, w.stations                          # picks (station, P/S, time) and station coordinates
+report = bepad.evaluate(predictions, test)       # score your associator
+```
+
+Share your associator and get on the leaderboard in four commands:
+
+```bash
+bepad new my-method                                            # template folder
+bepad run associators/my-method --split test --out results/my-method
+bepad submit results/my-method --code associators/my-method    # opens a Hugging Face pull request
+```
+
+| | |
+|---|---|
+| **What is in a window** | several overlapping earthquakes on one network, with **real**, **generated** (plausible extra arrivals) and **noise** picks; every pick labelled with its earthquake and a confidence `prob` (≥ 0.5 arrival, 0 noise) |
+| **Difficulty** | 23 scenarios: 1–20 earthquakes per window × 5 noise levels, plus earthquakes outside the network |
+| **Size** | per network and scenario: 1,000 train · 200 val · 100 test windows (23 scenarios) |
+| **Associator authors** | fill in `bepad.Associator`, run on `test`, submit to the leaderboard — see *Contribute your associator* |
+| **Deep learning** | `bepad.torch.ReleaseWindowDataset` — see *Train a model* |
+
+Code: https://github.com/ecastillot/BEPAD · Data: https://huggingface.co/datasets/ecastillot/BEPAD · Documentation: `docs/` (Read the Docs). Install: `pip install bepad` (or `pip install -e .` from
+this folder; `.[dl]` for PyTorch, `.[build]` to rebuild the dataset).
+
+## Repository
+
+```
+bepad/            load, evaluate, Associator, leaderboard, torch, plots  (what users need)
+bepad/build/      pool, labels, windows, release (how the dataset is built)
+configs/          pool.yaml (quality control, split, labels), v0.1.yaml, dev.yaml (scenarios, sizes)
+docs/             Read the Docs sources; docs/gallery/ = the code of every figure
+associators/      contributed associators, one folder each (pyocto = worked example)
+pipeline/         the recipe of a release (SLURM): pool -> release -> baselines (pipeline/README.md)
+hf/README.md      Hugging Face dataset card
+```
+
+## License and citation
+
+Code: GPL-3.0. Data: the terms of the original providers, as in UTDQuake. Please cite UTDQuake and
+the network operators of the data you use.

bepad-0.1.0.dist-info/RECORD

@@ -0,0 +1,23 @@
+bepad/__init__.py,sha256=whHazCC_dwmQI0QLibRwsC1t1Ab-IZaG6wHCOui6h00,829
+bepad/__main__.py,sha256=aNVf11co3xU3vDlQu7bhwOiRr7j9aKcC6J3PDW8GkMo,9414
+bepad/associator.py,sha256=xw0cQcU8PyRQuPptRAgqk4P9-l4lIBPZjm4Ke5_3T3w,14902
+bepad/data.py,sha256=Pm3ZmwrnxMtIScYyC5HQqENtghHttSsk0bXUHUQ4ge4,8178
+bepad/leaderboard.py,sha256=o_VUDWEsqddIZJrGj6NMNzZ-8mxh31qnO47d6B8vx2k,18525
+bepad/metrics.py,sha256=t0_cEt7cSncw3XjEHff-sLVUpqsrD4QqYY-Ju1dz15Q,15272
+bepad/plots.py,sha256=5J_ahnncE7tLMk0ebL1ldbQNhkGBVzffr1Gedy7MbVY,50491
+bepad/stats.py,sha256=O5NEGcznxb831AYNJgcSKFEZODKQAYv6tUmDy7UuY3c,7223
+bepad/torch.py,sha256=l1lHSYLU0rn5smLeDb3M7A_zw4xgPpbDNINDPpXWcJY,6989
+bepad/build/__init__.py,sha256=bs2kYx4znsNu4mI6yc5fg1hV17sOD6os4pq2M-QTeM0,513
+bepad/build/config.py,sha256=mmG04H-7MFQIVyZ0qyg4GabY_gXO6m8r7ZRu-DdMXv0,3891
+bepad/build/geo.py,sha256=ucen45viVT8u_qxw1R4cjCIN4qSCGzzobVyx0UbiBZI,2130
+bepad/build/labels.py,sha256=7d6C5R1-efcP1RoTtWHvZZvwlj05HECApZWBeJ6KkNg,15812
+bepad/build/pool.py,sha256=Ugz1md1bmndQZNuAEVWrXXo4eUZSscxX8_4wFZBeAlg,16731
+bepad/build/release.py,sha256=2oEtph0Jy9ljsscEVoN_slkEVg6dumQMvc-cDvBUeZw,11282
+bepad/build/utdquake.py,sha256=1PGuJYP1eqW0PRnqDFVSCR5kcGCi8CbS-mbkFGciFI4,2957
+bepad/build/windows.py,sha256=1Ml4gMn6-mIT48cnj8mIaMebkY2mLXUQrIcLdRaN0Og,21640
+bepad-0.1.0.dist-info/licenses/LICENSE,sha256=OXLcl0T2SZ8Pmy2_dmlvKuetivmyPd5m1q-Gyd-zaYY,35149
+bepad-0.1.0.dist-info/METADATA,sha256=mnEFOq40J3TlQlhITzQxy7BC29f1CJ9xWXqh34CUrEQ,4031
+bepad-0.1.0.dist-info/WHEEL,sha256=mOdXR9C3wMVov7cSY8PY-XCzo3SmmWCHvGQr5csJoao,91
+bepad-0.1.0.dist-info/entry_points.txt,sha256=IAY9QLOeY3bfXwKOylC9rJJ9usZHOu1uW-wyQvQKVLU,46
+bepad-0.1.0.dist-info/top_level.txt,sha256=ohUvSDWJNaGyN7oCBdk_OZOdF2qgo0AtV33vCRHWvqI,6
+bepad-0.1.0.dist-info/RECORD,,

bepad-0.1.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (84.2.1)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

bepad-0.1.0.dist-info/entry_points.txt

@@ -0,0 +1,2 @@
+[console_scripts]
+bepad = bepad.__main__:main