Diffusion Calibration¶

Probabilistic ODE filters quantify uncertainty by propagating a Gaussian covariance. The magnitude of that covariance is governed by a scalar diffusion sigma^2 on top of the prior's structural scale matrix Xi. The filter is correct in the mean even with sigma = 1, but the posterior uncertainty is scaled to match the residual magnitude only when sigma^2 is calibrated from the observed ODE residuals.

The ode_filters.calibration subpackage provides three estimators that all consume the predicted-observation marginal (m_z, P_z_sqr) already returned by every EKF step.

Workflows¶

Post-hoc MLE on a fixed-step run¶

Cheapest workflow: run any existing fixed-step filter, then rescale. posthoc_mle_sigma_sqr is the joint MLE of sigma^2 under the assumption that one constant scale explains every step's whitened residual -- the same Gaussian likelihood the per-step estimator maximises, but with a single parameter rather than one per step. The closed form makes posthoc_mle_sigma_sqr numerically identical to aggregate_sigma_sqr(seq, kind="mean") applied to the per-step trace.

import jax.numpy as np
from ode_filters import gaussian_filter
from ode_filters.calibration import posthoc_mle_sigma_sqr, rescale_sqr_seq
from ode_filters.measurement import ODEInformation
from ode_filters.priors import IWP, taylor_mode_initialization


def vf(x, *, t):
    return x * (1 - x)


prior = IWP(q=2, d=1)
mu_0, S0 = taylor_mode_initialization(vf, np.array([0.1]), q=2)
measure = ODEInformation(vf, prior.E0, prior.E1)

# Run an uncalibrated fixed-step filter.
result = gaussian_filter(mu_0, S0, prior, measure, (0.0, 5.0), N=50, calibration="none")

sigma_sqr_hat = posthoc_mle_sigma_sqr(result.mz, result.Pz_sqr)

# Rescale stored covariances (returns a new array; result.P_sqr is unchanged).
P_seq_sqr_calibrated = rescale_sqr_seq(result.P_sqr, sigma_sqr_hat)

This is appropriate for smooth problems where a single global scalar captures the residual size well.

Joint / latent-force priors: pass prior=

For JointPrior / PrecondJointPrior, the diffusion sigma^2 belongs to the ODE-state block only -- the hidden/input block must not be rescaled (Schmidt et al. 2021). The default rescale_sqr / rescale_sqr_seq scale the whole covariance, which is wrong for these priors. Pass the prior so the rescaling is routed through its block-aware apply_state_sigma_to_cov_sqr:

P_seq_sqr_calibrated = rescale_sqr_seq(result.P_sqr, sigma_sqr_hat, prior=prior)

For non-joint priors (IWP, MaternPrior, ...) passing prior= is optional and gives the same result as the default whole-covariance scaling. The online calibration path (calibration="dynamic"/"diagonal") is already block-aware, so this only matters for the post-hoc workflow above.

Per-step quasi-MLE inside the adaptive solver¶

gaussian_filter_adaptive runs the per-step quasi-MLE by default and applies it to the stored covariances on every accepted step. See Adaptive Step-Size Control.

from ode_filters import gaussian_filter_adaptive

result = gaussian_filter_adaptive(
    mu_0, S0, prior, measure, save_at=np.linspace(0.0, 5.0, 50),
    atol=1e-5, rtol=1e-3,
)

For the per-step sigma_sqr_seq diagnostics on the accepted-step trajectory, use the lower-level trajectory driver sqr_adaptive_loop from the submodule (see Adaptive Step-Size Control):

from ode_filters.filters.ode_filter_adaptive import sqr_adaptive_loop

traj = sqr_adaptive_loop(
    mu_0, S0, prior, measure, (0.0, 5.0), atol=1e-5, rtol=1e-3,
)
# traj.sigma_sqr_seq holds the per-step estimates.

Pass calibration="none" to record the estimates without scaling the stored covariances (diagnostics only).

Manual per-step estimation¶

If you want the per-step number outside the adaptive solver (e.g., for diagnostics on an existing fixed-step run):

from ode_filters.calibration import quasi_mle_sigma_sqr

sigma_sqr_per_step = [
    float(quasi_mle_sigma_sqr(mz, Pz_sqr))
    for mz, Pz_sqr in zip(result.mz, result.Pz_sqr, strict=True)
]

The estimator is a one-line solve_triangular on the existing square-root covariance.

Aggregation¶

A per-step trace can be turned into a single number in three ways:

from ode_filters.calibration import aggregate_sigma_sqr

mean_sigma = aggregate_sigma_sqr(seq, kind="mean")       # joint-MLE equivalent
last_sigma = aggregate_sigma_sqr(seq, kind="last")       # use the final estimate
running    = aggregate_sigma_sqr(seq, kind="running")    # cumulative mean trace

"running" returns an array of the same length as seq -- useful when you want to display the calibrated uncertainty intermediate during integration without waiting for the end.

Note on measurement noise¶

The quasi-MLE assumes the predicted-observation covariance is proportional to sigma^2. When the EKF measurement update adds a non-zero R (e.g. for genuine observations under data assimilation), the full output covariance is H P_pred H.T + R and feeding that into the quasi-MLE conflates the diffusion scale with the measurement noise (it biases sigma_hat^2 downward).

The adaptive solver handles this transparently: internally it recomputes the noise-free predicted residual covariance and uses that for calibration, matching what probnum and probdiffeq do. With the default ODEInformation (no measurements, R = 0) the distinction is moot.

For the post-hoc workflow on a fixed-step run with non-zero R, the caller must pass a noise-free Pz_sqr to posthoc_mle_sigma_sqr -- the covariances returned by the filter include R and must be recomputed from the saved (result.m_pred, result.P_pred_sqr) first.

What calibration cannot fix¶

Calibration scales the posterior to match the size of the residuals. It does not correct a prior with the wrong smoothness (q too low), a vector field that is miscoded, or a Jacobian that the EKF linearisation cannot approximate. To check that the calibrated filter is honest, plot the whitened residuals: for each step the standardized innovation z_n = P_z_sqr^{-T} m_z should look like standard normal noise, so ||z_n||^2 / d hovers near 1 on a well-specified problem (much larger means the uncertainty is too tight; much smaller means it is too loose).