phasenet 0.2.6.dev2__tar.gz → 0.2.6.dev3__tar.gz

{phasenet-0.2.6.dev2/phasenet.egg-info → phasenet-0.2.6.dev3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: phasenet
-Version: 0.2.6.dev2
+Version: 0.2.6.dev3
 Summary: A PyTorch implementation of PhaseNet for seismic and DAS phase picking
 Author: Weiqiang Zhu
 License-Expression: MIT

{phasenet-0.2.6.dev2 → phasenet-0.2.6.dev3}/phasenet/utils/inference.py

@@ -85,7 +85,8 @@ def predict_2d(model: torch.nn.Module,
         max_batch_size: Max patches per forward pass (memory safety valve).
 
     Returns:
-        Phase scores of shape (3, nx, nt) — [noise, P, S].
+        Phase scores of shape (3, nx, nt) — [noise, P, S]; with a polarity head (4, nx, nt), the fourth row the
+        probability that the first motion is up.
         If model outputs event_center/event_time, returns
         (scores, event_center, event_time) at their native resolution.
     """
@@ -167,6 +168,10 @@ def predict_2d(model: torch.nn.Module,
             with torch.autocast(device_type=dev.type, dtype=torch.bfloat16, enabled=bf16):
                 out = model({"data": batch})
             probs = torch.softmax(out["phase"].float(), dim=1)
+            if "polarity" in out:  # P(up), carried as a fourth row through the same overlap-discard
+                probs = torch.cat([probs, torch.sigmoid(out["polarity"].float())], dim=1)
+            if scores_sum.shape[0] != probs.shape[1]:
+                scores_sum = torch.zeros(probs.shape[1], nx, nt, dtype=torch.float32, device=acc_device)
             if not use_gpu:
                 probs = probs.cpu()
 
@@ -328,6 +333,8 @@ def detect_picks(scores: np.ndarray,
                     "phase_score": round(prob, 4),
                     "phase_type":  phase,
                     "dt_s":        dt_s,
+                    **({"phase_polarity": round(scores[3, ch_idx, tidx].item(), 4)}
+                       if phase == "P" and scores.shape[0] > 3 else {}),  # P(up) at the P pick
                 })
 
     return picks

{phasenet-0.2.6.dev2 → phasenet-0.2.6.dev3/phasenet.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: phasenet
-Version: 0.2.6.dev2
+Version: 0.2.6.dev3
 Summary: A PyTorch implementation of PhaseNet for seismic and DAS phase picking
 Author: Weiqiang Zhu
 License-Expression: MIT

{phasenet-0.2.6.dev2 → phasenet-0.2.6.dev3}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "phasenet"
-version = "0.2.6.dev2"
+version = "0.2.6.dev3"
 description = "A PyTorch implementation of PhaseNet for seismic and DAS phase picking"
 readme = "README.md"
 license = "MIT"

{phasenet-0.2.6.dev2 → phasenet-0.2.6.dev3}/scripts/predict_mseed.py

@@ -242,6 +242,8 @@ def run_inference(model, waveform, velocity, station_id, begin_iso, dt,
     waveform_3d = waveform[:, np.newaxis, :]  # (3, 1, nt)
     # whole trace or overlapping windows: predict_2d follows model.input
     scores = predict_2d(model, waveform_3d, device=DEVICE, max_batch_size=batch_size)
+    if isinstance(scores, tuple):  # a model with event heads also returns its event detections
+        scores = scores[0]
 
     picks = detect_picks(
         scores,

{phasenet-0.2.6.dev2 → phasenet-0.2.6.dev3}/tests/test_models.py

@@ -31,6 +31,20 @@ def test_window_nx_1():
     assert np.abs(single - patches).max() < 1e-5
 
 
+def test_polarity_at_picks():
+    """a polarity model's scores carry P(up) as a fourth row, equal to the head's sigmoid, and its P picks carry it"""
+    from phasenet.utils.inference import detect_picks
+    torch.manual_seed(0)
+    m = phasenet.models.build_model("phasenet_plus", input=replace(SPEC, window_nx=1, window_nt=None).asdict()).eval()
+    x = np.random.default_rng(0).standard_normal((3, 2, 2048)).astype(np.float32)
+    scores = predict_2d(m, x, max_batch_size=8)[0]
+    with torch.no_grad():
+        up = torch.sigmoid(m({"data": m.normalizer(torch.from_numpy(x)[None])})["polarity"])[0, 0].numpy()
+    assert scores.shape == (4, 2, 2048) and np.abs(scores[3] - up).max() < 1e-5
+    picks = detect_picks(scores, ["a", "b"], {}, None, vmin=0.0)
+    assert all(("phase_polarity" in p) == (p["phase_type"] == "P") for p in picks) and picks
+
+
 def test_time_downsample():
     """every family declares the length multiple a whole trace is padded to, and a trace of any length comes back
     scored at its own length (a strided stage added without updating time_downsample breaks this)"""