specmod 0.2.0__py3-none-any.whl

specmod/transforms/quadratic.py

@@ -0,0 +1,221 @@
+"""Quadratic multitaper estimation (Prieto et al. 2007).
+
+An ordinary multitaper estimate averages over ``K`` tapers, each of which
+smooths the spectrum across the inner band ``[-W, W]``. Where the true spectrum
+is curved, that smoothing does not average out: it pulls peaks down and fills
+troughs in, by an amount proportional to the second derivative. The quadratic
+inverse method estimates that curvature and subtracts the bias it causes.
+
+Measured on a pure line, where the true Fourier amplitude ``A*T/2`` is known
+exactly: the ordinary estimate recovers 0.87 of the peak, this one 1.02. It
+also separates two lines just past the resolution limit with about 1.4x the
+peak-to-trough contrast. On white noise — no curvature to correct — it agrees
+with the ordinary estimate to within a few percent, which is the control that
+says it is not simply sharpening everything.
+
+.. note::
+
+   **On a corner frequency it is only a marginal improvement.** Over 40
+   realisations of a true 4 Hz Brune corner: ordinary multitaper biases -0.087
+   and this +0.077, with the same scatter (IQR 0.133 against 0.139), so the two
+   are indistinguishable once both are counted. A lightly-tapered FFT beats
+   both (bias -0.020, IQR 0.055) and is around 100x faster. Reach for this
+   estimator when the feature of interest is a peak or a line, not to squeeze a
+   corner.
+
+.. warning::
+
+   ``qiinv`` builds cross-spectra from ``wt * yk`` and never divides by
+   ``sum(w**2)``, so **adaptive weights must be renormalised before they are
+   handed to it** — see :meth:`QuadraticMultitaperEstimator.estimate`. Omitting
+   that scales the estimate down by ``sum(w**2)/K`` wherever the weights bite,
+   which on a Brune spectrum is 0.57 in the tail and reads convincingly like a
+   curvature artefact. The tell is that it vanishes under flat weighting.
+
+   Upstream omits this and leans on a global variance rescaling that cannot
+   compensate, because the deficit varies with frequency — so results taken
+   from ``multitaper``'s own ``MTSpec.qiinv`` carry it.
+
+The cost is a per-frequency least-squares solve, so this is roughly two orders
+of magnitude slower than :class:`~specmod.transforms.multitaper.MultitaperEstimator`.
+It is not the estimator to reach for on a whole catalogue.
+
+Implementation
+--------------
+The numerical core is vendored from Prieto's ``multitaper`` package rather than
+reimplemented — see :mod:`specmod._vendor.qiinv` for the licence, the changes
+made, and why. This module owns the part that faces SpecMod: DPSS tapers,
+Thomson weights, and the units contract that every estimator here satisfies.
+
+References
+----------
+Prieto, G.A., Parker, R.L., Thomson, D.J., Vernon, F.L., Graham, R.L. (2007).
+Reducing the bias of multitaper spectrum estimates.
+*Geophysical Journal International* 171(3), 1269-1281.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+import numpy as np
+from numpy.typing import ArrayLike, NDArray
+from scipy.signal.windows import dpss
+
+from .._vendor.qiinv import qiinv
+from ..core.spectrum import Spectrum
+from ..core.units import AmplitudeKind, Motion
+from .base import build_spectrum, prepare_record
+from .multitaper import _adaptive_weights, center_on_energy_centroid
+
+__all__ = ["QuadraticMultitaperEstimator"]
+
+
+@dataclass(frozen=True)
+class QuadraticMultitaperEstimator:
+    """Curvature-corrected multitaper spectrum estimate.
+
+    Parameters
+    ----------
+    time_bandwidth, n_tapers
+        As :class:`~specmod.transforms.multitaper.MultitaperEstimator`. The
+        correction scales with ``W = NW/N``, so a larger time-bandwidth product
+        means more smoothing to undo and a correspondingly larger correction.
+    adaptive
+        Whether the eigencoefficients fed to the curvature fit are weighted by
+        Thomson's adaptive scheme. Applies to the input weights only; the
+        quadratic step itself is unweighted.
+    normalize_to_variance
+        Rescale so the spectrum integrates to the record variance, as ``mtspec``
+        and Prieto's package do. Off by default, for the reason given in
+        :class:`~specmod.transforms.multitaper.MultitaperEstimator`.
+
+        Note that Prieto's ``MTSpec.qiinv()`` applies this unconditionally, and
+        to an already-renormalised input, so reproducing that path needs it on.
+    center
+        Circularly shift the record so its energy centroid sits mid-window.
+        Same rationale and same wrap check as the ordinary estimator.
+    """
+
+    time_bandwidth: float = 3.0
+    n_tapers: int = 5
+    adaptive: bool = True
+    center: bool = False
+    center_edge_tolerance: float = 0.05
+    normalize_to_variance: bool = False
+    drop_dc: bool = True
+    name: str = "quadratic"
+
+    def __post_init__(self) -> None:
+        if self.time_bandwidth <= 0:
+            raise ValueError(
+                f"time_bandwidth must be positive, got {self.time_bandwidth}"
+            )
+        limit = int(2 * self.time_bandwidth - 1)
+        if self.n_tapers < 1:
+            raise ValueError(f"n_tapers must be at least 1, got {self.n_tapers}")
+        if self.n_tapers > limit:
+            raise ValueError(
+                f"n_tapers={self.n_tapers} exceeds 2*NW-1={limit} for "
+                f"time_bandwidth={self.time_bandwidth}. Tapers beyond that are "
+                f"poorly concentrated and add leakage rather than reducing "
+                f"variance; raise time_bandwidth or lower n_tapers."
+            )
+        if self.n_tapers < 2:
+            raise ValueError(
+                "The quadratic estimate fits three basis functions to the "
+                "cross-spectra of K tapers, so it needs at least 2 tapers "
+                "(K**2 = 4 equations) to be determined at all. Use "
+                "MultitaperEstimator for a single-taper estimate."
+            )
+
+    def estimate(
+        self,
+        data: ArrayLike,
+        dt: float,
+        *,
+        motion: Motion | str = Motion.VELOCITY,
+        meta: dict[str, Any] | None = None,
+    ) -> Spectrum:
+        x, n, duration = prepare_record(data, dt)
+        if self.center:
+            x = center_on_energy_centroid(x, self.center_edge_tolerance)
+
+        tapers, eigenvalues = dpss(
+            n, self.time_bandwidth, self.n_tapers, sym=False, return_ratios=True
+        )
+
+        # qiinv works on the full two-sided transform: it forms cross-spectra
+        # between tapers and needs the negative frequencies to do it.
+        yk = np.fft.fft(tapers * x, axis=-1)
+        eigenspectra = np.abs(yk) ** 2
+
+        if self.adaptive:
+            weights = _adaptive_weights(eigenspectra, eigenvalues, n)
+            # qiinv forms cross-spectra of ``wt * yk`` and never divides by
+            # ``sum(w**2)``, so its diagonal averages to ``(1/K) sum(w**2 |y|**2)``
+            # where the adaptive estimate is ``sum(w**2 |y|**2) / sum(w**2)``.
+            # Renormalising per frequency so ``sum(w**2) = K`` makes the two
+            # agree. Without it the estimate is scaled down by ``sum(w**2)/K``
+            # wherever the weights bite — 0.57 in the high-frequency tail of a
+            # Brune spectrum, which reads exactly like a curvature artefact and
+            # is not one. Upstream leaves this out and relies on its global
+            # variance renormalisation to hide it, which cannot: the deficit is
+            # frequency-dependent.
+            scale = self.n_tapers / np.maximum((weights**2).sum(axis=0), 1e-300)
+            weights = weights * np.sqrt(scale)
+        else:
+            weights = np.ones_like(eigenspectra)
+
+        qispec, _slope, curvature = qiinv(
+            yk.T, weights.T, tapers.T, eigenvalues, self.time_bandwidth
+        )
+
+        # qispec is in |FFT|**2; the same dt that makes an ordinary
+        # eigenspectrum a PSD applies here, and for the same reason.
+        psd = qispec * dt
+
+        # The correction is a subtraction and can in principle overshoot into
+        # negative power where curvature is large and poorly determined. The
+        # damping inside qiinv makes that rare, but a negative PSD is not a
+        # spectrum, so clip and say so rather than emit sqrt of a negative.
+        negative = int((psd < 0).sum())
+        if negative:
+            psd = np.maximum(psd, 0.0)
+
+        freq: NDArray[np.float64] = np.fft.rfftfreq(n, d=dt).astype(np.float64)
+        psd = psd[: freq.size].copy()
+        psd[1:] *= 2.0  # fold negative frequencies
+        if n % 2 == 0:
+            psd[-1] /= 2.0
+
+        if self.normalize_to_variance:
+            df = float(freq[1] - freq[0])
+            total = float(psd.sum() * df)
+            if total > 0:
+                psd = psd * (float(x.var()) / total)
+
+        if self.drop_dc:
+            freq, psd = freq[1:], psd[1:]
+
+        spectrum = build_spectrum(
+            freq,
+            psd,
+            kind=AmplitudeKind.PSD,
+            motion=motion,
+            duration=duration,
+            sampling_rate=1.0 / dt,
+            meta={
+                **(meta or {}),
+                "time_bandwidth": self.time_bandwidth,
+                "n_tapers": self.n_tapers,
+                "adaptive": self.adaptive,
+                "centered": self.center,
+                "normalize_to_variance": self.normalize_to_variance,
+                "curvature_rms": float(np.sqrt(np.mean(curvature**2))),
+                "clipped_bins": negative,
+            },
+            estimator=self.name,
+        )
+        return spectrum.to_kind(AmplitudeKind.FAS)

specmod/utils.py

@@ -0,0 +1,305 @@
+"""Reading picks and catalogues, and looking at waveforms.
+
+What is left after the migration. The rotation machinery that used to live here
+— ``find_rotation_angle``, ``find_rotation_angle_v2``, ``rotate``,
+``rotate_noise_full``, ``get_centroid_freq`` and ``non_lin_boost_noise_func`` —
+is now :mod:`specmod.core.noise`, where the two methods are registered models
+rather than an integer flag. The SAC-discovery scaffolding that used to sit
+between them (``DataSet`` and its channel-ranking helpers, ``cps``,
+``path_to_utc``) had no callers anywhere and was hardcoded to another study's
+station codes; it is replaced by the acquisition layer in REFACTOR_PLAN §5.2.
+"""
+
+from __future__ import annotations
+
+import contextlib
+import logging
+import warnings
+from typing import TYPE_CHECKING, Any
+
+import matplotlib
+import matplotlib.pyplot as plt
+import numpy as np
+import obspy
+import pandas as pd
+from matplotlib.dates import num2date
+from obspy.core.event import Catalog, Event, WaveformStreamID
+from obspy.core.event import Pick as ObsPyPick
+
+from specmod.picks import PickSet, SnufflerReader, from_catalog, select_event
+
+#: See the note in `specmod.fitting.event`: diagnostics go here, and nothing in
+#: this package configures logging on a caller's behalf.
+logger = logging.getLogger(__name__)
+
+if TYPE_CHECKING:  # pragma: no cover
+    from collections.abc import Sequence
+    from datetime import datetime
+    from os import PathLike
+
+    from obspy import Stream, UTCDateTime
+
+# Type 42 embeds fonts as editable text rather than outlines, so a figure can
+# be relabelled in a vector editor after the fact. Set at import because it is
+# a global, and this module is the one that draws.
+matplotlib.rcParams["ps.fonttype"] = 42
+matplotlib.rcParams["pdf.fonttype"] = 42
+matplotlib.rcParams["svg.fonttype"] = "none"
+
+
+def _to_datetime(value: Any) -> Any:
+    """``matplotlib.dates.num2date`` carries no annotations upstream.
+
+    Wrapped once rather than suppressed at each of its four call sites, so
+    there is one place to delete when matplotlib annotates it.
+    """
+    return num2date(value)  # type: ignore[no-untyped-call]
+
+
+def read_pyrocko_picks(path: str | PathLike[str]) -> PickSet:
+    """Snuffler marker file to a :class:`~specmod.picks.PickSet`.
+
+    Thin wrapper over :class:`specmod.picks.SnufflerReader`, which is where the
+    format lives.
+    """
+    return SnufflerReader().read(path)[0]
+
+
+def read_pyrocko(path: str | PathLike[str]) -> dict[str, dict[str, UTCDateTime]]:
+    """Snuffler marker file to ``{"NET.STA.LOC": {"P": UTCDateTime, ...}}``.
+
+    Keyed **per sensor**, not per channel. An arrival is an observation of one
+    sensor, and the component it was picked on is incidental — on the shipped
+    PNR file every station has P on ``HHZ`` and S on ``HHN``, so keying by full
+    SEED id would leave the horizontals with no S at all.
+
+    The location code stays in the key, because that is what distinguishes two
+    sensors at one site. A surface and a borehole instrument share a station
+    name and see genuinely different arrivals; collapsing them would silently
+    give one the other's picks.
+
+    A flat mapping cannot express a partial identity or a second event; prefer
+    :func:`read_pyrocko_picks` with :func:`specmod.picks.resolve`.
+    """
+    return read_pyrocko_picks(path).mapping()
+
+
+def read_quakeml_picks(
+    source: str | PathLike[str] | Catalog,
+) -> dict[str, dict[str, UTCDateTime]]:
+    """QuakeML picks to ``{"NET.STA.LOC": {"P": UTCDateTime, ...}}``.
+
+    The same shape :func:`read_pyrocko` returns, so the two are
+    interchangeable at the call site. ``source`` is a path or an
+    ``obspy.core.event.Catalog``.
+
+    Phase hints are normalised to ``P`` or ``S`` on their first letter, which
+    folds ``Pg``/``Pn``/``Pb`` and their S counterparts together. That matches
+    what this pipeline does with them — it wants a direct arrival, and the
+    branch distinction is not something the windowing uses.
+
+    Picks with an ``evaluation_status`` of ``rejected`` are dropped. Anything
+    else — reviewed, preliminary, unset — is kept, matching how the Snuffler
+    reader treats weights.
+
+    A multi-event source raises rather than merging. A flat mapping cannot
+    express a partial identity or a second event; prefer
+    :func:`specmod.picks.read` with :func:`specmod.picks.resolve`.
+    """
+    catalog = source if isinstance(source, Catalog) else obspy.read_events(str(source))
+    return select_event(from_catalog(catalog)).mapping()
+
+
+def picks_to_quakeml(
+    picks: dict[str, dict[str, UTCDateTime]],
+    *,
+    event_id: str | None = None,
+) -> Catalog:
+    """The inverse: a pick mapping as an ``obspy`` ``Catalog``.
+
+    Used to convert a Snuffler marker file to the standard format once, rather
+    than teaching every consumer both. The location code ``--`` is written back
+    as an empty string, which is what StationXML and miniSEED use.
+    """
+    event = Event(resource_id=event_id) if event_id else Event()
+    for key, phases in sorted(picks.items()):
+        net, sta, loc = key.split(".")
+        for phase, time in sorted(phases.items()):
+            event.picks.append(
+                ObsPyPick(
+                    time=time,
+                    phase_hint=phase,
+                    waveform_id=WaveformStreamID(
+                        network_code=net,
+                        station_code=sta,
+                        location_code="" if loc == "--" else loc,
+                    ),
+                )
+            )
+    return Catalog(events=[event])
+
+
+def plot_traces(
+    st: Stream,
+    *,
+    plot_theoreticals: bool = False,
+    plot_windows: bool = False,
+    conv: float | None = 1e-9,
+    bft: float = 1,
+    aftt: float = 60,
+    sig: Stream | None = None,
+    noise: Stream | None = None,
+    save: str | None = None,
+) -> None:
+    # `plot_windows` reads `sig[i]` and `noise[i]`. Without both it used to
+    # raise `TypeError: 'NoneType' object is not subscriptable` from inside the
+    # loop, naming neither the argument that was missing nor the option that
+    # required it.
+    if plot_windows and (sig is None or noise is None):
+        raise ValueError("plot_windows=True requires both sig and noise")
+
+    sharey = False
+    stc = stream_distance_sort(st)
+    if conv is None:
+        conv = 1
+        stc.normalize()
+        sharey = True
+
+    if sig is not None and noise is not None:
+        sig = stream_distance_sort(sig)
+        noise = stream_distance_sort(noise)
+
+    fig, drawn = plt.subplots(
+        len(stc), 1, sharex=True, sharey=sharey, figsize=(14, len(stc) * 3)
+    )
+    # `subplots` returns a bare `Axes` for a single row and an array otherwise.
+    # Kept as a separate name from what `subplots` returned: rebinding one
+    # variable to both shapes is what made the function unannotatable.
+    ax: Sequence[Any] = drawn.flatten() if len(stc) > 1 else [drawn]
+
+    for i, tr in enumerate(stc):
+        if plot_windows:
+            assert sig is not None  # both guarded above
+            assert noise is not None
+            # get window start and end times
+            sts = sig[i].stats["wstart"], sig[i].stats["wend"]
+            nts = noise[i].stats["wstart"], noise[i].stats["wend"]
+
+            tr.trim(nts[0] - bft, sts[1] + aftt)
+
+            for signal_edge, noise_edge in zip(sts, nts, strict=True):
+                ts: datetime = _to_datetime(signal_edge.matplotlib_date)
+                tn: datetime = _to_datetime(noise_edge.matplotlib_date)
+                ax[i].vlines(ts, tr.data.min() * conv, tr.data.max() * conv, color="k")
+                ax[i].vlines(
+                    tn, tr.data.min() * conv, tr.data.max() * conv, color="blue"
+                )
+
+        if plot_theoreticals:
+            try:
+                if not plot_windows:
+                    # Trim to a window around the P arrival, on the assumption
+                    # that the pick is trustworthy.
+                    tr.trim(tr.stats["p_time"] - bft, tr.stats["p_time"] + aftt)
+
+                p = _to_datetime(tr.stats["p_time"].matplotlib_date)
+                s = _to_datetime(tr.stats["s_time"].matplotlib_date)
+                ax[i].vlines(
+                    p,
+                    tr.data.min() * conv / 1.5,
+                    tr.data.max() * conv / 1.5,
+                    linestyles="dashed",
+                    color="blue",
+                    label="Pg",
+                )
+                ax[i].vlines(
+                    s,
+                    tr.data.min() * conv / 1.5,
+                    tr.data.max() * conv / 1.5,
+                    linestyles="dashed",
+                    color="red",
+                    label="Sg",
+                )
+            except KeyError:
+                pass
+
+        # `sample` rather than `i`: the comprehension used to reuse the name of
+        # the enclosing loop variable. Comprehensions have their own scope, so
+        # it was never a bug — but it reads like one, and would become one the
+        # moment this is turned back into a statement.
+        time = _to_datetime(
+            np.array(
+                [
+                    (
+                        tr.stats.starttime + (tr.stats.delta * (sample + 1))
+                    ).matplotlib_date
+                    for sample in range(len(tr.data))
+                ]
+            )
+        )
+
+        ax[i].plot(time, tr.data * conv, color="grey", label=tr.id, zorder=1)
+        # A trace with no geometry set still plots; it just goes untitled.
+        with contextlib.suppress(KeyError):
+            ax[i].set_title(
+                f"Repi: {tr.stats['repi']:.2f} km, Rhyp: {tr.stats['rhyp']:.2f} km"
+            )
+        ax[i].legend()
+    ax[-1].set_xlabel("Time (UTC)")
+    fig.suptitle(str(st[0].stats.otime))
+    fig.tight_layout()
+    if save is not None:
+        assert type(save) is str
+        fig.savefig(save)
+        fig.clear()
+        plt.close(fig)
+        logger.debug("closed the time-domain figure after saving to %s", save)
+
+
+def stream_distance_sort(st: Stream, dist_met: str = "repi") -> Stream:
+    """A copy of ``st`` ordered by distance. Not in place.
+
+    Returns the stream unsorted, with a warning, when the traces carry no
+    distance — which is why the copy is taken on the way out rather than only
+    on the sorted path.
+    """
+    try:
+        st = obspy.Stream(sorted(st, key=lambda x: x.stats[dist_met]))
+    except KeyError:
+        warnings.warn(
+            f"No {dist_met!r} on these traces, so the stream is returned "
+            "unsorted. Set distances with `preprocess.set_stream_distance`.",
+            stacklevel=2,
+        )
+
+    return st.copy()
+
+
+def read_cat(path: str | PathLike[str]) -> pd.DataFrame:
+    return pd.read_csv(path, sep=r"\s+")
+
+
+def cat2kstyle(row: Any) -> str:
+    """Catalogue date and time as dot-separated whole-number fields.
+
+    Sub-second precision is dropped, because :func:`keith2utc` parses every
+    field with :func:`int`. It used to be dropped with a fixed ``[:-3]`` slice,
+    which silently assumed exactly two decimal places on the seconds:
+
+    * ``"13:09:31.00"`` worked,
+    * ``"13:09:31.000"`` left a trailing separator and made ``keith2utc``
+      raise ``invalid literal for int()``,
+    * ``"13:09:31"`` — no decimals at all — lost the **seconds**, quietly
+      returning 13:09 as though it were 13:09:31.
+
+    The last is the one worth fixing: it is a wrong answer rather than an
+    error, and a catalogue written without fractional seconds is not unusual.
+    Splitting on the decimal point handles all three.
+    """
+    date = row["Date"].replace("/", ".")
+    time = row["Time"].split(".")[0].replace(":", ".")
+    return ".".join([date, time])
+
+
+def keith2utc(row: Any) -> UTCDateTime:
+    return obspy.UTCDateTime(*map(int, cat2kstyle(row).split(".")))