PhaseNet 0.1.0__tar.gz → 0.2.1__tar.gz

phasenet-0.2.1/.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.1/.gitignore

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

phasenet-0.2.1/PKG-INFO

@@ -0,0 +1,163 @@
+Metadata-Version: 2.4
+Name: phasenet
+Version: 0.2.1
+Summary: A PyTorch implementation of PhaseNet for seismic and DAS phase picking
+Author: Weiqiang Zhu
+License-Expression: MIT
+Project-URL: Homepage, https://github.com/AI4EPS/phasenet-pytorch
+Project-URL: Repository, https://github.com/AI4EPS/phasenet-pytorch
+Classifier: Programming Language :: Python :: 3
+Classifier: Operating System :: OS Independent
+Classifier: Topic :: Scientific/Engineering
+Requires-Python: >=3.8
+Description-Content-Type: text/markdown
+Requires-Dist: torch
+Requires-Dist: torchvision
+Requires-Dist: einops
+Requires-Dist: numpy
+Requires-Dist: scipy
+Requires-Dist: h5py
+Requires-Dist: matplotlib
+Requires-Dist: pandas
+Requires-Dist: tqdm
+Requires-Dist: fsspec
+Requires-Dist: obspy
+Requires-Dist: gcsfs
+Requires-Dist: datasets
+Requires-Dist: pyarrow
+Requires-Dist: wandb
+Requires-Dist: scikit-learn
+
+# PhaseNet-PyTorch
+
+PyTorch implementation of PhaseNet for seismic and DAS phase picking, event detection, and polarity classification.
+
+## Models
+
+| Model | Features | Data Type |
+|-------|----------|-----------|
+| `phasenet` | Phase (P/S) picking | Seismic 3-component |
+| `phasenet_tf` | Phase picking + STFT spectrogram | Seismic 3-component |
+| `phasenet_plus` | Phase + polarity + event detection | Seismic 3-component |
+| `phasenet_tf_plus` | Phase + polarity + event detection + STFT | Seismic 3-component |
+| `phasenet_das` | Phase picking | DAS single-channel |
+| `phasenet_das_plus` | Phase + event detection | DAS single-channel |
+
+The `_tf` variants add a Short-Time Fourier Transform (STFT) branch that extracts frequency features alongside the temporal waveform, improving performance on noisy data.
+
+## Prediction
+
+### CEED (Seismic) Prediction
+
+```bash
+# Demo: process a few events with plots
+python scripts/predict_ceed.py --n-events 5
+
+# Process all days for a year (saves parquets to results/ceed/)
+python scripts/predict_ceed.py --all --year 2025 --output-dir results/ceed
+```
+
+Output: one parquet per day file in `results/ceed/{region}/`, with columns:
+`event_id, station_id, waveform_index, origin_id, origin_index, origin_time, phase_index, phase_time, phase_score, phase_type, phase_polarity`
+
+### DAS Prediction
+
+```bash
+# Predict with base PhaseNet on a HuggingFace subset
+python scripts/predict_das.py --subset arcata --plot
+
+# Predict with a trained DAS model on local data
+python scripts/predict_das.py \
+    --data-dir data/quakeflow_das/arcata/data \
+    --model-type phasenet_das_plus \
+    --checkpoint output/train_das_arcata/checkpoint.pth \
+    --output-dir results/das/train_das_arcata/arcata \
+    --no-ema --plot
+
+# Predict from a file list
+python scripts/predict_das.py \
+    --file-list file_list.txt \
+    --model-type phasenet_das_plus \
+    --checkpoint output/model.pth \
+    --plot
+
+# Multi-GPU prediction
+bash scripts/predict_das.sh phasenet_das_plus arcata 8 output/model.pth
+```
+
+Output: one parquet per event in `results/das/{model_name}/{subset}/`, with columns:
+`event_id, channel_index, origin_id, origin_index, origin_time, phase_index, phase_time, phase_score, phase_type, dt_s, ps_center, ps_interval`
+
+By default, picks are associated using P-S pairing. Use `--use-event-head` to associate via the model's event detection head instead.
+
+## Training
+
+### CEED (Seismic) Training
+
+```bash
+python train.py \
+    --model phasenet_plus \
+    --dataset-type ceed \
+    --label-path results/ceed \
+    --nx 16 \
+    --max-iters 100000 \
+    --batch-size 8 \
+    --workers 4 \
+    --lr 3e-4 \
+    --eval-interval 5000 \
+    --output-dir output/train_ceed
+```
+
+### DAS Training
+
+```bash
+# Using the training script
+bash scripts/train_das.sh 0 arcata v26
+
+# Or directly
+python train.py \
+    --model phasenet_das_plus \
+    --dataset-type das \
+    --data-path data/quakeflow_das/arcata/data \
+    --label-path results/das/phasenet/arcata/picks \
+    --label-list results/das/phasenet/arcata/labels.txt \
+    --nx 2048 --nt 4096 \
+    --num-patch 16 \
+    --max-iters 50000 \
+    --batch-size 2 --workers 8 \
+    --lr 1e-4 --weight-decay 0.01 \
+    --model-ema --model-ema-decay 0.999 \
+    --eval-interval 1000 --save-interval 1000 \
+    --output-dir output/train_das_arcata_v26
+```
+
+### Key Training Options
+
+| Option | Description | Default |
+|--------|-------------|---------|
+| `--num-patch N` | Random crops per DAS sample (amortizes IO) | 2 |
+| `--model-ema` | Enable exponential moving average | off |
+| `--gradient-accumulation-steps N` | Accumulate gradients for larger effective batch | 1 |
+| `--clip-grad-norm V` | Gradient clipping | 1.0 |
+| `--compile` | Enable torch.compile | off |
+| `--resume --checkpoint PATH` | Resume from checkpoint | - |
+| `--reset-lr` | Reset LR schedule when resuming | off |
+
+## Semi-supervised Training
+
+Iterative self-training pipeline for DAS: predict → train on predictions → predict with new model → repeat.
+
+```bash
+# Start from PhaseNet (train DAS model from scratch)
+bash scripts/semisupervised_das.sh arcata 5 0
+
+# Start from a pretrained DAS model
+bash scripts/semisupervised_das.sh arcata 5 0 phasenet_das output/train_das_v26/checkpoint.pth
+```
+
+Arguments: `[subset] [num_iterations] [gpu] [start_from] [checkpoint]`
+
+- **From phasenet**: iteration 0 predicts with PhaseNet, iteration 1 trains from scratch (10k steps, warmup), iterations 2+ continue (1k steps, no warmup)
+- **From phasenet_das**: iteration 0 predicts with pretrained DAS model, iterations 1+ continue (1k steps, no warmup)
+
+Results are saved to `results/semisupervised_das/` and checkpoints to `output/semisupervised_das_v{N}/`.

phasenet-0.2.1/README.md

@@ -0,0 +1,133 @@
+# PhaseNet-PyTorch
+
+PyTorch implementation of PhaseNet for seismic and DAS phase picking, event detection, and polarity classification.
+
+## Models
+
+| Model | Features | Data Type |
+|-------|----------|-----------|
+| `phasenet` | Phase (P/S) picking | Seismic 3-component |
+| `phasenet_tf` | Phase picking + STFT spectrogram | Seismic 3-component |
+| `phasenet_plus` | Phase + polarity + event detection | Seismic 3-component |
+| `phasenet_tf_plus` | Phase + polarity + event detection + STFT | Seismic 3-component |
+| `phasenet_das` | Phase picking | DAS single-channel |
+| `phasenet_das_plus` | Phase + event detection | DAS single-channel |
+
+The `_tf` variants add a Short-Time Fourier Transform (STFT) branch that extracts frequency features alongside the temporal waveform, improving performance on noisy data.
+
+## Prediction
+
+### CEED (Seismic) Prediction
+
+```bash
+# Demo: process a few events with plots
+python scripts/predict_ceed.py --n-events 5
+
+# Process all days for a year (saves parquets to results/ceed/)
+python scripts/predict_ceed.py --all --year 2025 --output-dir results/ceed
+```
+
+Output: one parquet per day file in `results/ceed/{region}/`, with columns:
+`event_id, station_id, waveform_index, origin_id, origin_index, origin_time, phase_index, phase_time, phase_score, phase_type, phase_polarity`
+
+### DAS Prediction
+
+```bash
+# Predict with base PhaseNet on a HuggingFace subset
+python scripts/predict_das.py --subset arcata --plot
+
+# Predict with a trained DAS model on local data
+python scripts/predict_das.py \
+    --data-dir data/quakeflow_das/arcata/data \
+    --model-type phasenet_das_plus \
+    --checkpoint output/train_das_arcata/checkpoint.pth \
+    --output-dir results/das/train_das_arcata/arcata \
+    --no-ema --plot
+
+# Predict from a file list
+python scripts/predict_das.py \
+    --file-list file_list.txt \
+    --model-type phasenet_das_plus \
+    --checkpoint output/model.pth \
+    --plot
+
+# Multi-GPU prediction
+bash scripts/predict_das.sh phasenet_das_plus arcata 8 output/model.pth
+```
+
+Output: one parquet per event in `results/das/{model_name}/{subset}/`, with columns:
+`event_id, channel_index, origin_id, origin_index, origin_time, phase_index, phase_time, phase_score, phase_type, dt_s, ps_center, ps_interval`
+
+By default, picks are associated using P-S pairing. Use `--use-event-head` to associate via the model's event detection head instead.
+
+## Training
+
+### CEED (Seismic) Training
+
+```bash
+python train.py \
+    --model phasenet_plus \
+    --dataset-type ceed \
+    --label-path results/ceed \
+    --nx 16 \
+    --max-iters 100000 \
+    --batch-size 8 \
+    --workers 4 \
+    --lr 3e-4 \
+    --eval-interval 5000 \
+    --output-dir output/train_ceed
+```
+
+### DAS Training
+
+```bash
+# Using the training script
+bash scripts/train_das.sh 0 arcata v26
+
+# Or directly
+python train.py \
+    --model phasenet_das_plus \
+    --dataset-type das \
+    --data-path data/quakeflow_das/arcata/data \
+    --label-path results/das/phasenet/arcata/picks \
+    --label-list results/das/phasenet/arcata/labels.txt \
+    --nx 2048 --nt 4096 \
+    --num-patch 16 \
+    --max-iters 50000 \
+    --batch-size 2 --workers 8 \
+    --lr 1e-4 --weight-decay 0.01 \
+    --model-ema --model-ema-decay 0.999 \
+    --eval-interval 1000 --save-interval 1000 \
+    --output-dir output/train_das_arcata_v26
+```
+
+### Key Training Options
+
+| Option | Description | Default |
+|--------|-------------|---------|
+| `--num-patch N` | Random crops per DAS sample (amortizes IO) | 2 |
+| `--model-ema` | Enable exponential moving average | off |
+| `--gradient-accumulation-steps N` | Accumulate gradients for larger effective batch | 1 |
+| `--clip-grad-norm V` | Gradient clipping | 1.0 |
+| `--compile` | Enable torch.compile | off |
+| `--resume --checkpoint PATH` | Resume from checkpoint | - |
+| `--reset-lr` | Reset LR schedule when resuming | off |
+
+## Semi-supervised Training
+
+Iterative self-training pipeline for DAS: predict → train on predictions → predict with new model → repeat.
+
+```bash
+# Start from PhaseNet (train DAS model from scratch)
+bash scripts/semisupervised_das.sh arcata 5 0
+
+# Start from a pretrained DAS model
+bash scripts/semisupervised_das.sh arcata 5 0 phasenet_das output/train_das_v26/checkpoint.pth
+```
+
+Arguments: `[subset] [num_iterations] [gpu] [start_from] [checkpoint]`
+
+- **From phasenet**: iteration 0 predicts with PhaseNet, iteration 1 trains from scratch (10k steps, warmup), iterations 2+ continue (1k steps, no warmup)
+- **From phasenet_das**: iteration 0 predicts with pretrained DAS model, iterations 1+ continue (1k steps, no warmup)
+
+Results are saved to `results/semisupervised_das/` and checkpoints to `output/semisupervised_das_v{N}/`.

phasenet-0.2.1/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
+}