4D-Var: fitting a model trajectory to a window of observations

Four-dimensional variational assimilation (4D-Var) asks one question: which initial state, run forward through the model, best explains every observation in a time window, without straying far from what we believed before? This page is the background for the learned data assimilation page: the intuition, the mathematics, how the gradient is computed, three classic algorithms, and how it maps onto an ocean problem.

Intuition

A 3D-Var or optimal-interpolation analysis corrects one snapshot with the observations at that time. 4D-Var corrects a trajectory: it chooses the starting point, lets the model run, and scores the whole run against everything observed in the window. Two consequences follow.

• Observations spread information through the dynamics. An observation at day 7 constrains day 0 along exactly the directions the model carries from day 0 to day 7. Places never observed are still corrected, because the flow brings observed water to them or them to it.
• The model is a hard constraint. The analysed trajectory is a true model solution: balanced, smooth in time, consistent with the physics. That is also its weakness: if the model is wrong, the analysis has nowhere to put the error (the weak-constraint form relaxes this).

A mechanical picture: each observation is a spring pulling the trajectory toward it, with stiffness \(1/\sigma_\varepsilon^2\); the background is a spring on the initial state with stiffness \(\gmat{\Sigma}_x^{-1}\); the model is a rigid rod connecting all times. 4D-Var finds where the rod comes to rest.

The running example is Lorenz-96: 40 values around a latitude circle, chaotic, a standard data-assimilation test bed (the one vardax and somax use). The window is 10 observation times of 10 model steps each; every second site is observed with error \(\sigma_\varepsilon = 1\); the background error is \(\sigma_x = 1\) with a correlation length of 2 sites.

Where it comes from

In the terms of the six-step page, 4D-Var is step 2 for one particular set of densities. The prior on the initial state, the transition, and the likelihood are

The point-mass transition makes every state a function of \(\vect{x}_0\), so the posterior is over \(\vect{x}_0\) alone, and its negative log is, up to a constant, the 4D-Var cost. 4D-Var returns that posterior's mode (the MAP estimate). It is the right summary when the posterior has one dominant, near-Gaussian peak, and the Laplace approximation around the mode then gives its spread. Replacing the point mass by \(\mathcal{N}(f_k(\vect{x}_{k-1}), \gmat{\Sigma}_\eta)\) gives weak-constraint 4D-Var, in which the model errors become extra controls.

The mathematics

Notation

vardax also writes the state as \(\vect{x}\), so the letters here match its code; \(D_x\) and \(D_y\) are the state and per-time observation sizes.

The cost

With the change of variable \(\vect{x}_0 = \gvect{\mu}_x + \gmat{\Sigma}_x^{1/2}\gvect{\xi}\) the background term becomes \(\frac{1}{2}\Vert \gvect{\xi}\Vert^2\): all directions of \(\gvect{\xi}\) are equally likely a priori, the problem is far better conditioned, and \(\gmat{\Sigma}_x^{-1}\) is never needed. This "control-variable transform" is how operational 4D-Var is run.

The gradient: the adjoint recursion

Write \(\nabla_{\vect{x}_k}\mathcal{J}_{\text{o},k} = \mat{H}_k^\top \gmat{\Sigma}_{\varepsilon,k}^{-1}(h_k(\vect{x}_k) - \vect{y}_k)\) for the observation misfit at time \(k\), pulled back to the state. Starting from \(\vect{x}^\dagger_{T+1} = \vect{0}\), sweep backwards:

The adjoint carries each misfit backwards along the trajectory, through the transposed tangent-linear model, accumulating as it goes. Its cost is about one forward run, independent of the number of controls: that is why 4D-Var is feasible for millions of state values.

import einx
import jax
import jax.numpy as jnp
import lineax as lx
import optimistix as optx
import vardax

# Shapes: N = 40 state values (state); T = 10 observation times (time);
#         M = 20 observed values per time (obs)


def rollout(forward, initial, num_times):
    def interval(state, _):
        # user (forward.step: one observation interval, vardax's ForwardModel)
        # x_k = f_k(x_{k−1})
        state = forward.step(state, forward.dt)       # (N,) → (N,)
        return state, state

    # x_{1:T}: the model trajectory from x₀
    _, states = jax.lax.scan(interval, initial, None, length=num_times)  # → (T, N)
    return states


def cost(whitened, args):
    forward, background, prior_sqrt, observations, obs_index, obs_std = args
    # x₀ − μ_x = Σ_x^{1/2} ξ
    increment = prior_sqrt.mv(whitened)               # (N,) → (N,)
    # x₀ = μ_x + Σ_x^{1/2} ξ
    initial = background + increment                  # (N,)
    states = rollout(forward, initial, observations.shape[0])  # → (T, N)
    # h_k(x_k): the observed entries of each state
    predicted = states[:, obs_index]                  # (T, N) → (T, M)
    # y_k − h_k(x_k)
    misfit = observations - predicted                 # (T, M)
    # r_k = Σ_ε^{−1/2} (y_k − h_k(x_k)), diagonal Σ_ε = σ_ε² I
    residual = misfit / obs_std                       # (T, M)
    # ‖r_k‖² summed over the window
    misfit_sq = einx.sum("time obs ->", residual**2)  # (T, M) → ()
    # 𝒥_o = ½ ∑_k ‖y_k − h_k(x_k)‖²_{Σ_ε⁻¹}
    obs_cost = 0.5 * misfit_sq                        # → ()
    # 𝒥_b = ½ ‖x₀ − μ_x‖²_{Σ_x⁻¹} = ½ ‖ξ‖²
    background_cost = 0.5 * einx.dot("state, state ->", whitened, whitened)  # → ()
    # 𝒥(ξ) = 𝒥_b + 𝒥_o
    return background_cost + obs_cost

Curvature, convergence, and spread

The Gauss–Newton Hessian, in whitened variables, and the matching Laplace covariance of the analysis are

Its eigenvalues are 1 plus "how strongly the window's observations constrain each direction"; CG converges in about as many iterations as there are eigenvalues well above 1, which is why the whitened form converges fast. The same Hessian, through a few Lanczos vectors, gives the posterior spread. A last check needs no truth: at the minimum, \(2\mathcal{J}^\star / N_\text{obs} \approx 1\) for \(N_\text{obs}\) observations if \(\gmat{\Sigma}_x\) and \(\gmat{\Sigma}_\varepsilon\) are right. Here \(N_\text{obs} = 200\) and \(2\mathcal{J}^\star / N_\text{obs} = \) 1.01.

Computing the gradient

Every 4D-Var algorithm needs \(\nabla_{\vect{x}_0}\mathcal{J}\), or products with the tangent-linear model and its adjoint. There are several ways to get them, and they differ in exactness, memory and run time.

Two adjoints appear in a differentiable 4D-Var and should not be confused. The forward adjoint differentiates the model trajectory, to get \(\nabla_{\gvect{\xi}}\mathcal{J}\) for the optimiser. The minimiser adjoint differentiates the analysis \(\vect{x}_0^\star(\vect{y}, \gvect{\theta})\) with respect to its inputs, which is needed only when 4D-Var sits inside a larger trained system (learned priors, hyperparameter tuning). There, unrolling the optimiser is the expensive route; the implicit function theorem gives \(\partial \vect{x}_0^\star / \partial \gvect{\theta} = -(\nabla^2_{\vect{x}_0}\mathcal{J})^{-1} \nabla^2_{\vect{x}_0 \gvect{\theta}}\mathcal{J}\) at the solution, which is what optimistix.ImplicitAdjoint does.

Measured on Lorenz-96

Gradient of the cost over a 50-interval window (500 model steps), CPU, after compilation. Memory is the compiled temporary buffer size reported by XLA.

Checkpointing each interval cuts working memory by about 23× and returns the identical gradient; on this small CPU problem it is even faster, because the smaller working set stays in cache (on a GPU with a large state, expect roughly one extra forward run of recomputation instead). Forward mode needs one tangent run per control, 40 here, which is why it only pays when the control is small.

# Shapes: N state values (state); T observation times (time); M observed values (obs)

# --- (a) reverse mode through the scan: the discrete adjoint, exact ---
# ∇_ξ 𝒥: one forward run, one backward sweep, every state x_k stored
value, gradient = jax.value_and_grad(cost)(whitened, args)  # → (), (N,)

# --- (b) the same adjoint, written out with vjp ---
# ∇_{x_k} 𝒥_{o,k} = H_kᵀ Σ_ε⁻¹ (h_k(x_k) − y_k), scattered back to the state
obs_forcing = scatter_obs(-residual / obs_std, obs_index)  # (T, M) → (T, N)
previous = jnp.concatenate([initial[None], states[:-1]])  # → (T, N), x_{0:T−1}


def backward(carry, inputs):
    state_prev, forcing = inputs
    # x†_k = ∇_{x_k} 𝒥_{o,k} + F_{k+1}ᵀ x†_{k+1}   (carry holds F_{k+1}ᵀ x†_{k+1})
    lam = forcing + carry                             # (N,)
    # vjp of one interval at x_{k−1}: v ↦ F_kᵀ v
    _, interval_vjp = jax.vjp(lambda s: forward.step(s, forward.dt), state_prev)
    # F_kᵀ x†_k
    (carry,) = interval_vjp(lam)                      # (N,) → (N,)
    return carry, None


# backward in time, from k = T down to 1
grad_initial, _ = jax.lax.scan(                       # → (N,), ∇_{x₀} 𝒥_o
    backward, jnp.zeros(N), (previous, obs_forcing), reverse=True
)
# ∇_ξ 𝒥 = ξ + (Σ_x^{1/2})ᵀ ∇_{x₀} 𝒥_o
gradient_by_hand = whitened + prior_sqrt.T.mv(grad_initial)  # (N,)

# --- adjoint test: ⟨F δx, x†⟩ = ⟨δx, Fᵀ x†⟩ to round-off ---
_, tangent = jax.jvp(step_fn, (state,), (perturbation,))  # (N,) → (N,)
_, step_vjp = jax.vjp(step_fn, state)
(cotangent,) = step_vjp(dual)                         # (N,) → (N,)
mismatch = jnp.vdot(tangent, dual) - jnp.vdot(perturbation, cotangent)  # → ()

# --- (c) checkpointing: store only the interval boundaries, recompute inside ---
checkpointed_step = jax.checkpoint(step_fn)           # same signature, less memory

# --- (d) a diffrax model: pick the adjoint per solve ---
# today (diffrax 0.7.0)
# discretise-then-optimise, binomial checkpoints: exact, the default
adjoint = dfx.RecursiveCheckpointAdjoint(checkpoints=64)
# today (diffrax 0.7.0)
# continuous adjoint ODE solved backwards: O(1) memory, approximate gradient
adjoint = dfx.BacksolveAdjoint()
# today (diffrax 0.7.0)
# forward mode: one tangent solve per control direction
adjoint = dfx.ForwardMode()

# Shapes: N state values (state); T times (time)

# --- the minimisation, as vardax runs it ---
# not pinned (optimistix)
# ξ⋆ = argmin_ξ 𝒥(ξ), quasi-Newton with ∇𝒥 from the adjoint
solution = optx.minimise(                             # → (N,)
    cost, optx.BFGS(rtol=1e-6, atol=1e-6), jnp.zeros(N), args=args, max_steps=200
)
# x₀⋆ = μ_x + Σ_x^{1/2} ξ⋆
analysis = background + prior_sqrt.mv(solution.value)  # (N,)

# --- the same through vardax ---
# today (vardax)
strong = vardax.StrongFourDVar(
    forward=forward,  # .step(z, dt), .dt: one observation interval
    obs_op=vardax.MaskedIdentity(),
    prior_mean=background,                            # (N,)
    prior_cov_op=prior_cov,                           # (N, N) PSD operator
    obs_cov_op=obs_cov,                               # (N, N) PSD operator
    minimiser=optx.BFGS(rtol=1e-6, atol=1e-6),
)
# today (vardax)
# batch: Batch1D(input=(1, T+1, N), mask=(1, T+1, N))
analysis = strong(batch)                              # → (1, N)

# Shapes: N state values (state); T·M whitened observations, flattened (obs_flat)


def observe(whitened):
    # h_win(ξ) = Σ_ε^{−1/2} h(f(μ_x + Σ_x^{1/2} ξ)), flattened over the window
    initial = background + prior_sqrt.mv(whitened)    # (N,)
    states = rollout(forward, initial, num_times)     # → (T, N)
    observed = states[:, obs_index] / obs_std         # (T, N) → (T, M)
    return einx.rearrange("time obs -> (time obs)", observed)  # → (T·M,)


whitened = jnp.zeros(N)                               # (N,), the background
for _ in range(NUM_OUTER):
    # --- outer loop: one nonlinear run, linearised around it ---
    # h_win(ξ) and its tangent linear v ↦ H_win v
    predicted, tangent = jax.linearize(observe, whitened)  # → (T·M,)
    # adjoint w ↦ H_winᵀ w
    adjoint = jax.linear_transpose(tangent, whitened)
    # d = Σ_ε^{−1/2} y − h_win(ξ): whitened innovations
    innovation = observations_white - predicted       # (T·M,)

    def gauss_newton(v):
        # (I + H_winᵀ H_win) v: one tangent-linear and one adjoint run
        (projected,) = adjoint(tangent(v))            # (N,) → (N,)
        return v + projected

    hessian = lx.FunctionLinearOperator(
        gauss_newton, jax.ShapeDtypeStruct((N,), whitened.dtype),
        lx.positive_semidefinite_tag,
    )
    # b = H_winᵀ d − ξ  (= −∇_{δξ} 𝒥 at δξ = 0)
    (pulled,) = adjoint(innovation)                   # (T·M,) → (N,)
    rhs = pulled - whitened                           # (N,)
    # --- inner loop: CG on the quadratic, no nonlinear runs ---
    # not pinned (lineax)
    # δξ = (I + H_winᵀ H_win)⁻¹ b
    step = lx.linear_solve(                           # → (N,)
        hessian, rhs, solver=lx.CG(rtol=1e-8, atol=1e-8, max_steps=NUM_INNER)
    ).value
    # ξ ← ξ + δξ, then relinearise
    whitened = whitened + step                        # (N,)

# today (vardax)
# vardax's version solves in x-space with Σ_x⁻¹ applied by CG (no Σ_x^{1/2} transform)
incremental = vardax.IncrementalFourDVar(
    forward, vardax.MaskedIdentity(), background, prior_cov, obs_cov,
    config=vardax.IncrementalConfig(n_outer=4, n_inner=25),
)

# Shapes: N state values (state); P basis functions (basis), P ≪ N;
#         E background samples (sample)

# today (vardax)
# Φ: leading EOFs of an ensemble, Σ_{z_Φ} = diag(explained variance)
basis = vardax.eof_basis(ensemble, n_modes=16)        # (E, N) → Φ (N, P)


def cost_reduced(coeffs, args):
    forward, background, observations, obs_index, obs_std = args
    # δx₀ = Φ z_Φ
    increment = basis.operg(0.0, coeffs)              # (P,) → (N,)
    # x₀ = μ_x + Φ z_Φ
    initial = background + increment                  # (N,)
    states = rollout(forward, initial, observations.shape[0])  # → (T, N)
    # r_k = Σ_ε^{−1/2} (y_k − h_k(x_k))
    residual = (observations - states[:, obs_index]) / obs_std  # (T, M)
    # 𝒥_o = ½ ∑_k ‖r_k‖²
    obs_cost = 0.5 * einx.sum("time obs ->", residual**2)  # → ()
    # Σ_{z_Φ}⁻¹ z_Φ, diagonal
    scaled = basis.prior_inv(coeffs)                  # (P,)
    # 𝒥_b = ½ z_Φᵀ Σ_{z_Φ}⁻¹ z_Φ
    prior_cost = 0.5 * einx.dot("basis, basis ->", coeffs, scaled)  # → ()
    return prior_cost + obs_cost


# today (vardax)
# full nonlinear cost over z_Φ with BFGS; the name says "incremental", the method does not
reduced = vardax.Incremental4DVar(
    forward, vardax.MaskedIdentity(), background, prior_cov, obs_cov, basis=basis
)
# P = 16 controls: forward mode now costs only 16 tangent runs per gradient
gradient = jax.jacfwd(cost_reduced)(coeffs, args)     # (P,) → (P,)

The three side by side

Application: the Gulf Stream window

The ocean state estimation page applies the same cost to a 1.5-layer quasi-geostrophic model of the Gulf Stream, with ten daily SSH maps over a 1024 km box. The objects map one to one:

Notes on vardax, from reading the code

Checked against vardax main (3e0f34a) on 24 Oct 2026. Several docs and the code disagree; each point is filed as a vardax issue.

• Incremental4DVar is not incremental. It minimises the full nonlinear cost over basis coefficients with BFGS (the MASSH formulation). The Gauss–Newton method is IncrementalFourDVar. A rename (for example ReducedBasis4DVar) would remove the trap (#96).
• No control-variable transform yet. IncrementalFourDVar applies \(\gmat{\Sigma}_x^{-1}\) with an inner CG solve at every Hessian product. The whitened form on this page is the proposed fix; the IncrementalConfig(cvt=True) shown in the docs does not exist (#97).
• The docs' incremental right-hand side is incomplete. The code uses \(-\nabla_{\vect{x}_0}\mathcal{J}\) of the full nonlinear cost at the current outer iterate, which is the correct Gauss–Newton right-hand side. The docs give \(\sum_k (\mat{H}_k \mat{F}_{k \leftarrow 0})^\top \gmat{\Sigma}_\varepsilon^{-1} \vect{d}_k\), which drops the background term \(-\gmat{\Sigma}_x^{-1}(\vect{x}_0 - \gvect{\mu}_x)\); the two agree only on the first outer loop, where \(\vect{x}_0 = \gvect{\mu}_x\). The code is right and the docs need the fix (#98).
• The posterior adapters omit the model. LaplaceCovariance and GaussNewtonHessian build \(\gmat{\Sigma}_x^{-1} + \mat{H}^\top \gmat{\Sigma}_\varepsilon^{-1} \mat{H}\) without \(\mat{F}_k\): a 3D-Var Hessian, not the 4D-Var one their docstrings state (#94).
• forward_adjoint is stored but not used. StrongFourDVar rolls out with jax.lax.scan, so the gradient is plain reverse mode whatever adjoint is passed (#95).
• Missing names in the docs: IncrementalFourDVar.posterior, vardax.utils.multi_start, and a no-argument LaplaceCovariance() (#99).

Code markers: today (vardax) names ship on vardax main (3e0f34a); today (diffrax 0.7.0) checked in that release; not pinned marks optimistix and lineax, whose APIs were not checked for this page. Lorenz-96 numbers come from a JAX run of the same algorithms, in double precision.