Metadata-Version: 2.4
Name: phasenet
Version: 0.2.6.dev1
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.10
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.
Python 3.10 or later.

## 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.

The original PhaseNet-2018 weights, converted from TensorFlow, load through `load_phasenet2018_model` (and as `--model phasenet` in `predict_das.py`); every other model is trained here with `train.py`.

## Prediction

### Continuous mseed (seismic network)

```bash
# PhaseNet-2018 weights (the default), run as TF PhaseNet ran them: a 1 Hz highpass, accelerometers picked on velocity
python scripts/predict_mseed.py --data-list mseed_list.txt --inventory stations.xml --result-dir results/picks \
    --highpass 1.0 --integrate-acceleration

# a retrained model: its checkpoint brings its contract (normalization, windows); filtering is your choice
python scripts/predict_mseed.py --data-list mseed_list.txt --inventory stations.xml --result-dir results/picks \
    --checkpoint output/train_ceed/checkpoint.pth
```

`--data-list` has one station-instrument per line, its component files separated by `|`. `--inventory` is one
StationXML file for every station or a directory, whose `*.xml` files (flat or nested) are matched to each
station by `NET.STA`, then `STA`, then `NET_network`; it serves the `phase_amplitude` column and the channel
orientation, picking runs on counts, and a station without a match is still picked from its channel names. `--sampling-rate`
(100) is the working rate: the input is resampled to it and pick indices become times with it. `--workers`
(8) sets the reader processes; the reader, not the model, is the bottleneck.

Output: one CSV per station-instrument with `station_id, begin_time, phase_index, phase_time, phase_score, phase_type[, phase_amplitude]`.
This is the picker the QuakeFlow pipeline runs (`scripts/run_phasenet_pytorch.py` there).

### 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
```

The input contract is set here and saved in the checkpoint: `--norm-kind trailing --norm-window 1024 --norm-scope station --norm-stat std` are the defaults (see Design). Filtering is an augmentation, not a setting.

### 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 |
|--------|-------------|---------|
| `--norm-kind` | where the statistics come from: `window` (the whole input), `trailing` (before each sample), `centred` (around it) | trailing |
| `--norm-window N` | normalization window, samples (`trailing`, `centred`) | 1024 |
| `--norm-scope` | what shares one scale: `station`, `component`, `global` | station |
| `--norm-stat` | `std` or `l1` | std |
| `--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}/`.

## Design: one input contract per model

Three models share one code path; what differs between them is written down in each model's input
contract, not in the scripts that run them.

| | PhaseNet-2018 (TF port) | PhaseNet (retrain) | PhaseNet-DAS (retrain) |
|---|---|---|---|
| weights | converted TF checkpoint, no config of its own | `train.py` | `train.py` |
| backbone | `unet2018`, convolutions only | `unet2018` or `unet` | `unet` (attention over time and space) |
| contract | pinned in code, `PHASENET_2018_INPUT` | `model_config["input"]`, written by `train.py` | same |
| filtering | none of its own; the QuakeFlow pipeline applies a 1 Hz highpass and picks accelerometers on velocity, as TF did | trained with a random highpass (augmentation); at inference the caller's `--highpass` | trained with random highpass and common-mode removal (augmentations); at inference common-mode removal is on (`--no-common-mode`), the highpass off (`--highpass`) |
| normalization | TF `normalize_long` (`tf_interp`, 3000 samples) | trailing 1024, per station, std | trailing 1024, per channel (a channel is a one-component station) |
| inference cut | whole trace | whole trace for `unet2018`, else windows of `--nt`, half overlap, overlap-discard | windows of `--nt` x `--nx`, half overlap, overlap-discard |
| trained on | TF, 30 s windows | CEED event windows (`--dataset-type ceed`) | DAS HDF5 patches (`--dataset-type das`) |
| run on | continuous mseed (`scripts/predict_mseed.py`) | continuous mseed or CEED days (`scripts/predict_ceed.py`) | DAS HDF5 events (`scripts/predict_das.py`) |

### The contract

`InputSpec` (`phasenet/data/preprocess.py`) is saved in a checkpoint's `model_config["input"]` and
rebuilt by every loader; a checkpoint without one is refused. Its fields:

- `normalize`: a `NormalizeSpec` (`phasenet/normalize.py`): `kind` (`window`: the classic per-window
  statistics; `trailing`: the samples before each sample, causal, the default for the retrains;
  `centred`: the samples around each sample; `tf_interp`: the 2018 port's TF `normalize_long`),
  `scope` (`station`: one scale shared by a station's components, which keeps their relative
  amplitudes; `component`; `global`), `stat` (`std` or `l1`), `window` in samples, `compress`
  (`None` or `log1p`).
- `window_nt`: the longest input, in samples, the backbone takes at once. `None` means any length,
  which is true of a convolution-only backbone, so inference runs on the whole trace.
- `window_nx`: stations per patch, `None` for all at once. A model that couples channels (DAS) sets
  it and gets a taper along that axis.

The model knows samples, not a rate. The caller owns the working rate (`predict_mseed.py
--sampling-rate`, 100 Hz by default): it resamples the input to it, applies any Hz corner at it, and
turns pick indices into times with it. The same weights therefore apply to a network at 100 Hz or a
borehole at 500 Hz. Filtering is not in the contract: training applies a highpass (and, for DAS,
common-mode removal) at random as an augmentation, so the model works with or without it and the
choice belongs to whoever runs it.

### Preprocessing, shared with the dataset

Three functions live in `phasenet/data/preprocess.py` with identical copies in
`QuakeFlow/datasets/common.py`, asserted identical by `datasets/tests/test_preprocess.py`, so the
training data and the picker's input are prepared by the same code, in the same order: merge,
resample, orient. `merge_stream` joins a station's
fragments: overlaps resolved, every contiguous piece demeaned on its own, gaps filled with zeros that
continue each piece at its own level. `orient_from_inventory` recovers Z-up, N and E from
each channel's inventory azimuth and dip, whatever the channel is named, passing the usual E/N/Z
stations through untouched and recording `z_sign` where a vertical had to be negated; the picker
applies it when it has an inventory. `resample_trace` brings a trace to the working rate: polyphase
interpolation upward (an FIR lowpass at the old Nyquist, flat over the recorded band and free of edge
artifacts; obspy's FFT resampling puts a spike on the first sample, and its default Hann window is a
lowpass that removed 10-20 Hz from 40 Hz stations), a causal 4-pole Butterworth at 0.4 x the target
rate then decimation downward (causal, so nothing precedes an onset). On the demo the polyphase and
the FFT kernel give the same analyst-pick recall on every instrument class and differ in 0.06 % of the
picks. `QuakeFlow/datasets/common.py` holds an identical copy and
`datasets/tests/test_preprocess.py` asserts the two agree sample by sample.

`highpass` is the detrend, when one is wanted: TF PhaseNet's spline detrend gave the same picks at six
times the cost and follows the long-period motion of large earthquakes; the highpass is causal and
streams. `remove_common_mode` subtracts the median over DAS channels. Both are augmentations in
training (`RandomHighpass`, 0.5-2 Hz half the time; `RandomCommonMode`, half the time) and options
at inference: for DAS the common-mode removal is on by default (cheap), the highpass off (costly).

### What miniSEED data does, and what happens to it

A station-day is rarely three clean channels of one continuous record. Every case below is handled by
the same code on both sides (`merge_stream`, `orient_from_inventory`, `resample_trace`, the
normalization), so a training record and a picker input that came from the same kind of file look the
same to the model. The picker reports a station-day it cannot read as an error and moves on; the
dataset skips the station.

| the data | what is done | where |
|---|---|---|
| Overlapping fragments (re-sent or duplicated records) | obspy merge method 1: the later fragment wins, sample for sample, no blending. Duplicates agree exactly and leave no mark. | `merge_stream` |
| Timing tears (gaps up to 0.1 s between records, or at the join of two day files whose sample grids differ by a fraction of a sample) | bridged by linear interpolation: a tear is not missing data, and a zero there would be a spike | `merge_stream` |
| Gaps longer than 0.1 s (outages) | every contiguous piece is demeaned on its own, then the gap is filled with zeros, which continue each piece at its own level. The zeros are data from then on: the normalization does not infer gaps from values, a window straddling a gap is scaled by a std the zeros lower, and a filter applied afterwards rings a little into the gap. Nothing is tapered: real samples are never altered. | `merge_stream`; the normalization |
| A fragment with a wrong timestamp (months away) | fragments farther than two days from the station-day's median start are dropped before the merge | picker |
| Fragments of one channel at different sampling rates | obspy cannot merge them: the station-day fails with an error (picker) or the station is skipped (dataset) | `merge_stream` |
| A channel recorded below or above the working rate (BH at 40 Hz, borehole at 500 Hz) | resampled to the working rate, the same kernel on both sides | `resample_trace` |
| A missing channel (one or two of three) | the present ones are placed by name, the missing component is zeros. A component without variation has scale 0, read as 1, so it stays zero; with `station` scope the live components come out sqrt(3/live) larger, the 2018 model's `tf_interp` applies TF's 3/live gain | picker, dataset cut; the normalization |
| A dead channel (present, flat) | the same as missing: the dataset leaves a flat component zero and unlisted, the normalization gives it scale 0 | dataset cut; the normalization |
| Channels named 1/2/3 or X/Y/Z, or mislabeled E/N/Z (Z with dip +90, horizontals swapped or rotated) | with an inventory, Z-up, N and E are recovered from each channel's azimuth and dip whatever its name; `z_sign` records a negated vertical so polarity labels follow. Without an inventory the picker places channels by a name table | `orient_from_inventory`; picker |
| Nearly parallel channel vectors (a metadata error: two channels with the same azimuth) | no rotation can separate them, so the group is kept as named instead of averaging or amplifying the channels | `orient_from_inventory` |
| A single channel | at most one component is recovered (Z, or a horizontal on its axis); an oblique single horizontal gives nothing | `orient_from_inventory` |
| Several location codes or bands for one station | one station-day per location and instrument in the picker's file list; the dataset keeps the catalog's location and one band per record | picker, dataset cut |
| A record shorter than the model's window, or not a multiple of 16 samples | zero padding at the end, the zeros being data as above | `RandomCrop`, `predict_2d` |

Consistency is asserted, not assumed: `datasets/tests/test_preprocess.py` requires the three
functions to have identical source in both copies and `resample_trace` to give identical samples.

### Normalization

Training normalizes exactly as inference does, so the two cannot differ. For the local kinds
(`trailing`, `centred`) inference normalizes the whole record and cuts it afterwards; training
normalizes the whole loaded waveform and the random crop is a view of the result. For the
`window` kind the statistics are defined on the window, so each inference patch is normalized on
its own and the training crop comes before the normalization. `tests/test_normalize.py` asserts
that a training crop equals what `predict_2d` computes, for every kind.

`trailing`, the default for the retrains, scales every sample by the statistics of the `window`
samples ending at it (the first window's for the first window). It is causal by design: a centred
window lets a large event inflate the scale for half a window *before* it happens, so a small
event shortly before a larger one loses its amplitude; with a trailing window it keeps it
(measured on the Ridgecrest demo: 90 % of its isolated amplitude against 28 % under the centred TF
window). The statistics come from cumulative sums over every sample; a zero is a sample like any
other, and missing data is recognised by its variation, not its value: a window, station or record
with none has scale 0, read as 1, so a dead component comes out as zeros. The same `Normalizer` runs on a training waveform, a
whole day, or a stream with carried state. The model never normalizes in `forward`: the data
transforms and `predict_2d` call `model.normalizer`.

In the training pipelines, augmentations that shape amplitude or spectrum (highpass, bandpass,
common-mode removal, event stacking) run before the normalization; calibrated noise, flips,
masking and the median filter run after it, on normalized data, which is what the model meets at
inference.

The 2018 port keeps TF's `normalize_long` as the `tf_interp` branch: a centred window measured
every `window // 2` samples and interpolated, with TF's 3/live gain for stations missing
components. Those weights were validated against TF with exactly this; an exact centred window is
not equivalent (6.8 % of the demo picks change), so it is not a shortcut.

### Inference

`predict_2d(model, waveform, device)` is the one entry point. It normalizes the whole `(nch, nx,
nt)` input, then cuts by the contract: the whole trace as one window when `window_nt` is `None`
(patches only add seams, and the trace is padded to a multiple of 16 rather than the U-Net's
full 1024, since a wall of zeros inside the receptive field erases picks at the day boundary),
otherwise windows of `window_nt` overlapping by half, each sample taken from the window whose centre
is nearest (overlap-discard, so nothing is predicted within a quarter window of an edge), and along
stations `window_nx` patches assigned the same way when set. The three predict scripts differ only in how
they read data.

### Training

`train.py` builds the contract from `--norm-kind`, `--norm-window`, `--norm-scope`, `--norm-stat`,
`--norm-compress`, `--nt` and `--nx`, gives it to the data transforms and writes it into the
checkpoint. The default transforms run, in order: `RandomHighpass` (and `RandomCommonMode` for DAS),
a whole-sample `Normalize` that only equalizes samples for the augmentations,
event stacking, `InputNormalize` (before `RandomCrop` for the local kinds, after it for `window`),
then noise, flips and masking.

### Verification

Ten percent of CEED training samples are event-free: the 40.96 s of real noise before each record's
first P, through the same pipeline. `python tests/test_normalize.py` pins the contract: the `tf_interp` branch against a numpy transcription
of TF's `normalize_long`, the window, trailing and centred kinds against a sample-by-sample brute
force for every scope and statistic, the training pipelines against `predict_2d`, streamed chunks
against the whole trace, gaps and dead components, and the spec's round trip through `model_config`. On the
QuakeFlow Ridgecrest demo (45 station-days) the port reproduces TF PhaseNet: identical picks on the
100 Hz stations, analyst-pick recall 0.861 for both, the 40 Hz stations differing only by the
resampling kernel. `scripts/plot_normalization.py` shows what each candidate normalization does to
a station-day.

### Open issues

- **Station grouping.** With `--nx` above 1 the CEED loader builds one sample per event from its `nx`
  nearest stations, and the `unet` backbone attends across them. `scripts/predict_mseed.py` picks one
  station-day at a time, so it hands such a model a single station, a situation it was never trained on
  and in which its cross-station attention has nothing to attend to. `scripts/predict_ceed.py` builds the
  same event arrays as training and is not affected. Until resolved: train the model that picks
  continuous data with `--nx 1`, and `predict_mseed.py` should refuse a checkpoint whose `window_nx`
  is greater than 1 instead of feeding it one station. The full fix is a multi-station reader that
  assembles concurrent station-days into a `(3, nx, nt)` array with a spatial station order used in
  training as well (the training sort by distance to the event does not exist at inference), plus
  chunked inference along time to bound memory.
- **Crop edges.** Training crops are 4,096 samples while the `unet2018` receptive field is 3,331, so
  most training samples sit within half a receptive field of a zero-padded crop edge, which a whole
  day never shows the model. Longer crops (the record has 12,288 samples) or a loss mask near edges.
- **Trace-start rule in training.** A crop that overlaps the first 1,024 samples of its record is
  scaled by the first window's statistics; inference meets that once a day. Start crops at or after
  the normalization window.
- **DAS archives.** The training loader differentiates OptaSense continuous files to strain rate;
  the event files used for training and prediction are stored as strain rate already. A model run on
  continuous archives needs that step in the shared preprocessing.
- **Gap edges.** A gap is zeros after each piece's own demean (see the miniSEED table), and nothing
  is done about picks at its edges. What the model does there is a training question: a gap
  augmentation for CEED (DAS has `Masking`) and a pick filter near gap edges.
- **Shared preprocessing is a copy.** `QuakeFlow/datasets/common.py` holds identical copies of
  `merge_stream`, `orient_from_inventory` and `resample_trace`, kept equal by a test. One source,
  importable without torch (move `InputSpec` out of `preprocess.py`), would make the guarantee structural.
