Skip to content

Matched filter

Hyperspectral matched-filter detection pipeline. RTM stack.

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

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:plumax.assimilation.RadianceObservationModel via :func:jax.jvp (directional derivatives) or a finite-amplitude secant. No duplicate Beer–Lambert code.
  • I/O (io) — xarray wrappers for multi-scene workflows.

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.

validate_mf_inputs 

validate_mf_inputs(cov_op: LinearOperator, target: Float[Array, 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/plumax/matched_filter/core.py 
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
def validate_mf_inputs(cov_op: LinearOperator, target: Float[Array, 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[Array, B], mean: Float[Array, B], cov_op: LinearOperator, target: Float[Array, 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/plumax/matched_filter/core.py 
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
def apply_pixel(
    pixel: Float[Array, B],
    mean: Float[Array, B],
    cov_op: LinearOperator,
    target: Float[Array, 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(pixel - mean, w) / target_norm_sq

apply_image 

apply_image(cube: Float[Array, 'H W B'], mean: Float[Array, B], cov_op: LinearOperator, target: Float[Array, 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/plumax/matched_filter/core.py 
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
def apply_image(
    cube: Float[Array, "H W B"],
    mean: Float[Array, B],
    cov_op: LinearOperator,
    target: Float[Array, 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 = 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[Array, 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/plumax/matched_filter/core.py 
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
def matched_filter_snr(
    amplitude: float,
    cov_op: LinearOperator,
    target: Float[Array, 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[Array, 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/plumax/matched_filter/core.py 
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
def detection_threshold(
    false_alarm_rate: float,
    cov_op: LinearOperator,
    target: Float[Array, 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 ValueError(
            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:plumax.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.

estimate_mean 

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

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'.

Returns

np.ndarray Mean spectrum of length n_bands.

Source code in src/plumax/matched_filter/background.py 
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
def estimate_mean(
    cube: Float[Array, "H W B"] | np.ndarray,
    *,
    method: _RobustMethod = "mean",
    trim_proportion: float = 0.1,
    huber_c: float = 1.345,
) -> np.ndarray:
    """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'``.

    Returns
    -------
    np.ndarray
        Mean spectrum of length ``n_bands``.
    """
    cube_np = np.asarray(cube)
    if cube_np.ndim != 3:
        raise ValueError(
            f"estimate_mean: cube must be (H, W, n_bands), got shape {cube_np.shape}."
        )
    X = cube_np.reshape(-1, cube_np.shape[-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 ValueError(f"estimate_mean: unknown method {method!r}.")

estimate_cov_empirical 

estimate_cov_empirical(cube: Float[Array, 'H W B'] | ndarray, mean: ndarray | None = None, *, ridge: float = 0.0, mask: ndarray | 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/plumax/matched_filter/background.py 
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
def estimate_cov_empirical(
    cube: Float[Array, "H W B"] | np.ndarray,
    mean: np.ndarray | None = None,
    *,
    ridge: float = 0.0,
    mask: np.ndarray | 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[Array, 'H W B'] | ndarray, mean: ndarray | None = None, *, method: Literal['ledoit_wolf', 'oas'] = 'ledoit_wolf', mask: ndarray | 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/plumax/matched_filter/background.py 
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
def estimate_cov_shrunk(
    cube: Float[Array, "H W B"] | np.ndarray,
    mean: np.ndarray | None = None,
    *,
    method: Literal["ledoit_wolf", "oas"] = "ledoit_wolf",
    mask: np.ndarray | 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 ValueError(
            f"estimate_cov_shrunk: method must be 'ledoit_wolf' or 'oas'; got {method!r}."
        )
    return _dense_operator(est.covariance_)

estimate_cov_lowrank 

estimate_cov_lowrank(cube: Float[Array, 'H W B'] | ndarray, mean: ndarray | None = None, *, rank: int, tikhonov: float, random_state: int | Array | None = 0, n_oversamples: int = 10, mask: ndarray | 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/plumax/matched_filter/background.py 
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
def estimate_cov_lowrank(
    cube: Float[Array, "H W B"] | np.ndarray,
    mean: np.ndarray | None = None,
    *,
    rank: int,
    tikhonov: float,
    random_state: int | jax.Array | None = 0,
    n_oversamples: int = 10,
    mask: np.ndarray | 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 ValueError("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)

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:plumax.assimilation.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.

linear_target_from_obs 

linear_target_from_obs(obs: RadianceObservationModel, vmr_background: Float[Array, 'H W'], *, pattern: Pattern | Float[Array, '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 Output pixel (i, j) at which to read the signature. Defaults to the scene centre, or the impulse location if applicable. 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/plumax/matched_filter/target.py 
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
def linear_target_from_obs(
    obs: RadianceObservationModel,
    vmr_background: Float[Array, "H W"],
    *,
    pattern: Pattern | Float[Array, "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
        Output pixel ``(i, j)`` at which to read the signature. Defaults to
        the scene centre, or the impulse location if applicable.
    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 _extract_pixel(dL, pattern, pixel)

nonlinear_target_from_obs 

nonlinear_target_from_obs(obs: RadianceObservationModel, vmr_background: Float[Array, 'H W'], *, amplitude: float = 1.0, pattern: Pattern | Float[Array, '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/plumax/matched_filter/target.py 
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
def nonlinear_target_from_obs(
    obs: RadianceObservationModel,
    vmr_background: Float[Array, "H W"],
    *,
    amplitude: float = 1.0,
    pattern: Pattern | Float[Array, "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 _extract_pixel(L_perturbed - L_base, pattern, pixel)

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 the labels and a list of (mean, cov_operator) pairs, one 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 use scikit-learn / scipy under the hood; neither is on a JAX gradient path (the MF treats the background as a fixed input).

ClusterBackground dataclass 

ClusterBackground(labels: ndarray, means: list[ndarray], cov_operators: list[LinearOperator])

Per-cluster background model.

Attributes

labels Cluster assignment per pixel, shape (H, W) with values in [0, n_clusters). means List of per-cluster mean spectra. cov_operators List 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: ndarray, *, 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

Source code in src/plumax/matched_filter/cluster.py 
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
def gmm_cluster_background(
    cube: np.ndarray,
    *,
    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
    """
    arr = np.asarray(cube)
    if arr.ndim != 3:
        raise ValueError(
            f"gmm_cluster_background: cube must be (H, W, n_bands); got {arr.shape}."
        )
    h, w, _ = arr.shape
    X = arr.reshape(-1, arr.shape[-1])
    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 = 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 ValueError(
            f"gmm_cluster_background: unknown cov_estimator {cov_estimator!r}."
        )
    means: list[np.ndarray] = []
    ops: list[LinearOperator] = []
    for k in range(int(labels.max()) + 1):
        mask = labels == 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)
    return ClusterBackground(
        labels=labels.reshape(h, w),
        means=means,
        cov_operators=ops,
    )

adaptive_window_background 

adaptive_window_background(cube: ndarray, *, window_size: int, guard_size: int = 0) -> tuple[np.ndarray, np.ndarray]

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 variance, shape (H, W, n_bands).

Source code in src/plumax/matched_filter/cluster.py 
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
def adaptive_window_background(
    cube: np.ndarray,
    *,
    window_size: int,
    guard_size: int = 0,
) -> tuple[np.ndarray, np.ndarray]:
    """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 variance, shape ``(H, W, n_bands)``.
    """
    if window_size < 3 or window_size % 2 == 0:
        raise ValueError(
            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 ValueError(
            f"adaptive_window_background: guard_size must be 0 or an odd int; got {guard_size}."
        )
    if guard_size >= window_size:
        raise ValueError(
            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 ValueError(
            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
    n_ann = n_out - n_in

    def window_mean(plane: np.ndarray) -> np.ndarray:
        # Uniform (box) mean over the window, excluding the central guard region
        # when guard_size > 0: annulus_sum / annulus_count via box-sum difference.
        out_sum = uniform_filter(plane, size=window_size, mode="reflect") * n_out
        if guard_size > 0:
            in_sum = uniform_filter(plane, size=guard_size, mode="reflect") * n_in
            return (out_sum - in_sum) / n_ann
        return out_sum / n_out

    # Apply a 2-D moving average independently to each band — scipy handles
    # the reflect-boundary by default, matching how most MF codes handle
    # scene edges.
    mean = np.empty_like(arr)
    var = np.empty_like(arr)
    for b in range(arr.shape[-1]):
        m = window_mean(arr[..., b])
        m2 = window_mean(arr[..., b] ** 2)
        mean[..., b] = m
        # Population variance — E[x²] - E[x]². Clipped to zero for numerics.
        var[..., b] = np.maximum(m2 - m * m, 0.0)
    return mean, var

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 - 1).

.. [#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.

WelfordAccumulator dataclass 

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

Online mean and covariance 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 and :meth:covariance to finalise.

Attributes

n_bands Number of spectral bands. count Running sample count.

update 

update(batch: ndarray) -> WelfordAccumulator

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

Source code in src/plumax/matched_filter/streaming.py 
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
def update(self, batch: np.ndarray) -> WelfordAccumulator:
    """Ingest a batch of pixel spectra (shape ``(n_samples, 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 ValueError(
            f"WelfordAccumulator.update: expected shape (*, {self.n_bands}); "
            f"got {arr.shape}."
        )
    n_b = arr.shape[0]
    if n_b == 0:
        return self
    mean_b = arr.mean(axis=0)
    diff = arr - mean_b
    M2_b = diff.T @ diff
    n_a = self.count
    n = n_a + n_b
    delta = mean_b - self._mean
    self._mean = self._mean + delta * (n_b / n)
    self._M2 = self._M2 + M2_b + np.outer(delta, delta) * (n_a * n_b / n)
    self.count = n
    return self

merge 

merge(other: WelfordAccumulator) -> WelfordAccumulator

Combine with another accumulator on disjoint data.

Source code in src/plumax/matched_filter/streaming.py 
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
def merge(self, other: WelfordAccumulator) -> WelfordAccumulator:
    """Combine with another accumulator on disjoint data."""
    if other.n_bands != self.n_bands:
        raise ValueError(
            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.count = other.count
        return self
    n_a, n_b = self.count, other.count
    n = n_a + n_b
    delta = other._mean - self._mean
    self._mean = self._mean + delta * (n_b / n)
    self._M2 = self._M2 + other._M2 + np.outer(delta, delta) * (n_a * n_b / n)
    self.count = n
    return self

streaming_background 

streaming_background(batches: Iterable[ndarray], *, n_bands: int, ridge: float = 0.0, ddof: int = 0) -> tuple[np.ndarray, 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:~plumax.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:~plumax.matched_filter.core.matched_filter_snr and :func:~plumax.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.

Source code in src/plumax/matched_filter/streaming.py 
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
def streaming_background(
    batches: Iterable[np.ndarray],
    *,
    n_bands: int,
    ridge: float = 0.0,
    ddof: int = 0,
) -> tuple[np.ndarray, 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:`~plumax.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:`~plumax.matched_filter.core.matched_filter_snr` and
        :func:`~plumax.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.
    """
    acc = WelfordAccumulator(n_bands=n_bands)
    for batch in batches:
        acc.update(batch)
    mu = acc.mean()
    cov = acc.covariance(ddof=ddof)
    if ridge > 0.0:
        cov = cov + ridge * np.eye(n_bands)
    import jax.numpy as jnp

    return mu, lx.MatrixLinearOperator(
        jnp.asarray(cov),
        tags=frozenset({lx.symmetric_tag, lx.positive_semidefinite_tag}),
    )

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:~plumax.matched_filter.streaming.WelfordAccumulator.

apply_image_xarray 

apply_image_xarray(cube: DataArray, mean: ndarray | Float[Array, B], cov_op: LinearOperator, target: ndarray | Float[Array, 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:plumax.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/plumax/matched_filter/io.py 
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
def apply_image_xarray(
    cube: xr.DataArray,
    mean: np.ndarray | Float[Array, B],
    cov_op: LinearOperator,
    target: np.ndarray | Float[Array, 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:`plumax.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 plumax.matched_filter.core import apply_image

    if band_dim not in cube.dims:
        raise ValueError(
            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[np.ndarray]

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:~plumax.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/plumax/matched_filter/io.py 
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
def open_multi_scene(
    paths: str | Iterable[str],
    *,
    band_dim: str = "band",
    variable: str | None = None,
    **open_kwargs,
) -> Iterator[np.ndarray]:
    """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:`~plumax.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 FileNotFoundError(
                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 ValueError(
                    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 ValueError(
                    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()