3DVar — Three-Dimensional Variational Assimilation

3DVar extends Optimal Interpolation to nonlinear observation operators. When $H$ is linear the posterior is Gaussian and the analysis is a single matrix expression (the BLUE). When $H$ is nonlinear that closed form is lost, so instead of solving for the posterior exactly we compute its mode — the maximum a posteriori (MAP) estimate — by iteratively minimising a variational cost Lorenc, 1986Talagrand, 1997.

The “3D” denotes three spatial dimensions at a single time ($T = 0$, no temporal evolution); the dynamical extension is 4DVar. As before, the code is library-agnostic — jaxtyping for shapes, einx for array ops, JAX autodiff for the tangent-linear/adjoint, and a general optimiser (optimistix) for the minimisation.

From the MAP to the Variational Cost¶

The MAP estimate maximises the posterior, equivalently minimising its negative log:

The two terms trade off the background (prior) against the observations (via the observation model), each weighted by its inverse covariance.

Gradient and Hessian¶

Minimisation needs the gradient, and Gauss–Newton needs (an approximation of) the Hessian:

where $H'(\boldsymbol{x})$ is the tangent-linear of the observation operator (its adjoint $H'(\boldsymbol{x})^\top$ comes from autodiff — see the observation-model note).

The Gauss–Newton Inner Loop = Iterated BLUE¶

Each Gauss–Newton iteration solves the normal equations for an increment $\delta\boldsymbol{x}$ and updates the iterate:

where $H_k' = H'(\boldsymbol{x}_k)$. The increment is exactly the minimiser of the linearised cost about $\boldsymbol{x}_k$ — a linear-Gaussian problem with innovation $\boldsymbol{d}_k = \boldsymbol{y} - H(\boldsymbol{x}_k)$. In other words, 3DVar is a sequence of BLUE analyses, each one re-linearising $H$ about the current iterate.

Each inner solve uses the same matrix-free machinery as OI — never forming $\mathbf{B}$ or the Hessian, only their action on vectors:

import jax
from jax.scipy.sparse.linalg import cg
from jaxtyping import Array, Float

def gauss_newton_step(
    xk: Float[Array, "N"],
    xb: Float[Array, "N"],
    y:  Float[Array, "M"],
    H,             # nonlinear observation operator  x -> H(x)
    apply_B_inv,   # v -> B⁻¹ v
    apply_R_inv,   # v -> R⁻¹ v
) -> Float[Array, "N"]:
    Hk, Hk_lin = jax.linearize(H, xk)              # Hk = H(xk);  Hk_lin(v) = H'_k v
    Hk_adj     = jax.linear_transpose(Hk_lin, xk)  # (H'_k)ᵀ

    grad = apply_B_inv(xk - xb) - Hk_adj(apply_R_inv(y - Hk))[0]   # ∇J(xk)
    def hess(v: Float[Array, "N"]) -> Float[Array, "N"]:          # B⁻¹ + H'_kᵀ R⁻¹ H'_k
        return apply_B_inv(v) + Hk_adj(apply_R_inv(Hk_lin(v)))[0]

    dx, _ = cg(hess, -grad)                                       # solve normal equations
    return xk + dx

Preconditioning: the control-variable transform¶

The Hessian (4) contains $\mathbf{B}^{-1}$, which is badly conditioned for smooth priors. Operational 3DVar minimises in a whitened control variable $\boldsymbol{\chi}$ defined by $\boldsymbol{x} = \boldsymbol{x}_b + \mathbf{B}^{1/2}\boldsymbol{\chi}$ Bannister, 2008, which turns the background term into an identity and preconditions the problem:

The background Hessian contribution becomes $\mathbf{I}$, so the conjugate- gradient inner loop converges in far fewer iterations.

Implementation¶

The Gauss–Newton form is a nonlinear least-squares problem: stack the whitened residuals so that $\tfrac{1}{2}\|\boldsymbol{r}(\boldsymbol{x})\|^2 = J(\boldsymbol{x})$, and hand it to a least-squares solver.

import optimistix as optx
import jax.numpy as jnp
from jaxtyping import Array, Float

def threedvar(
    xb, y, H, whiten_B, whiten_R,            # whiten_B: v -> B^{-1/2} v ;  whiten_R: v -> R^{-1/2} v
    *, solver=optx.GaussNewton(rtol=1e-8, atol=1e-8),
):
    def residual(x: Float[Array, "N"], args) -> Float[Array, "N+M"]:
        return jnp.concatenate([whiten_B(x - xb), whiten_R(y - H(x))])
    sol = optx.least_squares(residual, solver, xb)    # ½‖r‖² = J(x); warm-start at the background
    return sol.value                                  # MAP estimate x*

Convergence is typically tens of iterations for well-conditioned problems. To differentiate through the optimum (e.g. to learn $\mathbf{B}$ or $\mathbf{R}$), use an implicit adjoint (optx.ImplicitAdjoint): the implicit function theorem gives the gradient at the solution in $O(1)$ memory, without unrolling the optimiser.

Posterior via the Laplace Approximation¶

For a Gaussian likelihood near the MAP, a second-order (Laplace) expansion of $J$ gives a Gaussian posterior whose covariance is the inverse Gauss–Newton Hessian:

This is exactly the OI posterior covariance with $H$ linearised at the MAP. Sampling and marginal variances are matvecs with $(\mathbf{P}^\star)^{1/2}$ — nothing is materialised. For strongly nonlinear operators the Laplace approximation under-reports uncertainty; reach for an ensemble covariance or an amortized posterior for better calibration.

When 3DVar Is Appropriate¶