bepad 0.1.0__tar.gz → 0.1.1__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: bepad
-Version: 0.1.0
+Version: 0.1.1
 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

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

@@ -6,7 +6,7 @@
     class MyMethod(bepad.Associator): ...           # template for reusable associators
 """
 
-__version__ = "0.1.0"
+__version__ = "0.1.1"
 
 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.1}/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.0 → bepad-0.1.1}/bepad/data.py

@@ -97,24 +97,36 @@ 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)
+        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_windows(f, columns, window_ids) 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."""
+    def load(self, split, networks=None, labels=True, scenarios=None, max_windows=None, window_ids=None):
+        """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.
+        """
         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"])
         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 +144,32 @@ 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 _read_windows(file, columns, window_ids, batch_size=200_000):
+    """Rows of one parquet file whose window is in ``window_ids``, read in batches (low memory)."""
+    import pyarrow as pa
+    import pyarrow.compute as pc
+    import pyarrow.parquet as pq
+
+    pf = pq.ParquetFile(file, pre_buffer=False, buffer_size=1 << 20)   # stream, do not read the file at once
+    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)
+    parts = [b.filter(pc.is_in(b.column("window_id"), value_set=ids))
+             for b in pf.iter_batches(batch_size=batch_size, columns=columns)]
+    return pa.Table.from_batches(parts, schema=schema).to_pandas()
+
+
 def _networks(networks):
     """``"core"`` in a list of networks stands for the core networks."""
     if networks is None:
@@ -164,7 +202,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):
     """Load one split as a :class:`Split` (iterate it to get windows).
 
     ``root``: a local release folder (default: the ``BEPAD_ROOT`` environment
@@ -172,11 +211,17 @@ 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).
     """
+    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)
 
 
 def predictions_frame(window_ids, pick_idx, pred_event):

{bepad-0.1.0 → bepad-0.1.1}/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.1}/bepad.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: bepad
-Version: 0.1.0
+Version: 0.1.1
 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

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

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "bepad"
-version = "0.1.0"
+version = "0.1.1"
 description = "BEPAD: a benchmark dataset for seismic phase association built from UTDQuake"
 readme = "README.md"
 license = "GPL-3.0-or-later"