PhaseNet 0.2.4__tar.gz → 0.2.6.dev0__tar.gz

phasenet-0.2.6.dev0/PKG-INFO

@@ -0,0 +1,375 @@
+Metadata-Version: 2.4
+Name: phasenet
+Version: 0.2.6.dev0
+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.