The Observation Model

The observation model is the bridge between the state we want to estimate and the data we actually measure. It answers a single question: given a state $\boldsymbol{x}$, what observations $\boldsymbol{y}$ should we expect to see? Everything downstream — the likelihood, the variational cost, the tangent-linear and adjoint operators in 4DVar — is built on top of the observation operator $H$.

The code below is deliberately library-agnostic: shapes are pinned with jaxtyping, array operations are written with einx, and the tangent-linear / adjoint operators come for free from JAX autodiff.

The Likelihood Model¶

We assume Gaussian observation errors,

which yields the likelihood $p(\boldsymbol{y} \mid \boldsymbol{x}) = \mathcal{N}(\boldsymbol{y}; H(\boldsymbol{x}), \mathbf{R})$. The negative log-likelihood (up to a constant) is the observation cost that appears in every variational problem,

import einx
from jaxtyping import Array, Float

def obs_cost(
    y:     Float[Array, "M"],
    Hx:    Float[Array, "M"],
    R_inv: Float[Array, "M M"],   # applied matrix-free in practice (see below)
) -> Float[Array, ""]:
    """Jₒᵦₛ = ½ (y - H(x))ᵀ R⁻¹ (y - H(x))."""
    r = y - Hx
    return 0.5 * einx.dot("M, M ->", r, einx.dot("M K, K -> M", R_inv, r))

The Observation Operator¶

An observation operator is any differentiable map from state space to observation space. That is the entire contract:

from typing import Protocol
from jaxtyping import Array, Float

class ObservationOperator(Protocol):
    def __call__(self, x: Float[Array, "N"]) -> Float[Array, "M"]:
        """Forward map H : state space (N) -> observation space (M)."""
        ...

Four Operator Families¶

def masked_identity(
    x:    Float[Array, "N"],
    mask: Float[Array, "N"],   # 0/1 weights
) -> Float[Array, "N"]:
    return einx.multiply("N, N -> N", mask, x)

def linear_obs(
    H: Float[Array, "M N"],
    x: Float[Array, "N"],
) -> Float[Array, "M"]:
    return einx.dot("M N, N -> M", H, x)

def averaging_kernel(
    x:   Float[Array, "N"],
    A:   Float[Array, "M N"],
    x_a: Float[Array, "N"],
    h:   Float[Array, "N"],
) -> Float[Array, "M"]:
    blended = einx.multiply("N, N -> N", h, x) + einx.multiply("N, N -> N", 1.0 - h, x_a)
    return einx.dot("M N, N -> M", A, blended)

from typing import Callable, NamedTuple

class Instrument(NamedTuple):
    H:      Callable[[Float[Array, "N"]], Float[Array, "M"]]
    y:      Float[Array, "M"]
    mask:   Float[Array, "M"]    # 0/1 quality weights
    R_inv:  Float[Array, "M M"]
    weight: float = 1.0          # αᵢ

def fusion_cost(
    x: Float[Array, "N"],
    instruments: list[Instrument],
) -> Float[Array, ""]:
    total = 0.0
    for ins in instruments:
        r      = einx.multiply("M, M -> M", ins.mask, ins.y - ins.H(x))
        Rinv_r = einx.dot("M K, K -> M", ins.R_inv, r)
        total += ins.weight * 0.5 * einx.dot("M, M ->", r, Rinv_r)
    return total

Tangent-Linear and Adjoint¶

Incremental 4DVar and the Gauss–Newton inner loop need the tangent-linear operator $H'(\boldsymbol{x})$ (forward) and its adjoint $H'(\boldsymbol{x})^\top$ (backward). Both are autodiff primitives — no hand-coding:

import jax

def tlm(H, x: Float[Array, "N"], dx: Float[Array, "N"]) -> Float[Array, "M"]:
    """Forward tangent-linear:  H'(x) · dx   (a JVP)."""
    _, dy = jax.jvp(H, (x,), (dx,))
    return dy

def adjoint(H, x: Float[Array, "N"], r: Float[Array, "M"]) -> Float[Array, "N"]:
    """Adjoint:  H'(x)ᵀ · r   (a VJP)."""
    _, vjp_fn = jax.vjp(H, x)
    (out,) = vjp_fn(r)
    return out

The adjoint test verifies the inner-product identity

up to floating-point tolerance. A failure means the gradient flowing back through $H$ is wrong, so downstream optimizers take incorrect steps.

import jax.numpy as jnp

def adjoint_test(H, x, u: Float[Array, "N"], v: Float[Array, "M"]) -> bool:
    lhs = jnp.vdot(tlm(H, x, u), v)        # ⟨H'u, v⟩
    rhs = jnp.vdot(u, adjoint(H, x, v))    # ⟨u, (H')ᵀv⟩
    return bool(jnp.allclose(lhs, rhs))

Observation-Error Covariance $\mathbf{R}$¶

$\mathbf{R}$ enters $J_{\text{obs}}$ only through $\mathbf{R}^{-1}$, which is applied lazily (e.g. via conjugate gradients) and never materialized. Common parameterizations:

• Diagonal — $\mathbf{R} = \operatorname{diag}(\boldsymbol{\sigma}^2)$. The default for heteroscedastic retrievals (per-pixel $\sigma$ from L2 metadata); $\mathbf{R}^{-1}\boldsymbol{r}$ is an elementwise divide.

• Block-diagonal — one block $\mathbf{R}_i$ per instrument when fusing.

• Structured — a spatial-correlation kernel (e.g. Matérn) when observation errors are correlated, as for high-resolution imaging spectrometers. Applied matrix-free.

def apply_R_inv_diag(
    variances: Float[Array, "M"],   # σ²
    r:         Float[Array, "M"],
) -> Float[Array, "M"]:
    """R⁻¹ r for diagonal R — no inverse materialized."""
    return einx.divide("M, M -> M", r, variances)

When the Observation Model Is Wrong¶

If $H$ is misspecified — wrong averaging kernel, missing instrument bias, ignored representation error — the MAP estimate $\boldsymbol{x}^\star$ converges to the wrong answer with a confident-looking posterior. Three diagnostics catch this:

• Innovation statistics — $\boldsymbol{d} = \boldsymbol{y} - H(\boldsymbol{x}_b)$ should be roughly mean-zero and consistent with $\sqrt{\operatorname{diag}(\mathbf{H}\mathbf{B}\mathbf{H}^\top + \mathbf{R})}$. Large systematic offsets indicate bias.

• Posterior-vs-prior shift — if the data move the posterior much further than $\sigma_{\text{post}}$ from the prior, the $\mathbf{B}$ or $\mathbf{R}$ assumptions are violated.

• Cross-instrument disagreement — per-instrument residual statistics should be consistent; large disagreements flag uncorrected bias and motivate joint bias estimation.