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:radtrax.slab 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. LayeredRadianceOperator maps the column through radtrax's layered ForwardModel instead of the slab, in JAX end to end, so jax.grad with respect to Q runs through plume and radiance together. The scalar-AK path above is unchanged.

• backends — the Tier II and Tier III transports behind the same forward (:class:TransportBackend, :class:LagrangianSource, :class:EulerianSource), through the shared column observation.

• quality — the per-instrument quality-flag aggregator (:class:QualityPolicy, :func:apply_quality, :func:contribution_report): flagged pixels are masked or have their variance inflated, and :func:stack_observations hands the result to vardax.

• inversion — the nonlinear end-to-end inversion (:func:solve_coupled) over the emission, source position, wind, background and biases, through the averaging kernel or a radiance operator per instrument: plumax builds the forward, prior and masks, vardax solves (ThreeDVar, L-BFGS) and gives the Laplace posterior. Needs the assim extra.

• k_sources — masked-K multi-source handling across instruments (:class:MaskedKPrior, :func:infer_masked_k_sources, :func:build_masked_k_problem, :func:saturation_rate), on Tier I's slot prior and posterior carrier.

Later build-order steps (the Q(t) stochastic process, 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.

EulerianSource dataclass

EulerianSource(forward: EulerianForward4DVar, background: float = 0.0)

Tier III transport: the les_fvm finite-volume run through the shared column.

r_i = A_i Σ_k Δz c_k(x_i, t_inst; Q = 1),

the field of a unit constant emission at the overpass time t_inst, read in the T-cell that holds each receptor. The forward is linear in the emission with a zero initial field, so Q · r is the run at constant emission Q.

Attributes:

Name	Type	Description
forward	EulerianForward4DVar	The differentiable transport (les_fvm.build_forward); its save times must include each instrument's overpass time.
background	float	Additive background column c_bg.

unit_history cached property

unit_history: Float[Array, 'time level y x']

Interior field history of a unit constant emission, one solve.

plume_grid property

plume_grid: PlumeGrid3D

The analysis grid the transport runs on.

time_index

time_index(instrument: Instrument) -> int

Index of the save time that equals the instrument's overpass time.

Raises:

Type	Description
GridMismatchError	If no save time matches it.

Source code in src/plumax/coupled/_src/backends.py

def time_index(self, instrument: Instrument) -> int:
    """Index of the save time that equals the instrument's overpass time.

    Raises:
        GridMismatchError: If no save time matches it.
    """
    save_times = np.asarray(self.forward.save_times, dtype=float)
    matches = np.flatnonzero(np.isclose(save_times, instrument.observation_time))
    if matches.size == 0:
        raise GridMismatchError(
            f"EulerianSource: overpass time {instrument.observation_time} s of "
            f"{instrument.name!r} is not one of the forward's save times; "
            "build the forward with it in `save_times`."
        )
    return int(matches[0])

column_response

column_response(instrument: Instrument) -> Float[Array, ' obs']

Column per unit emission at the instrument's receptors.

Parameters:

Name	Type	Description	Default
instrument	Instrument	The observing instrument.	required

Returns:

Type	Description
Float[Array, ' obs']	The response r_inst.

Raises:

Type	Description
GridMismatchError	If the overpass time is not a save time.
DomainError	If a receptor lies outside the grid.

Source code in src/plumax/coupled/_src/backends.py

def column_response(self, instrument: Instrument) -> Float[Array, " obs"]:
    """Column per unit emission at the instrument's receptors.

    Args:
        instrument: The observing instrument.

    Returns:
        The response r_inst.

    Raises:
        GridMismatchError: If the overpass time is not a save time.
        DomainError: If a receptor lies outside the grid.
    """
    index = self.time_index(instrument)
    observation = ColumnObservation.for_grid(self.plume_grid, instrument=instrument)
    return observation.from_field(self.unit_history[index])

LagrangianSource dataclass

LagrangianSource(emission_pattern: Float[Array, ' cell'] | Float[Array, 'cell param'], winds: tuple[GriddedWind | Callable[[Array], Array], ...], turbulence: TurbulenceModel, domain_x: tuple[float, float, int], domain_y: tuple[float, float, int], column_z: Float[Array, ' level'], n_particles: int = 1000, t_back: float = 600.0, dt: float = 1.0, pbl_height: float = 1000.0, pbl_fraction: float = 0.5, air_density: float = 1.2, seed: int = 0, background: float = 0.0)

Tier II transport: met-ensemble backward footprints through the shared column.

c_k(x_i) = ρ_air F̄(x_i, z_k) q, F̄ = mean_e F_e, r_i = A_i Σ_k h_k c_k(x_i),

with F_e the footprint of member e at receptor (x_i, z_k), q the surface flux of a unit source and h_k the trapezoid weights over column_z.

Attributes:

Name	Type	Description
emission_pattern	Float[Array, ' cell'] | Float[Array, 'cell param']	Surface flux per unit emission q [m⁻²] on the footprint grid, cells flattened row-major over (x, y). A second axis holds one pattern per emission parameter.
winds	tuple[GriddedWind | Callable[[Array], Array], ...]	One mean wind per met-ensemble member (a GriddedWind or a t -> (u, v, w) callable). One member is a deterministic run.
turbulence	TurbulenceModel	Turbulence model shared by the members.
domain_x	tuple[float, float, int]	(start, stop, n_cells) of the footprint grid along x.
domain_y	tuple[float, float, int]	(start, stop, n_cells) of the footprint grid along y.
column_z	Float[Array, ' level']	Receptor levels z_k [m] of the column, strictly increasing.
n_particles	int	Particles per receptor level and member.
t_back	float	Backward integration time [s].
dt	float	Particle time step [s].
pbl_height	float	Boundary-layer height [m].
pbl_fraction	float	Fraction of the boundary layer that sees the surface flux.
air_density	float	Air density ρ_air [kg m⁻³].
seed	int	Base PRNG seed of the particle draws.
background	float	Additive background column c_bg.

point classmethod

point(location: tuple[float, float], *, winds: Sequence[GriddedWind | Callable[[Array], Array]], turbulence: TurbulenceModel, domain_x: tuple[float, float, int], domain_y: tuple[float, float, int], column_z: Float[Array, ' level'], **settings: float | int) -> LagrangianSource

A point source: unit emission spread over the cell that holds it.

q_s = 1 / A_cell in the containing cell s, zero elsewhere.

Parameters:

Name	Type	Description	Default
location	tuple[float, float]	Source (x, y) [m].	required
winds	Sequence[GriddedWind | Callable[[Array], Array]]	One mean wind per met member.	required
turbulence	TurbulenceModel	Turbulence model.	required
domain_x	tuple[float, float, int]	(start, stop, n_cells) along x.	required
domain_y	tuple[float, float, int]	(start, stop, n_cells) along y.	required
column_z	Float[Array, ' level']	Receptor levels z_k [m].	required
**settings	float | int	Any other field of :class:LagrangianSource.	{}

Returns:

Type	Description
LagrangianSource	The source, with a one-cell emission pattern.

Raises:

Type	Description
DomainError	If the location lies outside the footprint grid.

Source code in src/plumax/coupled/_src/backends.py

@classmethod
def point(
    cls,
    location: tuple[float, float],
    *,
    winds: Sequence[GriddedWind | Callable[[jax.Array], jax.Array]],
    turbulence: TurbulenceModel,
    domain_x: tuple[float, float, int],
    domain_y: tuple[float, float, int],
    column_z: Float[Array, " level"],
    **settings: float | int,
) -> LagrangianSource:
    """A point source: unit emission spread over the cell that holds it.

    q_s = 1 / A_cell in the containing cell s, zero elsewhere.

    Args:
        location: Source ``(x, y)`` [m].
        winds: One mean wind per met member.
        turbulence: Turbulence model.
        domain_x: ``(start, stop, n_cells)`` along x.
        domain_y: ``(start, stop, n_cells)`` along y.
        column_z: Receptor levels z_k [m].
        **settings: Any other field of :class:`LagrangianSource`.

    Returns:
        The source, with a one-cell emission pattern.

    Raises:
        DomainError: If the location lies outside the footprint grid.
    """
    x_edges = np.linspace(domain_x[0], domain_x[1], int(domain_x[2]) + 1)
    y_edges = np.linspace(domain_y[0], domain_y[1], int(domain_y[2]) + 1)
    ix = int(np.searchsorted(x_edges, location[0], side="right")) - 1
    iy = int(np.searchsorted(y_edges, location[1], side="right")) - 1
    if not (0 <= ix < x_edges.size - 1 and 0 <= iy < y_edges.size - 1):
        raise DomainError(
            f"LagrangianSource.point: {location} lies outside the footprint grid."
        )
    # A_cell = Δx Δy
    cell_area = (x_edges[1] - x_edges[0]) * (y_edges[1] - y_edges[0])
    pattern = np.zeros((x_edges.size - 1, y_edges.size - 1))
    # q_s = 1 / A_cell in the containing cell
    pattern[ix, iy] = 1.0 / cell_area
    return cls(
        emission_pattern=jnp.asarray(pattern.reshape(-1)),
        winds=tuple(winds),
        turbulence=turbulence,
        domain_x=domain_x,
        domain_y=domain_y,
        column_z=jnp.asarray(column_z),
        **settings,
    )

footprints

footprints(instrument: Instrument) -> Float[Array, 'member point cell']

Member footprints F_e at every receptor level, level fastest.

Parameters:

Name	Type	Description	Default
instrument	Instrument	The observing instrument; its overpass time is the receptor time of the backward runs.	required

Returns:

Type	Description
Float[Array, 'member point cell']	One row per (receptor, level) pair, units s m² kg⁻¹.

Source code in src/plumax/coupled/_src/backends.py

def footprints(self, instrument: Instrument) -> Float[Array, "member point cell"]:
    """Member footprints F_e at every receptor level, level fastest.

    Args:
        instrument: The observing instrument; its overpass time is the
            receptor time of the backward runs.

    Returns:
        One row per (receptor, level) pair, units s m² kg⁻¹.
    """
    levels = np.asarray(self.column_z, dtype=float)
    receptors = np.asarray(instrument.receptors, dtype=float)
    # every receptor (x_i, y_i) at every level z_k, level fastest
    horizontal = np.repeat(receptors, levels.size, axis=0)
    heights = np.tile(levels, receptors.shape[0])
    points = np.column_stack([horizontal, heights])
    return ensemble_footprints(
        list(self.winds),
        jnp.asarray(points),
        self.turbulence,
        self.domain_x,
        self.domain_y,
        n_particles=self.n_particles,
        t_back=self.t_back,
        dt=self.dt,
        pbl_height=self.pbl_height,
        pbl_fraction=self.pbl_fraction,
        air_density=self.air_density,
        receptor_time=instrument.observation_time,
        seed=self.seed,
    )

column_response

column_response(instrument: Instrument) -> Float[Array, ' obs'] | Float[Array, 'obs param']

Column per unit emission at the instrument's receptors.

Algorithm

1. Backward footprints from every receptor level, per member.
2. Average over the members: F̄ = mean_e F_e.
3. Concentration of a unit source: c = ρ_air F̄ q.
4. Column and averaging kernel: r_i = A_i Σ_k h_k c_k(x_i).

Parameters:

Name	Type	Description	Default
instrument	Instrument	The observing instrument.	required

Returns:

Type	Description
Float[Array, ' obs'] | Float[Array, 'obs param']	The response r_inst, one column per emission pattern when the
Float[Array, ' obs'] | Float[Array, 'obs param']	pattern has a parameter axis.

Source code in src/plumax/coupled/_src/backends.py

def column_response(
    self, instrument: Instrument
) -> Float[Array, " obs"] | Float[Array, "obs param"]:
    """Column per unit emission at the instrument's receptors.

    Algorithm:
        1. Backward footprints from every receptor level, per member.
        2. Average over the members: F̄ = mean_e F_e.
        3. Concentration of a unit source: c = ρ_air F̄ q.
        4. Column and averaging kernel: r_i = A_i Σ_k h_k c_k(x_i).

    Args:
        instrument: The observing instrument.

    Returns:
        The response r_inst, one column per emission pattern when the
        pattern has a parameter axis.
    """
    n_levels = int(np.shape(self.column_z)[0])
    footprints = self.footprints(instrument)
    # F̄ = mean_e F_e
    mean_footprint = einx.mean("member point cell -> point cell", footprints)
    pattern = jnp.asarray(self.emission_pattern, dtype=mean_footprint.dtype)
    observation = ColumnObservation.trapezoid(
        jnp.asarray(self.column_z, dtype=mean_footprint.dtype),
        averaging_kernel=instrument.ak,
    )
    if pattern.ndim == 1:
        # F̄ q
        mixing_ratio = einx.dot(
            "point cell, cell -> point", mean_footprint, pattern
        )
        # c = ρ_air F̄ q
        concentration = self.air_density * mixing_ratio
        profiles = einx.rearrange(
            "(obs level) -> obs level", concentration, level=n_levels
        )
        return observation.column(profiles)
    # F̄ q_p, one per pattern
    mixing_ratio = einx.dot(
        "point cell, cell param -> point param", mean_footprint, pattern
    )
    # c_p = ρ_air F̄ q_p
    concentration = self.air_density * mixing_ratio
    profiles = einx.rearrange(
        "(obs level) param -> param obs level", concentration, level=n_levels
    )
    columns = jax.vmap(observation.column)(profiles)
    return einx.rearrange("param obs -> obs param", columns)

TransportBackend

Bases: Protocol

A transport tier as the coupled forward sees it.

r_inst = A_inst col_z(c(·; Q = 1)), the column per unit emission at each of the instrument's receptors.

Attributes:

Name	Type	Description
background	float	Additive background column c_bg, in the units of the observation.

column_response

column_response(instrument: Instrument) -> Float[Array, ' obs'] | Float[Array, 'obs param']

Column per unit emission at the instrument's receptors.

Parameters:

Name	Type	Description	Default
instrument	Instrument	The observing instrument (receptors and kernel).	required

Returns:

Type	Description
Float[Array, ' obs'] | Float[Array, 'obs param']	The response r_inst: one value per receptor for a scalar emission
Float[Array, ' obs'] | Float[Array, 'obs param']	rate, or one column per emission parameter.

Source code in src/plumax/coupled/_src/backends.py

def column_response(
    self, instrument: Instrument
) -> Float[Array, " obs"] | Float[Array, "obs param"]:
    """Column per unit emission at the instrument's receptors.

    Args:
        instrument: The observing instrument (receptors and kernel).

    Returns:
        The response r_inst: one value per receptor for a scalar emission
        rate, or one column per emission parameter.
    """
    ...

CoupledForward dataclass

CoupledForward(source: PlumeSource | TransportBackend, 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 | TransportBackend	Shared transport backend (PlumeSource, LagrangianSource or EulerianSource).
instruments	tuple[Instrument, ...]	The observing instruments (fusion list).
responses	tuple[Array, ...]	Per-instrument unit-Q column responses (one (n_obs,) array each, or an (n_obs, n_param) design matrix for a multi-parameter emission), 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], or the (n_param,) emission vector when the responses are design matrices.	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], or the ``(n_param,)``
            emission vector when the responses are design matrices.
        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 [
        _scale_response(resp, emission_rate) + 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)

column_response

column_response(instrument: Instrument) -> jax.Array

Unit-Q column response, the Tier I :class:TransportBackend method.

Delegates to the module-level :func:column_response.

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

def column_response(self, instrument: Instrument) -> jax.Array:
    """Unit-``Q`` column response, the Tier I :class:`TransportBackend` method.

    Delegates to the module-level :func:`column_response`.
    """
    return column_response(self, instrument)

FusionPosterior dataclass

FusionPosterior(emission_rate: Array, emission_std: Array, biases: Array, bias_std: Array, mean: Array, covariance: Array, instrument_names: tuple[str, ...], pixels_used: Array | None = None, prior_mean: Array | None = None, prior_covariance: Array | None = None)

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.

With design-matrix responses (n_param emission parameters, e.g. one per cell of a Tier II source) the emission fields are (n_param,) vectors and the state is [q₁ … q_P, bias₁ … bias_M].

Attributes:

Name	Type	Description
emission_rate	Array	Posterior mean emission Q* [kg/s], 0-d array ((n_param,) for design-matrix responses).
emission_std	Array	Posterior standard deviation of Q, 0-d array ((n_param,) for design-matrix responses).
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 (n_param + n_inst,).
covariance	Array	Full posterior covariance P*, shape (n_param + n_inst, n_param + n_inst).
instrument_names	tuple[str, ...]	Names aligned with the bias entries.
pixels_used	Array | None	Pixels each instrument contributed at full weight (unmasked, not inflated), shape (n_inst,), aligned with instrument_names.
prior_mean	Array | None	Prior mean state θ_b the fusion used (None when built by hand).
prior_covariance	Array | None	Prior covariance B (None when built by hand).

For a scalar Q it also carries the Tier V per-event payload (population.PerEventPosterior): Q's posterior marginal N(Q, P₀₀) and its prior marginal N(θ_b,0, B₀₀) are exact.

samples

samples(key: Array, n: int) -> jax.Array

n exact draws of Q from its posterior marginal N(Q, P₀₀).

Raises:

Type	Description
PayloadError	For design-matrix (vector) emissions.

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

def samples(self, key: jax.Array, n: int) -> jax.Array:
    """n exact draws of Q from its posterior marginal N(Q*, P*₀₀).

    Raises:
        PayloadError: For design-matrix (vector) emissions.
    """
    self._require_scalar("samples")
    return GaussianEmission(self.emission_rate, self.emission_std).sample(key, n)

prior_logpdf

prior_logpdf(q: Array) -> jax.Array

log π(Q) of Q's prior marginal N(θ_b,0, B₀₀).

Raises:

Type	Description
PayloadError	For vector emissions, or without the prior.

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

def prior_logpdf(self, q: jax.Array) -> jax.Array:
    """log π(Q) of Q's prior marginal N(θ_b,0, B₀₀).

    Raises:
        PayloadError: For vector emissions, or without the prior.
    """
    self._require_scalar("prior_logpdf")
    if self.prior_mean is None or self.prior_covariance is None:
        raise PayloadError(
            "FusionPosterior.prior_logpdf: this posterior carries no prior; "
            "build it with fuse_observations."
        )
    prior_std = jnp.sqrt(self.prior_covariance[0, 0])
    return GaussianEmission(self.prior_mean[0], prior_std).log_prob(q)

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, flags: Array | None = None, mask: Array | None = None, variance_inflation: Array | None = None)

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.
flags	Array | None	Per-pixel quality bit flags, shape (n_obs,), 0 = clear (see :class:~plumax.coupled.quality.QualityFlag). None is all-clear.
mask	Array | None	Per-pixel observation mask, shape (n_obs,): 1 keeps the pixel, 0 drops it from the likelihood. None keeps all.
variance_inflation	Array | None	Per-pixel factor ≥ 1 on the total observation variance, shape (n_obs,). None is no inflation.

n_obs property

n_obs: int

Number of receptor pixels.

ak property

ak: Array

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

observation_mask

observation_mask() -> jax.Array

Per-pixel mask as a float array, 1 = used (all ones if unset).

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

def observation_mask(self) -> jax.Array:
    """Per-pixel mask as a float array, ``1`` = used (all ones if unset)."""
    if self.mask is None:
        return jnp.ones(self.n_obs)
    return jnp.asarray(self.mask, dtype=float)

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. A variance_inflation multiplies the sum pixel by pixel.

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. A ``variance_inflation`` multiplies the
    sum pixel by pixel.

    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)."
        )
    if self.variance_inflation is not None:
        r = r * jnp.asarray(self.variance_inflation, dtype=float)
    return r

CoupledPrior dataclass

CoupledPrior(mean: Mapping[str, Float[Array, '*shape']], covariance: Float[Array, ' param'] | Float[Array, 'param param'])

Gaussian prior N(μ_θ, Σ_θ) on the stacked control, by named block.

Attributes:

Name	Type	Description
mean	Mapping[str, Float[Array, '*shape']]	Prior mean μ_θ of each block. Exactly one of emission_rate and log_emission_rate (a scalar, or one value per emission parameter); optionally location (x, y) [m], wind (speed [m/s], direction [deg from North]), background (scalar) and biases (one per instrument).
covariance	Float[Array, ' param'] | Float[Array, 'param param']	Σ_θ over the flattened control, in :attr:layout order: a vector of variances (diagonal) or a dense matrix.

Raises:

Type	Description
ValueError	On an unknown or missing block, a mis-shaped block or a covariance that does not match the control.

size property

size: int

Length D_θ of the flattened control.

layout property

layout: tuple[tuple[str, int], ...]

(block, length) pairs in the flattened order (sorted by name).

log_emission property

log_emission: bool

Whether the emission is lognormal (controlled through log Q).

from_std classmethod

from_std(mean: Mapping[str, Float[Array, '*shape'] | float], std: Mapping[str, Float[Array, '*shape'] | float]) -> CoupledPrior

A diagonal prior from per-block means and standard deviations.

Σ_θ = diag(σ²), each block's σ broadcast to the block's shape.

Parameters:

Name	Type	Description	Default
mean	Mapping[str, Float[Array, '*shape'] | float]	Prior mean μ_θ of each block.	required
std	Mapping[str, Float[Array, '*shape'] | float]	Prior standard deviation σ of each block, same keys.	required

Returns:

Type	Description
CoupledPrior	The prior.

Raises:

Type	Description
ValueError	If the keys differ or a deviation is not positive.

Source code in src/plumax/coupled/_src/inversion.py

@classmethod
def from_std(
    cls,
    mean: Mapping[str, Float[Array, "*shape"] | float],
    std: Mapping[str, Float[Array, "*shape"] | float],
) -> CoupledPrior:
    """A diagonal prior from per-block means and standard deviations.

    Σ_θ = diag(σ²), each block's σ broadcast to the block's shape.

    Args:
        mean: Prior mean μ_θ of each block.
        std: Prior standard deviation σ of each block, same keys.

    Returns:
        The prior.

    Raises:
        ValueError: If the keys differ or a deviation is not positive.
    """
    if set(mean) != set(std):
        raise ValueError("CoupledPrior.from_std: `mean` and `std` keys differ.")
    means = {name: jnp.asarray(value, dtype=float) for name, value in mean.items()}
    stds = {
        name: jnp.broadcast_to(jnp.asarray(std[name], dtype=float), value.shape)
        for name, value in means.items()
    }
    flat_std, _ = ravel_pytree(stds)
    if not bool(jnp.all(flat_std > 0.0)):
        raise ValueError("CoupledPrior.from_std: standard deviations must be > 0.")
    return cls(mean=means, covariance=flat_std**2)

flatten

flatten() -> tuple[Float[Array, ' param'], Callable[[jax.Array], dict[str, jax.Array]]]

μ_θ as one vector, and the map from such a vector back to blocks.

Source code in src/plumax/coupled/_src/inversion.py

def flatten(
    self,
) -> tuple[Float[Array, " param"], Callable[[jax.Array], dict[str, jax.Array]]]:
    """μ_θ as one vector, and the map from such a vector back to blocks."""
    means = {
        name: jnp.asarray(value, dtype=float) for name, value in self.mean.items()
    }
    return ravel_pytree(means)

cholesky

cholesky() -> Float[Array, 'param param']

L with Σ_θ = L Lᵀ (√σ² on the diagonal for a diagonal prior).

Source code in src/plumax/coupled/_src/inversion.py

def cholesky(self) -> Float[Array, "param param"]:
    """L with Σ_θ = L Lᵀ (√σ² on the diagonal for a diagonal prior)."""
    covariance = jnp.asarray(self.covariance, dtype=float)
    if covariance.ndim == 1:
        # L = diag(σ)
        return jnp.diag(jnp.sqrt(covariance))
    psd = frozenset({lx.symmetric_tag, lx.positive_semidefinite_tag})
    return gx.safe_cholesky(lx.MatrixLinearOperator(covariance, psd))

row_operator

row_operator() -> lx.AbstractLinearOperator

Σ_θ on a (1, D_θ) control row, the shape vardax applies it to.

Source code in src/plumax/coupled/_src/inversion.py

def row_operator(self) -> lx.AbstractLinearOperator:
    """Σ_θ on a ``(1, D_θ)`` control row, the shape vardax applies it to."""
    covariance = jnp.asarray(self.covariance, dtype=float)
    if covariance.ndim == 1:
        return lx.DiagonalLinearOperator(einx.rearrange("p -> 1 p", covariance))
    psd = frozenset({lx.symmetric_tag, lx.positive_semidefinite_tag})
    return lx.FunctionLinearOperator(
        lambda row: einx.dot("row p, p q -> row q", row, covariance),
        jax.ShapeDtypeStruct((1, covariance.shape[0]), covariance.dtype),
        psd,
    )

CoupledProblem dataclass

CoupledProblem(forward: CoupledForward, prior: CoupledPrior, observations: Mapping[str, Float[Array, '*shape']], radiance_ops: Mapping[str, RadianceOperator] = dict(), radiance_variance: Mapping[str, Float[Array, '*shape'] | float] = dict())

The coupled forward, its prior and every instrument's observations.

Column instruments are observed through their averaging kernel (the CoupledForward responses); an instrument named in radiance_ops is observed through that operator instead, applied to its column (bias included). Masks and variance inflation come from each instrument (coupled.quality.apply_quality).

Attributes:

Name	Type	Description
forward	CoupledForward	The coupled forward. location and wind blocks need the Tier I PlumeSource, whose response is recomputed at each control; other backends keep their precomputed responses.
prior	CoupledPrior	Prior on the control.
observations	Mapping[str, Float[Array, '*shape']]	Observations by instrument name: (n_obs,) columns, or (n_obs, n_channels) radiances.
radiance_ops	Mapping[str, RadianceOperator]	JAX radiance operators by instrument name.
radiance_variance	Mapping[str, Float[Array, '*shape'] | float]	Radiance error variance by instrument name, a scalar or (n_obs, n_channels).

Raises:

Type	Description
ValueError	On missing or extra observations, a met block without a PlumeSource, mis-shaped biases or a radiance operator without a variance.
TypeError	For a RadianceObservationOperator, which runs in NumPy and so cannot be differentiated.

responses

responses(control: Mapping[str, Array]) -> list[jax.Array]

Unit-emission response of each instrument at this control.

Source code in src/plumax/coupled/_src/inversion.py

def responses(self, control: Mapping[str, jax.Array]) -> list[jax.Array]:
    """Unit-emission response of each instrument at this control."""
    if not any(name in control for name in _MET_BLOCKS):
        return list(self.forward.responses)
    source = self.forward.source
    assert isinstance(source, PlumeSource)
    location = control.get("location", jnp.asarray(source.location[:2]))
    wind = control.get(
        "wind", jnp.asarray([source.wind_speed, source.wind_direction])
    )
    return [
        plume_response(source, instrument, location, wind)
        for instrument in self.forward.instruments
    ]

predict

predict(control: Mapping[str, Array]) -> dict[str, jax.Array]

h_i(θ) for every instrument, by name.

Parameters:

Name	Type	Description	Default
control	Mapping[str, Array]	The control, by block (as :meth:CoupledPrior.flatten unflattens it).	required

Returns:

Type	Description
dict[str, Array]	Columns (n_obs,), or radiances (n_obs, n_channels).

Source code in src/plumax/coupled/_src/inversion.py

def predict(self, control: Mapping[str, jax.Array]) -> dict[str, jax.Array]:
    """h_i(θ) for every instrument, by name.

    Args:
        control: The control, by block (as :meth:`CoupledPrior.flatten`
            unflattens it).

    Returns:
        Columns ``(n_obs,)``, or radiances ``(n_obs, n_channels)``.
    """
    if "log_emission_rate" in control:
        # A line-search trial far from the data can overflow exp, and
        # lineax raises on the non-finite residual; capping log Q at 50
        # (Q ≈ 5e21 kg/s) keeps the trial cost finite, so the search
        # backtracks instead.
        log_emission = jnp.minimum(control["log_emission_rate"], _LOG_EMISSION_CAP)
        # Q = exp(log Q)
        emission = jnp.exp(log_emission)
    else:
        emission = control["emission_rate"]
    background = control.get(
        "background", jnp.asarray(self.forward.source.background, dtype=float)
    )
    n_instruments = len(self.forward.instruments)
    biases = control.get("biases", jnp.zeros(n_instruments))
    predictions = {}
    for index, (instrument, response) in enumerate(
        zip(self.forward.instruments, self.responses(control), strict=True)
    ):
        # c_i = Q r_i + c_bg + bias_i
        column = _scale_response(response, emission) + background + biases[index]
        operator = self.radiance_ops.get(instrument.name)
        predictions[instrument.name] = (
            column if operator is None else operator.predict(column)
        )
    return predictions

predict_stacked

predict_stacked(control: Mapping[str, Array]) -> jax.Array

Every instrument's prediction, flattened and stacked, shape (D_y,).

Source code in src/plumax/coupled/_src/inversion.py

def predict_stacked(self, control: Mapping[str, jax.Array]) -> jax.Array:
    """Every instrument's prediction, flattened and stacked, shape ``(D_y,)``."""
    predictions = self.predict(control)
    return jnp.concatenate(
        [
            jnp.ravel(predictions[instrument.name])
            for instrument in self.forward.instruments
        ]
    )

stacked_observations

stacked_observations() -> StackedObservations

Observations, masks and Σ_ε stacked like :meth:predict_stacked.

A radiance pixel's mask and variance inflation apply to all its channels. Masked values are zero (jejjohnson/vardax#147).

Source code in src/plumax/coupled/_src/inversion.py

def stacked_observations(self) -> StackedObservations:
    """Observations, masks and Σ_ε stacked like :meth:`predict_stacked`.

    A radiance pixel's mask and variance inflation apply to all its
    channels. Masked values are zero (jejjohnson/vardax#147).
    """
    values, masks, variances = [], [], []
    for instrument in self.forward.instruments:
        observation = jnp.asarray(self.observations[instrument.name], dtype=float)
        mask = instrument.observation_mask()
        if instrument.name in self.radiance_ops:
            n_channels = observation.shape[-1]
            variance = jnp.broadcast_to(
                jnp.asarray(self.radiance_variance[instrument.name], dtype=float),
                observation.shape,
            )
            if instrument.variance_inflation is not None:
                inflation = jnp.asarray(instrument.variance_inflation, dtype=float)
                variance = einx.multiply(
                    "obs channel, obs -> obs channel", variance, inflation
                )
            mask = einx.rearrange("obs -> obs channel", mask, channel=n_channels)
        else:
            variance = instrument.observation_variance()
        if observation.shape != mask.shape:
            raise ValueError(
                f"CoupledProblem: observation for {instrument.name!r} has shape "
                f"{observation.shape}; expected {mask.shape}."
            )
        # y = 0 where m = 0
        values.append(jnp.ravel(jnp.where(mask > 0.0, observation, 0.0)))
        masks.append(jnp.ravel(mask))
        variances.append(jnp.ravel(variance))
    return StackedObservations(
        values=jnp.concatenate(values),
        mask=jnp.concatenate(masks),
        variance=jnp.concatenate(variances),
    )

CoupledResult dataclass

CoupledResult(map_estimate: dict[str, Array], std: dict[str, Array] | None, mean: Array, covariance: Array | None, emission_rate: Array, emission_std: Array, converged: bool, n_steps: int, cost: float, per_instrument_residuals: dict[str, Array], pixels_used: dict[str, int], prior: CoupledPrior | None = None)

MAP and Laplace posterior of the coupled inversion.

emission_rate and emission_std are the posterior mean and standard deviation of Q (population.PerEventPosterior): for a lognormal Q the lognormal moments of the Laplace marginal on log Q, E[Q] = exp(μ + σ²/2) and sd[Q] = E[Q] √(exp(σ²) − 1).

Attributes:

Name	Type	Description
map_estimate	dict[str, Array]	MAP θ* by block (log_emission_rate stays in logs).
std	dict[str, Array] | None	Laplace posterior standard deviation by block (None without the posterior).
mean	Array	MAP θ* flattened, in CoupledPrior.layout order.
covariance	Array | None	Laplace covariance Σ_θ|y, None without the posterior.
emission_rate	Array	Posterior mean of Q.
emission_std	Array	Posterior standard deviation of Q (NaN without the posterior).
converged	bool	optimistix's success flag for the L-BFGS solve.
n_steps	int	Iterations taken.
cost	float	𝒥 at the MAP, in the whitened control.
per_instrument_residuals	dict[str, Array]	y_i − h_i(θ*) by instrument, NaN where masked.
pixels_used	dict[str, int]	Pixels each instrument contributed at full weight.
prior	CoupledPrior | None	The prior the solve used (None when built by hand).

For a scalar Q it also carries the Tier V per-event payload (population.PerEventPosterior): draws of Q from the Laplace marginal of its block, N(Q, σ²) or log Q ~ N(log Q, σ²) for a lognormal Q, and the log-density of that block's prior marginal, both exact for the Gaussian laws on θ.

samples

samples(key: Array, n: int) -> Float[Array, ' member']

n draws of Q from the Laplace marginal of its block.

Raises:

Type	Description
PayloadError	Without the Laplace posterior, or for a vector Q.

Source code in src/plumax/coupled/_src/inversion.py

def samples(self, key: jax.Array, n: int) -> Float[Array, " member"]:
    """n draws of Q from the Laplace marginal of its block.

    Raises:
        PayloadError: Without the Laplace posterior, or for a vector Q.
    """
    block = _emission_block(self.map_estimate)
    if self.std is None:
        raise PayloadError(
            "CoupledResult.samples needs the Laplace posterior; solve with "
            "laplace=True."
        )
    center = self.map_estimate[block]
    spread = self.std[block]
    _require_scalar(center, "samples")
    if block == "log_emission_rate":
        return LognormalEmission(center, spread).sample(key, n)
    return GaussianEmission(center, spread).sample(key, n)

prior_logpdf

prior_logpdf(q: Array) -> jax.Array

log π(Q) of the prior marginal of Q's block.

Raises:

Type	Description
PayloadError	Without the prior, or for a vector Q.

Source code in src/plumax/coupled/_src/inversion.py

def prior_logpdf(self, q: jax.Array) -> jax.Array:
    """log π(Q) of the prior marginal of Q's block.

    Raises:
        PayloadError: Without the prior, or for a vector Q.
    """
    if self.prior is None:
        raise PayloadError(
            "CoupledResult.prior_logpdf: this result carries no prior; build "
            "it with solve_coupled."
        )
    block = _emission_block(self.prior.mean)
    center = jnp.asarray(self.prior.mean[block], dtype=float)
    _require_scalar(center, "prior_logpdf")
    # Q's diagonal entry of Σ_θ, at its offset in the flattened layout
    offset = 0
    for name, length in self.prior.layout:
        if name == block:
            break
        offset += length
    covariance = jnp.asarray(self.prior.covariance, dtype=float)
    if covariance.ndim == 1:
        variance = covariance[offset]
    else:
        variance = covariance[offset, offset]
    if block == "log_emission_rate":
        return LognormalEmission(center, jnp.sqrt(variance)).log_prob(q)
    return GaussianEmission(center, jnp.sqrt(variance)).log_prob(q)

MaskedKPrior dataclass

MaskedKPrior(k_max: int = 10, prior_active: float = 0.3, emission_mean: float = 0.1, emission_std: float = 0.05, background_std: float = 1e-06)

The slot prior: how many slots, how likely each is on, and its Q_k.

a_k ~ Bernoulli(p_active), Q_k ~ LogNormal moment-matched to the real-scale mean and standard deviation, as :func:plumax.gauss_plume.inference.multi_source_plume_model.

Attributes:

Name	Type	Description
k_max	int	K_max ≥ 1 (the design's v1 cap is 10).
prior_active	float	p_active in (0, 1]; 1 keeps every slot on, with no discrete sites.
emission_mean	float	Real-scale prior mean of each Q_k [kg/s].
emission_std	float	Real-scale prior standard deviation of each Q_k [kg/s].
background_std	float	HalfNormal scale of the shared background column.

Raises:

Type	Description
ValueError	On a parameter out of range.

FlagPolicy

Bases: StrEnum

What a rejected pixel becomes: dropped, or kept with an inflated variance.

QualityFlag

Bases: IntFlag

Per-pixel quality bits, 0 = clear.

The bits of the planned radtrax retrieval flags (jejjohnson/radtrax#110, which supersedes plumax#97). Any integer bitmask works; these name the common ones.

QualityPolicy dataclass

QualityPolicy(policy: FlagPolicy = FlagPolicy.MASK, inflate_factor: float = 10.0, reject: int = QualityFlag.CLOUD | QualityFlag.NOT_CONVERGED)

Which flags reject a pixel, and how a rejected pixel enters the fusion.

Attributes:

Name	Type	Description
policy	FlagPolicy	MASK drops rejected pixels; INFLATE scales their variance by inflate_factor.
inflate_factor	float	Variance factor κ ≥ 1 for INFLATE.
reject	int	Bitmask of the flags that reject a pixel.

Raises:

Type	Description
ParameterError	If inflate_factor is below 1 or reject is negative.

rejected

rejected(flags: Int[Array, ' obs']) -> Bool[Array, ' obs']

Whether each pixel carries a rejecting flag, (flags & reject) ≠ 0.

Source code in src/plumax/coupled/_src/quality.py

def rejected(self, flags: Int[Array, " obs"]) -> Bool[Array, " obs"]:
    """Whether each pixel carries a rejecting flag, (flags & reject) ≠ 0."""
    hits = jnp.bitwise_and(jnp.asarray(flags), int(self.reject))
    return hits != 0

StackedObservations

Bases: Module

Every instrument's observations, mask and variance on one stacked axis.

The form vardax takes: y with masked pixels filled by zeros, the mask m and the diagonal of R (inflation included), instrument after instrument.

Attributes:

Name	Type	Description
values	Float[Array, ' obs']	Observations y, zero where masked.
mask	Float[Array, ' obs']	Observation mask m, 1 = used.
variance	Float[Array, ' obs']	Diagonal of the observation-error covariance R.

batch

batch() -> vardax.Batch1D

One-window vardax.Batch1D of shape (1, 1, D_y).

Raises:

Type	Description
ImportError	Without the assim extra.

Source code in src/plumax/coupled/_src/quality.py

def batch(self) -> vardax.Batch1D:
    """One-window ``vardax.Batch1D`` of shape ``(1, 1, D_y)``.

    Raises:
        ImportError: Without the ``assim`` extra.
    """
    vdx = require_vardax("StackedObservations.batch")
    return vdx.Batch1D(
        input=einx.rearrange("obs -> 1 1 obs", self.values),
        mask=einx.rearrange("obs -> 1 1 obs", self.mask),
    )

obs_cov_op

obs_cov_op() -> lx.DiagonalLinearOperator

R as a diagonal operator on a (1, D_y) observation row.

A diagonal operator rather than a scaled identity, which vardax mis-solves (jejjohnson/vardax#146).

Source code in src/plumax/coupled/_src/quality.py

def obs_cov_op(self) -> lx.DiagonalLinearOperator:
    """R as a diagonal operator on a ``(1, D_y)`` observation row.

    A diagonal operator rather than a scaled identity, which vardax
    mis-solves (jejjohnson/vardax#146).
    """
    return lx.DiagonalLinearOperator(einx.rearrange("obs -> 1 obs", self.variance))

LayeredRadianceOperator

Bases: Module

Map a column-mass enhancement to band radiance through radtrax's layered model.

The JAX counterpart of :class:RadianceObservationOperator, on :class:radtrax.ForwardModel instead of the single slab::

column mass m [kg/m²]  →  ΔX = m / M_CH₄ [mol/m²]
    →  state.add_column(gas, ΔX, profile)  →  y = F(state, surface, geometry)

and, when normalise (the default), y(ΔX) / y(0), so zero enhancement gives ones as the slab operator's normalised radiance does. Every step is JAX, so jax.grad of a misfit flows through the plume model and the radiative transfer together, and predict is a jax.vmap over receptors. The layers remove the slab's error of broadening every line at one pressure; with one layer on a table node, an identity instrument and no shift it equals the slab operator.

Band values are divided after band integration, y_b(ΔX) / y_b(0), which is what a sensor's ratio is. The slab operator instead integrates exp(-Δτ) over each band, so with an SRF the two differ by the background spectrum's weighting within the band.

Attributes:

Name	Type	Description
model	ForwardModel	The layered forward model (opacity, solver, instrument).
state	AtmosphericState	The background atmosphere; it must carry gas.
surface	Surface	The surface reflectance.
geometry	ViewingGeometry	Sun and sensor directions.
profile	Float[Array, ' L'] | None	The plume's relative weight per layer, (L,); None puts the whole enhancement in the bottom layer.
gas	str	The gas the plume adds.
normalise	bool	Divide by the background band values y(0).

Examples:

>>> import jax.numpy as jnp
>>> from radtrax import ForwardModel
>>> from radtrax.atmosphere import ViewingGeometry, layers_from_levels
>>> from radtrax.instrument import BandInstrument
>>> from radtrax.solvers import BeerLambert
>>> from radtrax.spectroscopy import CrossSectionTable, TableOpacity
>>> from radtrax.surface import Lambertian
>>> nu = jnp.array([4300.0, 4300.1, 4300.2])
>>> table = CrossSectionTable(
...     jnp.log(jnp.array([1e4, 1e5])),
...     jnp.array([200.0, 300.0]),
...     jnp.full((2, 2, 3), 1e-21),
... )
>>> model = ForwardModel(
...     TableOpacity(nu, {"CH4": table}),
...     BeerLambert(),
...     BandInstrument(nu, jnp.full((1, 3), 1 / 3)),
... )
>>> atm = layers_from_levels(
...     jnp.array([1e5, 5e4]),
...     jnp.array([288.0, 255.0]),
...     {"CH4": jnp.full(2, 1.9e-6)},
... )
>>> op = LayeredRadianceOperator(
...     model, atm, Lambertian(0.3), ViewingGeometry(1.0, 1.0)
... )
>>> float(op.predict_single(0.0)[0])
1.0

n_channels property

n_channels: int

Number of output channels per receptor (the instrument's bands).

delta_column

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

Column mass [kg/m²] → column enhancement ΔX [mol/m²].

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

def delta_column(self, column_mass_kg_m2: jax.Array | float) -> jax.Array:
    """Column mass [kg/m²] → column enhancement ``ΔX`` [mol/m²]."""
    return jnp.asarray(column_mass_kg_m2, dtype=float) / CH4_MOLAR_MASS_KG_PER_MOL

predict_single

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

Band values for one receptor, shape (n_channels,).

Parameters:

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

Returns:

Type	Description
Array	y(ΔX) / y(0) when normalise, else y(ΔX).

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

def predict_single(self, column_mass_kg_m2: jax.Array | float) -> jax.Array:
    """Band values for one receptor, shape ``(n_channels,)``.

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

    Returns:
        ``y(ΔX) / y(0)`` when ``normalise``, else ``y(ΔX)``.
    """
    plume = self.state.add_column(
        self.gas, self.delta_column(column_mass_kg_m2), self.profile
    )
    y = self.model(plume, self.surface, self.geometry)
    if self.normalise:
        y = y / self.model(self.state, self.surface, self.geometry)
    return y

predict

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

Per-receptor band values, shape (n_obs, n_channels).

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

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

def predict(self, column_enhancement: jax.Array | np.ndarray) -> jax.Array:
    """Per-receptor band values, shape ``(n_obs, n_channels)``.

    Args:
        column_enhancement: Per-receptor column mass enhancement [kg/m²],
            shape ``(n_obs,)`` (e.g. ``Q · column_response``).
    """
    m = jnp.asarray(column_enhancement, dtype=float)
    if m.ndim != 1:
        raise ValueError(
            "LayeredRadianceOperator.predict: `column_enhancement` must be "
            f"1-D (n_obs,); got shape {m.shape}."
        )
    return jax.vmap(self.predict_single)(m)

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:radtrax.slab.forward_nonlinear_normalized. With srf supplied, the high-resolution normalised radiance exp(-Δτ(ν)) is band-integrated by the instrument :class:~radtrax.instrument.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:radtrax.spectroscopy.hitran.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 radtrax slab 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 radtrax slab 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 | TransportBackend, instruments: Sequence[Instrument]) -> CoupledForward

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

Parameters:

Name	Type	Description	Default
source	PlumeSource | TransportBackend	Shared transport backend: the Tier I PlumeSource, or a :class:~plumax.coupled.backends.LagrangianSource / :class:~plumax.coupled.backends.EulerianSource.	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 | TransportBackend, instruments: Sequence[Instrument]
) -> CoupledForward:
    """Assemble a :class:`CoupledForward`, precomputing unit-``Q`` responses.

    Args:
        source: Shared transport backend: the Tier I ``PlumeSource``, or a
            :class:`~plumax.coupled.backends.LagrangianSource` /
            :class:`~plumax.coupled.backends.EulerianSource`.
        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(source.column_response(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, then integrates over source.column_z (trapezoidal) and applies the instrument averaging kernel with the shared :class:~plumax.gauss_plume.observation.ColumnObservation (whose module docstring states the units). 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, then
    integrates over ``source.column_z`` (trapezoidal) and applies the instrument
    averaging kernel with the shared
    :class:`~plumax.gauss_plume.observation.ColumnObservation` (whose module
    docstring states the units). 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)
    observation = ColumnObservation.trapezoid(z, averaging_kernel=instrument.ak)
    return observation.column(conc)

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, quality: QualityPolicy | None = None) -> 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)::

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

solved as a vardax.OptimalInterpolation analysis in its state-space (information) form, which factorises only the small (n_param + n_inst) × (n_param + n_inst) posterior precision.

Each instrument's mask zeroes its pixels' weight in R⁻¹ and its variance_inflation scales their variance (see :mod:plumax.coupled.quality); a masked observation may hold anything, NaN included.

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 (n_param + n_inst,) with n_param = 1 for a scalar Q (biases typically prior-zero).	required
prior_covariance	Array	Prior covariance B, shape (n_param + n_inst, n_param + n_inst) — e.g. diag(σ_Q², σ_bias², …).	required
quality	QualityPolicy | None	Optional quality policy, applied to every instrument that carries flags (:func:~plumax.coupled.quality.apply_quality). Default None: use the instruments' masks and inflation as they are.	None

Returns:

Name	Type	Description
The	FusionPosterior	class:FusionPosterior.

Raises:

Type	Description
ValueError	on a per-instrument observation-count mismatch or a mis-shaped prior.
ImportError	Without the assim extra (vardax).

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,
    quality: QualityPolicy | None = None,
) -> 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``)::

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

    solved as a ``vardax.OptimalInterpolation`` analysis in its state-space
    (information) form, which factorises only the small
    ``(n_param + n_inst) × (n_param + n_inst)`` posterior precision.

    Each instrument's ``mask`` zeroes its pixels' weight in ``R⁻¹`` and its
    ``variance_inflation`` scales their variance (see
    :mod:`plumax.coupled.quality`); a masked observation may hold anything,
    NaN included.

    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
            ``(n_param + n_inst,)`` with ``n_param = 1`` for a scalar ``Q``
            (biases typically prior-zero).
        prior_covariance: Prior covariance ``B``, shape
            ``(n_param + n_inst, n_param + n_inst)`` — e.g.
            ``diag(σ_Q², σ_bias², …)``.
        quality: Optional quality policy, applied to every instrument that
            carries ``flags`` (:func:`~plumax.coupled.quality.apply_quality`).
            Default ``None``: use the instruments' masks and inflation as
            they are.

    Returns:
        The :class:`FusionPosterior`.

    Raises:
        ValueError: on a per-instrument observation-count mismatch or a
            mis-shaped prior.
        ImportError: Without the ``assim`` extra (vardax).
    """
    if quality is not None:
        screened = tuple(
            inst if inst.flags is None else apply_quality(inst, inst.flags, quality)
            for inst in forward.instruments
        )
        forward = dataclasses.replace(forward, instruments=screened)
    n_inst = len(forward.instruments)
    n_param = _n_emission_params(forward)
    n_state = n_param + 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 (n_param + n_inst = "
            f"{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, m, bg = _stack_design(forward)
    theta, posterior_cov = _information_form_analysis(
        g, y - bg, m, r, prior_mean=theta_b, prior_covariance=b
    )

    std = jnp.sqrt(jnp.clip(jnp.diag(posterior_cov), 0.0, None))
    names = tuple(inst.name for inst in forward.instruments)
    emission = theta[0] if n_param == 1 else theta[:n_param]
    emission_std = std[0] if n_param == 1 else std[:n_param]
    return FusionPosterior(
        emission_rate=emission,
        emission_std=emission_std,
        biases=theta[n_param:],
        bias_std=std[n_param:],
        mean=theta,
        covariance=posterior_cov,
        instrument_names=names,
        pixels_used=jnp.stack(
            [full_weight_pixels(inst) for inst in forward.instruments]
        ),
        prior_mean=theta_b,
        prior_covariance=b,
    )

solve_coupled

solve_coupled(problem: CoupledProblem, *, method: Literal['lbfgs', 'gauss_newton'] = 'lbfgs', laplace: bool = True, rtol: float = 1e-08, atol: float = 1e-10, max_steps: int = 200) -> CoupledResult

MAP and Laplace posterior of the coupled problem, on vardax.

θ = argmin 𝒥(θ), Σ_θ|y = (Σ_θ⁻¹ + J_hᵀ Σ_ε⁻¹ J_h)⁻¹ at θ.

Algorithm

1. Flatten the prior to μ_θ, factor Σ_θ = L Lᵀ.
2. Stack observations, masks and Σ_ε over the instruments.
3. vardax.ThreeDVar with the stacked forward as a CallableObs and the whitened control θ = μ_θ + L χ; minimise with L-BFGS.
4. vardax.LaplaceCovariance at θ*, materialised.

Parameters:

Name	Type	Description	Default
problem	CoupledProblem	The coupled problem.	required
method	Literal['lbfgs', 'gauss_newton']	"lbfgs". "gauss_newton" waits on vardax's Gauss–Newton minimiser (jejjohnson/vardax#72).	'lbfgs'
laplace	bool	Also compute the Laplace posterior.	True
rtol	float	L-BFGS relative tolerance.	1e-08
atol	float	L-BFGS absolute tolerance.	1e-10
max_steps	int	Iteration cap.	200

Returns:

Name	Type	Description
The	CoupledResult	class:CoupledResult.

Raises:

Type	Description
NotImplementedError	For method="gauss_newton".
ImportError	Without the assim extra.

Source code in src/plumax/coupled/_src/inversion.py

def solve_coupled(
    problem: CoupledProblem,
    *,
    method: Literal["lbfgs", "gauss_newton"] = "lbfgs",
    laplace: bool = True,
    rtol: float = 1e-8,
    atol: float = 1e-10,
    max_steps: int = 200,
) -> CoupledResult:
    """MAP and Laplace posterior of the coupled problem, on vardax.

    θ* = argmin 𝒥(θ),   Σ_θ|y = (Σ_θ⁻¹ + J_hᵀ Σ_ε⁻¹ J_h)⁻¹ at θ*.

    Algorithm:
        1. Flatten the prior to μ_θ, factor Σ_θ = L Lᵀ.
        2. Stack observations, masks and Σ_ε over the instruments.
        3. ``vardax.ThreeDVar`` with the stacked forward as a ``CallableObs``
           and the whitened control θ = μ_θ + L χ; minimise with L-BFGS.
        4. ``vardax.LaplaceCovariance`` at θ*, materialised.

    Args:
        problem: The coupled problem.
        method: ``"lbfgs"``. ``"gauss_newton"`` waits on vardax's
            Gauss–Newton minimiser (jejjohnson/vardax#72).
        laplace: Also compute the Laplace posterior.
        rtol: L-BFGS relative tolerance.
        atol: L-BFGS absolute tolerance.
        max_steps: Iteration cap.

    Returns:
        The :class:`CoupledResult`.

    Raises:
        NotImplementedError: For ``method="gauss_newton"``.
        ImportError: Without the ``assim`` extra.
    """
    if method == "gauss_newton":
        raise NotImplementedError(
            "solve_coupled: the Gauss–Newton minimiser waits on vardax "
            "(jejjohnson/vardax#72); use method='lbfgs'."
        )
    if method != "lbfgs":
        raise ValueError(f"solve_coupled: unknown method {method!r}.")
    import optimistix as optx

    vdx = require_vardax("solve_coupled")

    # --- prior and whitened control ---
    prior = problem.prior
    prior_mean, unflatten = prior.flatten()
    n_param = prior_mean.shape[0]
    # Σ_θ = L Lᵀ
    cholesky = prior.cholesky()
    # θ = μ_θ + L χ, ½ ‖χ‖²
    basis = vdx.linear_basis(cholesky, variance=1.0)
    control = vdx.LinearParameterization(basis, background=prior_mean[None])

    # --- observations ---
    stacked = problem.stacked_observations()

    def stacked_prediction(state: jax.Array) -> jax.Array:
        # h(θ) on a (1, D_θ) row
        return einx.rearrange(
            "obs -> 1 obs", problem.predict_stacked(unflatten(state[0]))
        )

    model = vdx.ThreeDVar(
        obs_op=vdx.CallableObs(fn=stacked_prediction, materialise=True),
        prior_mean=prior_mean[None],
        prior_cov_op=prior.row_operator(),
        obs_cov_op=stacked.obs_cov_op(),
        minimiser=vdx.OptimistixMinimiser(
            optx.LBFGS(rtol=rtol, atol=atol), max_steps=max_steps
        ),
        parameterization=control,
        cov_solver=lx.AutoLinearSolver(well_posed=True),
    )
    batch = stacked.batch()

    # --- MAP ---
    # θ* = argmin 𝒥(θ)
    analysis, outcome = model.analyse(batch)
    theta = analysis[0, 0]
    map_estimate = unflatten(theta)

    # --- Laplace posterior ---
    covariance = None
    std = None
    if laplace:
        # Σ_θ|y = (Σ_θ⁻¹ + J_hᵀ Σ_ε⁻¹ J_h)⁻¹
        posterior = vdx.LaplaceCovariance(solver=lx.AutoLinearSolver(well_posed=True))(
            theta[None], model, batch
        )
        columns = jax.vmap(
            lambda basis_vector: posterior.cov.mv(basis_vector[None])[0]
        )(jnp.eye(n_param, dtype=theta.dtype))
        covariance = gx.symmetrize(columns)
        marginal_std = jnp.sqrt(jnp.clip(jnp.diagonal(covariance), 0.0, None))
        std = unflatten(marginal_std)

    emission_rate, emission_std = _emission_moments(prior, map_estimate, std)

    # --- diagnostics ---
    predictions = problem.predict(map_estimate)
    residuals = {}
    for instrument in problem.forward.instruments:
        observation = jnp.asarray(problem.observations[instrument.name], dtype=float)
        residual = observation - predictions[instrument.name]
        mask = instrument.observation_mask()
        if residual.ndim == 2:
            mask = einx.rearrange("obs -> obs channel", mask, channel=1)
        residuals[instrument.name] = jnp.where(mask > 0.0, residual, jnp.nan)
    return CoupledResult(
        map_estimate=map_estimate,
        std=std,
        mean=theta,
        covariance=covariance,
        emission_rate=emission_rate,
        emission_std=emission_std,
        converged=bool(outcome.converged[0]),
        n_steps=int(outcome.num_steps[0]),
        cost=float(outcome.cost[0]),
        per_instrument_residuals=residuals,
        pixels_used=contribution_report(problem.forward.instruments),
        prior=prior,
    )

build_masked_k_problem

build_masked_k_problem(instruments: Sequence[Instrument], *, transport: PlumeSource | TransportBackend, prior: MaskedKPrior, observations: Mapping[str, Array], active: Bool[Array, ' source'], candidate_locations: Float[Array, 'source 2'] | None = None) -> CoupledProblem

The coupled problem of one active set, for the MAP and Laplace path.

y_i = Σ_{k active} Q_k r_i(x_k, y_k) + c_bg + ε_i over log Q_k and c_bg, under the same lognormal prior as the slots. The background is the Gaussian N(0, σ²_bg) here, against the sampler's HalfNormal(σ_bg): the block prior is Gaussian, and the two have the same curvature at zero. Hand it to :func:~plumax.coupled.solve_coupled, which solves on vardax. This is the operational path; the active set itself comes from :func:infer_masked_k_sources or from the caller.

Parameters:

Name	Type	Description	Default
instruments	Sequence[Instrument]	The observing instruments.	required
transport	PlumeSource | TransportBackend	The transport backend.	required
prior	MaskedKPrior	The slot prior.	required
observations	Mapping[str, Array]	Observations by instrument.	required
active	Bool[Array, ' source']	a_k, with at least one slot on.	required
candidate_locations	Float[Array, 'source 2'] | None	Fixed (x, y) per slot [m], for a PlumeSource.	None

Returns:

Type	Description
CoupledProblem	The problem, with a log_emission_rate block over the active slots and a
CoupledProblem	background block.

Raises:

Type	Description
ValueError	When active does not match the slots or has none on.

Source code in src/plumax/coupled/_src/k_sources.py

def build_masked_k_problem(
    instruments: Sequence[Instrument],
    *,
    transport: PlumeSource | TransportBackend,
    prior: MaskedKPrior,
    observations: Mapping[str, Array],
    active: Bool[Array, " source"],
    candidate_locations: Float[Array, "source 2"] | None = None,
) -> CoupledProblem:
    """The coupled problem of one active set, for the MAP and Laplace path.

    y_i = Σ_{k active} Q_k r_i(x_k, y_k) + c_bg + ε_i over log Q_k and c_bg, under the
    same lognormal prior as the slots. The background is the Gaussian N(0, σ²_bg) here,
    against the sampler's HalfNormal(σ_bg): the block prior is Gaussian, and the two
    have the same curvature at zero. Hand it to :func:`~plumax.coupled.solve_coupled`,
    which solves on vardax. This is the operational path; the active set itself comes
    from :func:`infer_masked_k_sources` or from the caller.

    Args:
        instruments: The observing instruments.
        transport: The transport backend.
        prior: The slot prior.
        observations: Observations by instrument.
        active: a_k, with at least one slot on.
        candidate_locations: Fixed ``(x, y)`` per slot [m], for a ``PlumeSource``.

    Returns:
        The problem, with a ``log_emission_rate`` block over the active slots and a
        ``background`` block.

    Raises:
        ValueError: When ``active`` does not match the slots or has none on.
    """
    mask = np.asarray(active, dtype=bool)
    matrices = slot_responses(transport, instruments, candidate_locations)
    if mask.shape != (prior.k_max,) or not mask.any():
        raise ValueError(
            f"build_masked_k_problem: active must have shape ({prior.k_max},) with "
            f"at least one slot on, got {mask.tolist()}."
        )
    emission = LognormalEmission.from_moments(prior.emission_mean, prior.emission_std)
    n_active = int(mask.sum())
    forward = CoupledForward(
        source=transport,
        instruments=tuple(instruments),
        responses=tuple(matrix[:, mask] for matrix in matrices),
    )
    coupled_prior = CoupledPrior.from_std(
        {
            "log_emission_rate": emission.loc * jnp.ones(n_active),
            "background": 0.0,
        },
        {
            "log_emission_rate": emission.scale * jnp.ones(n_active),
            "background": prior.background_std,
        },
    )
    return CoupledProblem(
        forward=forward, prior=coupled_prior, observations=dict(observations)
    )

infer_masked_k_sources

infer_masked_k_sources(observations: Mapping[str, Array], *, instruments: Sequence[Instrument], transport: PlumeSource | TransportBackend, prior: MaskedKPrior, candidate_locations: Float[Array, 'source 2'] | None = None, source_bounds: tuple[tuple[float, float], tuple[float, float]] | None = None, num_warmup: int = 500, num_samples: int = 1000, seed: int = 0) -> MultiSourcePosterior

Posterior over the active count and the per-slot Q_k, from every instrument.

Tier I's runner on :func:masked_k_coupled_model: NUTS for the continuous sites, a Gibbs step for the a_k.

Parameters:

Name	Type	Description	Default
observations	Mapping[str, Array]	Column observations by instrument.	required
instruments	Sequence[Instrument]	The observing instruments.	required
transport	PlumeSource | TransportBackend	The transport backend.	required
prior	MaskedKPrior	The slot prior.	required
candidate_locations	Float[Array, 'source 2'] | None	Fixed (x, y) per slot [m].	None
source_bounds	tuple[tuple[float, float], tuple[float, float]] | None	Box of free slots, instead of candidates.	None
num_warmup	int	NUTS warm-up steps.	500
num_samples	int	Posterior draws.	1000
seed	int	PRNG seed.	0

Returns:

Type	Description
MultiSourcePosterior	The posterior: count_probabilities() is P(K | y), and per_source(k)
MultiSourcePosterior	the slot's summary, which population.event_from_posterior accepts.

Raises:

Type	Description
ValueError	On missing observations or an invalid slot layout.

Source code in src/plumax/coupled/_src/k_sources.py

def infer_masked_k_sources(
    observations: Mapping[str, Array],
    *,
    instruments: Sequence[Instrument],
    transport: PlumeSource | TransportBackend,
    prior: MaskedKPrior,
    candidate_locations: Float[Array, "source 2"] | None = None,
    source_bounds: tuple[tuple[float, float], tuple[float, float]] | None = None,
    num_warmup: int = 500,
    num_samples: int = 1000,
    seed: int = 0,
) -> MultiSourcePosterior:
    """Posterior over the active count and the per-slot Q_k, from every instrument.

    Tier I's runner on :func:`masked_k_coupled_model`: NUTS for the continuous sites,
    a Gibbs step for the a_k.

    Args:
        observations: Column observations by instrument.
        instruments: The observing instruments.
        transport: The transport backend.
        prior: The slot prior.
        candidate_locations: Fixed ``(x, y)`` per slot [m].
        source_bounds: Box of free slots, instead of candidates.
        num_warmup: NUTS warm-up steps.
        num_samples: Posterior draws.
        seed: PRNG seed.

    Returns:
        The posterior: ``count_probabilities()`` is P(K | y), and ``per_source(k)``
        the slot's summary, which ``population.event_from_posterior`` accepts.

    Raises:
        ValueError: On missing observations or an invalid slot layout.
    """
    from plumax.gauss_plume.inference import _check_multi_source_layout, _run_masked_k

    names = {instrument.name for instrument in instruments}
    if set(observations) != names:
        raise ValueError(
            f"infer_masked_k_sources: observations for {sorted(observations)} but "
            f"the instruments are {sorted(names)}."
        )
    fixed_slots = source_bounds is None
    is_plume = isinstance(transport, PlumeSource)
    if is_plume or not fixed_slots:
        _check_multi_source_layout(
            prior.k_max, candidate_locations, source_bounds, prior.prior_active
        )
    if not fixed_slots and not is_plume:
        raise ValueError(
            "infer_masked_k_sources: free slots rerun the Tier I plume and need a "
            "PlumeSource."
        )
    fixed_responses = None
    if fixed_slots:
        fixed_responses = slot_responses(transport, instruments, candidate_locations)
        widths = {matrix.shape[-1] for matrix in fixed_responses}
        if widths != {prior.k_max}:
            raise ValueError(
                f"infer_masked_k_sources: the responses have {widths} slots but "
                f"k_max is {prior.k_max}."
            )
    layout = candidate_locations
    if layout is None and fixed_slots:
        # a backend's slots have no coordinates of their own
        layout = jnp.full((prior.k_max, 2), jnp.nan)
    return _run_masked_k(
        masked_k_coupled_model,
        observations,
        k_max=prior.k_max,
        candidate_locations=layout,
        prior_active=prior.prior_active,
        num_warmup=num_warmup,
        num_samples=num_samples,
        num_chains=1,
        seed=seed,
        progress_bar=False,
        instruments=instruments,
        transport=transport,
        source_bounds=source_bounds,
        emission_mean=prior.emission_mean,
        emission_std=prior.emission_std,
        background_std=prior.background_std,
        fixed_responses=fixed_responses,
    )

masked_k_coupled_model

masked_k_coupled_model(observations: Mapping[str, Array] | None, *, instruments: Sequence[Instrument], transport: PlumeSource | TransportBackend, k_max: int, candidate_locations: Float[Array, 'source 2'] | None = None, source_bounds: tuple[tuple[float, float], tuple[float, float]] | None = None, prior_active: float = 0.3, emission_mean: float = 0.1, emission_std: float = 0.05, background_std: float = 1e-06, fixed_responses: Sequence[Array] | None = None) -> None

NumPyro masked-K model over every instrument's columns.

y_i ~ Normal(Σ_k a_k Q_k r_i(x_k, y_k) + c_bg, Σ_ε,i), with Σ_ε,i and the mask of each :class:~plumax.coupled.Instrument, under Tier I's slot prior. With source_bounds the slots move, so a PlumeSource is rerun at each draw.

Parameters:

Name	Type	Description	Default
observations	Mapping[str, Array] | None	Column observations by instrument, or None to predict.	required
instruments	Sequence[Instrument]	The observing instruments.	required
transport	PlumeSource | TransportBackend	The transport backend.	required
k_max	int	K_max.	required
candidate_locations	Float[Array, 'source 2'] | None	Fixed (x, y) per slot [m].	None
source_bounds	tuple[tuple[float, float], tuple[float, float]] | None	Box of free, x-ordered slots; needs a PlumeSource.	None
prior_active	float	p_active.	0.3
emission_mean	float	Real-scale prior mean of each Q_k [kg/s].	0.1
emission_std	float	Real-scale prior standard deviation of each Q_k [kg/s].	0.05
background_std	float	HalfNormal scale of the background column.	1e-06
fixed_responses	Sequence[Array] | None	:func:slot_responses, when the slots are fixed; passed in so they are computed once rather than at every model evaluation.	None

Source code in src/plumax/coupled/_src/k_sources.py

def masked_k_coupled_model(
    observations: Mapping[str, Array] | None,
    *,
    instruments: Sequence[Instrument],
    transport: PlumeSource | TransportBackend,
    k_max: int,
    candidate_locations: Float[Array, "source 2"] | None = None,
    source_bounds: tuple[tuple[float, float], tuple[float, float]] | None = None,
    prior_active: float = 0.3,
    emission_mean: float = 0.1,
    emission_std: float = 0.05,
    background_std: float = 1e-6,
    fixed_responses: Sequence[Array] | None = None,
) -> None:
    """NumPyro masked-K model over every instrument's columns.

    y_i ~ Normal(Σ_k a_k Q_k r_i(x_k, y_k) + c_bg, Σ_ε,i), with Σ_ε,i and the mask of
    each :class:`~plumax.coupled.Instrument`, under Tier I's slot prior. With
    ``source_bounds`` the slots move, so a ``PlumeSource`` is rerun at each draw.

    Args:
        observations: Column observations by instrument, or ``None`` to predict.
        instruments: The observing instruments.
        transport: The transport backend.
        k_max: K_max.
        candidate_locations: Fixed ``(x, y)`` per slot [m].
        source_bounds: Box of free, x-ordered slots; needs a ``PlumeSource``.
        prior_active: p_active.
        emission_mean: Real-scale prior mean of each Q_k [kg/s].
        emission_std: Real-scale prior standard deviation of each Q_k [kg/s].
        background_std: HalfNormal scale of the background column.
        fixed_responses: :func:`slot_responses`, when the slots are fixed; passed in so
            they are computed once rather than at every model evaluation.
    """
    import numpyro
    import numpyro.distributions as dist

    from plumax.gauss_plume.inference import _masked_source_prior

    sources, background = _masked_source_prior(
        k_max=k_max,
        candidate_locations=candidate_locations,
        source_bounds=source_bounds,
        prior_emission_rate_mean=emission_mean,
        prior_emission_rate_std=emission_std,
        prior_active=prior_active,
        background_prior_std=background_std,
    )
    # a_k Q_k
    active_rates = sources.active * sources.emission_rate
    for index, instrument in enumerate(instruments):
        if fixed_responses is not None:
            design = fixed_responses[index]
        else:
            locations = jnp.stack([sources.x, sources.y], axis=-1)
            design = slot_responses(transport, [instrument], locations)[0]
        # Σ_k a_k Q_k r_i(x_k, y_k) + c_bg
        predicted = design @ active_rates + background
        noise_std = jnp.sqrt(instrument.observation_variance())
        keep = instrument.observation_mask() > 0.0
        likelihood = dist.Normal(predicted, noise_std).mask(keep)
        observed = None
        if observations is not None:
            # masked pixels may hold anything, and the mask drops them
            observed = jnp.where(keep, jnp.asarray(observations[instrument.name]), 0.0)
        numpyro.sample(f"obs_{instrument.name}", likelihood.to_event(1), obs=observed)

saturation_rate

saturation_rate(posteriors: Sequence[MultiSourcePosterior]) -> float

Fraction of overpasses whose most probable active count is K_max.

K̂ = argmax_k P(K = k | y), rate = #{K̂ = K_max} / N. Above 0.05 the cap is binding and the RJMCMC model is due (see the module docstring).

Parameters:

Name	Type	Description	Default
posteriors	Sequence[MultiSourcePosterior]	One masked-K posterior per overpass.	required

Returns:

Type	Description
float	The saturated fraction.

Raises:

Type	Description
ValueError	With no posteriors.

Examples:

One overpass at the cap, one below it:

>>> import numpy as np
>>> def posterior(active):
...     shape = np.shape(active)
...     return MultiSourcePosterior(
...         emission_rate=np.ones(shape),
...         x=np.zeros(shape),
...         y=np.zeros(shape),
...         active=np.asarray(active),
...         background=np.zeros(shape[0]),
...     )
>>> full = posterior([[True, True], [True, True]])
>>> partial = posterior([[True, False], [True, False]])
>>> saturation_rate([full, partial])
0.5

Source code in src/plumax/coupled/_src/k_sources.py

def saturation_rate(posteriors: Sequence[MultiSourcePosterior]) -> float:
    """Fraction of overpasses whose most probable active count is K_max.

    K̂ = argmax_k P(K = k | y), rate = #{K̂ = K_max} / N. Above 0.05 the cap is binding
    and the RJMCMC model is due (see the module docstring).

    Args:
        posteriors: One masked-K posterior per overpass.

    Returns:
        The saturated fraction.

    Raises:
        ValueError: With no posteriors.

    Examples:
        One overpass at the cap, one below it:

        >>> import numpy as np
        >>> def posterior(active):
        ...     shape = np.shape(active)
        ...     return MultiSourcePosterior(
        ...         emission_rate=np.ones(shape),
        ...         x=np.zeros(shape),
        ...         y=np.zeros(shape),
        ...         active=np.asarray(active),
        ...         background=np.zeros(shape[0]),
        ...     )
        >>> full = posterior([[True, True], [True, True]])
        >>> partial = posterior([[True, False], [True, False]])
        >>> saturation_rate([full, partial])
        0.5
    """
    if not posteriors:
        raise ValueError("saturation_rate: no posteriors given.")
    saturated = [
        int(np.argmax(posterior.count_probabilities())) == posterior.k_max
        for posterior in posteriors
    ]
    return float(np.mean(saturated))

slot_responses

slot_responses(transport: PlumeSource | TransportBackend, instruments: Sequence[Instrument], candidate_locations: Float[Array, 'source 2'] | None = None) -> list[Array]

Unit-emission column of every slot, as one design matrix per instrument.

R_i[n, k] = r_i(x_k, y_k) at receptor n. The Tier I plume is rerun at each candidate. Any other backend has its emission parameters as the slots already: its column_response must return the (obs, source) matrix.

Parameters:

Name	Type	Description	Default
transport	PlumeSource | TransportBackend	The transport backend.	required
instruments	Sequence[Instrument]	The observing instruments.	required
candidate_locations	Float[Array, 'source 2'] | None	(x, y) per slot [m], for a PlumeSource.	None

Returns:

Type	Description
list[Array]	One (obs, source) matrix per instrument.

Raises:

Type	Description
ValueError	For a PlumeSource without candidates, or a backend with candidates, or whose response is not a matrix.

Source code in src/plumax/coupled/_src/k_sources.py

def slot_responses(
    transport: PlumeSource | TransportBackend,
    instruments: Sequence[Instrument],
    candidate_locations: Float[Array, "source 2"] | None = None,
) -> list[Array]:
    """Unit-emission column of every slot, as one design matrix per instrument.

    R_i[n, k] = r_i(x_k, y_k) at receptor n. The Tier I plume is rerun at each
    candidate. Any other backend has its emission parameters as the slots already:
    its ``column_response`` must return the ``(obs, source)`` matrix.

    Args:
        transport: The transport backend.
        instruments: The observing instruments.
        candidate_locations: ``(x, y)`` per slot [m], for a ``PlumeSource``.

    Returns:
        One ``(obs, source)`` matrix per instrument.

    Raises:
        ValueError: For a ``PlumeSource`` without candidates, or a backend with
            candidates, or whose response is not a matrix.
    """
    if isinstance(transport, PlumeSource):
        if candidate_locations is None:
            raise ValueError(
                "slot_responses: a PlumeSource needs candidate_locations (x, y)."
            )
        locations = jnp.asarray(candidate_locations, dtype=float)
        wind = jnp.asarray([transport.wind_speed, transport.wind_direction])

        def design(instrument: Instrument) -> Array:
            response = jax.vmap(plume_response, (None, None, 0, None), 1)
            return response(transport, instrument, locations, wind)

        return [design(instrument) for instrument in instruments]
    if candidate_locations is not None:
        raise ValueError(
            "slot_responses: only a PlumeSource reruns at candidate_locations; another "
            "backend's emission parameters are its slots."
        )
    matrices = [transport.column_response(instrument) for instrument in instruments]
    if any(jnp.ndim(matrix) != 2 for matrix in matrices):
        raise ValueError(
            "slot_responses: the backend must return one column per slot, (obs, source)."
        )
    return [jnp.asarray(matrix, dtype=float) for matrix in matrices]

apply_quality

apply_quality(instrument: Instrument, flags: Int[Array, ' obs'] | None = None, policy: QualityPolicy | None = None) -> Instrument

The instrument with its flags and the policy's mask or inflation set.

m_i = 0 for a rejected pixel under MASK; R_ii ← κ R_ii under INFLATE. The geometry, kernel and error budget are unchanged, so the coupled forward's responses stay valid. The result depends only on the flags and the policy, so applying it twice is the same as once.

Parameters:

Name	Type	Description	Default
instrument	Instrument	The instrument to screen.	required
flags	Int[Array, ' obs'] | None	Per-pixel bit flags. Default: the instrument's own flags (all clear when it has none).	None
policy	QualityPolicy | None	The quality policy. Default QualityPolicy() (mask cloud and failed convergence).	None

Returns:

Type	Description
Instrument	A copy of instrument with flags, mask and
Instrument	variance_inflation set.

Examples:

Mask the cloudy pixel of three:

>>> import jax.numpy as jnp
>>> from plumax.coupled import Instrument
>>> instrument = Instrument("A", jnp.zeros((3, 2)))
>>> flags = jnp.array([0, QualityFlag.CLOUD, QualityFlag.AOD])
>>> apply_quality(instrument, flags).observation_mask().tolist()
[1.0, 0.0, 1.0]

Inflate instead, by 100:

>>> inflate = QualityPolicy(FlagPolicy.INFLATE, inflate_factor=100.0)
>>> apply_quality(instrument, flags, inflate).observation_variance().tolist()
[1.0, 100.0, 1.0]

Source code in src/plumax/coupled/_src/quality.py

def apply_quality(
    instrument: Instrument,
    flags: Int[Array, " obs"] | None = None,
    policy: QualityPolicy | None = None,
) -> Instrument:
    """The instrument with its flags and the policy's mask or inflation set.

    m_i = 0 for a rejected pixel under ``MASK``; R_ii ← κ R_ii under
    ``INFLATE``. The geometry, kernel and error budget are unchanged, so the
    coupled forward's responses stay valid. The result depends only on the
    flags and the policy, so applying it twice is the same as once.

    Args:
        instrument: The instrument to screen.
        flags: Per-pixel bit flags. Default: the instrument's own ``flags``
            (all clear when it has none).
        policy: The quality policy. Default ``QualityPolicy()`` (mask cloud
            and failed convergence).

    Returns:
        A copy of ``instrument`` with ``flags``, ``mask`` and
        ``variance_inflation`` set.

    Examples:
        Mask the cloudy pixel of three:

        >>> import jax.numpy as jnp
        >>> from plumax.coupled import Instrument
        >>> instrument = Instrument("A", jnp.zeros((3, 2)))
        >>> flags = jnp.array([0, QualityFlag.CLOUD, QualityFlag.AOD])
        >>> apply_quality(instrument, flags).observation_mask().tolist()
        [1.0, 0.0, 1.0]

        Inflate instead, by 100:

        >>> inflate = QualityPolicy(FlagPolicy.INFLATE, inflate_factor=100.0)
        >>> apply_quality(instrument, flags, inflate).observation_variance().tolist()
        [1.0, 100.0, 1.0]
    """
    policy = QualityPolicy() if policy is None else policy
    if flags is None:
        flags = instrument.flags
    if flags is None:
        flags = jnp.zeros(instrument.n_obs, dtype=int)
    rejected = policy.rejected(flags)
    if policy.policy is FlagPolicy.MASK:
        # m_i = 0 for a rejected pixel
        mask = jnp.where(rejected, 0.0, 1.0)
        return dataclasses.replace(
            instrument, flags=jnp.asarray(flags), mask=mask, variance_inflation=None
        )
    # κ_i = κ for a rejected pixel, 1 otherwise
    inflation = jnp.where(rejected, policy.inflate_factor, 1.0)
    return dataclasses.replace(
        instrument, flags=jnp.asarray(flags), mask=None, variance_inflation=inflation
    )

contribution_report

contribution_report(instruments: Sequence[Instrument]) -> dict[str, int]

Pixels each instrument contributes at full weight, by name.

Masked pixels and inflated pixels both count as not contributing in full, so the report is the effective per-instrument observation count the design asks to keep auditable.

Parameters:

Name	Type	Description	Default
instruments	Sequence[Instrument]	The instruments of a fusion.	required

Returns:

Type	Description
dict[str, int]	{name: pixel count}.

Raises:

Type	Description
ValueError	If two instruments share a name.

Examples:

>>> import jax.numpy as jnp
>>> from plumax.coupled import Instrument
>>> clear = Instrument("A", jnp.zeros((2, 2)))
>>> cloudy = apply_quality(
...     Instrument("B", jnp.zeros((2, 2))), jnp.full(2, QualityFlag.CLOUD)
... )
>>> contribution_report([clear, cloudy])
{'A': 2, 'B': 0}

Source code in src/plumax/coupled/_src/quality.py

def contribution_report(instruments: Sequence[Instrument]) -> dict[str, int]:
    """Pixels each instrument contributes at full weight, by name.

    Masked pixels and inflated pixels both count as not contributing in full,
    so the report is the effective per-instrument observation count the design
    asks to keep auditable.

    Args:
        instruments: The instruments of a fusion.

    Returns:
        ``{name: pixel count}``.

    Raises:
        ValueError: If two instruments share a name.

    Examples:
        >>> import jax.numpy as jnp
        >>> from plumax.coupled import Instrument
        >>> clear = Instrument("A", jnp.zeros((2, 2)))
        >>> cloudy = apply_quality(
        ...     Instrument("B", jnp.zeros((2, 2))), jnp.full(2, QualityFlag.CLOUD)
        ... )
        >>> contribution_report([clear, cloudy])
        {'A': 2, 'B': 0}
    """
    names = [instrument.name for instrument in instruments]
    if len(set(names)) != len(names):
        raise ValueError(f"contribution_report: instrument names repeat: {names}.")
    return {
        instrument.name: int(full_weight_pixels(instrument))
        for instrument in instruments
    }

stack_observations

stack_observations(instruments: Sequence[Instrument], observations: Sequence[Float[Array, ' _']]) -> StackedObservations

Stack the instruments' observations, masks and variances for vardax.

Parameters:

Name	Type	Description	Default
instruments	Sequence[Instrument]	The instruments, with any quality policy already applied.	required
observations	Sequence[Float[Array, ' _']]	One observation vector per instrument, aligned with instruments. Masked entries may hold anything, NaN included.	required

Returns:

Type	Description
StackedObservations	The stacked observations, masked entries set to zero.

Raises:

Type	Description
ValueError	On a count or shape mismatch.

Examples:

A NaN under the mask becomes a zero:

>>> import jax.numpy as jnp
>>> from plumax.coupled import Instrument
>>> instrument = Instrument("A", jnp.zeros((2, 2)), mask=jnp.array([1.0, 0.0]))
>>> stacked = stack_observations([instrument], [jnp.array([3.0, jnp.nan])])
>>> stacked.values.tolist(), stacked.mask.tolist()
([3.0, 0.0], [1.0, 0.0])

Source code in src/plumax/coupled/_src/quality.py

def stack_observations(
    instruments: Sequence[Instrument],
    observations: Sequence[Float[Array, " _"]],
) -> StackedObservations:
    """Stack the instruments' observations, masks and variances for vardax.

    Args:
        instruments: The instruments, with any quality policy already applied.
        observations: One observation vector per instrument, aligned with
            ``instruments``. Masked entries may hold anything, NaN included.

    Returns:
        The stacked observations, masked entries set to zero.

    Raises:
        ValueError: On a count or shape mismatch.

    Examples:
        A NaN under the mask becomes a zero:

        >>> import jax.numpy as jnp
        >>> from plumax.coupled import Instrument
        >>> instrument = Instrument("A", jnp.zeros((2, 2)), mask=jnp.array([1.0, 0.0]))
        >>> stacked = stack_observations([instrument], [jnp.array([3.0, jnp.nan])])
        >>> stacked.values.tolist(), stacked.mask.tolist()
        ([3.0, 0.0], [1.0, 0.0])
    """
    if len(observations) != len(instruments):
        raise ValueError(
            f"stack_observations: got {len(observations)} observation vectors for "
            f"{len(instruments)} instruments."
        )
    values, masks, variances = [], [], []
    for instrument, observation in zip(instruments, observations, strict=True):
        if np.shape(observation) != (instrument.n_obs,):
            raise ValueError(
                f"stack_observations: observation for {instrument.name!r} must be "
                f"({instrument.n_obs},); got {np.shape(observation)}."
            )
        mask = instrument.observation_mask()
        # y_i = 0 where m_i = 0, so a NaN never reaches the cost
        filled = jnp.where(mask > 0.0, jnp.asarray(observation, dtype=float), 0.0)
        values.append(filled)
        masks.append(mask)
        variances.append(instrument.observation_variance())
    return StackedObservations(
        values=jnp.concatenate(values),
        mask=jnp.concatenate(masks),
        variance=jnp.concatenate(variances),
    )

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 radtrax slab 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².)

1. 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:radtrax.slab.number_density_cm3) and L is the vertical path length [cm] over which the enhancement is distributed.

1. 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 radtrax slab 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 radtrax slab 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:`radtrax.slab.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 radtrax slab 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 ``radtrax.slab.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 | LayeredRadianceOperator) -> 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. For a :class:LayeredRadianceOperator the gain is jax.jacfwd of its predict_single at zero, already in column-mass units.

Parameters:

Name	Type	Description	Default
operator	RadianceObservationOperator | LayeredRadianceOperator	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 | LayeredRadianceOperator,
) -> 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. For a
    :class:`LayeredRadianceOperator` the gain is ``jax.jacfwd`` of its
    ``predict_single`` at zero, already in column-mass units.

    Args:
        operator: The configured nonlinear RTM observation operator.

    Returns:
        The matching :class:`LinearisedRadianceOperator`.
    """
    if isinstance(operator, LayeredRadianceOperator):
        zero = jnp.asarray(0.0)
        return LinearisedRadianceOperator(
            baseline=operator.predict_single(zero),
            gain=jax.jacfwd(operator.predict_single)(zero),
        )
    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 | LayeredRadianceOperator, emission_rate: Array | float) -> np.ndarray | jax.Array

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 | LayeredRadianceOperator	The configured :class:RadianceObservationOperator, or a :class:LayeredRadianceOperator, with which the whole map from Q to radiance is JAX and differentiable in Q.	required
emission_rate	Array | float	Emission rate Q [kg/s] (scalar).	required

Returns:

Type	Description
ndarray | Array	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 | LayeredRadianceOperator,
    emission_rate: jax.Array | float,
) -> np.ndarray | jax.Array:
    """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`, or a
            :class:`LayeredRadianceOperator`, with which the whole map from
            ``Q`` to radiance is JAX and differentiable in ``Q``.
        emission_rate: Emission rate ``Q`` [kg/s] (scalar).

    Returns:
        Per-receptor normalised radiance, shape ``(n_obs, n_channels)``.
    """
    response = column_response(source, instrument)
    if isinstance(operator, LayeredRadianceOperator):
        return operator.predict(jnp.asarray(emission_rate, dtype=float) * response)
    # Unit-Q column response is differentiable JAX; scale to the column mass and
    # hand concrete values to the NumPy-side RTM forward.
    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)