Nonlinear inversion

The end-to-end coupled inversion over the emission rate (Gaussian or lognormal), source position, wind, background and per-instrument biases, through each instrument's averaging kernel or a JAX radiance operator. plumax builds the forward, the prior and the quality masks; vardax solves it (ThreeDVar with L-BFGS in the whitened control) and gives the Laplace posterior (LaplaceCovariance). Needs the assim extra. Tier IV.

inversion

Nonlinear end-to-end coupled inversion on vardax (implementation in plumax.coupled._src.inversion).

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)

RadianceOperator

Bases: Protocol

A JAX radiance operator: column enhancement → per-receptor channels.

LinearisedRadianceOperator and LayeredRadianceOperator qualify.

predict

predict(column_enhancement: Array) -> jax.Array

Radiance per receptor and channel, shape (n_obs, n_channels).

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

def predict(self, column_enhancement: jax.Array) -> jax.Array:
    """Radiance per receptor and channel, shape ``(n_obs, n_channels)``."""
    ...

plume_response

plume_response(source: PlumeSource, instrument: Instrument, location: Float[Array, ' 2'], wind: Float[Array, ' 2']) -> Float[Array, ' obs']

Tier I unit-emission column at a traced source position and wind.

coupled.column_response with the horizontal position and the wind (speed, meteorological direction) as JAX arrays, so the coupled cost can differentiate through them. The source height, stability class and levels stay those of source.

Parameters:

Name	Type	Description	Default
source	PlumeSource	The Tier I source (height, stability, levels).	required
instrument	Instrument	The observing instrument.	required
location	Float[Array, ' 2']	Source (x, y) [m].	required
wind	Float[Array, ' 2']	(speed [m/s], direction [deg from North]).	required

Returns:

Type	Description
Float[Array, ' obs']	The unit-emission column through the averaging kernel.

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

def plume_response(
    source: PlumeSource,
    instrument: Instrument,
    location: Float[Array, " 2"],
    wind: Float[Array, " 2"],
) -> Float[Array, " obs"]:
    """Tier I unit-emission column at a traced source position and wind.

    ``coupled.column_response`` with the horizontal position and the wind
    (speed, meteorological direction) as JAX arrays, so the coupled cost can
    differentiate through them. The source height, stability class and levels
    stay those of ``source``.

    Args:
        source: The Tier I source (height, stability, levels).
        instrument: The observing instrument.
        location: Source ``(x, y)`` [m].
        wind: ``(speed [m/s], direction [deg from North])``.

    Returns:
        The unit-emission column through the averaging kernel.
    """
    levels = jnp.asarray(source.column_z, dtype=float)
    receptors = jnp.asarray(instrument.receptors, dtype=float)
    n_levels = levels.shape[0]
    # every receptor at every level, level fastest
    x = einx.rearrange("obs -> (obs level)", receptors[:, 0], level=n_levels)
    y = einx.rearrange("obs -> (obs level)", receptors[:, 1], level=n_levels)
    z = einx.rearrange("level -> (obs level)", levels, obs=receptors.shape[0])
    direction = jnp.deg2rad(wind[1])
    # (u, v) = −s (sin φ, cos φ)
    wind_u = -wind[0] * jnp.sin(direction)
    wind_v = -wind[0] * jnp.cos(direction)
    concentration = plume_concentration_vmap(
        x,
        y,
        z,
        location[0],
        location[1],
        source.location[2],
        wind_u,
        wind_v,
        1.0,
        get_dispersion_params(source.stability_class),
    )
    profiles = einx.rearrange("(obs level) -> obs level", concentration, level=n_levels)
    observation = ColumnObservation.trapezoid(levels, averaging_kernel=instrument.ak)
    return observation.column(profiles)

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,
    )