Learned data assimilation: what can be learned in the analysis

The taxonomy page took the model apart. This page takes the assimilation apart: the objective, its gradient, the solver, the output, and whether the answer is computed per window or learned across windows. Eleven ways to learn one of those pieces, on the same Gulf Stream QG example and the same 10-day window of SSH maps.

Notation

As on the other two pages. On this page \(\gvect{\nu}\) is the variational parameters; the QG viscosity does not appear. Losses \(\mathcal{L}\) are always minimised, so \(\mathcal{L}_\text{ELBO} = -\mathrm{ELBO}\).

The master equation

Weak-constraint 4D-Var over one window. With \(p(\vect{x} \mid \vect{y}) \propto \exp(-\mathcal{J}(\vect{x}))\), every type on this page edits this cost, its gradient, its solver, or replaces all three:

The classical solver iterates \(\vect{x}^{(i+1)} = \vect{x}^{(i)} - \gamma \mat{P} \nabla_{\vect{x}} \mathcal{J}(\vect{x}^{(i)})\) from \(\vect{x}^{(0)}\), with step size \(\gamma\), \(N \approx 50\)–\(100\) and \(\mat{P}\) built from \(\gmat{\Sigma}_x^{1/2}\). In code, the cost, the solver and the variational objective are three short functions. qg_day and QGParams are the ones defined on the taxonomy page:

import einx
import equinox as eqx
import finitevolx as fvx
import gaussx as gx
import jax
import jax.numpy as jnp
import lineax as lx
from jaxtyping import Array, Bool, Float, PRNGKeyArray

# Shapes: T = 10 maps (time); Z = T + 1 = 11 states x₀ … x₁₀ (day);
#         H, W = 130 × 130 grid with ghost ring (height, width);
#         N = H·W (cell); F = Z·N (flat); K = 16 (rank); S = 8 samples (member)

Op = lx.AbstractLinearOperator


class Window(eqx.Module):
    """One window: maps y, masks m, errors σ_ε², background μ_x, Σ_x, diagonal Σ_η."""

    maps: Float[Array, "time height width"]
    masks: Bool[Array, "time height width"]
    obs_var: Float[Array, "time height width"]
    model_err_var: Float[Array, "time height width"]
    background: Float[Array, "height width"]
    covariates: Float[Array, "variable height width"]
    forcing: Float[Array, "time height width"]
    prior_cov: Op
    grid: fvx.CartesianGrid2D


def weak_4dvar_cost(
    states: Float[Array, "day height width"],
    window: Window,
    params: QGParams,
) -> Float[Array, ""]:
    # --- background ---
    # δ₀ = x₀ − μ_x
    initial_departure = states[0] - window.background  # (H, W)
    departure_flat = einx.rearrange(                  # (H, W) → (N,)
        "height width -> (height width)", initial_departure
    )
    # today (gaussx)
    # δ₀ᵀ Σ_x⁻¹ δ₀
    background_quad = gx.quadratic_form(window.prior_cov, departure_flat)  # → ()
    # 𝒥_x = ½ δ₀ᵀ Σ_x⁻¹ δ₀
    background_cost = 0.5 * background_quad           # → ()

    # --- dynamics, all days at once ---
    previous = states[:-1]                            # (Z, Z) → (T, H, W)
    current = states[1:]                              # (Z, Z) → (T, H, W)
    batched_day = jax.vmap(qg_day, in_axes=(0, None, 0, None, None))
    # user (qg_day, taxonomy page)
    # f(x_{k−1}, u_k; θ) for every k, in parallel
    forecasts = batched_day(                          # (T, H, W) → (T, H, W)
        previous, window.covariates, window.forcing, params, window.grid
    )
    # η_k = x_k − f(x_{k−1}, u_k; θ)
    model_error = current - forecasts                 # (T, H, W)
    # η_k²
    squared_error = model_error**2                    # (T, H, W)
    # η_k² / σ_η²  (diagonal Σ_η)
    weighted_error = squared_error / window.model_err_var  # (T, H, W)
    # ∑_k η_kᵀ Σ_η⁻¹ η_k
    model_sum = einx.sum("time height width ->", weighted_error)  # → ()
    # 𝒥_η = ½ ∑_k ‖η_k‖²_{Σ_η⁻¹}
    model_cost = 0.5 * model_sum                      # → ()

    # --- observations ---
    # y_k − x_k
    misfit = window.maps - current                    # (T, H, W)
    # r_k = m_k ⊙ (y_k − x_k)
    residual = jnp.where(window.masks, misfit, 0.0)   # (T, H, W)
    # r_k²
    squared_residual = residual**2                    # (T, H, W)
    # r_k² / σ_{ε,k}²
    weighted_residual = squared_residual / window.obs_var  # (T, H, W)
    # ∑_k r_kᵀ Σ_{ε,k}⁻¹ r_k
    obs_sum = einx.sum("time height width ->", weighted_residual)  # → ()
    # 𝒥_ε = ½ ∑_k ‖m_k ⊙ (y_k − x_k)‖²_{Σ_{ε,k}⁻¹}
    obs_cost = 0.5 * obs_sum                          # → ()

    # 𝒥 = 𝒥_x + 𝒥_η + 𝒥_ε
    total_cost = background_cost + model_cost + obs_cost  # → ()
    return total_cost


def analyse(
    window: Window,
    params: QGParams,
    initial_states: Float[Array, "day height width"],
    *,
    step_size: float = 0.1,
    n_iterations: int = 100,
) -> Float[Array, "day height width"]:
    cost_gradient = jax.grad(weak_4dvar_cost)

    def iterate(states, _):
        # ∇_x 𝒥(x⁽ⁱ⁾), by reverse mode
        gradient = cost_gradient(states, window, params)  # → (Z, H, W)
        # user
        # d⁽ⁱ⁾ = P ∇_x 𝒥(x⁽ⁱ⁾)
        direction = precondition(gradient, window)    # (Z, H, W)
        # γ d⁽ⁱ⁾
        update = step_size * direction                # (Z, H, W)
        # x⁽ⁱ⁺¹⁾ = x⁽ⁱ⁾ − γ P ∇_x 𝒥(x⁽ⁱ⁾)
        states = states - update                      # (Z, H, W)
        return states, None

    analysis, _ = jax.lax.scan(                       # (Z, H, W) → (Z, H, W)
        iterate, initial_states, None, length=n_iterations
    )
    return analysis


class LowRankGaussian(eqx.Module):
    """q(x) = 𝒩(μ_q, L_q L_qᵀ + D_q) over a whole window, flattened."""

    mean: Float[Array, "flat"]
    factor: Float[Array, "flat rank"]
    log_diag: Float[Array, "flat"]


def neg_elbo(
    posterior: LowRankGaussian,
    window: Window,
    params: QGParams,
    *,
    n_samples: int = 8,
    key: PRNGKeyArray,
) -> Float[Array, ""]:
    height, width = window.background.shape

    # --- variational density ---
    # D_q = diag(exp(log d))
    diag = jnp.exp(posterior.log_diag)                # (F,) → (F,)
    # today (gaussx)
    # Σ_q = L_q L_qᵀ + D_q
    covariance = gx.low_rank_plus_diag(diag, posterior.factor)  # (F, F) PSD op

    # --- samples ---
    # today (gaussx)
    # x⁽ᵉ⁾ = μ_q + Σ_q^{1/2} ξ,  ξ ~ 𝒩(0, I)
    samples_flat = gx.sample_mvn(                     # → (S, F)
        posterior.mean, covariance, key=key, num_samples=n_samples
    )
    samples = einx.rearrange(                         # (S, F) → (S, Z, H, W)
        "member (day height width) -> member day height width",
        samples_flat,
        height=height,
        width=width,
    )

    # --- expected cost and entropy ---
    batched_cost = jax.vmap(weak_4dvar_cost, in_axes=(0, None, None))
    # 𝒥(x⁽ᵉ⁾) for each sample
    sample_costs = batched_cost(samples, window, params)  # → (S,)
    # 𝔼_q[𝒥] ≈ (1/S) ∑_e 𝒥(x⁽ᵉ⁾)
    expected_cost = einx.mean("member ->", sample_costs)  # (S,) → ()
    # today (gaussx)
    # ℍ[q] = ½ (N (1 + log 2π) + log det Σ_q)
    entropy = gx.gaussian_entropy(covariance)         # → ()

    # ℒ_ELBO = −ELBO = 𝔼_q[𝒥(x)] − ℍ[q],  since p(x | y) ∝ exp(−𝒥(x))
    loss = expected_cost - entropy                    # → ()
    return loss

Markers: today ships on main (checked 24 Oct 2026: finitevolX 831c69f, gaussx eeda235, somax 9cb19b3); today* ships, arguments sketched; user is defined on these pages; not pinned is a library whose API was not checked. Shape column letters are defined in the header.

No Gaussian variational family ships in gaussx or vardax, so LowRankGaussian is page-defined; gaussx supplies its covariance, samples and entropy.

Training signals

What a type can learn from decides where it can be used. Five signals, with what each needs and which statistic it returns:

The eleven types

Each card gives the equation, the operator signature with what is learned and what it trains on, and attribute rows. The training signal is a row, as on the taxonomy page.

Model type × assimilation type

Any model from the taxonomy page can sit inside any type that uses a model. A6 and A7 use a model only to make training data (OSSEs) or, for the ELBO, inside \(\log p(\vect{x})\).

● natural   ◐ works, with a caveat   ○ does not fit

The caveats: T6–T8 need a trustworthy Jacobian wherever a gradient flows through them (A2–A4, A8, A9); T9 has no one-day step, so it cannot enter a weak-constraint dynamics term or cycle in a filter; T10 moves the whole problem into the embedding, with observations entering through the decoder. A11 with the physical model only calibrates θ, which is A1's job done the hard way.

What the structure implies

Where each type lives