Strong-Constrained 4DVar

Strong-constraint 4DVar extends 3DVar from a single snapshot to a time window: observations $\boldsymbol{y}_t$ at times $t = 0, 1, \ldots, T$ are related to a single initial state $\boldsymbol{u}_0$ through known dynamics Le Dimet & Talagrand, 1986Talagrand & Courtier, 1987. The label strong-constraint means the dynamics are treated as exact — any $\boldsymbol{u}_0$ deterministically produces the whole trajectory, and model error is excluded. Systems that need a model-error term use weak-constraint 4DVar instead.

The control variable is the initial state $\boldsymbol{u}_0$ (and optionally the parameters $\boldsymbol{\theta}$); the trajectory follows from the dynamical model.

Problem Formulation¶

Under the strong constraint the trajectory is a deterministic function of the initial state, $\boldsymbol{u}_t = M_t(\boldsymbol{u}_0)$, where $M_t$ propagates $\boldsymbol{u}_0$ from $t_0$ to $t$. The dynamics enter the cost through the composition $H_t \circ M_t$ — the observation operator applied to the propagated state. The generative model is

with Gaussian components

The 4DVar Cost¶

The MAP estimate of $\boldsymbol{u}_0$ minimises the negative log-posterior — the 4DVar cost — a background term plus an observation term summed over the window:

Minimising requires the gradient, which by the chain rule pulls back through every $M_t$:

The Adjoint Problem¶

Assembling (4) directly would require the dense Jacobians $(M'_t)^\top$. Instead the gradient is computed by a single backward sweep of an adjoint variable $\boldsymbol{\lambda}_t$, initialised at the final time and accumulated as it propagates backward:

Algorithm — Continuous-Time Adjoint¶

When the dynamics are a continuous flow $\partial_t \boldsymbol{u} = \boldsymbol{f}(\boldsymbol{u}, t, \boldsymbol{\theta})$, the flow map is an ODE solve, $M_t(\boldsymbol{u}_0) = \mathrm{ODESolve}(\boldsymbol{f}, \boldsymbol{u}_0, \boldsymbol{\theta}, [t_0, t])$, and the discrete adjoint sweep (5) becomes a backward ODE.

Data¶

Dynamical and Adjoint Models¶

The forward dynamics and its adjoint (the transpose-Jacobian action), obtained for free from a VJP:

import jax
from typing import Callable

def init_f_adjoint(f: Callable) -> Callable:
    """Adjoint of the dynamics via VJP:  λ ↦ (∂f/∂u)ᵀ λ — no hand-coding."""
    def f_adj(u, lam):
        _, vjp = jax.vjp(f, u)
        return vjp(lam)[0]
    return f_adj

Initial Condition¶

We need an initial guess for the state and (optionally) the parameters:

Integrate the Forward Model¶

Build the time grid that matches the observation times and roll the state forward:

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

dt: float = 0.01
time_steps: Float[Array, "Nt"] = jnp.arange(0, num_steps) * dt

state_sol = ode_solve(f, u0, theta, time_steps)     # forward rollout
U: Float[Array, "Nt Du"] = state_sol.u              # state trajectory

so that

Integrate the Backward Adjoint Model¶

Run the solver in reverse over the flipped time grid, accumulating the adjoint (7):

time_steps_reverse: Float[Array, "Nt"] = time_steps[::-1]
# λ_{t+1} = f*(u_t, λ_t, y, θ)  — backward sweep accumulating the obs misfit

Gradient Steps¶

In modern code you rarely write the sweep by hand: define the cost (3), wrap the rollout in a differentiable ODE solve, and call jax.grad(J)(u0) — the dynamical-model adjoints compose the pullback automatically.

Computational Considerations¶

Long Assimilation Windows¶

For $T \gtrsim 100$ cycles, memory dominates. Pick a diffrax adjoint to trade memory against recompute:

• BacksolveAdjoint() — solves the adjoint ODE backward at $O(1)$ memory in the trajectory length (watch for reverse-time stiffness).

• RecursiveCheckpointAdjoint(checkpoints=...) — $O(\sqrt{T})$ memory, modest recompute; the safe default.

• Shorter windows with cycling — split a long window into chained shorter assimilations.

Nonlinear Dynamics¶

When $M_t$ is nonlinear the cost (3) is non-convex with multiple local minima. Mitigations: multi-start optimisation, warm-starting from the previous cycle, or switching to incremental 4DVar — which solves a sequence of linearised inner problems (each an OI/BLUE, exactly as in 3DVar’s Gauss–Newton loop) and is the more robust choice for production.

Posterior¶

A Laplace approximation gives a Gaussian posterior whose covariance is the inverse Gauss–Newton Hessian of (3), evaluated through the trajectory. Forming it exactly needs many Krylov iterations (tens of matrix–vector products, each a full forward+adjoint pass), so for production uncertainty quantification, incremental 4DVar or an ensemble / amortized posterior is cheaper.

When Strong-Constraint 4DVar Is Appropriate¶