Ocean state estimation, from the generative story to each box

The objects above define a state-space model, a hidden Markov chain with Gaussian noise. Read top to bottom, it is a recipe for simulating an ocean and the satellite data it would produce:

The recipe is equivalent to a joint density, which is everything step 2 needs:

The arrow from step 1 to steps 3 and 5 in the figure is this recipe run many times. In DA vocabulary that is an observing system simulation experiment (OSSE): a "nature run" \(\vect{x}^\text{true}_{0:T}\) plus synthetic maps \(\vect{y}_{1:T}\) made with the real \(\mat{H}_k\) and \(\gmat{\Sigma}_{\varepsilon,k}\). Each OSSE gives a pair \((\vect{x}_{0:T}, \vect{y}_{1:T})\) whose answer is known exactly.

Bayes' rule turns the joint density into the target. Everything in classical DA is a way of computing, approximating or summarising one of two posteriors:

The exact case: the Kalman filter

If \(\mathcal{F}\) were linear, \(\mathcal{F}[\vect{x}; \gvect{\theta}] = \mat{F}\vect{x}\), every density above would stay Gaussian and the filter would be exact. Write \(\vect{x}_{k \mid k-1}, \gmat{\Sigma}_{x,k \mid k-1}\) for the forecast mean and covariance and \(\vect{x}_{k \mid k}, \gmat{\Sigma}_{x,k \mid k}\) for the analysis:

Two facts kill it for the ocean. QG advection \(J(\psi,q)\) is quadratic, so \(\mathcal{F}\) is nonlinear. And \(\gmat{\Sigma}_{x,k \mid k-1}\) has \(D_x^2 \approx 2.7\times 10^8\) entries (about 1 GB in float32) and propagating it costs \(D_x\) tangent-linear runs per day. Every practical method below drops one of these two ingredients: 4D-Var drops the covariance, the EnKF drops the full rank.

3D-Var: one analysis as an optimisation

For a single time, the Kalman analysis mean is also the minimiser of a quadratic cost. Write \( \Vert \vect{r} \Vert^2_{\mat{A}} = \vect{r}^\top \mat{A} \vect{r}\):

This is \(-\log\) of prior times likelihood, so its minimiser is the MAP estimate. With \(\gvect{\mu}_x = \vect{x}_{k \mid k-1}\) and \(\gmat{\Sigma}_x = \gmat{\Sigma}_{x,k \mid k-1}\) it equals \(\vect{x}_{k \mid k}\) exactly. 3D-Var ignores time: it is the L2 → L3 optimal interpolation written as a cost function.

Strong-constraint 4D-Var: the whole window through the model

Assume the model is perfect, \(\gmat{\Sigma}_\eta = \mat{0}\). Then the trajectory is fixed by \(\vect{x}_0\) and \(\gvect{\theta}\), and the negative log posterior is a function of those alone:

The minimiser \((\vect{x}_0^\star,\gvect{\theta}^\star)\) is the MAP estimate, and running the model from it gives the analysed trajectory. The minimisation needs \(\nabla_{\vect{x}_0}\mathcal{J}\). Let \(\vect{d}_k = \gmat{\Sigma}_{\varepsilon,k}^{-1}(\vect{y}_k - \mat{H}_k\vect{x}_k)\) be the weighted misfit on day \(k\). Run the adjoint backwards, starting from \(\vect{x}^\dagger_{T+1} = \vect{0}\):

The backward recursion is reverse-mode automatic differentiation of the model. In JAX it is jax.grad of \(\mathcal{J}\); in a Fortran ocean model it is a separately maintained adjoint code. One gradient costs about one forward run plus one adjoint run, two to three forward runs in total. Adding \(\gvect{\theta}\) to the control vector, as above, is parameter estimation by the same machinery.

In practice the optimisation variable is changed to \(\vect{x}_0 = \gvect{\mu}_x + \gmat{\Sigma}_x^{1/2}\gvect{\xi}\), so the background term becomes \(\frac{1}{2} \Vert \gvect{\xi} \Vert^2\) and the problem is well conditioned. Applying \(\gmat{\Sigma}_x^{1/2}\) matrix-free (a Kronecker or spectral factor of the Matérn kernel) is a linear-algebra job, which is where gaussx sits.

import einx
import equinox as eqx
import finitevolx as fvx
import gaussx as gx
import jax
import jax.numpy as jnp
import lineax as lx
import optimistix as optx
from jaxtyping import Array, Bool, Float

# Shapes: T = 10 daily maps (time); H, W = 130 × 130 grid with ghost ring
#         (height, width); N = H·W = 16,900 cells (cell); V covariates
#         (variable); P = 2 parameters L_d, ν_visc (param)

Op = lx.AbstractLinearOperator


class Window(eqx.Module):
    """One 10-day window: maps y, masks m, errors σ², background μ_x, priors."""

    maps: Float[Array, "time height width"]
    masks: Bool[Array, "time height width"]
    obs_var: Float[Array, "time height width"]
    background: Float[Array, "cell"]
    covariates: Float[Array, "variable height width"]
    forcing: Float[Array, "time height width"]
    prior_cov: Op
    param_mean: Float[Array, "param"]
    param_cov: Op
    grid: fvx.CartesianGrid2D


def strong_4dvar_cost(
    control: tuple[Float[Array, "cell"], Float[Array, "param"]],
    window: Window,
) -> Float[Array, ""]:
    whitened, params = control
    height = window.maps.shape[1]
    # user (taxonomy page)
    qg_params = QGParams(deformation_radius=params[0], viscosity=params[1])

    # --- background ---
    # today* (gaussx)
    # Σ_x^{1/2}
    prior_sqrt = gx.sqrt(window.prior_cov)            # (N, N) PSD op
    # x₀ − μ_x = Σ_x^{1/2} ξ
    increment = prior_sqrt.mv(whitened)               # (N,) → (N,)
    # x₀ = μ_x + Σ_x^{1/2} ξ
    initial_flat = window.background + increment      # (N,) → (N,)
    initial_state = einx.rearrange(                   # (N,) → (H, W)
        "(height width) -> height width", initial_flat, height=height
    )
    # ‖ξ‖²
    whitened_norm = einx.dot("cell, cell ->", whitened, whitened)  # → ()
    # 𝒥_b = ½‖ξ‖²
    background_cost = 0.5 * whitened_norm             # → ()

    # --- parameters ---
    # θ − μ_θ
    param_departure = params - window.param_mean      # (P,) → (P,)
    # today (gaussx)
    # (θ − μ_θ)ᵀ Σ_θ⁻¹ (θ − μ_θ)
    param_quad = gx.quadratic_form(window.param_cov, param_departure)  # → ()
    # 𝒥_θ = ½‖θ − μ_θ‖²_{Σ_θ⁻¹}
    param_cost = 0.5 * param_quad                     # → ()

    # --- observations, one day per scan step ---
    def assimilate_day(state, day):
        observed_map, mask, obs_var, daily_forcing = day
        # user (qg_day, taxonomy page)
        # x_k = 𝓕[x_{k−1}, u^cov, u_k; θ], 48 RK4 steps
        state = qg_day(                               # (H, W) → (H, W)
            state, window.covariates, daily_forcing, qg_params, window.grid
        )
        # y_k − x_k
        misfit = observed_map - state                 # (H, W)
        # r_k = m_k ⊙ (y_k − x_k)
        residual = jnp.where(mask, misfit, 0.0)       # (H, W)
        # r_k²
        squared = residual**2                         # (H, W)
        # r_k² / σ_k²
        weighted = squared / obs_var                  # (H, W)
        # ½ ∑ r_k² / σ_k²
        day_sum = einx.sum("height width ->", weighted)  # → ()
        day_cost = 0.5 * day_sum                      # → ()
        return state, day_cost

    days = (window.maps, window.masks, window.obs_var, window.forcing)
    # 𝒥_o = ∑_k ½‖m_k ⊙ (y_k − x_k)‖²_{Σ_ε,k⁻¹}
    _, day_costs = jax.lax.scan(assimilate_day, initial_state, days)  # → (T,)
    obs_cost = einx.sum("time ->", day_costs)         # (T,) → ()

    # 𝒥 = 𝒥_b + 𝒥_θ + 𝒥_o
    total_cost = background_cost + param_cost + obs_cost  # → ()
    return total_cost


# not pinned (optimistix)
# (ξ⋆, θ⋆) = argmin 𝒥(ξ, θ); ∇𝒥 by reverse mode is the adjoint recursion
solution = optx.minimise(                             # → ((N,), (P,))
    strong_4dvar_cost,
    optx.BFGS(rtol=1e-6, atol=1e-6),
    (whitened_init, params_init),
    args=window,
)
whitened_star, params_star = solution.value

Markers: today ships on main (checked 24 Oct 2026: finitevolX 831c69f, gaussx eeda235, somax 9cb19b3); today* ships, arguments sketched; user is defined on these pages; not pinned is a library whose API was not checked. Shape column letters are defined in the header.

vardax ships a strong-constraint 4D-Var (StrongFourDVar) but no Σ_x^1/2 change of variables, so the sketch builds that step with gaussx.

Writing \(\mat{H}_k\) as the mask \(\vect{m}_k\) applied to the full grid keeps every day the same shape, (height, width), so the scan needs no padding. The hyperparameters \(\gvect{\alpha}\) live inside prior_cov, obs_var and param_cov: fixed here, tuned in step 6.

Weak-constraint 4D-Var: admitting the model is wrong

With \(\gmat{\Sigma}_\eta \neq \mat{0}\) the whole trajectory becomes the control variable, and the cost is exactly the negative log of the smoothing posterior:

Strong-constraint 4D-Var is the limit \(\gmat{\Sigma}_\eta \to \mat{0}\). The control vector grows from \(D_x\) to \((T+1)D_x\), but the terms decouple in time, so the days can be evaluated in parallel.

Uncertainty from 4D-Var

4D-Var returns a mode, not a distribution. The Laplace approximation recovers a Gaussian around it, with the Gauss–Newton Hessian. Let \(\mat{F}_{k \leftarrow 0} = \partial \vect{x}_k / \partial \vect{x}_0 = \mat{F}_k \ldots \mat{F}_1\):

It is never formed. Each Hessian–vector product is one tangent-linear plus one adjoint run, and Lanczos extracts its leading eigenpairs: the directions the ten maps actually constrain.

The ensemble Kalman filter: sampling the covariance

The EnKF keeps the Kalman equations and replaces \(\gmat{\Sigma}_{x,k \mid k-1}\) with a sample. Run \(N_{\mathcal{E}}\) members \(\vect{x}^{(e)}_k\), say \(N_{\mathcal{E}} = 32\), each through the full nonlinear model. Stack the scaled anomalies into \(\tilde{\mat{X}}_k \in \mathbb{R}^{D_x\times N_{\mathcal{E}}}\):

The gain becomes \(\mat{G}_k = \tilde{\mat{X}}_k(\mat{H}_k\tilde{\mat{X}}_k)^\top(\mat{H}_k\tilde{\mat{X}}_k(\mat{H}_k\tilde{\mat{X}}_k)^\top + \gmat{\Sigma}_{\varepsilon,k})^{-1}\), and each member is updated with it. The matrix to invert is \(D_{y,k}\times D_{y,k}\), or \(N_{\mathcal{E}}\times N_{\mathcal{E}}\) in the ensemble-transform form. The sample covariance has rank at most \(N_{\mathcal{E}} - 1 = 31\) against \(D_x = 16\ 384\), so two fixes are always needed. Localisation multiplies the covariance elementwise by a compact taper \(\mat{C}_\text{loc}\) (for example 150 km), \(\gmat{\Sigma}_{x,k \mid k-1} \approx \mat{C}_\text{loc} \odot \tilde{\mat{X}}\tilde{\mat{X}}^\top\), which removes spurious long-range correlations. Inflation scales the anomalies by \(\alpha_\text{infl} \gtrsim 1\) to stop the spread collapsing. Parameters are estimated by appending \(\gvect{\theta}\) to each member's state. The ensemble smoother (EnKS) gives the reanalysis.

Cost of the oracle

With 50 iterations, 4D-Var costs about 50 × 3 × 480 ≈ 7 × 10⁴ QG steps. At 128² that is seconds to minutes on one GPU. The "hours" on the figure describe a realistic primitive-equation model at 1/12°, where one model day can take minutes on hundreds of cores. For this example, step 2 is cheap enough to be the oracle for everything to its right.

Look at where the cost of step 2 lives: in \(\mathcal{F}\), evaluated hundreds of times per iteration, and in its adjoint. Step 3 replaces that one object. A network \(\mathcal{F}[\cdot; \gvect{\theta}_\text{emu}]\) with weights \(\gvect{\theta}_\text{emu}\) learns the one-day map from the step-1 nature runs:

With \(T = 1\) this is one-step training; with \(T > 1\) it is a rollout loss (somax's rollout_loss on a scanned cycle). Its reference is step 1: the emulator is scored against the simulator it replaces, not against data.

The figure's "skip if cheap" applies here. One emulator call replaces 48 QG steps, but 48 QG steps at 128² already take milliseconds in JAX, and the QG model is already differentiable. For this running example the emulator buys little. It pays off when \(\mathcal{F}\) is a large ocean model with no adjoint.

Step 4 is step 2 with one substitution, \(\mathcal{F}[\cdot; \gvect{\theta}] \to \mathcal{F}[\cdot; \gvect{\theta}_\text{emu}]\). Nothing else changes: same \(\gmat{\Sigma}_x\), \(\gmat{\Sigma}_{\varepsilon,k}\), \(\mat{H}_k\), same optimiser or ensemble.

Its reference is step 2: the emulator-based analysis \(\hat{\vect{x}}^\text{emu}_{0:T}\) is compared with the classical analysis \(\hat{\vect{x}}_{0:T}\) on the same windows and the same maps, and in an OSSE both are compared with \(\vect{x}^\text{true}\). Each query still solves an optimisation, so step 4 still needs the observations, the prior and a solver at run time. That is the difference from step 5.

Steps 2 and 4 solve a new optimisation for every window. Step 5 learns the answer as a function of the data. A conditional density \(q(\cdot; \gvect{\nu})\) with weights \(\gvect{\nu}\) is trained on step-1 OSSE pairs:

Minimising this expectation minimises \(\mathbb{E}_{\vect{y}}[\mathrm{KL}[p(\cdot \mid \vect{y}) \Vert q(\cdot \mid \vect{y}; \gvect{\nu})]]\), so a flexible \(q(\cdot; \gvect{\nu})\) converges to the true posterior. The simplest case is a Gaussian with fixed variance and learned mean \(\gvect{\mu}_{x \mid y}(\vect{y}; \gvect{\nu})\). The loss becomes mean squared error, whose minimiser is the posterior mean:

That is the "direct analysis": a network from ten SSH maps to the state, at the cost of one forward pass.

The formulation contains its own test. If steps 1 and 2 are right, the innovations, \(\vect{d}_k = \vect{y}_k - \mat{H}_k\vect{x}_{k \mid k-1}\), must be zero-mean with a known covariance, because they are the sum of two independent Gaussian errors:

Two scalar checks follow. With \(D_y = \sum_k D_{y,k}\) observations, the minimum of a correctly specified cost satisfies \(\mathbb{E}[2 \mathcal{J}(\vect{x}^\star)]\approx D_y\), so \(2\mathcal{J}^\star/D_y\) should be near 1. These checks are how the hyperparameters \(\gvect{\alpha}\) get tuned. The Desroziers identity checks \(\gmat{\Sigma}_\varepsilon\) on its own: \(\mathbb{E}[(\vect{y}_k - \mat{H}_k \vect{x}_{k \mid k})(\vect{y}_k - \mat{H}_k \vect{x}_{k \mid k-1})^\top] \approx \gmat{\Sigma}_{\varepsilon,k}\).

Each box downstream is improved against its reference. The formulation fixes those references precisely, because each step approximates one object defined earlier: