Coupled inversion¶

End-to-end coupled, multi-instrument source inversion. Tier IV.

coupled ¶

Tier IV — end-to-end coupled, multi-instrument source inversion.

Tier IV is assembly + multi-instrument fusion, not new physics (design Tier IV): it composes a transport tier with the observation operator (averaging kernel / RTM) and fuses several satellites' observations into one coherent posterior.

This v1 implements the design's build-order step 1 — Tier I (Gaussian plume) + averaging kernel + multi-instrument L2 fusion for a static emission rate, with the per-instrument bias as a first-class state element:

instrument — :class:Instrument: per-satellite receptor geometry, averaging kernel, and R = R_retr + R_repr + R_align error budget.
forward — the coupled Tier I → column → AK forward, kept per-instrument at native resolution (build_coupled_forward / CoupledForward); linear in the emission rate Q.
fusion — the closed-form joint posterior over (Q, bias_inst) (fuse_observations), exploiting that linearity (the design's linear-conditional-Gaussian limit).
rtm — additive RTM-based observation operators that map the plume column enhancement → gas ΔVMR → band-integrated normalised radiance via the :mod:plumax.radtran Beer–Lambert + SRF stack, as the build-order step toward L1-radiance fusion. Two complementary, well-known models: the exact nonlinear RadianceObservationOperator / radiance_response (L_norm = exp(-Δτ)), and its tangent-linear counterpart LinearisedRadianceOperator / radiance_response_linear (via linearise), which is linear in Q and pure-JAX so it feeds gradient -based and closed-form inversion. The scalar-AK path above is unchanged.

Later build-order steps (Tier II/III transport + RTM, the Q(t) stochastic process, trans-dimensional source count, coupled emulator, operational predictor) are future work; the fusion harness and likelihood structure here are designed to stay the same as those blocks are swapped in.

CoupledForward `dataclass` ¶

CoupledForward(source: PlumeSource, instruments: tuple[Instrument, ...], responses: tuple[Array, ...])

Multi-instrument coupled forward: a list of per-instrument observations.

Holds the shared source/met configuration and the instrument list, and pre-computes each instrument's unit-Q :func:column_response (the static, Q-independent part of the forward). Use :func:build_coupled_forward.

Attributes:

Name	Type	Description
`source`	`PlumeSource`	Shared source + met configuration.
`instruments`	`tuple[Instrument, ...]`	The observing instruments (fusion list).
`responses`	`tuple[Array, ...]`	Per-instrument unit-`Q` column responses (one `(n_obs,)` array each), aligned with `instruments`.

predict ¶

predict(emission_rate: Array | float, biases: Sequence[Array | float] | None = None) -> list[jax.Array]

Predicted observation vector per instrument (list, native resolution).

Parameters:

Name	Type	Description	Default
`emission_rate`	`Array \| float`	Emission rate `Q` [kg/s].	required
`biases`	`Sequence[Array \| float] \| None`	Per-instrument additive biases, one per instrument. Defaults to zero for every instrument.	`None`

Returns:

Type	Description
`list[Array]`	A list of `(n_obs_i,)` predictions, aligned with `instruments`.

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

def predict(
    self,
    emission_rate: jax.Array | float,
    biases: Sequence[jax.Array | float] | None = None,
) -> list[jax.Array]:
    """Predicted observation vector per instrument (list, native resolution).

    Args:
        emission_rate: Emission rate ``Q`` [kg/s].
        biases: Per-instrument additive biases, one per instrument. Defaults
            to zero for every instrument.

    Returns:
        A list of ``(n_obs_i,)`` predictions, aligned with ``instruments``.
    """
    if biases is None:
        biases = [0.0] * len(self.instruments)
    if len(biases) != len(self.instruments):
        raise ValueError(
            f"CoupledForward.predict: got {len(biases)} biases for "
            f"{len(self.instruments)} instruments."
        )
    return [
        emission_rate * resp + self.source.background + b
        for resp, b in zip(self.responses, biases, strict=True)
    ]

PlumeSource `dataclass` ¶

PlumeSource(location: tuple[float, float, float], wind_speed: float, wind_direction: float, stability_class: str, column_z: Array, background: float = 0.0)

Static source + meteorology for the Tier I coupled forward.

The emission rate Q is deliberately not stored here — it is the scalar the inversion solves for, so the forward is parameterised by everything except Q (location, wind, stability, background) and is linear in it.

Attributes:

Name	Type	Description
`location`	`tuple[float, float, float]`	`(x, y, z)` source coordinates [m].
`wind_speed`	`float`	Wind speed [m/s], `> 0`.
`wind_direction`	`float`	Meteorological direction [deg from North] (`270` = wind from the west, flowing east).
`stability_class`	`str`	Pasquill class `'A'`–`'F'`.
`column_z`	`Array`	1-D vertical grid [m] for the column integral, shape `(n_z,)`, strictly increasing.
`background`	`float`	Additive background column (same units as the AK output), applied per receptor. Default `0`.

wind_uv ¶

wind_uv() -> tuple[float, float]

Wind (u, v) [m/s] from speed + meteorological direction.

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

def wind_uv(self) -> tuple[float, float]:
    """Wind ``(u, v)`` [m/s] from speed + meteorological direction."""
    theta = np.deg2rad(self.wind_direction)
    u = -self.wind_speed * np.sin(theta)
    v = -self.wind_speed * np.cos(theta)
    return float(u), float(v)

FusionPosterior `dataclass` ¶

FusionPosterior(emission_rate: Array, emission_std: Array, biases: Array, bias_std: Array, mean: Array, covariance: Array, instrument_names: tuple[str, ...])

Closed-form joint posterior over (Q, bias₁ … bias_M).

All numeric fields are kept as JAX arrays (the scalars as 0-d arrays) so the whole posterior stays differentiable / jittable — the design advertises the estimator for gradient-based workflows, and forcing float(...) here would raise a concretization error under jax.grad / jax.jit. Use float(post.emission_rate) at the call site when a Python scalar is wanted.

Attributes:

Name	Type	Description
`emission_rate`	`Array`	Posterior mean emission `Q*` [kg/s], 0-d array.
`emission_std`	`Array`	Posterior standard deviation of `Q`, 0-d array.
`biases`	`Array`	Posterior mean per-instrument biases, shape `(n_inst,)`.
`bias_std`	`Array`	Posterior standard deviations of the biases, shape `(n_inst,)`.
`mean`	`Array`	Full posterior mean state `[Q, bias₁ … bias_M]`, shape `(1 + n_inst,)`.
`covariance`	`Array`	Full posterior covariance `P*`, shape `(1 + n_inst, 1 + n_inst)`.
`instrument_names`	`tuple[str, ...]`	Names aligned with the bias entries.

Instrument `dataclass` ¶

Instrument(name: str, receptors: Array, averaging_kernel: Array | None = None, retrieval_variance: Array | float = 1.0, representation_variance: Array | float = 0.0, alignment_variance: Array | float = 0.0, observation_time: float = 0.0)

Fixed observation geometry + error budget for one satellite instrument.

Attributes:

Name	Type	Description
`name`	`str`	Instrument identifier (e.g. `"TROPOMI"`).
`receptors`	`Array`	Ground sample locations `(n_obs, 2)` as `(x, y)` [m].
`averaging_kernel`	`Array \| None`	Per-receptor scalar AK `A_inst` applied to the column enhancement, shape `(n_obs,)`. Defaults to all-ones (identity AK).
`retrieval_variance`	`Array \| float`	Per-receptor L2 retrieval error `R_retr`, shape `(n_obs,)`.
`representation_variance`	`Array \| float`	Model-vs-footprint mismatch `R_repr` (scalar or `(n_obs,)`). Default `0`.
`alignment_variance`	`Array \| float`	Temporal-misalignment error `R_align` (scalar or `(n_obs,)`). Default `0`.
`observation_time`	`float`	Overpass time [s] (used by time-resolved transport tiers; ignored by the steady-state Tier I forward). Default `0`.

n_obs `property` ¶

n_obs: int

Number of receptor pixels.

ak `property` ¶

ak: Array

Averaging kernel as an array (identity ones if unset).

observation_variance ¶

observation_variance() -> jax.Array

Total per-receptor observation variance R_retr + R_repr + R_align.

Each term broadcasts to (n_obs,). The design flags omitting the representation / alignment terms as a source of overconfident posteriors, so all three are summed here.

Raises:

Type	Description
`ValueError`	if any resulting variance is non-positive (a zero/negative variance would blow up or invert the likelihood weight).

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

def observation_variance(self) -> jax.Array:
    """Total per-receptor observation variance ``R_retr + R_repr + R_align``.

    Each term broadcasts to ``(n_obs,)``. The design flags omitting the
    representation / alignment terms as a source of overconfident posteriors,
    so all three are summed here.

    Raises:
        ValueError: if any resulting variance is non-positive (a zero/negative
            variance would blow up or invert the likelihood weight).
    """
    n = self.n_obs
    r = (
        jnp.broadcast_to(jnp.asarray(self.retrieval_variance, dtype=float), (n,))
        + jnp.broadcast_to(
            jnp.asarray(self.representation_variance, dtype=float), (n,)
        )
        + jnp.broadcast_to(jnp.asarray(self.alignment_variance, dtype=float), (n,))
    )
    # Eager positivity check, skipped under jit / grad / vmap — the fusion
    # forward calls this from inside a traced cost, where ``r`` is a tracer
    # (every jnp op inside a trace is, even on concrete inputs) and a
    # host-bool conversion would raise. Guard on ``r`` itself, using the
    # shared ``_is_traced`` helper to match the inversion modules.
    if not _is_traced(r) and bool(jnp.any(r <= 0.0)):
        raise ValueError(
            f"Instrument {self.name!r}: total observation variance must be > 0 "
            "(check retrieval / representation / alignment variances)."
        )
    return r

LinearisedRadianceOperator `dataclass` ¶

LinearisedRadianceOperator(baseline: Array, gain: Array)

Tangent-linear radiance operator — the linearised counterpart of the RTM.

The full :class:RadianceObservationOperator is exponential in the column (L_norm = exp(-Δτ)), so it has no closed-form inversion. Linearising the normalised Beer–Lambert forward about the background (ΔVMR = 0, where L_norm = 1) gives the affine response

L_norm(ΔVMR)  ≈  1 + J₀ · ΔVMR,        J₀ = d(L_norm)/d(ΔVMR)|₀,

and since ΔVMR = slope · column_mass is itself linear (the constant slope of :func:column_mass_to_delta_vmr) and the Tier I column is linear in the emission rate Q, the whole radiance forward becomes linear in Q. That makes this the radiance analogue of the scalar averaging kernel: a per-channel linear gain J₀ · slope on the column mass, usable as a drop-in linear observation response (and, with one channel, in the closed-form fusion path the scalar AK uses).

This is the well-known linearised / optimal-estimation radiance model (first-order Beer–Lambert), complementary to the exact nonlinear operator; it is accurate for small enhancements and exact in the optically-thin limit.

Build with :func:linearise from a :class:RadianceObservationOperator so the two share the LUT, geometry and SRF configuration.

Attributes:

Name	Type	Description
`baseline`	`Array`	Background normalised radiance `L_norm(0)`, shape `(n_channels,)` (all ones for the normalised forward; kept explicit so the affine offset is auditable / overridable).
`gain`	`Array`	Per-channel linear response to column mass `d(L_norm)/d(mass)` = `J₀ · slope`, shape `(n_channels,)` [per kg/m²].

n_channels `property` ¶

n_channels: int

Number of output channels per receptor.

predict_single ¶

predict_single(column_mass_kg_m2: Array | float) -> jax.Array

Linearised normalised radiance baseline + gain · mass for one receptor.

Source code in src/plumax/coupled/rtm.py

def predict_single(self, column_mass_kg_m2: jax.Array | float) -> jax.Array:
    """Linearised normalised radiance ``baseline + gain · mass`` for one receptor."""
    m = jnp.asarray(column_mass_kg_m2, dtype=float)
    return self.baseline + self.gain * m

predict ¶

predict(column_enhancement: Array) -> jax.Array

Per-receptor linearised radiance, shape (n_obs, n_channels).

Fully differentiable / vectorised JAX (unlike the nonlinear operator's NumPy LUT path): baseline[None, :] + mass[:, None] · gain[None, :].

Source code in src/plumax/coupled/rtm.py

def predict(self, column_enhancement: jax.Array) -> jax.Array:
    """Per-receptor linearised radiance, shape ``(n_obs, n_channels)``.

    Fully differentiable / vectorised JAX (unlike the nonlinear operator's
    NumPy LUT path): ``baseline[None, :] + mass[:, None] · gain[None, :]``.
    """
    m = jnp.asarray(column_enhancement, dtype=float)
    if m.ndim != 1:
        raise ValueError(
            "LinearisedRadianceOperator.predict: `column_enhancement` must "
            f"be 1-D (n_obs,); got shape {m.shape}."
        )
    return self.baseline[None, :] + m[:, None] * self.gain[None, :]

RadianceObservationOperator `dataclass` ¶

RadianceObservationOperator(ds: Dataset, nu_obs: ndarray, T_K: float, p_atm: float, vmr_background: float, path_length_cm: float, amf: float, srf: SpectralResponseFunction | None = None, var: str = 'absorption_cross_section')

Map a per-receptor column-enhancement field to band radiances via the RTM.

The RTM analogue of the scalar averaging kernel: instead of y = A · column, this composes

column mass [kg/m²]  →  ΔVMR  →  L_norm(ν) = exp(-Δτ)  →  band radiance

using :func:column_mass_to_delta_vmr and the normalised Beer–Lambert forward :func:plumax.radtran.forward_nonlinear_normalized. With srf supplied, the high-resolution normalised radiance exp(-Δτ(ν)) is band-integrated by the instrument :class:~plumax.radtran.srf.SpectralResponseFunction; otherwise the full nu_obs spectrum is returned per receptor.

The forward is NumPy-side (the LUT interpolation is not JAX-traceable), matching :class:plumax.adapters.RadtranObservationOperator; the analytic d(L_norm)/d(ΔVMR) Jacobian is exposed through :meth:linearize.

Attributes:

Name	Type	Description
`ds`	`Dataset`	Absorption-cross-section LUT dataset (HAPI schema; see :func:`plumax.hapi_lut.build_lut_dataset`).
`nu_obs`	`ndarray`	Observation wavenumbers [cm⁻¹], shape `(n_nu,)`.
`T_K`	`float`	Layer temperature [K].
`p_atm`	`float`	Layer pressure [atm].
`vmr_background`	`float`	Background gas volume mixing ratio.
`path_length_cm`	`float`	Optical / absorbing-path length [cm].
`amf`	`float`	Air-mass factor.
`srf`	`SpectralResponseFunction \| None`	Optional band-integration operator. When `None` the operator returns the per-receptor high-resolution normalised radiance.
`var`	`str`	LUT cross-section variable name.

n_channels `property` ¶

n_channels: int

Number of output channels per receptor (bands if SRF, else n_nu).

delta_vmr ¶

delta_vmr(column_mass_kg_m2: Array | float) -> jax.Array

Column mass [kg/m²] → ΔVMR via :func:column_mass_to_delta_vmr.

Source code in src/plumax/coupled/rtm.py

def delta_vmr(self, column_mass_kg_m2: jax.Array | float) -> jax.Array:
    """Column mass [kg/m²] → ΔVMR via :func:`column_mass_to_delta_vmr`."""
    return column_mass_to_delta_vmr(
        column_mass_kg_m2,
        p_atm=self.p_atm,
        T_K=self.T_K,
        path_length_cm=self.path_length_cm,
    )

predict_single ¶

predict_single(column_mass_kg_m2: float) -> np.ndarray

Band-integrated normalised radiance for one receptor.

Parameters:

Name	Type	Description	Default
`column_mass_kg_m2`	`float`	Scalar column mass enhancement [kg/m²].	required

Returns:

Type	Description
`ndarray`	Normalised radiance, shape `(n_channels,)` (`exp(-Δτ)`, band-
`ndarray`	integrated when an SRF is configured).

Source code in src/plumax/coupled/rtm.py

def predict_single(self, column_mass_kg_m2: float) -> np.ndarray:
    """Band-integrated normalised radiance for one receptor.

    Args:
        column_mass_kg_m2: Scalar column mass enhancement [kg/m²].

    Returns:
        Normalised radiance, shape ``(n_channels,)`` (``exp(-Δτ)``, band-
        integrated when an SRF is configured).
    """
    dvmr = float(self.delta_vmr(column_mass_kg_m2))
    return self._band_integrate(self._forward_hr(dvmr).radiance)

predict ¶

predict(column_enhancement: Array | ndarray) -> np.ndarray

Per-receptor band-integrated normalised radiance (vectorised).

Parameters:

Name	Type	Description	Default
`column_enhancement`	`Array \| ndarray`	Per-receptor column mass enhancement [kg/m²], shape `(n_obs,)` (e.g. `Q · column_response`).	required

Returns:

Type	Description
`ndarray`	Per-receptor normalised radiance, shape `(n_obs, n_channels)`.

Source code in src/plumax/coupled/rtm.py

def predict(self, column_enhancement: jax.Array | np.ndarray) -> np.ndarray:
    """Per-receptor band-integrated normalised radiance (vectorised).

    Args:
        column_enhancement: Per-receptor column mass enhancement [kg/m²],
            shape ``(n_obs,)`` (e.g. ``Q · column_response``).

    Returns:
        Per-receptor normalised radiance, shape ``(n_obs, n_channels)``.
    """
    cols = np.asarray(column_enhancement, dtype=float)
    if cols.ndim != 1:
        raise ValueError(
            "RadianceObservationOperator.predict: `column_enhancement` must "
            f"be 1-D (n_obs,); got shape {cols.shape}."
        )
    rows = [self.predict_single(float(c)) for c in cols]
    return np.stack(rows, axis=0)

linearize_single ¶

linearize_single(column_mass_kg_m2: float) -> np.ndarray

Band-integrated d(L_norm)/d(ΔVMR) at one receptor.

Returns the radtran analytic Jacobian, band-integrated through the SRF when configured. This is the tangent-linear w.r.t. ΔVMR; chain with d(ΔVMR)/d(column mass) (the constant slope of :func:column_mass_to_delta_vmr) for the column-space Jacobian.

Parameters:

Name	Type	Description	Default
`column_mass_kg_m2`	`float`	Scalar column mass enhancement [kg/m²].	required

Returns:

Type	Description
`ndarray`	Jacobian, shape `(n_channels,)` (negative for an absorbing band).

Source code in src/plumax/coupled/rtm.py

def linearize_single(self, column_mass_kg_m2: float) -> np.ndarray:
    """Band-integrated ``d(L_norm)/d(ΔVMR)`` at one receptor.

    Returns the radtran analytic Jacobian, band-integrated through the SRF
    when configured. This is the tangent-linear w.r.t. ``ΔVMR``; chain with
    ``d(ΔVMR)/d(column mass)`` (the constant slope of
    :func:`column_mass_to_delta_vmr`) for the column-space Jacobian.

    Args:
        column_mass_kg_m2: Scalar column mass enhancement [kg/m²].

    Returns:
        Jacobian, shape ``(n_channels,)`` (negative for an absorbing band).
    """
    dvmr = float(self.delta_vmr(column_mass_kg_m2))
    return self._band_integrate(self._forward_hr(dvmr).jacobian)

build_coupled_forward ¶

build_coupled_forward(source: PlumeSource, instruments: Sequence[Instrument]) -> CoupledForward

Assemble a :class:CoupledForward, precomputing unit-Q responses.

Parameters:

Name	Type	Description	Default
`source`	`PlumeSource`	Shared source + met configuration.	required
`instruments`	`Sequence[Instrument]`	One or more observing instruments to fuse.	required

Returns:

Type	Description
`CoupledForward`	The configured :class:`CoupledForward`.

Raises:

Type	Description
`ValueError`	if `instruments` is empty.

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

def build_coupled_forward(
    source: PlumeSource, instruments: Sequence[Instrument]
) -> CoupledForward:
    """Assemble a :class:`CoupledForward`, precomputing unit-``Q`` responses.

    Args:
        source: Shared source + met configuration.
        instruments: One or more observing instruments to fuse.

    Returns:
        The configured :class:`CoupledForward`.

    Raises:
        ValueError: if ``instruments`` is empty.
    """
    insts = tuple(instruments)
    if not insts:
        raise ValueError("build_coupled_forward: need at least one instrument.")
    responses = tuple(column_response(source, inst) for inst in insts)
    return CoupledForward(source=source, instruments=insts, responses=responses)

column_response ¶

column_response(source: PlumeSource, instrument: Instrument) -> jax.Array

Per-receptor column enhancement at unit emission for one instrument.

Evaluates the Gaussian plume at unit Q on each receptor column, integrates over source.column_z (trapezoidal), and applies the instrument averaging kernel. The full observation at emission Q is then Q · column_response + bias (plus the background; see :func:predict_observation). Linearity in Q is what makes the coupled inversion closed-form.

Parameters:

Name	Type	Description	Default
`source`	`PlumeSource`	The static source + met configuration.	required
`instrument`	`Instrument`	The observing instrument (receptors + AK).	required

Returns:

Type	Description
`Array`	Unit-`Q` AK-weighted column response, shape `(n_obs,)`.

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

def column_response(source: PlumeSource, instrument: Instrument) -> jax.Array:
    """Per-receptor column enhancement at **unit emission** for one instrument.

    Evaluates the Gaussian plume at unit ``Q`` on each receptor column, integrates
    over ``source.column_z`` (trapezoidal), and applies the instrument averaging
    kernel. The full observation at emission ``Q`` is then
    ``Q · column_response + bias`` (plus the background; see
    :func:`predict_observation`). Linearity in ``Q`` is what makes the coupled
    inversion closed-form.

    Args:
        source: The static source + met configuration.
        instrument: The observing instrument (receptors + AK).

    Returns:
        Unit-``Q`` AK-weighted column response, shape ``(n_obs,)``.
    """
    z = jnp.asarray(source.column_z, dtype=float)
    receptors = jnp.asarray(instrument.receptors, dtype=float)
    n_obs = receptors.shape[0]
    n_z = z.shape[0]

    # Build (n_obs * n_z) receptor points: every (x, y) crossed with every z.
    x = jnp.repeat(receptors[:, 0], n_z)
    y = jnp.repeat(receptors[:, 1], n_z)
    z_tiled = jnp.tile(z, n_obs)

    sx, sy, sz = source.location
    u, v = source.wind_uv()
    dispersion = get_dispersion_params(source.stability_class)

    conc = plume_concentration_vmap(
        x, y, z_tiled, sx, sy, sz, u, v, 1.0, dispersion
    )  # unit Q
    conc = conc.reshape(n_obs, n_z)
    column = jnp.trapezoid(conc, x=z, axis=1)  # (n_obs,)
    return instrument.ak * column

predict_observation ¶

predict_observation(source: PlumeSource, instrument: Instrument, emission_rate: Array | float, bias: Array | float = 0.0) -> jax.Array

Predict one instrument's observation vector y_inst.

y = Q · column_response + background + bias — the affine forward the coupled cost compares against the observed enhancement.

Parameters:

Name	Type	Description	Default
`source`	`PlumeSource`	Static source + met configuration.	required
`instrument`	`Instrument`	The observing instrument.	required
`emission_rate`	`Array \| float`	Emission rate `Q` [kg/s] (scalar).	required
`bias`	`Array \| float`	Per-instrument additive offset (scalar). Default `0`.	`0.0`

Returns:

Type	Description
`Array`	Predicted observation, shape `(n_obs,)`.

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

def predict_observation(
    source: PlumeSource,
    instrument: Instrument,
    emission_rate: jax.Array | float,
    bias: jax.Array | float = 0.0,
) -> jax.Array:
    """Predict one instrument's observation vector ``y_inst``.

    ``y = Q · column_response + background + bias`` — the affine forward the
    coupled cost compares against the observed enhancement.

    Args:
        source: Static source + met configuration.
        instrument: The observing instrument.
        emission_rate: Emission rate ``Q`` [kg/s] (scalar).
        bias: Per-instrument additive offset (scalar). Default ``0``.

    Returns:
        Predicted observation, shape ``(n_obs,)``.
    """
    response = column_response(source, instrument)
    return emission_rate * response + source.background + bias

default_prior ¶

default_prior(*, n_instruments: int, emission_prior_mean: float, emission_prior_std: float, bias_prior_std: float) -> tuple[jax.Array, jax.Array]

Build a diagonal prior (θ_b, B) for (Q, bias₁ … bias_M).

A convenience for the common case: a Gaussian prior on Q and zero-mean Gaussian biases (the design's bias_inst ~ N(0, σ²_inst)).

Parameters:

Name	Type	Description	Default
`n_instruments`	`int`	Number of instruments (bias slots).	required
`emission_prior_mean`	`float`	Prior mean for `Q` [kg/s].	required
`emission_prior_std`	`float`	Prior std for `Q` (`> 0`).	required
`bias_prior_std`	`float`	Prior std for each per-instrument bias (`> 0`).	required

Returns:

Type	Description
`tuple[Array, Array]`	`(prior_mean, prior_covariance)` ready for :func:`fuse_observations`.

Source code in src/plumax/coupled/fusion.py

def default_prior(
    *,
    n_instruments: int,
    emission_prior_mean: float,
    emission_prior_std: float,
    bias_prior_std: float,
) -> tuple[jax.Array, jax.Array]:
    """Build a diagonal prior ``(θ_b, B)`` for ``(Q, bias₁ … bias_M)``.

    A convenience for the common case: a Gaussian prior on ``Q`` and zero-mean
    Gaussian biases (the design's ``bias_inst ~ N(0, σ²_inst)``).

    Args:
        n_instruments: Number of instruments (bias slots).
        emission_prior_mean: Prior mean for ``Q`` [kg/s].
        emission_prior_std: Prior std for ``Q`` (``> 0``).
        bias_prior_std: Prior std for each per-instrument bias (``> 0``).

    Returns:
        ``(prior_mean, prior_covariance)`` ready for :func:`fuse_observations`.
    """
    if emission_prior_std <= 0.0 or bias_prior_std <= 0.0:
        raise ValueError("default_prior: prior standard deviations must be > 0.")
    mean = jnp.concatenate([jnp.array([emission_prior_mean]), jnp.zeros(n_instruments)])
    variances = np.concatenate(
        [[emission_prior_std**2], np.full(n_instruments, bias_prior_std**2)]
    )
    return mean, jnp.diag(jnp.asarray(variances, dtype=float))

fuse_observations ¶

fuse_observations(forward: CoupledForward, observations: Sequence[Array], *, prior_mean: Array, prior_covariance: Array) -> FusionPosterior

Closed-form joint posterior over (Q, bias_inst) from fused observations.

Solves the BLUE update for the affine forward y = G θ + c_bg with a Gaussian prior θ ~ N(θ_b, B) and diagonal observation covariance R (the per-instrument R_retr + R_repr + R_align)::

θ* = θ_b + B Gᵀ (G B Gᵀ + R)⁻¹ (y − G θ_b − c_bg)
P* = B − B Gᵀ (G B Gᵀ + R)⁻¹ G B

solved in the (n_total_obs × n_total_obs) innovation form via Cholesky.

Parameters:

Name	Type	Description	Default
`forward`	`CoupledForward`	The assembled multi-instrument coupled forward.	required
`observations`	`Sequence[Array]`	Observed enhancement vector per instrument, aligned with `forward.instruments` (each `(n_obs_i,)`).	required
`prior_mean`	`Array`	Prior mean state `[Q_b, bias_b₁ … bias_b_M]`, shape `(1 + n_inst,)` (biases typically prior-zero).	required
`prior_covariance`	`Array`	Prior covariance `B`, shape `(1 + n_inst, 1 + n_inst)` — e.g. `diag(σ_Q², σ_bias², …)`.	required

Returns:

Name	Type	Description
`The`	`FusionPosterior`	class:`FusionPosterior`.

Raises:

Type	Description
`ValueError`	on a per-instrument observation-count mismatch or a mis-shaped prior.

Source code in src/plumax/coupled/fusion.py

def fuse_observations(
    forward: CoupledForward,
    observations: Sequence[jax.Array],
    *,
    prior_mean: jax.Array,
    prior_covariance: jax.Array,
) -> FusionPosterior:
    """Closed-form joint posterior over ``(Q, bias_inst)`` from fused observations.

    Solves the BLUE update for the affine forward ``y = G θ + c_bg`` with a
    Gaussian prior ``θ ~ N(θ_b, B)`` and diagonal observation covariance ``R``
    (the per-instrument ``R_retr + R_repr + R_align``)::

        θ* = θ_b + B Gᵀ (G B Gᵀ + R)⁻¹ (y − G θ_b − c_bg)
        P* = B − B Gᵀ (G B Gᵀ + R)⁻¹ G B

    solved in the ``(n_total_obs × n_total_obs)`` innovation form via Cholesky.

    Args:
        forward: The assembled multi-instrument coupled forward.
        observations: Observed enhancement vector per instrument, aligned with
            ``forward.instruments`` (each ``(n_obs_i,)``).
        prior_mean: Prior mean state ``[Q_b, bias_b₁ … bias_b_M]``, shape
            ``(1 + n_inst,)`` (biases typically prior-zero).
        prior_covariance: Prior covariance ``B``, shape
            ``(1 + n_inst, 1 + n_inst)`` — e.g. ``diag(σ_Q², σ_bias², …)``.

    Returns:
        The :class:`FusionPosterior`.

    Raises:
        ValueError: on a per-instrument observation-count mismatch or a
            mis-shaped prior.
    """
    n_inst = len(forward.instruments)
    n_state = 1 + n_inst
    theta_b = jnp.asarray(prior_mean, dtype=float)
    b = jnp.asarray(prior_covariance, dtype=float)
    if theta_b.shape != (n_state,):
        raise ValueError(
            f"fuse_observations: `prior_mean` must be (1 + n_inst = {n_state},); "
            f"got {theta_b.shape}."
        )
    if b.shape != (n_state, n_state):
        raise ValueError(
            f"fuse_observations: `prior_covariance` must be ({n_state}, {n_state}); "
            f"got {b.shape}."
        )
    if len(observations) != n_inst:
        raise ValueError(
            f"fuse_observations: got {len(observations)} observation vectors for "
            f"{n_inst} instruments."
        )
    obs_blocks = []
    for inst, obs in zip(forward.instruments, observations, strict=True):
        o = jnp.asarray(obs, dtype=float)
        if o.shape != (inst.n_obs,):
            raise ValueError(
                f"fuse_observations: observation for {inst.name!r} must be "
                f"({inst.n_obs},); got {o.shape}."
            )
        obs_blocks.append(o)
    y = jnp.concatenate(obs_blocks)

    g, r, bg = _stack_design(forward)
    innovation = y - (g @ theta_b + bg)

    # Information-form BLUE, solved in the tiny (1 + n_inst)-dim state space.
    # For fusion n_obs ≫ n_state, so the innovation form (n_obs × n_obs) would be
    # large *and* ill-conditioned (G B Gᵀ has rank ≤ n_state, so adding a small
    # R gives a near-singular system). The information form inverts only the
    # (n_state × n_state) posterior precision, which is well-conditioned:
    #     P*⁻¹ = B⁻¹ + Gᵀ R⁻¹ G,   θ* = θ_b + P* Gᵀ R⁻¹ (y − G θ_b − c_bg).
    r_inv = 1.0 / r
    gt_rinv = g.T * r_inv[None, :]  # Gᵀ R⁻¹, (n_state, n_obs)
    precision = jnp.linalg.inv(b) + gt_rinv @ g  # B⁻¹ + Gᵀ R⁻¹ G
    precision = 0.5 * (precision + precision.T)
    posterior_cov = jnp.linalg.inv(precision)
    posterior_cov = 0.5 * (posterior_cov + posterior_cov.T)
    increment = posterior_cov @ (gt_rinv @ innovation)
    theta = theta_b + increment

    std = jnp.sqrt(jnp.clip(jnp.diag(posterior_cov), 0.0, None))
    names = tuple(inst.name for inst in forward.instruments)
    return FusionPosterior(
        emission_rate=theta[0],
        emission_std=std[0],
        biases=theta[1:],
        bias_std=std[1:],
        mean=theta,
        covariance=posterior_cov,
        instrument_names=names,
    )

column_mass_to_delta_vmr ¶

column_mass_to_delta_vmr(column_mass_kg_m2: Array | float, *, p_atm: float, T_K: float, path_length_cm: float) -> jax.Array

Convert a plume column mass enhancement [kg/m²] to a gas ΔVMR.

The Tier I plume column integral gives a column mass enhancement m [kg/m²] of methane sitting above the pixel. The radtran Beer–Lambert forward instead parameterises absorption by the enhancement in volume mixing ratio ΔVMR over the absorbing path. The two are related by the column number densities of the gas and of dry air:

Gas column number density [molecules/cm²]::

N_gas = m · (N_A / M_CH₄) / CM2_PER_M2

(N_A = Avogadro, M_CH₄ = methane molar mass; dividing by CM2_PER_M2 = 1e4 converts the per-m² areal density to per-cm².)

Air column number density [molecules/cm²]::

N_air = n(p, T) · L

where n(p, T) is the ideal-gas number density [molecules/cm³] (:func:plumax.radtran.config.number_density_cm3) and L is the vertical path length [cm] over which the enhancement is distributed.

Volume mixing ratio enhancement (dimensionless)::

ΔVMR = N_gas / N_air

This is the standard "slant/vertical column density divided by the air column" relation used in trace-gas retrievals (cf. the partial-column-to-VMR conversion in e.g. TROPOMI / matched-filter retrievals); keeping L and (p, T) explicit makes the per-instrument geometry auditable.

Parameters:

Name	Type	Description	Default
`column_mass_kg_m2`	`Array \| float`	Column mass enhancement [kg/m²] (scalar or array; the Tier I plume column response is linear in the emission rate `Q`).	required
`p_atm`	`float`	Layer pressure [atm], `> 0`.	required
`T_K`	`float`	Layer temperature [K], `> 0`.	required
`path_length_cm`	`float`	Vertical absorbing-path length [cm], `> 0` — the depth over which the column enhancement is spread (use the same value passed to the radtran forward for consistency).	required

Returns:

Type	Description
`Array`	The dimensionless `ΔVMR` enhancement, same shape as
`Array`	`column_mass_kg_m2`.

Raises:

Type	Description
`ValueError`	if `p_atm`, `T_K` or `path_length_cm` is non-positive (eager check, skipped under JAX tracing via `_is_traced`).

Source code in src/plumax/coupled/rtm.py

def column_mass_to_delta_vmr(
    column_mass_kg_m2: jax.Array | float,
    *,
    p_atm: float,
    T_K: float,
    path_length_cm: float,
) -> jax.Array:
    """Convert a plume column **mass** enhancement [kg/m²] to a gas ``ΔVMR``.

    The Tier I plume column integral gives a column **mass enhancement**
    ``m`` [kg/m²] of methane sitting above the pixel. The radtran Beer–Lambert
    forward instead parameterises absorption by the *enhancement in volume
    mixing ratio* ``ΔVMR`` over the absorbing path. The two are related by the
    column number densities of the gas and of dry air:

    1. **Gas column number density** [molecules/cm²]::

           N_gas = m · (N_A / M_CH₄) / CM2_PER_M2

       (``N_A`` = Avogadro, ``M_CH₄`` = methane molar mass; dividing by
       ``CM2_PER_M2 = 1e4`` converts the per-m² areal density to per-cm².)

    2. **Air column number density** [molecules/cm²]::

           N_air = n(p, T) · L

       where ``n(p, T)`` is the ideal-gas number density [molecules/cm³]
       (:func:`plumax.radtran.config.number_density_cm3`) and ``L`` is the
       vertical path length [cm] over which the enhancement is distributed.

    3. **Volume mixing ratio enhancement** (dimensionless)::

           ΔVMR = N_gas / N_air

    This is the standard "slant/vertical column density divided by the air
    column" relation used in trace-gas retrievals (cf. the partial-column-to-VMR
    conversion in e.g. TROPOMI / matched-filter retrievals); keeping ``L`` and
    ``(p, T)`` explicit makes the per-instrument geometry auditable.

    Args:
        column_mass_kg_m2: Column mass enhancement [kg/m²] (scalar or array; the
            Tier I plume column response is linear in the emission rate ``Q``).
        p_atm: Layer pressure [atm], ``> 0``.
        T_K: Layer temperature [K], ``> 0``.
        path_length_cm: Vertical absorbing-path length [cm], ``> 0`` — the depth
            over which the column enhancement is spread (use the same value
            passed to the radtran forward for consistency).

    Returns:
        The dimensionless ``ΔVMR`` enhancement, same shape as
        ``column_mass_kg_m2``.

    Raises:
        ValueError: if ``p_atm``, ``T_K`` or ``path_length_cm`` is non-positive
            (eager check, skipped under JAX tracing via ``_is_traced``).
    """
    # Eager geometry validation, skipped under jit/grad/vmap (the scalars could
    # be tracers in a differentiable pipeline) — mirrors the inversion modules.
    if not _is_traced(p_atm, T_K, path_length_cm) and (
        p_atm <= 0.0 or T_K <= 0.0 or path_length_cm <= 0.0
    ):
        raise ValueError(
            "column_mass_to_delta_vmr: `p_atm`, `T_K` and `path_length_cm` "
            f"must be > 0 (got p_atm={p_atm!r}, T_K={T_K!r}, "
            f"path_length_cm={path_length_cm!r})."
        )
    m = jnp.asarray(column_mass_kg_m2, dtype=float)
    # Gas column number density [molecules/cm²].
    n_gas_cm2 = m * (AVOGADRO_PER_MOL / CH4_MOLAR_MASS_KG_PER_MOL) / CM2_PER_M2
    # Air number density [molecules/cm³] — the ideal-gas relation inlined with
    # ``jnp`` (rather than the NumPy ``radtran.config.number_density_cm3``) so a
    # traced ``p_atm`` / ``T_K`` flows through differentiably, honouring the
    # pure-JAX contract; identical formula and constants.
    p_atm_j = jnp.asarray(p_atm, dtype=float)
    T_K_j = jnp.asarray(T_K, dtype=float)
    n_air_cm3 = (p_atm_j * ATM_TO_PA) / (BOLTZMANN_J_PER_K * T_K_j) / CM3_PER_M3
    # Air column number density [molecules/cm²].
    n_air_cm2 = n_air_cm3 * jnp.asarray(path_length_cm, dtype=float)
    return n_gas_cm2 / n_air_cm2

linearise ¶

linearise(operator: RadianceObservationOperator) -> LinearisedRadianceOperator

Build the :class:LinearisedRadianceOperator tangent to operator at 0.

Evaluates the background radiance L_norm(0) and the analytic Jacobian J₀ = d(L_norm)/d(ΔVMR)|₀ (both band-integrated through the operator's SRF), and folds in the constant d(ΔVMR)/d(mass) slope of :func:column_mass_to_delta_vmr to give the per-channel gain in column-mass units. The result is exact to first order in the enhancement.

Parameters:

Name	Type	Description	Default
`operator`	`RadianceObservationOperator`	The configured nonlinear RTM observation operator.	required

Returns:

Type	Description
`LinearisedRadianceOperator`	The matching :class:`LinearisedRadianceOperator`.

Source code in src/plumax/coupled/rtm.py

def linearise(operator: RadianceObservationOperator) -> LinearisedRadianceOperator:
    """Build the :class:`LinearisedRadianceOperator` tangent to ``operator`` at 0.

    Evaluates the background radiance ``L_norm(0)`` and the analytic Jacobian
    ``J₀ = d(L_norm)/d(ΔVMR)|₀`` (both band-integrated through the operator's
    SRF), and folds in the constant ``d(ΔVMR)/d(mass)`` slope of
    :func:`column_mass_to_delta_vmr` to give the per-channel gain in column-mass
    units. The result is exact to first order in the enhancement.

    Args:
        operator: The configured nonlinear RTM observation operator.

    Returns:
        The matching :class:`LinearisedRadianceOperator`.
    """
    baseline = jnp.asarray(operator.predict_single(0.0), dtype=float)
    jac_dvmr = jnp.asarray(operator.linearize_single(0.0), dtype=float)
    # d(ΔVMR)/d(mass): the conversion is linear, so its slope is ΔVMR at unit mass.
    slope = float(operator.delta_vmr(1.0))
    return LinearisedRadianceOperator(baseline=baseline, gain=jac_dvmr * slope)

radiance_response ¶

radiance_response(source: PlumeSource, instrument: Instrument, operator: RadianceObservationOperator, emission_rate: Array | float) -> np.ndarray

Per-receptor band radiances for one instrument under the RTM operator.

The radiance-forward parallel to :func:plumax.coupled.forward.predict_observation: it evaluates the Tier I unit-Q column response (:func:plumax.coupled.forward.column_response), scales by the emission rate to get the per-receptor column mass enhancement Q · column_response [kg/m²], then maps each receptor through the RTM observation operator to a band-integrated normalised radiance.

Unlike the scalar-AK forward this is not linear in Q (Beer–Lambert is exponential in the column), so it has no closed-form fusion path; it is intended for radiance simulation and as the forward of a gradient-based / linearised L1 inversion (future build-order step), leaving the existing closed-form :func:plumax.coupled.fuse_observations AK path intact.

Parameters:

Name	Type	Description	Default
`source`	`PlumeSource`	Static source + met configuration.	required
`instrument`	`Instrument`	Observing instrument (receptors + AK). The RTM operator is the physical replacement for the scalar AK, so for a pure-RTM forward leave the instrument's `averaging_kernel` at its identity default; a non-identity AK is still applied to the column before the RTM (it then acts as a per-receptor column-scaling factor).	required
`operator`	`RadianceObservationOperator`	The configured :class:`RadianceObservationOperator`.	required
`emission_rate`	`Array \| float`	Emission rate `Q` [kg/s] (scalar).	required

Returns:

Type	Description
`ndarray`	Per-receptor normalised radiance, shape `(n_obs, n_channels)`.

Source code in src/plumax/coupled/rtm.py

def radiance_response(
    source: PlumeSource,
    instrument: Instrument,
    operator: RadianceObservationOperator,
    emission_rate: jax.Array | float,
) -> np.ndarray:
    """Per-receptor band radiances for one instrument under the RTM operator.

    The radiance-forward parallel to :func:`plumax.coupled.forward.predict_observation`:
    it evaluates the Tier I unit-``Q`` column response
    (:func:`plumax.coupled.forward.column_response`), scales by the emission rate
    to get the per-receptor column **mass** enhancement
    ``Q · column_response`` [kg/m²], then maps each receptor through the RTM
    observation operator to a band-integrated normalised radiance.

    Unlike the scalar-AK forward this is **not** linear in ``Q`` (Beer–Lambert
    is exponential in the column), so it has no closed-form fusion path; it is
    intended for radiance simulation and as the forward of a gradient-based /
    linearised L1 inversion (future build-order step), leaving the existing
    closed-form :func:`plumax.coupled.fuse_observations` AK path intact.

    Args:
        source: Static source + met configuration.
        instrument: Observing instrument (receptors + AK). The RTM operator is
            the physical replacement for the scalar AK, so for a pure-RTM forward
            leave the instrument's ``averaging_kernel`` at its identity default;
            a non-identity AK is still applied to the column before the RTM (it
            then acts as a per-receptor column-scaling factor).
        operator: The configured :class:`RadianceObservationOperator`.
        emission_rate: Emission rate ``Q`` [kg/s] (scalar).

    Returns:
        Per-receptor normalised radiance, shape ``(n_obs, n_channels)``.
    """
    # Unit-Q column response is differentiable JAX; scale to the column mass and
    # hand concrete values to the NumPy-side RTM forward.
    response = column_response(source, instrument)
    column_mass = np.asarray(emission_rate, dtype=float) * np.asarray(
        response, dtype=float
    )
    return operator.predict(column_mass)

radiance_response_linear ¶

radiance_response_linear(source: PlumeSource, instrument: Instrument, operator: LinearisedRadianceOperator, emission_rate: Array | float) -> jax.Array

Per-receptor linearised radiances for one instrument (differentiable).

The tangent-linear parallel of :func:radiance_response: it scales the Tier I unit-Q column response by Q and maps each receptor through the :class:LinearisedRadianceOperator. Unlike :func:radiance_response this is linear in Q and stays in pure JAX end-to-end, so it is jax.grad / jax.jit-friendly and (single-channel) feeds the closed-form :func:plumax.coupled.fuse_observations path.

Parameters:

Name	Type	Description	Default
`source`	`PlumeSource`	Static source + met configuration.	required
`instrument`	`Instrument`	Observing instrument (receptors + AK).	required
`operator`	`LinearisedRadianceOperator`	The linearised radiance operator (from :func:`linearise`).	required
`emission_rate`	`Array \| float`	Emission rate `Q` [kg/s].	required

Returns:

Type	Description
`Array`	Per-receptor linearised normalised radiance, shape `(n_obs, n_channels)`.

Source code in src/plumax/coupled/rtm.py

def radiance_response_linear(
    source: PlumeSource,
    instrument: Instrument,
    operator: LinearisedRadianceOperator,
    emission_rate: jax.Array | float,
) -> jax.Array:
    """Per-receptor linearised radiances for one instrument (differentiable).

    The tangent-linear parallel of :func:`radiance_response`: it scales the
    Tier I unit-``Q`` column response by ``Q`` and maps each receptor through the
    :class:`LinearisedRadianceOperator`. Unlike :func:`radiance_response` this is
    **linear in ``Q``** and stays in pure JAX end-to-end, so it is
    ``jax.grad`` / ``jax.jit``-friendly and (single-channel) feeds the
    closed-form :func:`plumax.coupled.fuse_observations` path.

    Args:
        source: Static source + met configuration.
        instrument: Observing instrument (receptors + AK).
        operator: The linearised radiance operator (from :func:`linearise`).
        emission_rate: Emission rate ``Q`` [kg/s].

    Returns:
        Per-receptor linearised normalised radiance, shape ``(n_obs, n_channels)``.
    """
    response = column_response(source, instrument)
    column_mass = jnp.asarray(emission_rate, dtype=float) * response
    return operator.predict(column_mass)

Coupled inversion¶

coupled ¶

CoupledForward dataclass ¶

predict ¶

PlumeSource dataclass ¶

wind_uv ¶

FusionPosterior dataclass ¶

Instrument dataclass ¶

n_obs property ¶

ak property ¶

observation_variance ¶

LinearisedRadianceOperator dataclass ¶

n_channels property ¶

predict_single ¶

predict ¶

RadianceObservationOperator dataclass ¶

n_channels property ¶

delta_vmr ¶

predict_single ¶

predict ¶

linearize_single ¶

build_coupled_forward ¶

column_response ¶

predict_observation ¶

default_prior ¶

fuse_observations ¶

column_mass_to_delta_vmr ¶

linearise ¶

radiance_response ¶

radiance_response_linear ¶

CoupledForward `dataclass` ¶

PlumeSource `dataclass` ¶

FusionPosterior `dataclass` ¶

Instrument `dataclass` ¶

n_obs `property` ¶

ak `property` ¶

LinearisedRadianceOperator `dataclass` ¶

n_channels `property` ¶

RadianceObservationOperator `dataclass` ¶

n_channels `property` ¶