bepad 0.1.0__tar.gz → 0.1.2__tar.gz

{bepad-0.1.0 → bepad-0.1.2}/PKG-INFO

@@ -1,8 +1,9 @@
 Metadata-Version: 2.4
 Name: bepad
-Version: 0.1.0
+Version: 0.1.2
 Summary: BEPAD: a benchmark dataset for seismic phase association built from UTDQuake
 License-Expression: GPL-3.0-or-later
+Project-URL: Documentation, https://bepad.readthedocs.io
 Project-URL: Repository, https://github.com/ecastillot/BEPAD
 Project-URL: Dataset, https://huggingface.co/datasets/ecastillot/BEPAD
 Requires-Python: >=3.10
@@ -30,6 +31,8 @@ Dynamic: license-file
 
 # BEPAD — a benchmark dataset for seismic phase association
 
+[![PyPI](https://img.shields.io/pypi/v/bepad?color=orange&label=PyPI&style=for-the-badge&logo=pypi&logoColor=white)](https://pypi.org/project/bepad/) [![Data v0.1](https://img.shields.io/badge/Data-v0.1-yellow?style=for-the-badge)](https://huggingface.co/datasets/ecastillot/BEPAD/tree/main/v0.1) [![Hugging Face Dataset](https://img.shields.io/badge/HuggingFace-Dataset-yellow?style=for-the-badge&logo=huggingface&logoColor=black)](https://huggingface.co/datasets/ecastillot/BEPAD) [![Read the Docs](https://img.shields.io/badge/Docs-Read%20the%20Docs-blue?style=for-the-badge&logo=read-the-docs)](https://bepad.readthedocs.io/en/latest/) [![GitHub](https://img.shields.io/badge/GitHub-BEPAD-black?style=for-the-badge&logo=github)](https://github.com/ecastillot/BEPAD)
+
 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.
@@ -37,12 +40,12 @@ thousands of association problems built from **real earthquakes and real picks**
 ![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.*
+its earthquake, (c) picks by type. The associator's job is to find which picks belong together.*
 
 ```python
 import bepad
 
-test = bepad.load("test", labels=False)          # download the test split from Hugging Face
+test = bepad.load("test", networks=["tx"], max_windows=20, labels=False)   # 20 Texas test windows
 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
@@ -64,8 +67,8 @@ bepad submit results/my-method --code associators/my-method    # opens a Hugging
 | **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).
+Documentation: https://bepad.readthedocs.io · Data: https://huggingface.co/datasets/ecastillot/BEPAD · Code: https://github.com/ecastillot/BEPAD.
+Install: `pip install bepad` (or `pip install -e .` from this folder; `.[dl]` for PyTorch, `.[build]` to rebuild the dataset).
 
 ## Repository
 

{bepad-0.1.0 → bepad-0.1.2}/README.md

@@ -1,5 +1,7 @@
 # BEPAD — a benchmark dataset for seismic phase association
 
+[![PyPI](https://img.shields.io/pypi/v/bepad?color=orange&label=PyPI&style=for-the-badge&logo=pypi&logoColor=white)](https://pypi.org/project/bepad/) [![Data v0.1](https://img.shields.io/badge/Data-v0.1-yellow?style=for-the-badge)](https://huggingface.co/datasets/ecastillot/BEPAD/tree/main/v0.1) [![Hugging Face Dataset](https://img.shields.io/badge/HuggingFace-Dataset-yellow?style=for-the-badge&logo=huggingface&logoColor=black)](https://huggingface.co/datasets/ecastillot/BEPAD) [![Read the Docs](https://img.shields.io/badge/Docs-Read%20the%20Docs-blue?style=for-the-badge&logo=read-the-docs)](https://bepad.readthedocs.io/en/latest/) [![GitHub](https://img.shields.io/badge/GitHub-BEPAD-black?style=for-the-badge&logo=github)](https://github.com/ecastillot/BEPAD)
+
 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.
@@ -7,12 +9,12 @@ thousands of association problems built from **real earthquakes and real picks**
 ![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.*
+its earthquake, (c) picks by type. The associator's job is to find which picks belong together.*
 
 ```python
 import bepad
 
-test = bepad.load("test", labels=False)          # download the test split from Hugging Face
+test = bepad.load("test", networks=["tx"], max_windows=20, labels=False)   # 20 Texas test windows
 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
@@ -34,8 +36,8 @@ bepad submit results/my-method --code associators/my-method    # opens a Hugging
 | **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).
+Documentation: https://bepad.readthedocs.io · Data: https://huggingface.co/datasets/ecastillot/BEPAD · Code: https://github.com/ecastillot/BEPAD.
+Install: `pip install bepad` (or `pip install -e .` from this folder; `.[dl]` for PyTorch, `.[build]` to rebuild the dataset).
 
 ## Repository
 

{bepad-0.1.0 → bepad-0.1.2}/bepad/__init__.py

@@ -6,7 +6,7 @@
     class MyMethod(bepad.Associator): ...           # template for reusable associators
 """
 
-__version__ = "0.1.0"
+__version__ = "0.1.2"
 
 from .data import CORE_NAME, CORE_NETWORKS, LABEL_COLUMNS, Release, Split, download, load, open, predictions_frame  # noqa: E402,A004
 from .metrics import evaluate  # noqa: E402

{bepad-0.1.0 → bepad-0.1.2}/bepad/associator.py

@@ -90,26 +90,23 @@ class Associator:
         raise NotImplementedError
 
     # --- provided ------------------------------------------------------------------
-    def run(self, split="test", networks=None, scenarios=None, max_windows=None, workers=1, release=None,
-            root=None, progress=True):
+    def run(self, split="test", networks=None, scenarios=None, max_windows=None, window_ids=None, workers=1,
+            release=None, root=None, progress=True):
         """Run on every window of a split and score it. Returns a :class:`Result`.
 
         ``workers``: number of CPU processes (windows are independent).
-        ``scenarios`` / ``max_windows`` run a subset (quick tests); such runs are
-        marked partial and are not ranked on the leaderboard.
+        ``scenarios`` / ``max_windows`` / ``window_ids`` run a subset (quick tests, as in
+        :func:`bepad.load`); such runs are marked partial and are not ranked on the leaderboard.
         """
         from .data import RELEASE, load
         from .metrics import evaluate
 
-        data = load(split, networks=networks, labels=False, release=release or RELEASE, root=root)
-        w = data.windows
-        keep = w
-        if scenarios is not None:
-            keep = keep[keep.scenario.isin(scenarios)]
-        if max_windows is not None:
-            keep = keep.groupby("network", group_keys=False).head(max_windows)
+        subset = scenarios is not None or max_windows is not None or window_ids is not None
+        data = load(split, networks=networks, labels=False, release=release or RELEASE, root=root,
+                    scenarios=scenarios, max_windows=max_windows, window_ids=window_ids)
+        keep = data.windows
         ids = set(keep.window_id)
-        todo = [(win.window_id, win.meta["network"], win.picks, win.stations) for win in data if win.window_id in ids]
+        todo = [(win.window_id, win.meta["network"], win.picks, win.stations) for win in data]
         rows, locs, secs, t0, n = [], [], [], time.time(), len(todo)
         _RUNNING[0] = self
         if workers > 1:
@@ -137,8 +134,8 @@ class Associator:
         meta = {
             "split": split, "release": data.release.path.name,
             "networks": sorted(keep.network.unique()),
-            "n_windows": n, "complete": n == len(w),
-            "scenarios": "all" if scenarios is None else list(scenarios),
+            "n_windows": n, "complete": not subset,
+            "scenarios": "all" if scenarios is None else ([scenarios] if isinstance(scenarios, str) else list(scenarios)),
             "workers": workers, "runtime_s": round(runtime, 2),
             "seconds_per_window": round(float(np.mean(secs)), 4) if secs else None,   # time inside associate()
             "hardware": f"{platform.processor() or platform.machine()} ({platform.system()})",

bepad-0.1.2/bepad/blocks.py

@@ -0,0 +1,170 @@
+"""Release files stored in blocks, so one window can be read without reading the whole file.
+
+The picks, events and window_stations tables of a release are sorted by ``window_id`` (keeping the
+order of the rows inside each window) and written in row groups ("blocks") of about
+:data:`BLOCK_ROWS` rows (per table), compressed with zstd. A block holds whole windows only. Parquet keeps the first and last
+``window_id`` of every block, so finding a window needs no separate index. The windows table keeps
+its own order (it is small and read whole).
+
+Files in the older layout (not sorted) still work: they are read whole.
+"""
+
+from __future__ import annotations
+
+import bisect
+import os
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+
+# Rows per block. Smaller blocks read one window faster but compress less; window_stations has two
+# short columns, so its blocks are larger (measured on tx and R07: about the old file sizes).
+BLOCK_ROWS = {"picks": 5_000, "events": 5_000, "window_stations": 20_000}
+COMPRESSION = "zstd"
+BLOCKED_TABLES = ("picks", "events", "window_stations")
+LAYOUT = "blocks-v1"          # recorded in manifest.json
+
+
+def write_blocked(df, path, block_rows=5_000):
+    """Write ``df`` sorted by window_id, in blocks of about ``block_rows`` rows of whole windows."""
+    import pyarrow as pa
+    import pyarrow.parquet as pq
+
+    path = Path(path)
+    df = df.sort_values("window_id", kind="stable", ignore_index=True)
+    ids = df["window_id"].to_numpy()
+    starts = np.flatnonzero(np.r_[True, ids[1:] != ids[:-1]]) if len(ids) else np.array([], int)
+    block = starts // block_rows                                   # block of each window, by its first row
+    bounds = list(starts[np.r_[True, block[1:] != block[:-1]]]) + [len(df)] if len(ids) else [0, 0]
+    schema = pa.Schema.from_pandas(df, preserve_index=False)
+    tmp = path.with_suffix(".tmp")
+    with pq.ParquetWriter(tmp, schema, compression=COMPRESSION) as w:
+        for a, b in zip(bounds[:-1], bounds[1:]):
+            w.write_table(pa.Table.from_pandas(df.iloc[a:b], schema=schema, preserve_index=False),
+                          row_group_size=max(b - a, 1))
+    os.replace(tmp, path)
+
+
+def _bounds(md):
+    """(first, last) window_id of every row group, or None without statistics."""
+    col = md.schema.to_arrow_schema().get_field_index("window_id")
+    out = []
+    for i in range(md.num_row_groups):
+        st = md.row_group(i).column(col).statistics
+        if st is None or not st.has_min_max:
+            return None
+        out.append((st.min, st.max))
+    return out
+
+
+def block_bounds(pf):
+    """(first, last) window_id per block if the file is in blocks (sorted, no window across two
+    blocks), else None."""
+    b = _bounds(pf.metadata)
+    if b is None or any(b[i][1] >= b[i + 1][0] for i in range(len(b) - 1)):
+        return None
+    return b
+
+
+def check_blocked(path):
+    """Problems of one file (empty list = in blocks): sorted rows, whole windows per block."""
+    import pyarrow.parquet as pq
+
+    pf = pq.ParquetFile(path)
+    if pf.metadata.num_rows == 0:
+        return []
+    problems = []
+    ids = pf.read(columns=["window_id"]).column("window_id").to_numpy(zero_copy_only=False)
+    if not (ids[1:] >= ids[:-1]).all():
+        problems.append(f"{path}: rows are not sorted by window_id")
+    if block_bounds(pf) is None:
+        problems.append(f"{path}: a window is split across two blocks (or no statistics)")
+    return problems
+
+
+def read_rows(path, window_ids, columns=None, batch_size=200_000):
+    """Rows of one file whose window is in ``window_ids``: only the blocks that hold them are read
+    (older files: read in batches, keeping the matching rows)."""
+    import pyarrow as pa
+    import pyarrow.compute as pc
+    import pyarrow.parquet as pq
+
+    pf = pq.ParquetFile(path, pre_buffer=False, buffer_size=1 << 20)
+    schema = pa.schema([pf.schema_arrow.field(c) for c in columns]) if columns else pf.schema_arrow
+    ids = pa.array(sorted(window_ids), type=schema.field("window_id").type)
+    b = block_bounds(pf)
+    if b is not None:
+        lo = [x[0] for x in b]
+        groups = sorted({g for w in ids.to_pylist() if (g := bisect.bisect_right(lo, w) - 1) >= 0 and w <= b[g][1]})
+        batches = pf.iter_batches(batch_size=batch_size, row_groups=groups, columns=columns) if groups else []
+    else:
+        batches = pf.iter_batches(batch_size=batch_size, columns=columns)
+    parts = [x.filter(pc.is_in(x.column("window_id"), value_set=ids)) for x in batches]
+    return pa.Table.from_batches(parts, schema=schema).to_pandas()
+
+
+class BlockReader:
+    """One file, read one window at a time. Keeps the last block in memory, so windows read in
+    order cost one block read per block. Safe in forked DataLoader workers (reopens per process).
+    An older file (not in blocks) is read whole once."""
+
+    def __init__(self, path, columns=None):
+        self.path, self.columns = Path(path), columns
+        self._pid = None
+
+    def __getstate__(self):                     # pickled DataLoader workers reopen the file
+        return {"path": self.path, "columns": self.columns, "_pid": None}
+
+    def _open(self):
+        import pyarrow.parquet as pq
+
+        self._pf = pq.ParquetFile(self.path)
+        b = block_bounds(self._pf)
+        self._lo = None if b is None else [x[0] for x in b]
+        self._hi = None if b is None else [x[1] for x in b]
+        self._key, self._df, self._rows = None, None, {}
+        self._pid = os.getpid()
+
+    def _load(self, key):
+        if key != self._key:
+            t = self._pf.read(columns=self.columns) if key == "all" else \
+                self._pf.read_row_group(key, columns=self.columns)
+            self._df = t.to_pandas()
+            self._rows = self._df.groupby("window_id", sort=False).indices
+            self._key = key
+
+    def window(self, wid):
+        """The rows of window ``wid`` (an empty frame if it has none)."""
+        if self._pid != os.getpid():
+            self._open()
+        if self._lo is None:
+            self._load("all")
+        else:
+            g = bisect.bisect_right(self._lo, wid) - 1
+            if g < 0 or wid > self._hi[g]:
+                return self._empty()
+            self._load(g)
+        rows = self._rows.get(wid)
+        return self._df.iloc[rows].reset_index(drop=True) if rows is not None else self._empty()
+
+    def _empty(self):
+        if self._df is None:
+            self._load(0 if self._lo else "all")
+        return self._df.iloc[:0].reset_index(drop=True)
+
+
+def reblock(release_dir, networks=None, splits=None, block_rows=None):
+    """Rewrite the blocked tables of an existing release in blocks (same rows, same content)."""
+    release_dir = Path(release_dir)
+    done = []
+    for split_dir in sorted(p for p in release_dir.iterdir() if (p / "windows").is_dir()):
+        if splits is not None and split_dir.name not in splits:
+            continue
+        for table in BLOCKED_TABLES:
+            for f in sorted((split_dir / table).glob("network=*.parquet")):
+                if networks is not None and f.stem.split("=", 1)[1] not in networks:
+                    continue
+                write_blocked(pd.read_parquet(f), f, block_rows or BLOCK_ROWS[table])
+                done.append(f)
+    return done

{bepad-0.1.0 → bepad-0.1.2}/bepad/build/release.py

@@ -8,6 +8,7 @@ Hugging Face dataset configs/splits:
     <split>/picks/network=<name>.parquet            one row per pick (label: event_idx, -1 noise)
     <split>/events/network=<name>.parquet           one row per event in a window
     <split>/window_stations/network=<name>.parquet  stations available in each window
+    (picks, events and window_stations are sorted by window_id and stored in blocks: bepad.blocks)
     _parts/<name>.json                              per-network build record
     --- written by finalize_release ---
     stations.parquet                                station coordinates
@@ -38,6 +39,7 @@ from pathlib import Path
 import pandas as pd
 
 from .. import __version__
+from ..blocks import BLOCK_ROWS, BLOCKED_TABLES, COMPRESSION, LAYOUT, check_blocked, write_blocked
 from .config import REPO_ROOT, WINDOW_DEFAULTS
 from .geo import stable_seed
 from .labels import LabelModel
@@ -138,7 +140,10 @@ def _build_network(network, cfg, scenarios, workers):
         for t, df in frames.items():
             d = out / split / t
             d.mkdir(parents=True, exist_ok=True)
-            df.to_parquet(d / f"network={network}.parquet", index=False)
+            if t in BLOCKED_TABLES:                    # one window can be read without the whole file
+                write_blocked(df, d / f"network={network}.parquet", BLOCK_ROWS[t])
+            else:
+                df.to_parquet(d / f"network={network}.parquet", index=False)
 
         w, e = frames["windows"], frames["events"]
         uses = e["src_event_id"].value_counts()
@@ -193,6 +198,10 @@ def finalize_release(cfg):
     networks = sorted(parts)
     if not networks:
         raise RuntimeError(f"nothing to finalize: no built networks in {out / '_parts'}")
+    problems = [q for t in BLOCKED_TABLES for f in sorted(out.glob(f"*/{t}/network=*.parquet"))
+                for q in check_blocked(f)]
+    if problems:
+        raise RuntimeError("files not in blocks (rebuild them, or bepad.blocks.reblock):\n  " + "\n  ".join(problems))
 
     pd.concat([load_network(cfg["data_root"], n)[2] for n in networks], ignore_index=True) \
         .to_parquet(out / "stations.parquet", index=False)
@@ -229,6 +238,7 @@ def finalize_release(cfg):
                                         "source_sha256", "n_events", "built")}
                           for n in networks},
         "counts": counts,
+        "layout": {"name": LAYOUT, "block_rows": BLOCK_ROWS, "compression": COMPRESSION},
         "skipped_splits": {n: p["skipped_splits"] for n, p in parts.items() if p["skipped_splits"]},
     }
     with open(out / "manifest.json", "w") as f:

{bepad-0.1.0 → bepad-0.1.2}/bepad/data.py

@@ -16,6 +16,8 @@ from pathlib import Path
 
 import pandas as pd
 
+from .blocks import BlockReader, read_rows
+
 #: Columns of picks.parquet that carry ground truth. Never give them to an associator.
 LABEL_COLUMNS = ("event_idx", "pick_type", "real", "prob", "sigma_s", "in_scorable_event",
                  "src_event_id", "src_pick_id")
@@ -70,6 +72,57 @@ class Split:
                 f"labels={self.labels}>")
 
 
+class LazySplit:
+    """The windows of one split, each read from disk when it is asked for.
+
+    Same use as :class:`Split` (``len``, ``split[i]``, ``split["<window_id>"]``, iteration, the
+    ``windows`` table), but only the window table is in memory, so it fits any split. Reading is
+    fast when the release files are in blocks (:mod:`bepad.blocks`); windows read in order share
+    their block reads.
+    """
+
+    def __init__(self, release, split, windows, labels=True):
+        self.release, self.name, self.labels = release, split, labels
+        self.windows = windows.reset_index(drop=True)
+        self._pos = dict(zip(self.windows["window_id"], range(len(self.windows))))
+        st = pd.read_parquet(release.path / "stations.parquet")
+        self._stations = {n: g.drop(columns="network") for n, g in st.groupby("network")}
+        self._readers = {}
+
+    def _reader(self, network, table):
+        key = (network, table)
+        if key not in self._readers:
+            cols = list(INPUT_COLUMNS) if table == "picks" and not self.labels else None
+            self._readers[key] = BlockReader(self.release.path / self.name / table / f"network={network}.parquet", cols)
+        return self._readers[key]
+
+    def __len__(self):
+        return len(self.windows)
+
+    def __iter__(self):
+        for i in range(len(self.windows)):
+            yield self[i]
+
+    def __getitem__(self, i):
+        if isinstance(i, str):
+            i = self._pos[i]
+        meta = self.windows.iloc[i]
+        wid, net = meta["window_id"], meta["network"]
+        picks = self._reader(net, "picks").window(wid)
+        ws = self._reader(net, "window_stations").window(wid)
+        st = ws.merge(self._stations[net], on="station", how="left")
+        st = st[["station", "latitude", "longitude", "elevation"]].reset_index(drop=True)
+        ev = None
+        if self.labels:
+            ev = self._reader(net, "events").window(wid)
+            ev = ev if len(ev) else None
+        return WindowData(wid, meta, picks, st, ev)
+
+    def __repr__(self):
+        nets = self.windows["network"].nunique() if len(self.windows) else 0
+        return f"<bepad.LazySplit {len(self)} windows, {nets} networks, labels={self.labels}>"
+
+
 class Release:
     """A BEPAD release folder (local copy of the Hugging Face dataset or a local build)."""
 
@@ -97,24 +150,51 @@ class Release:
     def scenarios(self):
         return json.loads((self.path / "scenarios.json").read_text())
 
-    def _read(self, split, table, networks, columns=None):
+    def _read(self, split, table, networks, columns=None, window_ids=None):
         files = sorted((self.path / split / table).glob("network=*.parquet"))
         if networks is not None:
             keep = {f"network={n}.parquet" for n in networks}
             files = [f for f in files if f.name in keep]
         if not files:
             raise FileNotFoundError(f"no {table} files for split={split!r}, networks={networks}")
-        return pd.concat([pd.read_parquet(f, columns=columns) for f in files], ignore_index=True)
-
-    def load(self, split, networks=None, labels=True):
-        """Load one split. ``labels=False`` returns only what an associator may use."""
+        if window_ids is None:
+            return pd.concat([pd.read_parquet(f, columns=columns) for f in files], ignore_index=True)
+        return pd.concat([read_rows(f, window_ids, columns) for f in files], ignore_index=True)
+
+    def table(self, split, name, networks=None, window_ids=None):
+        """One table as a DataFrame: ``"windows"`` (metadata), ``"window_stations"``, ``"events"``,
+        ``"picks"`` (with labels) of ``split``, or ``"stations"`` (coordinates, every split).
+        ``networks`` (``"core"`` = BEPAD-10) and ``window_ids`` limit the rows that are read."""
+        if name == "stations":
+            st = pd.read_parquet(self.path / "stations.parquet")
+            nets = _networks(networks)
+            return st if nets is None else st[st["network"].isin(nets)].reset_index(drop=True)
+        return self._read(split, name, _networks(networks), window_ids=None if window_ids is None else set(window_ids))
+
+    def load(self, split, networks=None, labels=True, scenarios=None, max_windows=None, window_ids=None,
+             lazy=False):
+        """Load one split. ``labels=False`` returns only what an associator may use.
+
+        To save memory, load only some windows: ``scenarios`` (names, e.g. ``"ev-med__noise-high"``),
+        ``max_windows`` (a fixed random sample of at most this many windows per network) or
+        ``window_ids``. Only the rows of those windows are read.
+        ``lazy=True``: keep only the window table in memory and read each window from disk when it
+        is asked for (:class:`LazySplit`; for training on many windows).
+        """
         cols = None if labels else list(INPUT_COLUMNS)
+        windows = self._read(split, "windows", networks)
+        ids = None
+        if scenarios is not None or max_windows is not None or window_ids is not None:
+            windows = _select(windows, scenarios, max_windows, window_ids)
+            ids = set(windows["window_id"])
+        if lazy:
+            return LazySplit(self, split, windows, labels)
         s = Split(
-            windows=self._read(split, "windows", networks),
-            picks=self._read(split, "picks", networks, columns=cols),
-            window_stations=self._read(split, "window_stations", networks),
+            windows=windows,
+            picks=self._read(split, "picks", networks, columns=cols, window_ids=ids),
+            window_stations=self._read(split, "window_stations", networks, window_ids=ids),
             stations=pd.read_parquet(self.path / "stations.parquet"),
-            events=self._read(split, "events", networks) if labels else None,
+            events=self._read(split, "events", networks, window_ids=ids) if labels else None,
             labels=labels,
         )
         s.release, s.name = self, split
@@ -132,6 +212,18 @@ CORE_NAME = "BEPAD-10"
 CORE_NETWORKS = ["R05", "R07", "R01", "nc", "tx", "JMA", "CM", "R09", "ak", "NOA"]
 
 
+def _select(windows, scenarios=None, max_windows=None, window_ids=None):
+    """The windows kept by :meth:`Release.load` (scenario names, a sample per network, or ids)."""
+    if window_ids is not None:
+        windows = windows[windows["window_id"].isin(set(window_ids))]
+    if scenarios is not None:
+        windows = windows[windows["scenario"].isin([scenarios] if isinstance(scenarios, str) else scenarios)]
+    if max_windows is not None:
+        windows = (windows.groupby("network", group_keys=False)
+                   .apply(lambda g: g.sample(min(len(g), max_windows), random_state=0).sort_index()))
+    return windows
+
+
 def _networks(networks):
     """``"core"`` in a list of networks stands for the core networks."""
     if networks is None:
@@ -164,7 +256,8 @@ def open(path):  # noqa: A001 -- mirrors the usual `dataset.open` idiom
     return Release(path)
 
 
-def load(split="test", networks=None, labels=True, release=RELEASE, root=None, repo_id=REPO_ID):
+def load(split="test", networks=None, labels=True, release=RELEASE, root=None, repo_id=REPO_ID,
+         scenarios=None, max_windows=None, window_ids=None, lazy=False):
     """Load one split as a :class:`Split` (iterate it to get windows).
 
     ``root``: a local release folder (default: the ``BEPAD_ROOT`` environment
@@ -172,11 +265,18 @@ def load(split="test", networks=None, labels=True, release=RELEASE, root=None, r
     (``repo_id``, ``release``) and cached.
     ``labels=False`` returns only what an associator may use. ``networks=["core"]``: the
     :data:`CORE_NETWORKS`.
+    Less memory: ``scenarios`` (e.g. ``"ev-med__noise-high"``), ``max_windows`` (a fixed random
+    sample of at most this many windows per network) or ``window_ids`` (their networks are then
+    the only ones downloaded and read). ``lazy=True``: read each window from disk when it is asked
+    for (:class:`LazySplit`).
     """
+    if networks is None and window_ids is not None:     # a window_id starts with its network: "tx.test..."
+        networks = sorted({w.split(".", 1)[0] for w in window_ids})
     networks = _networks(networks)
     root = root or os.environ.get("BEPAD_ROOT")        # a local copy, e.g. export BEPAD_ROOT=data/releases/v0.1
     path = Path(root) if root else download(release, splits=(split,), networks=networks, repo_id=repo_id)
-    return Release(path).load(split, networks=networks, labels=labels)
+    return Release(path).load(split, networks=networks, labels=labels, scenarios=scenarios,
+                              max_windows=max_windows, window_ids=window_ids, lazy=lazy)
 
 
 def predictions_frame(window_ids, pick_idx, pred_event):

{bepad-0.1.0 → bepad-0.1.2}/bepad/metrics.py

@@ -301,15 +301,13 @@ def evaluate(predictions, release, split="test", networks=None, min_shared=3, wi
     if not isinstance(predictions, pd.DataFrame):
         p = Path(predictions)
         predictions = pd.read_parquet(p) if p.suffix == ".parquet" else pd.read_csv(p)
+    keep = set(window_ids) if window_ids is not None else None     # a partial run: read only its rows
     windows = rel._read(split, "windows", networks)
-    picks = rel._read(split, "picks", networks, columns=["window_id", "pick_idx", "event_idx"])
-    events = rel._read(split, "events", networks, columns=["window_id", "event_idx", "scorable", "latitude",
-                                                           "longitude", "depth", "t0_rel_s"])
-    if window_ids is not None:
-        keep = set(window_ids)
+    if keep is not None:
         windows = windows[windows.window_id.isin(keep)]
-        picks = picks[picks.window_id.isin(keep)]
-        events = events[events.window_id.isin(keep)]
+    picks = rel._read(split, "picks", networks, columns=["window_id", "pick_idx", "event_idx"], window_ids=keep)
+    events = rel._read(split, "events", networks, columns=["window_id", "event_idx", "scorable", "latitude",
+                                                           "longitude", "depth", "t0_rel_s"], window_ids=keep)
     checks = check_predictions(predictions, picks)
     if checks["n_missing_picks"]:
         print(f"[evaluate] warning: {checks['n_missing_picks']} picks have no prediction -> counted as noise")

{bepad-0.1.0 → bepad-0.1.2}/bepad/plots.py

@@ -70,6 +70,12 @@ def _finish(fig, save_path, dpi=DPI):
     return fig
 
 
+def _window_id(ax, wid):
+    """The window_id in a corner of a panel, so the window can be loaded again: bepad.load(window_ids=[...])."""
+    ax.text(0.99, 0.02, wid, transform=ax.transAxes, ha="right", va="bottom", fontsize=4.5, color=INK2, zorder=10,
+            bbox=dict(boxstyle="square,pad=0.15", fc="white", ec="none", alpha=0.85))
+
+
 def _panel(ax, letter):
     ax.text(-0.02, 1.02, letter, transform=ax.transAxes, fontweight="bold", fontsize=8.5, ha="right", va="bottom",
             color=INK)
@@ -157,6 +163,8 @@ def plot_window(w, model=None, ax=None, show_bands=True, legend=True, title=None
     ax.set_ylabel("Hypocentral distance [km]")
     if title:
         ax.set_title(title, loc="left")
+    elif own:
+        ax.set_title(w.window_id, loc="left", fontsize=7, color=INK2)
     if legend:
         h = [Line2D([], [], marker="o", ls="", ms=3.5, color=INK2, label="real P"),
              Line2D([], [], marker="o", ls="", ms=3.5, mfc="white", color=INK2, label="real S"),
@@ -208,6 +216,8 @@ def plot_window_map(w, ax=None, all_stations=None, legend=True, save_path=None,
     ax.yaxis.set_major_locator(matplotlib.ticker.MaxNLocator(4))
     ax.set_xlabel("Longitude [°]")
     ax.set_ylabel("Latitude [°]")
+    if own:
+        ax.set_title(w.window_id, loc="left", fontsize=6, color=INK2)
     if legend:
         h = [Line2D([], [], marker="^", ls="", ms=4, color=INK2, label="station with picks"),
              Line2D([], [], marker="^", ls="", ms=4, color="#c9c8c1", label="station, no picks"),
@@ -321,6 +331,7 @@ def plot_event_labels(w, model, events=None, ncols=3, save_path=None, dpi=DPI):
          plt.Rectangle((0, 0), 1, 1, color="#b9d3f1", label="band: generation allowed"),
          Line2D([], [], color=SERIES[1], ls="--", lw=0.8, label="last station")]
     fig.legend(handles=h, loc="upper center", ncol=7, bbox_to_anchor=(0.45, 1.0), fontsize=5.5)
+    fig.suptitle(w.window_id, x=0.01, y=1.04, ha="left", fontsize=7, color=INK2)
     return _finish(fig, save_path, dpi)
 
 
@@ -466,6 +477,7 @@ def plot_scenario_grid(split, network, ev_levels=("sparse", "low", "med", "high"
             w = split[wid]
             plot_window(w, model=model, ax=ax, legend=False, show_bands=False)
             ax.set_title(f"{el} ({int(w.meta['n_events'])} ev.) · noise {nl}", loc="left", fontsize=6)
+            _window_id(ax, wid)
             if i < len(noise_levels) - 1:
                 ax.set_xlabel("")
             if j > 0:
@@ -497,6 +509,7 @@ def plot_gap_sweep(split, network, scenarios=("gap-000-090", "gap-090-180", "gap
         lo, hi = sc.split("-")[1:]
         axes[0][j].set_title(f"azimuthal gap {int(lo)}–{int(hi)}°", loc="left")
         plot_window(w, ax=axes[1][j], legend=False, show_bands=False)
+        _window_id(axes[1][j], wid)
         if j:
             axes[0][j].set_ylabel("")
             axes[1][j].set_ylabel("")
@@ -841,6 +854,7 @@ def plot_representative_windows(split, by="max_concurrent_events", quantiles=(0.
             plot_window(win, model=model, ax=ax, legend=False, show_bands=False)
             ax.set_title(f"{n} ({scales[n]}) · {names.get(q, q)} · {by.replace('_', ' ')} = {int(win.meta[by])}",
                          loc="left", fontsize=5.8)
+            _window_id(ax, wid)
             if i < len(quantiles) - 1:
                 ax.set_xlabel("")
             if j:

{bepad-0.1.0 → bepad-0.1.2}/bepad/torch.py

@@ -70,16 +70,22 @@ class ReleaseWindowDataset(Dataset):
 
     ``release``: a :class:`bepad.data.Release`, a local release folder, or None
     to download ``split`` from Hugging Face.
+    Each window is read from disk when the ``DataLoader`` asks for it (``lazy=True``, the default),
+    so any split fits in memory. ``lazy=False`` holds all chosen windows in memory (faster per
+    window, but 20 GB for R07 train alone). ``networks``, ``scenarios``, ``max_windows`` (per
+    network) and ``window_ids`` choose windows, as in :func:`bepad.load`.
     """
 
-    def __init__(self, release=None, split="train", networks=None):
+    def __init__(self, release=None, split="train", networks=None, scenarios=None, max_windows=None,
+                 window_ids=None, lazy=True):
         from .data import Release, load
 
+        sub = dict(scenarios=scenarios, max_windows=max_windows, window_ids=window_ids, lazy=lazy)
         if release is None:
-            self.split = load(split, networks=networks, labels=True)
+            self.split = load(split, networks=networks, labels=True, **sub)
         else:
             rel = release if isinstance(release, Release) else Release(release)
-            self.split = rel.load(split, networks=networks, labels=True)
+            self.split = rel.load(split, networks=networks, labels=True, **sub)
 
     def __len__(self):
         return len(self.split)
@@ -94,8 +100,10 @@ class GeneratedWindowDataset(Dataset):
 
     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.
+    window parameters, e.g. ``Release.scenarios()``, BEPAD's own) are
+    drawn uniformly per item, as is the network. Networks without earthquakes in ``split`` (test-only
+    networks, e.g. R06 for train) are left out; if a network cannot make the drawn scenario, the item
+    draws another network and scenario (still reproducible).
     """
 
     def __init__(self, data_root, networks, scenarios, split="train", length=10_000, base_seed=0):
@@ -105,8 +113,17 @@ class GeneratedWindowDataset(Dataset):
         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.index = {}
+        for n in networks:
+            ev, pk, st = load_network(data_root, n)
+            if (ev["split"] == split).any():
+                self.index[n] = NetworkIndex(ev, pk, st, network=n, label_model=model)
+        skipped = sorted(set(networks) - set(self.index))
+        if skipped:
+            print(f"[GeneratedWindowDataset] no {split} earthquakes, left out: {', '.join(skipped)}")
+        if not self.index:
+            raise ValueError(f"no network has {split} earthquakes")
+        self.networks = list(self.index)
         self.scenarios = [(k, {**WINDOW_DEFAULTS, **v}) for k, v in
                           (scenarios.items() if isinstance(scenarios, dict) else enumerate(scenarios))]
         self.split = split
@@ -125,14 +142,14 @@ class GeneratedWindowDataset(Dataset):
 
         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))
+        for draw in range(20):                          # a network may not be able to make a scenario
+            net = self.networks[rng.integers(len(self.networks))]
+            name, cfg = self.scenarios[rng.integers(len(self.scenarios))]
+            w = generate_window(self.index[net], self.split, cfg, seed + draw, 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}")
+            raise RuntimeError(f"could not generate window {i}: 20 draws of network and scenario failed")
         st = self.index[net].stations.loc[w.stations["station"]].reset_index()
         return window_to_tensors(w.window["window_id"], w.picks, st, w.events)
 

{bepad-0.1.0 → bepad-0.1.2}/bepad.egg-info/PKG-INFO

@@ -1,8 +1,9 @@
 Metadata-Version: 2.4
 Name: bepad
-Version: 0.1.0
+Version: 0.1.2
 Summary: BEPAD: a benchmark dataset for seismic phase association built from UTDQuake
 License-Expression: GPL-3.0-or-later
+Project-URL: Documentation, https://bepad.readthedocs.io
 Project-URL: Repository, https://github.com/ecastillot/BEPAD
 Project-URL: Dataset, https://huggingface.co/datasets/ecastillot/BEPAD
 Requires-Python: >=3.10
@@ -30,6 +31,8 @@ Dynamic: license-file
 
 # BEPAD — a benchmark dataset for seismic phase association
 
+[![PyPI](https://img.shields.io/pypi/v/bepad?color=orange&label=PyPI&style=for-the-badge&logo=pypi&logoColor=white)](https://pypi.org/project/bepad/) [![Data v0.1](https://img.shields.io/badge/Data-v0.1-yellow?style=for-the-badge)](https://huggingface.co/datasets/ecastillot/BEPAD/tree/main/v0.1) [![Hugging Face Dataset](https://img.shields.io/badge/HuggingFace-Dataset-yellow?style=for-the-badge&logo=huggingface&logoColor=black)](https://huggingface.co/datasets/ecastillot/BEPAD) [![Read the Docs](https://img.shields.io/badge/Docs-Read%20the%20Docs-blue?style=for-the-badge&logo=read-the-docs)](https://bepad.readthedocs.io/en/latest/) [![GitHub](https://img.shields.io/badge/GitHub-BEPAD-black?style=for-the-badge&logo=github)](https://github.com/ecastillot/BEPAD)
+
 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.
@@ -37,12 +40,12 @@ thousands of association problems built from **real earthquakes and real picks**
 ![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.*
+its earthquake, (c) picks by type. The associator's job is to find which picks belong together.*
 
 ```python
 import bepad
 
-test = bepad.load("test", labels=False)          # download the test split from Hugging Face
+test = bepad.load("test", networks=["tx"], max_windows=20, labels=False)   # 20 Texas test windows
 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
@@ -64,8 +67,8 @@ bepad submit results/my-method --code associators/my-method    # opens a Hugging
 | **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).
+Documentation: https://bepad.readthedocs.io · Data: https://huggingface.co/datasets/ecastillot/BEPAD · Code: https://github.com/ecastillot/BEPAD.
+Install: `pip install bepad` (or `pip install -e .` from this folder; `.[dl]` for PyTorch, `.[build]` to rebuild the dataset).
 
 ## Repository
 

{bepad-0.1.0 → bepad-0.1.2}/bepad.egg-info/SOURCES.txt

@@ -4,6 +4,7 @@ pyproject.toml
 bepad/__init__.py
 bepad/__main__.py
 bepad/associator.py
+bepad/blocks.py
 bepad/data.py
 bepad/leaderboard.py
 bepad/metrics.py

{bepad-0.1.0 → bepad-0.1.2}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "bepad"
-version = "0.1.0"
+version = "0.1.2"
 description = "BEPAD: a benchmark dataset for seismic phase association built from UTDQuake"
 readme = "README.md"
 license = "GPL-3.0-or-later"
@@ -20,6 +20,7 @@ dependencies = [
 ]
 
 [project.urls]
+Documentation = "https://bepad.readthedocs.io"
 Repository = "https://github.com/ecastillot/BEPAD"
 Dataset = "https://huggingface.co/datasets/ecastillot/BEPAD"