PhaseNet 0.2.0__tar.gz → 0.2.2__tar.gz

phasenet-0.2.2/.github/workflows/publish.yml

@@ -0,0 +1,28 @@
+name: Publish to PyPI
+
+on:
+  release:
+    types: [published]
+
+permissions:
+  id-token: write
+
+jobs:
+  publish:
+    runs-on: ubuntu-latest
+    environment: pypi
+    steps:
+      - uses: actions/checkout@v4
+
+      - uses: actions/setup-python@v5
+        with:
+          python-version: "3.10"
+
+      - name: Install build tools
+        run: pip install build
+
+      - name: Build package
+        run: python -m build
+
+      - name: Publish to PyPI
+        uses: pypa/gh-action-pypi-publish@release/v1

{phasenet-0.2.0 → phasenet-0.2.2}/.gitignore

@@ -1,12 +1,15 @@
+docs/notebooks/data/
 figures_backup/
 figures/
 results/
 output/
 __pycache__/
+wandb/
+ToDelete/
+phasenet.egg-info/
 CLAUDE.md
 *.pth
 *.png
 *.pdf
-wandb/
-ToDelete/
-
+*.whl
+*.gzdocs/notebooks/*_executed.ipynb

{phasenet-0.2.0 → phasenet-0.2.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: phasenet
-Version: 0.2.0
+Version: 0.2.2
 Summary: A PyTorch implementation of PhaseNet for seismic and DAS phase picking
 Author: Weiqiang Zhu
 License-Expression: MIT
@@ -26,6 +26,7 @@ Requires-Dist: gcsfs
 Requires-Dist: datasets
 Requires-Dist: pyarrow
 Requires-Dist: wandb
+Requires-Dist: scikit-learn
 
 # PhaseNet-PyTorch
 

phasenet-0.2.2/docs/notebooks/lab2a_phasenet_das_inference.ipynb

@@ -0,0 +1,285 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/AI4EPS/phasenet-pytorch/blob/main/docs/notebooks/lab2a_phasenet_das_inference.ipynb)\n\n# Lab 2a: PhaseNet-DAS Inference\n\nIn this lab, you will learn how to:\n- Load DAS (Distributed Acoustic Sensing) data from the USGS Arcata array\n- Apply PhaseNet for P/S phase picking on DAS data\n- Apply PhaseNet-DAS+, a DAS-specific model with event detection\n- Compare picking quality between a general seismic model and a DAS-specific model\n\n**References:**\n- Zhu, W. & Beroza, G. (2019). PhaseNet: a deep-neural-network-based seismic arrival-time picking method. *GJI*.\n- Zhu, W. et al. (2023). Seismic arrival-time picking on distributed acoustic sensing data using semi-supervised learning. *Nature Communications*, 14, 8192.\n- Zhu, W. et al. (2025). Towards end-to-end earthquake monitoring using a multitask deep learning model. [arXiv:2506.06939](https://arxiv.org/abs/2506.06939)\n- McGuire, J.J. et al. (2025). The GorDAS DAS Experiment Above the Cascadia Locked Zone and Subducted Gorda Slab. *SRL*, 96(4), 2489-2503.\n\n---"
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Setup"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": "# Install phasenet (only needed once)\n# !pip install phasenet"
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": "import os\nimport numpy as np\nimport pandas as pd\nimport matplotlib.pyplot as plt\nimport h5py\nimport torch\nfrom huggingface_hub import hf_hub_download\n\nimport phasenet.models\nfrom phasenet.utils.inference import predict_2d, detect_picks\nfrom phasenet.utils.postprocess import (\n    estimate_origins_from_ps_pairs,\n    associate_picks,\n)\n\nDEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\nprint(f\"Using device: {DEVICE}\")\n\nDS = 10  # downsample factor for plotting\n\ndef plot_das(ax, arr, nx, nt, dt_s, **kwargs):\n    \"\"\"Plot a downsampled DAS 2D array.\"\"\"\n    arr = np.asarray(arr)\n    ax.imshow(arr[::DS, ::DS], aspect=\"auto\", interpolation=\"bilinear\",\n              extent=[0, nt * dt_s, nx, 0], **kwargs)\n\ndef plot_picks(ax, df, dt_s):\n    \"\"\"Plot P and S picks.\"\"\"\n    for phase, color in [(\"P\", \"red\"), (\"S\", \"blue\")]:\n        sub = df[df[\"phase_type\"] == phase]\n        ax.scatter(sub[\"phase_index\"] * dt_s, sub[\"channel_index\"].astype(int),\n                   s=0.5, c=color, alpha=0.3,\n                   label=f\"{phase} ({len(sub)})\")"
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": "## 1. Load DAS Data\n\nWe use DAS data from the USGS Arcata array in Northern California (McGuire et al. 2025), deployed to monitor the 2022 M6.4 Ferndale aftershock sequence. The array records strain rate along a fiber-optic cable with ~2m channel spacing at 250 Hz sampling rate.\n\nWe chose this dataset because:\n1. It is **independent** from any data used in previous PhaseNet/PhaseNet-DAS development — allowing unbiased evaluation\n2. It is a **small, self-contained** dataset (~2,470 events) — ideal for demonstrating how semi-supervised learning can rapidly improve a model on a new DAS array\n\nEach event file is an H5 file containing a 2D array of shape `(nx, nt)` — number of spatial channels by number of time samples."
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "HF_REPO = \"AI4EPS/quakeflow_das\"\n",
+    "\n",
+    "# Download a strong event with clear P and S arrivals\n",
+    "EVENT_ID = \"20230106T122215Z\"\n",
+    "h5_path = hf_hub_download(\n",
+    "    HF_REPO, f\"arcata/data/{EVENT_ID}.h5\", repo_type=\"dataset\",\n",
+    "    local_dir=\"data/quakeflow_das\",\n",
+    ")\n",
+    "\n",
+    "with h5py.File(h5_path, \"r\") as fp:\n",
+    "    data = fp[\"data\"][:]\n",
+    "    attrs = dict(fp[\"data\"].attrs)\n",
+    "\n",
+    "dt_s = float(attrs.get(\"dt_s\", 0.008))\n",
+    "begin_time = str(attrs.get(\"begin_time\", \"\"))\n",
+    "nx, nt = data.shape\n",
+    "\n",
+    "print(f\"Event: {EVENT_ID}\")\n",
+    "print(f\"Shape: (nx={nx}, nt={nt})\")\n",
+    "print(f\"Sampling interval: {dt_s} s ({1/dt_s:.0f} Hz)\")\n",
+    "print(f\"Duration: {nt * dt_s:.1f} s\")\n",
+    "print(f\"Begin time: {begin_time}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": "# Plot raw DAS waveform\nfig, ax = plt.subplots(figsize=(12, 8))\nvmax = 2 * np.percentile(np.abs(data), 95)\nplot_das(ax, data, nx, nt, dt_s, cmap=\"seismic\", vmin=-vmax, vmax=vmax)\nax.set_xlabel(\"Time (s)\")\nax.set_ylabel(\"Channel index\")\nax.set_title(f\"Raw DAS waveform — {EVENT_ID}\")\nplt.tight_layout()\nplt.show()"
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 2. Run PhaseNet (Seismic Model)\n",
+    "\n",
+    "PhaseNet (Zhu & Beroza, 2019) is a deep learning model originally designed for 3-component seismic data. It can be applied to single-component DAS data by treating the DAS channel as the vertical (Z) component and zero-padding the other two components.\n",
+    "\n",
+    "PhaseNet processes each channel independently — it sees `(3, 1, nt)` per channel, treating each DAS channel as a single seismic station."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Prepare input: DAS → 3-component format (E=0, N=0, Z=data)\n",
+    "waveform_3ch = np.zeros((3, nx, nt), dtype=np.float32)\n",
+    "waveform_3ch[2] = data.astype(np.float32)  # Z-component\n",
+    "\n",
+    "print(f\"Input shape: {waveform_3ch.shape}  (3 channels, {nx} stations, {nt} samples)\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Load pretrained PhaseNet model\n",
+    "import urllib.request\n",
+    "\n",
+    "MODELS_REPO = \"AI4EPS/models\"\n",
+    "\n",
+    "def download_model(release, filename):\n",
+    "    \"\"\"Download a model from GitHub releases (cached locally).\"\"\"\n",
+    "    cache_dir = os.path.expanduser(f\"~/.cache/phasenet/{release}\")\n",
+    "    local_path = os.path.join(cache_dir, filename)\n",
+    "    if not os.path.exists(local_path):\n",
+    "        os.makedirs(cache_dir, exist_ok=True)\n",
+    "        url = f\"https://github.com/{MODELS_REPO}/releases/download/{release}/{filename}\"\n",
+    "        print(f\"Downloading {url}\")\n",
+    "        urllib.request.urlretrieve(url, local_path)\n",
+    "    return local_path\n",
+    "\n",
+    "# Load PhaseNet-2018\n",
+    "ckpt_path = download_model(\"PhaseNet-2018\", \"phasenet2018.pth\")\n",
+    "ckpt = torch.load(ckpt_path, map_location=\"cpu\", weights_only=False)\n",
+    "model_pn = phasenet.models.phasenet.build_model(\n",
+    "    backbone=\"unet2018\", **ckpt.get(\"model_config\", {})\n",
+    ")\n",
+    "model_pn.load_state_dict(ckpt[\"model\"])\n",
+    "model_pn.to(DEVICE).eval()\n",
+    "print(\"PhaseNet-2018 loaded\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Run inference\n",
+    "# PhaseNet uses a sliding window of (256 channels, 3000 samples)\n",
+    "NX_PN, NT_PN = 256, 3000\n",
+    "MIN_PROB = 0.3\n",
+    "\n",
+    "with torch.no_grad():\n",
+    "    scores_pn = predict_2d(\n",
+    "        model_pn, waveform_3ch, NX_PN, NT_PN,\n",
+    "        min_t_overlap=0.1, device=DEVICE,\n",
+    "        taper_nx=False, taper_nt=True,  # Independent channels → no nx taper\n",
+    "    )\n",
+    "\n",
+    "print(f\"Output scores shape: {scores_pn.shape}\")  # (3, nx, nt) — P, S, Noise"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Detect picks from probability scores using peak detection\n",
+    "channel_ids = [str(i) for i in range(nx)]\n",
+    "\n",
+    "picks_pn = detect_picks(\n",
+    "    scores_pn, channel_ids, EVENT_ID, begin_time,\n",
+    "    dt_s=dt_s, vmin=MIN_PROB, device=DEVICE,\n",
+    ")\n",
+    "\n",
+    "df_pn = pd.DataFrame(picks_pn)\n",
+    "n_p = (df_pn[\"phase_type\"] == \"P\").sum()\n",
+    "n_s = (df_pn[\"phase_type\"] == \"S\").sum()\n",
+    "print(f\"PhaseNet picks: {len(df_pn)} total ({n_p} P, {n_s} S)\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": "# Plot picks on waveform\nfig, ax = plt.subplots(figsize=(12, 8))\nplot_das(ax, data, nx, nt, dt_s, cmap=\"seismic\", vmin=-vmax, vmax=vmax)\nplot_picks(ax, df_pn, dt_s)\nax.set_xlabel(\"Time (s)\")\nax.set_ylabel(\"Channel index\")\nax.set_title(f\"PhaseNet picks — {EVENT_ID}\")\nax.legend(loc=\"upper right\", markerscale=10)\nplt.tight_layout()\nplt.show()"
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": "## 3. Run PhaseNet-DAS+\n\nPhaseNet-DAS+ (Zhu et al. 2023, 2025) is specifically designed for DAS data. Unlike PhaseNet which treats each channel independently, PhaseNet-DAS+ uses 2D convolutions across both spatial (channel) and temporal dimensions, learning the spatiotemporal patterns of seismic waves propagating along the fiber.\n\nThe model was trained using a semi-supervised approach: initial labels from PhaseNet picks were iteratively refined by retraining on associated picks.\n\nKey differences from PhaseNet:\n- **Input**: 1 channel (native DAS strain rate) instead of 3\n- **Architecture**: 2D convolutions across (nx, nt) instead of independent 1D per channel\n- **Output**: P/S scores + event detection head (PhaseNet+ multi-task framework)"
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Prepare input: native DAS format (1 channel)\n",
+    "waveform_1ch = data.astype(np.float32)[np.newaxis, :, :]  # (1, nx, nt)\n",
+    "print(f\"Input shape: {waveform_1ch.shape}  (1 channel, {nx} stations, {nt} samples)\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Load pretrained PhaseNet-DAS-Plus (Arcata)\n",
+    "ckpt_path = download_model(\"PhaseNet-DAS-Plus-Arcata-v1\", \"phasenet_das_plus_arcata_v1.pth\")\n",
+    "ckpt = torch.load(ckpt_path, map_location=\"cpu\", weights_only=False)\n",
+    "model_das = phasenet.models.phasenet_das_plus.build_model(**ckpt.get(\"model_config\", {}))\n",
+    "\n",
+    "# Use EMA (Exponential Moving Average) weights for better generalization\n",
+    "if \"model_ema_weights\" in ckpt:\n",
+    "    model_das.load_state_dict(ckpt[\"model_ema_weights\"])\n",
+    "else:\n",
+    "    model_das.load_state_dict(ckpt[\"model\"])\n",
+    "\n",
+    "model_das.to(DEVICE).eval()\n",
+    "print(\"PhaseNet-DAS-Plus loaded\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "source": "# Run inference\n# PhaseNet-DAS uses a larger window: (2048 channels, 4096 samples)\nNX_DAS, NT_DAS = 2048, 4096\n\nwith torch.no_grad():\n    result = predict_2d(\n        model_das, waveform_1ch, NX_DAS, NT_DAS,\n        min_t_overlap=0.1, device=DEVICE,\n        taper_nx=True, taper_nt=True,  # DAS model: taper both axes\n    )\n\n# PhaseNet-DAS+ returns (scores, event_detections)\nif isinstance(result, tuple):\n    scores_das, event_detections = result\nelse:\n    scores_das = result\n    event_detections = None\n\nprint(f\"Phase scores shape: {scores_das.shape}\")",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Detect picks\n",
+    "picks_das = detect_picks(\n",
+    "    scores_das, channel_ids, EVENT_ID, begin_time,\n",
+    "    dt_s=dt_s, vmin=MIN_PROB, device=DEVICE,\n",
+    ")\n",
+    "\n",
+    "df_das = pd.DataFrame(picks_das)\n",
+    "n_p = (df_das[\"phase_type\"] == \"P\").sum()\n",
+    "n_s = (df_das[\"phase_type\"] == \"S\").sum()\n",
+    "print(f\"PhaseNet-DAS picks: {len(df_das)} total ({n_p} P, {n_s} S)\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": "# Side-by-side comparison\nfig, axes = plt.subplots(1, 2, figsize=(20, 8))\n\nfor ax, df, title in zip(axes, [df_pn, df_das], [\"PhaseNet\", \"PhaseNet-DAS+\"]):\n    plot_das(ax, data, nx, nt, dt_s, cmap=\"seismic\", vmin=-vmax, vmax=vmax)\n    plot_picks(ax, df, dt_s)\n    ax.set_xlabel(\"Time (s)\")\n    ax.set_ylabel(\"Channel index\")\n    ax.set_title(f\"{title} — {len(df)} picks\")\n    ax.legend(loc=\"upper right\", markerscale=10)\n\nplt.tight_layout()\nplt.show()"
+  },
+  {
+   "cell_type": "markdown",
+   "source": "## 4. Compare at Scale\n\nWe have pre-computed picks for all 2,470 Arcata events using both models. Let's compare on specific examples.",
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": "# Visualize a weak event: PhaseNet fails, PhaseNet-DAS+ succeeds\nWEAK_EVENT = \"20230119T202515Z\"\n\nh5_path = hf_hub_download(\n    HF_REPO, f\"arcata/data/{WEAK_EVENT}.h5\", repo_type=\"dataset\",\n    local_dir=\"data/quakeflow_das\",\n)\nwith h5py.File(h5_path, \"r\") as fp:\n    weak_data = fp[\"data\"][:]\n\npn_path = hf_hub_download(\n    HF_REPO, f\"arcata/phasenet/picks/{WEAK_EVENT}.parquet\",\n    repo_type=\"dataset\", local_dir=\"data/quakeflow_das\",\n)\ndas_path = hf_hub_download(\n    HF_REPO, f\"arcata/phasenet_das/picks/{WEAK_EVENT}.parquet\",\n    repo_type=\"dataset\", local_dir=\"data/quakeflow_das\",\n)\n\npn_weak = pd.read_parquet(pn_path)\ndas_weak = pd.read_parquet(das_path)\n\nfig, axes = plt.subplots(1, 2, figsize=(20, 8))\nvmax_w = 2 * np.percentile(np.abs(weak_data), 95)\nnx_w, nt_w = weak_data.shape\n\nfor ax, df, title in zip(axes, [pn_weak, das_weak], [\"PhaseNet\", \"PhaseNet-DAS+\"]):\n    plot_das(ax, weak_data, nx_w, nt_w, dt_s, cmap=\"seismic\", vmin=-vmax_w, vmax=vmax_w)\n\n    assoc = df[df[\"origin_id\"] >= 0]\n    unassoc = df[df[\"origin_id\"] < 0]\n\n    if len(unassoc) > 0:\n        ax.scatter(unassoc[\"phase_index\"] * dt_s, unassoc[\"channel_index\"].astype(int),\n                   s=0.3, c=\"gray\", alpha=0.1, label=f\"Unassociated ({len(unassoc)})\")\n    if len(assoc) > 0:\n        for phase, color in [(\"P\", \"red\"), (\"S\", \"blue\")]:\n            a = assoc[assoc[\"phase_type\"] == phase]\n            ax.scatter(a[\"phase_index\"] * dt_s, a[\"channel_index\"].astype(int),\n                       s=1, c=color, alpha=0.5, label=f\"{phase} assoc ({len(a)})\")\n\n    ax.set_xlabel(\"Time (s)\")\n    ax.set_ylabel(\"Channel index\")\n    ax.set_title(f\"{title} — {WEAK_EVENT}\")\n    ax.legend(loc=\"upper right\", markerscale=10)\n\nplt.tight_layout()\nplt.show()",
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": "# Noise-only event: PhaseNet hallucinates picks, DAS+ correctly rejects\nNOISE_EVENT = \"20230103T220320Z\"\n\nh5_path = hf_hub_download(\n    HF_REPO, f\"arcata/data/{NOISE_EVENT}.h5\", repo_type=\"dataset\",\n    local_dir=\"data/quakeflow_das\",\n)\nwith h5py.File(h5_path, \"r\") as fp:\n    noise_data = fp[\"data\"][:]\n\npn_noise = pd.read_parquet(hf_hub_download(\n    HF_REPO, f\"arcata/phasenet/picks/{NOISE_EVENT}.parquet\",\n    repo_type=\"dataset\", local_dir=\"data/quakeflow_das\",\n))\n\nfig, axes = plt.subplots(1, 2, figsize=(20, 8))\nvmax_n = 2 * np.percentile(np.abs(noise_data), 95)\nnx_n, nt_n = noise_data.shape\n\nax = axes[0]\nplot_das(ax, noise_data, nx_n, nt_n, dt_s, cmap=\"seismic\", vmin=-vmax_n, vmax=vmax_n)\nax.scatter(pn_noise[\"phase_index\"] * dt_s, pn_noise[\"channel_index\"].astype(int),\n           s=0.3, c=\"orange\", alpha=0.1)\nax.set_title(f\"PhaseNet — {len(pn_noise)} false picks\")\nax.set_xlabel(\"Time (s)\")\nax.set_ylabel(\"Channel index\")\n\nax = axes[1]\nplot_das(ax, noise_data, nx_n, nt_n, dt_s, cmap=\"seismic\", vmin=-vmax_n, vmax=vmax_n)\nax.set_title(\"PhaseNet-DAS+ — 0 picks (correct)\")\nax.set_xlabel(\"Time (s)\")\nax.set_ylabel(\"Channel index\")\n\nplt.tight_layout()\nplt.show()"
+  },
+  {
+   "cell_type": "markdown",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": "## Summary\n\n| | PhaseNet (seismic) | PhaseNet-DAS+ |\n|---|---|---|\n| **Input** | 3-component (zero-padded DAS) | Native 1-component DAS |\n| **Architecture** | Independent per-channel 1D | Spatiotemporal 2D convolutions |\n| **Strong events** | Detects P/S arrivals | Detects P/S + event heads |\n| **Weak events** | Misses most arrivals | Detects using spatial coherence |\n| **Noise** | Many false picks | Correctly rejects noise |\n\nBrowse all 2,467 prediction figures: [Google Drive](https://drive.google.com/drive/folders/1natqPKvg0q07AxLWFOSEpTEcE1OTzt2h?usp=sharing)\n\n**Next:** [Lab 2b](lab2b_pick_association.ipynb) explores how raw picks are grouped into seismic events using P-S pairing and origin time estimation."
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.10.0"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}

phasenet-0.2.2/docs/notebooks/lab2b_pick_association.ipynb

@@ -0,0 +1,195 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/AI4EPS/phasenet-pytorch/blob/main/docs/notebooks/lab2b_pick_association.ipynb)\n\n# Lab 2b: Phase Association\n\nIn this lab, you will learn how to group raw P and S picks into seismic events:\n- Pair P and S arrivals per DAS channel\n- Estimate earthquake origin times from P-S differential times\n- Use histogram-based clustering to detect events\n- Use RANSAC moveout fitting to improve association quality\n- Tune association parameters interactively\n\n**References:**\n- Zhu, W. et al. (2022). Earthquake phase association using a Bayesian Gaussian mixture model. *JGR: Solid Earth*, 127(5), e2021JB023249.\n- Zhu, W. et al. (2025a). Towards end-to-end earthquake monitoring using a multitask deep learning model. [arXiv:2506.06939](https://arxiv.org/abs/2506.06939)\n- Zhu, W. et al. (2025b). Robust earthquake location using Random Sample Consensus (RANSAC). [arXiv:2502.10933](https://arxiv.org/abs/2502.10933)\n- Song, J. et al. (2026). Scalable earthquake phase association based on Wadati and Voronoi diagrams. SSA Annual Meeting.\n- Zhu, W. et al. (2023). Seismic arrival-time picking on DAS data using semi-supervised learning. *Nature Communications*, 14, 8192.\n\n---"
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Setup"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": "# !pip install phasenet"
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": "import numpy as np\nimport pandas as pd\nimport matplotlib.pyplot as plt\nimport h5py\nfrom huggingface_hub import hf_hub_download\n\nfrom phasenet.utils.postprocess import (\n    estimate_origins_from_ps_pairs,\n    associate_picks,\n)\n\nHF_REPO = \"AI4EPS/quakeflow_das\"\n\nDS = 10  # downsample factor for plotting\n\ndef plot_das(ax, arr, nx, nt, dt_s, **kwargs):\n    \"\"\"Plot a downsampled DAS 2D array.\"\"\"\n    arr = np.asarray(arr)\n    ax.imshow(arr[::DS, ::DS], aspect=\"auto\", interpolation=\"bilinear\",\n              extent=[0, nt * dt_s, nx, 0], **kwargs)"
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": "## 1. From Picks to Origins\n\nAfter running PhaseNet or PhaseNet-DAS, we have a list of P and S picks — each with a channel index, sample index, score, and phase type. But these are individual detections, not grouped into earthquakes.\n\nThe **association** problem: given thousands of P and S picks scattered across channels and time, which picks belong to the same earthquake? Various methods exist for seismic networks, including Bayesian Gaussian mixture models (GaMMA, Zhu et al. 2022) and Voronoi tessellation (VoRA, Song et al. 2026).\n\nOn a DAS array, the problem is simpler: all channels are on a single fiber, so we can exploit the P-S differential time directly. Following PhaseNet+ (Zhu et al. 2025), we pair P and S picks per channel, estimate origin times, and cluster them using histogram peak detection."
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Load a multi-event example\n",
+    "EVENT_ID = \"20241205T185220Z\"\n",
+    "\n",
+    "h5_path = hf_hub_download(\n",
+    "    HF_REPO, f\"arcata/data/{EVENT_ID}.h5\", repo_type=\"dataset\",\n",
+    "    local_dir=\"data/quakeflow_das\",\n",
+    ")\n",
+    "with h5py.File(h5_path, \"r\") as fp:\n",
+    "    data = fp[\"data\"][:]\n",
+    "    attrs = dict(fp[\"data\"].attrs)\n",
+    "\n",
+    "dt_s = float(attrs.get(\"dt_s\", 0.008))\n",
+    "nx, nt = data.shape\n",
+    "\n",
+    "# Load pre-computed picks from PhaseNet-DAS-Plus\n",
+    "das_path = hf_hub_download(\n",
+    "    HF_REPO, f\"arcata/phasenet_das/picks/{EVENT_ID}.parquet\",\n",
+    "    repo_type=\"dataset\", local_dir=\"data/quakeflow_das\",\n",
+    ")\n",
+    "df = pd.read_parquet(das_path)\n",
+    "\n",
+    "print(f\"Event: {EVENT_ID}\")\n",
+    "print(f\"Data shape: ({nx}, {nt}), dt={dt_s}s\")\n",
+    "print(f\"Picks: {len(df)} total\")\n",
+    "print(f\"Columns: {list(df.columns)}\")\n",
+    "df.head()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": "# Visualize all picks (before association)\nfig, ax = plt.subplots(figsize=(12, 8))\nvmax = 2 * np.percentile(np.abs(data), 95)\nplot_das(ax, data, nx, nt, dt_s, cmap=\"seismic\", vmin=-vmax, vmax=vmax)\n\np_picks = df[df[\"phase_type\"] == \"P\"]\ns_picks = df[df[\"phase_type\"] == \"S\"]\nax.scatter(p_picks[\"phase_index\"] * dt_s, p_picks[\"channel_index\"].astype(int),\n           s=0.5, c=\"red\", alpha=0.3, label=f\"P ({len(p_picks)})\")\nax.scatter(s_picks[\"phase_index\"] * dt_s, s_picks[\"channel_index\"].astype(int),\n           s=0.5, c=\"blue\", alpha=0.3, label=f\"S ({len(s_picks)})\")\n\nax.set_xlabel(\"Time (s)\")\nax.set_ylabel(\"Channel index\")\nax.set_title(f\"All picks (unassociated) — {EVENT_ID}\")\nax.legend(loc=\"upper right\", markerscale=10)\nplt.tight_layout()\nplt.show()"
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": "## 2. P-S Pairing\n\nThe first step is to pair P and S picks on the same channel. For each P pick, we find the first S pick that follows within a maximum P-S time window. This exploits the physical constraint that S-waves always arrive after P-waves.\n\nFrom each P-S pair, we estimate the earthquake origin time using the Wadati relation:\n\n$$t_0 = \\frac{t_P + t_S}{2} - \\frac{t_S - t_P}{2} \\cdot \\frac{V_P/V_S + 1}{V_P/V_S - 1}$$\n\nwhere $V_P/V_S \\approx 1.73$ is the regional P-to-S velocity ratio."
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Convert pre-computed picks back to dict format for the pairing function\n",
+    "picks = df.to_dict(\"records\")\n",
+    "\n",
+    "# Pair P and S picks per channel\n",
+    "pairs = estimate_origins_from_ps_pairs(\n",
+    "    picks,\n",
+    "    dt_s=dt_s,\n",
+    "    vp_vs=1.73,         # regional Vp/Vs ratio\n",
+    "    min_ps_dt=1.0,      # minimum P-S time (seconds) to avoid false pairs\n",
+    "    max_ps_dt=30.0,     # maximum P-S time (seconds)\n",
+    ")\n",
+    "\n",
+    "print(f\"Found {len(pairs)} P-S pairs\")\n",
+    "# Each pair: (channel_id, p_index, s_index, t0_seconds, weight, ps_center, ps_interval)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": "# Visualize P-S pairs on waveform\nfig, ax = plt.subplots(figsize=(12, 8))\nplot_das(ax, data, nx, nt, dt_s, cmap=\"seismic\", vmin=-vmax, vmax=vmax)\nfor ch, p_idx, s_idx, t0, w, pc, pi in pairs:\n    ch_int = int(ch)\n    ax.plot([p_idx * dt_s, s_idx * dt_s], [ch_int, ch_int],\n            c=\"green\", alpha=0.1, lw=0.5)\nax.set_xlabel(\"Time (s)\")\nax.set_ylabel(\"Channel index\")\nax.set_title(f\"P-S pairs ({len(pairs)})\")\nplt.tight_layout()\nplt.show()"
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": "## 3. Origin Time Estimation\n\nFrom each P-S pair, we estimate when the earthquake occurred (origin time):\n\n$$t_0 = \\frac{t_P + t_S}{2} - \\frac{t_S - t_P}{2} \\cdot \\frac{V_P/V_S + 1}{V_P/V_S - 1}$$\n\nIf multiple P-S pairs from the **same earthquake** all produce similar origin times, we can detect that earthquake by finding **peaks in the histogram of estimated origin times**.\n\nThis is the core idea behind the association algorithm (Zhu et al. 2025; Song et al. 2026): cluster origin time estimates to identify discrete events."
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": "# Extract origin time estimates from pairs\nt0_estimates = np.array([t0 for _, _, _, t0, _, _, _ in pairs])\nweights = np.array([w for _, _, _, _, w, _, _ in pairs])\n\n# Build weighted histogram of origin times\nbin_size = 0.1  # seconds\nt_max = nt * dt_s\nhist_edges = np.arange(0, t_max + bin_size, bin_size)\nhist_counts, _ = np.histogram(t0_estimates, bins=hist_edges, weights=weights)\nbin_centers = (hist_edges[:-1] + hist_edges[1:]) / 2\n\nfig, ax = plt.subplots(figsize=(14, 4))\nax.bar(bin_centers, hist_counts, width=bin_size, color=\"steelblue\", edgecolor=\"none\")\nax.set_xlabel(\"Estimated origin time (seconds)\")\nax.set_ylabel(\"Weighted count\")\nax.set_title(\"Histogram of estimated origin times\")\nax.set_xlim(0, t_max)\nax.set_yscale(\"log\")\nplt.tight_layout()\nplt.show()"
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": "# Run full association: pair → histogram → peak detection → link picks to events\nassoc = associate_picks(\n    pairs,\n    nt=nt,\n    dt_s=dt_s,\n    bin_size=0.1,          # histogram bin width (seconds)\n    mph=100,               # minimum peak height\n    mpd=50,                # minimum peak distance (bins)\n    origin_window=0.5,     # max distance from peak to link a pair (seconds)\n    moveout_window=0.5,    # RANSAC residual threshold (seconds)\n)\n\nprint(f\"Detected {len(assoc.event_times)} events\")\nprint(f\"Event origin times: {[f'{t:.1f}s' for t in assoc.event_times]}\")\nprint(f\"Associated picks: {len(assoc.pick_to_event)} / {len(picks)}\")"
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": "# Visualize association results\nfig, ax = plt.subplots(figsize=(12, 8))\nplot_das(ax, data, nx, nt, dt_s, cmap=\"seismic\", vmin=-vmax, vmax=vmax)\n\ncolors = plt.cm.Set1(np.linspace(0, 1, max(len(assoc.event_times), 1)))\n\nunassoc_mask = df.apply(\n    lambda r: (r[\"channel_index\"], int(r[\"phase_index\"]), r[\"phase_type\"]) not in assoc.associated_set,\n    axis=1,\n)\nunassoc = df[unassoc_mask]\nax.scatter(unassoc[\"phase_index\"] * dt_s, unassoc[\"channel_index\"].astype(int),\n           s=0.3, c=\"gray\", alpha=0.1, label=f\"Unassociated ({len(unassoc)})\")\n\nfor eid, t_origin in enumerate(assoc.event_times):\n    event_picks = [(ch, idx, ptype) for (ch, idx, ptype), e in assoc.pick_to_event.items() if e == eid]\n    if not event_picks:\n        continue\n    chs = [int(ch) for ch, _, _ in event_picks]\n    idxs = [idx * dt_s for _, idx, _ in event_picks]\n    ax.scatter(idxs, chs, s=1, c=[colors[eid % len(colors)]], alpha=0.5,\n               label=f\"Event {eid} (t={t_origin:.1f}s, {len(event_picks)} picks)\")\n\nax.set_xlabel(\"Time (s)\")\nax.set_ylabel(\"Channel index\")\nax.set_title(f\"Associated picks — {len(assoc.event_times)} events detected\")\nax.legend(loc=\"upper right\", markerscale=10, fontsize=8)\nplt.tight_layout()\nplt.show()"
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": "## 4. Event Head Detection\n\nPhaseNet-DAS+ (Zhu et al. 2025a) includes an **event detection head** that directly predicts where earthquakes occur in the spatiotemporal domain. This provides an alternative to purely pick-based association.\n\nThe event head outputs a probability map of event centers — peaks in this map indicate likely earthquake locations. These can be used to:\n1. Provide initial event locations for phase association\n2. Detect events that have too few picks for reliable P-S pairing\n3. Improve association in windows with overlapping events"
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Let's compare: single-event window vs multi-event window\n",
+    "# Strong single event\n",
+    "SINGLE_EVENT = \"20230106T122215Z\"\n",
+    "single_path = hf_hub_download(\n",
+    "    HF_REPO, f\"arcata/phasenet_das/picks/{SINGLE_EVENT}.parquet\",\n",
+    "    repo_type=\"dataset\", local_dir=\"data/quakeflow_das\",\n",
+    ")\n",
+    "df_single = pd.read_parquet(single_path)\n",
+    "\n",
+    "n_origins_single = df_single[df_single[\"origin_id\"] >= 0][\"origin_id\"].nunique()\n",
+    "n_origins_multi = df[df[\"origin_id\"] >= 0][\"origin_id\"].nunique()\n",
+    "\n",
+    "print(f\"Single-event window ({SINGLE_EVENT}): {n_origins_single} origin(s), \"\n",
+    "      f\"{(df_single['origin_id'] >= 0).sum()} associated picks\")\n",
+    "print(f\"Multi-event window ({EVENT_ID}): {n_origins_multi} origin(s), \"\n",
+    "      f\"{(df['origin_id'] >= 0).sum()} associated picks\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": "## 5. Tune Parameters\n\nThe association has several key parameters. Let's explore how `mph` (minimum peak height) and `origin_window`/`moveout_window` affect the results:"
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": "# Compare: mph and origin/moveout window\nconfigs = [\n    {\"mph\": 100, \"origin_window\": 0.2, \"moveout_window\": 0.2},\n    {\"mph\": 100, \"origin_window\": 0.5, \"moveout_window\": 0.5},\n    {\"mph\": 200, \"origin_window\": 0.2, \"moveout_window\": 0.2},\n    {\"mph\": 200, \"origin_window\": 0.5, \"moveout_window\": 0.5},\n]\n\nfig, axes = plt.subplots(2, 2, figsize=(18, 12))\n\nfor ax, cfg in zip(axes.flat, configs):\n    assoc_test = associate_picks(\n        pairs, nt=nt, dt_s=dt_s,\n        bin_size=0.1, mph=cfg[\"mph\"], mpd=50,\n        origin_window=cfg[\"origin_window\"],\n        moveout_window=cfg[\"moveout_window\"],\n    )\n    n_events = len(assoc_test.event_times)\n    n_assoc = len(assoc_test.pick_to_event)\n\n    plot_das(ax, data, nx, nt, dt_s, cmap=\"seismic\", vmin=-vmax, vmax=vmax)\n\n    colors_t = plt.cm.Set1(np.linspace(0, 1, max(n_events, 1)))\n    for eid in range(n_events):\n        event_picks = [(ch, idx, pt) for (ch, idx, pt), e in assoc_test.pick_to_event.items() if e == eid]\n        if event_picks:\n            chs = [int(ch) for ch, _, _ in event_picks]\n            idxs = [idx * dt_s for _, idx, _ in event_picks]\n            ax.scatter(idxs, chs, s=0.5, c=[colors_t[eid % len(colors_t)]], alpha=0.5)\n\n    ax.set_title(f\"mph={cfg['mph']}, window={cfg['origin_window']}s: \"\n                 f\"{n_events} events, {n_assoc} picks\")\n    ax.set_xlabel(\"Time (s)\")\n    ax.set_ylabel(\"Channel\")\n\nplt.suptitle(\"Effect of mph and origin/moveout window\", fontsize=14)\nplt.tight_layout()\nplt.show()"
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": "## Summary\n\nThe association pipeline:\n1. **P-S pairing** — match P and S picks on the same channel within a time window\n2. **Origin estimation** — compute origin time from each pair using the Wadati relation\n3. **Histogram clustering** — bin origin time estimates, find peaks → one peak per event\n4. **Pick linking** — assign each pair to the nearest peak within `origin_window`\n5. **RANSAC moveout fitting** — fit the P-S midpoint moveout curve using RANSAC (Zhu et al. 2025b) to reject outlier pairs that don't follow a consistent moveout pattern. The `moveout_window` parameter controls the maximum RANSAC residual.\n\nKey parameters to tune:\n- `mph`: trade-off between detecting weak events vs false detections\n- `origin_window`: how tightly pairs must cluster around a detected origin time\n- `moveout_window`: RANSAC residual threshold for moveout consistency\n\nThis simple histogram-based approach works well for single DAS arrays. For seismic networks with multiple stations, more sophisticated methods are needed: GaMMA (Zhu et al. 2022) uses a Bayesian Gaussian mixture model, while VoRA (Song et al. 2026) uses Voronoi tessellation for scalable association. The origin time estimation from P-S pairs is the fundamental building block shared across all these methods.\n\n**Next:** [Lab 2c](lab2c_semisupervised_training.ipynb) shows how associated picks become training labels for iteratively improving the PhaseNet-DAS model."
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.10.0"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}