Weak-constraint 4D-Var: when the model is allowed to be wrong

Strong-constraint 4D-Var treats the model as exact and corrects only the initial state. Weak-constraint 4D-Var admits that the model makes errors, and estimates them alongside the state. This page follows the 4D-Var primer: the intuition, the probability, the mathematics, how the gradient changes, three formulations, and how it maps onto the ocean problem.

Intuition

In the 4D-Var primer's mechanical picture the model was a rigid rod: the trajectory had to be an exact model solution, so every misfit had to be blamed on the initial state. Weak constraint cuts the rod into one segment per observation interval and joins the segments with springs of stiffness \(\gmat{\Sigma}_\eta^{-1}\). The trajectory may now jump at the end of each interval, at a price.

• When the model is right, the springs are not needed: the jumps stay near zero and the result is close to strong constraint.
• When the model is wrong, strong constraint distorts the initial state to compensate, and over a long window it cannot. Weak constraint puts the blame where it belongs, at the times and places the model drifts.
• The jumps are information. Averaged over many windows, \(\gvect{\eta}^\star\) is a map of the model's systematic error.

The running example is Lorenz-96 again, now with a deliberately wrong model: the truth uses forcing \(F = 8\), the model \(F = 7\). The window is 20 observation times (twice the primer's), every second site is observed with \(\sigma_\varepsilon = 0.5\), and the background is as before.

The last column is the strong-constraint primer's consistency check, with \(N_\text{obs} = 400\) observations. A value of 7.6 instead of about 1 is the classic signature of model error: the data cannot be fitted by any model trajectory.

Where it comes from

On the six-step page's terms, only the transition density changes. Strong constraint used a point mass; weak constraint uses a Gaussian around the model forecast:

The posterior is now over the whole trajectory, \(p(\vect{x}_{0:T} \mid \vect{y}_{1:T})\), and weak-constraint 4D-Var returns its mode. That is the smoothing problem: for a linear model the answer equals the Kalman smoother's mean. The two limits are worth keeping in mind. As \(\gmat{\Sigma}_\eta \to \mat{0}\) the jumps are forbidden and strong constraint returns. As \(\gmat{\Sigma}_\eta \to \infty\) consecutive states decouple and the method becomes a separate 3D-Var at each observation time, ignoring the dynamics.

The mathematics

Notation

The cost, in two equivalent forms

With the model errors as controls (the forcing formulation):

With the states as controls (the state formulation), substitute \(\gvect{\eta}_k = \vect{x}_k - f_k(\vect{x}_{k-1})\):

The map between the two sets of controls is one-to-one, so the two costs have the same minima. They are different optimisation problems all the same: different conditioning, different paths, and, for a nonlinear model, possibly different local minima from the same start.

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


def weak_cost(control, args):
    """Forcing formulation: control (ξ, w₁:T), with η_k = Σ_η^{1/2} w_k."""
    whitened_initial, whitened_errors = control       # (N,), (T, N)
    forward, background, prior_sqrt, error_sqrt, observations, obs_index, obs_std = args
    # x₀ = μ_x + Σ_x^{1/2} ξ
    initial = background + prior_sqrt.mv(whitened_initial)  # (N,)

    def interval(state, whitened_error):
        # user (forward.step: one observation interval)
        # f_k(x_{k−1})
        forecast = forward.step(state, forward.dt)    # (N,) → (N,)
        # η_k = Σ_η^{1/2} w_k
        model_error = error_sqrt.mv(whitened_error)   # (N,)
        # x_k = f_k(x_{k−1}) + η_k
        state = forecast + model_error                # (N,)
        return state, state

    # x₁:T, one model run with a kick after every interval
    _, states = jax.lax.scan(interval, initial, whitened_errors)  # → (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)  # → ()
    # 𝒥_b = ½ ‖ξ‖²
    background_cost = 0.5 * einx.sum("state ->", whitened_initial**2)  # → ()
    # 𝒥_η = ½ ∑_k ‖η_k‖²_{Σ_η⁻¹} = ½ ∑_k ‖w_k‖²
    error_cost = 0.5 * einx.sum("time state ->", whitened_errors**2)  # → ()
    return background_cost + error_cost + obs_cost

Choosing \(\gmat{\Sigma}_\eta\)

• Scale. \(\gmat{\Sigma}_\eta\) is the variance of what the model gets wrong in one observation interval, so it scales with the interval length. In the example the true error is a forcing deficit of 1 per unit time, which is 0.1 per interval.
• Structure. A diagonal \(\gmat{\Sigma}_\eta\) allows any pattern of jumps, including grid-scale noise. Spatial correlations (a smooth kernel, or a basis) restrict the jumps to plausible model-error patterns.
• Random versus systematic. \(\mathcal{N}(\vect{0}, \gmat{\Sigma}_\eta)\) independent per interval describes random error. A bias, an error repeated every interval, is better written as one parameter shared over time; that is the reduced formulation below.
• Identifiability. A misfit early in the window can be explained by \(\vect{x}_0\) or by \(\gvect{\eta}_1\). \(\gmat{\Sigma}_x\) and \(\gmat{\Sigma}_\eta\) set the exchange rate; if both are poorly known, the split between them is too.

The gradient

Forcing formulation. The adjoint recursion is the strong-constraint one, run on the kicked trajectory. Because \(\gvect{\eta}_k\) enters exactly where \(\vect{x}_k\) does, its gradient comes free from the same sweep:

State formulation. Each term couples only neighbouring states, so the gradient with respect to \(\vect{x}_k\) needs only intervals \(k\) and \(k+1\): \(\nabla_{\vect{x}_k}\mathcal{J} = \vect{a}_k - \mat{F}_{k+1}^\top \vect{a}_{k+1} + \nabla_{\vect{x}_k}\mathcal{J}_{\text{o},k}\), with \(\vect{a}_k = \gmat{\Sigma}_{\eta,k}^{-1}(\vect{x}_k - f_k(\vect{x}_{k-1}))\). All \(T\) model runs and adjoints are independent. Everything on the primer's gradient page still applies inside each interval: checkpointing, diffrax adjoints, the adjoint test.

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

# --- forcing formulation: one backward sweep gives every gradient ---
# x†_k = ∇_{x_k} 𝒥_{o,k} + F_{k+1}ᵀ x†_{k+1}, as in strong-constraint 4D-Var


def backward(carry, inputs):
    state_prev, obs_forcing = inputs
    # x†_k = ∇_{x_k} 𝒥_{o,k} + F_{k+1}ᵀ x†_{k+1}
    lam = obs_forcing + carry                         # (N,)
    _, interval_vjp = jax.vjp(lambda s: forward.step(s, forward.dt), state_prev)
    # F_kᵀ x†_k, handed to the previous time
    (carry,) = interval_vjp(lam)                      # (N,) → (N,)
    return carry, lam


grad_initial, lams = jax.lax.scan(                    # → (N,), (T, N)
    backward, jnp.zeros(N), (previous, obs_forcing), reverse=True
)
# ∇_{w_k} 𝒥 = w_k + (Σ_η^{1/2})ᵀ x†_k: the kick at time k feels the same x†_k as x_k
grad_errors = whitened_errors + jax.vmap(error_sqrt.T.mv)(lams)  # (T, N)
# ∇_ξ 𝒥 = ξ + (Σ_x^{1/2})ᵀ F_1ᵀ x†_1
grad_initial = whitened_initial + prior_sqrt.T.mv(grad_initial)  # (N,)

# --- state formulation: every interval differentiated at once ---
# f_k(x_{k−1}) for k = 1…T, all in parallel
forecasts, interval_vjps = jax.vmap(                  # (T, N) → (T, N)
    lambda s: jax.vjp(lambda u: forward.step(u, forward.dt), s)
)(previous)
# a_k = Σ_η⁻¹ (x_k − f_k(x_{k−1}))
scaled = jax.vmap(error_cov_inv.mv)(states - forecasts)  # (T, N)
# ∇_{x_k} 𝒥 ∋ a_k − F_{k+1}ᵀ a_{k+1} + ∇_{x_k} 𝒥_{o,k}: local in time, no sweep
(pulled,) = jax.vmap(lambda f, a: f(a))(interval_vjps, scaled)  # (T, N)

Even on one CPU the state formulation's gradient is about 2.6 times faster, because the 20 intervals run as one batched call instead of a chain. On several devices the gap grows with \(T\).

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

# not pinned (optimistix)
# (ξ⋆, w⋆) = argmin 𝒥(ξ, w): L-BFGS over N(T+1) controls
solution = optx.minimise(                             # → ((N,), (T, N))
    weak_cost,
    optx.BFGS(rtol=1e-6, atol=1e-6),
    (jnp.zeros(N), jnp.zeros((T, N))),
    args=args,
    max_steps=3000,
)
whitened_initial, whitened_errors = solution.value
# η_k⋆ = Σ_η^{1/2} w_k⋆: the model-error estimate, one field per interval
model_errors = jax.vmap(error_sqrt.mv)(whitened_errors)  # (T, N)

# --- the same through vardax ---
# today (vardax)
# control (x₀, η₁:T) in physical units; Σ_η⁻¹ applied by a linear solve per step
weak = vardax.WeakFourDVar(
    forward,
    vardax.MaskedIdentity(),
    background,                                       # (N,)
    prior_cov,                                        # (N, N) PSD operator, Σ_x
    obs_cov,                                          # (N, N) PSD operator, Σ_ε
    model_err_cov,                                    # (N, N) PSD operator, Σ_η
)
# today (vardax)
analysis_initial, model_errors = weak(batch)          # → (1, N), (1, T, N)

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


def state_cost(control, args):
    """State formulation: control (ξ, x₁:T); the model only scores consecutive pairs."""
    whitened_initial, states = control                # (N,), (T, N)
    forward, background, prior_sqrt, error_std, observations, obs_index, obs_std = args
    # x₀ = μ_x + Σ_x^{1/2} ξ
    initial = background + prior_sqrt.mv(whitened_initial)  # (N,)
    # (x₀, …, x_{T−1}): the start of every interval
    previous = einx.rearrange(                        # (N,), (T, N) → (T, N)
        "state, time state -> (1 + time) state", initial, states
    )[:-1]
    # f_k(x_{k−1}) for all k at once: no sequential dependence
    forecasts = jax.vmap(lambda s: forward.step(s, forward.dt))(previous)  # (T, N)
    # Σ_η^{−1/2} (x_k − f_k(x_{k−1})), diagonal Σ_η = σ_η² I
    mismatch = (states - forecasts) / error_std       # (T, N)
    # 𝒥_η = ½ ∑_k ‖x_k − f_k(x_{k−1})‖²_{Σ_η⁻¹}
    error_cost = 0.5 * einx.sum("time state ->", mismatch**2)  # → ()
    # r_k = Σ_ε^{−1/2} (y_k − H_k x_k)
    residual = (observations - states[:, obs_index]) / obs_std  # (T, M)
    obs_cost = 0.5 * einx.sum("time obs ->", residual**2)  # → ()
    background_cost = 0.5 * einx.sum("state ->", whitened_initial**2)  # → ()
    return background_cost + error_cost + obs_cost


# start from the background trajectory, which makes every mismatch zero
initial_states = rollout(forward, background, T)      # → (T, N)

# Shapes: N state values (state); T observation times (time);
#         P spatial basis functions (basis); C temporal basis functions (temporal)


def reduced_error(coeffs, spatial, temporal, scales):
    """η(s, t_k) = ∑_{c,p} ω_{cp} χ_c(t_k) φ_p(s): model error in C × P coefficients."""
    # ω = Σ_ω^{1/2} w_ω, whitened control
    weights = scales * coeffs                         # (C, P)
    # ∑_p ω_{cp} φ_p(s): one spatial field per temporal mode
    fields = einx.dot(                                # (C, P), (N, P) → (C, N)
        "temporal basis, state basis -> temporal state", weights, spatial
    )
    # η_k = ∑_c χ_c(t_k) fields_c
    return einx.dot(                                  # (T, C), (C, N) → (T, N)
        "time temporal, temporal state -> time state", temporal, fields
    )


def reduced_cost(control, args):
    whitened_initial, coeffs = control                # (N,), (C, P)
    (forward, background, prior_sqrt, spatial, temporal, scales,
     observations, obs_index, obs_std) = args
    initial = background + prior_sqrt.mv(whitened_initial)  # (N,)
    # η₁:T from C·P numbers instead of T·N
    model_errors = reduced_error(coeffs, spatial, temporal, scales)  # → (T, N)

    def interval(state, model_error):
        # x_k = f_k(x_{k−1}) + η_k
        state = forward.step(state, forward.dt) + model_error  # (N,)
        return state, state

    _, states = jax.lax.scan(interval, initial, model_errors)  # → (T, N)
    residual = (observations - states[:, obs_index]) / obs_std  # (T, M)
    obs_cost = 0.5 * einx.sum("time obs ->", residual**2)  # → ()
    # 𝒥_b = ½ ‖ξ‖²
    background_cost = 0.5 * einx.sum("state ->", whitened_initial**2)  # → ()
    # 𝒥_η = ½ ‖w_ω‖², the model-error prior, whitened
    error_cost = 0.5 * einx.sum("temporal basis ->", coeffs**2)  # (C, P) → ()
    prior_cost = background_cost + error_cost         # → ()
    return prior_cost + obs_cost


# the bias case: C = 1, χ ≡ 1 (constant in time), P = 1, φ = uniform
# then η is one number per interval, the forcing error times Δt

The three side by side

Application: the Gulf Stream window

The QG model of the other pages is a good candidate for weak constraint: it omits unresolved eddy fluxes, uses an imperfect wind forcing, and has no topography.

Notes on vardax

Checked against vardax main (3e0f34a), alongside issues #94–#99 filed from the 4D-Var primer.

• WeakFourDVar optimises \(\gvect{\eta}_k\) in physical units and applies \(\gmat{\Sigma}_\eta^{-1}\) by an inner linear solve at every step. For a small \(\gmat{\Sigma}_\eta\) that direction is stiff (curvature \(1/\sigma_\eta^2\)), so whitening \(\gvect{\eta}_k = \gmat{\Sigma}_\eta^{1/2}\vect{w}_k\), the weak-constraint counterpart of #97, would help the optimiser most.
• The docstring says \(\gmat{\Sigma}_\eta\) has a default ("identity scaled by a small variance if not supplied"), but model_err_cov_op is a required argument.
• No state formulation and no reduced model error. docs/07 points users to a custom eqx.Module; the reduced form above is about thirty lines on top of the forcing cost, and somax's forcing-basis note already specifies the basis side.
• No forward_adjoint slot, although the docs/12 table lists one (#95).

Code markers: today (vardax) names ship on vardax main (3e0f34a); not pinned marks optimistix. Lorenz-96 numbers come from a JAX run in double precision: truth F = 8, model F = 7, D_x = 40, 20 observation times of 10 RK4 steps, every second site observed with σ_ε = 0.5, Σ_x as in the primer.