specmod 0.2.0__py3-none-any.whl

specmod/transforms/base.py

@@ -0,0 +1,242 @@
+"""The estimator interface and the one place normalisation is defined.
+
+Every backend — FFT, Welch, multitaper, and later the CWT — returns a
+:class:`~specmod.core.spectrum.Spectrum` obeying the same contract, so a single
+test suite pins all of them. That is the point of the abstraction: the
+pre-refactor code called ``mtspec`` directly from ``Spectrum.__init__`` and
+threaded backend keyword arguments through three layers of public API, so there
+was nowhere for a shared contract to live.
+
+The contract
+------------
+Given a real record ``x`` of ``N`` samples at interval ``dt``:
+
+- The returned spectrum is **one-sided**, spanning ``(0, f_Nyquist]``.
+- Its default kind is :attr:`~specmod.core.units.AmplitudeKind.FAS`, in
+  ``[x] * s``.
+- ``duration`` is ``N * dt``, the **physical** record length. Normalisation is
+  keyed off it and never off ``len(freq)``, so zero-padding changes resolution
+  without changing amplitude.
+- Energy is preserved: ``spectrum.energy()`` recovers ``sum(x^2) * dt`` to
+  within the accuracy of the estimator.
+
+Taper correction
+----------------
+A taper attenuates the record, and the correction depends on what you are
+measuring. The two are not interchangeable:
+
+``energy`` (default)
+    Divide by ``sqrt(mean(w^2))``, preserving total power. Parseval then holds
+    exactly, which is what the shared test asserts, and it is the right choice
+    for a transient — which is what a seismic arrival is.
+
+``amplitude``
+    Divide by ``mean(w)``, preserving the peak of a coherent sinusoid so that
+    it reads ``A0 * T``. Right when measuring a monochromatic line.
+
+For a Tukey taper with ``alpha=0.05`` the two differ by well under a percent,
+but the choice is explicit rather than implied.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Literal, Protocol, runtime_checkable
+
+import numpy as np
+from numpy.typing import ArrayLike, NDArray
+from scipy.fft import next_fast_len
+from scipy.signal import get_window
+
+from ..core.spectrum import Spectrum
+from ..core.units import AmplitudeKind, Motion
+
+__all__ = [
+    "SpectralEstimator",
+    "TaperCorrection",
+    "make_window",
+    "prepare_record",
+    "window_correction",
+]
+
+TaperCorrection = Literal["energy", "amplitude"]
+
+
+@runtime_checkable
+class SpectralEstimator(Protocol):
+    """Anything that turns a real record into a :class:`Spectrum`."""
+
+    @property
+    def name(self) -> str:
+        """Short identifier, recorded in the spectrum's metadata.
+
+        Declared read-only so the frozen dataclasses implementing this protocol
+        satisfy it; a plain mutable attribute would not.
+        """
+        ...
+
+    def estimate(
+        self,
+        data: ArrayLike,
+        dt: float,
+        *,
+        motion: Motion | str = Motion.VELOCITY,
+        meta: dict[str, Any] | None = None,
+    ) -> Spectrum:
+        """Estimate the spectrum of ``data`` sampled at interval ``dt``."""
+        ...
+
+
+def prepare_record(
+    data: ArrayLike, dt: float
+) -> tuple[NDArray[np.float64], int, float]:
+    """Validate and demean a record, returning ``(x, n_samples, duration)``.
+
+    Demeaning is unconditional: a non-zero mean puts all of its energy in the
+    DC bin, which is then discarded, so leaving it in loses energy the Parseval
+    check would report as a failure.
+    """
+    x = np.ascontiguousarray(np.asarray(data, dtype=np.float64).squeeze())
+    if x.ndim != 1:
+        raise ValueError(f"Expected a 1-D record, got shape {x.shape}")
+    if x.size < 2:
+        raise ValueError(f"Record too short: {x.size} sample(s)")
+    if not np.all(np.isfinite(x)):
+        raise ValueError("Record contains NaN or Inf")
+    if dt <= 0:
+        raise ValueError(f"dt must be positive, got {dt}")
+    return x - x.mean(), int(x.size), float(x.size * dt)
+
+
+def make_window(kind: str, n: int, alpha: float = 0.05) -> NDArray[np.float64]:
+    """Build a taper of length ``n``."""
+    if kind == "boxcar":
+        return np.ones(n, dtype=np.float64)
+    spec: str | tuple[str, float] = ("tukey", alpha) if kind == "tukey" else kind
+    return np.asarray(get_window(spec, n, fftbins=False), dtype=np.float64)
+
+
+def window_correction(
+    window: NDArray[np.float64], correction: TaperCorrection
+) -> float:
+    """Scale factor that undoes a taper's attenuation.
+
+    See the module docstring for why there are two.
+    """
+    if correction == "energy":
+        return float(np.sqrt(np.mean(window**2)))
+    if correction == "amplitude":
+        return float(np.mean(window))
+    raise ValueError(
+        f"Unknown taper correction {correction!r}; expected 'energy' or 'amplitude'."
+    )
+
+
+def resolve_n_fft(n_fft: int | str | None, n_samples: int) -> int:
+    """Turn an ``n_fft`` setting into a transform length.
+
+    Accepts ``None`` for no padding, an explicit integer, or a strategy:
+
+    ``"fast"``
+        The next length ``scipy.fft`` factorises efficiently — a 5-smooth
+        composite. **The one to use.**
+    ``"pow2"``
+        The next power of two. Familiar, and the right answer for a
+        radix-2 implementation, but numpy's pocketfft is not one.
+
+    Both exist because padding to a fast length is purely a speed knob here:
+    the normalisation is keyed off the record duration, so padding changes the
+    frequency sampling and nothing else. Cut windows have arbitrary lengths and
+    a prime one is slow — measured across the 28 PNR S-windows, 1.77x.
+
+    ``"fast"`` is preferred over ``"pow2"`` on measurement, not taste. A power
+    of two overshoots: for a 65537-sample record it pads to 131072, doubling
+    the transform, where the next 5-smooth length is 65610. Over a range of
+    real and awkward lengths ``"fast"`` totalled 0.76 ms against 1.49 ms for
+    ``"pow2"`` and 11.0 ms unpadded.
+    """
+    if n_fft is None:
+        return n_samples
+    if isinstance(n_fft, str):
+        strategy = n_fft.lower()
+        if strategy == "fast":
+            return int(next_fast_len(n_samples))
+        if strategy == "pow2":
+            return 1 << int(np.ceil(np.log2(n_samples)))
+        raise ValueError(
+            f"n_fft must be an integer, None, 'fast' or 'pow2', got {n_fft!r}"
+        )
+    nfft = int(n_fft)
+    if nfft < n_samples:
+        raise ValueError(
+            f"n_fft ({nfft}) is shorter than the record ({n_samples} samples)"
+        )
+    return nfft
+
+
+def one_sided_fas(
+    x: NDArray[np.float64],
+    dt: float,
+    duration: float,
+    *,
+    window: NDArray[np.float64],
+    correction: TaperCorrection,
+    n_fft: int | str | None,
+    drop_dc: bool,
+) -> tuple[NDArray[np.float64], NDArray[np.float64]]:
+    """Core FFT path: taper, transform, fold, normalise.
+
+    Returns ``(freq, fas)`` with ``fas`` in ``[x] * s``.
+
+    The normalisation is ``dt`` per sample, doubled to fold the negative
+    frequencies, and divided by the taper correction. ``dt`` is used rather
+    than anything derived from the transform length, which is what makes
+    padding change only the frequency sampling.
+    """
+    n = x.size
+    nfft = resolve_n_fft(n_fft, n)
+
+    spec = np.fft.rfft(x * window, n=nfft)
+    freq: NDArray[np.float64] = np.fft.rfftfreq(nfft, d=dt).astype(np.float64)
+
+    fas = np.abs(spec) * dt / window_correction(window, correction)
+    # Fold: every bin except DC and, for even nfft, Nyquist has a negative twin.
+    fas[1:] *= 2.0
+    if nfft % 2 == 0:
+        fas[-1] /= 2.0
+
+    # Deliberately no rescaling for zero-padding. The DFT sums over the N
+    # non-zero samples whatever nfft is, so it approximates the same continuous
+    # transform `dt * sum(x_n exp(-2 pi i f t_n))` and merely evaluates it on a
+    # finer grid. Normalising by `dt` alone is what makes that true; scaling by
+    # any length ratio would reintroduce exactly the padding sensitivity that
+    # made the pre-refactor `psd_to_amp` wrong (§2.2).
+    del duration, n  # normalisation depends on neither; kept for the signature
+
+    if drop_dc:
+        return freq[1:], fas[1:]
+    return freq, fas
+
+
+def build_spectrum(
+    freq: NDArray[np.float64],
+    amp: NDArray[np.float64],
+    *,
+    kind: AmplitudeKind,
+    motion: Motion | str,
+    duration: float,
+    sampling_rate: float,
+    meta: dict[str, Any] | None,
+    estimator: str,
+) -> Spectrum:
+    """Assemble a Spectrum, recording which estimator produced it."""
+    info = dict(meta or {})
+    info.setdefault("estimator", estimator)
+    return Spectrum(
+        freq=freq,
+        amp=amp,
+        motion=Motion(motion),
+        kind=kind,
+        duration=duration,
+        sampling_rate=sampling_rate,
+        meta=info,
+    )

specmod/transforms/cwt.py

@@ -0,0 +1,219 @@
+"""Continuous wavelet transform on an L2-normalised Morlet.
+
+The reason this is written here rather than delegated to PyWavelets: the
+normalisation is the entire difficulty. CWT coefficients carry units of
+``[signal] * sqrt(time)`` under L2 and ``[signal]`` under L1, and neither is the
+``[signal] * s`` of a Fourier amplitude spectrum. A naive ``|W|`` plotted
+against scale-derived frequency looks entirely reasonable and gives the right
+corner frequency with the wrong ``Omega`` — so the wrong seismic moment, with
+nothing anywhere to catch it. ``pywt.cwt``'s conventions are not documented to
+the precision that requires, so using it would mean reverse-engineering them.
+
+The bridge back to amplitude lives in
+:meth:`specmod.core.scalogram.Scalogram.time_average`, and
+``tests/test_cwt.py`` is its specification: the CWT has to agree with the FFT
+and multitaper estimators on a sinusoid of known amplitude and on the energy of
+a record. The derivation is not trusted; the test is.
+
+References
+----------
+Torrence, C. and Compo, G.P. (1998). A practical guide to wavelet analysis.
+*Bulletin of the American Meteorological Society* 79(1), 61-78.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from types import MappingProxyType
+from typing import Any
+
+import numpy as np
+from numpy.typing import ArrayLike, NDArray
+
+from ..core.scalogram import Scalogram
+from ..core.spectrum import Spectrum
+from ..core.units import Motion
+from .base import prepare_record
+
+__all__ = ["CWTEstimator"]
+
+
+def _morlet_fourier_factor(omega0: float) -> float:
+    """Scale-to-Fourier-period conversion, T&C Table 1.
+
+    ``period = factor * scale``. This is what lets the output frequency axis
+    mean the same thing as an FFT's rather than being "scale, roughly".
+    """
+    return float(4.0 * np.pi / (omega0 + np.sqrt(2.0 + omega0**2)))
+
+
+@dataclass(frozen=True)
+class CWTEstimator:
+    """Continuous wavelet transform, time-averaged to an amplitude spectrum.
+
+    Produces both outputs from one transform: :meth:`scalogram` returns the
+    full time-frequency surface, and :meth:`estimate` returns its time average
+    as an ordinary :class:`~specmod.core.spectrum.Spectrum`, so the fitting
+    pipeline does not know the difference. The surface is what you look at when
+    a fit comes out wrong.
+
+    Parameters
+    ----------
+    omega0
+        Morlet central frequency, dimensionless. Trades time resolution against
+        frequency resolution; 6 is the conventional choice and the value for
+        which the analytic scale-frequency relation is usually quoted.
+    dj
+        Scale spacing in octaves. Smaller resolves the frequency axis more
+        finely at proportionally more cost. The normalisation carries an
+        explicit ``dj`` factor, so the recovered energy does not depend on it.
+    mask_coi
+        Exclude the cone of influence from the time average. On by default:
+        without it a short window reads low at low frequency, which is the band
+        that constrains ``Omega``.
+    f_min, f_max
+        Frequency range to cover. Defaults span ``1/T`` — the longest period
+        the record can represent — up to the Nyquist frequency.
+    """
+
+    omega0: float = 6.0
+    dj: float = 0.125
+    mask_coi: bool = True
+    f_min: float | None = None
+    f_max: float | None = None
+    name: str = "cwt"
+
+    def __post_init__(self) -> None:
+        if self.omega0 <= 0:
+            raise ValueError(f"omega0 must be positive, got {self.omega0}")
+        if self.dj <= 0:
+            raise ValueError(f"dj must be positive, got {self.dj}")
+        if (
+            self.f_min is not None
+            and self.f_max is not None
+            and self.f_min >= self.f_max
+        ):
+            raise ValueError(f"f_min={self.f_min} must be below f_max={self.f_max}")
+
+    def _scales(self, n: int, dt: float) -> NDArray[np.float64]:
+        factor = _morlet_fourier_factor(self.omega0)
+        f_max = self.f_max if self.f_max is not None else 0.5 / dt
+        f_min = self.f_min if self.f_min is not None else 1.0 / (n * dt)
+        # scale = period / factor, and period = 1/f.
+        s_min = 1.0 / (f_max * factor)
+        s_max = 1.0 / (f_min * factor)
+        n_scales = int(np.floor(np.log2(s_max / s_min) / self.dj)) + 1
+        return np.asarray(s_min * 2.0 ** (np.arange(n_scales) * self.dj), dtype=float)
+
+    def _c_delta(self, dt: float) -> float:
+        """Reconstruction constant, computed against *this* scale grid.
+
+        Two reasons not to use the tabulated value, and the second is the one
+        that matters.
+
+        First, T&C tabulate 0.776 for ``omega0=6`` only, and this class lets
+        ``omega0`` vary. The constant moves with it — 1.05 at 4, 0.55 at 8 —
+        so a hardcoded 0.776 silently mis-normalises every non-default choice.
+
+        Second, and less obvious: the value derived here is deliberately *not*
+        0.776 even at ``omega0=6``. It comes out near 0.72 and drifts with
+        ``dj``, because the reconstruction sum is a discrete approximation to a
+        continuous integral over scale, and this measures the approximation
+        actually being used rather than its limit. That is what keeps the
+        Parseval contract: with the computed constant, recovered energy is
+        0.97-1.07 of the truth; with 0.776 it is systematically ~7% low, and
+        the error grows as ``dj`` coarsens.
+
+        So the ``dj`` dependence is the point, not a defect — do not "fix" it
+        by substituting the published number. ``tests/test_cwt.py`` pins this.
+        """
+        n = 1024
+        delta = np.zeros(n)
+        delta[n // 2] = 1.0
+        scales = self._scales(n, dt)
+        coefficients = self._transform(delta, dt, scales)
+        # T&C eq. 11 rearranged: the delta reconstructs to 1 at its own sample.
+        summed = float(np.sum(np.real(coefficients[:, n // 2]) / np.sqrt(scales)))
+        psi0_at_zero = np.pi**-0.25
+        return float(self.dj * np.sqrt(dt) / psi0_at_zero * summed)
+
+    def _transform(
+        self, x: NDArray[np.float64], dt: float, scales: NDArray[np.float64]
+    ) -> NDArray[np.complex128]:
+        """Wavelet coefficients, L2-normalised, computed in the frequency domain."""
+        n = x.size
+        x_hat = np.fft.fft(x)
+        omega = 2.0 * np.pi * np.fft.fftfreq(n, d=dt)
+
+        # Morlet is analytic: it has no support at negative frequency, which is
+        # what makes |W| an envelope rather than an oscillation.
+        heaviside = (omega > 0).astype(float)
+        scaled = scales[:, None] * omega[None, :]
+        psi_hat = (
+            np.pi**-0.25
+            * heaviside[None, :]
+            * np.exp(-((scaled - self.omega0) ** 2) / 2.0)
+        )
+        # T&C eq. 6: the sqrt(2*pi*s/dt) factor is the L2 normalisation, and is
+        # what makes the coefficients comparable across scales.
+        psi_hat = psi_hat * np.sqrt(2.0 * np.pi * scales[:, None] / dt)
+
+        coefficients: NDArray[np.complex128] = np.fft.ifft(
+            x_hat[None, :] * np.conjugate(psi_hat), axis=-1
+        )
+        return coefficients
+
+    def scalogram(
+        self,
+        data: ArrayLike,
+        dt: float,
+        *,
+        motion: Motion | str = Motion.VELOCITY,
+        meta: dict[str, Any] | None = None,
+    ) -> Scalogram:
+        """The full time-frequency surface."""
+        x, n, duration = prepare_record(data, dt)
+        scales = self._scales(n, dt)
+        coefficients = self._transform(x, dt, scales)
+
+        factor = _morlet_fourier_factor(self.omega0)
+        freq = 1.0 / (factor * scales)
+
+        # T&C section 3g: the e-folding time for Morlet is sqrt(2)*s, so the
+        # longest resolvable period at a given sample grows with its distance
+        # from the nearest edge.
+        edge_distance = np.minimum(np.arange(n), n - 1 - np.arange(n)) * dt
+        coi = factor * np.sqrt(2.0) * edge_distance
+
+        return Scalogram(
+            time=(np.arange(n) * dt).astype(np.float64),
+            freq=freq,
+            power=np.abs(coefficients) ** 2,
+            scales=scales,
+            coi=coi,
+            c_delta=self._c_delta(dt),
+            dj=self.dj,
+            dt=dt,
+            motion=Motion(motion),
+            meta=MappingProxyType(
+                {
+                    **(meta or {}),
+                    "estimator": self.name,
+                    "omega0": self.omega0,
+                    "dj": self.dj,
+                    "duration": duration,
+                }
+            ),
+        )
+
+    def estimate(
+        self,
+        data: ArrayLike,
+        dt: float,
+        *,
+        motion: Motion | str = Motion.VELOCITY,
+        meta: dict[str, Any] | None = None,
+    ) -> Spectrum:
+        """Time-averaged amplitude spectrum, interchangeable with the others."""
+        surface = self.scalogram(data, dt, motion=motion, meta=meta)
+        return surface.time_average(mask_coi=self.mask_coi)

specmod/transforms/fft.py

@@ -0,0 +1,157 @@
+"""Plain FFT and Welch estimators.
+
+``FFTEstimator`` is the direct route: taper, transform, fold, normalise. Paired
+with a smoother (:mod:`specmod.smoothing`) it is the conventional
+engineering-seismology approach and the fastest option here.
+
+``WelchEstimator`` averages over overlapping segments, trading frequency
+resolution for variance. That makes it a good default for *noise* windows,
+where a stable level matters more than resolving structure.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+import numpy as np
+from numpy.typing import ArrayLike
+from scipy.signal import welch
+
+from ..core.spectrum import Spectrum
+from ..core.units import AmplitudeKind, Motion
+from .base import (
+    TaperCorrection,
+    build_spectrum,
+    make_window,
+    one_sided_fas,
+    prepare_record,
+)
+
+__all__ = ["FFTEstimator", "WelchEstimator"]
+
+
+@dataclass(frozen=True)
+class FFTEstimator:
+    """One-sided Fourier amplitude spectrum via :func:`numpy.fft.rfft`.
+
+    Parameters
+    ----------
+    taper, taper_alpha
+        Window applied before transforming. ``tukey`` with a small ``alpha``
+        suppresses edge discontinuities while leaving the body of the record
+        untouched.
+    taper_correction
+        ``"energy"`` preserves Parseval, ``"amplitude"`` preserves the peak of a
+        coherent sinusoid. See :mod:`specmod.transforms.base`.
+    n_fft
+        Transform length: ``None`` for no padding, an integer, or a strategy —
+        ``"fast"`` for the next efficiently-factorised length, ``"pow2"`` for
+        the next power of two. See
+        :func:`~specmod.transforms.base.resolve_n_fft`.
+
+        Padding refines the frequency grid without changing amplitude — the
+        property the pre-refactor normalisation got wrong by keying off
+        ``len(freq)``. So it buys two things and neither is leakage
+        suppression, which is the taper's job: it removes scalloping loss
+        (36% worst case on a line falling between bins, unpadded), and it
+        avoids the slow path for an awkward record length.
+
+        ``"fast"`` is the one to reach for. Cut windows are not round numbers —
+        of the 28 PNR S-windows, 17 are odd and several are prime — and a prime
+        length costs 1.77x across those. ``"pow2"`` is offered because it is
+        what people expect, but it overshoots: numpy's pocketfft handles
+        5-smooth lengths, so padding 65537 to 131072 does twice the work of
+        padding it to 65610.
+    drop_dc
+        Discard the zero-frequency bin.
+    """
+
+    taper: str = "tukey"
+    taper_alpha: float = 0.05
+    taper_correction: TaperCorrection = "energy"
+    n_fft: int | str | None = None
+    drop_dc: bool = True
+    name: str = "fft"
+
+    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)
+        window = make_window(self.taper, x.size, self.taper_alpha)
+        freq, fas = one_sided_fas(
+            x,
+            dt,
+            duration,
+            window=window,
+            correction=self.taper_correction,
+            n_fft=self.n_fft,
+            drop_dc=self.drop_dc,
+        )
+        return build_spectrum(
+            freq,
+            fas,
+            kind=AmplitudeKind.FAS,
+            motion=motion,
+            duration=duration,
+            sampling_rate=1.0 / dt,
+            meta=meta,
+            estimator=self.name,
+        )
+
+
+@dataclass(frozen=True)
+class WelchEstimator:
+    """Segment-averaged PSD via :func:`scipy.signal.welch`, returned as FAS.
+
+    Averaging reduces variance at the cost of frequency resolution, which is
+    the right trade for a noise window. Note that ``duration`` remains the full
+    record length: it is the physical property of the record, not of the
+    segments, and every kind conversion depends on it.
+    """
+
+    segment_length: int | None = None
+    overlap: float = 0.5
+    taper: str = "hann"
+    drop_dc: bool = True
+    name: str = "welch"
+
+    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)
+        nperseg = min(self.segment_length or n, n)
+        freq, psd = welch(
+            x,
+            fs=1.0 / dt,
+            window=self.taper,
+            nperseg=nperseg,
+            noverlap=int(nperseg * self.overlap),
+            detrend=False,
+            return_onesided=True,
+            scaling="density",
+        )
+        # scipy's density is one-sided [x]^2/Hz, matching AmplitudeKind.PSD.
+        if self.drop_dc:
+            freq, psd = freq[1:], psd[1:]
+        spectrum = build_spectrum(
+            np.asarray(freq, dtype=np.float64),
+            np.asarray(psd, dtype=np.float64),
+            kind=AmplitudeKind.PSD,
+            motion=motion,
+            duration=duration,
+            sampling_rate=1.0 / dt,
+            meta=meta,
+            estimator=self.name,
+        )
+        return spectrum.to_kind(AmplitudeKind.FAS)