Skip to content

Incremental 4DVar with Control-Variable Transform

Operational 4DVar at ECMWF, NCEP, JMA, and UKMO all follow a common template: Gauss-Newton outer iterations on the full nonlinear cost, CG (or Lanczos) inner iterations on a linearised quadratic subproblem, with a control-variable transform that preconditions the inner solve. IncrementalFourDVar packages this pattern.

It is functionally equivalent to StrongFourDVar (chapter 6) — same problem, same answer — but the specialised inner solver is dramatically faster when the prior is structured (Matérn, Kronecker) and the window is long.

The outer / inner decomposition

Strong-constraint 4DVar minimises

\[ J(x_0) = \tfrac{1}{2} \|x_0 - x_b\|^2_{B^{-1}} + \tfrac{1}{2} \sum_t \|y_t - H_t(M_t(x_0))\|^2_{R_t^{-1}}. \]

Write \(x^{(k)}\) for the outer iterate, starting from \(x^{(0)} = x_b\). The background \(x_b\) stays fixed across outer loops; only \(x^{(k)}\) moves. Linearise around \(x^{(k)}\):

\[ M_t(x^{(k)} + \delta x) \approx M_t(x^{(k)}) + M'_t \delta x, \]
\[ H_t(z + \delta z) \approx H_t(z) + H'_t \delta z, \]

where \(M'_t = \partial M_t / \partial x|_{x^{(k)}}\) from jax.linearize and \(H'_t\) is the tangent of \(H_t\) at \(M_t(x^{(k)})\). With the innovation \(d_t^{(k)} = y_t - H_t(M_t(x^{(k)}))\), the linearised (quadratic) cost of an increment \(\delta x\) is

\[ \delta J(\delta x) = \tfrac{1}{2} \|x^{(k)} + \delta x - x_b\|^2_{B^{-1}} + \tfrac{1}{2} \sum_t \|d_t^{(k)} - H'_t M'_t \delta x\|^2_{R_t^{-1}}. \]

The background term measures the total departure \(x^{(k)} + \delta x - x_b\), not \(\delta x\) alone. The Gauss-Newton Hessian is

\[ \mathbf{H}_\text{GN} = B^{-1} + \sum_t (H'_t M'_t)^\top R_t^{-1} (H'_t M'_t), \]

and setting \(\nabla \delta J = 0\) gives the inner system that the CG / Lanczos solver solves for \(\delta x^*\):

\[ \mathbf{H}_\text{GN}\, \delta x^* = -B^{-1}\big(x^{(k)} - x_b\big) + \sum_t (H'_t M'_t)^\top R_t^{-1} d_t^{(k)} = -\nabla J\big(x^{(k)}\big). \]

The right-hand side is the negative gradient of the full nonlinear cost at the current iterate. On the first outer loop \(x^{(0)} = x_b\), so the background term vanishes and only the innovation term is left. On later loops it does not vanish: dropping it would pull every increment toward \(x^{(k)}\) instead of \(x_b\), and the iteration would no longer converge to the 4DVar minimum.

Outer update: \(x^{(k+1)} = x^{(k)} + \delta x^*\). Relinearise.

This pattern converges in a handful of outer iterations (typically 3–5) regardless of nonlinearity strength, because each outer step sees a fresh tangent linearisation.

Control-variable transform (CVT)

The Hessian \(\mathbf{H}_\text{GN}\) inherits the condition number of \(B^{-1}\): a smooth (long-correlation) \(B\) has tiny eigenvalues, which become huge ones of \(B^{-1}\). In state space every Hessian product also needs \(B^{-1}v\), which IncrementalFourDVar would have to compute by a nested CG solve.

The CVT changes variables. Let \(S\) be a square root of \(B\), \(S S^\top = B\), and write the departure from the background as

\[ x_0 - x_b = S \chi . \]

The background term becomes \(\tfrac{1}{2}\|x_0 - x_b\|^2_{B^{-1}} = \tfrac{1}{2}\|\chi\|^2\), and the inner system at the outer iterate \(\chi^{(k)}\) becomes

\[ \big(I + S^\top G'^\top R^{-1} G' S\big)\, \delta\chi^* = -\nabla_\chi J\big(\chi^{(k)}\big) = -\chi^{(k)} + S^\top G'^\top R^{-1} d^{(k)}, \]

where \(G'\) stacks \(H'_t M'_t\) over the window. Its eigenvalues are all \(\ge 1\), close to 1 in directions the observations barely see, and no \(B^{-1}\) appears anywhere. The outer update is \(\chi^{(k+1)} = \chi^{(k)} + \delta\chi^*\), and the analysis is \(x_0^* = x_b + S\chi^*\).

IncrementalConfig(cvt=True) (the default) uses it. \(S\) is prior_sqrt_op= when given (a Cholesky factor, a structured operator's own square root, or gaussx.sqrt(B, lanczos_order=k) for a large unstructured \(B\)), else gaussx.sqrt(prior_cov_op): closed form for diagonal, Kronecker and block-diagonal \(B\), a dense eigendecomposition otherwise. cvt=False keeps the state-space solve, with \(B^{-1}\) and \(R^{-1}\) applied by CG at tolerance covariance_tol.

How much it helps. A linear problem with 20 states, 4 times, every third site observed, a Gaussian-kernel \(B\) of length scale 3 sites plus a diagonal jitter that sets its condition number; inner CG iterations to reach \(10^{-4}\) relative error on the exact analysis:

\(\kappa(B)\) CVT state space
\(7 \times 10^1\) 10 8
\(7 \times 10^2\) 9 14
\(7 \times 10^3\) 9 20

The CVT count does not grow with \(\kappa(B)\), and each state-space iteration also runs a nested CG for \(B^{-1}v\). For a well-conditioned \(B\) the two are comparable.

Implementation in vardax

import gaussx as gx
import lineax as lx
from vardax import AveragingKernel, IncrementalConfig, IncrementalFourDVar

# separable B on an (ny, nx) grid: B = B_y ⊗ B_x, whose square root
# gaussx.sqrt computes factor by factor
B_op = gx.Kronecker(B_y_op, B_x_op)

model = IncrementalFourDVar(
    forward=somax_model,
    obs_op=AveragingKernel(kernel_matrix=A, prior_profile=xa, pressure_weights=h),
    prior_mean=x_b,
    prior_cov_op=B_op,
    obs_cov_op=lx.DiagonalLinearOperator(obs_uncertainty),
    config=IncrementalConfig(n_outer=3, n_inner=20, cvt=True),
)
x_star = model(batch)

# Gauss–Newton posterior at the analysis, with the model's B and R
posterior = model.posterior(batch)

IncrementalConfig:

  • n_outer — number of Gauss-Newton outer iterations (typical: 3)
  • n_inner — max CG iterations per outer (typical: 20–50); the inner solve is truncated there
  • cg_atol, cg_rtol — inner CG convergence tolerances
  • cvt — solve over \(\chi\), \(x_0 = x_b + B^{1/2}\chi\) (default True)
  • covariance_tol — tolerance of the CG solves that apply \(R^{-1}\) (and \(B^{-1}\) when cvt=False)

Algorithm

Inputs:
   x_b           — background mean
   B_op          — prior covariance; B^{1/2} from gaussx.sqrt or prior_sqrt_op
   R_op          — observation-error covariance
   forward       — pipekit_cycle.ForwardModel
   obs_op        — pipekit_cycle.ObservationOperator
   batch         — observations y_t for t = 0, …, T
   config        — IncrementalConfig

x ← x_b                       # x⁽⁰⁾ = x_b; x_b itself never changes
for k in range(config.n_outer):

    # 1. Tangent of the observed trajectory G(x) = [H_t(M_t(x))]_t at x⁽ᵏ⁾
    #    G'δx = [H'_t M'_t δx]_t   (jax.linearize; Gᵀ via jax.vjp)
    G_x, G_lin = jax.linearize(observed_trajectory, x)

    # 2. Innovations at the current iterate
    #    d_tᵏ = y_t − H_t(M_t(x⁽ᵏ⁾))
    d = batch.y - G_x

    # 3. Right-hand side: the negative gradient of the full cost at x⁽ᵏ⁾
    #    −∇J(x⁽ᵏ⁾) = −B⁻¹(x⁽ᵏ⁾ − x_b) + G'ᵀ R⁻¹ d
    rhs = -B_inv @ (x - x_b) + G_lin.T @ (R_inv @ d)

    # 4. Inner solve via lineax.CG
    if config.cvt:
        B_half = gaussx.sqrt(B_op)    # or prior_sqrt_op
        # (I + B^{1/2}ᵀ G'ᵀ R⁻¹ G' B^{1/2}) δχ = B^{1/2}ᵀ (−∇J(x⁽ᵏ⁾))
        preconditioned_hessian = identity_op + B_half.T @ GtRG @ B_half
        chi = lineax.linear_solve(
            preconditioned_hessian, B_half.T @ rhs,
            solver=lineax.CG(atol=config.cg_atol, rtol=config.cg_rtol,
                              max_steps=config.n_inner),
        )
        # δx = B^{1/2} δχ
        dx = B_half @ chi
    else:
        # (B⁻¹ + G'ᵀ R⁻¹ G') δx = −∇J(x⁽ᵏ⁾)
        dx = lineax.linear_solve(B_inv + GtRG, rhs, solver=lineax.CG(...))

    # 5. Outer update: x⁽ᵏ⁺¹⁾ = x⁽ᵏ⁾ + δx*
    x = x + dx

return x

Why it's the operational fast path

Three multiplicative speedups over StrongFourDVar with NonlinearCG:

  1. CVT preconditioning — CG converges in \(O(10)\) iterations instead of \(O(\sqrt{\kappa(B)})\) where \(\kappa(B)\) is the prior condition number.
  2. Linearisation reuse — the tangent linear \(M'\) is computed once per outer iteration and reused across all \(n_\text{inner}\) inner CG steps. NonlinearCG on the full nonlinear cost recomputes the gradient (which requires a full forward + adjoint pass) every step.
  3. Cheap posterior — the posterior precision is the same Gauss–Newton operator the inner loop solves with, so model.posterior(batch) costs one more linearisation of the rollout at the analysis.

For a typical operational problem (\(T = 12\) hours, \(n = 10^7\) grid points, \(m = 10^5\) observations), IncrementalFourDVar converges in minutes; StrongFourDVar with NonlinearCG takes hours.

Linear-Gaussian agreement

For \(T = 0\) and linear \(H\), IncrementalFourDVar reduces to OptimalInterpolation. The outer iteration converges in one step (no relinearisation needed), and the inner CG produces the closed-form Kalman update:

import jax
import jax.numpy as jnp
import lineax as lx
import vardax


class Static:
    """M(x) = x: a single-time window reduces 4DVar to OI."""

    dt = 1.0

    def step(self, x, dt):
        return x


def test_incremental_recovers_oi():
    n = 4
    x_b = jnp.zeros(n)
    B_op = lx.DiagonalLinearOperator(jnp.full(n, 2.0))
    R_op = lx.DiagonalLinearOperator(jnp.full(n, 0.5))
    batch = vardax.Batch1D(
        input=jnp.arange(1.0, n + 1)[None, None], mask=jnp.ones((1, 1, n))
    )
    oi = vardax.OptimalInterpolation(
        obs_op=vardax.MaskedIdentity(), prior_mean=x_b[None],
        prior_cov_op=lx.DiagonalLinearOperator(jnp.full((1, n), 2.0)),
        obs_cov_op=lx.DiagonalLinearOperator(jnp.full((1, n), 0.5)),
    )
    inc = vardax.IncrementalFourDVar(
        forward=Static(), obs_op=vardax.MaskedIdentity(),
        prior_mean=x_b, prior_cov_op=B_op, obs_cov_op=R_op,
        config=vardax.IncrementalConfig(n_outer=2, n_inner=50, cvt=True),
    )
    assert jnp.allclose(oi(batch)[0, 0], inc(batch)[0], atol=1e-3)

Part of the Decision D14 invariant.

Strong- vs weak-constraint

IncrementalFourDVar is strong-constraint by default. For weak-constraint incremental 4DVar (the augmented control vector with \(\eta_t\)), use WeakFourDVar directly — the incremental machinery for the augmented problem is heavier and not always worth the complexity. The strong / weak choice is structural; the inner-solver choice (general vs incremental) is performance.

Posterior covariance

The posterior precision at the analysis is the Gauss–Newton operator \(B^{-1} + \sum_t (H'_t M'_t)^\top R_t^{-1} (H'_t M'_t)\), the inner loop's Hessian. model.posterior(batch) relinearises the rollout once at the analysis and returns its lazy inverse (chapter 13):

posterior = model.posterior(batch)              # batch: one window
# Posterior(mean=x_star, cov=lazy (B⁻¹ + G'ᵀR⁻¹G')⁻¹, samples=None, provenance={...})

Pass adapter=vardax.GaussNewtonHessian(n_krylov=...) to set the Krylov budget of each covariance mat-vec.

When to switch from IncrementalFourDVar to StrongFourDVar

Almost never, for operational work. IncrementalFourDVar is faster and produces the same answer. Use StrongFourDVar only when:

  • \(B\) is large, unstructured and has no cheap square root, so gaussx.sqrt would fall back to a dense eigendecomposition. Pass a Lanczos square root (prior_sqrt_op=gaussx.sqrt(B, lanczos_order=k)) or use cvt=False before giving up on the incremental solver.
  • You want to swap in a custom optimistix minimiser (e.g. optx.LevenbergMarquardt for very nonlinear cases). Incremental uses a fixed GN+CG; the rest of the optimistix zoo isn't compatible with the incremental machinery.

For most users, IncrementalFourDVar is the right default.

See also

  • Chapter 6 — StrongFourDVar (equivalent problem, slower solver)
  • Chapter 12 — adjoint composition (CVT relies on gaussx factorisation)
  • Chapter 13 — posterior covariance (reuses GN Hessian)
  • Design doc: design/decisions.md#d11

References

  • Courtier, P., Thépaut, J.-N., & Hollingsworth, A. (1994). A strategy for operational implementation of 4D-Var, using an incremental approach. QJRMS 120(519).
  • Lorenc, A. C. (1997). Development of an operational variational assimilation scheme. JMSJ 75(1B).
  • Bannister, R. N. (2017). A review of operational methods of variational and ensemble-variational data assimilation. QJRMS 143(703).