bepad 0.1.0__py3-none-any.whl

bepad/__init__.py

@@ -0,0 +1,17 @@
+"""BEPAD: a benchmark dataset for seismic phase association, built from UTDQuake.
+
+    import bepad
+    test = bepad.load("test", labels=False)         # download + read the test split
+    report = bepad.evaluate(predictions, test)      # score an associator
+    class MyMethod(bepad.Associator): ...           # template for reusable associators
+"""
+
+__version__ = "0.1.0"
+
+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
+from .associator import Associator, load_associator  # noqa: E402
+
+__all__ = ["CORE_NAME", "CORE_NETWORKS", "load", "download", "open", "evaluate", "Associator", "load_associator", "predictions_frame",
+           "Release", "Split",
+           "LABEL_COLUMNS", "__version__"]

bepad/__main__.py

@@ -0,0 +1,206 @@
+"""Command line interface.
+
+Users
+    bepad download --split test [--networks tx CM] [--release v0.1]
+    bepad evaluate predictions.parquet --release-dir <dir> [--split test] [--out report/]
+
+Associator authors
+    bepad new my-method                                   create associators/my-method/ from the template
+    bepad run associators/my-method --split val           run + score (tune on val)
+    bepad run associators/my-method --split test --out results/my-method
+    bepad submit results/my-method --code associators/my-method   check, then open a Hugging Face pull request
+
+Maintainers (building the dataset)
+    python -m bepad pool add tx CM [--force]     QC networks into data/pool
+    python -m bepad pool sync [--dry-run]         add every UTDQuake network not in the pool
+    python -m bepad pool list
+    python -m bepad pool fit-labels [--networks NEW --keep-global]
+    python -m bepad build configs/v0.1.yaml [--networks tx] [--workers N] [--no-finalize] [--release NAME]
+    python -m bepad finalize configs/v0.1.yaml
+    python -m bepad figures configs/v0.1.yaml
+"""
+
+from __future__ import annotations
+
+import argparse
+import sys
+
+
+def main(argv=None):
+    from .data import RELEASE
+
+    ap = argparse.ArgumentParser(prog="bepad", description="BEPAD phase association benchmark")
+    sub = ap.add_subparsers(dest="cmd", required=True)
+
+    dl = sub.add_parser("download", help="download part of a release from Hugging Face")
+    dl.add_argument("--split", nargs="*", default=["test"])
+    dl.add_argument("--networks", nargs="*")
+    dl.add_argument("--release", default=RELEASE)
+    dl.add_argument("--local-dir")
+
+    ev = sub.add_parser("evaluate", help="score predictions on a release split")
+    ev.add_argument("predictions")
+    ev.add_argument("--release-dir", required=True, help="local release folder (see `bepad download`)")
+    ev.add_argument("--split", default="test")
+    ev.add_argument("--networks", nargs="*")
+    ev.add_argument("--min-shared", type=int, default=3)
+    ev.add_argument("--out", help="folder for the report (csv + summary.json)")
+
+    nw = sub.add_parser("new", help="create an associator folder from the template")
+    nw.add_argument("name")
+    nw.add_argument("--root", default="associators")
+
+    rn = sub.add_parser("run", help="run an associator folder on a split and score it")
+    rn.add_argument("associator", help="folder with associator.py (e.g. associators/pyocto)")
+    rn.add_argument("--split", default="test")
+    rn.add_argument("--networks", nargs="*")
+    rn.add_argument("--scenarios", nargs="*", help="only these scenarios (partial run)")
+    rn.add_argument("--max-windows", type=int, help="at most N windows per network (partial run)")
+    rn.add_argument("--workers", type=int, default=1, help="CPU processes (windows run in parallel)")
+    rn.add_argument("--release", default=RELEASE)
+    rn.add_argument("--release-dir", help="local release folder instead of downloading")
+    rn.add_argument("--out", help="save predictions, report and submission.json here")
+
+    sb = sub.add_parser("submit", help="check a result folder and open a leaderboard pull request")
+    sb.add_argument("result")
+    sb.add_argument("--code", help="the associator folder, published with the results (recommended)")
+    sb.add_argument("--release-dir", help="local release folder for the check (otherwise downloaded)")
+    sb.add_argument("--dry-run", action="store_true", help="check only; print the manual steps")
+
+    vf = sub.add_parser("verify", help="re-check a result folder (scores recomputed from predictions)")
+    vf.add_argument("result")
+    vf.add_argument("--code")
+    vf.add_argument("--release-dir")
+
+    lb = sub.add_parser("leaderboard", help="[maintainers] build tables / review pull requests")
+    lb.add_argument("action", choices=["build", "review"])
+    lb.add_argument("--dir", default="leaderboard", help="local leaderboard folder (build)")
+    lb.add_argument("--release-dir")
+
+    pool = sub.add_parser("pool", help="[build] QC'd event pool")
+    pool.add_argument("action", choices=["add", "sync", "list", "remove", "fit-labels"])
+    pool.add_argument("networks", nargs="*")
+    pool.add_argument("--config", default="configs/pool.yaml")
+    pool.add_argument("--force", action="store_true")
+    pool.add_argument("--dry-run", action="store_true")
+    pool.add_argument("--keep-global", action="store_true",
+                      help="fit-labels: keep q(M,r) and sigma(r), only add curves of the given networks")
+
+    for name, hlp in (("build", "[build] generate release windows"), ("finalize", "[build] manifest + metadata"),
+                      ("figures", "[build] release figures")):
+        p = sub.add_parser(name, help=hlp)
+        p.add_argument("config")
+        p.add_argument("--release", help="override the release name of the config")
+        if name == "build":
+            p.add_argument("--networks", nargs="*")
+            p.add_argument("--workers", type=int)
+            p.add_argument("--no-finalize", action="store_true")
+
+    a = ap.parse_args(argv)
+
+    if a.cmd == "download":
+        from .data import download
+        print(download(a.release, splits=a.split, networks=a.networks, local_dir=a.local_dir))
+        return
+
+    if a.cmd == "new":
+        from .associator import new_associator
+        d = new_associator(a.name, a.root)
+        print(f"created {d}/ -- fill in associate() in {d}/associator.py, then: bepad run {d} --split val")
+        return
+
+    if a.cmd == "run":
+        from .associator import load_associator
+        assoc = load_associator(a.associator)
+        result = assoc.run(a.split, networks=a.networks, scenarios=a.scenarios, max_windows=a.max_windows,
+                           workers=a.workers, release=a.release, root=a.release_dir)
+        print(result)
+        if a.out:
+            print(f"[run] saved to {result.save(a.out)}")
+        return
+
+    if a.cmd in ("submit", "verify"):
+        from . import leaderboard as L
+        if a.cmd == "submit":
+            L.submit(a.result, code_dir=a.code, root=a.release_dir, dry_run=a.dry_run)
+        else:
+            probs = L.verify(a.result, root=a.release_dir, code_dir=a.code)
+            print("OK" if not probs else "Problems:\n  - " + "\n  - ".join(probs))
+        return
+
+    if a.cmd == "leaderboard":
+        from . import leaderboard as L
+        if a.action == "build":
+            from pathlib import Path
+            print(L.build(a.dir, out_md=Path(a.dir) / "README.md", out_csv=Path(a.dir) / "leaderboard.csv",
+                          out_html=Path(a.dir) / "leaderboard.html"))
+        else:
+            L.review(root=a.release_dir)
+        return
+
+    if a.cmd == "evaluate":
+        from .metrics import evaluate
+        report = evaluate(a.predictions, a.release_dir, split=a.split, networks=a.networks, min_shared=a.min_shared)
+        print(report)
+        if a.out:
+            print(f"[evaluate] report written to {report.save(a.out)}")
+        return
+
+    from .build.config import load_pool_config, load_release_config
+
+    if a.cmd == "pool":
+        import json
+
+        from .build import pool as P
+        from .build.labels import LabelModel, fit_label_model
+
+        cfg = load_pool_config(a.config)
+        if a.action == "add":
+            for n in a.networks:
+                P.build_network(n, cfg, force=a.force)
+        elif a.action == "sync":
+            todo = P.missing_networks(cfg)
+            print(f"[pool] networks not in pool: {todo}")
+            if not a.dry_run:
+                failed = []
+                for n in todo:
+                    try:
+                        P.build_network(n, cfg)
+                    except Exception as exc:  # keep going; report at the end
+                        failed.append(n)
+                        print(f"[pool] {n} FAILED: {exc!r}", file=sys.stderr)
+                if failed:
+                    print(f"[pool] failed networks: {failed}", file=sys.stderr)
+        elif a.action == "remove":
+            for n in a.networks:
+                P.remove_network(n, cfg)
+        elif a.action == "fit-labels":
+            path = P.pool_dir(cfg["data_root"]) / "label_model.json"
+            prev = json.loads(path.read_text()) if path.exists() else None
+            nets = a.networks or P.list_networks(cfg["data_root"])
+            fit_label_model(cfg["data_root"], nets, cfg["labels"], previous=prev,
+                            refit_global=not a.keep_global).save(path)
+            print(f"[labels] wrote {path}")
+            LabelModel.load(path)
+        m = P.read_manifest(cfg["data_root"]) or {"networks": {}}
+        for n, e in sorted(m["networks"].items()):
+            print(f"  {n:<16} {e['n_events']:>8} events {e['n_picks']:>9} picks  {e['split_counts']}")
+        return
+
+    cfg = load_release_config(a.config)
+    if a.release:
+        cfg["release"] = a.release
+    if a.cmd == "build":
+        from .build.release import build_release
+        build_release(cfg, networks=a.networks, workers=a.workers, finalize=not a.no_finalize)
+    elif a.cmd == "finalize":
+        from .build.release import finalize_release
+        finalize_release(cfg)
+    elif a.cmd == "figures":
+        from .build.release import release_dir
+        from .plots import release_gallery
+        print(release_gallery(release_dir(cfg), cfg, networks=cfg.get("gallery_networks")))
+
+
+if __name__ == "__main__":
+    main()

bepad/associator.py

@@ -0,0 +1,339 @@
+"""The associator template: fill in one method, get runs, reports and leaderboard submissions.
+
+    import bepad
+
+    class MyAssociator(bepad.Associator):
+        name = "my-method"
+        version = "1.0"
+
+        def associate(self, picks, stations):
+            ...                      # one integer per pick: event number, or -1 for noise
+            return labels
+
+    result = MyAssociator(threshold=0.3).run("test", networks=["tx"])
+    result.report                    # scores
+    result.save("results/my-method")  # predictions + report + submission.json
+
+A contributed associator is a folder (``bepad new <name>`` creates it)::
+
+    associators/<name>/
+        associator.py      the class
+        config.yaml        its parameters (passed to __init__)
+        README.md          what it is, paper, license
+        requirements.txt   extra packages it needs
+"""
+
+from __future__ import annotations
+
+import datetime as dt
+import importlib.util
+import json
+import platform
+import time
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+
+METHOD_TYPES = ("classical", "deep learning", "hybrid")
+TRAINING = ("none", "per network", "all networks")
+
+
+class Associator:
+    """Base class. Subclasses set the metadata attributes and implement :meth:`associate`."""
+
+    #: short unique name (letters, digits, dashes), used for folders and the leaderboard
+    name: str = "unnamed"
+    version: str = "0.1"
+    authors: str = ""
+    paper: str = ""          # DOI or URL
+    code: str = ""           # repository URL
+    method_type: str = "classical"   # "classical" | "deep learning" | "hybrid"
+    training: str = "none"           # "none" | "per network" | "all networks"
+    description: str = ""
+
+    def __init__(self, **config):
+        nets = config.get("networks") or {}
+        if not isinstance(nets, dict) or not all(isinstance(v, dict) for v in nets.values()):
+            raise ValueError("config 'networks' must map a network to its settings, e.g. {tx: {p_velocity: 6.0}}")
+        self._base_config = config
+        self.network = None          # the network of the current window (set before each associate call)
+        self.path = None             # the associator folder (set by load_associator), for model files
+        self.config = self.network_config(None)
+
+    def network_config(self, network):
+        """``config.yaml`` for one network: the general settings, overridden by its ``networks:`` entry."""
+        base = {k: v for k, v in self._base_config.items() if k != "networks"}
+        return {**base, **(self._base_config.get("networks") or {}).get(network, {})}
+
+    def use_network(self, network):
+        """Called by BEPAD before each window: sets ``self.network`` and ``self.config``."""
+        if network != self.network:
+            self.network, self.config = network, self.network_config(network)
+
+    # --- to implement --------------------------------------------------------------
+    def fit(self, train, val):
+        """Optional: learn or tune on the ``train`` / ``val`` splits (never on test)."""
+
+    def associate(self, picks: pd.DataFrame, stations: pd.DataFrame):
+        """Associate the picks of ONE window.
+
+        picks:    pick_idx, station, phase ("P"/"S"), t_rel_s (s from window start), time (UTC)
+        stations: station, latitude, longitude, elevation -- every station of the window
+        returns:  one integer per pick, in the same order: event number (>= 0) or -1 (noise)
+
+        Associators that also locate can return ``labels, events`` instead, where
+        ``events`` has one row per event number: event, latitude, longitude,
+        depth_km, t0_rel_s (origin time in seconds from the window start); any of
+        the location columns may be left out. Location errors are then reported too.
+        """
+        raise NotImplementedError
+
+    # --- provided ------------------------------------------------------------------
+    def run(self, split="test", networks=None, scenarios=None, max_windows=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.
+        """
+        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)
+        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]
+        rows, locs, secs, t0, n = [], [], [], time.time(), len(todo)
+        _RUNNING[0] = self
+        if workers > 1:
+            import multiprocessing as mp
+            pool = mp.get_context("fork").Pool(workers)
+            results = pool.imap(_associate_one, todo, chunksize=2)
+        else:
+            pool, results = None, map(_associate_one, todo)
+        try:
+            for k, (df, loc, sec) in enumerate(results):
+                rows.append(df)
+                secs.append(sec)
+                if loc is not None:
+                    locs.append(loc)
+                if progress and (k + 1) % 50 == 0:
+                    print(f"[{self.name}] {k + 1}/{n} windows, {time.time() - t0:.0f} s")
+        finally:
+            if pool is not None:
+                pool.close()
+                pool.join()
+        runtime = time.time() - t0
+        predictions = pd.concat(rows, ignore_index=True)
+        locations = pd.concat(locs, ignore_index=True) if locs else None
+        report = evaluate(predictions, data, window_ids=ids, locations=locations)
+        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),
+            "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()})",
+        }
+        return Result(self, predictions, report, meta, locations)
+
+    def metadata(self):
+        if self.method_type not in METHOD_TYPES:
+            raise ValueError(f"method_type must be one of {METHOD_TYPES}")
+        if self.training not in TRAINING:
+            raise ValueError(f"training must be one of {TRAINING}")
+        return {"name": self.name, "version": self.version, "authors": self.authors, "paper": self.paper,
+                "code": self.code, "method_type": self.method_type, "training": self.training,
+                "description": self.description, "config": self._base_config}
+
+
+_RUNNING = [None]          # the associator being run (inherited by forked workers)
+
+
+def _associate_one(item):
+    window_id, network, picks, stations = item
+    assoc = _RUNNING[0]
+    assoc.use_network(network)
+    tic = time.perf_counter()
+    out = assoc.associate(picks, stations)
+    elapsed = time.perf_counter() - tic
+    loc = None
+    if isinstance(out, tuple):
+        out, loc = out
+        loc = pd.DataFrame(loc).rename(columns={"pred_event": "event"})
+        if "event" not in loc:
+            raise ValueError(f"{assoc.name}: the events table needs an 'event' column (the event numbers)")
+        keep = ["event"] + [c for c in ("latitude", "longitude", "depth_km", "t0_rel_s") if c in loc]
+        loc = loc[keep].assign(window_id=window_id)
+    labels = np.asarray(out)
+    if labels.shape != (len(picks),) or not np.issubdtype(labels.dtype, np.integer):
+        raise ValueError(f"{assoc.name}.associate must return {len(picks)} integers for {window_id}, "
+                         f"got shape {labels.shape} dtype {labels.dtype}")
+    return pd.DataFrame({"window_id": window_id, "pick_idx": picks.pick_idx.to_numpy(),
+                         "pred_event": labels.astype(np.int64)}), loc, elapsed
+
+
+class Result:
+    """Predictions + report of one run. ``save`` writes everything a submission needs."""
+
+    def __init__(self, associator, predictions, report, meta, locations=None):
+        self.associator, self.predictions, self.report, self.meta = associator, predictions, report, meta
+        self.locations = locations      # one row per predicted event (only if the associator locates)
+
+    def __repr__(self):
+        return f"<bepad.Result {self.associator.name} {self.meta['split']} {self.meta['n_windows']} windows>\n" \
+               + repr(self.report).split("\n", 1)[1]
+
+    def submission(self):
+        """The leaderboard record (a dict, saved as submission.json)."""
+        from . import __version__
+
+        r = self.report
+        cols = ["event_precision", "event_recall", "event_f1", "pick_precision", "pick_recall", "pick_f1",
+                "noise_rejection", "split_rate", "merge_rate", "ari_mean", "n_windows"]
+        cols += [c for c in ("located_frac", "epi_median_km", "epi_p90_km", "depth_median_km", "depth_p90_km",
+                             "t0_median_s", "t0_p90_s") if c in r.overall]
+        by_net = r.by_network.set_index("network")[cols]
+        return {
+            "bepad_version": __version__,
+            "created": dt.datetime.now().isoformat(timespec="seconds"),
+            "associator": self.associator.metadata(),
+            "run": self.meta,
+            "scores": {"overall": {k: _num(v) for k, v in r.overall.iloc[0][cols].items()},
+                       "by_network": {n: {k: _num(v) for k, v in row.items()} for n, row in by_net.iterrows()}},
+        }
+
+    def save(self, out_dir):
+        """Write ``submission.json``, ``predictions/network=<name>.parquet``, ``locations/`` (if any)
+        and the report tables."""
+        out = Path(out_dir)
+        (out / "predictions").mkdir(parents=True, exist_ok=True)
+        net = self.predictions.window_id.str.split(".", n=1).str[0]
+        for n, df in self.predictions.groupby(net):
+            df.to_parquet(out / "predictions" / f"network={n}.parquet", index=False, compression="zstd")
+        if self.locations is not None:
+            (out / "locations").mkdir(exist_ok=True)
+            lnet = self.locations.window_id.str.split(".", n=1).str[0]
+            for n, df in self.locations.groupby(lnet):
+                df.to_parquet(out / "locations" / f"network={n}.parquet", index=False)
+        self.report.save(out / "report")
+        (out / "submission.json").write_text(json.dumps(self.submission(), indent=2))
+        return out
+
+
+def _num(v):
+    v = float(v)
+    return None if v != v else round(v, 6)
+
+
+# ----------------------------------------------------------------------
+def load_associator(path, **overrides):
+    """Load an associator folder (or .py file): its Associator subclass + config.yaml."""
+    import yaml
+
+    path = Path(path)
+    file = path / "associator.py" if path.is_dir() else path
+    spec = importlib.util.spec_from_file_location(f"bepad_assoc_{file.parent.name}", file)
+    mod = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(mod)
+    classes = [c for c in vars(mod).values()
+               if isinstance(c, type) and issubclass(c, Associator) and c is not Associator]
+    if len(classes) != 1:
+        raise ValueError(f"{file} must define exactly one bepad.Associator subclass, found {len(classes)}")
+    cfg_file = file.parent / "config.yaml"
+    config = (yaml.safe_load(cfg_file.read_text()) or {}) if cfg_file.exists() else {}
+    obj = classes[0](**{**config, **overrides})
+    obj.path = file.parent
+    return obj
+
+
+TEMPLATE = {
+    "associator.py": '''"""{name}: a BEPAD associator.
+
+Fill in ``associate`` (required) and ``fit`` (optional), set the metadata, then:
+    bepad run associators/{name} --split val        # tune
+    bepad run associators/{name} --split test --out results/{name}
+    bepad submit results/{name}
+"""
+
+import numpy as np
+
+import bepad
+
+
+class Associator(bepad.Associator):
+    name = "{name}"
+    version = "0.1"
+    authors = "Your Name"
+    paper = ""                   # DOI or URL of the paper describing the method
+    code = ""                    # URL of the original code, if any
+    method_type = "classical"    # "classical" | "deep learning" | "hybrid"
+    training = "none"            # "none" | "per network" | "all networks"
+    description = "One sentence: what the method does."
+
+    def fit(self, train, val):
+        """Optional: learn or tune on train / val. Parameters are in self.config."""
+
+    def associate(self, picks, stations):
+        """One window.
+
+        picks:    pick_idx, station, phase ("P"/"S"), t_rel_s (s from window start), time (UTC)
+        stations: station, latitude, longitude, elevation
+        return:   one integer per pick (same order): event number, or -1 for noise
+
+        self.network is the window's network (e.g. "tx") and self.config its settings
+        (config.yaml, with that network's overrides). Files of the folder: self.path / "models/...".
+        """
+        labels = np.full(len(picks), -1)          # <- replace with your method
+        return labels
+''',
+    "config.yaml": '''# Parameters of {name}; passed to the class as self.config.
+# Every value here is saved with the results, so others can reproduce them.
+example_threshold: 0.5
+
+# Optional: settings for some networks, overriding the ones above (e.g. a 1D velocity model).
+# Model files go in models/velocity/ (velocity models) and models/weights/ (trained models).
+# networks:
+#   tx:  {{velocity_model: models/velocity/tx_1d.csv}}
+#   JMA: {{velocity_model: models/velocity/jma2001.csv}}
+''',
+    "README.md": '''# {name}
+
+What the method does, in two or three sentences.
+
+- **Paper:** <DOI or URL>
+- **Original code:** <URL>
+- **License:** <license of this code>
+- **Requirements:** see `requirements.txt`
+
+## Parameters
+
+| Parameter | Meaning |
+|---|---|
+| `example_threshold` | ... |
+''',
+    "requirements.txt": "# extra packages this associator needs, one per line\n",
+}
+
+
+def new_associator(name, root="associators"):
+    """Create ``<root>/<name>/`` from the template."""
+    import re
+
+    if not re.fullmatch(r"[a-z0-9][a-z0-9_-]*", name):
+        raise ValueError("name: lowercase letters, digits, dashes and underscores (e.g. 'my-method')")
+    d = Path(root) / name
+    if d.exists():
+        raise FileExistsError(f"{d} already exists")
+    d.mkdir(parents=True)
+    for fname, text in TEMPLATE.items():
+        (d / fname).write_text(text.format(name=name))
+    return d

bepad/build/__init__.py

@@ -0,0 +1,11 @@
+"""Building BEPAD (maintainers only).
+
+Users never need this subpackage: they download the finished dataset with
+``bepad.load``. The build runs in four steps:
+
+    pool.py     UTDQuake network -> quality-controlled events, picks, stations
+    labels.py   fit the label model: detectability q(M, r), timing sigma(r),
+                per-network travel-time curves and S/P ratios
+    windows.py  pool + label model + seed -> one window
+    release.py  thousands of windows per network/split/scenario -> Parquet
+"""