Radiative transfer

Band-integrated Beer–Lambert forward model, normalised-brightness LUT, and matched-filter retrieval. RTM stack.

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

radtran

Band-integrated radiative transfer + matched-filter retrieval.

Bridges the line-by-line :mod:plumax.hapi_lut absorption-LUT family to the multispectral / hyperspectral retrieval world by providing:

• :mod:.config — :class:ObservationGeometry, :class:InstrumentSpec.
• :mod:.srf — :class:SpectralResponseFunction (gaussian / rectangular / triangular / custom) for band integration.
• :mod:.forward — exact / Maclaurin / Taylor Beer-Lambert forward models and their normalised (differential) variants, each returning (radiance, Jacobian, transmittance).
• :mod:.nb_lut — pre-tabulated normalised-brightness LUT nB_b(ΔX) plus :func:inject_plume for scene-injection.
• :mod:.target — matched-filter target-spectrum constructors (nonlinear + linearised).
• :mod:.background — robust trimmed-mean background and low-rank covariance estimators (depends on :mod:scipy.stats).
• :mod:.matched_filter — the core ε̂ = (x−μ)ᵀΣ⁻¹t / (tᵀΣ⁻¹t) at pixel and image scope, plus an SNR helper.
• :mod:.gaussx_solve — structured-operator matched filter using :mod:gaussx's LowRankUpdate + Woodbury-backed solve, for hyperspectral retrievals where the O(n_bands³) dense inverse in :mod:.matched_filter becomes the bottleneck.

The input LUT is expected to be the :class:xarray.Dataset produced by :func:plumax.hapi_lut.build_lut_dataset.

Notes

radtran is an aggregator: it re-exports symbols from its submodules but imports them lazily via :pep:562. This means

• import plumax (which itself imports radtran) does not pull in :mod:scipy or :mod:gaussx until you actually touch :mod:.background / :mod:.gaussx_solve.
• from plumax.radtran import build_nb_lut triggers exactly one submodule import (radtran.nb_lut) and stops there.
• The __all__ list still drives star-imports and IDE auto-complete, so the lazy loading is invisible to consumers.

config

Observation-geometry and instrument configuration for radtran.

The forward models in :mod:plumax.radtran.forward need three groups of inputs: atmospheric state (T, P, background VMR), a gas-species LUT (σ(ν, T, P), from :mod:plumax.hapi_lut), and the viewing geometry and instrument response.

This module factors the latter two into frozen dataclasses:

• :class:ObservationGeometry — solar/viewing zenith angles, vertical path length, optional AMF override.
• :class:InstrumentSpec — band centres and widths for a synthetic multispectral instrument. Real instruments should build a :class:~plumax.radtran.srf.SpectralResponseFunction from measured SRFs; InstrumentSpec is a convenient constructor for the tutorial notebooks.

BOLTZMANN_J_PER_K and ATM_TO_PA mirror the constants in :mod:plumax.hapi_lut.beers so the two forward-model stacks agree on the ideal-gas number density computation.

ObservationGeometry dataclass

ObservationGeometry(sza_deg: float = 30.0, vza_deg: float = 0.0, path_length_cm: float = 840000.0, amf: float | None = None)

Viewing geometry for a single pixel.

Attributes

sza_deg : float Solar zenith angle [degrees], 0 (overhead sun) to ~90 (grazing). vza_deg : float Viewing zenith angle [degrees], 0 (nadir) to ~90 (limb). path_length_cm : float Vertical atmospheric path [cm] traversed by the line of sight before AMF correction. For a single-slab column over the full atmosphere, a common choice is L = p_surf / (ρ · g) ≈ 8.4 km = 8.4e5 cm (scale height of dry air at 1 atm, 288 K). amf : float, optional Air-mass factor. If None the plane-parallel two-way form 1/cos(SZA) + 1/cos(VZA) is used.

Notes

The plane-parallel AMF breaks down past ~75° zenith; beyond that a spherical correction (Chapman function) is needed. For the tutorial regime (SZA, VZA < 60°) the plane-parallel form is within a few percent of the exact value.

air_mass_factor property

air_mass_factor: float

Plane-parallel two-way AMF, or the explicit override if provided.

InstrumentSpec dataclass

InstrumentSpec(name: str = 'synthetic', band_centers_nm: ndarray = (lambda: np.array([1610.0, 2190.0], dtype=float))(), band_widths_nm: ndarray = (lambda: np.array([90.0, 180.0], dtype=float))(), band_names: tuple[str, ...] = ('B11', 'B12'), srf_type: str = 'gaussian')

Multispectral instrument description.

Attributes

name : str Free-text instrument identifier for dataset metadata. band_centers_nm : np.ndarray Band centre wavelengths [nm], shape (n_bands,). band_widths_nm : np.ndarray Full-width at half-maximum of each band [nm], shape (n_bands,). band_names : tuple of str Short names for each band (e.g. ('B11', 'B12')). Length matches band_centers_nm. srf_type : str Shape of the synthetic SRF built by :meth:InstrumentSpec.make_srf — one of 'gaussian', 'rectangular', 'triangular'.

Examples

Sentinel-2 B11/B12-like synthetic instrument::

InstrumentSpec(
    name='S2A-like',
    band_centers_nm=np.array([1610.0, 2190.0]),
    band_widths_nm=np.array([90.0, 180.0]),
    band_names=('B11', 'B12'),
)

n_bands property

n_bands: int

Number of bands.

make_srf

make_srf(wavelengths_nm: ndarray) -> SpectralResponseFunction

Construct a :class:SpectralResponseFunction over wavelengths_nm.

Imported locally to avoid a circular import at module load.

Source code in src/plumax/radtran/config.py

def make_srf(self, wavelengths_nm: np.ndarray) -> SpectralResponseFunction:
    """Construct a :class:`SpectralResponseFunction` over ``wavelengths_nm``.

    Imported locally to avoid a circular import at module load.
    """
    from plumax.radtran.srf import SpectralResponseFunction

    return SpectralResponseFunction(
        wavelengths_hr_nm=np.asarray(wavelengths_nm, dtype=float),
        band_centers_nm=np.asarray(self.band_centers_nm, dtype=float),
        band_widths_nm=np.asarray(self.band_widths_nm, dtype=float),
        band_names=tuple(self.band_names),
        srf_type=self.srf_type,
    )

number_density_cm3

number_density_cm3(p_atm: float, T_K: float) -> float

Ideal-gas total number density [molecules / cm³].

Source code in src/plumax/radtran/config.py

def number_density_cm3(p_atm: float, T_K: float) -> float:
    """Ideal-gas total number density [molecules / cm³]."""
    if not (T_K > 0.0 and p_atm > 0.0):
        raise ValueError(
            "number_density_cm3: `T_K` and `p_atm` must be > 0 "
            f"(got T_K={T_K!r}, p_atm={p_atm!r})"
        )
    p_pa = p_atm * ATM_TO_PA
    n_per_m3 = p_pa / (BOLTZMANN_J_PER_K * T_K)
    return float(n_per_m3 / CM3_PER_M3)

forward

Beer-Lambert forward models for atmospheric retrieval.

Three flavours, each returning radiance + Jacobian + transmittance:

• :func:forward_nonlinear — exact Beer-Lambert L = L₀ · exp(-τ).
• :func:forward_maclaurin — polynomial in VMR via Maclaurin of exp(-τ) around VMR = 0. Order 1 is the classical "linear" retrieval.
• :func:forward_taylor — Taylor-linearised around a background state VMR_bg (the state used in 3D-/4D-Var inner loops).

Plus normalised variants that divide by the background radiance:

• :func:forward_nonlinear_normalized — L_norm = exp(-Δτ); cancels surface reflectance, solar irradiance, and common aerosol slope.
• :func:forward_maclaurin_normalized / :func:forward_taylor_normalized.

All functions take an xarray.Dataset LUT carrying an absorption_cross_section variable with dims (wavenumber, temperature, pressure) — the output of :func:plumax.hapi_lut.build_lut_dataset.

Ported and adapted from jej_vc_snippets/methane_retrieval/lut_model_beers.py — the Jacobian is now always returned alongside the radiance so the same functions can drive variational retrievals (H = dL/dVMR) and matched-filter target-spectrum construction (t = H · ΔVMR).

ForwardResult dataclass

ForwardResult(radiance: ndarray, jacobian: ndarray, transmittance: ndarray)

Bundle of (radiance, Jacobian, transmittance) returned by the forward models.

Using a frozen dataclass rather than a plain tuple gives the caller named access — result.jacobian — which matters because several functions in :mod:plumax.radtran.target and :mod:plumax.radtran.matched_filter consume only the Jacobian or only the transmittance.

Attributes

radiance : np.ndarray Simulated radiance at each wavenumber, shape (n_nu,). jacobian : np.ndarray dL/dVMR evaluated at the supplied state, shape (n_nu,). transmittance : np.ndarray Atmospheric transmittance exp(-τ) (non-normalised variants) or exp(-Δτ) (normalised variants), shape (n_nu,).

forward_nonlinear

forward_nonlinear(ds: Dataset, nu_obs: ndarray, *, T_K: float, p_atm: float, vmr: float, path_length_cm: float, amf: float, surface_reflectance: float = 1.0, solar_irradiance: float = 1.0, var: str = 'absorption_cross_section') -> ForwardResult

Exact Beer-Lambert forward model.

Radiance: L(ν) = (F₀ R / π) · exp(-τ(ν, VMR)) Optical τ: τ = σ · N_total · VMR · L · AMF Jacobian: dL/dVMR = -L · (dτ/dVMR), dτ/dVMR = σ · N_total · L · AMF

Defaults surface_reflectance = solar_irradiance = 1 collapse the prefactor to 1/π, which is convenient for toy tests where the absolute radiance scale does not matter.

Source code in src/plumax/radtran/forward.py

def forward_nonlinear(
    ds: xr.Dataset,
    nu_obs: np.ndarray,
    *,
    T_K: float,
    p_atm: float,
    vmr: float,
    path_length_cm: float,
    amf: float,
    surface_reflectance: float = 1.0,
    solar_irradiance: float = 1.0,
    var: str = "absorption_cross_section",
) -> ForwardResult:
    """Exact Beer-Lambert forward model.

    Radiance:    ``L(ν) = (F₀ R / π) · exp(-τ(ν, VMR))``
    Optical τ:   ``τ = σ · N_total · VMR · L · AMF``
    Jacobian:    ``dL/dVMR = -L · (dτ/dVMR)``, ``dτ/dVMR = σ · N_total · L · AMF``

    Defaults ``surface_reflectance = solar_irradiance = 1`` collapse the
    prefactor to ``1/π``, which is convenient for toy tests where the
    absolute radiance scale does not matter.
    """
    sigma = _interp_sigma(ds, nu_obs, T_K, p_atm, var)
    N_total = number_density_cm3(p_atm, T_K)
    tau = sigma * N_total * vmr * path_length_cm * amf
    transmittance = np.exp(-tau)
    L0 = solar_irradiance * surface_reflectance / np.pi
    radiance = L0 * transmittance
    dtau_dvmr = sigma * N_total * path_length_cm * amf
    jacobian = -radiance * dtau_dvmr
    return ForwardResult(
        radiance=radiance, jacobian=jacobian, transmittance=transmittance
    )

forward_nonlinear_normalized

forward_nonlinear_normalized(ds: Dataset, nu_obs: ndarray, *, T_K: float, p_atm: float, vmr_background: float, delta_vmr: float, path_length_cm: float, amf: float, var: str = 'absorption_cross_section') -> ForwardResult

Exact Beer-Lambert normalised by the background radiance.

L_norm = L(VMR_bg + ΔVMR) / L(VMR_bg) = exp(-Δτ), which cancels any multiplicative surface/solar/aerosol factors that appear equally in the plume and background pixels.

Returns

ForwardResult radiance is the normalised transmittance exp(-Δτ); transmittance is the same quantity (kept for API symmetry with :func:forward_nonlinear); jacobian is d(L_norm)/d(ΔVMR) = -exp(-Δτ) · (σ · N · L · AMF).

Source code in src/plumax/radtran/forward.py

def forward_nonlinear_normalized(
    ds: xr.Dataset,
    nu_obs: np.ndarray,
    *,
    T_K: float,
    p_atm: float,
    vmr_background: float,
    delta_vmr: float,
    path_length_cm: float,
    amf: float,
    var: str = "absorption_cross_section",
) -> ForwardResult:
    """Exact Beer-Lambert *normalised* by the background radiance.

    ``L_norm = L(VMR_bg + ΔVMR) / L(VMR_bg) = exp(-Δτ)``, which cancels any
    multiplicative surface/solar/aerosol factors that appear equally in the
    plume and background pixels.

    Returns
    -------
    ForwardResult
        ``radiance`` is the normalised transmittance ``exp(-Δτ)``;
        ``transmittance`` is the same quantity (kept for API symmetry with
        :func:`forward_nonlinear`); ``jacobian`` is
        ``d(L_norm)/d(ΔVMR) = -exp(-Δτ) · (σ · N · L · AMF)``.
    """
    sigma = _interp_sigma(ds, nu_obs, T_K, p_atm, var)
    N_total = number_density_cm3(p_atm, T_K)
    dtau_d_dvmr = sigma * N_total * path_length_cm * amf
    delta_tau = dtau_d_dvmr * delta_vmr
    L_norm = np.exp(-delta_tau)
    jacobian = -L_norm * dtau_d_dvmr
    return ForwardResult(radiance=L_norm, jacobian=jacobian, transmittance=L_norm)

forward_maclaurin

forward_maclaurin(ds: Dataset, nu_obs: ndarray, *, T_K: float, p_atm: float, vmr: float, path_length_cm: float, amf: float, surface_reflectance: float = 1.0, solar_irradiance: float = 1.0, order: int = 1, var: str = 'absorption_cross_section') -> ForwardResult

Maclaurin-series forward model: expand exp(-τ(VMR)) around VMR = 0.

With a = σ · N · L · AMF so τ = a · VMR:

• order 1: T ≈ 1 − a·VMR (linear in VMR).
• order 2: T ≈ 1 − a·VMR + ½ (a·VMR)².
• order 3: T ≈ 1 − a·VMR + ½ (a·VMR)² − ⅙ (a·VMR)³.

Accurate when the total optical depth a·VMR is ≪ 1 — the regime of classical linear retrievals for thin absorbers.

Source code in src/plumax/radtran/forward.py

def forward_maclaurin(
    ds: xr.Dataset,
    nu_obs: np.ndarray,
    *,
    T_K: float,
    p_atm: float,
    vmr: float,
    path_length_cm: float,
    amf: float,
    surface_reflectance: float = 1.0,
    solar_irradiance: float = 1.0,
    order: int = 1,
    var: str = "absorption_cross_section",
) -> ForwardResult:
    """Maclaurin-series forward model: expand ``exp(-τ(VMR))`` around VMR = 0.

    With ``a = σ · N · L · AMF`` so ``τ = a · VMR``:

    - order 1: ``T ≈ 1 − a·VMR`` (linear in VMR).
    - order 2: ``T ≈ 1 − a·VMR + ½ (a·VMR)²``.
    - order 3: ``T ≈ 1 − a·VMR + ½ (a·VMR)² − ⅙ (a·VMR)³``.

    Accurate when the total optical depth ``a·VMR`` is ≪ 1 — the regime of
    classical linear retrievals for thin absorbers.
    """
    if order not in (1, 2, 3):
        raise ValueError(
            f"forward_maclaurin: `order` must be 1, 2 or 3 (got {order!r})"
        )
    sigma = _interp_sigma(ds, nu_obs, T_K, p_atm, var)
    N_total = number_density_cm3(p_atm, T_K)
    a = sigma * N_total * path_length_cm * amf
    a_vmr = a * vmr

    if order == 1:
        transmittance = 1.0 - a_vmr
        dtrans_dvmr = -a
    elif order == 2:
        transmittance = 1.0 - a_vmr + 0.5 * a_vmr**2
        dtrans_dvmr = -a + (a**2) * vmr
    else:  # order == 3
        transmittance = 1.0 - a_vmr + 0.5 * a_vmr**2 - (1.0 / 6.0) * a_vmr**3
        dtrans_dvmr = -a + (a**2) * vmr - 0.5 * (a**3) * vmr**2

    L0 = solar_irradiance * surface_reflectance / np.pi
    radiance = L0 * transmittance
    jacobian = L0 * dtrans_dvmr
    return ForwardResult(
        radiance=radiance, jacobian=jacobian, transmittance=transmittance
    )

forward_maclaurin_normalized

forward_maclaurin_normalized(ds: Dataset, nu_obs: ndarray, *, T_K: float, p_atm: float, delta_vmr: float, path_length_cm: float, amf: float, order: int = 1, var: str = 'absorption_cross_section') -> ForwardResult

Maclaurin expansion of exp(-Δτ) around ΔVMR = 0.

At order 1 this reduces to L_norm ≈ 1 − σ · N · ΔVMR · L · AMF — the canonical linearised retrieval signal.

Source code in src/plumax/radtran/forward.py

def forward_maclaurin_normalized(
    ds: xr.Dataset,
    nu_obs: np.ndarray,
    *,
    T_K: float,
    p_atm: float,
    delta_vmr: float,
    path_length_cm: float,
    amf: float,
    order: int = 1,
    var: str = "absorption_cross_section",
) -> ForwardResult:
    """Maclaurin expansion of ``exp(-Δτ)`` around ``ΔVMR = 0``.

    At order 1 this reduces to ``L_norm ≈ 1 − σ · N · ΔVMR · L · AMF`` — the
    canonical linearised retrieval signal.
    """
    if order not in (1, 2, 3):
        raise ValueError("forward_maclaurin_normalized: `order` must be 1, 2 or 3.")
    sigma = _interp_sigma(ds, nu_obs, T_K, p_atm, var)
    N_total = number_density_cm3(p_atm, T_K)
    a = sigma * N_total * path_length_cm * amf
    a_dvmr = a * delta_vmr
    if order == 1:
        L_norm = 1.0 - a_dvmr
        jac = -a
    elif order == 2:
        L_norm = 1.0 - a_dvmr + 0.5 * a_dvmr**2
        jac = -a + (a**2) * delta_vmr
    else:
        L_norm = 1.0 - a_dvmr + 0.5 * a_dvmr**2 - (1.0 / 6.0) * a_dvmr**3
        jac = -a + (a**2) * delta_vmr - 0.5 * (a**3) * delta_vmr**2
    return ForwardResult(radiance=L_norm, jacobian=jac, transmittance=L_norm)

forward_taylor

forward_taylor(ds: Dataset, nu_obs: ndarray, *, T_K: float, p_atm: float, vmr: float, vmr_background: float, path_length_cm: float, amf: float, surface_reflectance: float = 1.0, solar_irradiance: float = 1.0, var: str = 'absorption_cross_section') -> ForwardResult

First-order Taylor expansion around vmr_background.

L(VMR) ≈ L_b + H · (VMR − VMR_b) with H = dL/dVMR |_{VMR_b}. This is the linearisation used in the inner loops of 3D-/4D-Var.

Source code in src/plumax/radtran/forward.py

def forward_taylor(
    ds: xr.Dataset,
    nu_obs: np.ndarray,
    *,
    T_K: float,
    p_atm: float,
    vmr: float,
    vmr_background: float,
    path_length_cm: float,
    amf: float,
    surface_reflectance: float = 1.0,
    solar_irradiance: float = 1.0,
    var: str = "absorption_cross_section",
) -> ForwardResult:
    """First-order Taylor expansion around ``vmr_background``.

    ``L(VMR) ≈ L_b + H · (VMR − VMR_b)`` with ``H = dL/dVMR |_{VMR_b}``. This
    is the linearisation used in the inner loops of 3D-/4D-Var.
    """
    sigma = _interp_sigma(ds, nu_obs, T_K, p_atm, var)
    N_total = number_density_cm3(p_atm, T_K)
    dtau_dvmr = sigma * N_total * path_length_cm * amf
    tau_b = dtau_dvmr * vmr_background
    L0 = solar_irradiance * surface_reflectance / np.pi
    L_b = L0 * np.exp(-tau_b)
    H = -L_b * dtau_dvmr
    radiance = L_b + H * (vmr - vmr_background)
    transmittance = np.exp(-tau_b) + (-dtau_dvmr * np.exp(-tau_b)) * (
        vmr - vmr_background
    )
    return ForwardResult(radiance=radiance, jacobian=H, transmittance=transmittance)

forward_taylor_normalized

forward_taylor_normalized(ds: Dataset, nu_obs: ndarray, *, T_K: float, p_atm: float, delta_vmr: float, path_length_cm: float, amf: float, var: str = 'absorption_cross_section') -> ForwardResult

Taylor expansion of L_norm around ΔVMR = 0.

Since L_norm(0) = 1 and dL_norm/dΔVMR(0) = -σ · N · L · AMF, the first-order Taylor expansion coincides with the order-1 Maclaurin expansion. Kept as a named function for API symmetry with :func:forward_taylor.

Source code in src/plumax/radtran/forward.py

def forward_taylor_normalized(
    ds: xr.Dataset,
    nu_obs: np.ndarray,
    *,
    T_K: float,
    p_atm: float,
    delta_vmr: float,
    path_length_cm: float,
    amf: float,
    var: str = "absorption_cross_section",
) -> ForwardResult:
    """Taylor expansion of ``L_norm`` around ``ΔVMR = 0``.

    Since ``L_norm(0) = 1`` and ``dL_norm/dΔVMR(0) = -σ · N · L · AMF``,
    the first-order Taylor expansion coincides with the order-1 Maclaurin
    expansion. Kept as a named function for API symmetry with
    :func:`forward_taylor`.
    """
    return forward_maclaurin_normalized(
        ds,
        nu_obs,
        T_K=T_K,
        p_atm=p_atm,
        delta_vmr=delta_vmr,
        path_length_cm=path_length_cm,
        amf=amf,
        order=1,
        var=var,
    )

nb_lut

Normalised-brightness LUT for fast per-pixel plume injection.

The band-integrated normalised brightness for a pixel with methane column enhancement ΔX is

nB_b(ΔX) = ∫ f_b(λ) · exp(-AMF · σ(λ) · ΔX) dλ  /  ∫ f_b(λ) dλ

where f_b is the instrument SRF for band b and σ is the target gas cross-section in cm²/molecule. This integral is the natural scalar "plume throughput" for band b once surface reflectance and solar irradiance have been divided out (differential Beer-Lambert).

Evaluating the integral per pixel is expensive, so we follow the Eucalyptus radtran() trick: pre-tabulate nB over a 1-D grid of ΔX values, then look up per-pixel via piecewise-linear interpolation (np.interp). A 4-k-point LUT covering [0, 200] mol/m² costs ~1 ms to build and makes plume injection a single np.interp call per band.

Public surface

• :func:build_nb_lut — build the (n_bands, n_delta) LUT.
• :func:lookup_nb — look up nB per band for a 2-D plume map.
• :func:inject_plume — apply nB to a clean multispectral scene.

NBLookup dataclass

NBLookup(delta_column: ndarray, nB: ndarray, band_names: tuple[str, ...])

Pre-computed normalised-brightness table, indexed by band and ΔX.

Attributes

delta_column : np.ndarray Column-enhancement grid [mol/m²], shape (n_delta,). Strictly increasing, starts at 0. nB : np.ndarray Normalised brightness per band at each grid point, shape (n_bands, n_delta). nB[:, 0] == 1 always. band_names : tuple of str Band labels, length n_bands.

to_dataset

to_dataset() -> xr.Dataset

Wrap as an :class:xarray.Dataset for easy serialisation.

Source code in src/plumax/radtran/nb_lut.py

def to_dataset(self) -> xr.Dataset:
    """Wrap as an :class:`xarray.Dataset` for easy serialisation."""
    return xr.Dataset(
        data_vars={
            "nB": (["band", "delta_column"], self.nB),
        },
        coords={
            "band": list(self.band_names),
            "delta_column": (
                "delta_column",
                self.delta_column,
                {"units": "mol/m^2"},
            ),
        },
        attrs={
            "long_name": "Band-integrated normalised brightness",
            "description": (
                "nB_b(ΔX) = ∫ f_b(λ) exp(-AMF σ(λ) ΔX) dλ / ∫ f_b(λ) dλ"
            ),
        },
    )

build_nb_lut

build_nb_lut(ds: Dataset, srf: SpectralResponseFunction, *, T_K: float, p_atm: float, amf: float, max_delta_column: float = 200.0, n_grid: int = 4001, var: str = 'absorption_cross_section') -> NBLookup

Pre-tabulate band-integrated nB over a column-enhancement grid.

Parameters

ds : xr.Dataset HAPI absorption-cross-section LUT, output of :func:plumax.hapi_lut.build_lut_dataset. srf : SpectralResponseFunction Instrument SRF mapping the LUT wavelength grid to bands. T_K, p_atm : float Slab atmospheric temperature [K] and pressure [atm], at which the cross-section σ(λ; T, P) is interpolated. ΔX is a column already, so the slab path length does not enter the integrand. amf : float Two-way air-mass factor multiplying the slant column. max_delta_column : float Upper bound of the ΔX grid [mol/m²]. Default 200. n_grid : int Number of grid points (default 4001 gives 0.05 mol/m² resolution over [0, 200] — more than enough for 1-ppm-level retrievals). var : str Cross-section variable name (default 'absorption_cross_section', matching :func:plumax.hapi_lut.build_lut_dataset).

Returns

NBLookup nB[b, k] is the band-integrated normalised brightness at delta_column[k] for band b.

Source code in src/plumax/radtran/nb_lut.py

def build_nb_lut(
    ds: xr.Dataset,
    srf: SpectralResponseFunction,
    *,
    T_K: float,
    p_atm: float,
    amf: float,
    max_delta_column: float = 200.0,
    n_grid: int = 4001,
    var: str = "absorption_cross_section",
) -> NBLookup:
    """Pre-tabulate band-integrated nB over a column-enhancement grid.

    Parameters
    ----------
    ds : xr.Dataset
        HAPI absorption-cross-section LUT, output of
        :func:`plumax.hapi_lut.build_lut_dataset`.
    srf : SpectralResponseFunction
        Instrument SRF mapping the LUT wavelength grid to bands.
    T_K, p_atm : float
        Slab atmospheric temperature [K] and pressure [atm], at which the
        cross-section ``σ(λ; T, P)`` is interpolated. ΔX is a column
        already, so the slab path length does not enter the integrand.
    amf : float
        Two-way air-mass factor multiplying the slant column.
    max_delta_column : float
        Upper bound of the ΔX grid [mol/m²]. Default 200.
    n_grid : int
        Number of grid points (default 4001 gives 0.05 mol/m² resolution
        over [0, 200] — more than enough for 1-ppm-level retrievals).
    var : str
        Cross-section variable name (default ``'absorption_cross_section'``,
        matching :func:`plumax.hapi_lut.build_lut_dataset`).

    Returns
    -------
    NBLookup
        ``nB[b, k]`` is the band-integrated normalised brightness at
        ``delta_column[k]`` for band ``b``.
    """
    if var not in ds:
        raise KeyError(f"build_nb_lut: {var!r} not in dataset.")
    if n_grid < 2:
        raise ValueError(f"build_nb_lut: `n_grid` must be ≥ 2 (got {n_grid!r})")
    if not (max_delta_column > 0.0):
        raise ValueError(
            f"build_nb_lut: `max_delta_column` must be > 0 (got {max_delta_column!r})"
        )

    # Cross-section on the LUT wavenumber grid, at (T, P).
    sigma_nu = ds[var].interp(temperature=T_K, pressure=p_atm, method="linear")

    # Convert to wavelength (nm) to match the SRF. Drop the unused coord to
    # avoid `interp` warning about non-dim-coord changes.
    sigma_lambda = sigma_nu.assign_coords(
        lam_nm=("wavenumber", np.asarray(1e7 / ds["wavenumber"].values, dtype=float))
    )
    # Interpolate σ onto the SRF's HR wavelength grid (decreasing in lam is fine
    # for np.interp after flipping).
    lam_src = sigma_lambda["lam_nm"].values
    sort = np.argsort(lam_src)
    # Flag bands whose SRF response reaches past the LUT wavelength coverage;
    # np.interp zero-fills σ there, producing a silently-attenuated nB (≡ 1).
    check_srf_lut_coverage(srf, lam_src, context="build_nb_lut")
    sigma_hr = np.interp(
        srf.wavelengths_hr_nm,
        lam_src[sort],
        sigma_lambda.values[sort],
        left=0.0,
        right=0.0,
    )

    # Convert ΔX [mol/m²] to the per-wavelength optical depth prefactor.
    #   τ_λ(ΔX) = σ_λ [cm²/molec] · ΔX [mol/m²] · N_A [molec/mol] · 1e-4 [m²/cm²] · AMF
    # No path_length_cm appears here — ΔX already integrates the column.
    N_A = 6.02214076e23
    optical_coef = sigma_hr * N_A * 1e-4 * amf  # (n_lambda,)

    delta_column = np.linspace(0.0, float(max_delta_column), int(n_grid))
    # Per-(band, ΔX) brightness: integrate exp(-optical_coef · ΔX) over the SRF.
    tau_grid = optical_coef[None, :] * delta_column[:, None]  # (n_delta, n_lambda)
    transmittance = np.exp(-tau_grid)  # (n_delta, n_lambda)
    nB = srf.apply(transmittance).T  # (n_bands, n_delta)
    return NBLookup(delta_column=delta_column, nB=nB, band_names=srf.band_names)

lookup_nb

lookup_nb(plume_column: ndarray, nb_lookup: NBLookup) -> np.ndarray

Look up per-band nB for every pixel of a 2-D column map.

Values outside the LUT range are clipped. Clipping at the top end is physically conservative (saturates the absorption); negative inputs are clipped to zero since they are not physical.

Parameters

plume_column : np.ndarray Column enhancement [mol/m²], shape (ny, nx) or any broadcastable shape. NaNs are preserved as NaNs in the output. nb_lookup : NBLookup Pre-built LUT.

Returns

nb_map : np.ndarray Shape (n_bands, *plume_column.shape). Dimensionless multiplier in [0, 1] (1 at ΔX=0).

Source code in src/plumax/radtran/nb_lut.py

def lookup_nb(
    plume_column: np.ndarray,
    nb_lookup: NBLookup,
) -> np.ndarray:
    """Look up per-band nB for every pixel of a 2-D column map.

    Values outside the LUT range are clipped. Clipping at the top end is
    physically conservative (saturates the absorption); negative inputs are
    clipped to zero since they are not physical.

    Parameters
    ----------
    plume_column : np.ndarray
        Column enhancement [mol/m²], shape ``(ny, nx)`` or any broadcastable
        shape. NaNs are preserved as NaNs in the output.
    nb_lookup : NBLookup
        Pre-built LUT.

    Returns
    -------
    nb_map : np.ndarray
        Shape ``(n_bands, *plume_column.shape)``. Dimensionless multiplier
        in ``[0, 1]`` (1 at ΔX=0).
    """
    x = np.asarray(plume_column, dtype=float)
    grid = nb_lookup.delta_column
    nan_mask = np.isnan(x)
    x_clipped = np.clip(np.where(nan_mask, 0.0, x), grid[0], grid[-1])

    # Piecewise-linear interp per band (avoids nearest-neighbour stair-step
    # artefacts that the Eucalyptus `lookup_nB` shows on smooth plumes).
    out_shape = (nb_lookup.n_bands, *x.shape)
    out = np.empty(out_shape, dtype=float)
    for b in range(nb_lookup.n_bands):
        out[b] = np.interp(x_clipped, grid, nb_lookup.nB[b])
    if nan_mask.any():
        out[:, nan_mask] = np.nan
    return out

inject_plume

inject_plume(scene: ndarray, plume_column: ndarray, nb_lookup: NBLookup, *, band_axis: int = 0) -> np.ndarray

Multiplicative plume injection: scene_plume = scene · nB(ΔX).

Parameters

scene : np.ndarray Clean multispectral scene, shape (n_bands, ny, nx) by default (band_axis=0). plume_column : np.ndarray ΔX column map [mol/m²], shape (ny, nx). nb_lookup : NBLookup Pre-built LUT. nb_lookup.n_bands must match scene.shape[band_axis]. band_axis : int Axis of scene corresponding to bands. Default 0.

Returns

scene_plume : np.ndarray Same shape as scene.

Source code in src/plumax/radtran/nb_lut.py

def inject_plume(
    scene: np.ndarray,
    plume_column: np.ndarray,
    nb_lookup: NBLookup,
    *,
    band_axis: int = 0,
) -> np.ndarray:
    """Multiplicative plume injection: ``scene_plume = scene · nB(ΔX)``.

    Parameters
    ----------
    scene : np.ndarray
        Clean multispectral scene, shape ``(n_bands, ny, nx)`` by default
        (``band_axis=0``).
    plume_column : np.ndarray
        ΔX column map [mol/m²], shape ``(ny, nx)``.
    nb_lookup : NBLookup
        Pre-built LUT. ``nb_lookup.n_bands`` must match ``scene.shape[band_axis]``.
    band_axis : int
        Axis of ``scene`` corresponding to bands. Default 0.

    Returns
    -------
    scene_plume : np.ndarray
        Same shape as ``scene``.
    """
    scene_arr = np.asarray(scene)
    if scene_arr.shape[band_axis] != nb_lookup.n_bands:
        raise ValueError(
            f"inject_plume: scene has {scene_arr.shape[band_axis]} bands along "
            f"axis {band_axis}, but LUT has {nb_lookup.n_bands}."
        )
    nb_map = lookup_nb(plume_column, nb_lookup)  # (n_bands, ny, nx)
    # Move `scene`'s band axis to the front for broadcasting, then move back.
    scene_bf = np.moveaxis(scene_arr, band_axis, 0)
    if scene_bf.shape[1:] != nb_map.shape[1:]:
        raise ValueError(
            f"inject_plume: spatial shape mismatch — scene {scene_bf.shape[1:]} "
            f"vs. plume {nb_map.shape[1:]}"
        )
    return np.moveaxis(scene_bf * nb_map, 0, band_axis)

srf

Spectral Response Function — band integration of hyperspectral radiance.

The HAPI LUTs in :mod:plumax.hapi_lut tabulate σ(ν, T, P) on a fine line-by-line wavenumber grid. Real instruments observe band-integrated radiance

L_b = ∫ f_b(λ) · L(λ) dλ          / ∫ f_b(λ) dλ

where f_b(λ) is the instrument's Spectral Response Function for band b. This module builds the (n_bands, n_λ) SRF matrix once — each row is a normalised response profile — and applies it as a plain matrix multiplication against a high-resolution radiance spectrum (1-D) or radiance cube (3-D).

The class mirrors the API in the Orbio Project Eucalyptus tutorial (SpectralFilter + get_sensor_spectral_response_function) but exposes a linear-operator view: apply (forward), jacobian (tangent action — identical to forward since SRF is linear), and adjoint (Sᵀ for reverse-mode / VJP use in 3DVar). These three methods match the operator interface used by the observation-operator classes in the original jej_vc_snippets/methane_retrieval/lut_obs_op.py.

Three parametric SRF shapes are supported:

• 'gaussian' : exp[-(λ-λ_c)² / (2 σ²)] with FWHM = 2.355 σ.
• 'rectangular' : 1 inside [λ_c − W/2, λ_c + W/2], 0 outside.
• 'triangular' : linear ramp peaking at λ_c, zero at the edges.

Plus 'custom' which accepts an arbitrary (n_bands, n_λ) matrix (each row gets normalised to sum to 1).

SpectralResponseFunction dataclass

SpectralResponseFunction(wavelengths_hr_nm: ndarray, band_centers_nm: ndarray, band_widths_nm: ndarray, band_names: tuple[str, ...] = ('B0',), srf_type: SRFType = 'gaussian', custom_srfs: ndarray | None = None)

Band-integration operator for a multispectral / hyperspectral instrument.

Attributes

wavelengths_hr_nm : np.ndarray High-resolution wavelength grid [nm], shape (n_lambda,). Must be monotone increasing. band_centers_nm : np.ndarray Band centre wavelengths [nm], shape (n_bands,). band_widths_nm : np.ndarray Band FWHM (for 'gaussian'), full-width (for 'rectangular'/'triangular') [nm], shape (n_bands,). band_names : tuple of str Band labels, length n_bands. srf_type : str 'gaussian' | 'rectangular' | 'triangular' | 'custom'. custom_srfs : np.ndarray or None Used only when srf_type == 'custom'. Shape (n_bands, n_lambda).

Notes

Rows of :attr:matrix are L1-normalised so that a flat input spectrum L(λ) = c maps to L_b = c for every band; this matches the "dot-product-in-frequency" convention in the Eucalyptus tutorial.

n_bands property

n_bands: int

Number of bands.

n_lambda property

n_lambda: int

Number of high-resolution wavelength samples.

apply

apply(radiance_hr: ndarray) -> np.ndarray

Forward operator: map HR spectrum/cube to band-integrated values.

Parameters

radiance_hr : np.ndarray Input radiance. Shape (n_lambda,) for a single spectrum, or (..., n_lambda) for a cube (spatial + spectral).

Returns

radiance_bands : np.ndarray Band-integrated radiance with the final axis replaced by n_bands.

Source code in src/plumax/radtran/srf.py

def apply(self, radiance_hr: np.ndarray) -> np.ndarray:
    """Forward operator: map HR spectrum/cube to band-integrated values.

    Parameters
    ----------
    radiance_hr : np.ndarray
        Input radiance. Shape ``(n_lambda,)`` for a single spectrum, or
        ``(..., n_lambda)`` for a cube (spatial + spectral).

    Returns
    -------
    radiance_bands : np.ndarray
        Band-integrated radiance with the final axis replaced by
        ``n_bands``.
    """
    x = np.asarray(radiance_hr, dtype=float)
    if x.shape[-1] != self.n_lambda:
        raise ValueError(
            f"SpectralResponseFunction.apply: last axis size {x.shape[-1]} "
            f"must equal n_lambda={self.n_lambda}"
        )
    # ``S @ xᵀ``: contract last axis of `x` (wavelength) with second axis
    # of `matrix`, so `...l, bl -> ...b`.
    return np.einsum("bl, ...l -> ...b", self.matrix, x)

jacobian

jacobian(tangent_hr: ndarray) -> np.ndarray

Jacobian-vector product S · v. Linear operator → same as apply.

Source code in src/plumax/radtran/srf.py

def jacobian(self, tangent_hr: np.ndarray) -> np.ndarray:
    """Jacobian-vector product ``S · v``. Linear operator → same as ``apply``."""
    return self.apply(tangent_hr)

adjoint

adjoint(cotangent_bands: ndarray) -> np.ndarray

Adjoint operator Sᵀ: map band-space cotangents back to HR.

Used for reverse-mode / VJP gradient calculations.

Parameters

cotangent_bands : np.ndarray Band-space cotangent, shape (..., n_bands).

Returns

cotangent_hr : np.ndarray HR-space cotangent, shape (..., n_lambda).

Source code in src/plumax/radtran/srf.py

def adjoint(self, cotangent_bands: np.ndarray) -> np.ndarray:
    """Adjoint operator ``Sᵀ``: map band-space cotangents back to HR.

    Used for reverse-mode / VJP gradient calculations.

    Parameters
    ----------
    cotangent_bands : np.ndarray
        Band-space cotangent, shape ``(..., n_bands)``.

    Returns
    -------
    cotangent_hr : np.ndarray
        HR-space cotangent, shape ``(..., n_lambda)``.
    """
    u = np.asarray(cotangent_bands, dtype=float)
    if u.shape[-1] != self.n_bands:
        raise ValueError(
            f"SpectralResponseFunction.adjoint: last axis size {u.shape[-1]} "
            f"must equal n_bands={self.n_bands}"
        )
    return np.einsum("bl, ...b -> ...l", self.matrix, u)

check_srf_lut_coverage

check_srf_lut_coverage(srf: SpectralResponseFunction, lut_wavelengths_nm: ndarray, *, context: str, tol: float = 0.01) -> np.ndarray

Warn when SRF bands respond outside the LUT's wavelength coverage.

Resampling a cross-section / target spectrum onto the SRF grid zero-fills wavelengths beyond the LUT range (np.interp(..., left=0, right=0)), silently attenuating any band whose response reaches there — a band whose support lies wholly outside gets an exactly-zero (and misleading) value. This computes each band's fraction of SRF mass outside the LUT coverage and warns for bands exceeding tol.

Parameters

srf Instrument SRF. lut_wavelengths_nm Wavelengths [nm] covered by the LUT (any order); coverage is taken as [min, max]. context Caller name, prepended to the warning message. tol Out-of-coverage SRF-mass fraction above which a band is flagged.

Returns

frac_out : np.ndarray Per-band fraction of SRF response outside the LUT coverage.

Source code in src/plumax/radtran/srf.py

def check_srf_lut_coverage(
    srf: SpectralResponseFunction,
    lut_wavelengths_nm: np.ndarray,
    *,
    context: str,
    tol: float = 1e-2,
) -> np.ndarray:
    """Warn when SRF bands respond outside the LUT's wavelength coverage.

    Resampling a cross-section / target spectrum onto the SRF grid zero-fills
    wavelengths beyond the LUT range (``np.interp(..., left=0, right=0)``),
    silently attenuating any band whose response reaches there — a band whose
    support lies wholly outside gets an exactly-zero (and misleading) value.
    This computes each band's fraction of SRF mass *outside* the LUT coverage
    and warns for bands exceeding ``tol``.

    Parameters
    ----------
    srf
        Instrument SRF.
    lut_wavelengths_nm
        Wavelengths [nm] covered by the LUT (any order); coverage is taken as
        ``[min, max]``.
    context
        Caller name, prepended to the warning message.
    tol
        Out-of-coverage SRF-mass fraction above which a band is flagged.

    Returns
    -------
    frac_out : np.ndarray
        Per-band fraction of SRF response outside the LUT coverage.
    """
    lam = np.asarray(lut_wavelengths_nm, dtype=float)
    lam_min, lam_max = float(lam.min()), float(lam.max())
    hr = srf.wavelengths_hr_nm
    in_range = (hr >= lam_min) & (hr <= lam_max)
    total = srf.matrix.sum(axis=1)
    covered = srf.matrix[:, in_range].sum(axis=1)
    frac_out = np.where(total > 0.0, 1.0 - covered / total, 0.0)
    for b in np.nonzero(frac_out > tol)[0]:
        warnings.warn(
            f"{context}: band {srf.band_names[b]!r} has {frac_out[b] * 100:.1f}% "
            f"of its SRF response outside the LUT wavelength coverage "
            f"[{lam_min:.1f}, {lam_max:.1f}] nm; its band-integrated value is "
            "silently zero-filled there.",
            stacklevel=3,
        )
    return frac_out

instrument

JAX point-spread + ground-sampling instrument operators.

These two operators model the deterministic part of an imaging-spectrometer forward chain that sits between the radiative-transfer output and the band-integration step:

HR radiance cube  ──PSF──►  blurred HR cube  ──GSD──►  LR cube  ──SRF──►  bands

The :class:~plumax.radtran.srf.SpectralResponseFunction already covers the SRF stage; this module fills in the PSF (spatial blur) and GSD (detector down-sampling). Both operators here are:

• JAX-native — every method is jit-able and vmap-able; gradients flow through apply/jacobian/adjoint via reverse-mode AD without manual bookkeeping. This is the prerequisite for the variational retrieval module :mod:plumax.assimilation.
• Linear — jacobian == apply (state-independent), and adjoint implements the Euclidean transpose so the inner-product identity ⟨A x, y⟩ = ⟨x, Aᵀ y⟩ holds to machine precision.
• Cube-shaped (band-last) — inputs are (ny, nx, n_lambda); channels are processed independently. This matches the convention in :mod:plumax.radtran.srf and plays nicely with jax.vmap along the spectral axis.

The numpy reference for the PSF/GSD logic lives in jej_vc_snippets/methane_retrieval/lut_obs_op.py; this is the JAX rewrite that lets the same operators participate in autodiff-based 3D-Var solves.

PointSpreadFunction dataclass

PointSpreadFunction(kernel: ndarray)

2D PSF blur applied independently per spectral channel.

The PSF is a fixed convolution kernel (kernel_size × kernel_size, L1-normalised) so the operator is linear and shift-invariant. The Jacobian is the operator itself; the adjoint is convolution with the 180°-flipped kernel (real-valued case).

The class is frozen=True, eq=False so the default object.__eq__ (identity comparison) and object.__hash__ (id-based) are retained. An auto-generated dataclass __eq__ would do element-wise equality on the kernel: np.ndarray field — ambiguous in bool context — and the corresponding __hash__ would crash because numpy arrays are unhashable. Identity semantics are what jax.jit static-arg caching actually wants: a PSF instance is re-used as-is or replaced wholesale; it never needs structural equality.

Attributes

kernel : np.ndarray L1-normalised 2-D kernel, shape (kernel_size, kernel_size). Stored as numpy for transparent inspection; the JAX copies used by :meth:apply / :meth:adjoint are cached as private fields below.

gaussian classmethod

gaussian(fwhm_pixels: float, kernel_size: int = 9) -> PointSpreadFunction

Build an isotropic Gaussian PSF.

FWHM = 2.355 σ; the kernel is sampled on the discrete pixel grid and L1-normalised in __post_init__.

Source code in src/plumax/radtran/instrument.py

@classmethod
def gaussian(cls, fwhm_pixels: float, kernel_size: int = 9) -> PointSpreadFunction:
    """Build an isotropic Gaussian PSF.

    ``FWHM = 2.355 σ``; the kernel is sampled on the discrete pixel grid
    and L1-normalised in ``__post_init__``.
    """
    if fwhm_pixels <= 0.0:
        raise ValueError("PointSpreadFunction.gaussian: `fwhm_pixels` must be > 0.")
    c = kernel_size // 2
    y, x = np.ogrid[-c : c + 1, -c : c + 1]
    r2 = x * x + y * y
    sigma = fwhm_pixels / 2.355
    k = np.exp(-r2 / (2.0 * sigma * sigma))
    return cls(kernel=k)

airy classmethod

airy(fwhm_pixels: float, kernel_size: int = 9) -> PointSpreadFunction

Build an Airy-disk-like PSF as a diffraction-ish surrogate.

fwhm_pixels is the kernel's full width at half maximum, matching :meth:gaussian. The Airy argument is scaled so the half-max radius (x = _AIRY_HALF_MAX_X) falls at r = fwhm_pixels / 2; the earlier 1.22·fwhm (first-zero) scaling produced a kernel ≈ 0.63× the requested FWHM.

Source code in src/plumax/radtran/instrument.py

@classmethod
def airy(cls, fwhm_pixels: float, kernel_size: int = 9) -> PointSpreadFunction:
    """Build an Airy-disk-like PSF as a diffraction-ish surrogate.

    ``fwhm_pixels`` is the kernel's full width at half maximum, matching
    :meth:`gaussian`. The Airy argument is scaled so the half-max radius
    (``x = _AIRY_HALF_MAX_X``) falls at ``r = fwhm_pixels / 2``; the earlier
    ``1.22·fwhm`` (first-zero) scaling produced a kernel ≈ 0.63× the
    requested FWHM.
    """
    from scipy.special import j1

    if fwhm_pixels <= 0.0:
        raise ValueError("PointSpreadFunction.airy: `fwhm_pixels` must be > 0.")
    c = kernel_size // 2
    y, x = np.ogrid[-c : c + 1, -c : c + 1]
    r = np.sqrt(x * x + y * y)
    xarg = np.where(r == 0.0, 1e-12, 2.0 * _AIRY_HALF_MAX_X * r / fwhm_pixels)
    k = (2.0 * j1(xarg) / xarg) ** 2
    return cls(kernel=k)

apply

apply(radiance: HRCube) -> HRCube

Forward blur: y = K * x per channel.

Source code in src/plumax/radtran/instrument.py

def apply(self, radiance: HRCube) -> HRCube:
    """Forward blur: ``y = K * x`` per channel."""
    x = jnp.asarray(radiance)
    if x.ndim != 3:
        raise ValueError(
            f"PointSpreadFunction.apply: input must be 3-D (ny, nx, n_lambda), "
            f"got shape {x.shape}."
        )
    return _convolve_channels(x, self._kernel_jax)

jacobian

jacobian(tangent: HRCube) -> HRCube

JVP — same as :meth:apply because the operator is linear.

Source code in src/plumax/radtran/instrument.py

def jacobian(self, tangent: HRCube) -> HRCube:
    """JVP — same as :meth:`apply` because the operator is linear."""
    return self.apply(tangent)

adjoint

adjoint(cotangent: HRCube) -> HRCube

Adjoint: convolution with the flipped kernel.

For a real-valued L1-normalised kernel K the discrete adjoint of the per-channel cross-correlation y = K ⋆ x is Kᵀ ⋆ u = K[::-1] ⋆ u. Symmetric kernels (e.g. Gaussian) are self-adjoint.

Source code in src/plumax/radtran/instrument.py

def adjoint(self, cotangent: HRCube) -> HRCube:
    """Adjoint: convolution with the flipped kernel.

    For a real-valued L1-normalised kernel ``K`` the discrete adjoint of
    the per-channel cross-correlation `y = K ⋆ x` is `Kᵀ ⋆ u = K[::-1] ⋆ u`.
    Symmetric kernels (e.g. Gaussian) are self-adjoint.
    """
    u = jnp.asarray(cotangent)
    if u.ndim != 3:
        raise ValueError(
            f"PointSpreadFunction.adjoint: input must be 3-D, got shape {u.shape}."
        )
    return _convolve_channels(u, self._kernel_adj)

GroundSamplingDistance dataclass

GroundSamplingDistance(downsample_factor: int)

Block-mean down-sampling operator for a fixed integer factor.

Maps an HR cube (ny, nx, n_lambda) to LR (ny // f, nx // f, n_lambda) by averaging non-overlapping f × f spatial blocks. If (ny, nx) is not divisible by f, the input is truncated to (ny // f * f, nx // f * f) (top-left crop).

The adjoint distributes each LR value uniformly across the corresponding HR block with weight 1 / f² so that ⟨A x, y⟩ = ⟨x, Aᵀ y⟩ holds exactly under the standard Euclidean inner product.

Attributes

downsample_factor : int Integer block size f.

Notes

The classmethod :meth:from_optics derives f from the standard pinhole-camera GSD formula GSD [m/px] = sensor_width [mm] · altitude [m] · 1000 / (focal_length [mm] · image_width [px]) and rounds to the nearest integer. For non-integer ratios you should pre-resample the HR cube — block-mean is only well-defined for integer factors.

from_optics classmethod

from_optics(sensor_width_mm: float, focal_length_mm: float, image_width_px: int, altitude_m: float, pixel_size_hr_m: float) -> GroundSamplingDistance

Build a GSD operator from camera + altitude + HR pixel size.

Dimensional analysis:

GSD [m/px] = (sensor_pixel_size [m/px]) · (altitude [m]) / (focal_length [m])
           = (sensor_width_mm / image_width_px) · altitude_m / focal_length_mm

The mm units on numerator and denominator cancel, so no unit-conversion factor is needed — the numpy reference snippet we ported from had a spurious × 1000 that inflated every realistic camera's GSD by a factor of 1000.

Source code in src/plumax/radtran/instrument.py

@classmethod
def from_optics(
    cls,
    sensor_width_mm: float,
    focal_length_mm: float,
    image_width_px: int,
    altitude_m: float,
    pixel_size_hr_m: float,
) -> GroundSamplingDistance:
    """Build a GSD operator from camera + altitude + HR pixel size.

    Dimensional analysis:

        GSD [m/px] = (sensor_pixel_size [m/px]) · (altitude [m]) / (focal_length [m])
                   = (sensor_width_mm / image_width_px) · altitude_m / focal_length_mm

    The ``mm`` units on numerator and denominator cancel, so no unit-conversion
    factor is needed — the numpy reference snippet we ported from had a spurious
    ``× 1000`` that inflated every realistic camera's GSD by a factor of 1000.
    """
    gsd_m_per_px = (sensor_width_mm * altitude_m) / (
        focal_length_mm * image_width_px
    )
    f = round(gsd_m_per_px / pixel_size_hr_m)
    return cls(downsample_factor=max(1, f))

lr_shape

lr_shape(hr_shape: tuple[int, int, int]) -> tuple[int, int, int]

LR shape implied by an HR shape and the current factor.

Source code in src/plumax/radtran/instrument.py

def lr_shape(self, hr_shape: tuple[int, int, int]) -> tuple[int, int, int]:
    """LR shape implied by an HR shape and the current factor."""
    ny_hr, nx_hr, n_lambda = hr_shape
    f = self.downsample_factor
    return (ny_hr // f, nx_hr // f, n_lambda)

apply

apply(radiance: HRCube) -> LRCube

Forward down-sample by block mean.

Source code in src/plumax/radtran/instrument.py

def apply(self, radiance: HRCube) -> LRCube:
    """Forward down-sample by block mean."""
    x = jnp.asarray(radiance)
    if x.ndim != 3:
        raise ValueError(
            f"GroundSamplingDistance.apply: input must be 3-D, got shape {x.shape}."
        )
    ny_hr, nx_hr, n_lambda = x.shape
    f = self.downsample_factor
    ny_lr, nx_lr = ny_hr // f, nx_hr // f
    if ny_lr == 0 or nx_lr == 0:
        raise ValueError(
            f"GroundSamplingDistance.apply: input shape {x.shape[:2]} too small "
            f"for factor {f}."
        )
    x = x[: ny_lr * f, : nx_lr * f, :]
    return x.reshape(ny_lr, f, nx_lr, f, n_lambda).mean(axis=(1, 3))

jacobian

jacobian(tangent: HRCube) -> LRCube

JVP — same as :meth:apply.

Source code in src/plumax/radtran/instrument.py

def jacobian(self, tangent: HRCube) -> LRCube:
    """JVP — same as :meth:`apply`."""
    return self.apply(tangent)

adjoint

adjoint(cotangent: LRCube, hr_shape: tuple[int, int, int]) -> HRCube

Adjoint: spread each LR value uniformly across its HR block.

Parameters

cotangent : Array LR cotangent, shape (ny_lr, nx_lr, n_lambda). hr_shape : tuple Target HR shape (ny_hr, nx_hr, n_lambda). Truncated entries from :meth:apply (when ny_hr or nx_hr are not multiples of f) are filled with zeros — the adjoint of the truncation is the inclusion (... → ny_lr·f rows).

Source code in src/plumax/radtran/instrument.py

def adjoint(self, cotangent: LRCube, hr_shape: tuple[int, int, int]) -> HRCube:
    """Adjoint: spread each LR value uniformly across its HR block.

    Parameters
    ----------
    cotangent : Array
        LR cotangent, shape ``(ny_lr, nx_lr, n_lambda)``.
    hr_shape : tuple
        Target HR shape ``(ny_hr, nx_hr, n_lambda)``. Truncated entries
        from :meth:`apply` (when ``ny_hr`` or ``nx_hr`` are not multiples
        of ``f``) are filled with zeros — the adjoint of the truncation
        is the inclusion ``(... → ny_lr·f rows)``.
    """
    u = jnp.asarray(cotangent)
    if u.ndim != 3:
        raise ValueError(
            f"GroundSamplingDistance.adjoint: input must be 3-D, got shape {u.shape}."
        )
    ny_hr, nx_hr, n_lambda = hr_shape
    f = self.downsample_factor
    ny_lr, nx_lr = ny_hr // f, nx_hr // f
    if u.shape != (ny_lr, nx_lr, n_lambda):
        raise ValueError(
            f"GroundSamplingDistance.adjoint: cotangent shape {u.shape} does not "
            f"match expected {(ny_lr, nx_lr, n_lambda)} for hr_shape={hr_shape}."
        )
    # Spread + scale: each LR value contributes 1/f² to every HR pixel
    # in its block (the gradient of a mean reduction).
    block = jnp.repeat(jnp.repeat(u, f, axis=0), f, axis=1) / (f * f)
    if ny_lr * f == ny_hr and nx_lr * f == nx_hr:
        return block
    # Pad with zeros along the truncated rows/cols.
    pad_y = ny_hr - ny_lr * f
    pad_x = nx_hr - nx_lr * f
    return jnp.pad(block, ((0, pad_y), (0, pad_x), (0, 0)))

target

Target-spectrum construction for matched-filter retrievals.

The matched filter projects an observed pixel spectrum onto the expected plume signature, normalised by the background covariance. That signature — the "target spectrum" — is the radiance (or normalised radiance) perturbation a reference plume enhancement would produce given the current atmospheric state.

This module offers three flavours, all returning NumPy arrays shaped (n_channels,) where n_channels is either the LUT wavenumber count (hyperspectral) or the number of SRF bands (multispectral) — see :func:target_bands for the band-integrated variant.

• :func:target_spectrum_normalized_nonlinear — exact exp(-Δτ) - 1. Use when Δτ can approach unity (strong-plume regime).
• :func:target_spectrum_normalized_linear — first-order Maclaurin, t ≈ -σ·N·ΔVMR·L·AMF. Classical linearised matched filter.
• :func:target_bands — integrate any hi-res target through an SRF to produce a band-space target.

All three are ported with simplifications from jej_vc_snippets/methane_retrieval/matched_filter_beerslaw.py — dropping the duplicate "reverse order" Taylor/Maclaurin variants since they agree with the standard-order linearisation at first order.

target_spectrum_normalized_nonlinear

target_spectrum_normalized_nonlinear(ds: Dataset, nu_obs: ndarray, *, T_K: float, p_atm: float, vmr_background: float, delta_vmr: float, path_length_cm: float, amf: float, var: str = 'absorption_cross_section') -> np.ndarray

Exact normalised-radiance target t(ν) = exp(-Δτ(ν)) − 1.

Negative for absorption. Magnitude tracks the cross-section profile and scales with ΔVMR.

vmr_background is accepted for API symmetry with the linearised form but is not used here: the normalised nonlinear model depends only on ΔVMR. Keeping the kwarg means callers can swap forward models without editing the call site.

Source code in src/plumax/radtran/target.py

def target_spectrum_normalized_nonlinear(
    ds: xr.Dataset,
    nu_obs: np.ndarray,
    *,
    T_K: float,
    p_atm: float,
    vmr_background: float,
    delta_vmr: float,
    path_length_cm: float,
    amf: float,
    var: str = "absorption_cross_section",
) -> np.ndarray:
    """Exact normalised-radiance target ``t(ν) = exp(-Δτ(ν)) − 1``.

    Negative for absorption. Magnitude tracks the cross-section profile and
    scales with ``ΔVMR``.

    ``vmr_background`` is accepted for API symmetry with the linearised
    form but is not used here: the *normalised* nonlinear model depends
    only on ``ΔVMR``. Keeping the kwarg means callers can swap forward
    models without editing the call site.
    """
    _ = vmr_background  # see docstring
    result = forward_nonlinear_normalized(
        ds,
        nu_obs,
        T_K=T_K,
        p_atm=p_atm,
        vmr_background=vmr_background,
        delta_vmr=delta_vmr,
        path_length_cm=path_length_cm,
        amf=amf,
        var=var,
    )
    return result.radiance - 1.0

target_spectrum_normalized_linear

target_spectrum_normalized_linear(ds: Dataset, nu_obs: ndarray, *, T_K: float, p_atm: float, delta_vmr: float, path_length_cm: float, amf: float, var: str = 'absorption_cross_section') -> np.ndarray

First-order linearised target t(ν) = -σ · N · ΔVMR · L · AMF.

Equivalent to target_spectrum_normalized_nonlinear in the thin-plume limit Δτ ≪ 1 and considerably faster — no exponential per pixel. This is the canonical matched-filter target for weak retrievals.

Source code in src/plumax/radtran/target.py

def target_spectrum_normalized_linear(
    ds: xr.Dataset,
    nu_obs: np.ndarray,
    *,
    T_K: float,
    p_atm: float,
    delta_vmr: float,
    path_length_cm: float,
    amf: float,
    var: str = "absorption_cross_section",
) -> np.ndarray:
    """First-order linearised target ``t(ν) = -σ · N · ΔVMR · L · AMF``.

    Equivalent to ``target_spectrum_normalized_nonlinear`` in the thin-plume
    limit ``Δτ ≪ 1`` and considerably faster — no exponential per pixel.
    This is the canonical matched-filter target for weak retrievals.
    """
    result = forward_maclaurin_normalized(
        ds,
        nu_obs,
        T_K=T_K,
        p_atm=p_atm,
        delta_vmr=delta_vmr,
        path_length_cm=path_length_cm,
        amf=amf,
        order=1,
        var=var,
    )
    return result.radiance - 1.0

target_bands

target_bands(target_hr: ndarray, srf: SpectralResponseFunction, wavelengths_nm: ndarray) -> np.ndarray

Band-integrate a hi-res target spectrum through an SRF.

Parameters

target_hr : np.ndarray Target spectrum on the LUT wavenumber grid, shape (n_nu,). srf : SpectralResponseFunction Instrument SRF defined on srf.wavelengths_hr_nm. wavelengths_nm : np.ndarray Wavelength grid 1e7 / nu_obs [nm] at which target_hr is evaluated. Must be the same length as target_hr but may be in decreasing order (we sort internally).

Returns

target_b : np.ndarray Shape (n_bands,).

Source code in src/plumax/radtran/target.py

def target_bands(
    target_hr: np.ndarray,
    srf: SpectralResponseFunction,
    wavelengths_nm: np.ndarray,
) -> np.ndarray:
    """Band-integrate a hi-res target spectrum through an SRF.

    Parameters
    ----------
    target_hr : np.ndarray
        Target spectrum on the LUT wavenumber grid, shape ``(n_nu,)``.
    srf : SpectralResponseFunction
        Instrument SRF defined on ``srf.wavelengths_hr_nm``.
    wavelengths_nm : np.ndarray
        Wavelength grid ``1e7 / nu_obs`` [nm] at which ``target_hr`` is
        evaluated. Must be the same length as ``target_hr`` but may be in
        decreasing order (we sort internally).

    Returns
    -------
    target_b : np.ndarray
        Shape ``(n_bands,)``.
    """
    if target_hr.shape != wavelengths_nm.shape:
        raise ValueError(
            f"target_bands: `target_hr` {target_hr.shape} and "
            f"`wavelengths_nm` {wavelengths_nm.shape} must match."
        )
    sort_idx = np.argsort(wavelengths_nm)
    lam = np.asarray(wavelengths_nm)[sort_idx]
    t = np.asarray(target_hr)[sort_idx]
    # Flag bands whose SRF response reaches past the LUT wavelength coverage;
    # np.interp zero-fills there and would silently attenuate the target.
    check_srf_lut_coverage(srf, lam, context="target_bands")
    t_on_srf_grid = np.interp(srf.wavelengths_hr_nm, lam, t, left=0.0, right=0.0)
    return srf.apply(t_on_srf_grid)

background

Robust background statistics for matched-filter retrievals.

The matched filter needs a background mean spectrum μ and an inverse covariance Σ⁻¹ estimated from the clean pixels in the scene. Real scenes have contaminants — plumes, clouds, shadows, bright targets — that would bias a naive mean and inflate a naive covariance, so both estimators here are written to be robust against outlier pixels.

• :func:trimmed_mean_spectrum — per-channel trimmed mean. Cheap, robust to a few percent of outliers, matches a Gaussian mean in the clean limit.
• :func:robust_lowrank_covariance — low-rank plus diagonal-regularisation covariance via truncated SVD on the mean-subtracted, trimmed pixel stack. Returns (Σ, Σ⁻¹) ready for the matched filter.

Both operate on a pixel × band array (or an xarray.DataArray with band as the last dim). The trim_frac parameter is a two-sided fraction in [0, 0.5) — 0.1 removes the brightest and darkest 10% of pixels per channel before averaging.

Heavily simplified from jej_vc_snippets/methane_retrieval/matched_filter_{mean,covariance}.py (GMM-based background estimators and dynamic-mode rejection dropped — the trimmed mean + truncated SVD covers the demo regime with far less tuning).

trimmed_mean_spectrum

trimmed_mean_spectrum(radiance: ndarray, *, trim_frac: float = 0.1, band_axis: int = 0) -> np.ndarray

Per-channel trimmed mean μ_b.

Parameters

radiance : np.ndarray Radiance cube, shape (n_bands, ny, nx) by default (band_axis=0), or any array where one axis indexes the band. trim_frac : float Fraction trimmed from each tail per channel, in [0, 0.5). Default 0.1 removes the top and bottom 10% per band. band_axis : int Band axis of radiance. Default 0.

Returns

mu : np.ndarray Trimmed mean per band, shape (n_bands,).

Source code in src/plumax/radtran/background.py

def trimmed_mean_spectrum(
    radiance: np.ndarray,
    *,
    trim_frac: float = 0.1,
    band_axis: int = 0,
) -> np.ndarray:
    """Per-channel trimmed mean ``μ_b``.

    Parameters
    ----------
    radiance : np.ndarray
        Radiance cube, shape ``(n_bands, ny, nx)`` by default (``band_axis=0``),
        or any array where one axis indexes the band.
    trim_frac : float
        Fraction trimmed from *each* tail per channel, in ``[0, 0.5)``.
        Default 0.1 removes the top and bottom 10% per band.
    band_axis : int
        Band axis of ``radiance``. Default 0.

    Returns
    -------
    mu : np.ndarray
        Trimmed mean per band, shape ``(n_bands,)``.
    """
    if not (0.0 <= trim_frac < 0.5):
        raise ValueError(
            f"trimmed_mean_spectrum: `trim_frac` must be in [0, 0.5) (got {trim_frac!r})"
        )
    flat = _flatten_pixels(radiance, band_axis)  # (n_pixels, n_bands)
    return np.asarray(trim_mean(flat, trim_frac, axis=0), dtype=float)

robust_lowrank_covariance

robust_lowrank_covariance(radiance: ndarray, *, rank: int | None = None, trim_frac: float = 0.1, regularization: float | None = None, band_axis: int = 0) -> tuple[np.ndarray, np.ndarray]

Robust low-rank covariance + inverse (Σ, Σ⁻¹).

The pixel stack is centred with a trimmed mean, the extreme 2·trim_frac pixels (by total energy) are removed, and the covariance is estimated as

Σ = U_k · S_k² · U_kᵀ / N + λ · I

where U_k are the top-k left singular vectors of the centred matrix. The added diagonal term λ · I stabilises the inverse when the covariance is ill-conditioned. Returns both Σ and Σ⁻¹ since the matched filter needs the inverse and diagnostics sometimes want the forward covariance too.

Parameters

radiance : np.ndarray Shape (n_bands, ny, nx) by default. rank : int or None Truncation rank. None uses min(n_bands - 1, 16) — the practical sweet spot for multispectral data where most of the variance is in a handful of modes. trim_frac : float Per-pixel energy-based trim fraction in [0, 0.5). Default 0.1. regularization : float or None Diagonal Tikhonov floor λ on the discarded subspace. None (default) uses the noise-matched floor — the mean of the discarded covariance eigenvalues mean(S[rank:]²)/N — so Σ⁻¹ weights that subspace at the real sensor-noise level instead of an arbitrary constant. A floor too small over-weights the noise subspace (inflating false alarms); too large leaks signal into it. Pass a float to override. band_axis : int Band axis of radiance. Default 0.

Returns

Sigma, Sigma_inv : np.ndarray Symmetric PSD matrices, shape (n_bands, n_bands).

Source code in src/plumax/radtran/background.py

def robust_lowrank_covariance(
    radiance: np.ndarray,
    *,
    rank: int | None = None,
    trim_frac: float = 0.1,
    regularization: float | None = None,
    band_axis: int = 0,
) -> tuple[np.ndarray, np.ndarray]:
    """Robust low-rank covariance + inverse ``(Σ, Σ⁻¹)``.

    The pixel stack is centred with a trimmed mean, the extreme 2·``trim_frac``
    pixels (by total energy) are removed, and the covariance is estimated as

        Σ = U_k · S_k² · U_kᵀ / N + λ · I

    where ``U_k`` are the top-``k`` left singular vectors of the centred
    matrix. The added diagonal term ``λ · I`` stabilises the inverse when
    the covariance is ill-conditioned. Returns both ``Σ`` and ``Σ⁻¹`` since
    the matched filter needs the inverse and diagnostics sometimes want the
    forward covariance too.

    Parameters
    ----------
    radiance : np.ndarray
        Shape ``(n_bands, ny, nx)`` by default.
    rank : int or None
        Truncation rank. ``None`` uses ``min(n_bands - 1, 16)`` — the
        practical sweet spot for multispectral data where most of the
        variance is in a handful of modes.
    trim_frac : float
        Per-pixel energy-based trim fraction in ``[0, 0.5)``. Default 0.1.
    regularization : float or None
        Diagonal Tikhonov floor ``λ`` on the discarded subspace. ``None``
        (default) uses the **noise-matched** floor — the mean of the discarded
        covariance eigenvalues ``mean(S[rank:]²)/N`` — so ``Σ⁻¹`` weights that
        subspace at the real sensor-noise level instead of an arbitrary
        constant. A floor too small over-weights the noise subspace (inflating
        false alarms); too large leaks signal into it. Pass a float to override.
    band_axis : int
        Band axis of ``radiance``. Default 0.

    Returns
    -------
    Sigma, Sigma_inv : np.ndarray
        Symmetric PSD matrices, shape ``(n_bands, n_bands)``.
    """
    mu, U, d, lam = _estimate_lowrank_model(
        radiance,
        rank=rank,
        trim_frac=trim_frac,
        regularization=regularization,
        band_axis=band_axis,
        context="robust_lowrank_covariance",
    )
    n_bands = mu.shape[0]
    # Σ = U diag(d) Uᵀ + λ I (d = S_k² / N are the retained covariance eigenvalues).
    Sigma = (U * d) @ U.T + lam * np.eye(n_bands)

    # Woodbury inverse for efficiency: (λI + Uᵀ D U)⁻¹ where
    #   U = V_k (rank, n_bands), D = diag(S_k² / N) (rank, rank).
    # But since n_bands is small for multispectral, direct inversion is fine
    # and keeps the code obvious.
    Sigma_inv = np.linalg.inv(Sigma)
    return Sigma, Sigma_inv

matched_filter

Matched-filter retrieval for multispectral / hyperspectral methane detection.

The matched filter is the optimal linear detector for a known target signature embedded in additive Gaussian noise with known covariance. For a pixel spectrum x ∈ ℝ^B, background mean μ, covariance Σ, and target t, the maximum-likelihood abundance estimate is

ε̂ = (x − μ)ᵀ Σ⁻¹ t  /  (tᵀ Σ⁻¹ t) .

Positive values of ε̂ indicate the target is present; ε̂ ≈ ΔVMR / ΔVMR_ref when t is the target spectrum produced by :func:plumax.radtran.target.target_spectrum_normalized_linear at delta_vmr = ΔVMR_ref. Multiplying ε̂ back by ΔVMR_ref gives an estimate of the actual pixel VMR enhancement.

This module exposes the raw operation at three scopes:

• :func:matched_filter_pixel — one pixel (scalar output).
• :func:matched_filter_image — vectorised over spatial axes.
• :func:matched_filter_snr — per-pixel detection SNR.

Ported from jej_vc_snippets/methane_retrieval/matched_filter.py with the dask/chunked path dropped (callers who need it can chunk the input themselves and map :func:matched_filter_image over chunks).

matched_filter_pixel

matched_filter_pixel(spectrum: ndarray, mean_spectrum: ndarray, cov_inv: ndarray, target: ndarray) -> float

Matched-filter abundance for a single pixel.

Parameters

spectrum : np.ndarray Pixel spectrum, shape (n_bands,). mean_spectrum : np.ndarray Background mean, shape (n_bands,). cov_inv : np.ndarray Inverse background covariance Σ⁻¹, shape (n_bands, n_bands). target : np.ndarray Target signature, shape (n_bands,).

Returns

abundance : float Scalar retrieval in the same units as target.

Source code in src/plumax/radtran/matched_filter.py

def matched_filter_pixel(
    spectrum: np.ndarray,
    mean_spectrum: np.ndarray,
    cov_inv: np.ndarray,
    target: np.ndarray,
) -> float:
    """Matched-filter abundance for a single pixel.

    Parameters
    ----------
    spectrum : np.ndarray
        Pixel spectrum, shape ``(n_bands,)``.
    mean_spectrum : np.ndarray
        Background mean, shape ``(n_bands,)``.
    cov_inv : np.ndarray
        Inverse background covariance ``Σ⁻¹``, shape ``(n_bands, n_bands)``.
    target : np.ndarray
        Target signature, shape ``(n_bands,)``.

    Returns
    -------
    abundance : float
        Scalar retrieval in the same units as ``target``.
    """
    cov_inv_target, target_norm = _precompute_constants(target, cov_inv)
    innovation = np.asarray(spectrum, dtype=float) - np.asarray(
        mean_spectrum, dtype=float
    )
    return float(np.einsum("i, i ->", cov_inv_target, innovation) / target_norm)

matched_filter_image

matched_filter_image(radiance: ndarray, mean_spectrum: ndarray, cov_inv: ndarray, target: ndarray, *, band_axis: int = 0) -> np.ndarray

Matched-filter abundance map, vectorised over spatial axes.

Parameters

radiance : np.ndarray Radiance cube, shape (n_bands, ny, nx) by default. mean_spectrum, cov_inv, target : np.ndarray As in :func:matched_filter_pixel. band_axis : int Band axis of radiance. Default 0.

Returns

abundance : np.ndarray Shape radiance.shape minus the band axis — e.g. (ny, nx) for the default layout.

Source code in src/plumax/radtran/matched_filter.py

def matched_filter_image(
    radiance: np.ndarray,
    mean_spectrum: np.ndarray,
    cov_inv: np.ndarray,
    target: np.ndarray,
    *,
    band_axis: int = 0,
) -> np.ndarray:
    """Matched-filter abundance map, vectorised over spatial axes.

    Parameters
    ----------
    radiance : np.ndarray
        Radiance cube, shape ``(n_bands, ny, nx)`` by default.
    mean_spectrum, cov_inv, target : np.ndarray
        As in :func:`matched_filter_pixel`.
    band_axis : int
        Band axis of ``radiance``. Default 0.

    Returns
    -------
    abundance : np.ndarray
        Shape ``radiance.shape`` minus the band axis — e.g. ``(ny, nx)`` for
        the default layout.
    """
    cov_inv_target, target_norm = _precompute_constants(target, cov_inv)
    arr = np.asarray(radiance, dtype=float)
    bands_first = np.moveaxis(arr, band_axis, 0)
    n_bands = bands_first.shape[0]
    if n_bands != target.size:
        raise ValueError(
            f"matched_filter_image: radiance has {n_bands} bands but target "
            f"has {target.size}."
        )
    # Broadcasting subtraction: (n_bands, ...) - (n_bands,)
    innovation = bands_first - np.asarray(mean_spectrum, dtype=float).reshape(
        (n_bands,) + (1,) * (bands_first.ndim - 1)
    )
    # Project along the band axis. Reshape cov_inv_target for broadcasting,
    # collapse leading band axis with sum; this matches the einops
    # "b, b ... -> ..." contraction without the einops dependency.
    broadcastable = cov_inv_target.reshape((n_bands,) + (1,) * (innovation.ndim - 1))
    projected = (broadcastable * innovation).sum(axis=0)
    return projected / target_norm

matched_filter_snr

matched_filter_snr(abundance: ndarray, cov_inv: ndarray, target: ndarray) -> np.ndarray

Detection SNR for a matched-filter abundance map.

For the matched filter, SNR(x) = (target energy) · ε̂ = ε̂ · √(tᵀ Σ⁻¹ t).

Parameters

abundance : np.ndarray Output of :func:matched_filter_image, shape (..., ny, nx). cov_inv : np.ndarray Inverse covariance used to build the filter. target : np.ndarray Target used to build the filter.

Returns

snr : np.ndarray Same shape as abundance.

Source code in src/plumax/radtran/matched_filter.py

def matched_filter_snr(
    abundance: np.ndarray,
    cov_inv: np.ndarray,
    target: np.ndarray,
) -> np.ndarray:
    """Detection SNR for a matched-filter abundance map.

    For the matched filter,
        SNR(x) = (target energy) · ε̂ = ε̂ · √(tᵀ Σ⁻¹ t).

    Parameters
    ----------
    abundance : np.ndarray
        Output of :func:`matched_filter_image`, shape ``(..., ny, nx)``.
    cov_inv : np.ndarray
        Inverse covariance used to build the filter.
    target : np.ndarray
        Target used to build the filter.

    Returns
    -------
    snr : np.ndarray
        Same shape as ``abundance``.
    """
    _, target_norm = _precompute_constants(target, cov_inv)
    return np.asarray(abundance, dtype=float) * np.sqrt(target_norm)

gaussx_solve

Gaussx-backed structured solves for the matched filter.

The default matched_filter_* functions in :mod:.matched_filter take a materialised inverse covariance Σ⁻¹ ∈ ℝ^{B×B} and compute

ε̂ = (x - μ)ᵀ Σ⁻¹ t / (tᵀ Σ⁻¹ t)

directly. That's fine when B is small (multispectral) but wasteful for hyperspectral retrievals with B ≳ 200 because:

1. The O(B³) dense inverse is entirely avoidable — the empirical covariance is naturally represented as a low-rank update Σ = λ I + U diag(d) Uᵀ.
2. Every solve can then run in O(B·k) via the Woodbury identity, where k is the retained rank (typically 8–32).

This module wraps :mod:gaussx's :class:gaussx.LowRankUpdate operator and structural :func:gaussx.solve so the matched filter can stay inversion-free on the hot path:

• :func:build_lowrank_covariance_operator — construct a :class:gaussx.LowRankUpdate from the same SVD the dense estimator uses.
• :func:matched_filter_pixel_op / :func:matched_filter_image_op — operator-backed matched filter that calls gaussx.solve(cov, target) once and then dots the result into every pixel's innovation.

The functions accept NumPy inputs and return NumPy outputs — the gaussx internals run on JAX arrays for us, but callers don't need to know.

.. important:: This path requires 64-bit JAX (jax_enable_x64). With x64 disabled (JAX's default) the float64 SVD factors are silently truncated to float32 and the ill-conditioned Woodbury capacitance solve loses precision (up to ~1e-5 relative, worse for hyperspectral B ≳ 200). Rather than return a quietly-wrong result, every entry point raises a clear error when x64 is off — enable it once at import time: import jax; jax.config.update("jax_enable_x64", True).

build_lowrank_covariance_operator

build_lowrank_covariance_operator(radiance: ndarray, *, rank: int | None = None, trim_frac: float = 0.1, regularization: float | None = None, band_axis: int = 0) -> tuple[LowRankUpdateOp, np.ndarray]

Build a :class:gaussx.LowRankUpdate covariance operator from a scene.

Parameters

radiance : np.ndarray Radiance cube, shape (n_bands, ny, nx) by default. rank, trim_frac, regularization, band_axis Same semantics as :func:.background.robust_lowrank_covariance — see that function for the trimming and rank-selection details.

Returns

cov : gaussx.LowRankUpdate Structured operator λ I + U diag(d) Uᵀ (symmetric + PSD tagged) where d = S_k² / N are the squared singular values of the trimmed, mean-subtracted pixel stack. mu : np.ndarray Background mean spectrum, shape (n_bands,). Returned here (rather than a separate call) so callers build μ and Σ from the same trimmed pixel stack — avoiding a subtle bias where μ drops outliers that Σ's trim step kept.

Source code in src/plumax/radtran/gaussx_solve.py

def build_lowrank_covariance_operator(
    radiance: np.ndarray,
    *,
    rank: int | None = None,
    trim_frac: float = 0.1,
    regularization: float | None = None,
    band_axis: int = 0,
) -> tuple[LowRankUpdateOp, np.ndarray]:
    """Build a :class:`gaussx.LowRankUpdate` covariance operator from a scene.

    Parameters
    ----------
    radiance : np.ndarray
        Radiance cube, shape ``(n_bands, ny, nx)`` by default.
    rank, trim_frac, regularization, band_axis
        Same semantics as :func:`.background.robust_lowrank_covariance` —
        see that function for the trimming and rank-selection details.

    Returns
    -------
    cov : gaussx.LowRankUpdate
        Structured operator ``λ I + U diag(d) Uᵀ`` (symmetric + PSD tagged)
        where ``d = S_k² / N`` are the squared singular values of the
        trimmed, mean-subtracted pixel stack.
    mu : np.ndarray
        Background mean spectrum, shape ``(n_bands,)``. Returned here
        (rather than a separate call) so callers build μ and Σ from the
        *same* trimmed pixel stack — avoiding a subtle bias where μ drops
        outliers that Σ's trim step kept.
    """
    gx, jnp, _ = _import_gaussx_stack()

    from plumax.radtran.background import _estimate_lowrank_model

    # Shared trim/SVD model — identical (μ, U, d, λ) to the dense
    # `robust_lowrank_covariance` path, so the two cannot drift apart.
    mu, U, d, lam = _estimate_lowrank_model(
        radiance,
        rank=rank,
        trim_frac=trim_frac,
        regularization=regularization,
        band_axis=band_axis,
        context="build_lowrank_covariance_operator",
    )
    n_bands = mu.shape[0]
    # λ I + U diag(d) Uᵀ via the gaussx SVD constructor: it infers the
    # symmetric/PSD tags and takes the orthonormal-factor solve path (U's columns
    # are SVD directions), so we don't hand-maintain tags or a diagonal base.
    cov = gx.svd_low_rank_plus_diag(
        jnp.full(n_bands, lam), jnp.asarray(U), jnp.asarray(d), jnp.asarray(U)
    )
    return cov, mu

matched_filter_pixel_op

matched_filter_pixel_op(spectrum: ndarray, mean_spectrum: ndarray, cov: LowRankUpdateOp, target: ndarray) -> float

Scalar matched filter using a structured covariance operator.

Computes Σ⁻¹ t via :func:gaussx.solve once (O(B·k) via Woodbury), then projects a single pixel's innovation onto it.

Source code in src/plumax/radtran/gaussx_solve.py

def matched_filter_pixel_op(
    spectrum: np.ndarray,
    mean_spectrum: np.ndarray,
    cov: LowRankUpdateOp,
    target: np.ndarray,
) -> float:
    """Scalar matched filter using a structured covariance operator.

    Computes ``Σ⁻¹ t`` via :func:`gaussx.solve` once (O(B·k) via Woodbury),
    then projects a single pixel's innovation onto it.
    """
    gx, jnp, _ = _import_gaussx_stack()

    cov_inv_target = np.asarray(gx.solve(cov, jnp.asarray(target)))
    target_norm = float(np.dot(target, cov_inv_target))
    if not (target_norm > 0.0):
        raise ValueError(
            "matched_filter_pixel_op: `tᵀ Σ⁻¹ t` must be positive; check "
            "that Σ is PD and the target is non-trivial."
        )
    innovation = np.asarray(spectrum, dtype=float) - np.asarray(
        mean_spectrum,
        dtype=float,
    )
    return float(np.dot(cov_inv_target, innovation) / target_norm)

matched_filter_image_op

matched_filter_image_op(radiance: ndarray, mean_spectrum: ndarray, cov: LowRankUpdateOp, target: ndarray, *, band_axis: int = 0) -> np.ndarray

Vectorised matched filter using a structured covariance operator.

Parameters

radiance : np.ndarray Radiance cube, shape (n_bands, ny, nx) by default. mean_spectrum : np.ndarray Background mean, shape (n_bands,). cov : gaussx.LowRankUpdate Structured covariance operator (from :func:build_lowrank_covariance_operator). target : np.ndarray Target signature, shape (n_bands,). band_axis : int Band axis of radiance. Default 0.

Returns

abundance : np.ndarray Pixel-wise retrieval, shape radiance.shape minus the band axis.

Source code in src/plumax/radtran/gaussx_solve.py

def matched_filter_image_op(
    radiance: np.ndarray,
    mean_spectrum: np.ndarray,
    cov: LowRankUpdateOp,
    target: np.ndarray,
    *,
    band_axis: int = 0,
) -> np.ndarray:
    """Vectorised matched filter using a structured covariance operator.

    Parameters
    ----------
    radiance : np.ndarray
        Radiance cube, shape ``(n_bands, ny, nx)`` by default.
    mean_spectrum : np.ndarray
        Background mean, shape ``(n_bands,)``.
    cov : gaussx.LowRankUpdate
        Structured covariance operator (from
        :func:`build_lowrank_covariance_operator`).
    target : np.ndarray
        Target signature, shape ``(n_bands,)``.
    band_axis : int
        Band axis of ``radiance``. Default 0.

    Returns
    -------
    abundance : np.ndarray
        Pixel-wise retrieval, shape ``radiance.shape`` minus the band axis.
    """
    gx, jnp, _ = _import_gaussx_stack()

    arr = np.asarray(radiance, dtype=float)
    bands_first = np.moveaxis(arr, band_axis, 0)
    n_bands = bands_first.shape[0]
    if n_bands != target.size:
        raise ValueError(
            f"matched_filter_image_op: radiance has {n_bands} bands but "
            f"target has {target.size}."
        )

    cov_inv_target = np.asarray(gx.solve(cov, jnp.asarray(target)))  # (n_bands,)
    target_norm = float(np.dot(target, cov_inv_target))
    if not (target_norm > 0.0):
        raise ValueError(
            "matched_filter_image_op: `tᵀ Σ⁻¹ t` must be positive; check "
            "that Σ is PD and the target is non-trivial."
        )

    # Subtract μ along the band axis, project onto (Σ⁻¹ t), divide by norm.
    innovation = bands_first - np.asarray(mean_spectrum, dtype=float).reshape(
        (n_bands,) + (1,) * (bands_first.ndim - 1),
    )
    projector = cov_inv_target.reshape(
        (n_bands,) + (1,) * (bands_first.ndim - 1),
    )
    return (projector * innovation).sum(axis=0) / target_norm

matched_filter_snr_op

matched_filter_snr_op(abundance: ndarray, cov: LowRankUpdateOp, target: ndarray) -> np.ndarray

Detection SNR when the covariance is a gaussx operator.

Mirrors :func:plumax.radtran.matched_filter.matched_filter_snr but uses :func:gaussx.solve to avoid materialising Σ⁻¹.

Source code in src/plumax/radtran/gaussx_solve.py

def matched_filter_snr_op(
    abundance: np.ndarray,
    cov: LowRankUpdateOp,
    target: np.ndarray,
) -> np.ndarray:
    """Detection SNR when the covariance is a gaussx operator.

    Mirrors :func:`plumax.radtran.matched_filter.matched_filter_snr`
    but uses :func:`gaussx.solve` to avoid materialising ``Σ⁻¹``.
    """
    gx, jnp, _ = _import_gaussx_stack()

    cov_inv_target = np.asarray(gx.solve(cov, jnp.asarray(target)))
    target_norm = float(np.dot(target, cov_inv_target))
    return np.asarray(abundance, dtype=float) * np.sqrt(target_norm)