Variational data assimilation, written as probabilistic programs

1 · 3D-Var

Intuition. At each observation time, correct one state: pull the forecast toward the observations, by an amount set by \(\gmat{\Sigma}_x\) against \(\sigma_\varepsilon^2\), and spread the correction to unobserved sites through \(\gmat{\Sigma}_x\)'s correlations. The model only carries the analysis forward; it plays no part in the fit.

# Shapes: N = 40 sites (state); M = 20 observed sites (obs)


def three_dvar_readable(
    forecast: Float[Array, "state"],
    prior_cov: lx.AbstractLinearOperator,
    obs_index: Int[Array, "obs"],
    obs_std: float,
    obs: Float[Array, "obs"] | None = None,
    obs_mask: Bool[Array, "obs"] | bool = True,
):
    # today (gaussx)
    # x ~ 𝒩(x_k|k−1, Σ_x)
    state = numpyro.sample(                                 # event (N,)
        "state", gx.MultivariateNormal(forecast, prior_cov)
    )
    # y ~ 𝒩(H x, σ_ε² I), masked where m = 0
    likelihood = dist.Normal(state[obs_index], obs_std)     # batch (M,)
    numpyro.sample("obs", likelihood.mask(obs_mask).to_event(1), obs=obs)


def three_dvar(
    forecast: Float[Array, "state"],
    prior_sqrt: lx.AbstractLinearOperator,
    obs_index: Int[Array, "obs"],
    obs_std: float,
    obs: Float[Array, "obs"] | None = None,
    obs_mask: Bool[Array, "obs"] | bool = True,
):
    # ξ ~ 𝒩(0, I)
    whitened = numpyro.sample(                              # event (N,)
        "whitened", dist.Normal(0.0, 1.0).expand(forecast.shape).to_event(1)
    )
    # x = x_k|k−1 + Σ_x^{1/2} ξ, the control-variable transform
    state = numpyro.deterministic("state", forecast + prior_sqrt.mv(whitened))  # (N,)
    # y ~ 𝒩(H x, σ_ε² I)
    likelihood = dist.Normal(state[obs_index], obs_std)     # batch (M,)
    numpyro.sample("obs", likelihood.mask(obs_mask).to_event(1), obs=obs)

Both functions define the same posterior; on the first analysis they differ by at most 0.0014 at any site. The readable one, a gaussx MultivariateNormal site in physical units, took 1,952 L-BFGS iterations per analysis; the whitened one took 27. Cycled over the window, 3D-Var reaches a window error of 1.80 (background: 3.50).

2 · Reduced 3D-Var

Intuition. Correct the forecast only in a few smooth patterns. The same reduction as reduced 4D-Var, applied to one time: the increment is \(\mat{L} \vect{z}_L\), with \(N_\text{feat}\) coefficients instead of \(D_x\) values. It is cheap and well conditioned, but whatever the features cannot express is never corrected, and in a cycled system those uncorrected scales keep growing from one analysis to the next.

# Shapes: N = 40 sites (state); M = 20 observed sites (obs);
#         R = 2K + 1 features (basis)


def reduced_three_dvar(
    forecast: Float[Array, "state"],
    features: Float[Array, "state basis"],
    obs_index: Int[Array, "obs"],
    obs_std: float,
    obs: Float[Array, "obs"] | None = None,
    obs_mask: Bool[Array, "obs"] | bool = True,
):
    # z_L ~ 𝒩(0, I_R)
    whitened = numpyro.sample(                              # event (R,)
        "whitened", dist.Normal(0.0, 1.0).expand([features.shape[1]]).to_event(1)
    )
    # δx = L z_L, so δx ~ 𝒩(0, L Lᵀ) ≈ 𝒩(0, Σ_x)
    increment = einx.dot("state basis, basis -> state", features, whitened)  # (N,)
    # x = x_k|k−1 + L z_L
    state = numpyro.deterministic("state", forecast + increment)  # (N,)
    # y ~ 𝒩(H x, σ_ε² I)
    likelihood = dist.Normal(state[obs_index], obs_std)     # batch (M,)
    numpyro.sample("obs", likelihood.mask(obs_mask).to_event(1), obs=obs)

Cycling makes the truncation costly: the small scales the features leave out are never corrected, the chaotic model amplifies them, and the window error is 2.82 with 9 features and 2.39 with 17, against 1.80 for full 3D-Var.

3 · Strong-constraint 4D-Var

Intuition. Trust the model completely. Then the trajectory is a function of its starting point, and the only unknown is \(\vect{x}_0\): every observation in the window pulls on it through the dynamics. In the program, the rollout is a numpyro.deterministic site, so the posterior lives on \(\gvect{\xi}\) alone.

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


def rollout(model_step, initial, errors=None, num_times=None):
    """x₁:T from x₀, with an optional kick η_k after every interval."""

    def interval(state, error):
        # user (model_step: one observation interval of the dynamics)
        # x_k = f_k(x_{k−1}) (+ η_k)
        state = model_step(state) if error is None else model_step(state) + error
        return state, state

    _, states = jax.lax.scan(interval, initial, errors, length=num_times)  # → (T, N)
    return states


def strong_4dvar(
    model_step: Callable,
    background: Float[Array, "state"],
    prior_sqrt: lx.AbstractLinearOperator,
    num_times: int,
    obs_index: Int[Array, "obs"],
    obs_std: float,
    obs: Float[Array, "time obs"] | None = None,
    obs_mask: Bool[Array, "time obs"] | bool = True,
):
    # ξ ~ 𝒩(0, I)
    whitened = numpyro.sample(                              # event (N,)
        "whitened", dist.Normal(0.0, 1.0).expand(background.shape).to_event(1)
    )
    # x₀ = μ_x + Σ_x^{1/2} ξ
    initial = numpyro.deterministic(                        # (N,)
        "initial", background + prior_sqrt.mv(whitened)
    )
    # x_k = f_{k←0}(x₀): deterministic, so the posterior lives on ξ alone
    states = numpyro.deterministic(                         # (T, N)
        "states", rollout(model_step, initial, num_times=num_times)
    )
    # y_k ~ 𝒩(H x_k, σ_ε² I) for every k
    likelihood = dist.Normal(states[:, obs_index], obs_std)  # batch (T, M)
    numpyro.sample("obs", likelihood.mask(obs_mask).to_event(1), obs=obs)

Analysis: \(\mathcal{J}^\star = \) 509, \(2\mathcal{J}^\star/N_\text{obs} = \) 2.55 (expected about 1 for a correct model), window error 0.64. The Laplace spread of \(\vect{x}_0\) is 0.12, but its actual error is 0.85: the posterior is confident and wrong, because the model it trusts is biased. A ratio of error to spread far above 1 is the same alarm as \(2\mathcal{J}^\star/N_\text{obs}\gg 1\).

4 · Reduced strong-constraint 4D-Var

Intuition. Correct the initial state only in a few smooth patterns. The control shrinks from \(D_x\) values to \(N_\text{feat}\) coefficients; the price is that variance outside the features cannot be corrected. Here the features are the truncated Fourier series of the background kernel, so the reduced prior is a truncation of \(\gmat{\Sigma}_x\) itself.

# Shapes: N = 40 sites (state); T = 20 observation times (time);
#         M = 20 observed sites (obs); R = 2K + 1 features (basis)


def reduced_strong_4dvar(
    model_step: Callable,
    background: Float[Array, "state"],
    features: Float[Array, "state basis"],
    num_times: int,
    obs_index: Int[Array, "obs"],
    obs_std: float,
    obs: Float[Array, "time obs"] | None = None,
    obs_mask: Bool[Array, "time obs"] | bool = True,
):
    # z_L ~ 𝒩(0, I_R)
    whitened = numpyro.sample(                              # event (R,)
        "whitened", dist.Normal(0.0, 1.0).expand([features.shape[1]]).to_event(1)
    )
    # δx₀ = L z_L
    increment = einx.dot("state basis, basis -> state", features, whitened)  # (N,)
    # x₀ = μ_x + L z_L
    initial = numpyro.deterministic("initial", background + increment)  # (N,)
    states = numpyro.deterministic(                         # (T, N)
        "states", rollout(model_step, initial, num_times=num_times)
    )
    likelihood = dist.Normal(states[:, obs_index], obs_std)  # batch (T, M)
    numpyro.sample("obs", likelihood.mask(obs_mask).to_event(1), obs=obs)


def lowrank_strong_4dvar(
    model_step: Callable,
    background: Float[Array, "state"],
    features: Float[Array, "state basis"],
    nugget: float,
    num_times: int,
    obs_index: Int[Array, "obs"],
    obs_std: float,
    obs: Float[Array, "time obs"] | None = None,
    obs_mask: Bool[Array, "time obs"] | bool = True,
):
    # today (gaussx)
    # Σ_x ≈ σ_nug² I + L Lᵀ: solves and log-determinants by Woodbury
    prior_cov = gx.low_rank_plus_diag(                      # (N, N) PSD op
        nugget * jnp.ones_like(background), features, psd=True
    )
    # today (gaussx)
    # x₀ ~ 𝒩(μ_x, σ_nug² I + L Lᵀ): every value is a control again
    initial = numpyro.sample(                               # event (N,)
        "initial", gx.MultivariateNormal(background, prior_cov)
    )
    states = numpyro.deterministic(                         # (T, N)
        "states", rollout(model_step, initial, num_times=num_times)
    )
    likelihood = dist.Normal(states[:, obs_index], obs_std)  # batch (T, M)
    numpyro.sample("obs", likelihood.mask(obs_mask).to_event(1), obs=obs)

The second function keeps every value as a control but gives \(\vect{x}_0\) a low_rank_plus_diag prior: the 9 features plus a white nugget carrying the remaining 0.15 of variance per site. gaussx solves and log-determines it with the Woodbury identity. It is a different prior, so its \(\mathcal{J}^\star\) (*) is not comparable with the rows above; its window error is.

5 · Incremental 4D-Var

Intuition. Strong-constraint 4D-Var is hard because the model is nonlinear. Incremental 4D-Var takes that difficulty out of the inner solve: run the nonlinear model once from the current analysis, linearise it around that run, and solve the resulting linear-Gaussian problem exactly. Then rerun the model from the new analysis and repeat. It is Gauss–Newton on the strong-constraint cost: it targets the same MAP, and its Laplace spread uses the Gauss–Newton Hessian. In numpyro terms, the inner problem is the strong-constraint program with the rollout replaced by its tangent-linear approximation, which jax.linearize provides.

The inner cost is quadratic in \(\gvect{\xi}\), so the linearised problem's posterior is exactly Gaussian and its Laplace approximation is exact; for the nonlinear problem it is the Gauss–Newton approximation at the analysis. The step length \(\alpha\) is halved until the nonlinear cost decreases: on this chaotic window, plain Gauss–Newton steps overshoot. Without the halving, the cost over the last four outer loops went 514.2 → 511.0 → 512.0 → 512.2 for the full increment and 1362 → 1413 → 1371 → 1442 with 9 features.

# Shapes: N = 40 sites (state); T = 20 observation times (time);
#         M = 20 observed sites (obs); R = 2K + 1 features (basis);
#         D = N or R controls (latent)


def incremental_inner(
    reference_initial: Float[Array, "state"],
    reference_states: Float[Array, "time state"],
    tangent: Callable,
    initial_of: Callable,
    num_controls: int,
    obs_index: Int[Array, "obs"],
    obs_std: float,
    obs: Float[Array, "time obs"] | None = None,
    obs_mask: Bool[Array, "time obs"] | bool = True,
):
    """The inner problem: the strong-constraint program with the model linearised."""
    # z ~ 𝒩(0, I): ξ for the full increment, z_L for the reduced one
    whitened = numpyro.sample(                              # event (D,)
        "whitened", dist.Normal(0.0, 1.0).expand([num_controls]).to_event(1)
    )
    # x₀ = μ_x + Σ_x^{1/2} ξ, or μ_x + L z_L
    initial = numpyro.deterministic("initial", initial_of(whitened))  # (N,)
    # δx₀ = x₀ − x₀⁽ⁱ⁾, the departure from the outer loop's reference
    increment = initial - reference_initial                 # (N,)
    # x_k ≈ x_k⁽ⁱ⁾ + F_{k←0} δx₀: tangent-linear, so x is linear in z
    states = numpyro.deterministic(                         # (T, N)
        "states", reference_states + tangent(increment)
    )
    # y_k ~ 𝒩(H x_k, σ_ε² I): linear-Gaussian, so the posterior is Gaussian
    likelihood = dist.Normal(states[:, obs_index], obs_std)  # batch (T, M)
    numpyro.sample("obs", likelihood.mask(obs_mask).to_event(1), obs=obs)


def nonlinear_cost(model_step, initial_of, whitened, num_times, obs_index, obs_std,
                   obs, obs_mask=True):
    """𝒥(z) = ½‖z‖² + ½ Σ_k ‖y_k − H f_{k←0}(x₀(z))‖²/σ_ε², the outer-loop cost."""
    states = rollout(model_step, initial_of(whitened), num_times=num_times)  # (T, N)
    # r_k = (y_k − H x_k) / σ_ε, zero where m = 0
    residual = (obs - states[:, obs_index]) / obs_std       # (T, M)
    residual = jnp.where(obs_mask, residual, 0.0)           # (T, M)
    return 0.5 * jnp.sum(whitened**2) + 0.5 * jnp.sum(residual**2)  # → ()


def incremental_4dvar(
    model_step: Callable,
    initial_of: Callable,
    num_controls: int,
    num_times: int,
    num_outer: int,
    obs_index: Int[Array, "obs"],
    obs_std: float,
    obs: Float[Array, "time obs"],
    obs_mask: Bool[Array, "time obs"] | bool = True,
    init_values: dict | None = None,
):
    """Gauss–Newton: relinearise, solve the linear-Gaussian inner program, repeat."""
    values = init_values or {"whitened": jnp.zeros(num_controls)}
    data = (num_times, obs_index, obs_std, obs, obs_mask)
    cost = nonlinear_cost(model_step, initial_of, values["whitened"], *data)
    inner_steps = 0
    for _ in range(num_outer):
        # x₀⁽ⁱ⁾, the current analysis
        reference_initial = initial_of(values["whitened"])  # (N,)

        def trajectory(initial):
            return rollout(model_step, initial, num_times=num_times)  # (N,) → (T, N)

        # x_k⁽ⁱ⁾ = f_{k←0}(x₀⁽ⁱ⁾) and δx₀ ↦ F_{k←0} δx₀: one nonlinear run per outer loop
        reference_states, tangent = jax.linearize(trajectory, reference_initial)
        inner, solution, potential_fn, unravel = fit_map(
            incremental_inner,
            values,
            reference_initial,
            reference_states,
            tangent,
            initial_of,
            num_controls,
            obs_index,
            obs_std,
            obs=obs,
            obs_mask=obs_mask,
        )
        inner_steps += int(solution.stats["num_steps"])
        # Δz, the Gauss–Newton step
        step = inner["whitened"] - values["whitened"]       # (D,)
        # backtrack: halve Δz until the nonlinear 𝒥 decreases
        for _ in range(6):
            trial = values["whitened"] + step               # (D,)
            trial_cost = nonlinear_cost(model_step, initial_of, trial, *data)
            if trial_cost < cost:
                values, cost = {"whitened": trial}, trial_cost
                break
            step = 0.5 * step                               # (D,)
    return values, cost, inner_steps


# full increment: ξ ∈ ℝᴺ, x₀ = μ_x + Σ_x^{1/2} ξ
def full_initial(whitened):
    return background + prior_sqrt.mv(whitened)             # (N,) → (N,)


# reduced increment: z_L ∈ ℝᴿ, x₀ = μ_x + L z_L
def reduced_initial(whitened):
    return background + einx.dot(                           # (R,) → (N,)
        "state basis, basis -> state", prior_features, whitened
    )

With the quasi-static schedule (two outer loops per window stage, four on the full window), incremental 4D-Var reaches the same minimum as the direct solve; the cost over the last four outer loops was 514.2 → 511.0 → 509.5 → 509.4. From the background in one go, eight outer loops had only brought \(\mathcal{J}\) down to 971. Each outer loop costs one nonlinear run; every inner iteration costs only a tangent-linear run and its adjoint, which is the point of the method at scale.

6 · Reduced incremental 4D-Var

Intuition. The same outer loop, with the inner problem restricted to the features: the increment is \(\mat{L}(\vect{z}_L - \vect{z}_L^{(i)})\), so each inner problem has only \(N_\text{feat}\) unknowns and an \(N_\text{feat}\times N_\text{feat}\) Hessian. The outer loop still runs the full nonlinear model. Reduced-order incremental 4D-Var in the literature usually takes the basis from EOFs of a model run; here kernellib's features take that place.

The snippet above covers both: reduced_initial replaces full_initial, and nothing else changes.

7 · Weak-constraint 4D-Var

Intuition. Stop trusting the model completely: after every interval the trajectory may jump by \(\gvect{\eta}_k\), at a price set by \(\gmat{\Sigma}_\eta\). The controls become the whole trajectory's worth of values. With \(\gmat{\Sigma}_\eta = \gmat{\Sigma}_{\eta,t}\otimes\gmat{\Sigma}_{\eta,s}\) the jumps can be correlated in time (a persistent error) and in space (a smooth one), and the whitened errors are mapped through the Kronecker square root.

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


def weak_4dvar(
    model_step: Callable,
    background: Float[Array, "state"],
    prior_sqrt: lx.AbstractLinearOperator,
    error_sqrt: lx.AbstractLinearOperator,
    num_times: int,
    obs_index: Int[Array, "obs"],
    obs_std: float,
    obs: Float[Array, "time obs"] | None = None,
    obs_mask: Bool[Array, "time obs"] | bool = True,
):
    num_state = background.shape[0]
    # ξ ~ 𝒩(0, I)
    whitened = numpyro.sample(                              # event (N,)
        "whitened", dist.Normal(0.0, 1.0).expand([num_state]).to_event(1)
    )
    # x₀ = μ_x + Σ_x^{1/2} ξ
    initial = numpyro.deterministic(                        # (N,)
        "initial", background + prior_sqrt.mv(whitened)
    )
    # w ~ 𝒩(0, I_{T·N})
    whitened_errors = numpyro.sample(                       # event (T·N,)
        "whitened_errors",
        dist.Normal(0.0, 1.0).expand([num_times * num_state]).to_event(1),
    )
    # (η₁, …, η_T) = (Σ_η,t ⊗ Σ_η,s)^{1/2} w
    stacked_errors = error_sqrt.mv(whitened_errors)         # (T·N,)
    errors = einx.id("(t n) -> t n", stacked_errors, t=num_times)  # (T·N,) → (T, N)
    numpyro.deterministic("errors", errors)
    # x_k = f_k(x_{k−1}) + η_k
    states = numpyro.deterministic(                         # (T, N)
        "states", rollout(model_step, initial, errors)
    )
    likelihood = dist.Normal(states[:, obs_index], obs_std)  # batch (T, M)
    numpyro.sample("obs", likelihood.mask(obs_mask).to_event(1), obs=obs)


def weak_4dvar_states(
    model_step: Callable,
    background: Float[Array, "state"],
    prior_sqrt: lx.AbstractLinearOperator,
    error_std: float,
    obs_index: Int[Array, "obs"],
    obs_std: float,
    obs: Float[Array, "time obs"],
):
    whitened = numpyro.sample(                              # event (N,)
        "whitened", dist.Normal(0.0, 1.0).expand(background.shape).to_event(1)
    )
    initial = numpyro.deterministic(                        # (N,)
        "initial", background + prior_sqrt.mv(whitened)
    )

    def transition(previous, obs_k):
        # x_k ~ 𝒩(f_k(x_{k−1}), σ_η² I): every state is a control
        state = numpyro.sample(                             # event (N,)
            "state", dist.Normal(model_step(previous), error_std).to_event(1)
        )
        # y_k ~ 𝒩(H x_k, σ_ε² I)
        likelihood = dist.Normal(state[obs_index], obs_std)  # batch (M,)
        numpyro.sample("obs", likelihood.to_event(1), obs=obs_k)
        return state, None

    # today (numpyro)
    scan(transition, initial, obs)

The time-correlated prior fits the window equally well, but attributes far more of the misfit to a persistent error (the mean kick corresponds to \(\Delta F\) = 0.83 instead of 0.40; truth 1.0). The second function is the state formulation as a numpyro scan: each \(\vect{x}_k\) is a sample site with a Gaussian transition. With \(\sigma_\eta = 0.3\) it converged to the same minimum as the forcing form (\(\mathcal{J}^\star\) = 146.4) in 133 iterations from the background trajectory, without the quasi-static schedule.

8 · Reduced weak-constraint 4D-Var

Intuition. If the model's error has a shape you can name (a bias, a slow drift, a large-scale pattern), estimate the few numbers of that shape instead of a kick at every point and time. A temporal basis \(\chi_c(t)\) (Legendre polynomials over the window) and a spatial basis \(\varphi_j(s)\) carry a small coefficient matrix \(\mat{W}\) with a Kronecker prior. Below, \(\gmat{\chi}\in\mathbb{R}^{T\times N_\chi}\) and \(\gmat{\Phi}\in\mathbb{R}^{D_x \times N_\Phi}\) are the basis matrices, and \(\op{vec}_r\) stacks a matrix row by row, as the code does.

# Shapes: N = 40 sites (state); T = 20 observation times (time);
#         M = 20 observed sites (obs); P spatial modes (basis);
#         C temporal modes (temporal)

# today (geonnax)
# χ_c(t): Legendre polynomials 1, t, … on the window mapped to [−1, 1]
window_time = jnp.linspace(-1.0, 1.0, num_times)            # (T,)
temporal_basis = legendre_polynomials(window_time, num_temporal - 1)  # (T,) → (T, C)
# today (geonnax)
# φ_j(s): 1/√(2b), cos(kπs/b)/√b, sin(kπs/b)/√b, orthonormal on [−b, b]
spatial_basis = fourier_basis_1d(                           # (N, P)
    positions, num_spatial, half_width, boundary="periodic"
)
# today (geonnax)
# κ_j² = (kπ / b)², the Laplacian eigenvalue of each mode
eigvals = fourier_eigenvalues_1d(                           # (P,)
    num_spatial, half_width, boundary="periodic", dtype=positions.dtype
)
# today (kernellib)
# Σ_ω,s = diag(S_k(κ_j)): the error kernel's spectral density at each mode
error_rbf = kl.RBF(variance=error_coeff_var, lengthscale=error_coeff_length)
frequencies = einx.id("p -> p 1", jnp.sqrt(eigvals))        # (P,) → (P, 1)
error_spectrum = error_rbf.spectral_density(frequencies)    # (P, 1) → (P,)
# today (gaussx)
# vec_r W ~ 𝒩(0, Σ_ω,t ⊗ Σ_ω,s)
coeff_cov = gx.Kronecker(                                   # (C·P, C·P) PSD op
    lx.DiagonalLinearOperator(temporal_var), lx.DiagonalLinearOperator(error_spectrum)
)
# today (gaussx)
coeff_sqrt = gx.sqrt(coeff_cov)                             # (C·P, C·P) op


def reduced_weak_4dvar(
    model_step: Callable,
    background: Float[Array, "state"],
    prior_sqrt: lx.AbstractLinearOperator,
    spatial_basis: Float[Array, "state basis"],
    temporal_basis: Float[Array, "time temporal"],
    coeff_sqrt: lx.AbstractLinearOperator,
    num_times: int,
    obs_index: Int[Array, "obs"],
    obs_std: float,
    obs: Float[Array, "time obs"] | None = None,
    obs_mask: Bool[Array, "time obs"] | bool = True,
):
    num_temporal = temporal_basis.shape[1]
    num_spatial = spatial_basis.shape[1]
    whitened = numpyro.sample(                              # event (N,)
        "whitened", dist.Normal(0.0, 1.0).expand(background.shape).to_event(1)
    )
    initial = numpyro.deterministic(                        # (N,)
        "initial", background + prior_sqrt.mv(whitened)
    )
    # ω ~ 𝒩(0, I_{C·P})
    whitened_coeffs = numpyro.sample(                       # event (C·P,)
        "whitened_coeffs",
        dist.Normal(0.0, 1.0).expand([num_temporal * num_spatial]).to_event(1),
    )
    # vec_r W = (Σ_ω,t ⊗ Σ_ω,s)^{1/2} ω
    stacked_coeffs = coeff_sqrt.mv(whitened_coeffs)         # (C·P,)
    coeffs = einx.id("(c p) -> c p", stacked_coeffs, c=num_temporal)  # → (C, P)
    numpyro.deterministic("coeffs", coeffs)
    # η(s_n, t_k) = Σ_c Σ_j χ_c(t_k) W_cj φ_j(s_n)
    # (T, C), (C, P), (N, P) → (T, N)
    errors = einx.dot(
        "t c, c p, n p -> t n", temporal_basis, coeffs, spatial_basis
    )
    numpyro.deterministic("errors", errors)
    states = numpyro.deterministic(                         # (T, N)
        "states", rollout(model_step, initial, errors)
    )
    likelihood = dist.Normal(states[:, obs_index], obs_std)  # batch (T, M)
    numpyro.sample("obs", likelihood.mask(obs_mask).to_event(1), obs=obs)

One coefficient recovers the bias to within its Laplace standard deviation, and the forecast with it (1.47) beats a forecast from the true state with the uncorrected model (1.81). The richer basis fits the window as well but spends its freedom on a time trend; extrapolating that trend past the window hurts the forecast.