The six-step cycle, written in probabilities

Step 1 · The simple model is a joint distribution

A state-space model is four conditional densities. Their product is the joint distribution of everything, observed or not:

  α ─────────► θ
               │
               ▼
  x₀ ──► x₁ ──► x₂ ──► … ──► x_T        transition  p(xₖ | xₖ₋₁, θ, uᶜᵒᵛ, uₖ)
         │      │             │
         ▼      ▼             ▼
         y₁     y₂            y_T        likelihood  p(yₖ | xₖ, mₖ)

"Physics you can simulate" means you can draw from this joint by ancestral sampling: draw each variable given its parents, in order. That is all step 1 has to provide, and it is all steps 3 and 5 need for training. The prior predictive \(p(\vect{y}) = \iint p(\gvect{\theta}) p(\vect{x} \mid \gvect{\theta}) p(\vect{y} \mid \vect{x}) \mathrm{d}\vect{x} \mathrm{d}\gvect{\theta}\) is what the model expects to observe before seeing any data; drawing from it and comparing with real maps is the first check of a model.

import einx
import equinox as eqx
import jax
from jaxtyping import Array, Float, PRNGKeyArray

# Shapes: T days (time); N state values (state); M observations per day (obs);
#         P parameters (param); D any density's event shape

# not pinned (distribution library: numpyro / distrax / tfp-style)
# Dist: an object with .sample(key) and .log_prob(value)


class StateSpaceModel(eqx.Module):
    """p(θ) p(x₀ | θ) ∏ₖ p(xₖ | xₖ₋₁, θ, uᶜᵒᵛ, uₖ) p(yₖ | xₖ, mₖ): four conditionals."""

    # p(θ)
    def prior_params(self) -> Dist: ...

    # p(x₀ | θ)
    def prior_state(self, params) -> Dist: ...

    # p(xₖ | xₖ₋₁, θ, uᶜᵒᵛ, uₖ): physics, a hybrid, or an emulator
    def transition(self, state, params, covariates, control) -> Dist: ...

    # p(yₖ | xₖ, mₖ)
    def likelihood(self, state, mask) -> Dist: ...


def sample_joint(
    model: StateSpaceModel,
    covariates: Float[Array, "variable state"],
    controls: Float[Array, "time state"],
    masks: Float[Array, "time obs state"],
    *,
    key: PRNGKeyArray,
) -> tuple[Float[Array, "param"], Float[Array, "time state"], Float[Array, "time obs"]]:
    params_key, state_key, time_key = jax.random.split(key, 3)
    # θ ~ p(θ)
    params = model.prior_params().sample(params_key)  # → (P,)
    # x₀ ~ p(x₀ | θ)
    state = model.prior_state(params).sample(state_key)  # → (N,)

    def day(state, inputs):
        control, mask, day_key = inputs
        transition_key, obs_key = jax.random.split(day_key)
        # xₖ ~ p(xₖ | xₖ₋₁, θ, uᶜᵒᵛ, uₖ)
        state = model.transition(state, params, covariates, control).sample(
            transition_key
        )                                             # (N,) → (N,)
        # yₖ ~ p(yₖ | xₖ, mₖ)
        obs = model.likelihood(state, mask).sample(obs_key)  # (N,) → (M,)
        return state, (state, obs)

    day_keys = jax.random.split(time_key, controls.shape[0])  # → (T, 2)
    # ancestral sampling: one draw of (θ, x₁:T, y₁:T) from the joint
    _, (states, observations) = jax.lax.scan(         # → (T, N), (T, M)
        day, state, (controls, masks, day_keys)
    )
    return params, states, observations

Markers: not pinned marks the density interface (.sample, .log_prob), which any of the JAX distribution libraries provides; user code is defined on this page. Shapes follow the code-sketch standard.

Step 2 · Model-based inference is the posterior of that joint

Condition the joint on what was observed. Bayes' rule gives the target and its normaliser, the evidence:

Two marginals are usually what is wanted: the state, \(p(\vect{x} \mid \vect{y}) = \int p(\gvect{\theta}, \vect{x} \mid \vect{y}) \mathrm{d}\gvect{\theta}\), and the parameters, \(p(\gvect{\theta} \mid \vect{y})\propto p(\gvect{\theta}) p(\vect{y} \mid \gvect{\theta})\), whose likelihood \(p(\vect{y} \mid \gvect{\theta}) = \int p(\vect{x} \mid \gvect{\theta}) p(\vect{y} \mid \vect{x}) \mathrm{d}\vect{x}\) integrates the state out. For a state-space model the integrals over \(\vect{x}\) split day by day into two recursions that hold for any densities:

Filtering is predict-then-update, forward in time; smoothing (reanalysis) runs backward once the window is filtered. Every algorithm in data assimilation is a way of evaluating these integrals, or the posterior they define, approximately. All of them start from the log joint:

# Shapes: T days (time); N state values (state); M observations per day (obs)


def log_joint(
    model, params, initial, states, observations, covariates, controls, masks
):
    # log p(θ)
    log_prior_params = model.prior_params().log_prob(params)  # → ()
    # log p(x₀ | θ)
    log_prior_state = model.prior_state(params).log_prob(initial)  # → ()
    previous = einx.rearrange(                        # (N,), (T, N) → (T, N)
        "state, time state -> (1 + time) state", initial, states
    )[:-1]

    def day_terms(prev, state, control, mask, obs):
        # log p(xₖ | xₖ₋₁, θ, uᶜᵒᵛ, uₖ)
        log_trans = model.transition(prev, params, covariates, control).log_prob(state)
        # log p(yₖ | xₖ, mₖ)
        log_lik = model.likelihood(state, mask).log_prob(obs)
        return log_trans, log_lik

    log_trans, log_lik = jax.vmap(day_terms)(         # → (T,), (T,)
        previous, states, controls, masks, observations
    )
    # log p(θ, x₀:T, y₁:T)
    #   = log p(θ) + log p(x₀|θ) + ∑ₖ log p(xₖ|xₖ₋₁, θ) + ∑ₖ log p(yₖ|xₖ)
    total = log_prior_params + log_prior_state        # → ()
    total = total + einx.sum("time ->", log_trans)    # → ()
    total = total + einx.sum("time ->", log_lik)      # → ()
    return total

The most general way to evaluate the recursions is to represent each density by samples: the bootstrap particle filter. It assumes nothing about linearity or Gaussianity, and it also returns an estimate of \(\log p(\vect{y} \mid \gvect{\theta})\):

# Shapes: T days (time); N state values (state); M obs per day (obs);
#         P_n = 1,000 particles (particle)


def particle_filter(model, params, observations, covariates, controls, masks, *, key):
    """Bootstrap filter: p(xₖ | y₁:ₖ) as weighted samples, and log p(y₁:T | θ)."""
    init_key, time_key = jax.random.split(key)
    init_keys = jax.random.split(init_key, NUM_PARTICLES)  # → (P_n, 2)
    # x₀⁽ⁱ⁾ ~ p(x₀ | θ)
    particles = jax.vmap(                             # → (P_n, N)
        lambda k: model.prior_state(params).sample(k)
    )(init_keys)

    def day(particles, inputs):
        obs, control, mask, day_key = inputs
        move_key, resample_key = jax.random.split(day_key)
        move_keys = jax.random.split(move_key, NUM_PARTICLES)  # → (P_n, 2)
        # predict: xₖ⁽ⁱ⁾ ~ p(xₖ | xₖ₋₁⁽ⁱ⁾, θ)  (Monte Carlo Chapman–Kolmogorov)
        particles = jax.vmap(                         # (P_n, N) → (P_n, N)
            lambda z, k: model.transition(z, params, covariates, control).sample(k)
        )(particles, move_keys)
        # update: log wₖ⁽ⁱ⁾ = log p(yₖ | xₖ⁽ⁱ⁾)
        log_weights = jax.vmap(                       # (P_n, N) → (P_n,)
            lambda z: model.likelihood(z, mask).log_prob(obs)
        )(particles)
        # log p(yₖ | y₁:ₖ₋₁) ≈ log (1/P_n) ∑ᵢ wₖ⁽ⁱ⁾
        log_sum = einx.logsumexp("[particle]", log_weights)  # (P_n,) → ()
        log_increment = log_sum - jnp.log(NUM_PARTICLES)  # → ()
        # resample: i ~ Categorical(wₖ / ∑ w)
        ancestors = jax.random.categorical(           # → (P_n,)
            resample_key, log_weights, shape=(NUM_PARTICLES,)
        )
        particles = particles[ancestors]              # (P_n, N)
        return particles, (particles, log_increment)

    day_keys = jax.random.split(time_key, observations.shape[0])  # → (T, 2)
    _, (filtered, log_increments) = jax.lax.scan(     # → (T, P_n, N), (T,)
        day, particles, (observations, controls, masks, day_keys)
    )
    # log p(y₁:T | θ) = ∑ₖ log p(yₖ | y₁:ₖ₋₁)
    log_evidence = einx.sum("time ->", log_increments)  # (T,) → ()
    return filtered, log_evidence

It is exact as the number of particles grows, but in high dimension the weights collapse onto one particle unless that number grows exponentially. That is why ocean data assimilation uses the cheaper families below, each with its own assumption.

Step 2 is the oracle not because an algorithm is exact, but because the posterior it targets is the correct answer under the model: everything in steps 3–5 is judged by how close it gets to it.

Why the linear case closes and the nonlinear case does not

The scalar example shows all three effects. With a speed-squared observation, the posterior has a branch for each sign of the current, weighted by the prior, and its shape changes with \(y\):

Two patterns to keep. The linearised standard deviation is 0.24 for every \(y\): that is the linear-Gaussian property that the posterior covariance ignores the data, carried into a problem where it is false (the exact sd runs from 0.25 to 1.19). And the Laplace approximation becomes more confident as the evidence for two branches grows: a single Gaussian cannot say "east or west". (For a truly linear observation, \(y = x + \varepsilon\), the grid and the closed-form Kalman update agree to four decimals.)

Step 3 · The emulator is a learned conditional density

The expensive object in step 2 is the transition: it is evaluated or sampled for every particle, ensemble member or iteration. An emulator replaces that one density with a learned one, keeping the same interface:

Train it on transitions drawn by running step 1, by maximum likelihood. That is the same as minimising the KL divergence from the simulator's transition to the emulator's, averaged over the states and parameters the simulations visit:

# Shapes: S simulated transitions (sim); N state values (state); P parameters (param)


class EmulatedTransition(eqx.Module):
    """q(xₖ | xₖ₋₁, θ, uᶜᵒᵛ, uₖ; θ_emu): a learned conditional density, same interface as p."""

    network: eqx.Module

    def __call__(self, state, params, covariates, control) -> Dist:
        # user
        # (μ, log σ) = network(xₖ₋₁, θ, uᶜᵒᵛ, uₖ; θ_emu)
        mean, log_scale = self.network(state, params, covariates, control)  # → (N,) ×2
        # not pinned (distribution library)
        # q(xₖ | ·; θ_emu) = 𝒩(μ, diag σ²); a flow or a diffusion model also fits here
        return Normal(mean, jnp.exp(log_scale))


def emulator_loss(emulator, previous, params, covariates, controls, nexts):
    # pairs (xₖ₋₁, θ, uₖ, xₖ) drawn by running step 1: xₖ ~ p(xₖ | xₖ₋₁, θ, uᶜᵒᵛ, uₖ)
    def nll(prev, theta, control, nxt):
        # −log q(xₖ | xₖ₋₁, θ, uᶜᵒᵛ, uₖ; θ_emu)
        return -emulator(prev, theta, covariates, control).log_prob(nxt)

    losses = jax.vmap(nll)(previous, params, controls, nexts)  # → (S,)
    # 𝔼_p[−log q(·; θ_emu)] = 𝔼_{xₖ₋₁,θ}[ KL(p(xₖ|·) ‖ q(xₖ|·; θ_emu)) ] + const
    return einx.mean("sim ->", losses)                # (S,) → ()

The deterministic emulator is the case \(q(\cdot; \gvect{\theta}_\text{emu}) = \delta(\vect{x}_k - f(\vect{x}_{k-1},\gvect{\theta}; \gvect{\theta}_\text{emu}))\). Its log-likelihood is undefined, so it is trained by mean squared error instead, which is the Gaussian log-likelihood with a fixed variance. What is lost is exactly the spread: a deterministic emulator claims to be right everywhere, and step 4 then trusts it fully. A learned \(\gvect{\sigma}\) reports where the emulator is unsure, and that uncertainty flows into the posterior like model error.

The expectation is over \(p(\vect{x}_{k-1},\gvect{\theta})\): the states and parameters the training simulations visit. An emulator is only trained where step 1 went. Step 4 will use it where the posterior puts mass, which may be a different place.

In the scalar example the "simulator" is the observation \(p(y \mid x) = \mathcal{N}(x^2, 0.25^2)\). An emulator \(q(y \mid x; \gvect{\theta}_\text{emu})\) trained only on the prior's typical range \(\vert x \vert \lesssim 2\) is unconstrained beyond it; a posterior that reaches past it, such as the two branches near \(\pm2.2\) after observing \(y = 5\), rests on the emulator's extrapolation.

Step 4 · Emulator-based inference: the same posterior, one density swapped

Put \(q(\cdot; \gvect{\theta}_\text{emu})\) in place of the simulator's transition and run step 2 unchanged. The target becomes

In code this is one argument: pass the emulator as the model's transition to the same log_joint or particle_filter. Dividing the two posteriors shows exactly how emulator error enters:

The posterior is distorted by the emulator's density ratio, renormalised. Errors where the posterior has no mass do not matter; errors where it does matter fully. So the emulator's training distribution (step 3's expectation) should cover the regions that posteriors visit, which is why emulators are often refined with simulations drawn near the posterior.

Step 5 · Amortized inference: learn the posterior as a function of the data

Steps 2 and 4 solve a new problem for every window of observations. Step 5 learns a conditional density that answers for all of them at once:

The forward KL, averaged over the data the model can produce, needs only samples of the joint from step 1. The intractable posterior drops out because it does not depend on \(\gvect{\nu}\):

That is simulation-based inference: simulate \((\gvect{\theta}, \vect{x}, \vect{y})\) triples, fit a conditional density, done. The alternative is the reverse KL per window, the ELBO, which needs the likelihood but no simulated truth:

Forward KL tends to spread mass over every mode (it is penalised for missing one); reverse KL tends to pick one (it is penalised for putting mass where the posterior has none). The amortization gap is what one network loses against solving each window separately.

# Shapes: S simulated windows (sim); P parameters (param); T, N, M as before;
#         L posterior draws (draw)


def amortized_loss(posterior_net, params, states, observations, masks):
    # (θ, x, y) ~ p(θ, x, y): S draws of sample_joint, nothing else is needed
    def nll(theta, z, y, m):
        # user
        # q(θ, x | y, m; ν): a conditional flow, diffusion or mixture
        density = posterior_net(y, m)
        # −log q(θ, x | y; ν)
        return -density.log_prob((theta, z))

    losses = jax.vmap(nll)(params, states, observations, masks)  # → (S,)
    # 𝔼_{p(y)}[ KL(p(θ, x | y) ‖ q(θ, x | y; ν)) ] = 𝔼_{p(θ,x,y)}[−log q(·; ν)] + const
    return einx.mean("sim ->", losses)                # (S,) → ()


def calibration_rank(posterior_net, model, *, key, num_draws=99):
    """Simulation-based calibration: the rank of the true θ among posterior draws."""
    sim_key, draw_key = jax.random.split(key)
    # θ_true ~ p(θ), y ~ p(y | θ_true)
    true_params, _, observations = sample_joint(model, *INPUTS, key=sim_key)
    draw_keys = jax.random.split(draw_key, num_draws)  # → (L, 2)
    # θ⁽ˡ⁾ ~ q(θ | y; ν)
    draws = jax.vmap(                                 # → (L, P)
        lambda k: posterior_net(observations, MASKS).sample_params(k)
    )(draw_keys)
    below = draws < true_params                       # (L, P), (P,) → (L, P)
    # rank = #{l : θ⁽ˡ⁾ < θ_true}, uniform on 0…L when q(·; ν) is calibrated
    return einx.sum("draw param -> param", below)     # (L, P) → (P,)

Using a posterior: estimate, forecast, predict

Training produces \(\gvect{\theta}_\text{emu}\) and \(\gvect{\nu}\); inference produces a posterior. Neither is the end product. What is used is a handful of integrals against the posterior, and with posterior draws \((\gvect{\theta}, \vect{x})^{(l)}\), \(l = 1,\ldots,L\), every one of them is an average:

# Shapes: L posterior draws (draw); N state values (state); J forecast days (lead);
#         M observations per day (obs)


def use_posterior(model, posterior_draws, covariates, controls, mask, *, key):
    """From posterior samples of (θ, x_T) to estimates, a forecast and a predictive."""
    params, final_states = posterior_draws            # (L, P), (L, N)

    # --- estimate: summaries of p(x_T | y) ---
    # 𝔼[x_T | y] ≈ (1/L) ∑ₗ x_T⁽ˡ⁾
    mean = einx.mean("draw state -> state", final_states)  # (L, N) → (N,)
    per_state = einx.rearrange("draw state -> state draw", final_states)  # → (N, L)
    # 5 % and 95 % quantiles of each state value
    bounds = jax.vmap(lambda v: jnp.quantile(v, LEVELS))(per_state)  # → (N, 2)

    # --- forecast: p(x_{T+j} | y) = ∫ p(x_{T+j} | x_T, θ) p(x_T, θ | y) dx_T dθ ---
    def rollout(theta, state, rollout_key):
        def day(state, inputs):
            control, day_key = inputs
            # x_{T+j} ~ p(x_{T+j} | x_{T+j−1}, θ): simulator, or emulator q(·; θ_emu)
            state = model.transition(state, theta, covariates, control).sample(day_key)
            return state, state

        day_keys = jax.random.split(rollout_key, controls.shape[0])  # → (J, 2)
        _, path = jax.lax.scan(day, state, (controls, day_keys))  # → (J, N)
        return path

    forecast_key, predictive_key = jax.random.split(key)
    rollout_keys = jax.random.split(forecast_key, final_states.shape[0])  # → (L, 2)
    # one forecast path per posterior draw: an ensemble with honest spread
    forecast = jax.vmap(rollout)(params, final_states, rollout_keys)  # → (L, J, N)

    # --- predict: p(y_* | y) = ∫ p(y_* | x) p(x | y) dx, for a new observation ---
    predictive_keys = jax.random.split(predictive_key, final_states.shape[0])
    # y_*⁽ˡ⁾ ~ p(y_* | x⁽ˡ⁾)
    predictive = jax.vmap(                            # (L, N) → (L, M)
        lambda z, k: model.likelihood(z, mask).sample(k)
    )(forecast[:, 0], predictive_keys)

    # --- event probability: P(A | y) = 𝔼[𝟙_A(x) | y] ---
    # user
    # 𝟙_A(x): e.g. "the jet sits north of 40° N", per draw
    in_event = event_indicator(forecast[:, -1])       # (L, N) → (L,)
    probability = einx.mean("draw ->", in_event)      # (L,) → ()
    return mean, bounds, forecast, predictive, probability

The scalar example, used

Take the posterior at \(y = 1.2\) and a simple transition for tomorrow's current, which persists and decays: \(p(x' \mid x) = \mathcal{N}(x'; 0.8 x, 0.3^2)\). Compare the exact posterior with the Laplace approximation from step 2:

Three lessons that hold well beyond the toy. The posterior mean is a fine summary of a unimodal posterior and a misleading one here: it sits in the valley between the branches, at a current the water almost certainly does not have. The forecast inherits the posterior's shape: pushing both branches forward keeps a 26 % chance of a westward current tomorrow, which the single-Gaussian forecast reduces to 0.3 %. And every use above needs draws, not a single field: a deterministic analysis can be estimated from, but it cannot be forecast with honest spread or turned into a probability.

Step 6 · Improve: compare distributions, then change the structure

Each learned density has a reference: the distribution it was built to approximate. The whole model has one more reference, the data.

The posterior predictive is the model's forecast of new data after seeing the old: \(p(\vect{y}_* \mid \vect{y}) = \int p(\vect{y}_* \mid \vect{x}) p(\vect{x} \mid \vect{y}) \mathrm{d}\vect{x}\). If held-out observations fall in its tails more or less often than its stated probabilities, something is wrong. If the misfit is random, the inner loop (re-estimate \(\vect{x}\) and \(\gvect{\theta}\)) is enough. If it has structure, the outer loop changes the model \(\mathcal{M}\): a different transition family, a richer likelihood, other priors. Competing structures are compared by their evidence: