Matched filter

The one matched-filter core: background estimators (empirical, shrunk, low-rank, cluster, streaming), the kernel on gaussx operators, targets from the observation model, and xarray I/O. Theory: matched filter.

matched_filter

Matched filter for hyperspectral methane retrieval.

The matched filter (MF) is the maximum-SNR linear detector for a known target signature t(ν) in additive Gaussian noise with covariance Σ. It is also the one-CG-iteration limit of the dual / PSAS 3D-Var solver with a flat (identity) prior — see notebooks/matched_filter/00_mf_derivation.md.

Score formula (per pixel)::

y(x) = (x − μ)ᵀ Σ⁻¹ t / (tᵀ Σ⁻¹ t)

where x is the observed spectrum, μ is the background mean, and t is the target signature — typically the radiance Jacobian ∂L/∂VMR from the Beer–Lambert forward model at the background state.

Architecture

• Estimation (background, cluster, streaming) — statistical estimation uses scikit-learn (EmpiricalCovariance, LedoitWolf, TruncatedSVD, GaussianMixture). Not differentiable — that's fine, MF doesn't need it.
• Application (core) — the per-pixel score evaluation runs in JAX and dispatches on gaussx operator structure. Covariance estimators return a :class:lineax.AbstractLinearOperator (dense, low-rank + Tikhonov via :class:gaussx.LowRankUpdate, etc.) so the kernel is single-implementation.
• Target signatures (target) — derived from an existing :class:radtrax.retrieval.var3d.RadianceObservationModel via :func:jax.jvp (directional derivatives) or a finite-amplitude secant, or per pixel from the layered :class:radtrax.ForwardModel (unit_absorption_target, radtrax#37). No duplicate Beer–Lambert code.
• I/O (io) — xarray wrappers for multi-scene workflows.

The derivation notebook is ported with the docs (#24).

Ported from plumax 0.1.2 at 28f257e (plumax.matched_filter); only import paths changed.

The package resolves its public names lazily (PEP 562), so they are documented here from the submodules that define them.

core

Matched-filter kernel — one implementation for any covariance structure.

The classical matched-filter score is

.. math:: y(x) = \frac{(x - \mu)^\top \Sigma^{-1} t}{t^\top \Sigma^{-1} t}

with the special case :math:y(\mu + \alpha t) = \alpha (so y is an unbiased linear estimator of the plume enhancement amplitude). The per-pixel evaluation cost is dominated by :math:w = \Sigma^{-1} t, a single solve that can be precomputed once per scene; after that each pixel is a dot product of length n_bands.

This module wraps that into a single kernel that works for any :class:lineax.AbstractLinearOperator covariance — dense, diagonal, low-rank + Tikhonov (Woodbury via :class:gaussx.LowRankUpdate), or Kronecker. The solve is dispatched via :func:gaussx.solve, so the MF caller never sees the structure.

Ported from plumax 0.1.2 at 28f257e (src/plumax/matched_filter/core.py); only import paths changed.

validate_mf_inputs

validate_mf_inputs(cov_op: LinearOperator, target: Float[ArrayLike, ' B']) -> None

Public fail-fast check for (Σ, t) compatibility.

Runs the same tᵀ Σ⁻¹ t > 0 test that :func:apply_pixel and :func:apply_image perform automatically, but without the rest of the MF work — useful for catching user errors before jitting a pipeline or batching over many pixels.

Raises

ValueError If target is all zero, or if the target energy tᵀ Σ⁻¹ t is not strictly positive (singular or indefinite covariance along the target direction).

Source code in src/radtrax/retrieval/matched_filter/core.py

def validate_mf_inputs(cov_op: LinearOperator, target: Float[ArrayLike, " B"]) -> None:
    """Public fail-fast check for ``(Σ, t)`` compatibility.

    Runs the same ``tᵀ Σ⁻¹ t > 0`` test that :func:`apply_pixel` and
    :func:`apply_image` perform automatically, but without the rest of the
    MF work — useful for catching user errors *before* jitting a pipeline
    or batching over many pixels.

    Raises
    ------
    ValueError
        If ``target`` is all zero, or if the target energy ``tᵀ Σ⁻¹ t`` is
        not strictly positive (singular or indefinite covariance along the
        target direction).
    """
    w = gx.solve(cov_op, target)
    target_norm_sq = jnp.dot(target, w)
    _check_target_energy(target, target_norm_sq)

apply_pixel

apply_pixel(pixel: Float[ArrayLike, ' B'], mean: Float[ArrayLike, ' B'], cov_op: LinearOperator, target: Float[ArrayLike, ' B']) -> Float[Array, '']

Matched-filter score for a single pixel spectrum.

Parameters

pixel Observed radiance spectrum at one pixel, shape (n_bands,). mean Background mean spectrum, shape (n_bands,). cov_op Background covariance as a :class:lineax.AbstractLinearOperator. Any structure is fine — :func:gaussx.solve dispatches on operator type (Woodbury for :class:gaussx.LowRankUpdate, Cholesky for dense, etc.). target Target signature t(ν), shape (n_bands,) — typically the radiance Jacobian ∂L/∂VMR at the background state.

Returns

Float[Array, ""] Scalar score — an unbiased linear estimator of the plume amplitude under the additive-Gaussian-noise model.

Source code in src/radtrax/retrieval/matched_filter/core.py

def apply_pixel(
    pixel: Float[ArrayLike, " B"],
    mean: Float[ArrayLike, " B"],
    cov_op: LinearOperator,
    target: Float[ArrayLike, " B"],
) -> Float[Array, ""]:
    """Matched-filter score for a single pixel spectrum.

    Parameters
    ----------
    pixel
        Observed radiance spectrum at one pixel, shape ``(n_bands,)``.
    mean
        Background mean spectrum, shape ``(n_bands,)``.
    cov_op
        Background covariance as a :class:`lineax.AbstractLinearOperator`. Any
        structure is fine — :func:`gaussx.solve` dispatches on operator type
        (Woodbury for :class:`gaussx.LowRankUpdate`, Cholesky for dense, etc.).
    target
        Target signature ``t(ν)``, shape ``(n_bands,)`` — typically the
        radiance Jacobian ``∂L/∂VMR`` at the background state.

    Returns
    -------
    Float[Array, ""]
        Scalar score — an unbiased linear estimator of the plume amplitude
        under the additive-Gaussian-noise model.
    """
    w, target_norm_sq = _precompute_filter(cov_op, target)
    return jnp.dot(jnp.asarray(pixel) - mean, w) / target_norm_sq

apply_image

apply_image(cube: Float[ArrayLike, 'H W B'], mean: Float[ArrayLike, ' B'], cov_op: LinearOperator, target: Float[ArrayLike, ' B']) -> Float[Array, 'H W']

Matched-filter score map for a hyperspectral cube.

Precomputes w = Σ⁻¹ t once, then evaluates (cube − μ) · w / (tᵀ w) in a single einsum — no per-pixel Python loop.

Parameters

cube Radiance cube with shape (H, W, n_bands). mean, cov_op, target See :func:apply_pixel.

Returns

Float[Array, "H W"] Score image. Pixel (i, j) is the estimated plume amplitude at that location.

Source code in src/radtrax/retrieval/matched_filter/core.py

def apply_image(
    cube: Float[ArrayLike, "H W B"],
    mean: Float[ArrayLike, " B"],
    cov_op: LinearOperator,
    target: Float[ArrayLike, " B"],
) -> Float[Array, "H W"]:
    """Matched-filter score map for a hyperspectral cube.

    Precomputes ``w = Σ⁻¹ t`` once, then evaluates
    ``(cube − μ) · w / (tᵀ w)`` in a single einsum — no per-pixel Python loop.

    Parameters
    ----------
    cube
        Radiance cube with shape ``(H, W, n_bands)``.
    mean, cov_op, target
        See :func:`apply_pixel`.

    Returns
    -------
    Float[Array, "H W"]
        Score image. Pixel ``(i, j)`` is the estimated plume amplitude at
        that location.
    """
    w, target_norm_sq = _precompute_filter(cov_op, target)
    residual = jnp.asarray(cube) - mean
    # Contract the last (spectral) axis with the precomputed filter
    # direction. jnp.einsum keeps the other two axes intact, matching the
    # "H W B, B -> H W" einops signature we replaced.
    scores = jnp.einsum("ijk,k->ij", residual, w)
    return scores / target_norm_sq

matched_filter_snr

matched_filter_snr(amplitude: float, cov_op: LinearOperator, target: Float[ArrayLike, ' B']) -> Float[Array, '']

Theoretical SNR of the matched filter at a given plume amplitude.

For x = μ + α t + ε with ε ∼ N(0, Σ), the MF score has mean = α and variance = 1 / (tᵀ Σ⁻¹ t). So the detection SNR is

.. math:: \text{SNR} = \alpha \sqrt{t^\top \Sigma^{-1} t}.

Parameters

amplitude The plume amplitude α at which to evaluate SNR. cov_op, target See :func:apply_pixel.

Returns

Float[Array, ""] Scalar SNR (dimensionless).

Source code in src/radtrax/retrieval/matched_filter/core.py

def matched_filter_snr(
    amplitude: float,
    cov_op: LinearOperator,
    target: Float[ArrayLike, " B"],
) -> Float[Array, ""]:
    """Theoretical SNR of the matched filter at a given plume amplitude.

    For ``x = μ + α t + ε`` with ``ε ∼ N(0, Σ)``, the MF score has
    ``mean = α`` and ``variance = 1 / (tᵀ Σ⁻¹ t)``. So the detection SNR is

    .. math::
        \\text{SNR} = \\alpha \\sqrt{t^\\top \\Sigma^{-1} t}.

    Parameters
    ----------
    amplitude
        The plume amplitude ``α`` at which to evaluate SNR.
    cov_op, target
        See :func:`apply_pixel`.

    Returns
    -------
    Float[Array, ""]
        Scalar SNR (dimensionless).
    """
    w = gx.solve(cov_op, target)
    target_norm_sq = jnp.dot(target, w)
    _check_target_energy(target, target_norm_sq)
    return jnp.asarray(amplitude) * jnp.sqrt(target_norm_sq)

detection_threshold

detection_threshold(false_alarm_rate: float, cov_op: LinearOperator, target: Float[ArrayLike, ' B']) -> Float[Array, '']

Gaussian detection threshold on the MF score for a given FAR.

Under the null (no plume), the score is :math:N(0, 1/(t^\top \Sigma^{-1} t)). The threshold above which the false-alarm rate does not exceed FAR is

.. math:: \tau = \Phi^{-1}(1 - \text{FAR}) / \sqrt{t^\top \Sigma^{-1} t}.

Parameters

false_alarm_rate Target FAR, in (0, 1) — e.g. 1e-6 for one false alarm per million pixels under Gaussian assumptions. cov_op, target See :func:apply_pixel.

Returns

Float[Array, ""] Threshold in score units (same units as amplitude in :func:matched_filter_snr).

Source code in src/radtrax/retrieval/matched_filter/core.py

def detection_threshold(
    false_alarm_rate: float,
    cov_op: LinearOperator,
    target: Float[ArrayLike, " B"],
) -> Float[Array, ""]:
    """Gaussian detection threshold on the MF score for a given FAR.

    Under the null (no plume), the score is :math:`N(0, 1/(t^\\top \\Sigma^{-1} t))`.
    The threshold above which the false-alarm rate does not exceed ``FAR`` is

    .. math::
        \\tau = \\Phi^{-1}(1 - \\text{FAR}) / \\sqrt{t^\\top \\Sigma^{-1} t}.

    Parameters
    ----------
    false_alarm_rate
        Target FAR, in ``(0, 1)`` — e.g. ``1e-6`` for one false alarm per
        million pixels under Gaussian assumptions.
    cov_op, target
        See :func:`apply_pixel`.

    Returns
    -------
    Float[Array, ""]
        Threshold in score units (same units as ``amplitude`` in
        :func:`matched_filter_snr`).
    """
    if not 0.0 < false_alarm_rate < 1.0:
        raise ParameterError(
            f"detection_threshold: false_alarm_rate must be in (0, 1); "
            f"got {false_alarm_rate}."
        )
    # ndtri is the inverse Gaussian CDF (aka Φ⁻¹ / probit).
    z = jax.scipy.special.ndtri(1.0 - false_alarm_rate)
    w = gx.solve(cov_op, target)
    target_norm_sq = jnp.dot(target, w)
    _check_target_energy(target, target_norm_sq)
    return z / jnp.sqrt(target_norm_sq)

background

Background mean and covariance estimators for the matched filter.

Estimation is not on the gradient path — the matched filter treats the background :math:(\mu, \Sigma) as known inputs. That lets us lean on scikit-learn's mature estimators (which have careful numerics, shrinkage, robust variants, etc.) without any JAX/autodiff constraints. The outputs are wrapped as :mod:lineax / :mod:gaussx linear operators so that :func:radtrax.retrieval.matched_filter.core.apply_image can use the same gaussx.solve dispatch path for all covariance flavours:

============================== ==================================================== Estimator Operator returned ============================== ==================================================== estimate_cov_empirical :class:lineax.MatrixLinearOperator (dense) estimate_cov_shrunk :class:lineax.MatrixLinearOperator (dense, shrunk) estimate_cov_lowrank :class:gaussx.LowRankUpdate (λI + UDUᵀ) ============================== ====================================================

The low-rank path uses scikit-learn's randomised :class:TruncatedSVD, which is O(n · k) on a (n_samples, n_bands) matrix and is typically 10–100× faster than a full SVD for the small k a matched filter needs.

Ported from plumax 0.1.2 at 28f257e (src/plumax/matched_filter/background.py); only import paths changed.

estimate_mean

estimate_mean(cube: Float[ArrayLike, 'H W B'], *, method: _RobustMethod = 'mean', trim_proportion: float = 0.1, huber_c: float = 1.345) -> Float[ndarray, ' B']

Estimate the per-band background mean μ from a scene.

Parameters

cube Hyperspectral scene of shape (H, W, n_bands). method - 'mean' — arithmetic mean (sensitive to anomalies). - 'median' — per-band median (very robust). - 'trimmed' — two-sided trimmed mean (scipy.stats.trim_mean). - 'huber' — Huber M-estimator (sklearn's HuberRegressor on a constant design). trim_proportion Fraction discarded from each tail for the trimmed mean. Ignored otherwise. huber_c Huber influence threshold in units of the per-band MAD. Ignored unless method='huber'.

Pixels with a non-finite band (NaN fill, saturated or missing data) are left out, as in every estimator of this module.

Returns

np.ndarray Mean spectrum of length n_bands.

Source code in src/radtrax/retrieval/matched_filter/background.py

def estimate_mean(
    cube: Float[ArrayLike, "H W B"],
    *,
    method: _RobustMethod = "mean",
    trim_proportion: float = 0.1,
    huber_c: float = 1.345,
) -> Float[np.ndarray, " B"]:
    """Estimate the per-band background mean ``μ`` from a scene.

    Parameters
    ----------
    cube
        Hyperspectral scene of shape ``(H, W, n_bands)``.
    method
        - ``'mean'``    — arithmetic mean (sensitive to anomalies).
        - ``'median'``  — per-band median (very robust).
        - ``'trimmed'`` — two-sided trimmed mean (``scipy.stats.trim_mean``).
        - ``'huber'``   — Huber M-estimator (sklearn's ``HuberRegressor`` on
          a constant design).
    trim_proportion
        Fraction discarded from each tail for the trimmed mean. Ignored
        otherwise.
    huber_c
        Huber influence threshold in units of the per-band MAD. Ignored
        unless ``method='huber'``.

    Pixels with a non-finite band (NaN fill, saturated or missing data) are
    left out, as in every estimator of this module.

    Returns
    -------
    np.ndarray
        Mean spectrum of length ``n_bands``.
    """
    X = _flatten_cube_masked(cube, None, min_samples=1)
    if method == "mean":
        return X.mean(axis=0)
    if method == "median":
        return np.median(X, axis=0)
    if method == "trimmed":
        from scipy.stats import trim_mean

        return trim_mean(X, proportiontocut=trim_proportion, axis=0)
    if method == "huber":
        return _huber_mean(X, c=huber_c)
    raise ParameterError(f"estimate_mean: unknown method {method!r}.")

estimate_cov_empirical

estimate_cov_empirical(cube: Float[ArrayLike, 'H W B'], mean: Float[ArrayLike, ' B'] | None = None, *, ridge: float = 0.0, mask: Bool[ArrayLike, '*pixels'] | None = None) -> LinearOperator

Empirical covariance via :class:sklearn.covariance.EmpiricalCovariance.

Returns a dense :class:lineax.MatrixLinearOperator. Optionally adds a diagonal ridge for numerical PD-ness when n_samples ≲ n_bands. mask (boolean exclude-pixels) drops flagged pixels — e.g. a detected plume — so Σ is estimated from background only, consistent with a robust masked μ.

Source code in src/radtrax/retrieval/matched_filter/background.py

def estimate_cov_empirical(
    cube: Float[ArrayLike, "H W B"],
    mean: Float[ArrayLike, " B"] | None = None,
    *,
    ridge: float = 0.0,
    mask: Bool[ArrayLike, "*pixels"] | None = None,
) -> LinearOperator:
    """Empirical covariance via :class:`sklearn.covariance.EmpiricalCovariance`.

    Returns a dense :class:`lineax.MatrixLinearOperator`. Optionally adds a
    diagonal ``ridge`` for numerical PD-ness when ``n_samples ≲ n_bands``.
    ``mask`` (boolean exclude-pixels) drops flagged pixels — e.g. a detected
    plume — so ``Σ`` is estimated from background only, consistent with a robust
    masked ``μ``.
    """
    from sklearn.covariance import EmpiricalCovariance

    X = _flatten_cube_masked(cube, mask)
    est = EmpiricalCovariance(store_precision=False).fit(X)
    cov = est.covariance_
    if mean is not None:
        # sklearn uses the empirical mean internally; re-centre if the caller
        # supplied a robust mean so (μ, Σ) are consistent.
        Xc = X - mean
        cov = (Xc.T @ Xc) / X.shape[0]
    if ridge > 0.0:
        cov = cov + ridge * np.eye(cov.shape[0])
    return _dense_operator(cov)

estimate_cov_shrunk

estimate_cov_shrunk(cube: Float[ArrayLike, 'H W B'], mean: Float[ArrayLike, ' B'] | None = None, *, method: Literal['ledoit_wolf', 'oas'] = 'ledoit_wolf', mask: Bool[ArrayLike, '*pixels'] | None = None) -> LinearOperator

Shrinkage covariance estimator.

Uses Ledoit–Wolf (method='ledoit_wolf', default) or OAS ('oas') — both shrink the sample covariance toward a scaled identity and are PD by construction even when n_samples < n_bands. mask (boolean exclude-pixels) drops flagged pixels before fitting.

Source code in src/radtrax/retrieval/matched_filter/background.py

def estimate_cov_shrunk(
    cube: Float[ArrayLike, "H W B"],
    mean: Float[ArrayLike, " B"] | None = None,
    *,
    method: Literal["ledoit_wolf", "oas"] = "ledoit_wolf",
    mask: Bool[ArrayLike, "*pixels"] | None = None,
) -> LinearOperator:
    """Shrinkage covariance estimator.

    Uses Ledoit–Wolf (``method='ledoit_wolf'``, default) or OAS
    (``'oas'``) — both shrink the sample covariance toward a scaled identity
    and are PD by construction even when ``n_samples < n_bands``. ``mask``
    (boolean exclude-pixels) drops flagged pixels before fitting.
    """
    X = _flatten_cube_masked(cube, mask)
    if mean is not None:
        X = X - mean
        assume_centered = True
    else:
        assume_centered = False
    if method == "ledoit_wolf":
        from sklearn.covariance import LedoitWolf

        est = LedoitWolf(assume_centered=assume_centered).fit(X)
    elif method == "oas":
        from sklearn.covariance import OAS

        est = OAS(assume_centered=assume_centered).fit(X)
    else:
        raise ParameterError(
            f"estimate_cov_shrunk: method must be 'ledoit_wolf' or 'oas'; got {method!r}."
        )
    return _dense_operator(est.covariance_)

shrink_covariance

shrink_covariance(empirical: Float[ArrayLike, 'B B'], *, method: Literal['ledoit_wolf', 'oas'], n_samples: int, fourth_moment: float | None = None) -> Float[ndarray, 'B B']

Shrink a 1/n empirical covariance toward a scaled identity.

Returns (1 − s)·Sₙ + s·μ·I with μ = tr(Sₙ)/p, where Sₙ is the maximum-likelihood (1/n) scatter of n samples in p bands — the normalisation every estimator in this module uses — and the intensity s ∈ [0, 1] is estimated from the data. Both methods reproduce scikit-learn's :class:~sklearn.covariance.LedoitWolf and :class:~sklearn.covariance.OAS exactly. This is the form a streaming pass needs (:func:~radtrax.retrieval.matched_filter.streaming.streaming_background), where only the moments, not the samples, are kept.

• "ledoit_wolf" (Ledoit & Wolf 2004), with dₖ = xₖ − x̄ and m₄ = (1/n) Σₖ ‖dₖ‖⁴::

  δ² = ‖Sₙ − μ I‖²_F β̄² = (1/n²) Σₖ ‖dₖdₖᵀ − Sₙ‖²_F = (m₄ − ‖Sₙ‖²_F) / n s = min(β̄², δ²) / δ² (0 when δ² = 0)

β̄² needs the fourth moment m₄, which Sₙ alone does not carry.

• "oas" (Chen et al. 2010, as scikit-learn implements it, without the 2/p terms of their Eq. 23)::

  s = min(1, (tr(Sₙ²) + tr(Sₙ)²) / ((n + 1)·(tr(Sₙ²) − tr(Sₙ)²/p)))

Parameters

empirical 1/n empirical covariance Σₖ dₖdₖᵀ / n, shape (B, B). method Shrinkage-intensity estimator. n_samples Number of samples behind empirical. fourth_moment m₄ over the same centred samples; required for "ledoit_wolf", ignored by "oas". :meth:~radtrax.retrieval.matched_filter.streaming.WelfordAccumulator.fourth_moment tracks it.

Examples

import numpy as np x = np.random.default_rng(0).normal(size=(30, 4)) d = x - x.mean(axis=0) s_n = d.T @ d / 30 m4 = float(np.mean(np.sum(d * d, axis=1) ** 2)) shrink_covariance( ... s_n, method="ledoit_wolf", n_samples=30, fourth_moment=m4 ... ).shape (4, 4)

Source code in src/radtrax/retrieval/matched_filter/background.py

def shrink_covariance(
    empirical: Float[ArrayLike, "B B"],
    *,
    method: Literal["ledoit_wolf", "oas"],
    n_samples: int,
    fourth_moment: float | None = None,
) -> Float[np.ndarray, "B B"]:
    """Shrink a ``1/n`` empirical covariance toward a scaled identity.

    Returns ``(1 − s)·Sₙ + s·μ·I`` with ``μ = tr(Sₙ)/p``, where ``Sₙ`` is the
    maximum-likelihood (``1/n``) scatter of ``n`` samples in ``p`` bands — the
    normalisation every estimator in this module uses — and the intensity
    ``s ∈ [0, 1]`` is estimated from the data. Both methods reproduce
    scikit-learn's :class:`~sklearn.covariance.LedoitWolf` and
    :class:`~sklearn.covariance.OAS` exactly. This is the form a streaming
    pass needs (:func:`~radtrax.retrieval.matched_filter.streaming.streaming_background`),
    where only the moments, not the samples, are kept.

    - ``"ledoit_wolf"`` (Ledoit & Wolf 2004), with ``dₖ = xₖ − x̄`` and
      ``m₄ = (1/n) Σₖ ‖dₖ‖⁴``::

          δ² = ‖Sₙ − μ I‖²_F
          β̄² = (1/n²) Σₖ ‖dₖdₖᵀ − Sₙ‖²_F = (m₄ − ‖Sₙ‖²_F) / n
          s  = min(β̄², δ²) / δ²                    (0 when δ² = 0)

      β̄² needs the fourth moment ``m₄``, which ``Sₙ`` alone does not carry.

    - ``"oas"`` (Chen et al. 2010, as scikit-learn implements it, without
      the ``2/p`` terms of their Eq. 23)::

          s = min(1, (tr(Sₙ²) + tr(Sₙ)²) / ((n + 1)·(tr(Sₙ²) − tr(Sₙ)²/p)))

    Parameters
    ----------
    empirical
        ``1/n`` empirical covariance ``Σₖ dₖdₖᵀ / n``, shape ``(B, B)``.
    method
        Shrinkage-intensity estimator.
    n_samples
        Number of samples behind ``empirical``.
    fourth_moment
        ``m₄`` over the same centred samples; required for
        ``"ledoit_wolf"``, ignored by ``"oas"``.
        :meth:`~radtrax.retrieval.matched_filter.streaming.WelfordAccumulator.fourth_moment`
        tracks it.

    Examples
    --------
    >>> import numpy as np
    >>> x = np.random.default_rng(0).normal(size=(30, 4))
    >>> d = x - x.mean(axis=0)
    >>> s_n = d.T @ d / 30
    >>> m4 = float(np.mean(np.sum(d * d, axis=1) ** 2))
    >>> shrink_covariance(
    ...     s_n, method="ledoit_wolf", n_samples=30, fourth_moment=m4
    ... ).shape
    (4, 4)
    """
    cov = np.asarray(empirical, dtype=float)
    if cov.ndim != 2 or cov.shape[0] != cov.shape[1]:
        raise ShapeError(
            f"shrink_covariance: empirical must be square, got shape {cov.shape}."
        )
    if n_samples < 1:
        raise ParameterError(
            f"shrink_covariance: n_samples must be ≥ 1, got {n_samples}."
        )
    p = cov.shape[0]
    mu = float(np.trace(cov)) / p
    fro_sq = float(np.sum(cov * cov))
    if method == "oas":
        trace_sq = float(np.trace(cov)) ** 2
        denom = (n_samples + 1) * (fro_sq - trace_sq / p)
        shrinkage = 1.0 if denom <= 0 else min(1.0, (fro_sq + trace_sq) / denom)
    elif method == "ledoit_wolf":
        if fourth_moment is None:
            raise ParameterError(
                "shrink_covariance: 'ledoit_wolf' needs fourth_moment = "
                "mean(‖xₖ − x̄‖⁴) over the samples (WelfordAccumulator.fourth_moment())."
            )
        diff = cov - mu * np.eye(p)
        d_sq = float(np.sum(diff * diff))
        b_bar_sq = max((float(fourth_moment) - fro_sq) / n_samples, 0.0)
        b_sq = min(b_bar_sq, d_sq)
        shrinkage = 0.0 if b_sq == 0 else b_sq / d_sq
    else:
        raise ParameterError(
            f"shrink_covariance: method must be 'ledoit_wolf' or 'oas'; got {method!r}."
        )
    return (1.0 - shrinkage) * cov + shrinkage * mu * np.eye(p)

estimate_cov_lowrank

estimate_cov_lowrank(cube: Float[ArrayLike, 'H W B'], mean: Float[ArrayLike, ' B'] | None = None, *, rank: int, tikhonov: float, random_state: int | PRNGKeyArray | None = 0, n_oversamples: int = 10, mask: Bool[ArrayLike, '*pixels'] | None = None, rank_rtol: float = 1e-10) -> LinearOperator

Low-rank + Tikhonov covariance: Σ = λI + V D Vᵀ.

Uses :func:gaussx.randomized_svd (Halko–Martinsson–Tropp, JAX, 5 power iterations) to find the top rank spectral directions of the sample covariance. Returned as a :class:gaussx.LowRankUpdate so that :func:gaussx.solve routes through the Woodbury identity — the MF precompute cost is O(n_bands · rank + rank³) instead of O(n_bands³).

Parameters

cube Hyperspectral scene (H, W, n_bands). mean Background mean. If None, subtracts the sample mean. rank Number of leading components to keep. Clamped to min(rank, n_samples - 1, n_bands - 1). tikhonov Diagonal floor λ > 0 — ensures strict PD. random_state Seed for the randomised SVD: an int (converted with jax.random.key), a JAX PRNG key, or None (seed 0). n_oversamples Extra random directions for the Halko sampling — higher is slower but better-conditioned. mask Boolean exclude-pixels mask; flagged pixels (e.g. a detected plume) are dropped before the SVD so Σ reflects background only. rank_rtol Components whose covariance eigenvalue d_i is below rank_rtol · max(d) are dropped before building the operator. On a rank-deficient scene (constant regions, duplicated spectra) those d_i ≈ 0 would make gaussx's Woodbury capacitance diag(1/d) + VᵀL⁻¹U blow up to inf/NaN; dropping them keeps every score finite.

Source code in src/radtrax/retrieval/matched_filter/background.py

def estimate_cov_lowrank(
    cube: Float[ArrayLike, "H W B"],
    mean: Float[ArrayLike, " B"] | None = None,
    *,
    rank: int,
    tikhonov: float,
    random_state: int | PRNGKeyArray | None = 0,
    n_oversamples: int = 10,
    mask: Bool[ArrayLike, "*pixels"] | None = None,
    rank_rtol: float = 1e-10,
) -> LinearOperator:
    """Low-rank + Tikhonov covariance: ``Σ = λI + V D Vᵀ``.

    Uses :func:`gaussx.randomized_svd` (Halko–Martinsson–Tropp, JAX, 5 power
    iterations) to find the top ``rank`` spectral directions of
    the sample covariance. Returned as a :class:`gaussx.LowRankUpdate` so that
    :func:`gaussx.solve` routes through the Woodbury identity — the MF
    precompute cost is ``O(n_bands · rank + rank³)`` instead of
    ``O(n_bands³)``.

    Parameters
    ----------
    cube
        Hyperspectral scene ``(H, W, n_bands)``.
    mean
        Background mean. If ``None``, subtracts the sample mean.
    rank
        Number of leading components to keep. Clamped to
        ``min(rank, n_samples - 1, n_bands - 1)``.
    tikhonov
        Diagonal floor ``λ > 0`` — ensures strict PD.
    random_state
        Seed for the randomised SVD: an ``int`` (converted with
        ``jax.random.key``), a JAX PRNG key, or ``None`` (seed 0).
    n_oversamples
        Extra random directions for the Halko sampling — higher is slower but
        better-conditioned.
    mask
        Boolean exclude-pixels mask; flagged pixels (e.g. a detected plume) are
        dropped before the SVD so ``Σ`` reflects background only.
    rank_rtol
        Components whose covariance eigenvalue ``d_i`` is below
        ``rank_rtol · max(d)`` are dropped before building the operator. On a
        rank-deficient scene (constant regions, duplicated spectra) those
        ``d_i ≈ 0`` would make gaussx's Woodbury capacitance ``diag(1/d) +
        VᵀL⁻¹U`` blow up to inf/NaN; dropping them keeps every score finite.
    """
    if tikhonov <= 0.0:
        raise ParameterError("estimate_cov_lowrank: tikhonov must be > 0.")
    X = _flatten_cube_masked(cube, mask)
    mu = X.mean(axis=0) if mean is None else np.asarray(mean)
    Xc = X - mu
    n_samples, n_bands = Xc.shape
    rank = max(1, min(int(rank), n_samples - 1, n_bands - 1))
    key = _as_key(random_state)
    # Same MLE normalization s_i² / n_samples as the empirical / LedoitWolf /
    # OAS paths in this module (using n_samples - 1 here would make
    # matched_filter_snr and detection_threshold depend on the estimator).
    V, d = lowrank_factors(
        jnp.asarray(Xc), rank, n_oversamples=n_oversamples, n_power_iter=5, key=key
    )
    V, d = np.asarray(V), np.asarray(d)
    diag = jnp.asarray(tikhonov * np.ones(n_bands, dtype=float))
    # Rank-deficiency guard: drop near-zero eigen-directions so the Woodbury
    # capacitance stays finite. If none survive (e.g. a constant scene), Σ is
    # just the strictly-PD Tikhonov floor λI.
    d_max = float(d.max()) if d.size else 0.0
    keep = d > rank_rtol * d_max if d_max > 0.0 else np.zeros(d.shape, dtype=bool)
    V, d = V[:, keep], d[keep]
    if d.size == 0:
        return lx.DiagonalLinearOperator(diag)
    # λ I + U diag(d) Uᵀ via the gaussx SVD constructor (auto-inferred
    # symmetric/PSD tags, orthonormal-factor solve path).
    U = jnp.asarray(V)  # (n_bands, n_kept) — spectral directions as columns
    return gx.svd_low_rank_plus_diag(diag, U, jnp.asarray(d), U)

cluster

Cluster-based and spatially-adaptive background estimators.

When a scene is heterogeneous (land/water, bright/dark surfaces, different land-cover classes) a single global (μ, Σ) smears together statistically distinct populations and the MF loses power. Two standard fixes:

• GMM clustering — fit a Gaussian mixture on the per-pixel spectra, label each pixel, and run the MF with the local cluster's (μ_k, Σ_k). :func:gmm_cluster_background returns a :class:ClusterBackground: the labels and one (mean, cov_operator) pair per cluster.
• Adaptive spatial windows — for each pixel, estimate (μ_{ij}, \sigma^2_{ij}) from a local neighbourhood. Only supports a diagonal local covariance (keeps memory at O(H·W·n_bands)) — for full local covariance, prefer the GMM route.

Both are estimated with scikit-learn / scipy, off the JAX gradient path. The fitted :class:ClusterBackground is an Equinox module, so the filter can still be differentiated with respect to its means and covariances.

Ported from plumax 0.1.2 at 28f257e (src/plumax/matched_filter/cluster.py); its frozen dataclass became an Equinox module in radtrax#59.

ClusterBackground

Bases: Module

Per-cluster background model.

An Equinox module (radtrax#59): the labels, means and covariance operators are leaves. It passes through eqx.filter_jit with the image, and through eqx.filter_vmap over a batch of images with in_axes=(None, 0) (the default would map the background's own arrays). Two backgrounds share one compilation when they have the same number of clusters, operator types and dtypes; the number of clusters depends on the data, not on n_clusters.

Under-populated clusters reuse the global (μ, Σ), so one array or operator can appear at several positions. Each position is still its own leaf: eqx.tree_at and gradients act per position, so the gradient with respect to the shared global mean is the sum over those positions.

Attributes

labels Cluster assignment per pixel, shape (H, W) with values in [0, n_clusters); -1 marks a pixel with a non-finite band, left out of the fit (:func:apply_cluster_mf scores it NaN). means Tuple of per-cluster mean spectra. cov_operators Tuple of per-cluster covariance operators (same structure across clusters — dense from :func:estimate_cov_empirical or :func:estimate_cov_shrunk).

gmm_cluster_background

gmm_cluster_background(cube: Float[ArrayLike, 'H W B'], *, n_clusters: int, cov_estimator: str = 'ledoit_wolf', random_state: int | None = 0, bayesian: bool = False) -> ClusterBackground

Fit a Gaussian mixture and build per-cluster background models.

Uses :class:sklearn.mixture.GaussianMixture (or :class:BayesianGaussianMixture when bayesian=True) on the per-pixel spectra. Each cluster gets its own shrunk covariance estimate so that clusters with few pixels remain well-conditioned.

Parameters

cube Hyperspectral cube, shape (H, W, n_bands). n_clusters Number of mixture components (upper bound for the Bayesian variant). cov_estimator 'ledoit_wolf' (default), 'oas', or 'empirical'. random_state Passed through to sklearn for reproducibility. bayesian If True, use Dirichlet-process-style :class:BayesianGaussianMixture, which prunes unused components automatically.

Returns

ClusterBackground The mixture is fitted on the pixels whose bands are all finite; the rest are labelled -1.

Source code in src/radtrax/retrieval/matched_filter/cluster.py

def gmm_cluster_background(
    cube: Float[ArrayLike, "H W B"],
    *,
    n_clusters: int,
    cov_estimator: str = "ledoit_wolf",
    random_state: int | None = 0,
    bayesian: bool = False,
) -> ClusterBackground:
    """Fit a Gaussian mixture and build per-cluster background models.

    Uses :class:`sklearn.mixture.GaussianMixture` (or
    :class:`BayesianGaussianMixture` when ``bayesian=True``) on the per-pixel
    spectra. Each cluster gets its own shrunk covariance estimate so that
    clusters with few pixels remain well-conditioned.

    Parameters
    ----------
    cube
        Hyperspectral cube, shape ``(H, W, n_bands)``.
    n_clusters
        Number of mixture components (upper bound for the Bayesian variant).
    cov_estimator
        ``'ledoit_wolf'`` (default), ``'oas'``, or ``'empirical'``.
    random_state
        Passed through to sklearn for reproducibility.
    bayesian
        If True, use Dirichlet-process-style :class:`BayesianGaussianMixture`,
        which prunes unused components automatically.

    Returns
    -------
    ClusterBackground
        The mixture is fitted on the pixels whose bands are all finite; the
        rest are labelled ``-1``.
    """
    arr = np.asarray(cube)
    if arr.ndim != 3:
        raise ShapeError(
            f"gmm_cluster_background: cube must be (H, W, n_bands); got {arr.shape}."
        )
    h, w, _ = arr.shape
    X_all = arr.reshape(-1, arr.shape[-1])
    valid = _finite_pixels(X_all)
    X = X_all[valid]
    if X.shape[0] < max(2, n_clusters):
        raise DegenerateInputError(
            f"gmm_cluster_background: {X.shape[0]} finite pixel(s) cannot fit "
            f"{n_clusters} cluster(s)."
        )
    if bayesian:
        from sklearn.mixture import BayesianGaussianMixture

        gmm = BayesianGaussianMixture(
            n_components=n_clusters,
            covariance_type="full",
            random_state=random_state,
        )
    else:
        from sklearn.mixture import GaussianMixture

        gmm = GaussianMixture(
            n_components=n_clusters,
            covariance_type="full",
            random_state=random_state,
        )
    labels_valid = gmm.fit_predict(X)
    # Pre-compute a global (μ, Σ) fallback once so under-populated clusters
    # get a *consistent* pair — the old code returned the per-cluster mean
    # paired with a scene-wide Ledoit–Wolf Σ, which biases scores.
    global_mean = X.mean(axis=0)
    if cov_estimator == "empirical":
        global_cov_op = estimate_cov_empirical(arr, mean=global_mean)
    elif cov_estimator in {"ledoit_wolf", "oas"}:
        global_cov_op = estimate_cov_shrunk(arr, mean=global_mean, method=cov_estimator)  # type: ignore[arg-type]
    else:
        raise ParameterError(
            f"gmm_cluster_background: unknown cov_estimator {cov_estimator!r}."
        )
    means: list[np.ndarray] = []
    ops: list[LinearOperator] = []
    for k in range(int(labels_valid.max()) + 1):
        mask = labels_valid == k
        if mask.sum() < 2:
            # Under-populated cluster: return the *same* (μ, Σ) pair for both
            # attributes so callers that use (means[k], cov_operators[k])
            # together see a self-consistent global-scene fallback.
            means.append(global_mean)
            ops.append(global_cov_op)
            continue
        # Reshape the flat cluster subset into a fake (n_k, 1, n_bands) cube
        # so the cube-shaped estimators accept it without dedicated per-cluster
        # code paths.
        fake_cube = X[mask].reshape(-1, 1, X.shape[1])
        mu_k = X[mask].mean(axis=0)
        if cov_estimator == "empirical":
            cov_op = estimate_cov_empirical(fake_cube, mean=mu_k)
        else:
            cov_op = estimate_cov_shrunk(fake_cube, mean=mu_k, method=cov_estimator)  # type: ignore[arg-type]
        means.append(mu_k)
        ops.append(cov_op)
    labels = np.full(X_all.shape[0], -1, dtype=int)
    labels[valid] = labels_valid
    # Tuples of JAX arrays: an immutable structure (a list could be appended
    # to after construction, changing the compile key) with one leaf type.
    return ClusterBackground(
        labels=jnp.asarray(labels.reshape(h, w)),
        means=tuple(jnp.asarray(m) for m in means),
        cov_operators=tuple(ops),
    )

adaptive_window_background

adaptive_window_background(cube: Float[ArrayLike, 'H W B'], *, window_size: int, guard_size: int = 0) -> tuple[Float[ndarray, 'H W B'], Float[ndarray, 'H W B']]

Per-pixel local mean and per-band variance over a square window.

For each pixel (i, j), estimate μ_{ij} and a diagonal local covariance diag(σ²_{ij,b}) using a window_size × window_size uniform window centred on the pixel (reflecting at boundaries).

Full local covariances would be (H · W · n_bands²) — too big for realistic scenes. For a dense local covariance, cluster with GMM instead.

Parameters

cube Hyperspectral cube, shape (H, W, n_bands). window_size Edge length of the averaging window. Must be odd and >= 3. guard_size Edge length of a central guard region excluded from each pixel's window (a guard annulus). Must be odd and 0 <= guard_size < window_size. The default 0 includes the centre pixel — but then a plume pixel contaminates its own local background; guard_size=1 excludes just the centre pixel, larger values exclude a bigger core.

Returns

local_mean Per-pixel mean spectrum, shape (H, W, n_bands). local_variance Per-pixel per-band (population) variance, shape (H, W, n_bands). Window statistics use valid pixels only (a pixel with any non-finite band carries no weight); invalid pixels, and pixels whose window holds fewer than two valid pixels (variance) or none (mean), are NaN.

Source code in src/radtrax/retrieval/matched_filter/cluster.py

def adaptive_window_background(
    cube: Float[ArrayLike, "H W B"],
    *,
    window_size: int,
    guard_size: int = 0,
) -> tuple[Float[np.ndarray, "H W B"], Float[np.ndarray, "H W B"]]:
    """Per-pixel local mean and per-band variance over a square window.

    For each pixel ``(i, j)``, estimate ``μ_{ij}`` and a *diagonal* local
    covariance ``diag(σ²_{ij,b})`` using a ``window_size × window_size``
    uniform window centred on the pixel (reflecting at boundaries).

    Full local covariances would be ``(H · W · n_bands²)`` — too big for
    realistic scenes. For a dense local covariance, cluster with GMM instead.

    Parameters
    ----------
    cube
        Hyperspectral cube, shape ``(H, W, n_bands)``.
    window_size
        Edge length of the averaging window. Must be odd and ``>= 3``.
    guard_size
        Edge length of a central **guard region** excluded from each pixel's
        window (a guard annulus). Must be odd and ``0 <= guard_size <
        window_size``. The default ``0`` includes the centre pixel — but then a
        plume pixel contaminates its *own* local background; ``guard_size=1``
        excludes just the centre pixel, larger values exclude a bigger core.

    Returns
    -------
    local_mean
        Per-pixel mean spectrum, shape ``(H, W, n_bands)``.
    local_variance
        Per-pixel per-band (population) variance, shape ``(H, W, n_bands)``.
        Window statistics use valid pixels only (a pixel with any non-finite
        band carries no weight); invalid pixels, and pixels whose window
        holds fewer than two valid pixels (variance) or none (mean), are
        ``NaN``.
    """
    if window_size < 3 or window_size % 2 == 0:
        raise ParameterError(
            f"adaptive_window_background: window_size must be an odd int ≥ 3; got {window_size}."
        )
    if guard_size < 0 or (guard_size > 0 and guard_size % 2 == 0):
        raise ParameterError(
            f"adaptive_window_background: guard_size must be 0 or an odd int; got {guard_size}."
        )
    if guard_size >= window_size:
        raise ParameterError(
            f"adaptive_window_background: guard_size ({guard_size}) must be < "
            f"window_size ({window_size})."
        )
    arr = np.asarray(cube, dtype=float)
    if arr.ndim != 3:
        raise ShapeError(
            f"adaptive_window_background: cube must be (H, W, n_bands); got {arr.shape}."
        )
    from scipy.ndimage import uniform_filter

    # Annulus counts: outer window minus the excluded central guard region.
    n_out = window_size**2
    n_in = guard_size**2 if guard_size > 0 else 0
    size_out = (window_size, window_size, 1)
    size_in = (guard_size, guard_size, 1)

    def window_sum(field: np.ndarray) -> np.ndarray:
        # Box sum over the annulus (outer window minus the guard region),
        # band by band; reflecting at the scene edges.
        total = uniform_filter(field, size=size_out, mode="reflect") * n_out
        if guard_size > 0:
            total = total - uniform_filter(field, size=size_in, mode="reflect") * n_in
        return total

    # Normalised convolution: invalid pixels (any band non-finite) get zero
    # weight, so a NaN fill never leaks into its neighbours' statistics.
    valid = _finite_pixels(arr)[..., None]
    weights = np.broadcast_to(valid, arr.shape).astype(float)
    # Shift each band by a reference level before forming E[x²] − E[x]²: the
    # variance is shift-invariant, and without the shift a radiance offset
    # large against its spread cancels catastrophically (4e-4 relative error
    # at mean 1e4, std 1e-2).
    ref = np.median(arr[valid[..., 0]], axis=0) if valid.any() else 0.0
    filled = np.where(valid, arr - ref, 0.0)
    count = np.rint(window_sum(weights))
    with np.errstate(invalid="ignore", divide="ignore"):
        mean = window_sum(filled) / count
        mean_sq = window_sum(filled * filled) / count
    # Population variance — E[x²] − E[x]², clipped at zero for round-off.
    var = np.maximum(mean_sq - mean * mean, 0.0)
    mean = mean + ref
    mean = np.where((count >= 1) & valid, mean, np.nan)
    var = np.where((count >= 2) & valid, var, np.nan)
    return mean, var

apply_cluster_mf

apply_cluster_mf(cube: Float[ArrayLike, 'H W B'], cluster: ClusterBackground, target: Float[ArrayLike, ' B']) -> Float[Array, 'H W']

Matched-filter scores with each pixel's own cluster background.

Pixel (i, j) with label k scores (x − μ_k)ᵀ Σ_k⁻¹ t / (tᵀ Σ_k⁻¹ t), the :func:~radtrax.retrieval.matched_filter.core.apply_image score with the cluster's (μ_k, Σ_k). One solve per cluster; the dispatch is a jnp.where per cluster, so the function runs under eqx.filter_jit and is differentiable in the cube, the target and the background.

Parameters

cube Radiance cube (H, W, n_bands), the same spatial shape as cluster.labels. cluster Background from :func:gmm_cluster_background. target Target signature (n_bands,).

Returns

Float[Array, "H W"] Score map; NaN where the label is outside [0, n_clusters) (the -1 of a pixel with a non-finite band).

Raises

ValueError If the label map does not match the cube, or (outside jit) a cluster's tᵀ Σ_k⁻¹ t is not positive.

Source code in src/radtrax/retrieval/matched_filter/cluster.py

def apply_cluster_mf(
    cube: Float[ArrayLike, "H W B"],
    cluster: ClusterBackground,
    target: Float[ArrayLike, " B"],
) -> Float[Array, "H W"]:
    """Matched-filter scores with each pixel's own cluster background.

    Pixel ``(i, j)`` with label ``k`` scores
    ``(x − μ_k)ᵀ Σ_k⁻¹ t / (tᵀ Σ_k⁻¹ t)``, the
    :func:`~radtrax.retrieval.matched_filter.core.apply_image` score with
    the cluster's ``(μ_k, Σ_k)``. One solve per cluster; the dispatch is a
    ``jnp.where`` per cluster, so the function runs under
    ``eqx.filter_jit`` and is differentiable in the cube, the target and the
    background.

    Parameters
    ----------
    cube
        Radiance cube ``(H, W, n_bands)``, the same spatial shape as
        ``cluster.labels``.
    cluster
        Background from :func:`gmm_cluster_background`.
    target
        Target signature ``(n_bands,)``.

    Returns
    -------
    Float[Array, "H W"]
        Score map; ``NaN`` where the label is outside ``[0, n_clusters)``
        (the ``-1`` of a pixel with a non-finite band).

    Raises
    ------
    ValueError
        If the label map does not match the cube, or (outside ``jit``) a
        cluster's ``tᵀ Σ_k⁻¹ t`` is not positive.
    """
    cube = jnp.asarray(cube)
    if cube.shape[:2] != tuple(cluster.labels.shape):
        raise ShapeError(
            f"apply_cluster_mf: labels shape {tuple(cluster.labels.shape)} does not "
            f"match the cube's spatial shape {cube.shape[:2]}."
        )
    # Zero the residual of a pixel with a non-finite band before scoring: a
    # NaN there would reach the gradient in w (0 · NaN) even though the
    # where() below never selects that pixel's score.
    finite = jnp.isfinite(cube).all(axis=-1, keepdims=True)
    scores = jnp.full(cube.shape[:2], jnp.nan, dtype=cube.dtype)
    for k, (mean_k, cov_k) in enumerate(
        zip(cluster.means, cluster.cov_operators, strict=True)
    ):
        # A non-finite pixel is replaced by the mean: a zero residual.
        score_k = apply_image(jnp.where(finite, cube, mean_k), mean_k, cov_k, target)
        scores = jnp.where((cluster.labels == k) & finite[..., 0], score_k, scores)
    return scores

apply_adaptive_mf

apply_adaptive_mf(cube: Float[ArrayLike, 'H W B'], mean: Float[ArrayLike, 'H W B'], variance: Float[ArrayLike, 'H W B'], target: Float[ArrayLike, ' B'], *, ridge: float | Float[Array, ''] = 0.0) -> Float[Array, 'H W']

Matched filter against a local diagonal-covariance background.

Scores every pixel against its own window statistics from :func:adaptive_window_background, μ(r) and Λ(r) = diag(σ²(r) + ridge), with the amplitude normalisation of :func:~radtrax.retrieval.matched_filter.core.apply_image::

α(r) = tᵀ Λ(r)⁻¹ (x(r) − μ(r)) / (tᵀ Λ(r)⁻¹ t)
     = Σ_b t_b (x_b − μ_b) / σ²_b  /  Σ_b t_b² / σ²_b

the per-pixel maximum-likelihood amplitude under an independent-band Gaussian background. Runs in JAX and is differentiable; a band with σ² + ridge ≤ 0 makes its pixel unscoreable without putting a NaN in the gradient of the others (double-where).

Parameters

cube Radiance cube (H, W, n_bands). mean, variance Local statistics of the same shape, from :func:adaptive_window_background. target Target signature (n_bands,); must not be all zero. ridge Non-negative variance floor added to every σ²; use it when flat windows give zero variance.

Returns

Float[Array, "H W"] Score map; NaN where the pixel, its local mean or its local variance is non-finite, or a band's σ² + ridge is not positive.

Raises

ValueError If the shapes disagree, ridge is negative, or (outside jit) the target is all zero.

Source code in src/radtrax/retrieval/matched_filter/cluster.py

def apply_adaptive_mf(
    cube: Float[ArrayLike, "H W B"],
    mean: Float[ArrayLike, "H W B"],
    variance: Float[ArrayLike, "H W B"],
    target: Float[ArrayLike, " B"],
    *,
    ridge: float | Float[Array, ""] = 0.0,
) -> Float[Array, "H W"]:
    """Matched filter against a local diagonal-covariance background.

    Scores every pixel against its own window statistics from
    :func:`adaptive_window_background`, ``μ(r)`` and
    ``Λ(r) = diag(σ²(r) + ridge)``, with the amplitude normalisation of
    :func:`~radtrax.retrieval.matched_filter.core.apply_image`::

        α(r) = tᵀ Λ(r)⁻¹ (x(r) − μ(r)) / (tᵀ Λ(r)⁻¹ t)
             = Σ_b t_b (x_b − μ_b) / σ²_b  /  Σ_b t_b² / σ²_b

    the per-pixel maximum-likelihood amplitude under an independent-band
    Gaussian background. Runs in JAX and is differentiable; a band with
    ``σ² + ridge ≤ 0`` makes its pixel unscoreable without putting a NaN in
    the gradient of the others (double-``where``).

    Parameters
    ----------
    cube
        Radiance cube ``(H, W, n_bands)``.
    mean, variance
        Local statistics of the same shape, from
        :func:`adaptive_window_background`.
    target
        Target signature ``(n_bands,)``; must not be all zero.
    ridge
        Non-negative variance floor added to every ``σ²``; use it when flat
        windows give zero variance.

    Returns
    -------
    Float[Array, "H W"]
        Score map; ``NaN`` where the pixel, its local mean or its local
        variance is non-finite, or a band's ``σ² + ridge`` is not positive.

    Raises
    ------
    ValueError
        If the shapes disagree, ``ridge`` is negative, or (outside ``jit``)
        the target is all zero.
    """
    # ridge is a configuration value: checked when concrete; a traced ridge
    # (jax.jit with ridge as an argument) skips the check.
    if isinstance(ridge, (int, float)) and ridge < 0:
        raise ParameterError(f"apply_adaptive_mf: ridge must be ≥ 0, got {ridge}.")
    cube, mean, variance = (jnp.asarray(a) for a in (cube, mean, variance))
    target = jnp.asarray(target)
    for name, stat in (("mean", mean), ("variance", variance)):
        if stat.shape != cube.shape:
            raise ShapeError(
                f"apply_adaptive_mf: background {name} shape {stat.shape} does not "
                f"match the cube shape {cube.shape}."
            )
    if target.shape != cube.shape[-1:]:
        raise ShapeError(
            f"apply_adaptive_mf: target shape {target.shape} does not match "
            f"{cube.shape[-1]} bands."
        )
    _check_target_energy(target, jnp.asarray(1.0))
    v = variance + ridge
    residual = cube - mean
    # Double where: unscoreable bands get a safe 1/1 and a zero residual, so
    # neither the value nor the gradient of any pixel picks up a NaN.
    ok = (v > 0) & jnp.isfinite(residual)  # v > 0 is False for NaN too
    inv_var = jnp.where(ok, 1.0 / jnp.where(ok, v, 1.0), 0.0)
    numer = jnp.einsum("ijb,b->ij", jnp.where(ok, residual, 0.0) * inv_var, target)
    denom = jnp.einsum("ijb,b->ij", inv_var, target * target)
    scoreable = ok.all(axis=-1) & (denom > 0)
    return jnp.where(scoreable, numer / jnp.where(scoreable, denom, 1.0), jnp.nan)

streaming

Streaming background estimation via Welford's algorithm.

For multi-scene MF workflows (e.g. aggregating background statistics across a campaign of flights), the full (n_samples, n_bands) stack does not fit in memory. The parallel Welford update [#welford1962] [#chan1979] maintains a mean and (scaled) covariance in one pass, merges partial results from separate chunks exactly, and is numerically stable:

.. math:: n \leftarrow n_A + n_B \delta = \mu_B - \mu_A \mu \leftarrow \mu_A + \delta \cdot n_B / n M_2 \leftarrow M_{2,A} + M_{2,B} + \delta \delta^\top \cdot n_A n_B / n

Covariance is then M_2 / (n - ddof); the accumulator also merges the third and fourth moments exactly, so a stream can be shrunk (Ledoit–Wolf, OAS) without a second pass.

.. [#welford1962] Welford (1962), Note on a method for calculating corrected sums of squares and products. .. [#chan1979] Chan, Golub & LeVeque (1979), Updating formulae and a pairwise algorithm for computing sample variances.

Ported from plumax 0.1.2 at 28f257e (src/plumax/matched_filter/streaming.py); the third and fourth moments, non-finite skipping and streaming shrinkage come from geotoolz (radtrax#26).

WelfordAccumulator dataclass

WelfordAccumulator(n_bands: int, count: int = 0)

Online mean, covariance and fourth moment for (n_samples, n_bands) streams.

Use :meth:update to feed one batch at a time; :meth:merge to combine two accumulators computed on disjoint subsets; :meth:mean, :meth:covariance and :meth:fourth_moment to finalise. Rows with a non-finite band are skipped, as in the batch estimators.

Besides M₂ = Σ dₖdₖᵀ (dₖ = xₖ − x̄) it keeps the third-moment vector M₃ = Σ ‖dₖ‖² dₖ and the scalar M₄ = Σ ‖dₖ‖⁴, so a stream can be shrunk with Ledoit–Wolf (:func:shrink_covariance needs M₄ / n). Merging parts X with offsets δ_X = x̄_X − x̄ from the pooled mean is exact::

M₂ = Σ_X [M₂_X + n_X δ_X δ_Xᵀ]
M₃ = Σ_X [M₃_X + tr(M₂_X) δ_X + 2 M₂_X δ_X + n_X ‖δ_X‖² δ_X]
M₄ = Σ_X [M₄_X + 4 δ_Xᵀ M₃_X + 4 δ_Xᵀ M₂_X δ_X
          + 2 ‖δ_X‖² tr(M₂_X) + n_X ‖δ_X‖⁴]

(expand ‖d + δ‖² with d = x − x̄_X, Σ d = 0). Ported from geotoolz's accumulator (radtrax#26).

Attributes

n_bands Number of spectral bands. count Running sample count.

update

update(batch: Float[ArrayLike, '*batch B']) -> WelfordAccumulator

Ingest a batch of pixel spectra (shape (n_samples, n_bands) or (H, W, n_bands)).

Source code in src/radtrax/retrieval/matched_filter/streaming.py

def update(self, batch: Float[ArrayLike, "*batch B"]) -> WelfordAccumulator:
    """Ingest a batch of pixel spectra (shape ``(n_samples, n_bands)`` or ``(H, W, n_bands)``)."""
    arr = np.asarray(batch, dtype=float)
    if arr.ndim == 3:
        arr = arr.reshape(-1, arr.shape[-1])
    if arr.ndim != 2 or arr.shape[1] != self.n_bands:
        raise ShapeError(
            f"WelfordAccumulator.update: expected shape (*, {self.n_bands}); "
            f"got {arr.shape}."
        )
    finite = _finite_pixels(arr)
    if not finite.all():
        arr = arr[finite]
    if arr.shape[0] == 0:
        return self
    part = WelfordAccumulator(n_bands=self.n_bands)
    part.count = arr.shape[0]
    part._mean = arr.mean(axis=0)
    d = arr - part._mean
    sq = np.sum(d * d, axis=1)
    part._M2 = d.T @ d
    part._M3 = sq @ d
    part._M4 = float(sq @ sq)
    return self.merge(part)

merge

merge(other: WelfordAccumulator) -> WelfordAccumulator

Combine with another accumulator on disjoint data (exact).

Source code in src/radtrax/retrieval/matched_filter/streaming.py

def merge(self, other: WelfordAccumulator) -> WelfordAccumulator:
    """Combine with another accumulator on disjoint data (exact)."""
    if other.n_bands != self.n_bands:
        raise ShapeError(
            f"WelfordAccumulator.merge: n_bands mismatch "
            f"({self.n_bands} vs {other.n_bands})."
        )
    if other.count == 0:
        return self
    if self.count == 0:
        self._mean = other._mean.copy()
        self._M2 = other._M2.copy()
        self._M3 = other._M3.copy()
        self._M4 = other._M4
        self.count = other.count
        return self
    n = self.count + other.count
    mean = self._mean + (other._mean - self._mean) * (other.count / n)
    M2 = np.zeros_like(self._M2)
    M3 = np.zeros_like(self._M3)
    M4 = 0.0
    for part in (self, other):
        delta = part._mean - mean
        delta_sq = float(delta @ delta)
        M2_delta = part._M2 @ delta
        tr_M2 = float(np.trace(part._M2))
        M2 += part._M2 + part.count * np.outer(delta, delta)
        M3 += (
            part._M3
            + tr_M2 * delta
            + 2.0 * M2_delta
            + part.count * delta_sq * delta
        )
        M4 += (
            part._M4
            + 4.0 * float(part._M3 @ delta)
            + 4.0 * float(delta @ M2_delta)
            + 2.0 * delta_sq * tr_M2
            + part.count * delta_sq**2
        )
    self._mean, self._M2, self._M3, self._M4 = mean, M2, M3, M4
    self.count = n
    return self

mean

mean() -> Float[ndarray, ' B']

Mean spectrum of the samples seen so far.

Source code in src/radtrax/retrieval/matched_filter/streaming.py

def mean(self) -> Float[np.ndarray, " B"]:
    """Mean spectrum of the samples seen so far."""
    if self.count == 0:
        raise DegenerateInputError("WelfordAccumulator.mean: no data yet.")
    return self._mean.copy()

covariance

covariance(ddof: int = 1) -> Float[ndarray, 'B B']

Covariance M₂ / (count − ddof).

Source code in src/radtrax/retrieval/matched_filter/streaming.py

def covariance(self, ddof: int = 1) -> Float[np.ndarray, "B B"]:
    """Covariance ``M₂ / (count − ddof)``."""
    if self.count <= ddof:
        raise DegenerateInputError(
            f"WelfordAccumulator.covariance: need count > ddof; got count={self.count}."
        )
    return self._M2 / (self.count - ddof)

fourth_moment

fourth_moment() -> float

m₄ = M₄ / count = (1/n) Σ ‖xₖ − x̄‖⁴, the Ledoit–Wolf input.

Source code in src/radtrax/retrieval/matched_filter/streaming.py

def fourth_moment(self) -> float:
    """``m₄ = M₄ / count = (1/n) Σ ‖xₖ − x̄‖⁴``, the Ledoit–Wolf input."""
    if self.count == 0:
        raise DegenerateInputError("WelfordAccumulator.fourth_moment: no data yet.")
    return self._M4 / self.count

streaming_background

streaming_background(batches: Iterable[Float[ArrayLike, '*_batch B']], *, n_bands: int, ridge: float = 0.0, ddof: int = 0, cov_kind: Literal['empirical', 'ledoit_wolf', 'oas'] = 'empirical') -> tuple[Float[ndarray, ' B'], LinearOperator]

Streamed (μ, Σ) from an iterable of batches.

Convenience wrapper for the common pattern of looping over scenes/tiles. Returns the final mean and a dense lineax operator for the covariance — pass the covariance through :func:~radtrax.retrieval.matched_filter.core.apply_image exactly like any other cov_op.

Parameters

batches Iterable yielding arrays of shape (n_samples, n_bands) or (H, W, n_bands). n_bands Number of spectral bands (declared up front since the iterable can be lazy). ridge Optional diagonal ridge added to the final covariance. ddof Delta degrees of freedom for the covariance denominator: M_2 / (n - ddof). Defaults to 0 (MLE, matches sklearn's EmpiricalCovariance, LedoitWolf, OAS and our :func:estimate_cov_lowrank), so :func:~radtrax.retrieval.matched_filter.core.matched_filter_snr and :func:~radtrax.retrieval.matched_filter.core.detection_threshold give the same scaling whether the covariance came from a streaming pass or a batch estimator. Set ddof=1 to match np.cov / sample-covariance conventions instead. cov_kind "empirical" (default), or "ledoit_wolf" / "oas" to shrink the streamed covariance toward a scaled identity with :func:~radtrax.retrieval.matched_filter.background.shrink_covariance. With ddof=0 the shrunk matrix equals scikit-learn's estimator on the concatenated batches; a non-zero ddof rescales it by n / (n − ddof) (the intensity is scale-free).

Source code in src/radtrax/retrieval/matched_filter/streaming.py

def streaming_background(
    batches: Iterable[Float[ArrayLike, "*_batch B"]],
    *,
    n_bands: int,
    ridge: float = 0.0,
    ddof: int = 0,
    cov_kind: Literal["empirical", "ledoit_wolf", "oas"] = "empirical",
) -> tuple[Float[np.ndarray, " B"], LinearOperator]:
    """Streamed ``(μ, Σ)`` from an iterable of batches.

    Convenience wrapper for the common pattern of looping over scenes/tiles.
    Returns the final mean and a dense lineax operator for the covariance —
    pass the covariance through :func:`~radtrax.retrieval.matched_filter.core.apply_image`
    exactly like any other ``cov_op``.

    Parameters
    ----------
    batches
        Iterable yielding arrays of shape ``(n_samples, n_bands)`` or
        ``(H, W, n_bands)``.
    n_bands
        Number of spectral bands (declared up front since the iterable can be
        lazy).
    ridge
        Optional diagonal ridge added to the final covariance.
    ddof
        Delta degrees of freedom for the covariance denominator:
        ``M_2 / (n - ddof)``. Defaults to ``0`` (MLE, matches sklearn's
        ``EmpiricalCovariance``, ``LedoitWolf``, ``OAS`` and our
        :func:`estimate_cov_lowrank`), so
        :func:`~radtrax.retrieval.matched_filter.core.matched_filter_snr` and
        :func:`~radtrax.retrieval.matched_filter.core.detection_threshold`
        give the same scaling whether the covariance came from a streaming
        pass or a batch estimator. Set ``ddof=1`` to match ``np.cov`` /
        sample-covariance conventions instead.
    cov_kind
        ``"empirical"`` (default), or ``"ledoit_wolf"`` / ``"oas"`` to shrink
        the streamed covariance toward a scaled identity with
        :func:`~radtrax.retrieval.matched_filter.background.shrink_covariance`.
        With ``ddof=0`` the shrunk matrix equals scikit-learn's estimator on
        the concatenated batches; a non-zero ``ddof`` rescales it by
        ``n / (n − ddof)`` (the intensity is scale-free).
    """
    acc = WelfordAccumulator(n_bands=n_bands)
    for batch in batches:
        acc.update(batch)
    mu = acc.mean()
    cov = acc.covariance(ddof=ddof)
    if cov_kind in ("ledoit_wolf", "oas"):
        shrunk = shrink_covariance(
            acc.covariance(ddof=0),
            method=cov_kind,
            n_samples=acc.count,
            fourth_moment=acc.fourth_moment(),
        )
        cov = shrunk * (acc.count / (acc.count - ddof))
    elif cov_kind != "empirical":
        raise ParameterError(
            "streaming_background: cov_kind must be 'empirical', 'ledoit_wolf' "
            f"or 'oas'; got {cov_kind!r}."
        )
    if ridge > 0.0:
        cov = cov + ridge * np.eye(n_bands)
    return mu, _dense_operator(cov)

target

Target signature construction from the radiance forward model.

The matched-filter target t(ν) is the spectral signature that a unit plume enhancement would leave in the observed radiance. Because we already have a JAX :class:radtrax.retrieval.var3d.RadianceObservationModel, we do not reimplement the Beer–Lambert derivation from scratch (1379 lines in the original snippet). Instead:

• Linear target is a directional derivative ∂L/∂VMR at the background state, obtained by a single :func:jax.jvp call — robust, cheap, and exactly consistent with the forward model used by the 3D-Var retrieval.

• Nonlinear target is the finite-amplitude response L(x_b + α·dx) − L(x_b), which captures Beer–Lambert saturation for strong plumes. For small α this reduces to α · t_linear and the MF score becomes scale-invariant to that choice.

• Unit-absorption target :func:unit_absorption_target is ∂ ln y / ∂ΔX per pixel on the layered :class:radtrax.ForwardModel (radtrax#37): the scene's own target, replacing a MODTRAN table.

Ported from plumax 0.1.2 at 28f257e (src/plumax/matched_filter/target.py); the read-out pixel follows a GSD since radtrax#80.

linear_target_from_obs

linear_target_from_obs(obs: RadianceObservationModel, vmr_background: Float[Array, 'H W'], *, pattern: Pattern | Float[ArrayLike, 'H W'] = 'uniform', pixel: tuple[int, int] | None = None, linear_forward: bool = False) -> Float[Array, ' B']

Matched-filter target via the tangent-linear of the forward operator.

Computes

.. math:: t = \frac{\partial}{\partial \alpha} H(x_b + \alpha \, \delta x) \Big|_{\alpha=0}

at a single output pixel. For the default pattern='uniform' this captures the spectral signature of a uniform VMR enhancement across the scene (appropriate when the plume fills several PSF footprints). For pattern='impulse' it captures the instrument's point response to a single-pixel enhancement (appropriate when the plume is subpixel or compact compared to the PSF).

Parameters

obs The JAX observation model. vmr_background Background VMR field, shape (H, W). Typically jnp.full((H, W), x_ref) where x_ref == obs.vmr_reference. pattern 'uniform' (default), 'impulse', or an explicit 2-D perturbation pattern of the same shape as vmr_background. pixel Pixel (i, j) of the VMR grid: the impulse location, and the place the signature is read — at the output pixel containing it, (i // f, j // f) for a GSD of factor f. Defaults to the scene centre (or a custom pattern's peak). pattern and pixel must be concrete (static, or closed over under jax.jit). linear_forward If True, use the linearised forward 1 - Δτ instead of exp(-Δτ). For small vmr_background the two agree to within Beer–Lambert curvature.

Returns

Float[Array, "B"] Target signature of length obs.n_bands.

Source code in src/radtrax/retrieval/matched_filter/target.py

def linear_target_from_obs(
    obs: RadianceObservationModel,
    vmr_background: Float[Array, "H W"],
    *,
    pattern: Pattern | Float[ArrayLike, "H W"] = "uniform",
    pixel: tuple[int, int] | None = None,
    linear_forward: bool = False,
) -> Float[Array, " B"]:
    """Matched-filter target via the tangent-linear of the forward operator.

    Computes

    .. math::
        t = \\frac{\\partial}{\\partial \\alpha}
            H(x_b + \\alpha \\, \\delta x) \\Big|_{\\alpha=0}

    at a single output pixel. For the default ``pattern='uniform'`` this
    captures the spectral signature of a uniform VMR enhancement across the
    scene (appropriate when the plume fills several PSF footprints). For
    ``pattern='impulse'`` it captures the instrument's point response to a
    single-pixel enhancement (appropriate when the plume is subpixel or
    compact compared to the PSF).

    Parameters
    ----------
    obs
        The JAX observation model.
    vmr_background
        Background VMR field, shape ``(H, W)``. Typically
        ``jnp.full((H, W), x_ref)`` where ``x_ref == obs.vmr_reference``.
    pattern
        ``'uniform'`` (default), ``'impulse'``, or an explicit 2-D perturbation
        pattern of the same shape as ``vmr_background``.
    pixel
        Pixel ``(i, j)`` **of the VMR grid**: the impulse location, and the
        place the signature is read — at the output pixel containing it,
        ``(i // f, j // f)`` for a GSD of factor ``f``. Defaults to the scene
        centre (or a custom pattern's peak). ``pattern`` and ``pixel`` must be
        concrete (static, or closed over under ``jax.jit``).
    linear_forward
        If True, use the linearised forward ``1 - Δτ`` instead of
        ``exp(-Δτ)``. For small ``vmr_background`` the two agree to within
        Beer–Lambert curvature.

    Returns
    -------
    Float[Array, "B"]
        Target signature of length ``obs.n_bands``.
    """
    dx = _build_pattern(vmr_background, pattern, pixel)
    fwd = obs.make_forward(linear=linear_forward)
    _, dL = jax.jvp(fwd, (vmr_background,), (dx,))
    return _read_out(dL, _vmr_pixel(vmr_background, pattern, pixel), obs)

nonlinear_target_from_obs

nonlinear_target_from_obs(obs: RadianceObservationModel, vmr_background: Float[Array, 'H W'], *, amplitude: float = 1.0, pattern: Pattern | Float[ArrayLike, 'H W'] = 'uniform', pixel: tuple[int, int] | None = None) -> Float[Array, ' B']

Matched-filter target via a finite-amplitude difference.

Computes t = H(x_b + α·δx) − H(x_b) without linearising Beer–Lambert. This is the "nonlinear" variant from the original snippets; it captures the saturation curvature that matters for strong plumes (α · σ · N · ΔX ≳ 0.1 optical depth).

For small amplitude this reduces to amplitude × linear_target.

Source code in src/radtrax/retrieval/matched_filter/target.py

def nonlinear_target_from_obs(
    obs: RadianceObservationModel,
    vmr_background: Float[Array, "H W"],
    *,
    amplitude: float = 1.0,
    pattern: Pattern | Float[ArrayLike, "H W"] = "uniform",
    pixel: tuple[int, int] | None = None,
) -> Float[Array, " B"]:
    """Matched-filter target via a finite-amplitude difference.

    Computes ``t = H(x_b + α·δx) − H(x_b)`` without linearising Beer–Lambert.
    This is the "nonlinear" variant from the original snippets; it captures
    the saturation curvature that matters for strong plumes
    (``α · σ · N · ΔX ≳ 0.1`` optical depth).

    For small ``amplitude`` this reduces to ``amplitude × linear_target``.
    """
    dx = _build_pattern(vmr_background, pattern, pixel)
    fwd = obs.make_forward(linear=False)
    L_perturbed = fwd(vmr_background + amplitude * dx)
    L_base = fwd(vmr_background)
    return _read_out(
        L_perturbed - L_base, _vmr_pixel(vmr_background, pattern, pixel), obs
    )

unit_absorption_target

unit_absorption_target(model: ForwardModel, state: AtmosphericState, surface: Surface, geometry: ViewingGeometry, *, gas: str = 'CH4', profile: Float[ArrayLike, ' L'] | None = None) -> Float[Array, ' B']

Matched-filter target t = ∂ ln y / ∂ΔX at ΔX = 0 [per mol/m²].

One forward-mode derivative of the band log-radiance with respect to a column enhancement ΔX of gas spread over the layers by profile (default: the bottom layer, a plume at the surface; see :meth:~radtrax.atmosphere.AtmosphericState.add_column). The scene enters through state (water vapour, surface pressure, temperature), geometry and surface, so t is the scene's own unit absorption spectrum. It is differentiable in all of them.

Feed it to :func:apply_image with a cube of ln y, or as μ ⊙ t with a cube of y and background mean μ (∂y/∂ΔX = y t). :func:linear_target_from_obs and :func:nonlinear_target_from_obs are the targets of the single-slab :class:~radtrax.retrieval.var3d.RadianceObservationModel (2-D fields, PSF and GSD); this one is per pixel on the layered :class:~radtrax.ForwardModel.

Parameters:

Name	Type	Description	Default
model	ForwardModel	The forward model.	required
state	AtmosphericState	The background atmosphere.	required
surface	Surface	The surface.	required
geometry	ViewingGeometry	Sun and sensor directions.	required
gas	str	The plume's gas.	'CH4'
profile	Float[ArrayLike, ' L'] | None	Relative weight of the plume per layer, (L,).	None

Returns:

Type	Description
Float[Array, ' B']	t, one value per band [1/(mol/m²)]; negative where the gas
Float[Array, ' B']	absorbs.

Source code in src/radtrax/retrieval/matched_filter/target.py

def unit_absorption_target(
    model: ForwardModel,
    state: AtmosphericState,
    surface: Surface,
    geometry: ViewingGeometry,
    *,
    gas: str = "CH4",
    profile: Float[ArrayLike, " L"] | None = None,
) -> Float[Array, " B"]:
    """Matched-filter target ``t = ∂ ln y / ∂ΔX`` at ``ΔX = 0`` [per mol/m²].

    One forward-mode derivative of the band log-radiance with respect to a
    column enhancement ``ΔX`` of ``gas`` spread over the layers by
    ``profile`` (default: the bottom layer, a plume at the surface; see
    :meth:`~radtrax.atmosphere.AtmosphericState.add_column`). The scene
    enters through ``state`` (water vapour, surface pressure, temperature),
    ``geometry`` and ``surface``, so ``t`` is the scene's own unit
    absorption spectrum. It is differentiable in all of them.

    Feed it to :func:`apply_image` with a cube of ``ln y``, or as ``μ ⊙ t``
    with a cube of ``y`` and background mean ``μ`` (``∂y/∂ΔX = y t``).
    :func:`linear_target_from_obs` and :func:`nonlinear_target_from_obs`
    are the targets of the single-slab
    :class:`~radtrax.retrieval.var3d.RadianceObservationModel` (2-D fields,
    PSF and GSD); this one is per pixel on the layered
    :class:`~radtrax.ForwardModel`.

    Args:
        model: The forward model.
        state: The background atmosphere.
        surface: The surface.
        geometry: Sun and sensor directions.
        gas: The plume's gas.
        profile: Relative weight of the plume per layer, ``(L,)``.

    Returns:
        ``t``, one value per band [1/(mol/m²)]; negative where the gas
        absorbs.
    """

    def log_bands(delta: Float[Array, ""]) -> Float[Array, " B"]:
        plume = state.add_column(gas, delta, profile)
        return jnp.log(model(plume, surface, geometry))

    # ΔX = 0 in the state's precision (float32 when x64 is off).
    return jax.jacfwd(log_bands)(jnp.zeros((), dtype=state.vmr[gas].dtype))

io

xarray glue: apply the matched filter to DataArrays and multi-file stores.

The MF kernel in :mod:.core is pure JAX and expects plain arrays. For real data workflows the scene is usually an :class:xarray.DataArray with named dims (y, x, band), coordinates, and possibly lazy Dask backing. These helpers thread the xarray metadata through while the compute stays in JAX:

• :func:apply_image_xarray — wraps :func:apply_image in :func:xarray.apply_ufunc so band-last DataArrays produce a score DataArray with (y, x) dims and preserved spatial coordinates. Dask-backed inputs are evaluated lazily with chunk-wise parallelism.
• :func:open_multi_scene — iterate over a sequence of scene files with :func:xarray.open_dataset, yielding (n_samples, n_bands) numpy batches with band_dim moved to the last axis. Each dataset is closed as soon as its batch has been yielded (or when the generator is stopped early), so file handles do not leak even when iteration is partial. Feeds naturally into :class:~radtrax.retrieval.matched_filter.streaming.WelfordAccumulator.

Ported from plumax 0.1.2 at 28f257e (src/plumax/matched_filter/io.py); only import paths changed.

apply_image_xarray

apply_image_xarray(cube: DataArray, mean: Float[ArrayLike, ' B'], cov_op: LinearOperator, target: Float[ArrayLike, ' B'], *, band_dim: str = 'band', dask: str = 'parallelized') -> DataArray

Apply the matched filter to an xarray DataArray.

Parameters

cube Radiance cube with a band-last convention (one dim named band_dim). Other dims are preserved on the output — typically y, x but any additional dims (e.g. time) pass through unchanged. mean, cov_op, target Same semantics as :func:radtrax.retrieval.matched_filter.core.apply_image. band_dim Name of the spectral dimension. Default 'band'. dask Forwarded to :func:xarray.apply_ufunc: 'parallelized' (default — chunk-wise) or 'allowed'. When 'parallelized' is used, the spectral axis is transparently rechunked into a single chunk per block (allow_rechunk=True) because the MF kernel takes the full spectrum at each pixel; this lets common zarr/netCDF files with a chunked spectral axis flow through without manual rechunking.

Returns

xr.DataArray Score DataArray with band_dim collapsed.

Source code in src/radtrax/retrieval/matched_filter/io.py

def apply_image_xarray(
    cube: xr.DataArray,
    mean: Float[ArrayLike, " B"],
    cov_op: LinearOperator,
    target: Float[ArrayLike, " B"],
    *,
    band_dim: str = "band",
    dask: str = "parallelized",
) -> xr.DataArray:
    """Apply the matched filter to an xarray DataArray.

    Parameters
    ----------
    cube
        Radiance cube with a band-last convention (one dim named ``band_dim``).
        Other dims are preserved on the output — typically ``y, x`` but any
        additional dims (e.g. ``time``) pass through unchanged.
    mean, cov_op, target
        Same semantics as :func:`radtrax.retrieval.matched_filter.core.apply_image`.
    band_dim
        Name of the spectral dimension. Default ``'band'``.
    dask
        Forwarded to :func:`xarray.apply_ufunc`:
        ``'parallelized'`` (default — chunk-wise) or ``'allowed'``. When
        ``'parallelized'`` is used, the spectral axis is transparently
        rechunked into a single chunk per block (``allow_rechunk=True``)
        because the MF kernel takes the full spectrum at each pixel; this
        lets common zarr/netCDF files with a chunked spectral axis flow
        through without manual rechunking.

    Returns
    -------
    xr.DataArray
        Score DataArray with ``band_dim`` collapsed.
    """
    import xarray as xr

    from radtrax.retrieval.matched_filter.core import apply_image

    if band_dim not in cube.dims:
        raise ShapeError(
            f"apply_image_xarray: DataArray has no dim named {band_dim!r}; "
            f"dims are {cube.dims}."
        )

    def _kernel(arr: np.ndarray) -> np.ndarray:
        # arr has shape (..., n_bands) where ... is the non-band dims. We
        # reshape once to 3-D (H, W, B) with H = product of non-band dims, W=1
        # so apply_image's einsum works uniformly.
        flat = arr.reshape(-1, 1, arr.shape[-1])
        scores = apply_image(flat, mean=mean, cov_op=cov_op, target=target)
        return np.asarray(scores).reshape(arr.shape[:-1])

    # allow_rechunk=True lets xarray merge adjacent chunks along band_dim
    # into a single chunk per block. Without this, apply_ufunc on a
    # spectrally-chunked Dask array raises "dimension 'band' is chunked but
    # the kernel requires it to be contiguous". The MF kernel needs the
    # full spectrum at each pixel, so rechunking is always safe here.
    return xr.apply_ufunc(
        _kernel,
        cube,
        input_core_dims=[[band_dim]],
        output_core_dims=[[]],
        exclude_dims={band_dim},
        dask=dask,
        output_dtypes=[np.float64],
        dask_gufunc_kwargs={"allow_rechunk": True},
    )

open_multi_scene

open_multi_scene(paths: str | Iterable[str], *, band_dim: str = 'band', variable: str | None = None, **open_kwargs) -> Iterator[Float[ndarray, 'N B']]

Yield (n_pixels, n_bands) batches from a sequence of scene files.

Iterates one file at a time with :func:xarray.open_dataset, closes each dataset as soon as its scene has been yielded (try/finally-bracketed so handles are released even if the generator is stopped early), and yields a flattened 2-D numpy batch suitable for :class:~radtrax.retrieval.matched_filter.streaming.WelfordAccumulator.update.

Note

Each yielded batch is a materialised numpy array — the current scene's values are loaded eagerly. This is fine for scenes that individually fit in memory (the Welford aggregator's usual operating regime); for truly out-of-core per-scene processing, use the lazy DataArray directly with :func:apply_image_xarray (which supports Dask chunks) instead of this helper.

Parameters

paths A glob string (e.g. "scenes/*.nc") or an iterable of paths. band_dim Name of the spectral dimension on the stored DataArray. Does not have to be the last axis in the file — each scene is transposed so band_dim becomes the last axis before flattening. variable If a file contains multiple variables, select this one. If None, uses the first data variable. **open_kwargs Forwarded to :func:xarray.open_dataset.

Source code in src/radtrax/retrieval/matched_filter/io.py

def open_multi_scene(
    paths: str | Iterable[str],
    *,
    band_dim: str = "band",
    variable: str | None = None,
    **open_kwargs,
) -> Iterator[Float[np.ndarray, "N B"]]:
    """Yield ``(n_pixels, n_bands)`` batches from a sequence of scene files.

    Iterates one file at a time with :func:`xarray.open_dataset`, closes each
    dataset as soon as its scene has been yielded (``try/finally``-bracketed
    so handles are released even if the generator is stopped early), and
    yields a flattened 2-D numpy batch suitable for
    :class:`~radtrax.retrieval.matched_filter.streaming.WelfordAccumulator.update`.

    Note
    ----
    Each yielded batch is a *materialised* numpy array — the current scene's
    values are loaded eagerly. This is fine for scenes that individually fit
    in memory (the Welford aggregator's usual operating regime); for truly
    out-of-core per-scene processing, use the lazy DataArray directly with
    :func:`apply_image_xarray` (which supports Dask chunks) instead of this
    helper.

    Parameters
    ----------
    paths
        A glob string (e.g. ``"scenes/*.nc"``) or an iterable of paths.
    band_dim
        Name of the spectral dimension on the stored DataArray. Does *not*
        have to be the last axis in the file — each scene is transposed so
        ``band_dim`` becomes the last axis before flattening.
    variable
        If a file contains multiple variables, select this one. If ``None``,
        uses the first data variable.
    **open_kwargs
        Forwarded to :func:`xarray.open_dataset`.
    """
    import glob as _glob

    import xarray as xr

    if isinstance(paths, str):
        expanded = sorted(_glob.glob(paths))
        if not expanded:
            raise MissingFileError(
                f"open_multi_scene: glob {paths!r} matched zero files."
            )
    else:
        expanded = list(paths)

    for path in expanded:
        ds = xr.open_dataset(path, **open_kwargs)
        try:
            if variable is None:
                var_name = next(iter(ds.data_vars))
                da = ds[var_name]
            else:
                da = ds[variable]
            if band_dim not in da.dims:
                raise ShapeError(
                    f"open_multi_scene: variable in {path!r} has no dim "
                    f"{band_dim!r}; dims are {da.dims}."
                )
            # Put band_dim last so the subsequent reshape always targets the
            # spectral axis — guards against files stored as (band, y, x) or
            # with any other dim ordering.
            other_dims = [d for d in da.dims if d != band_dim]
            da = da.transpose(*other_dims, band_dim)
            scene = da.values
            if scene.ndim < 2:
                raise ShapeError(
                    f"open_multi_scene: scene in {path!r} has unexpected "
                    f"shape {scene.shape}; need ≥ 2 dims."
                )
            yield scene.reshape(-1, scene.shape[-1])
        finally:
            ds.close()

3D-Var

Variational methane retrieval: the radiance observation operator, background covariances, control transforms, solvers and diagnostics. Theory: 3D-Var.

var3d

Variational data assimilation (3D-Var) for methane VMR retrieval.

Modules

• :mod:.obs_operator — JAX forward operator H = SRF ∘ GSD ∘ PSF ∘ exp(−a·ΔVMR) built from a HAPI cross-section LUT.
• :mod:.background — covariance B constructors: diagonal, Kronecker (separable AR(1) spatial), low-rank update from a sample stack.
• :mod:.control — control-variable transforms: identity vs whitening (Cholesky / CVT) — the latter dramatically improves Hessian conditioning.
• :mod:.cost — cost+grad+HVP bundles in either model space (δx) or whitened space (ξ). Reverse-mode AD handles the obs-term adjoint.
• :mod:.solve — three solver paths: primal LBFGS, Gauss-Newton with a custom :mod:lineax linear solver, and the dual / PSAS reformulation.
• :mod:.diagnostics — reduced χ², degrees-of-freedom-for-signal, posterior covariance proxy via CG on the Hessian.

The math derivation lives in notebooks/assimilation/00_3dvar_derivation.md and the end-to-end demo in notebooks/assimilation/06_3dvar_methane_retrieval.ipynb.

Notes

Submodules are imported lazily (PEP 562), only on first attribute access.

Where this lives (decision D2 of the radtrax backlog): the retrieval moves here with its radiance observation operator, which is radiative transfer. Its solvers are generic variational data assimilation, which is vardax's job; whether they move there waits for the rule on whether GeoModels may import GeoML packages (package-boundaries open questions 1 and 5).

Ported from plumax 0.1.2 at 28f257e (plumax.assimilation); only import paths changed.

obs_operator

JAX observation operator for the methane VMR retrieval.

The full forward map factors as

x  ──┐
     │ ΔVMR = x − x_ref
     ▼
    Δτ(ν, i, j) = a(ν) · ΔVMR(i, j)         a(ν) = σ(ν, T, P) · N · ΔX · AMF
     │
     ▼
    L_hr(i, j, ν) = exp(−Δτ)                 (or 1 − Δτ for linear)
     │
     ▼ PSF (per channel)
     ▼ GSD (per channel)
     ▼ BandInstrument (band integration, per pixel)
    y_obs(I, J, b)

Everything below the LUT lookup is JAX, so

• forward runs end-to-end under :func:jax.jit,
• jax.grad(loss∘forward) produces the discrete adjoint for free,
• jax.jvp(forward, …) produces the tangent-linear,

without us writing or maintaining any explicit adjoint. The reason radtran/instrument.py does expose hand-rolled adjoint methods is to let unit tests verify the inner-product identity independently — JAX's autodiff inherits the correctness, not the other way around.

The factory :meth:RadianceObservationModel.from_lut precomputes a(ν) once from the HAPI cross-section LUT so the per-iteration cost in 3D-Var is dominated by the cube-shaped operations (PSF, GSD, SRF).

The band integration is :class:~radtrax.instrument.BandInstrument, the one JAX band product in radtrax, mapped over pixels with eqx.filter_vmap; its spectral shift is a leaf, so a 3D-Var retrieval can estimate it too (radtrax#95). The high-resolution axis is the instrument's wavenumber grid, ascending.

Ported from plumax 0.1.2 at 28f257e (src/plumax/assimilation/obs_operator.py); its frozen dataclasses became Equinox modules in radtrax#58.

RadianceObservationModel

RadianceObservationModel(absorption_coeff_hr: Float[ArrayLike, ' Ghr'], instrument: BandInstrument | Float[ArrayLike, 'B Ghr'], vmr_reference: Float[ArrayLike, ''] = 0.0, psf: PointSpreadFunction | None = None, gsd: GroundSamplingDistance | None = None)

Bases: Module

Bundle of precomputed pieces that defines the forward operator H.

Attributes

absorption_coeff_hr : jnp.ndarray Per-wavenumber absorption coefficient a(ν) = σ(ν,T,P) · N · ΔX · AMF on instrument.nu_cm (ascending wavenumber), shape (n_lambda,). Captures everything except the VMR field, so the forward map is just L_hr(i,j,ν) = exp(-a(ν) · ΔVMR(i,j)). instrument : BandInstrument The band integration on the same wavenumber grid, (n_lambda,) → (n_bands,), applied per pixel; its shift_cm (0 by default) is a differentiable leaf. vmr_reference : Array Scalar background VMR, stored as an array so it is traced, not static, under eqx.filter_jit and differentiable with eqx.filter_grad. Background VMR x_ref around which radiances are normalised. Setting vmr_reference = 0 recovers absolute Beer-Lambert; setting it to the scene's true background produces the differential signal that matched filter retrievals consume. psf, gsd : optional JAX-side instrument operators from :mod:radtrax.instrument.psf. None skips that stage — convenient for hyperspectral tests where the spectral cube is already at sensor resolution.

Source code in src/radtrax/retrieval/var3d/obs_operator.py

def __init__(
    self,
    absorption_coeff_hr: Float[ArrayLike, " Ghr"],
    instrument: BandInstrument | Float[ArrayLike, "B Ghr"],
    vmr_reference: Float[ArrayLike, ""] = 0.0,
    psf: PointSpreadFunction | None = None,
    gsd: GroundSamplingDistance | None = None,
) -> None:
    a = jnp.asarray(absorption_coeff_hr)
    if a.ndim != 1:
        raise ShapeError(
            f"RadianceObservationModel: absorption_coeff_hr must be 1-D, "
            f"got shape {a.shape}."
        )
    if not isinstance(instrument, BandInstrument):
        instrument, a = _from_srf_matrix(instrument, a)
    if instrument.response.shape[1] != a.size:
        raise ShapeError(
            f"RadianceObservationModel: the instrument's response "
            f"{instrument.response.shape} does not align with "
            f"absorption_coeff_hr (n_lambda={a.size})."
        )
    self.absorption_coeff_hr = a
    self.instrument = instrument
    # An array, not a Python float: a float field is static under
    # eqx.filter_jit, so a new value would recompile and get no gradient.
    self.vmr_reference = jnp.asarray(vmr_reference, dtype=a.dtype)
    self.psf = psf
    self.gsd = gsd

n_lambda property

n_lambda: int

Number of high-resolution spectral samples.

n_bands property

n_bands: int

Number of sensor bands.

srf_matrix property

srf_matrix: Float[Array, 'B Ghr']

Deprecated: the response in ascending wavelength, instrument.response[:, ::-1].

from_lut classmethod

from_lut(lut: Dataset, *, srf: SpectralResponseFunction, T_K: float, p_atm: float, path_length_cm: float, amf: float, vmr_reference: float = 0.0, psf: PointSpreadFunction | None = None, gsd: GroundSamplingDistance | None = None, var: str = 'absorption_cross_section', shift_cm: float = 0.0) -> RadianceObservationModel

Build the operator from a HAPI σ(ν,T,P) LUT and an SRF.

Pre-collapses σ · N · ΔX · AMF into a single a(ν) vector so the forward call only multiplies by ΔVMR per pixel.

The HR spectral grid is taken from the SRF: its wavelengths srf.wavelengths_hr_nm, reversed into ascending wavenumber ν = 10⁷ / λ, are the grid of both a(ν) and the :class:~radtrax.instrument.BandInstrument built by :meth:BandInstrument.from_srf <radtrax.instrument.BandInstrument.from_srf>, so the band integration contracts aligned spectral bins by construction. Earlier revisions accepted a separate nu_obs argument; when callers passed a wavenumber-ordered grid while the SRF was wavelength-ordered (common whenever wl = 1e7/ν was sorted), the product silently mixed non-corresponding bins — self-consistent for twin retrievals but physically wrong for anything cross-validated.

shift_cm is the instrument's spectral shift δ [cm⁻¹], 0 by default (the band integration of earlier revisions).

Source code in src/radtrax/retrieval/var3d/obs_operator.py

@classmethod
def from_lut(
    cls,
    lut: xr.Dataset,
    *,
    srf: SpectralResponseFunction,
    T_K: float,
    p_atm: float,
    path_length_cm: float,
    amf: float,
    vmr_reference: float = 0.0,
    psf: PointSpreadFunction | None = None,
    gsd: GroundSamplingDistance | None = None,
    var: str = "absorption_cross_section",
    shift_cm: float = 0.0,
) -> RadianceObservationModel:
    """Build the operator from a HAPI σ(ν,T,P) LUT and an SRF.

    Pre-collapses ``σ · N · ΔX · AMF`` into a single ``a(ν)`` vector so the
    forward call only multiplies by ``ΔVMR`` per pixel.

    The HR spectral grid is **taken from the SRF**: its wavelengths
    ``srf.wavelengths_hr_nm``, reversed into ascending wavenumber
    ``ν = 10⁷ / λ``, are the grid of both ``a(ν)`` and the
    :class:`~radtrax.instrument.BandInstrument` built by
    :meth:`BandInstrument.from_srf
    <radtrax.instrument.BandInstrument.from_srf>`, so the band
    integration contracts aligned spectral bins by construction. Earlier
    revisions accepted a separate ``nu_obs`` argument; when callers
    passed a wavenumber-ordered grid while the SRF was wavelength-ordered
    (common whenever ``wl = 1e7/ν`` was sorted), the product silently
    mixed non-corresponding bins — self-consistent for twin retrievals
    but physically wrong for anything cross-validated.

    ``shift_cm`` is the instrument's spectral shift ``δ`` [cm⁻¹], 0 by
    default (the band integration of earlier revisions).
    """
    if var not in lut:
        raise MissingKeyError(
            f"RadianceObservationModel.from_lut: variable {var!r} not in "
            f"dataset (have {list(lut.data_vars)})."
        )
    import xarray as xr_mod

    # Source-of-truth HR grid: the SRF's own wavelength axis, reversed into
    # ascending wavenumber, the grid BandInstrument.from_srf puts the
    # response on.
    wl_hr_nm = np.asarray(srf.wavelengths_hr_nm, dtype=float)
    nu_hr = 1e7 / wl_hr_nm[::-1]
    nu_da = xr_mod.DataArray(nu_hr, dims=["obs_nu"])
    sigma = lut[var].interp(
        wavenumber=nu_da, temperature=T_K, pressure=p_atm, method="linear"
    )
    sigma_np = np.asarray(sigma.values, dtype=float)
    N_total = number_density_cm3(p_atm, T_K)
    a = sigma_np * N_total * path_length_cm * amf
    return cls(
        absorption_coeff_hr=jnp.asarray(a),
        instrument=BandInstrument.from_srf(srf, nu_hr, shift_cm=shift_cm),
        vmr_reference=vmr_reference,
        psf=psf,
        gsd=gsd,
    )

forward

forward(vmr_field: Float[Array, 'H W'], *, linear: bool = False) -> Float[Array, 'h w B']

Map a 2-D VMR field to a band-integrated obs cube.

Parameters

vmr_field : Array Per-pixel VMR, shape (ny_hr, nx_hr). linear : bool If True, replace exp(-Δτ) by its Maclaurin-1 approximation 1 - Δτ. The tangent-linear of the nonlinear forward at vmr_field == vmr_reference coincides with this; the matched filter is exactly this case with δx propagated through.

Returns

Array Sensor-space radiance cube, shape (ny_lr, nx_lr, n_bands) (ny_lr = ny_hr // f, nx_lr = nx_hr // f if a GSD is set).

Source code in src/radtrax/retrieval/var3d/obs_operator.py

def forward(
    self, vmr_field: Float[Array, "H W"], *, linear: bool = False
) -> Float[Array, "h w B"]:
    """Map a 2-D VMR field to a band-integrated obs cube.

    Parameters
    ----------
    vmr_field : Array
        Per-pixel VMR, shape ``(ny_hr, nx_hr)``.
    linear : bool
        If True, replace ``exp(-Δτ)`` by its Maclaurin-1 approximation
        ``1 - Δτ``. The tangent-linear of the nonlinear forward at
        ``vmr_field == vmr_reference`` coincides with this; the matched
        filter is exactly this case with ``δx`` propagated through.

    Returns
    -------
    Array
        Sensor-space radiance cube, shape ``(ny_lr, nx_lr, n_bands)``
        (``ny_lr = ny_hr // f``, ``nx_lr = nx_hr // f`` if a GSD is set).
    """
    if vmr_field.ndim != 2:
        raise ShapeError(
            f"RadianceObservationModel.forward: vmr_field must be 2-D, got shape {vmr_field.shape}."
        )
    delta_vmr = vmr_field - self.vmr_reference
    # Outer product (ny, nx) ⊗ (n_lambda,) → (ny, nx, n_lambda).
    delta_tau = delta_vmr[..., None] * self.absorption_coeff_hr
    # `linear` is a Python bool (static); branch on it directly so JAX
    # doesn't trace the unused branch — `exp(huge δτ)` would otherwise
    # overflow on stress-tests like the adjoint inner-product check.
    L_hr = (1.0 - delta_tau) if linear else jnp.exp(-delta_tau)
    if self.psf is not None:
        L_hr = self.psf.apply(L_hr)
    if self.gsd is not None:
        L_hr = self.gsd.apply(L_hr)
    # Band integration per pixel: (h, w, n_lambda) → (h, w, n_bands).
    return eqx.filter_vmap(eqx.filter_vmap(self.instrument))(L_hr)

make_forward

make_forward(*, linear: bool = False) -> Callable[[Float[Array, 'H W']], Float[Array, 'h w B']]

Return a closure vmr_field → obs with the model captured.

The closure composes directly with jax.value_and_grad and the optimisers, which differentiate with respect to the state only. To differentiate with respect to the model itself (its SRF, its PSF kernel), pass the model as an argument instead: eqx.filter_grad(lambda m, x: loss(m.forward(x))).

jax.jit(model.forward) no longer works: the model holds arrays, so it is not hashable (it was a hashable frozen dataclass in plumax). Use eqx.filter_jit(model.forward) or jax.jit(model.make_forward()).

Source code in src/radtrax/retrieval/var3d/obs_operator.py

def make_forward(
    self, *, linear: bool = False
) -> Callable[[Float[Array, "H W"]], Float[Array, "h w B"]]:
    """Return a closure ``vmr_field → obs`` with the model captured.

    The closure composes directly with ``jax.value_and_grad`` and the
    optimisers, which differentiate with respect to the state only. To
    differentiate with respect to the model itself (its SRF, its PSF
    kernel), pass the model as an argument instead:
    ``eqx.filter_grad(lambda m, x: loss(m.forward(x)))``.

    ``jax.jit(model.forward)`` no longer works: the model holds arrays, so
    it is not hashable (it was a hashable frozen dataclass in plumax). Use
    ``eqx.filter_jit(model.forward)`` or ``jax.jit(model.make_forward())``.
    """
    linear_flag = bool(linear)

    def _fwd(vmr_field: Float[Array, "H W"]) -> Float[Array, "h w B"]:
        return self.forward(vmr_field, linear=linear_flag)

    return _fwd

background

Background-error covariance B constructors for 3D-Var.

The prior term ½ δxᵀ B⁻¹ δx in the 3D-Var cost function is dominated, at the operator level, by how you write down B. We support three flavours that cover most of what a methane-retrieval workflow ever needs:

============= ====================================================== ================= Flavour Operator structure Cost of B⁻¹ ============= ====================================================== ================= diagonal B = diag(σ²) O(N) Kronecker B = ρ_y ⊗ ρ_x (separable, e.g. AR(1) in each axis) O(N · √N) low-rank B = λI + U diag(d) Uᵀ (Woodbury via gaussx) O(N · k) ============= ====================================================== =================

All three return a :class:lineax.AbstractLinearOperator (via gaussx) so the preconditioning transform :class:~radtrax.retrieval.var3d.control.WhiteningTransform and the cost gradient can both call :func:gaussx.solve and :func:gaussx.cholesky uniformly without caring about structure.

The Kronecker constructor uses an exponential (AR(1)-style) correlation [exp(-|i-j|/L)] along each axis — simple, positive-definite by construction, and matches the prior families used in operational variational assimilation systems for spatially-smooth tracer fields.

Ported from plumax 0.1.2 at 28f257e (src/plumax/assimilation/background.py); only import paths changed.

build_diagonal_background

build_diagonal_background(variance: Float[ArrayLike, ''] | Float[ArrayLike, ' N'], n_pixels: int | None = None) -> LinearOperator

Build B = diag(σ²) as a :class:lineax.DiagonalLinearOperator.

Parameters

variance : float or array Per-pixel prior variance. Scalar broadcasts to n_pixels. n_pixels : int, optional Required when variance is a scalar.

Returns

LinearOperator B as a diagonal operator. gaussx.solve(B, r) is then a per-element division.

Source code in src/radtrax/retrieval/var3d/background.py

def build_diagonal_background(
    variance: Float[ArrayLike, ""] | Float[ArrayLike, " N"], n_pixels: int | None = None
) -> LinearOperator:
    """Build ``B = diag(σ²)`` as a :class:`lineax.DiagonalLinearOperator`.

    Parameters
    ----------
    variance : float or array
        Per-pixel prior variance. Scalar broadcasts to ``n_pixels``.
    n_pixels : int, optional
        Required when ``variance`` is a scalar.

    Returns
    -------
    LinearOperator
        ``B`` as a diagonal operator. ``gaussx.solve(B, r)`` is then a
        per-element division.
    """
    var = np.asarray(variance, dtype=float).ravel()
    if var.size == 1:
        if n_pixels is None:
            raise ParameterError(
                "build_diagonal_background: scalar variance requires `n_pixels`."
            )
        var = np.full(int(n_pixels), float(var.item()))
    if np.any(var <= 0.0):
        raise ParameterError("build_diagonal_background: variances must be > 0.")
    return lx.DiagonalLinearOperator(jnp.asarray(var))

build_kronecker_background

build_kronecker_background(*, ny: int, nx: int, variance: float, length_scale_y: float, length_scale_x: float) -> LinearOperator

Build B = σ² · (ρ_y ⊗ ρ_x) using :class:gaussx.Kronecker.

The Kronecker structure makes B⁻¹ an O(ny³ + nx³) operation rather than O((ny·nx)³), which is the difference between feasible and not for realistic scene sizes (e.g. 256×256 → 65 k state vector vs. 256³ + 256³ ≈ 30 M flops vs. 280 T flops).

The convention here is B = σ² · K_y ⊗ K_x with row-major flattening δx[i·nx + j] = δx(i, j) — matches vmr_field.reshape(-1) order.

Source code in src/radtrax/retrieval/var3d/background.py

def build_kronecker_background(
    *,
    ny: int,
    nx: int,
    variance: float,
    length_scale_y: float,
    length_scale_x: float,
) -> LinearOperator:
    """Build ``B = σ² · (ρ_y ⊗ ρ_x)`` using :class:`gaussx.Kronecker`.

    The Kronecker structure makes ``B⁻¹`` an ``O(ny³ + nx³)`` operation rather
    than ``O((ny·nx)³)``, which is the difference between feasible and not for
    realistic scene sizes (e.g. 256×256 → 65 k state vector vs.
    256³ + 256³ ≈ 30 M flops vs. 280 T flops).

    The convention here is ``B = σ² · K_y ⊗ K_x`` with row-major flattening
    ``δx[i·nx + j] = δx(i, j)`` — matches ``vmr_field.reshape(-1)`` order.
    """
    if variance <= 0.0:
        raise ParameterError("build_kronecker_background: `variance` must be > 0.")
    Ky = _ar1_correlation_matrix(ny, length_scale_y)
    Kx = _ar1_correlation_matrix(nx, length_scale_x)
    # Spread the σ² factor evenly between the two Kronecker factors so each is
    # individually well-scaled. Mathematically (σ a) ⊗ (σ b) = σ² (a ⊗ b).
    sigma = float(np.sqrt(variance))
    op_y = lx.MatrixLinearOperator(
        jnp.asarray(sigma * Ky),
        tags=frozenset({lx.symmetric_tag, lx.positive_semidefinite_tag}),
    )
    op_x = lx.MatrixLinearOperator(
        jnp.asarray(sigma * Kx),
        tags=frozenset({lx.symmetric_tag, lx.positive_semidefinite_tag}),
    )
    return gx.Kronecker(op_y, op_x)

build_lowrank_background

build_lowrank_background(*, samples: Float[ArrayLike, 'S N'], rank: int, regularization: float = 1e-06, random_state: int = 0) -> LinearOperator

Estimate B from a stack of state samples and wrap as λI + U D Uᵀ.

Parameters

samples : np.ndarray Sample matrix, shape (n_samples, n_state) — e.g. flattened VMR snapshots from a NWP ensemble or a climatology. rank : int Number of leading SVD components to keep. Costs O(n_state · k) per solve. regularization : float Diagonal floor λ so the resulting operator is strictly PD even when samples is rank-deficient. Must be > 0. random_state : int Seed for :func:gaussx.randomized_svd, used when rank < 0.5 * min(n_samples, n_state) (otherwise a full SVD runs).

Returns

LinearOperator gaussx.LowRankUpdate tagged symmetric + PSD; gaussx.solve(B, r) uses Woodbury internally.

Source code in src/radtrax/retrieval/var3d/background.py

def build_lowrank_background(
    *,
    samples: Float[ArrayLike, "S N"],
    rank: int,
    regularization: float = 1e-6,
    random_state: int = 0,
) -> LinearOperator:
    """Estimate ``B`` from a stack of state samples and wrap as ``λI + U D Uᵀ``.

    Parameters
    ----------
    samples : np.ndarray
        Sample matrix, shape ``(n_samples, n_state)`` — e.g. flattened VMR
        snapshots from a NWP ensemble or a climatology.
    rank : int
        Number of leading SVD components to keep. Costs ``O(n_state · k)``
        per solve.
    regularization : float
        Diagonal floor ``λ`` so the resulting operator is strictly PD even
        when ``samples`` is rank-deficient. Must be > 0.
    random_state : int
        Seed for :func:`gaussx.randomized_svd`, used when
        ``rank < 0.5 * min(n_samples, n_state)`` (otherwise a full SVD runs).

    Returns
    -------
    LinearOperator
        ``gaussx.LowRankUpdate`` tagged symmetric + PSD; ``gaussx.solve(B, r)``
        uses Woodbury internally.
    """
    if samples.ndim != 2:
        raise ShapeError(
            f"build_lowrank_background: samples must be 2-D (n_samples, n_state); "
            f"got {samples.shape}."
        )
    n_samples, n_state = samples.shape
    if n_samples < 2:
        raise DegenerateInputError("build_lowrank_background: need ≥ 2 samples.")
    rank = max(1, min(int(rank), n_state, n_samples))
    if regularization <= 0.0:
        raise ParameterError("build_lowrank_background: `regularization` must be > 0.")
    centred = samples - samples.mean(axis=0, keepdims=True)
    if rank < 0.5 * min(n_samples, n_state):
        # Few components wanted: randomized range finder, O(rank) passes over
        # the samples instead of a full O(N·n·min(N,n)) SVD.
        _, S_k, V_k = gx.randomized_svd(
            lx.MatrixLinearOperator(jnp.asarray(centred)),
            rank,
            n_power_iter=5,
            key=jax.random.key(random_state),
        )
    else:
        _, S, Vt = np.linalg.svd(centred, full_matrices=False)
        S_k = S[:rank]
        V_k = Vt[:rank]
    U = jnp.asarray(V_k.T)  # (n_state, rank)
    d = jnp.asarray(S_k**2 / max(n_samples - 1, 1))
    base = lx.DiagonalLinearOperator(
        jnp.asarray(regularization * np.ones(n_state, dtype=float))
    )
    return gx.LowRankUpdate(
        base,
        U,
        d,
        tags=frozenset({lx.symmetric_tag, lx.positive_semidefinite_tag}),
    )

control

Control-variable transforms for variational data assimilation.

3D-Var minimises a cost J(δx) in some "control space". The two choices worth keeping straight are:

1. Model space — control variable is the increment δx = x − x_b. The prior term is ½ δxᵀ B⁻¹ δx; one B⁻¹-solve per gradient evaluation. Hessian conditioning = cond(B⁻¹ + H'ᵀ R⁻¹ H') blows up when B is poorly scaled.

2. Whitened (CVT) space — control variable is ξ with δx = U ξ and B = U Uᵀ. The prior collapses to ½ ξᵀ ξ and the Hessian becomes I + (HU)ᵀ R⁻¹ (HU) — a small, dense identity-plus-low-rank matrix that even simple LBFGS solves in a handful of iterations regardless of how ill-conditioned B was. This is the single biggest practical lever you have in 3D-Var.

Both transforms expose the same minimal interface:

• apply(xi) → δx (control-space → model-space)
• apply_inverse(δx) → xi (model-space → control-space; needed once at init if the user gives an x-space prior mean)
• project_gradient(g_x) → g_xi (chain rule for the prior gradient when you're working with the cost-in-x form; not used in JAX autodiff paths)

We rely on :func:gaussx.cholesky to produce U as a structured operator, so the whitening cost stays O(B-solve) rather than O(N²) for any of the structured B flavours in :mod:.background.

Ported from plumax 0.1.2 at 28f257e (src/plumax/assimilation/control.py); its frozen dataclasses became Equinox modules in radtrax#58.

IdentityTransform

Bases: Module

Pass-through transform. Use when the prior is already well-conditioned or when you want the textbook ½ δxᵀ B⁻¹ δx form for didactic clarity.

WhiteningTransform

Bases: Module

δx = U ξ where B = U Uᵀ.

Constructed once via :meth:from_background; afterwards each call is a single structured matvec / triangular solve from gaussx.

Notes

Whitening also defines a natural sampler from the prior: δx = U ε with ε ∼ 𝒩(0, I) is exactly distributed as 𝒩(0, B). We don't use that here but it's the same operator.

from_background classmethod

from_background(background: LinearOperator) -> WhiteningTransform

Factor B = U Uᵀ via :func:gaussx.cholesky.

gaussx returns a structured triangular operator (preserving Kronecker / low-rank-update structure), so subsequent matvecs and inverse-solves stay in the cheap regime.

Source code in src/radtrax/retrieval/var3d/control.py

@classmethod
def from_background(cls, background: LinearOperator) -> WhiteningTransform:
    """Factor ``B = U Uᵀ`` via :func:`gaussx.cholesky`.

    gaussx returns a structured triangular operator (preserving Kronecker /
    low-rank-update structure), so subsequent matvecs and inverse-solves
    stay in the cheap regime.
    """
    U = gx.cholesky(background)
    return cls(cholesky_op=U)

apply

apply(xi: Float[ArrayLike, ' N']) -> Float[Array, ' N']

ξ → U ξ.

Source code in src/radtrax/retrieval/var3d/control.py

def apply(self, xi: Float[ArrayLike, " N"]) -> Float[Array, " N"]:
    """``ξ → U ξ``."""
    return self.cholesky_op.mv(jnp.asarray(xi))

apply_inverse

apply_inverse(delta_x: Float[ArrayLike, ' N']) -> Float[Array, ' N']

δx → U⁻¹ δx via :func:gaussx.solve for structured dispatch.

gaussx.cholesky returns operators that preserve their underlying structure (a Cholesky of a Kronecker is itself a Kronecker of Cholesky factors); gaussx.solve knows how to dispatch on those, whereas lineax.AutoLinearSolver would try to detect generic triangularity and raise NotImplementedError.

Source code in src/radtrax/retrieval/var3d/control.py

def apply_inverse(self, delta_x: Float[ArrayLike, " N"]) -> Float[Array, " N"]:
    """``δx → U⁻¹ δx`` via :func:`gaussx.solve` for structured dispatch.

    ``gaussx.cholesky`` returns operators that preserve their underlying
    structure (a Cholesky of a ``Kronecker`` is itself a ``Kronecker`` of
    Cholesky factors); ``gaussx.solve`` knows how to dispatch on those,
    whereas ``lineax.AutoLinearSolver`` would try to detect generic
    triangularity and raise ``NotImplementedError``.
    """
    return gx.solve(self.cholesky_op, jnp.asarray(delta_x))

project_gradient

project_gradient(g_x: Float[ArrayLike, ' N']) -> Float[Array, ' N']

∂J/∂ξ = Uᵀ ∂J/∂δx — the chain rule for the explicit-cost form.

Not needed when you build the cost in ξ-space and let jax.grad differentiate end-to-end (which is what :mod:.cost does).

Source code in src/radtrax/retrieval/var3d/control.py

def project_gradient(self, g_x: Float[ArrayLike, " N"]) -> Float[Array, " N"]:
    """``∂J/∂ξ = Uᵀ ∂J/∂δx`` — the chain rule for the explicit-cost form.

    Not needed when you build the cost in ξ-space and let ``jax.grad``
    differentiate end-to-end (which is what :mod:`.cost` does).
    """
    return self.cholesky_op.transpose().mv(jnp.asarray(g_x))

cost

3D-Var cost functions and gradients.

The standard incremental 3D-Var minimises

J(δx) = ½ δxᵀ B⁻¹ δx + ½ ‖y − H(x_b + δx)‖²_R⁻¹

over the increment δx = x − x_b. Two versions live here:

• :func:build_cost_x — model-space form. The prior term is evaluated via :func:gaussx.solve(B, δx) so any structured B from :mod:.background works without materialisation.
• :func:build_cost_xi — whitened (CVT) form. The prior term collapses to ½ ξᵀ ξ and the obs term sees δx = U ξ via the supplied :class:~radtrax.retrieval.var3d.control.WhiteningTransform. This is the version we use in practice — gradients are well-scaled and the Hessian is identity-plus-low-rank.

In both cases the obs term is differentiated through the JAX forward map via reverse-mode AD (jax.grad), so we don't need to maintain an explicit adjoint of the observation operator. The :class:Cost bundle bundles J, ∇J and a Hessian-vector product Hv so the optimiser / posterior-covariance machinery downstream can pick whichever it needs.

Ported from plumax 0.1.2 at 28f257e (src/plumax/assimilation/cost.py); its frozen dataclasses became Equinox modules in radtrax#58.

Cost

Bases: Module

Bundle of (value, gradient, Hessian-vector product) for an optimiser.

The three callables share the same input signature — a 1-D state vector — so the optimiser doesn't have to know which control-space we're working in.

The observation, background and forward model live inside the closures, not as array leaves, so a Cost is a static bundle of functions: one compilation per build_cost_* call, and not differentiable as data. Differentiate with respect to the model through make_forward's model argument instead (see RadianceObservationModel.make_forward).

build_cost_x

build_cost_x(*, forward_fn: ForwardFn, background_op: AbstractLinearOperator, obs_inv_variance: ObsInvVariance, background_state: Float[ArrayLike, 'H W'], observation: Float[ArrayLike, 'h w B'], state_shape: tuple[int, ...]) -> Cost

Build J(δx) = ½ δxᵀ B⁻¹ δx + ½ (y − H(x_b + δx))ᵀ R⁻¹ (y − H(x_b + δx)).

Parameters

forward_fn : Callable H : 2-D VMR field → obs cube. Must be JAX-traceable. background_op : lineax operator B (any structured form from :mod:.background). The cost calls gaussx.solve(B, δx_flat) once per gradient evaluation. obs_inv_variance : float or Array R⁻¹. Pass a scalar for R = σ² I or an array broadcastable against the obs cube for diagonal heteroscedastic noise. background_state, observation, state_shape x_b, y_obs, and the shape of the VMR field (e.g. (ny, nx)).

Source code in src/radtrax/retrieval/var3d/cost.py

def build_cost_x(
    *,
    forward_fn: ForwardFn,
    background_op: lx.AbstractLinearOperator,
    obs_inv_variance: ObsInvVariance,
    background_state: Float[ArrayLike, "H W"],
    observation: Float[ArrayLike, "h w B"],
    state_shape: tuple[int, ...],
) -> Cost:
    """Build ``J(δx) = ½ δxᵀ B⁻¹ δx + ½ (y − H(x_b + δx))ᵀ R⁻¹ (y − H(x_b + δx))``.

    Parameters
    ----------
    forward_fn : Callable
        ``H : 2-D VMR field → obs cube``. Must be JAX-traceable.
    background_op : lineax operator
        ``B`` (any structured form from :mod:`.background`). The cost calls
        ``gaussx.solve(B, δx_flat)`` once per gradient evaluation.
    obs_inv_variance : float or Array
        ``R⁻¹``. Pass a scalar for ``R = σ² I`` or an array broadcastable
        against the obs cube for diagonal heteroscedastic noise.
    background_state, observation, state_shape
        ``x_b``, ``y_obs``, and the shape of the VMR field (e.g. ``(ny, nx)``).
    """
    x_b = jnp.asarray(background_state)
    y = jnp.asarray(observation)
    R_inv = jnp.asarray(obs_inv_variance)

    def value(delta_x_flat: jax.Array) -> jax.Array:
        delta_x = jnp.reshape(delta_x_flat, state_shape)
        x = x_b + delta_x
        residual = y - forward_fn(x)
        # Prior term: ½ δxᵀ B⁻¹ δx. gx.quadratic_form is xᵀ B⁻¹ x via a single
        # structurally-dispatched solve — the same op gx.solve does, one call.
        prior = 0.5 * gx.quadratic_form(background_op, delta_x_flat)
        obs = 0.5 * jnp.sum(R_inv * residual * residual)
        return prior + obs

    val_and_grad = jax.value_and_grad(value)
    grad = lambda x: val_and_grad(x)[1]

    def hvp(x: jax.Array, v: jax.Array) -> jax.Array:
        # ∂²J/∂x² · v via forward-over-reverse — the standard JAX recipe.
        return jax.jvp(grad, (x,), (v,))[1]

    return Cost(value=value, grad=grad, value_and_grad=val_and_grad, hvp=hvp)

build_cost_xi

build_cost_xi(*, forward_fn: ForwardFn, whitening: WhiteningTransform | IdentityTransform, obs_inv_variance: ObsInvVariance, background_state: Float[ArrayLike, 'H W'], observation: Float[ArrayLike, 'h w B'], state_shape: tuple[int, ...]) -> Cost

Build J(ξ) = ½ ξᵀ ξ + ½ (y − H(x_b + U ξ))ᵀ R⁻¹ (…).

The whitening transform turns the prior into a clean ½‖ξ‖², so the Hessian is I + (HU)ᵀ R⁻¹ (HU) — a small perturbation of the identity. Even vanilla LBFGS converges in a handful of iterations regardless of how badly conditioned B was.

Pass an :class:~radtrax.retrieval.var3d.control.IdentityTransform here to recover an unpreconditioned model-space cost with implicit B = I — useful for unit tests of the obs term in isolation.

Source code in src/radtrax/retrieval/var3d/cost.py

def build_cost_xi(
    *,
    forward_fn: ForwardFn,
    whitening: WhiteningTransform | IdentityTransform,
    obs_inv_variance: ObsInvVariance,
    background_state: Float[ArrayLike, "H W"],
    observation: Float[ArrayLike, "h w B"],
    state_shape: tuple[int, ...],
) -> Cost:
    """Build ``J(ξ) = ½ ξᵀ ξ + ½ (y − H(x_b + U ξ))ᵀ R⁻¹ (…)``.

    The whitening transform turns the prior into a clean ``½‖ξ‖²``, so the
    Hessian is ``I + (HU)ᵀ R⁻¹ (HU)`` — a small perturbation of the identity.
    Even vanilla LBFGS converges in a handful of iterations regardless of how
    badly conditioned ``B`` was.

    Pass an :class:`~radtrax.retrieval.var3d.control.IdentityTransform`
    here to recover an unpreconditioned model-space cost with implicit
    ``B = I`` — useful for unit tests of the obs term in isolation.
    """
    x_b = jnp.asarray(background_state)
    y = jnp.asarray(observation)
    R_inv = jnp.asarray(obs_inv_variance)

    def value(xi_flat: jax.Array) -> jax.Array:
        delta_x_flat = whitening.apply(xi_flat)
        delta_x = jnp.reshape(delta_x_flat, state_shape)
        x = x_b + delta_x
        residual = y - forward_fn(x)
        prior = 0.5 * jnp.dot(xi_flat, xi_flat)
        obs = 0.5 * jnp.sum(R_inv * residual * residual)
        return prior + obs

    val_and_grad = jax.value_and_grad(value)
    grad = lambda xi: val_and_grad(xi)[1]

    def hvp(xi: jax.Array, v: jax.Array) -> jax.Array:
        return jax.jvp(grad, (xi,), (v,))[1]

    return Cost(value=value, grad=grad, value_and_grad=val_and_grad, hvp=hvp)

finite_difference_grad

finite_difference_grad(fn: Callable[[Float[Array, '*shape']], Float[Array, '']], x: Float[ArrayLike, '*shape'], *, eps: float = 1e-06) -> Float[ndarray, '*shape']

Central finite differences — used only by tests to validate jax.grad.

O(N) function evaluations, so keep x.size small.

Source code in src/radtrax/retrieval/var3d/cost.py

def finite_difference_grad(
    fn: Callable[[Float[Array, "*shape"]], Float[Array, ""]],
    x: Float[ArrayLike, "*shape"],
    *,
    eps: float = 1e-6,
) -> Float[np.ndarray, "*shape"]:
    """Central finite differences — used only by tests to validate ``jax.grad``.

    O(N) function evaluations, so keep ``x.size`` small.
    """
    x_np = np.asarray(x, dtype=float).copy()
    g = np.zeros_like(x_np)
    for i in range(x_np.size):
        orig = x_np.flat[i]
        x_np.flat[i] = orig + eps
        f_plus = float(fn(jnp.asarray(x_np)))
        x_np.flat[i] = orig - eps
        f_minus = float(fn(jnp.asarray(x_np)))
        g.flat[i] = (f_plus - f_minus) / (2.0 * eps)
        x_np.flat[i] = orig
    return g

solve

Optimisers for 3D-Var: primal LBFGS, Gauss-Newton, and a dual PSAS path.

Three entry points, each tuned to a different conditioning story:

• :func:run_lbfgs — black-box minimiser. Pass any cost+grad bundle from :mod:.cost. Robust, but iteration count tracks cond(Hessian), so use this with a :class:~radtrax.retrieval.var3d.control.WhiteningTransform.
• :func:run_gauss_newton — assumes the cost is a sum of squared residuals (which 3D-Var is, by construction). Each outer step solves a linear system (JᵀJ) δ = -Jᵀr via a user-chosen :mod:lineax solver (CG by default, so structured B keeps Woodbury / Kronecker on the hot path).
• :func:run_dual_psas — the observation-space reformulation. When dim(y) ≪ dim(x) (typical for compact plumes), solving (H'BH'ᵀ + R) λ = d and recovering δx = B H'ᵀ λ is dramatically cheaper than the primal — and exactly equivalent for the linear forward.

The dual path is also the "matched filter, generalised": the matched-filter target spectrum t falls out as H' applied to a unit increment, and the PSAS solution δx̂ = B H'ᵀ (H'BH'ᵀ + R)⁻¹ d is the BLUE estimator that the pixel-wise matched filter approximates with B = ∞ I, R = Σ_pixel.

Ported from plumax 0.1.2 at 28f257e (src/plumax/assimilation/solve.py); its frozen dataclasses became Equinox modules in radtrax#58.

SolveResult

Bases: Module

Common result type for all three solvers.

Attributes

state : np.ndarray Optimal control variable (ξ for whitened, δx for model-space, δx for the dual after the back-transform). n_steps : int Number of iterations consumed (-1 if the solver doesn't report). cost_history : np.ndarray Cost trajectory; for solvers that don't expose it, this is empty.

run_lbfgs

run_lbfgs(cost: Cost, initial_state: Float[ArrayLike, ' N'], *, rtol: float = 1e-06, atol: float = 1e-09, max_steps: int = 200) -> SolveResult

Minimise cost.value via :class:optimistix.LBFGS.

The cost is wrapped to match optimistix's (y, args) → scalar signature. The implicit-function adjoint is the default — fine for this use case because we don't differentiate through the solve.

Source code in src/radtrax/retrieval/var3d/solve.py

def run_lbfgs(
    cost: Cost,
    initial_state: Float[ArrayLike, " N"],
    *,
    rtol: float = 1e-6,
    atol: float = 1e-9,
    max_steps: int = 200,
) -> SolveResult:
    """Minimise ``cost.value`` via :class:`optimistix.LBFGS`.

    The cost is wrapped to match optimistix's ``(y, args) → scalar`` signature.
    The implicit-function adjoint is the default — fine for this use case
    because we don't differentiate *through* the solve.
    """
    fn = lambda y, _args: cost.value(y)
    solver = optx.LBFGS(rtol=rtol, atol=atol)
    sol = optx.minimise(fn, solver, initial_state, max_steps=max_steps, throw=False)
    return SolveResult(
        state=np.asarray(sol.value),
        n_steps=int(sol.stats.get("num_steps", -1)),
        cost_history=np.asarray([]),
    )

run_gauss_newton

run_gauss_newton(*, residual_fn: Callable[[Float[Array, ' N']], Float[Array, ' M']], initial_state: Float[ArrayLike, ' N'], linear_solver: AbstractLinearSolver | None = None, rtol: float = 1e-06, atol: float = 1e-09, max_steps: int = 50) -> SolveResult

Run :class:optimistix.GaussNewton for min ½ ‖r(y)‖².

residual_fn should return a flat residual vector [B^{-½}(δx); R^{-½}(y - H(x_b + δx))] (the "augmented" 3D-Var residual); optimistix internally constructs the Jacobian and solves the normal equations (JᵀJ) δ = -Jᵀr using linear_solver. Pass lineax.CG(...) for a matrix-free inner solve (useful when the Jacobian has Kronecker / low-rank structure that full Cholesky would materialise), or lineax.Cholesky() for small dense systems.

Source code in src/radtrax/retrieval/var3d/solve.py

def run_gauss_newton(
    *,
    residual_fn: Callable[[Float[Array, " N"]], Float[Array, " M"]],
    initial_state: Float[ArrayLike, " N"],
    linear_solver: lx.AbstractLinearSolver | None = None,
    rtol: float = 1e-6,
    atol: float = 1e-9,
    max_steps: int = 50,
) -> SolveResult:
    """Run :class:`optimistix.GaussNewton` for ``min ½ ‖r(y)‖²``.

    ``residual_fn`` should return a flat residual vector
    ``[B^{-½}(δx); R^{-½}(y - H(x_b + δx))]`` (the "augmented" 3D-Var residual);
    optimistix internally constructs the Jacobian and solves the normal
    equations ``(JᵀJ) δ = -Jᵀr`` using ``linear_solver``. Pass
    ``lineax.CG(...)`` for a matrix-free inner solve (useful when the
    Jacobian has Kronecker / low-rank structure that full Cholesky would
    materialise), or ``lineax.Cholesky()`` for small dense systems.
    """
    if linear_solver is None:
        linear_solver = lx.AutoLinearSolver(well_posed=None)
    fn = lambda y, _args: residual_fn(y)
    solver = optx.GaussNewton(rtol=rtol, atol=atol, linear_solver=linear_solver)
    sol = optx.least_squares(
        fn, solver, initial_state, max_steps=max_steps, throw=False
    )
    return SolveResult(
        state=np.asarray(sol.value),
        n_steps=int(sol.stats.get("num_steps", -1)),
        cost_history=np.asarray([]),
    )

run_dual_psas

run_dual_psas(*, forward_fn: ForwardFn, background_op: AbstractLinearOperator, obs_inv_variance: float, background_state: Float[ArrayLike, 'H W'], observation: Float[ArrayLike, 'h w B'], state_shape: tuple[int, ...], cg_rtol: float = 1e-06, cg_atol: float = 1e-09, cg_max_steps: int = 200) -> SolveResult

Solve the dual 3D-Var system in observation space.

For a forward operator that is linear in the state perturbation the optimum of J(δx) = ½ δxᵀ B⁻¹ δx + ½(y - H(x_b + δx))ᵀ R⁻¹(y - H(x_b + δx)) satisfies

(H' B H'ᵀ + R) λ = d,           with d = y − H(x_b)
δx̂ = B H'ᵀ λ.

When dim(y) ≪ dim(x) this is much cheaper than primal — and it's exactly the estimator that the pixel-wise matched filter approximates.

Parameters

forward_fn : Callable Full-state forward x → y (i.e. the same object returned by :meth:RadianceObservationModel.make_forward). For the dual to be correct, H must be linear in δx = x − x_b. We extract the tangent-linear action by differencing at x_b once (see base_obs below) and the adjoint via :func:jax.vjp. Callers typically pass model.make_forward(linear=True). background_op : lineax operator B. obs_inv_variance : float Scalar R⁻¹ (heteroscedastic R can be folded into forward_fn at the cost of clarity — keep it scalar for the demo). background_state, observation, state_shape Self-explanatory. cg_rtol, cg_atol, cg_max_steps Tolerances for the CG solve of the dual system.

Source code in src/radtrax/retrieval/var3d/solve.py

def run_dual_psas(
    *,
    forward_fn: ForwardFn,
    background_op: lx.AbstractLinearOperator,
    obs_inv_variance: float,
    background_state: Float[ArrayLike, "H W"],
    observation: Float[ArrayLike, "h w B"],
    state_shape: tuple[int, ...],
    cg_rtol: float = 1e-6,
    cg_atol: float = 1e-9,
    cg_max_steps: int = 200,
) -> SolveResult:
    """Solve the **dual** 3D-Var system in observation space.

    For a forward operator that is linear in the state perturbation the
    optimum of ``J(δx) = ½ δxᵀ B⁻¹ δx + ½(y - H(x_b + δx))ᵀ R⁻¹(y - H(x_b + δx))``
    satisfies

        (H' B H'ᵀ + R) λ = d,           with d = y − H(x_b)
        δx̂ = B H'ᵀ λ.

    When ``dim(y) ≪ dim(x)`` this is much cheaper than primal — and it's
    exactly the estimator that the pixel-wise matched filter approximates.

    Parameters
    ----------
    forward_fn : Callable
        **Full-state forward** ``x → y`` (i.e. the same object returned by
        :meth:`RadianceObservationModel.make_forward`). For the dual to be
        correct, ``H`` must be linear in ``δx = x − x_b``. We extract the
        tangent-linear action by differencing at ``x_b`` once (see
        ``base_obs`` below) and the adjoint via :func:`jax.vjp`. Callers
        typically pass ``model.make_forward(linear=True)``.
    background_op : lineax operator
        ``B``.
    obs_inv_variance : float
        Scalar ``R⁻¹`` (heteroscedastic R can be folded into ``forward_fn``
        at the cost of clarity — keep it scalar for the demo).
    background_state, observation, state_shape
        Self-explanatory.
    cg_rtol, cg_atol, cg_max_steps
        Tolerances for the CG solve of the dual system.
    """
    x_b = jnp.asarray(background_state)
    y = jnp.asarray(observation)
    R_inv = float(obs_inv_variance)

    # Evaluate H(x_b) once — used both for the innovation and as the affine
    # anchor in the tangent-linear probe. Hoisting this out of the matvec
    # avoids a redundant forward pass per CG iteration.
    base_obs = forward_fn(x_b)
    d = (y - base_obs).reshape(-1)
    n_obs = d.size
    n_state = int(np.prod(state_shape))

    def H_op(delta_x_flat: jax.Array) -> jax.Array:
        """``H' · δx``: linearity ⇒ ``H(x_b + δx) − H(x_b)`` is exact."""
        delta_x = jnp.reshape(delta_x_flat, state_shape)
        return (forward_fn(x_b + delta_x) - base_obs).reshape(-1)

    # Build the VJP closure *once*, outside the CG matvec, so every iteration
    # reuses the traced pullback instead of re-tracing jax.vjp(H_op, 0) each
    # time.  This is materially faster — was the P2 perf finding in review.
    _, _Ht_pullback = jax.vjp(H_op, jnp.zeros(n_state))

    def Ht_op(lam: jax.Array) -> jax.Array:
        (out,) = _Ht_pullback(lam)
        return out

    def dual_matvec(lam: jax.Array) -> jax.Array:
        """``(H' B H'ᵀ + R) λ``, with ``R λ = λ / R_inv`` for scalar R."""
        BHt_lam = background_op.mv(Ht_op(lam))
        return H_op(BHt_lam) + lam / R_inv

    dual_op = lx.FunctionLinearOperator(
        dual_matvec,
        input_structure=jax.eval_shape(lambda: jnp.zeros(n_obs)),
        tags=frozenset({lx.symmetric_tag, lx.positive_semidefinite_tag}),
    )
    sol = lx.linear_solve(
        dual_op,
        d,
        solver=lx.CG(rtol=cg_rtol, atol=cg_atol, max_steps=cg_max_steps),
        throw=False,  # surface non-convergence in `sol.result` instead of raising
    )
    lam_star = sol.value
    delta_x_flat = background_op.mv(Ht_op(lam_star))
    return SolveResult(
        state=np.asarray(delta_x_flat),
        n_steps=int(sol.stats.get("num_steps", -1)) if hasattr(sol, "stats") else -1,
        cost_history=np.asarray([]),
    )

diagnostics

Posterior diagnostics for a converged 3D-Var solution.

After the optimiser has produced x̂, three quick scalars tell you whether the result is statistically consistent with the prior + likelihood you specified:

• :func:reduced_chi_squared — χ² = ‖y − H(x̂)‖²_R⁻¹ / dim(y). ≈ 1 means the posterior fits the observations to the noise level. Much smaller → over-fitting (R too generous); much larger → systematic errors not in R.
• :func:degrees_of_freedom_for_signal — DFS = trace(I − A) with the averaging kernel A = Bₐ B⁻¹ (Rodgers, 2000). For low-rank B this is cheap; for the dense form we use Hutchinson's stochastic trace estimator.
• :func:posterior_covariance_proxy — the inverse Hessian at the optimum. Returned as a callable that takes a probe vector v and returns Bₐ v (i.e. you can evaluate variances per-pixel without forming the full Bₐ matrix).

These are diagnostics, not exact statistics — for the linear-Gaussian case they coincide with the textbook posterior; for the nonlinear case they're Laplace approximations around the optimum.

Ported from plumax 0.1.2 at 28f257e (src/plumax/assimilation/diagnostics.py); only import paths changed.

reduced_chi_squared

reduced_chi_squared(*, forward_fn: ForwardFn, estimated_state: Float[ArrayLike, 'H W'], observation: Float[ArrayLike, 'h w B'], obs_inv_variance: ObsInvVariance) -> float

χ²_red = (y - H(x̂))ᵀ R⁻¹ (y - H(x̂)) / dim(y).

Values near 1 indicate the posterior is consistent with the observation error model; ≪ 1 → noise overestimated; ≫ 1 → noise underestimated or systematic forward-model error.

Source code in src/radtrax/retrieval/var3d/diagnostics.py

def reduced_chi_squared(
    *,
    forward_fn: ForwardFn,
    estimated_state: Float[ArrayLike, "H W"],
    observation: Float[ArrayLike, "h w B"],
    obs_inv_variance: ObsInvVariance,
) -> float:
    """``χ²_red = (y - H(x̂))ᵀ R⁻¹ (y - H(x̂)) / dim(y)``.

    Values near 1 indicate the posterior is consistent with the observation
    error model; ≪ 1 → noise overestimated; ≫ 1 → noise underestimated or
    systematic forward-model error.
    """
    y = jnp.asarray(observation)
    R_inv = jnp.asarray(obs_inv_variance)
    residual = y - forward_fn(jnp.asarray(estimated_state))
    chi2 = float(jnp.sum(R_inv * residual * residual))
    return chi2 / int(y.size)

degrees_of_freedom_for_signal

degrees_of_freedom_for_signal(*, hessian_vector_product: Callable[[StateVector], StateVector], state_size: int, background_op: AbstractLinearOperator | None = None, whitened: bool = False, n_probes: int = 32, seed: int = 0, cg_rtol: float = 1e-06, cg_atol: float = 1e-10, cg_max_steps: int = 2000) -> float

Estimate DFS = trace(I − Bₐ B⁻¹) (Rodgers 2000, §2.5).

For the linear-Gaussian case the averaging kernel is A = I − Bₐ B⁻¹ and DFS = trace(A) measures how many independent pieces of information the observations contributed (DFS = 0 → no information; DFS = N → observations fully pin the state). The posterior covariance is the inverse Hessian at the optimum, Bₐ ≈ Hess⁻¹. Because DFS is invariant under the control-variable change δx = U ξ, the estimate must be computed in the space the supplied hessian_vector_product lives in — mixing an ξ-space Hessian with a model-space B gives a wrong answer for any B ≠ I.

Two modes, one per cost builder:

• Model space (whitened=False, the default) — pass the Hessian from :func:radtrax.retrieval.var3d.cost.build_cost_x and background_op = B. Each Hutchinson probe needs one B⁻¹ z via :func:gaussx.solve (structured dispatch) and one Hess⁻¹ w via CG: DFS ≈ E[zᵀz − zᵀ Hess⁻¹ B⁻¹ z].
• Whitened space (whitened=True) — pass the Hessian from :func:radtrax.retrieval.var3d.cost.build_cost_xi (Hess_ξ = I + (HU)ᵀR⁻¹HU). The whitened prior is the identity, so B⁻¹ drops out: DFS ≈ E[zᵀz − zᵀ Hess_ξ⁻¹ z] and background_op must be omitted.

Cost per probe: O(n_CG · forward) — dominated by the CG solve.

Earlier versions computed trace(B · Hess), which is not DFS: in the zero-information limit (Hess = B⁻¹) it returns N instead of the correct 0. See PR-review thread on diagnostics.py.

Source code in src/radtrax/retrieval/var3d/diagnostics.py

def degrees_of_freedom_for_signal(
    *,
    hessian_vector_product: Callable[[StateVector], StateVector],
    state_size: int,
    background_op: lx.AbstractLinearOperator | None = None,
    whitened: bool = False,
    n_probes: int = 32,
    seed: int = 0,
    cg_rtol: float = 1e-6,
    cg_atol: float = 1e-10,
    cg_max_steps: int = 2000,
) -> float:
    """Estimate ``DFS = trace(I − Bₐ B⁻¹)`` (Rodgers 2000, §2.5).

    For the linear-Gaussian case the averaging kernel is ``A = I − Bₐ B⁻¹``
    and ``DFS = trace(A)`` measures how many independent pieces of information
    the observations contributed (``DFS = 0`` → no information; ``DFS = N`` →
    observations fully pin the state). The posterior covariance is the
    inverse Hessian at the optimum, ``Bₐ ≈ Hess⁻¹``. Because DFS is invariant
    under the control-variable change ``δx = U ξ``, the estimate must be
    computed in the space the supplied ``hessian_vector_product`` lives in —
    mixing an ξ-space Hessian with a model-space ``B`` gives a wrong answer for
    any ``B ≠ I``.

    Two modes, one per cost builder:

    - **Model space** (``whitened=False``, the default) — pass the Hessian from
      :func:`radtrax.retrieval.var3d.cost.build_cost_x` and ``background_op = B``.
      Each Hutchinson probe needs one ``B⁻¹ z`` via :func:`gaussx.solve`
      (structured dispatch) and one ``Hess⁻¹ w`` via CG:
      ``DFS ≈ E[zᵀz − zᵀ Hess⁻¹ B⁻¹ z]``.
    - **Whitened space** (``whitened=True``) — pass the Hessian from
      :func:`radtrax.retrieval.var3d.cost.build_cost_xi`
      (``Hess_ξ = I + (HU)ᵀR⁻¹HU``). The whitened prior is the identity, so
      ``B⁻¹`` drops out: ``DFS ≈ E[zᵀz − zᵀ Hess_ξ⁻¹ z]`` and ``background_op``
      must be omitted.

    Cost per probe: ``O(n_CG · forward)`` — dominated by the CG solve.

    Earlier versions computed ``trace(B · Hess)``, which is not DFS: in the
    zero-information limit (``Hess = B⁻¹``) it returns ``N`` instead of the
    correct ``0``. See PR-review thread on diagnostics.py.
    """
    if whitened:
        if background_op is not None:
            raise ParameterError(
                "degrees_of_freedom_for_signal: `whitened=True` uses the "
                "whitened prior B_ξ = I; do not pass `background_op` (the ξ-space "
                "Hessian already absorbs B via δx = U ξ)."
            )

        def b_inv(z: jax.Array) -> jax.Array:
            return z
    else:
        if background_op is None:
            raise ParameterError(
                "degrees_of_freedom_for_signal: model-space mode "
                "(`whitened=False`) requires `background_op = B`; pass the "
                "Hessian from build_cost_x, or set `whitened=True` for a "
                "build_cost_xi (ξ-space) Hessian."
            )
        import gaussx as gx

        def b_inv(z: jax.Array) -> jax.Array:
            return gx.solve(background_op, z)

    # Build a CG-solvable operator for ``Hess⁻¹``.
    hess_op = lx.FunctionLinearOperator(
        lambda v: hessian_vector_product(jnp.asarray(v)),
        input_structure=jax.eval_shape(lambda: jnp.zeros(state_size)),
        tags=frozenset({lx.symmetric_tag, lx.positive_semidefinite_tag}),
    )
    cg_solver = lx.CG(rtol=cg_rtol, atol=cg_atol, max_steps=cg_max_steps)

    rng = np.random.default_rng(seed)
    z_batch = rng.choice([-1.0, 1.0], size=(n_probes, state_size))
    total = 0.0
    for z in z_batch:
        z_j = jnp.asarray(z, dtype=jnp.float64)
        # w = B⁻¹ z (model space) or z (whitened, since B_ξ = I).
        w = b_inv(z_j)
        # u = Hess⁻¹ w  (matrix-free CG)
        u = lx.linear_solve(hess_op, w, solver=cg_solver, throw=False).value
        # zᵀ (I − Bₐ B⁻¹) z ≈ zᵀ z − zᵀ (Hess⁻¹ B⁻¹) z
        total += float(jnp.dot(z_j, z_j) - jnp.dot(z_j, u))
    return total / n_probes

posterior_covariance_proxy

posterior_covariance_proxy(*, hessian_vector_product: Callable[[StateVector], StateVector], state_size: int, whitening: WhiteningTransform | None = None, cg_rtol: float = 1e-06, cg_atol: float = 1e-09, cg_max_steps: int = 200) -> Callable[[StateVector], StateVector]

Return a callable v → Bₐ v (model-space posterior covariance) via CG.

Bₐ ≈ Hess⁻¹ is the Laplace-approximation posterior covariance. We expose it as a matvec (rather than materialising) so per-pixel variance estimates remain O(state · cg_iters) — fine for moderate scenes.

Pass the Hessian that matches whitening:

• whitening=None (default) — a model-space Hessian (from :func:radtrax.retrieval.var3d.cost.build_cost_x); the CG solve returns Hess⁻¹ v = Bₐ v directly.
• a :class:~radtrax.retrieval.var3d.control.WhiteningTransform — a whitened Hessian (from :func:radtrax.retrieval.var3d.cost.build_cost_xi), whose inverse is the ξ-space covariance, not Bₐ. The model-space covariance is recovered as Bₐ = U Hess_ξ⁻¹ Uᵀ, so each apply un-whitens on both sides: v → U (Hess_ξ⁻¹ (Uᵀ v)).

Source code in src/radtrax/retrieval/var3d/diagnostics.py

def posterior_covariance_proxy(
    *,
    hessian_vector_product: Callable[[StateVector], StateVector],
    state_size: int,
    whitening: WhiteningTransform | None = None,
    cg_rtol: float = 1e-6,
    cg_atol: float = 1e-9,
    cg_max_steps: int = 200,
) -> Callable[[StateVector], StateVector]:
    """Return a callable ``v → Bₐ v`` (model-space posterior covariance) via CG.

    ``Bₐ ≈ Hess⁻¹`` is the Laplace-approximation posterior covariance. We
    expose it as a matvec (rather than materialising) so per-pixel variance
    estimates remain ``O(state · cg_iters)`` — fine for moderate scenes.

    Pass the Hessian that matches ``whitening``:

    - ``whitening=None`` (default) — a **model-space** Hessian (from
      :func:`radtrax.retrieval.var3d.cost.build_cost_x`); the CG solve returns
      ``Hess⁻¹ v = Bₐ v`` directly.
    - a :class:`~radtrax.retrieval.var3d.control.WhiteningTransform` — a **whitened**
      Hessian (from :func:`radtrax.retrieval.var3d.cost.build_cost_xi`), whose
      inverse is the ξ-space covariance, *not* ``Bₐ``. The model-space
      covariance is recovered as ``Bₐ = U Hess_ξ⁻¹ Uᵀ``, so each apply
      un-whitens on both sides: ``v → U (Hess_ξ⁻¹ (Uᵀ v))``.
    """

    def matvec(v: jax.Array) -> jax.Array:
        return hessian_vector_product(jnp.asarray(v))

    op = lx.FunctionLinearOperator(
        matvec,
        input_structure=jax.eval_shape(lambda: jnp.zeros(state_size)),
        tags=frozenset({lx.symmetric_tag, lx.positive_semidefinite_tag}),
    )

    def apply(v: jax.Array) -> jax.Array:
        v = jnp.asarray(v)
        # Whitened Hessian: Bₐ = U Hess_ξ⁻¹ Uᵀ, so map into ξ-space first.
        rhs = whitening.project_gradient(v) if whitening is not None else v
        # ``throw=False`` returns the best partial solution at max_steps rather
        # than raising — for diagnostics we'd rather get a noisy estimate than
        # crash the whole notebook on a stiff probe direction.
        sol = lx.linear_solve(
            op,
            rhs,
            solver=lx.CG(rtol=cg_rtol, atol=cg_atol, max_steps=cg_max_steps),
            throw=False,
        ).value
        return whitening.apply(sol) if whitening is not None else sol

    return apply

MBMP

The Sentinel-2 multi-band multi-pass methane signal, radtrax · 2's multi-pass family with per-pass scale factors (SBMP, MBSP, MBMP) and the single-pass SWIR contrast, as differentiable JAX functions; the GeoTensor operators are MBMP, SBMP and MBSP in radtrax.geo.

The multi-band multi-pass (MBMP) methane signal, and a single-pass prefilter.

Sentinel-2's band 12 (~2200 nm) sits in strong methane absorption and band 11 (~1600 nm) in weak; a plume depresses band 12 relative to band 11. With a clean-air reference acquisition of the same ground (a second pass), the log band ratio of each pass removes the surface, and their difference removes what the passes share:

.. math:: s = \ln\frac{R_{11} + ε}{R_{12} + ε} - \ln\frac{R'{11} + ε}{R'} + ε

where :math:R_b is band :math:b's radiance or reflectance on the observed pass and :math:R'_b on the reference pass. This is the MBMP log signal of Varon et al. (2021) without their per-pass scale factors: a scene-wide change between passes (a water-vapour or illumination change) stays in it as a constant offset (radtrax · 2, the GeoStack design note, found 1 % water vapour reads as such an offset). geotoolz shipped it as plume.SBMP; radtrax · 2 shows it is an MBMP, so it is named that here.

radtrax · 2 ("Three ways to remove the surface") scales each pass by a factor fitted through the origin on background pixels, c = Σ x y / Σ x² (:func:scale_factor), and defines the family (radtrax#111):

• :func:sbmp_signal, one band, two passes: −ln(c_s R₁₂ / R′₁₂), with R′₁₂ ≈ c_s R₁₂;
• :func:mbsp_signal, two bands, one pass: −ln(c R₁₂ / R₁₁), with R₁₁ ≈ c R₁₂;
• :func:mbmp_signal, two bands, two passes: the MBSP signal of the observed pass minus that of the reference, which is the log signal above minus ln(c / c′); without scale factors it is the log signal.

Each is an observable g to invert through a band table (:class:radtrax.retrieval.band_table.BandTable: τ₁₂ for SBMP, τ₁₂ − τ₁₁ for the others). Fit the factors on background pixels: a fit over a plume biases them (radtrax · 2: 11 % of the mass), and :func:background_mask grows a first-pass detection into one.

Without a reference pass, :func:swir_contrast is the normalised contrast :math:(R_{11} - R_{12}) / (R_{11} + R_{12} + ε), a cheap detection prefilter (Ehret et al., 2022), not a column retrieval.

Negative values (noisy radiances) are clipped to zero before either formula. Both are JAX and differentiable away from that clip; with ε > 0 gradients stay finite at zero too (~1/ε). A band at zero saturates the log signal at about ±ln(R/ε) (±22 for ε = 1e-10), so callers should mask such pixels (:class:radtrax.geo.MBMP does); ε = 0 gives NaN there.

mbmp_log_signal

mbmp_log_signal(swir1: Float[ArrayLike, '*#shape'], swir2: Float[ArrayLike, '*#shape'], ref_swir1: Float[ArrayLike, '*#shape'], ref_swir2: Float[ArrayLike, '*#shape'], *, eps: float = 1e-10) -> Float[Array, '*#shape']

MBMP log signal ln(R₁₁/R₁₂) − ln(R′₁₁/R′₁₂), no scale factors.

Parameters:

Name	Type	Description	Default
swir1	Float[ArrayLike, '*#shape']	Band 11 (weakly absorbing) on the observed pass.	required
swir2	Float[ArrayLike, '*#shape']	Band 12 (strongly absorbing) on the observed pass.	required
ref_swir1	Float[ArrayLike, '*#shape']	Band 11 on the clean-air reference pass.	required
ref_swir2	Float[ArrayLike, '*#shape']	Band 12 on the reference pass.	required
eps	float	Guard added to every band before the logs.	1e-10

Returns:

Type	Description
Float[Array, '*#shape']	The signal, the broadcast shape of the inputs; positive where band 12
Float[Array, '*#shape']	is depressed relative to the reference (methane), zero where nothing
Float[Array, '*#shape']	changed between passes.

Examples:

>>> import jax.numpy as jnp
>>> float(mbmp_log_signal(jnp.e, 1.0, 1.0, 1.0, eps=0.0))
1.0

Source code in src/radtrax/retrieval/mbmp.py

def mbmp_log_signal(
    swir1: Float[ArrayLike, "*#shape"],
    swir2: Float[ArrayLike, "*#shape"],
    ref_swir1: Float[ArrayLike, "*#shape"],
    ref_swir2: Float[ArrayLike, "*#shape"],
    *,
    eps: float = 1e-10,
) -> Float[Array, "*#shape"]:
    """MBMP log signal ``ln(R₁₁/R₁₂) − ln(R′₁₁/R′₁₂)``, no scale factors.

    Args:
        swir1: Band 11 (weakly absorbing) on the observed pass.
        swir2: Band 12 (strongly absorbing) on the observed pass.
        ref_swir1: Band 11 on the clean-air reference pass.
        ref_swir2: Band 12 on the reference pass.
        eps: Guard added to every band before the logs.

    Returns:
        The signal, the broadcast shape of the inputs; positive where band 12
        is depressed relative to the reference (methane), zero where nothing
        changed between passes.

    Examples:
        >>> import jax.numpy as jnp
        >>> float(mbmp_log_signal(jnp.e, 1.0, 1.0, 1.0, eps=0.0))
        1.0
    """
    return _log_ratio(swir1, swir2, eps) - _log_ratio(ref_swir1, ref_swir2, eps)

swir_contrast

swir_contrast(swir1: Float[ArrayLike, '*#shape'], swir2: Float[ArrayLike, '*#shape'], *, eps: float = 1e-10) -> Float[Array, '*#shape']

Single-pass normalised SWIR contrast (R₁₁ − R₁₂) / (R₁₁ + R₁₂ + ε).

The normalised difference of bands 11 and 12 (the NBR2 index, as in geoops.indices.nbr2) after clipping at zero, in JAX. A detection prefilter, not a column retrieval.

Examples:

>>> float(swir_contrast(3.0, 1.0, eps=0.0))
0.5

Source code in src/radtrax/retrieval/mbmp.py

def swir_contrast(
    swir1: Float[ArrayLike, "*#shape"],
    swir2: Float[ArrayLike, "*#shape"],
    *,
    eps: float = 1e-10,
) -> Float[Array, "*#shape"]:
    """Single-pass normalised SWIR contrast ``(R₁₁ − R₁₂) / (R₁₁ + R₁₂ + ε)``.

    The normalised difference of bands 11 and 12 (the NBR2 index, as in
    ``geoops.indices.nbr2``) after clipping at zero, in JAX. A detection
    prefilter, not a column retrieval.

    Examples:
        >>> float(swir_contrast(3.0, 1.0, eps=0.0))
        0.5
    """
    r1, r2 = _clip(swir1), _clip(swir2)
    return (r1 - r2) / (r1 + r2 + eps)

scale_factor

scale_factor(x: Float[ArrayLike, '*shape'], y: Float[ArrayLike, '*shape'], mask: Bool[ArrayLike, '*shape'] | None = None) -> Float[Array, '']

The least-squares c with y ≈ c x through the origin, Σ x y / Σ x².

Parameters:

Name	Type	Description	Default
x	Float[ArrayLike, '*shape']	The scaled band or pass.	required
y	Float[ArrayLike, '*shape']	The band or pass it is fitted to, the same shape.	required
mask	Bool[ArrayLike, '*shape'] | None	Pixels to fit on (True = background); all of them when None.	None

Returns:

Type	Description
Float[Array, '']	c, a scalar.

Examples:

>>> import jax.numpy as jnp
>>> float(scale_factor(jnp.array([1.0, 2.0]), jnp.array([2.0, 4.0])))
2.0

Source code in src/radtrax/retrieval/mbmp.py

def scale_factor(
    x: Float[ArrayLike, "*shape"],
    y: Float[ArrayLike, "*shape"],
    mask: Bool[ArrayLike, "*shape"] | None = None,
) -> Float[Array, ""]:
    """The least-squares ``c`` with ``y ≈ c x`` through the origin, ``Σ x y / Σ x²``.

    Args:
        x: The scaled band or pass.
        y: The band or pass it is fitted to, the same shape.
        mask: Pixels to fit on (``True`` = background); all of them when
            ``None``.

    Returns:
        ``c``, a scalar.

    Examples:
        >>> import jax.numpy as jnp
        >>> float(scale_factor(jnp.array([1.0, 2.0]), jnp.array([2.0, 4.0])))
        2.0
    """
    x, y = jnp.asarray(x, dtype=float), jnp.asarray(y, dtype=float)
    keep = jnp.ones(x.shape, bool) if mask is None else jnp.asarray(mask, dtype=bool)
    # Zero the masked pixels before the products (double-where): a masked NaN
    # or sentinel stays out of the sums and out of the gradient.
    xs, ys = jnp.where(keep, x, 0.0), jnp.where(keep, y, 0.0)
    return jnp.sum(xs * ys) / jnp.sum(xs * xs)

sbmp_signal

sbmp_signal(swir2: Float[ArrayLike, '*shape'], ref_swir2: Float[ArrayLike, '*shape'], *, mask: Bool[ArrayLike, '*shape'] | None = None, eps: float = 1e-10) -> Float[Array, '*shape']

Single-band multi-pass signal −ln(c_s R₁₂ / R′₁₂), R′₁₂ ≈ c_s R₁₂ on mask.

Invert through τ₁₂ alone. The factor removes a scene-wide change between the passes, not a local one (a wetted field reads as methane).

Examples:

>>> import jax.numpy as jnp
>>> r = jnp.array([1.0, 2.0, 3.0])
>>> [
...     round(float(v), 6) + 0.0 for v in sbmp_signal(0.9 * r, r)
... ]  # dimmer, no plume
[0.0, 0.0, 0.0]

Source code in src/radtrax/retrieval/mbmp.py

def sbmp_signal(
    swir2: Float[ArrayLike, "*shape"],
    ref_swir2: Float[ArrayLike, "*shape"],
    *,
    mask: Bool[ArrayLike, "*shape"] | None = None,
    eps: float = 1e-10,
) -> Float[Array, "*shape"]:
    """Single-band multi-pass signal ``−ln(c_s R₁₂ / R′₁₂)``, ``R′₁₂ ≈ c_s R₁₂`` on ``mask``.

    Invert through ``τ₁₂`` alone. The factor removes a scene-wide change
    between the passes, not a local one (a wetted field reads as methane).

    Examples:
        >>> import jax.numpy as jnp
        >>> r = jnp.array([1.0, 2.0, 3.0])
        >>> [
        ...     round(float(v), 6) + 0.0 for v in sbmp_signal(0.9 * r, r)
        ... ]  # dimmer, no plume
        [0.0, 0.0, 0.0]
    """
    c = scale_factor(swir2, ref_swir2, mask)
    return _log_ratio(ref_swir2, swir2, eps) - jnp.log(c)

mbsp_signal

mbsp_signal(swir1: Float[ArrayLike, '*shape'], swir2: Float[ArrayLike, '*shape'], *, mask: Bool[ArrayLike, '*shape'] | None = None, eps: float = 1e-10) -> Float[Array, '*shape']

Multi-band single-pass signal −ln(c R₁₂ / R₁₁), R₁₁ ≈ c R₁₂ on mask.

Invert through τ₁₂ − τ₁₁. A surface whose band ratio departs from the scene's reads as methane.

Examples:

>>> import jax.numpy as jnp
>>> r = jnp.array([1.0, 2.0, 3.0])
>>> [
...     round(float(v), 6) + 0.0 for v in mbsp_signal(r, 0.8 * r)
... ]  # one ratio, no plume
[0.0, 0.0, 0.0]

Source code in src/radtrax/retrieval/mbmp.py

def mbsp_signal(
    swir1: Float[ArrayLike, "*shape"],
    swir2: Float[ArrayLike, "*shape"],
    *,
    mask: Bool[ArrayLike, "*shape"] | None = None,
    eps: float = 1e-10,
) -> Float[Array, "*shape"]:
    """Multi-band single-pass signal ``−ln(c R₁₂ / R₁₁)``, ``R₁₁ ≈ c R₁₂`` on ``mask``.

    Invert through ``τ₁₂ − τ₁₁``. A surface whose band ratio departs from the
    scene's reads as methane.

    Examples:
        >>> import jax.numpy as jnp
        >>> r = jnp.array([1.0, 2.0, 3.0])
        >>> [
        ...     round(float(v), 6) + 0.0 for v in mbsp_signal(r, 0.8 * r)
        ... ]  # one ratio, no plume
        [0.0, 0.0, 0.0]
    """
    c = scale_factor(swir2, swir1, mask)
    return _log_ratio(swir1, swir2, eps) - jnp.log(c)

mbmp_signal

mbmp_signal(swir1: Float[ArrayLike, '*shape'], swir2: Float[ArrayLike, '*shape'], ref_swir1: Float[ArrayLike, '*shape'], ref_swir2: Float[ArrayLike, '*shape'], *, mask: Bool[ArrayLike, '*shape'] | None = None, scale_factors: bool = True, eps: float = 1e-10) -> Float[Array, '*shape']

Multi-band multi-pass signal −ln((1 + δ) / (1 + δ′)), δ = c R₁₂ / R₁₁ − 1.

With scale_factors (each pass's c fitted on mask) it is :func:mbmp_log_signal minus ln(c / c′), which removes a scene-wide change between the passes; without, it is :func:mbmp_log_signal itself. Invert through τ₁₂ − τ₁₁.

Examples:

>>> import jax.numpy as jnp
>>> r = jnp.array([1.0, 2.0, 3.0])
>>> s = mbmp_signal(r, 0.8 * r, 1.1 * r, 0.88 * 0.99 * r)  # 1 % H₂O change
>>> [round(float(v), 6) + 0.0 for v in s]
[0.0, 0.0, 0.0]

Source code in src/radtrax/retrieval/mbmp.py

def mbmp_signal(
    swir1: Float[ArrayLike, "*shape"],
    swir2: Float[ArrayLike, "*shape"],
    ref_swir1: Float[ArrayLike, "*shape"],
    ref_swir2: Float[ArrayLike, "*shape"],
    *,
    mask: Bool[ArrayLike, "*shape"] | None = None,
    scale_factors: bool = True,
    eps: float = 1e-10,
) -> Float[Array, "*shape"]:
    """Multi-band multi-pass signal ``−ln((1 + δ) / (1 + δ′))``, ``δ = c R₁₂ / R₁₁ − 1``.

    With ``scale_factors`` (each pass's ``c`` fitted on ``mask``) it is
    :func:`mbmp_log_signal` minus ``ln(c / c′)``, which removes a scene-wide
    change between the passes; without, it is :func:`mbmp_log_signal`
    itself. Invert through ``τ₁₂ − τ₁₁``.

    Examples:
        >>> import jax.numpy as jnp
        >>> r = jnp.array([1.0, 2.0, 3.0])
        >>> s = mbmp_signal(r, 0.8 * r, 1.1 * r, 0.88 * 0.99 * r)  # 1 % H₂O change
        >>> [round(float(v), 6) + 0.0 for v in s]
        [0.0, 0.0, 0.0]
    """
    signal = mbmp_log_signal(swir1, swir2, ref_swir1, ref_swir2, eps=eps)
    if not scale_factors:
        return signal
    c = scale_factor(swir2, swir1, mask)
    c_ref = scale_factor(ref_swir2, ref_swir1, mask)
    return signal - jnp.log(c / c_ref)

background_mask

background_mask(first_pass: Float[ArrayLike, 'H W'], *, threshold: float = 0.03, sigma_px: float = 2.0, level: float = 0.05) -> Bool[ndarray, 'H W']

Background pixels: outside a first-pass detection grown by a Gaussian.

radtrax · 2's mask: pixels whose first-pass ΔΩ exceeds threshold, blurred by a Gaussian of sigma_px pixels, and every pixel the blur lifts above level removed. NumPy (SciPy's gaussian_filter), not differentiable.

Parameters:

Name	Type	Description	Default
first_pass	Float[ArrayLike, 'H W']	A first retrieval, e.g. with factors fitted on every pixel, (H, W).	required
threshold	float	Detection threshold, in the units of first_pass.	0.03
sigma_px	float	Gaussian growth [pixels].	2.0
level	float	Blurred detection above which a pixel is not background.	0.05

Returns:

Type	Description
Bool[ndarray, 'H W']	True on background pixels, (H, W); a NaN pixel (unscoreable
Bool[ndarray, 'H W']	in the first pass) is never background.

Source code in src/radtrax/retrieval/mbmp.py

def background_mask(
    first_pass: Float[ArrayLike, "H W"],
    *,
    threshold: float = 0.03,
    sigma_px: float = 2.0,
    level: float = 0.05,
) -> Bool[np.ndarray, "H W"]:
    """Background pixels: outside a first-pass detection grown by a Gaussian.

    *radtrax · 2*'s mask: pixels whose first-pass ``ΔΩ`` exceeds
    ``threshold``, blurred by a Gaussian of ``sigma_px`` pixels, and every
    pixel the blur lifts above ``level`` removed. NumPy (SciPy's
    ``gaussian_filter``), not differentiable.

    Args:
        first_pass: A first retrieval, e.g. with factors fitted on every
            pixel, ``(H, W)``.
        threshold: Detection threshold, in the units of ``first_pass``.
        sigma_px: Gaussian growth [pixels].
        level: Blurred detection above which a pixel is not background.

    Returns:
        ``True`` on background pixels, ``(H, W)``; a ``NaN`` pixel (unscoreable
        in the first pass) is never background.
    """
    from scipy.ndimage import gaussian_filter

    first_pass = np.asarray(first_pass)
    detected = (first_pass > threshold).astype(float)
    return ~(gaussian_filter(detected, sigma_px) > level) & np.isfinite(first_pass)

Band table

A band observable g(m[, h], ΔΩ) (one band's τ₁₂, or τ₁₂ − τ₁₁) tabulated against the enhancement and inverted by interpolation, as radtrax · 1 and · 2 do, from any observable or from the layered forward model.

Band-table inversion: a band observable g tabulated against the enhancement.

radtrax · 1 (stage 4) and radtrax · 2 invert a band observable through a table rather than a linear slope, which the curve of growth biases low:

• one band, g = τ₁₂(ΔΩ) = −ln(T₁₂(Ω_bg + ΔΩ) / T₁₂(Ω_bg)) (SBMP);
• two bands, g = τ₁₂ − τ₁₁ (MBSP and MBMP).

:class:BandTable holds g on a grid of two-way air-mass factors m (and, optionally, water-vapour scalings h) against the enhancement ΔΩ. It interpolates linearly in m and h and piecewise-linearly in ΔΩ, so the inverse is exact on the same nodes: invert(table(ΔΩ, m), m) = ΔΩ to rounding. Noise makes g negative off a plume, so g is extended as an odd function of ΔΩ (as both notes do). Everything is JAX and differentiable in g, m, h and the table.

Build a table from any observable with :meth:BandTable.from_function (the notes' monochromatic band transmittances) or from the layered forward model with :func:band_table_from_model (ΔΩ then in mol/m², h scales the H₂O profile).

Shapes: M air masses, H water-vapour scalings, D enhancements.

BandTable

BandTable(m_grid: Float[ArrayLike, ' M'], delta_grid: Float[ArrayLike, ' D'], g: Float[ArrayLike, 'M *H D'], h2o_grid: Float[ArrayLike, ' H'] | None = None)

Bases: Module

A band observable g(m[, h], ΔΩ), monotone in ΔΩ, and its inverse.

Attributes:

Name	Type	Description
m_grid	Float[Array, ' M']	Two-way air-mass factors, increasing, (M,); one value makes a table for a single geometry.
delta_grid	Float[Array, ' D']	Enhancements, increasing from 0, (D,), in the units the table was built in (kg/m² in the notes, mol/m² from the forward model).
g	Float[Array, 'M *H D']	The observable, strictly increasing in ΔΩ, (M, D), or (M, H, D) with a water-vapour axis.
h2o_grid	Float[Array, ' H'] | None	Water-vapour scalings, increasing, (H,), or None.

Examples:

>>> import jax.numpy as jnp
>>> delta = jnp.linspace(0.0, 0.4, 41)
>>> table = BandTable(
...     jnp.array([2.0, 3.0]), delta, jnp.stack([0.3 * delta, 0.45 * delta])
... )
>>> round(float(table(0.2, 2.5)), 6)  # g at m = 2.5, halfway
0.075
>>> round(float(table.invert(0.075, 2.5)), 6)
0.2

Source code in src/radtrax/retrieval/band_table.py

def __init__(
    self,
    m_grid: Float[ArrayLike, " M"],
    delta_grid: Float[ArrayLike, " D"],
    g: Float[ArrayLike, "M *H D"],
    h2o_grid: Float[ArrayLike, " H"] | None = None,
) -> None:
    self.m_grid = jnp.asarray(m_grid, dtype=float)
    self.delta_grid = jnp.asarray(delta_grid, dtype=float)
    self.g = jnp.asarray(g, dtype=float)
    self.h2o_grid = None if h2o_grid is None else jnp.asarray(h2o_grid, dtype=float)

row

row(m: Float[ArrayLike, ''], h2o: Float[ArrayLike, ''] | None = None) -> Float[Array, ' D']

g against ΔΩ at m (and h), linear between the nodes, (D,).

Source code in src/radtrax/retrieval/band_table.py

def row(
    self, m: Float[ArrayLike, ""], h2o: Float[ArrayLike, ""] | None = None
) -> Float[Array, " D"]:
    """``g`` against ``ΔΩ`` at ``m`` (and ``h``), linear between the nodes, ``(D,)``."""
    g = _lerp(self.g, self.m_grid, m)
    if self.h2o_grid is not None:
        if h2o is None:
            raise ParameterError(
                "BandTable: this table has a water-vapour axis; pass h2o."
            )
        g = _lerp(g, self.h2o_grid, h2o)
    return g

invert

invert(g: Float[ArrayLike, '*shape'], m: Float[ArrayLike, ''], h2o: Float[ArrayLike, ''] | None = None) -> Float[Array, '*shape']

ΔΩ = g⁻¹(g) at m (and h), odd in g; held at the table's end beyond it.

Source code in src/radtrax/retrieval/band_table.py

def invert(
    self,
    g: Float[ArrayLike, "*shape"],
    m: Float[ArrayLike, ""],
    h2o: Float[ArrayLike, ""] | None = None,
) -> Float[Array, "*shape"]:
    """``ΔΩ = g⁻¹(g)`` at ``m`` (and ``h``), odd in ``g``; held at the table's end beyond it."""
    row = self.row(m, h2o)
    return _odd(
        lambda x: jnp.interp(x, row, self.delta_grid), jnp.asarray(g, dtype=float)
    )

from_function classmethod

from_function(fn: Callable[..., Float[ArrayLike, ' D']], m_grid: Float[ArrayLike, ' M'], delta_grid: Float[ArrayLike, ' D'], h2o_grid: Float[ArrayLike, ' H'] | None = None) -> BandTable

Tabulate fn(m, delta_grid) (or fn(m, delta_grid, h)), each (D,).

Source code in src/radtrax/retrieval/band_table.py

@classmethod
def from_function(
    cls,
    fn: Callable[..., Float[ArrayLike, " D"]],
    m_grid: Float[ArrayLike, " M"],
    delta_grid: Float[ArrayLike, " D"],
    h2o_grid: Float[ArrayLike, " H"] | None = None,
) -> BandTable:
    """Tabulate ``fn(m, delta_grid)`` (or ``fn(m, delta_grid, h)``), each ``(D,)``."""
    m_grid, delta_grid = (
        np.asarray(m_grid, dtype=float),
        np.asarray(delta_grid, dtype=float),
    )
    if h2o_grid is None:
        g = np.stack([np.asarray(fn(m, delta_grid), dtype=float) for m in m_grid])
    else:
        h2o_grid = np.asarray(h2o_grid, dtype=float)
        g = np.stack(
            [
                np.stack(
                    [
                        np.asarray(fn(m, delta_grid, h), dtype=float)
                        for h in h2o_grid
                    ]
                )
                for m in m_grid
            ]
        )
    return cls(m_grid, delta_grid, g, h2o_grid)

band_table_from_model

band_table_from_model(model: ForwardModel, state: AtmosphericState, surface: Surface, *, bands: Sequence[int], m_grid: Float[ArrayLike, ' M'], delta_grid: Float[ArrayLike, ' D'], h2o_grid: Float[ArrayLike, ' H'] | None = None, gas: str = 'CH4', h2o_gas: str = 'H2O', mu: float = 1.0) -> BandTable

A :class:BandTable from the layered forward model.

τ_b(ΔΩ) = −ln(y_b(ΔΩ) / y_b(0)) with ΔΩ [mol/m²] of gas added by :meth:~radtrax.atmosphere.AtmosphericState.add_column; g = τ_{b₀} for one band, g = τ_{b₀} − τ_{b₁} for two. Each m is a geometry with the sensor at mu and the sun at μ₀ = 1 / (m − 1/μ); each h scales the state's h2o_gas profile.

Parameters:

Name	Type	Description	Default
model	ForwardModel	The forward model.	required
state	AtmosphericState	The background state.	required
surface	Surface	The surface.	required
bands	Sequence[int]	One band index (SBMP) or two, strong then weak (MBSP, MBMP).	required
m_grid	Float[ArrayLike, ' M']	Two-way air-mass factors, each at least 1 + 1/mu, (M,).	required
delta_grid	Float[ArrayLike, ' D']	Enhancements [mol/m²], increasing from 0, (D,).	required
h2o_grid	Float[ArrayLike, ' H'] | None	Water-vapour scalings, (H,), or None.	None
gas	str	The enhanced gas.	'CH4'
h2o_gas	str	The water-vapour gas the scalings apply to.	'H2O'
mu	float	Cosine of the viewing zenith angle.	1.0

Returns:

Type	Description
BandTable	The table.

Source code in src/radtrax/retrieval/band_table.py

def band_table_from_model(
    model: ForwardModel,
    state: AtmosphericState,
    surface: Surface,
    *,
    bands: Sequence[int],
    m_grid: Float[ArrayLike, " M"],
    delta_grid: Float[ArrayLike, " D"],
    h2o_grid: Float[ArrayLike, " H"] | None = None,
    gas: str = "CH4",
    h2o_gas: str = "H2O",
    mu: float = 1.0,
) -> BandTable:
    """A :class:`BandTable` from the layered forward model.

    ``τ_b(ΔΩ) = −ln(y_b(ΔΩ) / y_b(0))`` with ``ΔΩ`` [mol/m²] of ``gas`` added
    by :meth:`~radtrax.atmosphere.AtmosphericState.add_column`;
    ``g = τ_{b₀}`` for one band, ``g = τ_{b₀} − τ_{b₁}`` for two. Each
    ``m`` is a geometry with the sensor at ``mu`` and the sun at
    ``μ₀ = 1 / (m − 1/μ)``; each ``h`` scales the state's ``h2o_gas``
    profile.

    Args:
        model: The forward model.
        state: The background state.
        surface: The surface.
        bands: One band index (SBMP) or two, strong then weak (MBSP, MBMP).
        m_grid: Two-way air-mass factors, each at least ``1 + 1/mu``, ``(M,)``.
        delta_grid: Enhancements [mol/m²], increasing from 0, ``(D,)``.
        h2o_grid: Water-vapour scalings, ``(H,)``, or ``None``.
        gas: The enhanced gas.
        h2o_gas: The water-vapour gas the scalings apply to.
        mu: Cosine of the viewing zenith angle.

    Returns:
        The table.
    """
    if len(bands) not in (1, 2):
        raise ParameterError("band_table_from_model: give one band or two.")
    m_grid = np.asarray(m_grid, dtype=float)
    if np.any(m_grid < 1.0 + 1.0 / mu):
        raise ParameterError(
            f"band_table_from_model: every m must be at least 1 + 1/mu = {1 + 1 / mu}."
        )
    delta_grid = jnp.asarray(delta_grid, dtype=float)

    @eqx.filter_jit
    def observables(
        base: AtmosphericState, geometry: ViewingGeometry, deltas: Float[Array, " D"]
    ) -> Float[Array, " D"]:
        # One compilation for every (m, h): the state and geometry are arguments.
        y0 = model(base, surface, geometry)

        def tau(delta: Float[Array, ""]) -> Float[Array, ""]:
            t = -jnp.log(model(base.add_column(gas, delta), surface, geometry) / y0)
            return t[bands[0]] - t[bands[1]] if len(bands) == 2 else t[bands[0]]

        return jax.vmap(tau)(deltas)

    def observable(m: float, h: float | None) -> np.ndarray:
        base = (
            state
            if h is None
            else eqx.tree_at(lambda s: s.vmr[h2o_gas], state, h * state.vmr[h2o_gas])
        )
        geometry = ViewingGeometry(1.0 / (m - 1.0 / mu), mu)
        return np.asarray(observables(base, geometry, delta_grid))

    if h2o_grid is None:
        return BandTable.from_function(
            lambda m, _: observable(m, None), m_grid, delta_grid
        )
    return BandTable.from_function(
        lambda m, _, h: observable(m, h), m_grid, delta_grid, h2o_grid
    )

The absorption-shape test

A plume-agnostic retrieval: given two sets of pixels, does the ratio of their paired mean radiances have the shape of a gas's transmittance? One-to-one pairs on bands the gas does not absorb, the ratio spectrum, its fit to the forward model's transmittance, the fitted column ΔX̂ and the misfit D_norm (Xiang et al. 2024). plumax.vetting runs it on plume candidates. Theory: the absorption-shape test.

absorption

The absorption-shape test: does a ratio spectrum have a gas's shape?

A plume-agnostic retrieval on arrays. Given two sets of pixels of one radiance cube (in-pixels and out-pixels, as flat indices), it pairs them one to one on bands the gas does not absorb, divides the pairs' mean radiances into a ratio spectrum, fits that ratio with a smooth continuum times the gas's transmittance from radtrax.ForwardModel, and returns the fitted column ΔX̂ and the misfit D_norm (Xiang et al.'s eq. 11): how far the ratio's shape is from the transmittance's. It knows nothing of plumes or candidates; plumax.vetting chooses the pixel sets from plume candidates and decides on the scores. Theory: the absorption-shape test page.

• window: AbsorptionWindow, a gas's fitting and pairing bands and its matched-filter background cut, with EMIT's CH₄ and CO presets (#155).
• transmittance: transmittance_table, the band transmittance of a column enhancement from ForwardModel, exact or constant-coefficient (#154).
• pairing: pair_pixels and ratio_spectrum, one-to-one pairs of two flat pixel-index sets on the pairing bands and the ratio of their mean radiances (#149).
• fit: TransmittanceFit, fit_transmittance and d_norm, the profile fit of the ratio to a Legendre continuum times the transmittance, and its misfit score (#150).

window

A gas's absorption window: its fitting bands, its pairing bands and its background cut.

A window is configuration, not state, so it is a frozen dataclass, as radtrax.spectroscopy.hitran's configs are. Its masks are NumPy booleans: they select bands before anything is traced, and nothing differentiates through them.

AbsorptionWindow dataclass

AbsorptionWindow(gas: str, fit_nm: Ranges, absorbing_nm: Ranges, mf_background_mol_m2: float)

A gas's fitting bands, the bands kept out of pairing, and the background cut.

Ranges are in nm and inclusive at both ends, as in the reference code's configuration (pv/config/config.yaml). Pairing uses every band outside absorbing_nm, so a window whose fitting ranges lie inside its absorbing ranges has disjoint fitting and pairing masks.

Attributes:

Name	Type	Description
gas	str	The gas whose enhancement is fitted, as named in the forward model's state ("CH4").
fit_nm	Ranges	The (lo, hi) ranges, in nm, of the bands the transmittance is fitted on; at least one.
absorbing_nm	Ranges	The (lo, hi) ranges, in nm, where the gas absorbs; pairing uses only bands outside all of them.
mf_background_mol_m2	float	The matched-filter score, in mol/m², below which a pixel counts as background; the absorption-shape test does not use it, the code that chooses the out-pixels does (plumax.vetting's background pool). The presets convert EMIT's ppm·m with 1 ppm·m ≈ 4.03 × 10⁻⁵ mol/m², the molar density of air in the bottom layer of the US standard atmosphere, 954.8 hPa and 285.0 K, times 1 m.

Raises:

Type	Description
ParameterError	If fit_nm is empty, a range is not finite or has lo >= hi, or the threshold is negative or not finite.

Examples:

>>> import numpy as np
>>> from radtrax.retrieval.absorption import AbsorptionWindow
>>> window = AbsorptionWindow.ch4_emit()
>>> window.fit_nm, window.absorbing_nm
(((2100.0, 2440.0),), ((1640.0, 1690.0), (2100.0, 2440.0)))
>>> print(f"{window.mf_background_mol_m2:.3e} mol/m²")
1.209e-03 mol/m²
>>> wl = np.array([1500.0, 1650.0, 2100.0, 2300.0, 2440.0, 2450.0])
>>> window.fit_bands(wl)
array([False, False,  True,  True,  True, False])
>>> window.pairing_bands(wl)
array([ True, False, False, False, False,  True])

ch4_emit classmethod

ch4_emit() -> AbsorptionWindow

EMIT's CH₄ window, as the reference code's configuration sets it.

Fits 2100–2440 nm; keeps CH₄'s 1640–1690 and 2100–2440 nm bands out of pairing; counts pixels below 30 ppm·m as background.

Source code in src/radtrax/retrieval/absorption/window.py

@classmethod
def ch4_emit(cls) -> AbsorptionWindow:
    """EMIT's CH₄ window, as the reference code's configuration sets it.

    Fits 2100–2440 nm; keeps CH₄'s 1640–1690 and 2100–2440 nm bands out
    of pairing; counts pixels below 30 ppm·m as background.
    """
    return cls(
        gas="CH4",
        fit_nm=((2100.0, 2440.0),),
        absorbing_nm=((1640.0, 1690.0), (2100.0, 2440.0)),
        mf_background_mol_m2=30.0 * _MOL_M2_PER_PPM_M,
    )

co_emit classmethod

co_emit() -> AbsorptionWindow

EMIT's CO window, from the source report's Table S1 (Růžička et al.).

Fits 2278–2441 nm and keeps the same range out of pairing; counts pixels below 300 ppm·m as background.

Source code in src/radtrax/retrieval/absorption/window.py

@classmethod
def co_emit(cls) -> AbsorptionWindow:
    """EMIT's CO window, from the source report's Table S1 (Růžička et al.).

    Fits 2278–2441 nm and keeps the same range out of pairing; counts
    pixels below 300 ppm·m as background.
    """
    return cls(
        gas="CO",
        fit_nm=((2278.0, 2441.0),),
        absorbing_nm=((2278.0, 2441.0),),
        mf_background_mol_m2=300.0 * _MOL_M2_PER_PPM_M,
    )

fit_bands

fit_bands(wavelength_nm: Float[ArrayLike, ' B']) -> Bool[ndarray, ' B']

The bands inside any fitting range.

Parameters:

Name	Type	Description	Default
wavelength_nm	Float[ArrayLike, ' B']	Band centres, in nm, (B,).	required

Returns:

Type	Description
Bool[ndarray, ' B']	A boolean mask, (B,).

Source code in src/radtrax/retrieval/absorption/window.py

def fit_bands(
    self, wavelength_nm: Float[ArrayLike, " B"]
) -> Bool[np.ndarray, " B"]:
    """The bands inside any fitting range.

    Args:
        wavelength_nm: Band centres, in nm, ``(B,)``.

    Returns:
        A boolean mask, ``(B,)``.
    """
    return _inside(wavelength_nm, self.fit_nm)

pairing_bands

pairing_bands(wavelength_nm: Float[ArrayLike, ' B']) -> Bool[ndarray, ' B']

The bands outside every absorbing range, which pairing compares.

Parameters:

Name	Type	Description	Default
wavelength_nm	Float[ArrayLike, ' B']	Band centres, in nm, (B,).	required

Returns:

Type	Description
Bool[ndarray, ' B']	A boolean mask, (B,).

Source code in src/radtrax/retrieval/absorption/window.py

def pairing_bands(
    self, wavelength_nm: Float[ArrayLike, " B"]
) -> Bool[np.ndarray, " B"]:
    """The bands outside every absorbing range, which pairing compares.

    Args:
        wavelength_nm: Band centres, in nm, ``(B,)``.

    Returns:
        A boolean mask, ``(B,)``.
    """
    return ~_inside(wavelength_nm, self.absorbing_nm)

transmittance

The gas's transmittance: a scene's band transmittance on a grid of columns.

E_b(ΔX) = y_b(x + ΔX e) / y_b(x) from the same :class:radtrax.ForwardModel the matched filter's target comes from (:func:radtrax.unit_absorption_target), so the fit and the filter see the same gas and the curve of growth is built in. It replaces the reference code's MODTRAN tables. Theory: the absorption-shape test page, section 2.

Shapes: D column enhancements, B bands, L layers.

transmittance_table

transmittance_table(model: ForwardModel, state: AtmosphericState, geometry: ViewingGeometry, delta_grid: Float[ArrayLike, ' D'], *, gas: str = 'CH4', profile: Float[ArrayLike, ' L'] | None = None, linear: bool = False) -> Float[Array, 'D B']

The band transmittance E_b(ΔX) of a column enhancement, on a grid.

.. math:: E_b(\Delta X) = \frac{y_b(x + \Delta X\, e)}{y_b(x)}, \qquad E^{\text{lin}}_b(\Delta X) = e^{-\varepsilon_b \Delta X}, \quad \varepsilon_b = -\left.\frac{\partial \ln y_b}{\partial \Delta X} \right|_0 ,

with y the model's band values for a unit Lambertian albedo, ΔX added by :meth:~radtrax.atmosphere.AtmosphericState.add_column with weights e from profile, and ε_b = −t_b from :func:radtrax.unit_absorption_target. Beer–Lambert radiance is linear in a Lambertian albedo, so the albedo cancels in the ratio exactly per wavenumber, and after band integration to the extent the albedo is flat across a band's response. The background y(x) is its own forward model run at ΔX = 0, so the grid need not contain 0; a row with ΔX = 0 is exactly 1. E is decreasing and convex in ΔX, and E^lin ≤ E (Jensen): a constant-coefficient fit of a strong enhancement reads its column low.

One forward model per grid point and one for the background, jax.vmap over them. It is differentiable in the model's leaves (the spectral shift, say), the state and the geometry.

Parameters:

Name	Type	Description	Default
model	ForwardModel	The forward model.	required
state	AtmosphericState	The background atmosphere.	required
geometry	ViewingGeometry	Sun and sensor directions.	required
delta_grid	Float[ArrayLike, ' D']	Column enhancements ΔX [mol/m²], strictly increasing, (D,); checked when concrete. It need not start at 0 (BandTable and fit_transmittance's ΔX ≥ 0 want it to).	required
gas	str	The enhanced gas.	'CH4'
profile	Float[ArrayLike, ' L'] | None	Relative weight of the enhancement per layer, (L,); the default puts it all in the bottom layer.	None
linear	bool	Return the constant-coefficient model E^lin instead.	False

Returns:

Type	Description
Float[Array, 'D B']	E, one row per enhancement, (D, B).

Raises:

Type	Description
MissingKeyError	If the state carries no gas.
ShapeError	If delta_grid is not (D,).
GridError	If a concrete delta_grid does not strictly increase.

Examples:

>>> import jax.numpy as jnp
>>> from radtrax import ForwardModel
>>> from radtrax.atmosphere import ViewingGeometry, afgl_us_standard
>>> from radtrax.instrument import BandInstrument
>>> from radtrax.solvers import BeerLambert
>>> from radtrax.spectroscopy import CrossSectionTable, TableOpacity
>>> nu = jnp.array([4300.0, 4300.1, 4300.2])
>>> sigma = jnp.array([1e-21, 4e-21, 1e-21])  # cm²/molecule
>>> table = CrossSectionTable(
...     jnp.log(jnp.array([1e3, 1.2e5])),
...     jnp.array([200.0, 310.0]),
...     jnp.broadcast_to(sigma, (2, 2, 3)),
... )
>>> model = ForwardModel(
...     TableOpacity(nu, {"CH4": table}),
...     BeerLambert(),
...     BandInstrument(nu, jnp.array([[1.0, 0.0, 0.0], [0.0, 0.5, 0.5]])),
... )
>>> state = afgl_us_standard(z_top_km=3.0, gases=["CH4"])
>>> geometry = ViewingGeometry(1.0, 1.0)
>>> grid = jnp.array([0.0, 0.1, 0.2])  # mol/m²
>>> exact = transmittance_table(model, state, geometry, grid)
>>> lin = transmittance_table(model, state, geometry, grid, linear=True)
>>> exact.shape, bool(jnp.all(exact[0] == 1.0))
((3, 2), True)
>>> # One wavenumber: E is exactly exponential; two: the curve bends.
>>> float(jnp.max(jnp.abs(exact[:, 0] - lin[:, 0]))) < 1e-12
True
>>> bool(jnp.all(exact[1:, 1] > lin[1:, 1]))
True

Source code in src/radtrax/retrieval/absorption/transmittance.py

def transmittance_table(
    model: ForwardModel,
    state: AtmosphericState,
    geometry: ViewingGeometry,
    delta_grid: Float[ArrayLike, " D"],
    *,
    gas: str = "CH4",
    profile: Float[ArrayLike, " L"] | None = None,
    linear: bool = False,
) -> Float[Array, "D B"]:
    """The band transmittance ``E_b(ΔX)`` of a column enhancement, on a grid.

    .. math::
        E_b(\\Delta X) = \\frac{y_b(x + \\Delta X\\, e)}{y_b(x)},
        \\qquad
        E^{\\text{lin}}_b(\\Delta X) = e^{-\\varepsilon_b \\Delta X},
        \\quad
        \\varepsilon_b = -\\left.\\frac{\\partial \\ln y_b}{\\partial \\Delta X}
        \\right|_0 ,

    with ``y`` the model's band values for a unit Lambertian albedo, ``ΔX``
    added by :meth:`~radtrax.atmosphere.AtmosphericState.add_column` with
    weights ``e`` from ``profile``, and ``ε_b = −t_b`` from
    :func:`radtrax.unit_absorption_target`. Beer–Lambert radiance is linear
    in a Lambertian albedo, so the albedo cancels in the ratio exactly per
    wavenumber, and after band integration to the extent the albedo is flat
    across a band's response. The background ``y(x)`` is its own forward
    model run at ``ΔX = 0``, so the grid need not contain 0; a row with
    ``ΔX = 0`` is exactly 1. ``E`` is decreasing and
    convex in ``ΔX``, and ``E^lin ≤ E`` (Jensen): a constant-coefficient fit
    of a strong enhancement reads its column low.

    One forward model per grid point and one for the background,
    ``jax.vmap`` over them. It is
    differentiable in the model's leaves (the spectral shift, say), the
    state and the geometry.

    Args:
        model: The forward model.
        state: The background atmosphere.
        geometry: Sun and sensor directions.
        delta_grid: Column enhancements ``ΔX`` [mol/m²], strictly
            increasing, ``(D,)``; checked when concrete. It need not start
            at 0 (``BandTable`` and ``fit_transmittance``'s ``ΔX ≥ 0`` want
            it to).
        gas: The enhanced gas.
        profile: Relative weight of the enhancement per layer, ``(L,)``; the
            default puts it all in the bottom layer.
        linear: Return the constant-coefficient model ``E^lin`` instead.

    Returns:
        ``E``, one row per enhancement, ``(D, B)``.

    Raises:
        MissingKeyError: If the state carries no ``gas``.
        ShapeError: If ``delta_grid`` is not ``(D,)``.
        GridError: If a concrete ``delta_grid`` does not strictly increase.

    Examples:
        >>> import jax.numpy as jnp
        >>> from radtrax import ForwardModel
        >>> from radtrax.atmosphere import ViewingGeometry, afgl_us_standard
        >>> from radtrax.instrument import BandInstrument
        >>> from radtrax.solvers import BeerLambert
        >>> from radtrax.spectroscopy import CrossSectionTable, TableOpacity
        >>> nu = jnp.array([4300.0, 4300.1, 4300.2])
        >>> sigma = jnp.array([1e-21, 4e-21, 1e-21])  # cm²/molecule
        >>> table = CrossSectionTable(
        ...     jnp.log(jnp.array([1e3, 1.2e5])),
        ...     jnp.array([200.0, 310.0]),
        ...     jnp.broadcast_to(sigma, (2, 2, 3)),
        ... )
        >>> model = ForwardModel(
        ...     TableOpacity(nu, {"CH4": table}),
        ...     BeerLambert(),
        ...     BandInstrument(nu, jnp.array([[1.0, 0.0, 0.0], [0.0, 0.5, 0.5]])),
        ... )
        >>> state = afgl_us_standard(z_top_km=3.0, gases=["CH4"])
        >>> geometry = ViewingGeometry(1.0, 1.0)
        >>> grid = jnp.array([0.0, 0.1, 0.2])  # mol/m²
        >>> exact = transmittance_table(model, state, geometry, grid)
        >>> lin = transmittance_table(model, state, geometry, grid, linear=True)
        >>> exact.shape, bool(jnp.all(exact[0] == 1.0))
        ((3, 2), True)
        >>> # One wavenumber: E is exactly exponential; two: the curve bends.
        >>> float(jnp.max(jnp.abs(exact[:, 0] - lin[:, 0]))) < 1e-12
        True
        >>> bool(jnp.all(exact[1:, 1] > lin[1:, 1]))
        True
    """
    if gas not in state.vmr:
        raise MissingKeyError(
            f"transmittance_table: the state has no mixing ratio for {gas!r} "
            f"(have {sorted(state.vmr)})."
        )
    delta = jnp.asarray(delta_grid, dtype=state.vmr[gas].dtype)
    _check_delta_grid(delta)
    surface = Lambertian(1.0)
    if linear:
        # t = ∂ ln y / ∂ΔX |₀ = −ε;  E^lin = exp(ΔX t)                     (D, B)
        t = unit_absorption_target(
            model, state, surface, geometry, gas=gas, profile=profile
        )
        return jnp.exp(jnp.outer(delta, t))

    def bands(d: Float[Array, ""]) -> Float[Array, " B"]:
        return model(state.add_column(gas, d, profile), surface, geometry)

    # y(0) and y(ΔX_d) for every d, one vmap: row 0 is the background
    # whatever the grid holds                                      (1 + D, B)
    y = jax.vmap(bands)(jnp.concatenate([jnp.zeros_like(delta[:1]), delta]))
    # E_d = y(ΔX_d) / y(0), set to 1 exactly where ΔX_d = 0 (XLA may divide
    # by a reciprocal, which leaves y_0 / y_0 one ulp off)               (D, B)
    return jnp.where((delta == 0)[:, None], 1.0, y[1:] / y[0])

pairing

One-to-one spectral pairs and the ratio of their mean radiances.

The pairing and ratio steps of the absorption-shape test (steps 3–4 of Xiang et al.'s method, as the reference code runs them, pv/src/pv/pv.py): given two pixel sets, each in-pixel is paired with one out-pixel, the assignment minimising the total Euclidean distance between their radiances on bands the gas does not absorb, and the best-matched half of the pairs is kept; the ratio spectrum divides the kept pairs' mean in-pixel radiance by their mean out-pixel radiance. Which pixels form the two sets (a plume candidate's seeds and its background pool, say) is the caller's choice: plumax.vetting makes it.

Pairing is NumPy and SciPy, off the gradient path (an assignment is discrete); the distance it minimises is the one private step _pairing_distance. The ratio is JAX and differentiable in the cube.

Shapes: n_in in-pixels, n_out out-pixels, N kept pairs.

pair_pixels

pair_pixels(cube: Float[ArrayLike, 'H W B'], in_pixels: Int[ArrayLike, ' n_in'], out_pixels: Int[ArrayLike, ' n_out'], pairing_bands: Bool[ArrayLike, ' B'], *, keep: float = 0.5) -> tuple[Int[ndarray, ' N'], Int[ndarray, ' N']]

One-to-one pairs of in-pixels and out-pixels, the best-matched kept.

With s(p) pixel p's radiance on the pairing bands, the in-pixels I and the out-pixels O, the pairing is the injective assignment

π* = argmin_{π: I → O injective} ∑_{i ∈ I} ‖s(i) − s(π(i))‖₂,

solved exactly by scipy.optimize.linear_sum_assignment; the N = int(keep · |I|) pairs of smallest distance are kept. One to one, so no out-pixel's noise or surface mismatch enters two pairs.

The two sets are flat indices into the image, r · W + c for row r and column c; a boolean (H, W) mask converts with np.flatnonzero(mask). Nothing here knows what the sets are.

Parameters:

Name	Type	Description	Default
cube	Float[ArrayLike, 'H W B']	Radiance, (H, W, B).	required
in_pixels	Int[ArrayLike, ' n_in']	The set I (the pixels whose absorption is tested), distinct flat indices into H · W, (n_in,).	required
out_pixels	Int[ArrayLike, ' n_out']	The set O (the reference pixels they are paired from), distinct flat indices into H · W and disjoint from in_pixels, (n_out,) with n_out ≥ n_in.	required
pairing_bands	Bool[ArrayLike, ' B']	The bands compared, (B,) (AbsorptionWindow.pairing_bands: those outside every absorbing range).	required
keep	float	Fraction of the pairs kept, in (0, 1]; the reference code's SPECTRAL_SIMILARITY_FRACTION_RETAINED, 0.5.	0.5

Returns:

Name	Type	Description
	Int[ndarray, ' N']	(in_idx, out_idx), flat pixel indices into (H, W), (N,)
each	Int[ndarray, ' N']	pair n is in-pixel in_idx[n] and out-pixel
	tuple[Int[ndarray, ' N'], Int[ndarray, ' N']]	out_idx[n], ordered from the best-matched pair to the worst
	tuple[Int[ndarray, ' N'], Int[ndarray, ' N']]	kept. No index repeats in either.

Raises:

Type	Description
ParameterError	If keep is outside (0, 1], or an index set is not integer, has an index outside [0, H · W) or repeats a pixel, or the two sets share a pixel.
ShapeError	If cube is not 3-D, an index set is not 1-D, or pairing_bands is not (B,).
DegenerateInputError	If there are no in-pixels, there are fewer out-pixels than in-pixels, no band is a pairing band, keep keeps no pair, or an in- or out-pixel has a non-finite radiance on a pairing band.

Examples:

>>> import numpy as np
>>> from radtrax.retrieval.absorption import pair_pixels
>>> cube = np.array([[[1.0, 9.0], [3.0, 9.0], [3.1, 0.0], [0.8, 0.0]]])
>>> in_pixels, out_pixels = np.array([0, 1]), np.array([2, 3])
>>> bands = np.array([True, False])
>>> pair_pixels(cube, in_pixels, out_pixels, bands, keep=1.0)
(array([1, 0]), array([2, 3]))
>>> pair_pixels(cube, in_pixels, out_pixels, bands)
(array([1]), array([2]))

Source code in src/radtrax/retrieval/absorption/pairing.py

def pair_pixels(
    cube: Float[ArrayLike, "H W B"],
    in_pixels: Int[ArrayLike, " n_in"],
    out_pixels: Int[ArrayLike, " n_out"],
    pairing_bands: Bool[ArrayLike, " B"],
    *,
    keep: float = 0.5,
) -> tuple[Int[np.ndarray, " N"], Int[np.ndarray, " N"]]:
    """One-to-one pairs of in-pixels and out-pixels, the best-matched kept.

    With ``s(p)`` pixel ``p``'s radiance on the pairing bands, the in-pixels
    ``I`` and the out-pixels ``O``, the pairing is the injective assignment

    ``π* = argmin_{π: I → O injective} ∑_{i ∈ I} ‖s(i) − s(π(i))‖₂``,

    solved exactly by ``scipy.optimize.linear_sum_assignment``; the
    ``N = int(keep · |I|)`` pairs of smallest distance are kept. One to one,
    so no out-pixel's noise or surface mismatch enters two pairs.

    The two sets are flat indices into the image, ``r · W + c`` for row
    ``r`` and column ``c``; a boolean ``(H, W)`` mask converts with
    ``np.flatnonzero(mask)``. Nothing here knows what the sets are.

    Args:
        cube: Radiance, ``(H, W, B)``.
        in_pixels: The set ``I`` (the pixels whose absorption is tested),
            distinct flat indices into ``H · W``, ``(n_in,)``.
        out_pixels: The set ``O`` (the reference pixels they are paired
            from), distinct flat indices into ``H · W`` and disjoint from
            ``in_pixels``, ``(n_out,)`` with ``n_out ≥ n_in``.
        pairing_bands: The bands compared, ``(B,)``
            (``AbsorptionWindow.pairing_bands``: those outside every
            absorbing range).
        keep: Fraction of the pairs kept, in (0, 1]; the reference code's
            ``SPECTRAL_SIMILARITY_FRACTION_RETAINED``, 0.5.

    Returns:
        ``(in_idx, out_idx)``, flat pixel indices into ``(H, W)``, ``(N,)``
        each: pair ``n`` is in-pixel ``in_idx[n]`` and out-pixel
        ``out_idx[n]``, ordered from the best-matched pair to the worst
        kept. No index repeats in either.

    Raises:
        ParameterError: If ``keep`` is outside (0, 1], or an index set is
            not integer, has an index outside ``[0, H · W)`` or repeats a
            pixel, or the two sets share a pixel.
        ShapeError: If ``cube`` is not 3-D, an index set is not 1-D, or
            ``pairing_bands`` is not ``(B,)``.
        DegenerateInputError: If there are no in-pixels, there are fewer
            out-pixels than in-pixels, no band is a pairing band, ``keep``
            keeps no pair, or an in- or out-pixel has a non-finite radiance
            on a pairing band.

    Examples:
        >>> import numpy as np
        >>> from radtrax.retrieval.absorption import pair_pixels
        >>> cube = np.array([[[1.0, 9.0], [3.0, 9.0], [3.1, 0.0], [0.8, 0.0]]])
        >>> in_pixels, out_pixels = np.array([0, 1]), np.array([2, 3])
        >>> bands = np.array([True, False])
        >>> pair_pixels(cube, in_pixels, out_pixels, bands, keep=1.0)
        (array([1, 0]), array([2, 3]))
        >>> pair_pixels(cube, in_pixels, out_pixels, bands)
        (array([1]), array([2]))
    """
    if not 0.0 < keep <= 1.0:
        raise ParameterError(f"pair_pixels: keep must be in (0, 1], got {keep}")
    cube = np.asarray(cube, dtype=np.float64)
    i_in = np.asarray(in_pixels)
    i_out = np.asarray(out_pixels)
    pairing_bands = np.asarray(pairing_bands, dtype=bool)
    if cube.ndim != 3:
        raise ShapeError(f"pair_pixels: cube must be (H, W, B), got {cube.shape}")
    n_pixels = cube.shape[0] * cube.shape[1]
    _check_pixel_set("in_pixels", i_in, n_pixels)
    _check_pixel_set("out_pixels", i_out, n_pixels)
    if pairing_bands.shape != cube.shape[2:]:
        raise ShapeError(
            f"pair_pixels: pairing_bands has shape {pairing_bands.shape}, the "
            f"cube has {cube.shape[2]} bands"
        )
    if np.intersect1d(i_in, i_out).size:
        raise ParameterError(
            "pair_pixels: in_pixels and out_pixels share a pixel; they must be disjoint"
        )
    n_kept = int(keep * i_in.size)
    if i_in.size == 0:
        raise DegenerateInputError("pair_pixels: there are no in-pixels")
    if i_out.size < i_in.size:
        raise DegenerateInputError(
            f"pair_pixels: there are {i_out.size} out-pixels, fewer than the "
            f"{i_in.size} in-pixels; one-to-one pairing needs at least as many"
        )
    if not pairing_bands.any():
        raise DegenerateInputError("pair_pixels: no band is a pairing band")
    if n_kept == 0:
        raise DegenerateInputError(
            f"pair_pixels: keep={keep} of {i_in.size} in-pixels keeps no pair"
        )
    flat = cube.reshape(-1, cube.shape[-1])
    s_in = flat[i_in][:, pairing_bands]  # (n_in, B_o)
    s_out = flat[i_out][:, pairing_bands]  # (n_out, B_o)
    if not (np.isfinite(s_in).all() and np.isfinite(s_out).all()):
        raise DegenerateInputError(
            "pair_pixels: an in- or out-pixel has a non-finite radiance on a "
            "pairing band"
        )

    from scipy.optimize import linear_sum_assignment

    distance = _pairing_distance(s_in, s_out)  # (n_in, n_out)
    rows, cols = linear_sum_assignment(distance)  # rows = 0..n_in-1
    best = np.argsort(distance[rows, cols], kind="stable")[:n_kept]
    return i_in[rows[best]].astype(np.intp), i_out[cols[best]].astype(np.intp)

ratio_spectrum

ratio_spectrum(cube: Float[ArrayLike, 'H W B'], in_idx: Int[ArrayLike, ' N'], out_idx: Int[ArrayLike, ' N']) -> Float[Array, ' B']

The ratio of the pairs' mean in-pixel and mean out-pixel radiances.

R(λ) = ∑_n w_n(λ) L_in,n(λ) / ∑_n w_n(λ) L_out,n(λ), with w_n(λ) = 1[L_in,n(λ) and L_out,n(λ) finite, L_out,n(λ) > 0]: a pair is left out of a band where either of its radiances is unusable (a NaN or infinite one, or an out-pixel radiance that is not positive). When every pair is usable, R is the ratio of the means, (1/N) ∑_n L_in,n / (1/N) ∑_n L_out,n. Averaging before dividing keeps the ratio's noise bias at the size of the mean's, not of one pixel's.

Differentiable in cube (the weights are piecewise constant; a radiance left out has a zero gradient, and does not make its partner's NaN) and traceable under jax.jit in cube; the indices are NumPy, checked eagerly, and enter as constants.

Parameters:

Name	Type	Description	Default
cube	Float[ArrayLike, 'H W B']	Radiance, (H, W, B).	required
in_idx	Int[ArrayLike, ' N']	Flat indices of the pairs' in-pixels, (N,), from pair_pixels.	required
out_idx	Int[ArrayLike, ' N']	Flat indices of their out-pixels, (N,).	required

Returns:

Type	Description
Float[Array, ' B']	R, (B,); NaN in a band where no pair is usable (its gradient
Float[Array, ' B']	there is zero, not NaN).

Raises:

Type	Description
ShapeError	If cube is not 3-D or the indices are not two (N,) arrays of one length.
DegenerateInputError	If there are no pairs.
ParameterError	If an index is outside [0, H · W).

Examples:

>>> import numpy as np
>>> from radtrax.retrieval.absorption import ratio_spectrum
>>> cube = np.array([[[1.0, 2.0], [3.0, 4.0], [2.0, 8.0], [2.0, 0.0]]])
>>> print(ratio_spectrum(cube, np.array([0, 1]), np.array([2, 3])))
[1.   0.25]

Source code in src/radtrax/retrieval/absorption/pairing.py

def ratio_spectrum(
    cube: Float[ArrayLike, "H W B"],
    in_idx: Int[ArrayLike, " N"],
    out_idx: Int[ArrayLike, " N"],
) -> Float[Array, " B"]:
    """The ratio of the pairs' mean in-pixel and mean out-pixel radiances.

    ``R(λ) = ∑_n w_n(λ) L_in,n(λ) / ∑_n w_n(λ) L_out,n(λ)``, with
    ``w_n(λ) = 1[L_in,n(λ) and L_out,n(λ) finite, L_out,n(λ) > 0]``: a pair
    is left out of a band where either of its radiances is unusable (a NaN
    or infinite one, or an out-pixel radiance that is not positive). When
    every pair is usable, ``R`` is the ratio of the means,
    ``(1/N) ∑_n L_in,n / (1/N) ∑_n L_out,n``. Averaging before dividing
    keeps the ratio's noise bias at the size of the mean's, not of one
    pixel's.

    Differentiable in ``cube`` (the weights are piecewise constant; a
    radiance left out has a zero gradient, and does not make its partner's
    NaN) and
    traceable under ``jax.jit`` in ``cube``; the indices are NumPy, checked
    eagerly, and enter as constants.

    Args:
        cube: Radiance, ``(H, W, B)``.
        in_idx: Flat indices of the pairs' in-pixels, ``(N,)``, from
            [`pair_pixels`][radtrax.retrieval.absorption.pair_pixels].
        out_idx: Flat indices of their out-pixels, ``(N,)``.

    Returns:
        ``R``, ``(B,)``; NaN in a band where no pair is usable (its gradient
        there is zero, not NaN).

    Raises:
        ShapeError: If ``cube`` is not 3-D or the indices are not two
            ``(N,)`` arrays of one length.
        DegenerateInputError: If there are no pairs.
        ParameterError: If an index is outside ``[0, H · W)``.

    Examples:
        >>> import numpy as np
        >>> from radtrax.retrieval.absorption import ratio_spectrum
        >>> cube = np.array([[[1.0, 2.0], [3.0, 4.0], [2.0, 8.0], [2.0, 0.0]]])
        >>> print(ratio_spectrum(cube, np.array([0, 1]), np.array([2, 3])))
        [1.   0.25]
    """
    cube = jnp.asarray(cube)
    in_idx = np.asarray(in_idx)
    out_idx = np.asarray(out_idx)
    if cube.ndim != 3:
        raise ShapeError(f"ratio_spectrum: cube must be (H, W, B), got {cube.shape}")
    if in_idx.ndim != 1 or in_idx.shape != out_idx.shape:
        raise ShapeError(
            f"ratio_spectrum: in_idx and out_idx must be (N,) of one length, got "
            f"{in_idx.shape} and {out_idx.shape}"
        )
    if in_idx.size == 0:
        raise DegenerateInputError("ratio_spectrum: there are no pairs")
    n_pixels = cube.shape[0] * cube.shape[1]
    for name, idx in (("in_idx", in_idx), ("out_idx", out_idx)):
        if idx.min() < 0 or idx.max() >= n_pixels:
            raise ParameterError(
                f"ratio_spectrum: {name} must lie in [0, {n_pixels}), got "
                f"[{idx.min()}, {idx.max()}]"
            )
    flat = cube.reshape(-1, cube.shape[-1])
    l_in = flat[in_idx]  # (N, B)
    l_out = flat[out_idx]  # (N, B)
    # A pair counts in a band only where both radiances are finite and the
    # out-pixel's is positive: a NaN in-pixel would otherwise reach the sum,
    # and through it every out-pixel's gradient                      (N, B)
    w = jnp.isfinite(l_in) & jnp.isfinite(l_out) & (l_out > 0)
    numerator = jnp.sum(jnp.where(w, l_in, 0.0), axis=0)  # (B,)
    denominator = jnp.sum(jnp.where(w, l_out, 0.0), axis=0)  # (B,)
    usable = jnp.any(w, axis=0)  # (B,)
    # Double where: a safe denominator where no pair is usable, so the NaN
    # returned there does not reach the gradient.
    safe = jnp.where(usable, denominator, 1.0)
    return jnp.where(usable, numerator / safe, jnp.nan)

fit

The profile fit of the ratio spectrum and its score, D_norm.

The fit and the score of the absorption-shape test (step 5 of Xiang et al.'s method): fit the pair ratio with a smooth continuum times the gas's transmittance, R(λ) ≈ P(λ; c) E(λ; ΔX), and score the misfit with their eq. 11. The continuum is a Legendre series on the window (the basis of radtrax.surface.PolynomialAlbedo); for fixed ΔX the model is linear in c, so ĉ(ΔX) is a least-squares solve (lineax's QR) and ΔX̂ the grid point of least residual, refined between the nodes by a parabola through the residuals: no starting point, the global minimum over the grid. Theory: the absorption-shape test page, section 4.

The fit weights each band by _band_weights: 1 where the ratio is finite, 0 where it is not (ratio_spectrum returns NaN in a band with no usable pair). That is ordinary least squares on the usable bands, and the one step a noise model for the ratio (radtrax#165) would replace.

Shapes: B bands of the fitting window, D column enhancements of the grid, K continuum coefficients (degree K − 1).

TransmittanceFit

Bases: Module

The profile fit of R ≈ P(λ; c) E(λ; ΔX).

Attributes:

Name	Type	Description
delta	Float[Array, '']	ΔX̂, the column of least residual [mol/m²], refined between the grid's nodes, ().
coefficients	Float[Array, ' K']	ĉ, the continuum's Legendre coefficients (dimensionless), (K,).
continuum	Float[Array, ' B']	P(λ; ĉ) = Σ_k ĉ_k P_k(x_λ), (B,).
transmittance	Float[Array, ' B']	E(λ; ΔX̂), ln E interpolated linearly between the table's two rows around ΔX̂ (the row itself on a node), (B,).
residual_ss	Float[Array, '']	Σ_λ w_λ (R(λ) − P(λ; ĉ) E(λ; ΔX̂))², the weighted residual sum of squares at the fit, ().

fit_transmittance

fit_transmittance(ratio: Float[ArrayLike, ' B'], table: Float[ArrayLike, 'D B'], delta_grid: Float[ArrayLike, ' D'], wavelength_nm: Float[ArrayLike, ' B'], *, degree: int = 2) -> TransmittanceFit

Fit a Legendre continuum times the transmittance, profiling over ΔX.

With x_λ = 2 (λ − λ_min) / (λ_max − λ_min) − 1 the band centres mapped onto [−1, 1], V_λk = P_k(x_λ) the Legendre polynomials (K = degree + 1), A(ΔX) = diag(E(ΔX)) V and w_λ the band weights (1 where R is finite, else 0),

ĉ(ΔX) = argmin_c Σ_λ w_λ (R(λ) − (A(ΔX) c)_λ)², solved by QR (lineax.QR) at each grid column;

k = argmin_d RSS(ΔX_d), the grid column of least residual, then refined between the nodes: ΔX̂ is the vertex of the parabola through RSS at three neighbouring nodes around k (the first or last three at a grid end), clipped to their span; E(ΔX̂) interpolates ln E linearly between the two rows around it (so E > 0), and ĉ is solved again there. The refined fit replaces the node's only where the parabola is convex and the refit lowers the residual by more than rounding (10⁻¹² of ‖R‖²), so a fit exact on a node stays there. For RSS ≈ a u² + b u³ near its minimum (u the distance to it) and spacing h, the vertex is off by (b / 2a)(3δ² − h²), δ the middle node's distance to the minimum: at most |b| h² / a, against h / 2 for the node alone.

The grid should start at 0 (as transmittance_table's does by default), which imposes ΔX ≥ 0. Bands where ratio is NaN are left out of the fit; continuum and transmittance are still evaluated there.

Piecewise differentiable: ΔX̂ and the fields evaluated at it are differentiable in ratio and table while the node k and the choice between the node and the refined fit stay put (they jump where k changes), and their gradient in a NaN band is zero.

Parameters:

Name	Type	Description	Default
ratio	Float[ArrayLike, ' B']	The ratio spectrum R on the fitting bands, (B,), from ratio_spectrum masked with AbsorptionWindow.fit_bands; NaN in a band with no usable pair.	required
table	Float[ArrayLike, 'D B']	E(λ; ΔX_d) on the same bands, (D, B), from transmittance_table; finite and positive.	required
delta_grid	Float[ArrayLike, ' D']	The columns ΔX_d of the table's rows [mol/m²], (D,).	required
wavelength_nm	Float[ArrayLike, ' B']	The band centres [nm], (B,); their range sets the window the continuum's basis is orthogonal on.	required
degree	int	The continuum's degree, K − 1; 2 by default.	2

Returns:

Type	Description
TransmittanceFit	The TransmittanceFit
TransmittanceFit	at ΔX̂.

Raises:

Type	Description
ParameterError	If degree < 0 or degree + 1 > B.
ShapeError	If ratio or wavelength_nm is not (B,), table is not (D, B) with ratio's B, or delta_grid is not (D,) with table's D.
DegenerateInputError	If fewer than K bands of ratio are finite. Under jit this check runs with eqx.error_if and raises a RuntimeError instead (equinox.EquinoxRuntimeError under eqx.filter_jit).

Examples:

>>> import jax.numpy as jnp
>>> from radtrax.retrieval.absorption import fit_transmittance
>>> wl = jnp.linspace(2200.0, 2400.0, 6)
>>> grid = jnp.array([0.0, 0.1, 0.2])
>>> table = jnp.exp(
...     -grid[:, None] * jnp.array([1.0, 3.0, 0.5, 4.0, 1.0, 2.0])
... )
>>> ratio = (1.0 + 0.1 * jnp.linspace(-1.0, 1.0, 6)) * table[1]
>>> fit = fit_transmittance(ratio, table, grid, wl, degree=1)
>>> print(f"{fit.delta:.1f}", jnp.round(fit.coefficients, 6))
0.1 [1.  0.1]

Source code in src/radtrax/retrieval/absorption/fit.py

def fit_transmittance(
    ratio: Float[ArrayLike, " B"],
    table: Float[ArrayLike, "D B"],
    delta_grid: Float[ArrayLike, " D"],
    wavelength_nm: Float[ArrayLike, " B"],
    *,
    degree: int = 2,
) -> TransmittanceFit:
    """Fit a Legendre continuum times the transmittance, profiling over ``ΔX``.

    With ``x_λ = 2 (λ − λ_min) / (λ_max − λ_min) − 1`` the band centres
    mapped onto ``[−1, 1]``, ``V_λk = P_k(x_λ)`` the Legendre polynomials
    (``K = degree + 1``), ``A(ΔX) = diag(E(ΔX)) V`` and ``w_λ`` the band
    weights (1 where ``R`` is finite, else 0),

    ``ĉ(ΔX) = argmin_c Σ_λ w_λ (R(λ) − (A(ΔX) c)_λ)²``, solved by QR
    (``lineax.QR``) at each grid column;

    ``k = argmin_d RSS(ΔX_d)``, the grid column of least residual, then
    refined between the nodes: ``ΔX̂`` is the vertex of the parabola through
    ``RSS`` at three neighbouring nodes around ``k`` (the first or last three
    at a grid end), clipped to their span; ``E(ΔX̂)`` interpolates ``ln E``
    linearly between the two rows around it (so ``E > 0``), and ``ĉ`` is
    solved again there. The refined fit replaces the node's only where the
    parabola is convex and the refit lowers the residual by more than
    rounding (10⁻¹² of ``‖R‖²``), so a fit exact on a node stays there. For ``RSS ≈ a u² + b u³`` near its minimum (``u``
    the distance to it) and spacing ``h``, the vertex is off by
    ``(b / 2a)(3δ² − h²)``, ``δ`` the middle node's distance to the
    minimum: at most ``|b| h² / a``, against ``h / 2`` for the node alone.

    The grid should start at 0 (as `transmittance_table`'s does by
    default), which imposes ``ΔX ≥ 0``. Bands where ``ratio`` is NaN are
    left out of the fit; ``continuum`` and ``transmittance`` are still
    evaluated there.

    Piecewise differentiable: ``ΔX̂`` and the fields evaluated at it are
    differentiable in ``ratio`` and ``table`` while the node ``k`` and the
    choice between the node and the refined fit stay put (they jump where
    ``k`` changes), and their gradient in a NaN band is zero.

    Args:
        ratio: The ratio spectrum ``R`` on the fitting bands, ``(B,)``, from
            [`ratio_spectrum`][radtrax.retrieval.absorption.ratio_spectrum]
            masked with ``AbsorptionWindow.fit_bands``; NaN in a band with no
            usable pair.
        table: ``E(λ; ΔX_d)`` on the same bands, ``(D, B)``, from
            [`transmittance_table`][radtrax.retrieval.absorption.transmittance_table];
            finite and positive.
        delta_grid: The columns ``ΔX_d`` of the table's rows [mol/m²],
            ``(D,)``.
        wavelength_nm: The band centres [nm], ``(B,)``; their range sets the
            window the continuum's basis is orthogonal on.
        degree: The continuum's degree, ``K − 1``; 2 by default.

    Returns:
        The [`TransmittanceFit`][radtrax.retrieval.absorption.TransmittanceFit]
        at ``ΔX̂``.

    Raises:
        ParameterError: If ``degree < 0`` or ``degree + 1 > B``.
        ShapeError: If ``ratio`` or ``wavelength_nm`` is not ``(B,)``,
            ``table`` is not ``(D, B)`` with ``ratio``'s ``B``, or
            ``delta_grid`` is not ``(D,)`` with ``table``'s ``D``.
        DegenerateInputError: If fewer than ``K`` bands of ``ratio`` are
            finite. Under ``jit`` this check runs with ``eqx.error_if`` and
            raises a ``RuntimeError`` instead
            (``equinox.EquinoxRuntimeError`` under ``eqx.filter_jit``).

    Examples:
        >>> import jax.numpy as jnp
        >>> from radtrax.retrieval.absorption import fit_transmittance
        >>> wl = jnp.linspace(2200.0, 2400.0, 6)
        >>> grid = jnp.array([0.0, 0.1, 0.2])
        >>> table = jnp.exp(
        ...     -grid[:, None] * jnp.array([1.0, 3.0, 0.5, 4.0, 1.0, 2.0])
        ... )
        >>> ratio = (1.0 + 0.1 * jnp.linspace(-1.0, 1.0, 6)) * table[1]
        >>> fit = fit_transmittance(ratio, table, grid, wl, degree=1)
        >>> print(f"{fit.delta:.1f}", jnp.round(fit.coefficients, 6))
        0.1 [1.  0.1]
    """
    ratio = jnp.asarray(ratio)
    table = jnp.asarray(table)
    delta_grid = jnp.asarray(delta_grid)
    wavelength_nm = jnp.asarray(wavelength_nm)
    _check_fit_inputs(ratio, table, delta_grid, wavelength_nm, degree)

    usable = _require_bands(
        _usable(ratio),
        degree + 1,
        f"fit_transmittance: fewer than degree + 1 = {degree + 1} bands of the "
        "ratio are finite",
    )
    return _profile_fit(ratio, usable, table, delta_grid, wavelength_nm, degree)

d_norm

d_norm(ratio: Float[ArrayLike, ' B'], fit: TransmittanceFit) -> Float[Array, '']

Xiang et al.'s eq. 11: the fit's misfit relative to the ratio's spread.

With q = R / P(ĉ) the ratio with the continuum divided out,

D_norm = mean_λ |q(λ) − E(λ; ΔX̂)| / mean_λ |q(λ) − q̄|,

means over the bands where R is finite (q̄ is their mean). A perfect fit gives 0; with no gas (ΔX̂ = 0, E = 1) it is close to 1, and it can exceed 1. Differentiable in ratio (zero gradient in a NaN band); fit is held fixed.

Parameters:

Name	Type	Description	Default
ratio	Float[ArrayLike, ' B']	The ratio spectrum R the fit was made to, (B,); NaN in bands left out.	required
fit	TransmittanceFit	Its TransmittanceFit.	required

Returns:

Type	Description
Float[Array, '']	D_norm (dimensionless), (); NaN if no band is finite or
Float[Array, '']	q is constant (a noise-free fit at ΔX̂ = 0).

Raises:

Type	Description
ShapeError	If ratio and fit.continuum disagree in shape.

Examples:

>>> import jax.numpy as jnp
>>> from radtrax.retrieval.absorption import TransmittanceFit, d_norm
>>> e = jnp.array([1.0, 0.9, 0.8, 0.9])
>>> fit = TransmittanceFit(
...     jnp.array(0.1),
...     jnp.array([2.0]),
...     jnp.full(4, 2.0),
...     e,
...     jnp.array(0.0),
... )
>>> print(f"{d_norm(2.0 * e, fit):.3f}")
0.000
>>> print(f"{d_norm(jnp.array([2.0, 1.8, jnp.nan, 2.0]), fit):.3f}")
0.750

Source code in src/radtrax/retrieval/absorption/fit.py

def d_norm(ratio: Float[ArrayLike, " B"], fit: TransmittanceFit) -> Float[Array, ""]:
    """Xiang et al.'s eq. 11: the fit's misfit relative to the ratio's spread.

    With ``q = R / P(ĉ)`` the ratio with the continuum divided out,

    ``D_norm = mean_λ |q(λ) − E(λ; ΔX̂)| / mean_λ |q(λ) − q̄|``,

    means over the bands where ``R`` is finite (``q̄`` is their mean). A
    perfect fit gives 0; with no gas (``ΔX̂ = 0``, ``E = 1``) it is close to
    1, and it can exceed 1. Differentiable in ``ratio`` (zero gradient in a
    NaN band); ``fit`` is held fixed.

    Args:
        ratio: The ratio spectrum ``R`` the fit was made to, ``(B,)``; NaN in
            bands left out.
        fit: Its [`TransmittanceFit`][radtrax.retrieval.absorption.TransmittanceFit].

    Returns:
        ``D_norm`` (dimensionless), ``()``; NaN if no band is finite or
        ``q`` is constant (a noise-free fit at ``ΔX̂ = 0``).

    Raises:
        ShapeError: If ``ratio`` and ``fit.continuum`` disagree in shape.

    Examples:
        >>> import jax.numpy as jnp
        >>> from radtrax.retrieval.absorption import TransmittanceFit, d_norm
        >>> e = jnp.array([1.0, 0.9, 0.8, 0.9])
        >>> fit = TransmittanceFit(
        ...     jnp.array(0.1),
        ...     jnp.array([2.0]),
        ...     jnp.full(4, 2.0),
        ...     e,
        ...     jnp.array(0.0),
        ... )
        >>> print(f"{d_norm(2.0 * e, fit):.3f}")
        0.000
        >>> print(f"{d_norm(jnp.array([2.0, 1.8, jnp.nan, 2.0]), fit):.3f}")
        0.750
    """
    ratio = jnp.asarray(ratio)
    if ratio.shape != fit.continuum.shape:
        raise ShapeError(
            f"d_norm: ratio has shape {ratio.shape}, the fit's continuum "
            f"{fit.continuum.shape}"
        )
    usable = _usable(ratio)  # (B,)
    n = jnp.sum(usable)
    q = jnp.where(usable, ratio, 0.0) / fit.continuum  # (B,), double where
    q_mean = jnp.sum(jnp.where(usable, q, 0.0)) / n
    misfit = jnp.sum(jnp.where(usable, jnp.abs(q - fit.transmittance), 0.0)) / n
    spread = jnp.sum(jnp.where(usable, jnp.abs(q - q_mean), 0.0)) / n
    return misfit / spread