phasenet 0.2.0__tar.gz

phasenet-0.2.0/.gitignore

@@ -0,0 +1,12 @@
+figures_backup/
+figures/
+results/
+output/
+__pycache__/
+CLAUDE.md
+*.pth
+*.png
+*.pdf
+wandb/
+ToDelete/
+

phasenet-0.2.0/PKG-INFO

@@ -0,0 +1,162 @@
+Metadata-Version: 2.4
+Name: phasenet
+Version: 0.2.0
+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
+
+# 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.0/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.0/examples/ceed_training_plots.py

@@ -0,0 +1,181 @@
+#!/usr/bin/env python
+"""Generate training label plots from CEED prediction results.
+
+Loads a prediction parquet, joins with original waveforms from GCS,
+and plots top-N events with training labels.
+
+Usage:
+    python examples/ceed_training_plots.py \
+        --parquet results/ceed/phasenet2018/NC/2025_001.parquet \
+        --region NC --year 2025 --day 001
+"""
+import argparse
+import json
+import os
+import sys
+
+sys.path.insert(0, ".")
+
+import gcsfs
+import numpy as np
+import pandas as pd
+import pyarrow.parquet as pq
+
+from phasenet.data.ceed import (
+    records_to_sample, default_train_transforms, plot_demo, Target,
+)
+
+BUCKETS = {
+    "NC": "quakeflow_dataset/NCEDC",
+    "SC": "quakeflow_dataset/SCEDC",
+}
+GCS_CRED_PATH = os.path.expanduser("~/.config/gcloud/application_default_credentials.json")
+
+
+def get_storage_options():
+    if os.path.exists(GCS_CRED_PATH):
+        with open(GCS_CRED_PATH) as f:
+            return {"token": json.load(f)}
+    return {}
+
+
+def load_day_waveforms(fs, bucket, year, day):
+    """Return {waveform_index: ndarray (3, nt)} for a day parquet."""
+    path = f"{bucket}/waveform_parquet/{year}/{day}.parquet"
+    pf = pq.ParquetFile(fs.open(path))
+    non_wave_cols = [c for c in pf.schema_arrow.names if c != "waveform"]
+    result = {}
+    row_offset = 0
+    for batch in pf.iter_batches(batch_size=256):
+        n = len(batch)
+        col = batch.column("waveform")
+        flat = col.flatten().flatten()
+        buf = flat.buffers()[1]
+        waveforms = np.frombuffer(
+            buf, dtype=np.float32, offset=flat.offset * 4
+        ).reshape(n, 3, -1).copy()
+        meta = batch.select(non_wave_cols).to_pydict()
+        for i in range(n):
+            eid = meta["event_id"][i] or ""
+            if eid:
+                result[row_offset + i] = waveforms[i]
+        row_offset += n
+    return result
+
+
+def generate_training_plots(parquet_path, region, fs, year, day,
+                            n_events=3, figure_dir="figures/ceed"):
+    """Load label parquet, join with original waveforms, plot training labels."""
+    print(f"\n=== Training label plots for {region} ===")
+
+    df = pd.read_parquet(parquet_path)
+    print(f"  {len(df)} pick rows, "
+          f"{df.groupby(['event_id','origin_index']).ngroups} detected origins")
+
+    bucket = BUCKETS[region]
+    print("  Loading waveforms from original parquet...")
+    wave_by_idx = load_day_waveforms(fs, bucket, year, day)
+    print(f"  {len(wave_by_idx)} waveforms loaded")
+
+    raw_pq = pq.read_table(
+        fs.open(f"{bucket}/waveform_parquet/{year}/{day}.parquet"),
+        columns=["event_id", "network", "station", "distance_km", "begin_time"],
+    ).to_pydict()
+    n_raw = len(raw_pq["event_id"])
+    meta_by_idx = {i: {k: raw_pq[k][i] for k in raw_pq} for i in range(n_raw)}
+
+    top_eids = (
+        df.groupby("event_id")["phase_index"].count()
+        .nlargest(n_events).index
+    )
+
+    transforms = default_train_transforms(crop_length=4096, enable_stacking=False)
+    fig_dir = os.path.join(figure_dir, region)
+
+    for src_eid in top_eids:
+        ev_df = df[df["event_id"] == src_eid]
+        linked = ev_df[ev_df["origin_index"].notna()]
+        valid_widxs = set(linked["waveform_index"].unique())
+
+        if not valid_widxs:
+            print(f"  [WARN] {src_eid}: no stations with model picks")
+            continue
+
+        recs_tagged = []
+        for widx in valid_widxs:
+            wave = wave_by_idx.get(widx)
+            meta = meta_by_idx.get(widx, {})
+            if wave is None:
+                continue
+            recs_tagged.append({
+                "event_id": src_eid,
+                "network": meta.get("network", ""),
+                "station": meta.get("station", ""),
+                "distance_km": meta.get("distance_km"),
+                "begin_time": meta.get("begin_time"),
+                "waveform": wave,
+                "_widx": widx,
+            })
+
+        if not recs_tagged:
+            print(f"  [WARN] {src_eid}: no waveforms found")
+            continue
+
+        recs_tagged.sort(key=lambda r: r.get("distance_km") or 0.0)
+        widx_to_sta = {r["_widx"]: i for i, r in enumerate(recs_tagged)}
+        recs = [{k: v for k, v in r.items() if k != "_widx"} for r in recs_tagged]
+
+        p_picks_all, s_picks_all = [], []
+        ps_centers_all, ps_intervals_all = [], []
+
+        for widx, sta_grp in linked.groupby("waveform_index"):
+            sta_idx = widx_to_sta.get(widx)
+            if sta_idx is None:
+                continue
+            for pidx in sta_grp.loc[sta_grp["phase_type"] == "P", "phase_index"].dropna():
+                p_picks_all.append((sta_idx, float(pidx)))
+            for sidx in sta_grp.loc[sta_grp["phase_type"] == "S", "phase_index"].dropna():
+                s_picks_all.append((sta_idx, float(sidx)))
+            for _, o_grp in sta_grp.groupby("origin_index"):
+                p_o = o_grp.loc[o_grp["phase_type"] == "P", "phase_index"].dropna()
+                s_o = o_grp.loc[o_grp["phase_type"] == "S", "phase_index"].dropna()
+                if p_o.empty or s_o.empty:
+                    continue
+                p_i, s_i = float(p_o.iloc[0]), float(s_o.iloc[0])
+                ps_centers_all.append((sta_idx, (p_i + s_i) / 2.0))
+                ps_intervals_all.append((sta_idx, s_i - p_i))
+
+        try:
+            sample = records_to_sample(recs)
+            if p_picks_all or s_picks_all:
+                sample.targets = [Target(
+                    p_picks=p_picks_all, s_picks=s_picks_all,
+                    ps_centers=ps_centers_all, ps_intervals=ps_intervals_all,
+                    event_id=src_eid,
+                )]
+            n_origins = int(linked["origin_index"].nunique())
+            print(f"  {src_eid}: {len(recs)} stations  origins={n_origins}  "
+                  f"waveform {sample.waveform.shape}")
+            plot_demo(sample, transforms, event_id=src_eid,
+                      output_dir=fig_dir, n_augmented=2)
+        except Exception as e:
+            print(f"  [WARN] {src_eid}: {e}")
+
+
+def main():
+    parser = argparse.ArgumentParser(description="CEED training label plots")
+    parser.add_argument("--parquet", required=True, help="Path to prediction parquet")
+    parser.add_argument("--region", default="NC", help="Region (NC or SC)")
+    parser.add_argument("--year", type=int, required=True)
+    parser.add_argument("--day", required=True, help="Day string (e.g. 001)")
+    parser.add_argument("--n-events", type=int, default=3)
+    parser.add_argument("--figure-dir", default="figures/ceed")
+    args = parser.parse_args()
+
+    fs = gcsfs.GCSFileSystem(**get_storage_options())
+    generate_training_plots(args.parquet, args.region, fs,
+                            args.year, args.day, args.n_events, args.figure_dir)
+
+
+if __name__ == "__main__":
+    main()