specmod 0.2.0__py3-none-any.whl

specmod/fitting/spectrum.py

@@ -0,0 +1,330 @@
+"""Fitting a source model to a single spectrum.
+
+The lmfit surface used here is declared in ``stubs/lmfit``; see
+``stubs/README.md``. lmfit ships no annotations, so without those a
+``ModelResult`` is `Any` and nothing checks that ``result.redchi`` exists or
+that ``Parameter.stderr`` can be `None` — which it is under every minimiser
+that estimates no covariance matrix, including the shipped default.
+"""
+
+from __future__ import annotations
+
+from copy import deepcopy
+from typing import TYPE_CHECKING, Any
+
+import lmfit as lm
+import matplotlib.pyplot as plt
+import numpy as np
+from matplotlib.ticker import NullFormatter, StrMethodFormatter
+
+from .. import config as cfg
+from .. import sources
+from .base import REQUIRED_SPECTRUM_ATTRIBUTES, Spectrumish
+from .guess import selected_band
+
+if TYPE_CHECKING:  # pragma: no cover
+    from collections.abc import Mapping
+
+    from matplotlib.axes import Axes
+    from numpy.typing import NDArray
+
+__all__ = ["FitSpectrum"]
+
+
+class FitSpectrum:
+    """Fit a source model to one spectrum with lmfit.
+
+    Takes anything carrying :data:`REQUIRED_SPECTRUM_ATTRIBUTES` — in practice
+    a :class:`~specmod.core.collection.FittableView` from
+    :func:`fittable_signal`.
+    """
+
+    #: Declarations, not defaults. These were class attributes carrying `None`
+    #: and `{}`, which meant two things at once: every read had to cope with a
+    #: `None` that `__init__` had in fact replaced, and `meta = {}` was one
+    #: dictionary shared by every instance ever constructed. `__init__` assigns
+    #: all of them, so the type is what it is after construction — and the
+    #: shared-mutable-default hazard is gone rather than merely unreached.
+    sig: Spectrumish
+    mod: lm.Model
+    params: lm.Parameters
+    #: `None` until :meth:`fit_mod` runs. This one really is optional, and
+    #: callers test it — see :func:`specmod.plotting.plot_pair`.
+    result: lm.ModelResult | None
+    mod_freq: NDArray[np.float64]
+    mod_amp: NDArray[np.float64]
+    pass_fitting: bool
+    fit_bins: bool
+    meta: dict[str, Any]
+    #: The :class:`specmod.sources.SpectralModel` behind the fit, when there is
+    #: one. ``None`` if a bare callable was supplied.
+    spectral_model: sources.SpectralModel | None
+
+    def __init__(
+        self,
+        signal: Spectrumish,
+        model: Any = None,
+        fit_bins: bool = False,
+        **params: float,
+    ) -> None:
+        self.result = None
+        self.pass_fitting = True
+        self.meta = {}
+        self.spectral_model = None
+        self.mod_freq = np.array([])
+        self.mod_amp = np.array([])
+        self.fit_bins = fit_bins
+        self.set_signal(signal)
+        self.set_model(model, **params)
+
+    def fit_mod(self, **kwargs: Any) -> None:
+        """Fit, judge the result, then record it — in that order.
+
+        The judgement used to be made *after* the recording, so the
+        ``pass_fitting`` column of every flat file held the value from before
+        the fit ran — ``True``, the class default, on a fresh `FitSpectrum`.
+        The attribute and the table disagreed, and the table is what gets
+        written out and regressed on.
+        """
+        self.result = self.mod.fit(self.mod_amp, self.params, f=self.mod_freq, **kwargs)
+        self.__determine_pass_or_fail()
+        self.__set_results_to_meta()
+
+    def set_signal(self, signal: Spectrumish) -> None:
+        if self.__check_input(signal):
+            self.sig = signal
+            self.__set_meta(signal.meta)
+        # if setting a new signal - assess and adjust the freq bounds
+        self.__set_mod_amp_freq()
+
+    def set_model(self, model: Any = None, **params: float) -> None:
+        """Set the model to fit.
+
+        Accepts a :class:`specmod.sources.SpectralModel`, a bare callable, or
+        ``None`` — in which case the model is whatever ``[model]`` in the
+        configuration asks for. That default is the point: before it existed,
+        ``config.model.source`` was read by nothing and the caller had to pass
+        the right function by hand, so a study file saying
+        ``source = "boatwright"`` silently got Brune.
+
+        A bare callable still works, because fitting an ad-hoc shape is a
+        legitimate thing to want. It simply carries no provenance:
+        :attr:`spectral_model` is ``None`` and nothing can report what was fitted.
+        """
+        if model is None:
+            model = sources.from_config()
+
+        if isinstance(model, sources.SpectralModel):
+            self.spectral_model = model
+            model = model.as_callable()
+        else:
+            self.spectral_model = None
+
+        self.mod = lm.Model(model)
+        # whenever a model is set the inital params must be set also
+        self.__init_params(**params)
+
+    @property
+    def fitted(self) -> lm.ModelResult:
+        """The fit result, or a message saying it has not been fitted.
+
+        Every private reader below went straight through ``self.result``,
+        which is ``None`` until :meth:`fit_mod` runs — so calling
+        :meth:`quick_vis` on an unfitted spectrum raised ``AttributeError:
+        'NoneType' object has no attribute 'best_fit'``, from a line that
+        names neither the station nor the missing step.
+        """
+        if self.result is None:
+            raise RuntimeError(
+                f"{getattr(self.sig, 'id', 'this spectrum')} has not been "
+                "fitted yet; call fit_mod() first"
+            )
+        return self.result
+
+    def describe_model(self) -> str | None:
+        """What is being fitted, or ``None`` for a bare callable."""
+        return None if self.spectral_model is None else self.spectral_model.describe()
+
+    def set_const(self, pname: str, value: float) -> None:
+        self.params[pname].value = value
+        self.params[pname].vary = False
+
+    def set_bounds(
+        self, pname: str, min: float | None = None, max: float | None = None
+    ) -> None:
+        if min is not None:
+            self.params[pname].min = min
+        if max is not None:
+            self.params[pname].max = max
+
+    def __set_meta(self, meta: Mapping[str, Any]) -> None:
+        self.meta = deepcopy(dict(meta))
+
+    def __init_params(self, **params: Any) -> None:
+        """Seed the parameters, and floor ``t*`` where the configuration says.
+
+        ``fitting.t_star_min`` existed and was read by nothing. The tutorial
+        did ``fits.set_bounds("ts", min=0.0001)`` by hand and the config value
+        is 1e-4 — the same number — so the setting was a written-down record of
+        something every caller had to remember. Applied here, forgetting it is
+        no longer possible.
+
+        The same applies to ``fc``, and the legacy code knew it — the line
+        ``# self.set_bounds('fc', min=0)`` sat commented out here. It is not a
+        poor fit but an unphysical one: a negative ``t*`` says the wave gained
+        energy travelling, and a corner frequency below zero says nothing at
+        all. lmfit returns either if the misfit surface leans that way, and
+        with the shipped multitaper default it returned ``fc = -4.45 Hz`` on
+        one PNR station while ``pass_fitting`` reported success — because a
+        parameter with no bound cannot be *at* its bound.
+        """
+        # `**params: Any` rather than `float`, because lmfit's `make_params`
+        # takes a leading `verbose` argument: a model with a parameter of that
+        # name would have its seed swallowed as a flag. Not a hazard for any
+        # source model here, and not one this package can fix.
+        self.params = self.mod.make_params(**params)
+        fitting = cfg.load_config().config.fitting
+        for name, floor in (
+            ("ts", fitting.t_star_min),
+            ("fc", fitting.corner_frequency_min),
+        ):
+            if name in self.params and floor is not None:
+                self.set_bounds(name, min=floor)
+
+    def reset(self) -> None:
+        for par in self.params.values():
+            par.vary = True
+            par.min = -np.inf
+            par.max = np.inf
+
+    def __check_input(self, signal: Spectrumish) -> bool:
+        """Accept anything carrying what the fit reads, not one named class.
+
+        This used to be ``isinstance(signal, spectral.Signal)``, which is the
+        coupling that kept the container holding the legacy pair — nothing
+        could be handed to the fitter unless it *was* that class. What the fit
+        actually needs is the six attributes below, so that is what is checked.
+
+        Named explicitly rather than left to fail at first use: a missing
+        ``bamp`` should say so here, not surface as an AttributeError from
+        inside a band selection three calls later.
+        """
+        missing = [
+            name for name in REQUIRED_SPECTRUM_ATTRIBUTES if not hasattr(signal, name)
+        ]
+        if missing:
+            raise ValueError(
+                f"{type(signal).__name__} cannot be fitted: missing "
+                f"{', '.join(missing)}. A fittable spectrum needs "
+                f"{', '.join(REQUIRED_SPECTRUM_ATTRIBUTES)}."
+            )
+        return True
+
+    def __set_mod_amp_freq(self) -> None:
+        """
+        Only fit between signal limits if they are specified.
+        """
+
+        if self.fit_bins:
+            freq = self.sig.bfreq
+            amp = self.sig.bamp
+        else:
+            freq = self.sig.freq
+            amp = self.sig.amp
+
+        band = selected_band(self.sig)
+        if band is not None:
+            inds = np.where((freq >= band[0]) & (freq <= band[1]))
+            self.mod_freq = freq[inds]
+            self.mod_amp = amp[inds]
+        else:
+            self.mod_freq = freq
+            self.mod_amp = amp
+
+        self.mod_amp = np.log10(self.mod_amp)
+
+    def __param_string(self) -> str:
+        """``name: value+/-2sigma`` per parameter, or ``name: value`` alone.
+
+        The old version computed ``2 * k.stderr`` unconditionally inside a
+        bare ``except Exception``. ``stderr`` is ``None`` whenever the
+        minimiser estimated no covariance matrix — which Powell, the shipped
+        default, never does — so this raised ``TypeError`` on every fit made
+        with the default configuration, swallowed it, and titled the plot
+        ``NaN``. A missing uncertainty is a property of the method, not a
+        failed fit, so the value is still worth printing.
+        """
+        parts = []
+        for k in self.fitted.params.values():
+            if k.stderr is None:
+                parts.append(f"{k.name}: {k.value:.3f}")
+            else:
+                parts.append(f"{k.name}: {k.value:.3f}+/-{2 * k.stderr:.3f}")
+        return ", ".join(parts)
+
+    def quick_vis(self, ax: Axes | None = None) -> Axes:
+        if ax is None:
+            _fig, ax = plt.subplots(1, 1)
+
+        ax.loglog(self.mod_freq, 10**self.mod_amp, color="grey", label=self.sig.id)
+        ax.loglog(self.mod_freq, 10**self.fitted.best_fit, "k--", label="model")
+        ax.xaxis.set_major_formatter(StrMethodFormatter("{x:.2f}"))
+        ax.xaxis.set_minor_formatter(NullFormatter())
+        ax.set_title(self.__param_string())
+        ax.set_xlabel("freq [Hz]")
+        ax.set_ylabel("spectral amp")
+        ax.legend()
+        return ax
+
+    def __get_pars(self) -> dict[str, Any]:
+        p: dict[str, Any] = {}
+        for k in self.fitted.params.values():
+            p.update({k.name: k.value})
+            p.update({k.name + "-stderr": k.stderr})
+        return p
+
+    def __get_fit_stats(self) -> dict[str, float]:
+        res = self.fitted
+        s: dict[str, float] = {}
+        s.update({"aic": res.aic})
+        s.update({"bic": res.bic})
+        s.update({"chisqr": res.chisqr})
+        s.update({"redchi": res.redchi})
+        return s
+
+    def __get_test_results(self) -> dict[str, bool]:
+        t: dict[str, bool] = {}
+        t.update({"pass_fitting": self.pass_fitting})
+        return t
+
+    def __set_results_to_meta(self) -> None:
+        self.meta.update(self.__get_pars())
+        self.meta.update(self.__get_fit_stats())
+        self.meta.update(self.__get_test_results())
+
+    def __determine_pass_or_fail(self) -> None:
+        """A fit fails when a parameter is pinned against one of its bounds.
+
+        Which is the useful question: a corner frequency resting on its floor
+        is the minimiser saying "lower, if you would let me", and the value it
+        reports is the bound rather than a measurement.
+
+        Reset first. ``pass_fitting`` starts as a class attribute and was only
+        ever set *False*, so a `FitSpectrum` that failed once could never pass
+        again however many times it was refitted.
+
+        **Where there is no uncertainty, the value itself is compared.** The
+        old version treated a missing ``stderr`` as a failure, which would mark
+        every fit failed under the shipped configuration: Powell does not
+        estimate a covariance matrix, so lmfit has no uncertainties to report.
+        That is a property of the minimiser, not a fault in the fit. Asking
+        whether the value sits on the bound is the same question with the
+        error bar removed.
+        """
+        self.pass_fitting = True
+        for _par, vals in self.fitted.params.items():
+            if not vals.vary:
+                continue
+            spread = vals.stderr if vals.stderr is not None else 0.0
+            if (vals.value - spread <= vals.min) or (vals.value + spread >= vals.max):
+                self.pass_fitting = False

specmod/io.py

@@ -0,0 +1,241 @@
+"""Writing an event's spectra to disk and reading them back.
+
+Replaces ``spectral.Spectra.write_spectra``/``read_spectra``, which pickled the
+container. Pickle stores the import path of every class it holds, so a stored
+result stops loading the moment a class is renamed or moved — which is exactly
+what happened to the shipped ``Tutorial/Spectra/*.spec``: unreadable since the
+``Spectral.py`` → ``spectral.py`` rename, years before the classes were
+deleted. **A format that breaks when you refactor is a cache, not a format.**
+
+HDF5, per ``REFACTOR_PLAN`` §4.6. The four rules there are lessons from how
+pickle failed rather than general good practice, and each is visible in the
+layout below:
+
+*Never store class identity.*
+    Plain arrays and a documented layout. Nothing here names a Python type, so
+    nothing here can be broken by renaming one.
+
+*Every file carries a format version.*
+    :data:`FORMAT_VERSION`, on the root group. A reader checks it and fails
+    naming both versions. The absence of this is the whole problem with what
+    came before.
+
+*Self-describing units.*
+    ``motion``, ``kind``, ``duration`` and ``sampling_rate`` are stored
+    attributes, not conventions the reader has to know. This is the typing of
+    :class:`~specmod.core.Spectrum` expressed on disk, and it is what stops a
+    file being silently misread as displacement.
+
+*One file per event, one group per channel.*
+    Which matches how the science is done and how it is re-examined, and
+    sidesteps HDF5's single-writer limitation if the workflow is ever
+    parallelised across events.
+
+**Nothing here can pickle.** ``h5py`` has no mechanism to; trace metadata is
+stored as a JSON string attribute rather than as an object array, which is how
+``numpy`` would otherwise smuggle pickling back in.
+
+The tables — fit results — go to Parquet instead, in :mod:`specmod.tables`.
+Arrays and tables are used in genuinely different ways (random access into one
+event, versus a columnar scan over every event ever fitted), and one format for
+both would be worse at each.
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+import numpy as np
+
+from .core import (
+    AmplitudeKind,
+    BinnedSpectrum,
+    Motion,
+    Spectrum,
+    SpectrumPair,
+    SpectrumSet,
+)
+
+if TYPE_CHECKING:  # pragma: no cover
+    from collections.abc import Mapping
+
+__all__ = ["FORMAT_VERSION", "load", "save"]
+
+#: Bumped when the layout changes in a way an older reader cannot handle.
+FORMAT_VERSION = 1
+
+#: Compress datasets at least this large, in bytes; store smaller ones raw.
+#:
+#: Measured, not guessed. Compression in HDF5 requires *chunked* storage, which
+#: costs a chunk index per dataset. On this event the arrays are tiny — median
+#: 91 float64, 728 bytes — and gzip took the file from 370 KB to **867 KB**
+#: against a raw payload of 250 KB: the indices cost more than the data they
+#: describe. Above the threshold the ratio flips and compression pays, which is
+#: the case the plan's "chunked, compressed" argument was really about (a
+#: scalogram is ~1 MB per trace).
+_COMPRESS_ABOVE_BYTES = 16 * 1024
+
+
+def _require_h5py() -> Any:
+    """Import ``h5py``, or say what to install.
+
+    Deferred rather than imported at module scope so that the message names the
+    extra. A bare ``ModuleNotFoundError: h5py`` in the middle of a save tells a
+    user nothing about how to fix it.
+    """
+    try:
+        import h5py  # noqa: PLC0415
+    except ImportError as exc:  # pragma: no cover - depends on the environment
+        raise ImportError(
+            "reading and writing spectra needs h5py. Install it with "
+            "`pip install specmod[io]`."
+        ) from exc
+    return h5py
+
+
+def _dataset(group: Any, name: str, data: Any) -> None:
+    """One array, compressed only where compression pays. See
+    :data:`_COMPRESS_ABOVE_BYTES`."""
+    data = np.asarray(data)
+    if data.nbytes >= _COMPRESS_ABOVE_BYTES:
+        group.create_dataset(name, data=data, compression="gzip")
+    else:
+        group.create_dataset(name, data=data)
+
+
+def _write_spectrum(group: Any, name: str, spectrum: Spectrum) -> None:
+    """One spectrum: two datasets, and the attributes that make them readable.
+
+    The units go on as attributes rather than being implied, because a
+    frequency and an amplitude array alone do not say whether they are a
+    velocity magnitude or a displacement power spectral density — and the two
+    differ by factors that would pass unnoticed.
+    """
+    # Only the amplitude. The frequency axis is written once per group, by
+    # `save`: `compare` interpolates the noise onto the signal's axis and bins
+    # both against the same edges, so the two spectra *share* an axis by
+    # construction. Storing it twice would be waste, and worse, would make a
+    # file expressible in which they disagree — a state the containers cannot
+    # represent and no reader could act on.
+    _dataset(group, f"{name}_amp", spectrum.amp)
+    attrs = group.attrs
+    attrs[f"{name}_motion"] = str(spectrum.motion)
+    attrs[f"{name}_kind"] = str(spectrum.kind)
+    attrs[f"{name}_duration"] = float(spectrum.duration)
+    attrs[f"{name}_sampling_rate"] = float(spectrum.sampling_rate)
+    attrs[f"{name}_meta"] = json.dumps(dict(spectrum.meta))
+
+
+def _read_spectrum(group: Any, name: str) -> Spectrum:
+    attrs = group.attrs
+    return Spectrum(
+        freq=np.asarray(group["freq"]),
+        amp=np.asarray(group[f"{name}_amp"]),
+        # Through the enums rather than as bare strings: constructing them is
+        # what turns a corrupted or hand-edited attribute into an error naming
+        # the value, instead of a spectrum that claims a unit nothing else
+        # recognises.
+        motion=Motion(str(attrs[f"{name}_motion"])),
+        kind=AmplitudeKind(str(attrs[f"{name}_kind"])),
+        duration=float(attrs[f"{name}_duration"]),
+        sampling_rate=float(attrs[f"{name}_sampling_rate"]),
+        meta=json.loads(str(attrs[f"{name}_meta"])),
+    )
+
+
+def save(path: str | Path, spectra: SpectrumSet) -> Path:
+    """Write an event's spectra to ``path``.
+
+    One group per trace id, keyed by the id itself — HDF5 group names may
+    contain dots, so ``LV.L001..HHE`` needs no mangling and the file browses
+    with the same names the pipeline uses.
+
+    The parent directory is created if absent. The legacy version raised
+    ``FileNotFoundError`` from inside ``open``, naming the file rather than the
+    directory that did not exist.
+    """
+    h5py = _require_h5py()
+
+    path = Path(path)
+    # Appended, never substituted. An event id is an ISO timestamp —
+    # `2019-08-26T07:49:24.200000Z` — and `Path.with_suffix` reads `.200000Z`
+    # as an extension and replaces it, silently truncating the name to
+    # `2019-08-26T07:49:24.h5`. Two events in the same second would then
+    # overwrite each other.
+    if path.suffix != ".h5":
+        path = path.with_name(path.name + ".h5")
+    path.parent.mkdir(parents=True, exist_ok=True)
+
+    with h5py.File(path, "w") as handle:
+        handle.attrs["specmod_format_version"] = FORMAT_VERSION
+        handle.attrs["event"] = spectra.event
+        handle.attrs["meta"] = json.dumps(dict(spectra.meta))
+
+        for id in spectra.ids():
+            pair = spectra[id]
+            group = handle.create_group(id)
+            # The two shared axes, written once each.
+            _dataset(group, "freq", pair.signal.freq)
+            _dataset(group, "binned_freq", pair.binned_signal.freq)
+
+            _write_spectrum(group, "signal", pair.signal)
+            _write_spectrum(group, "noise", pair.noise)
+            _dataset(group, "binned_signal_amp", pair.binned_signal.amp)
+            _dataset(group, "binned_noise_amp", pair.binned_noise.amp)
+            _dataset(group, "snr", pair.snr)
+            group.attrs["resolution_floor"] = float(pair.resolution_floor)
+            group.attrs["meta"] = json.dumps(dict(pair.meta))
+            # `band` is absent rather than sentinel-valued when no band
+            # survived. A NaN pair would read as a band, and "no usable
+            # bandwidth" is a different claim from "a band at nowhere".
+            if pair.band is not None:
+                group.attrs["band"] = np.asarray(pair.band, dtype=np.float64)
+
+    return path
+
+
+def load(path: str | Path) -> SpectrumSet:
+    """Read back what :func:`save` wrote.
+
+    The arrays come back read-only, as they went in: a spectrum loaded from
+    disk gives the same immutability guarantee as one just computed, which is
+    what makes reload-and-refit safe to do in a loop.
+    """
+    h5py = _require_h5py()
+    path = Path(path)
+
+    with h5py.File(path, "r") as handle:
+        version = handle.attrs.get("specmod_format_version")
+        if version != FORMAT_VERSION:
+            raise ValueError(
+                f"{path} declares format version {version}, and this is "
+                f"version {FORMAT_VERSION}. Re-run the pipeline, or read it "
+                f"with the version of specmod that wrote it."
+            )
+
+        pairs: dict[str, SpectrumPair] = {}
+        for id, group in handle.items():
+            band = group.attrs.get("band")
+            pairs[id] = SpectrumPair(
+                signal=_read_spectrum(group, "signal"),
+                noise=_read_spectrum(group, "noise"),
+                binned_signal=BinnedSpectrum(
+                    freq=np.asarray(group["binned_freq"]),
+                    amp=np.asarray(group["binned_signal_amp"]),
+                ),
+                binned_noise=BinnedSpectrum(
+                    freq=np.asarray(group["binned_freq"]),
+                    amp=np.asarray(group["binned_noise_amp"]),
+                ),
+                snr=np.asarray(group["snr"]),
+                resolution_floor=float(group.attrs["resolution_floor"]),
+                band=None if band is None else (float(band[0]), float(band[1])),
+                meta=json.loads(str(group.attrs["meta"])),
+            )
+
+        event = str(handle.attrs.get("event", ""))
+        meta: Mapping[str, Any] = json.loads(str(handle.attrs.get("meta", "{}")))
+
+    return SpectrumSet(pairs=pairs, event=event, meta=dict(meta))