Quickstart¶
This notebook solves an ODE with a probabilistic solver: instead of a single trajectory you get a Gaussian posterior over the solution — a mean and a calibrated uncertainty.
(ODE = ordinary differential equation; IVP = initial value problem — an ODE together with a known starting value.)
The recipe is always the same four ingredients:
- a prior describing how smooth the solution is (an integrated Wiener process),
- a measurement model asserting that the solution's derivative matches the vector field (the ODE-information operator),
- an extended Kalman filter (EKF) forward pass, and
- an RTS smoother (Rauch-Tung-Striebel) backward pass that refines the forward estimate into the full smoothed posterior.
New to the idea? Read What is a probabilistic ODE solver? for the intuition first, then come back here.
We import the four public building blocks plus the Taylor-mode initializer,
and set the shared documentation plot style. JAX is imported as np so the code
reads like NumPy while running on JAX (which gives us JIT compilation later).
import jax
import jax.numpy as np
import matplotlib.pyplot as plt
from ode_filters import (
IWP,
ODEInformation,
gaussian_filter,
rts_smoother,
taylor_mode_initialization,
)
jax.config.update("jax_enable_x64", True)
plt.style.use("ode_filters.mplstyle")
Example problem: the logistic ODE¶
We solve the logistic equation, a classic model of constrained growth:
$$\dot{x}(t) = x(t)\,(1-x(t)), \qquad x(0) = x_0 = 0.01, \qquad t \in [0, 10].$$
It starts near zero, grows fastest in the middle, and saturates at $x = 1$. We pick it because it has a known closed-form solution, so at the end we can check the probabilistic estimate against the truth.
1. Define the initial value problem¶
The vector field vf returns the time-derivative $\dot{x}$ given the current
state. The state is a length-1 array (d = 1), x0 is the initial value, and
tspan is the time interval. Time is passed as a keyword-only argument, matching
the library convention.
def vf(x, *, t):
return x * (1 - x)
x0 = np.array([0.01])
tspan = (0, 10)
d = x0.shape[0]
2. Select a prior¶
The prior is a Gaussian process that encodes our smoothness assumptions about the solution before we look at the ODE. We use an integrated Wiener process (IWP):
qis the number of derivatives the prior tracks. The state it models is the stack $[x, \dot{x}, \ddot{x}, \dots]$ up to the $q$-th derivative. Higherqassumes a smoother solution and is more accurate when that assumption holds;q = 2(value, velocity, acceleration) is a solid default.Xiis the diffusion scale of the prior. Here (d = 1) it is a scalar controlling how much the highest derivative is allowed to wander — larger means a looser, more flexible prior. In generalXiis ad-by-dpositive-(semi)definite matrix entering the process noise askron(Q_scalar(h), Xi), so a diagonalXisets independent per-component scales and off-diagonal entries add cross-component correlation.
taylor_mode_initialization then turns the IVP into a consistent starting Gaussian:
it computes the initial mean for the full stacked state $[x, \dot{x}, \ddot{x}]$ by
repeatedly differentiating the vector field at x0 (Taylor-mode initialization), and
pairs it with a zero (Dirac) covariance — a zero-uncertainty start, so the filter
begins from a point estimate that already satisfies the ODE.
q = 2
Xi = 0.5 * np.eye(d)
prior = IWP(q, d, Xi=Xi)
mu_0, P0_sqr = taylor_mode_initialization(vf, x0, q)
3. Define the measurement model¶
A probabilistic ODE solver treats the ODE itself as data. At each time step it forms the residual $\dot{x}(t) - f(x(t), t)$ and conditions on it being zero — this is the "ODE-as-data" idea.
ODEInformation is exactly this residual operator. It needs two projection
matrices supplied by the prior:
E0selects the solution value $x$ from the stacked state, andE1selects the first derivative $\dot{x}$.
With those it can assert $\dot{x}(t) = f(x(t), t)$ at every grid point.
measure = ODEInformation(vf, prior.E0, prior.E1)
4. Run the filter and smoother¶
The forward pass is gaussian_filter. We JIT-compile it for speed. The
static_argnums=(2, 3, 4, 5) marks the prior, measurement model, time span, and
step count N as static — they define the structure of the computation rather
than numeric inputs, so JAX bakes them into the compiled code. (tspan is a tuple
precisely so it is hashable and usable as a static argument.)
We keep the default "dynamic" calibration, which rescales the prior diffusion at
each step from the ODE residuals via a scalar per-step quasi-MLE (Bosch et al. 2021).
This adapts the posterior covariance to the observed misfit instead of using a fixed
diffusion, so the band reflects the solver's estimated uncertainty rather than an
arbitrary fixed scale. It is an approximate estimator (conditioned on the EK1
linearization, and assuming the residual covariance scales linearly in the
diffusion), not an exact calibration guarantee.
jit_filter = jax.jit(gaussian_filter, static_argnums=(2, 3, 4, 5))
jit_smoother = jax.jit(rts_smoother, static_argnums=(0,))
# A deliberately coarse grid, so the uncertainty band is clearly visible
# below. On a fine grid the solver is so confident the band vanishes.
N = 10
result = jit_filter(mu_0, P0_sqr, prior, measure, tspan, N)
gaussian_filter returns a FilterResult whose fields are stacked sequences,
one entry per grid point. The two we use are:
result.m— the posterior mean of the stacked state, andresult.P_sqr— its square-root covariance, from which the full covariance isP = P_sqr.T @ P_sqr.
It also carries result.log_likelihood, a scalar marginal log-likelihood of the
ODE-information residuals (after diffusion calibration). It can be used for model
comparison, but because calibration rescales the per-step process noise, these
values are only comparable between runs that use the same calibration mode (we
don't need it here).
The filter is a single forward sweep. rts_smoother then runs a backward pass over
the same result, refining every estimate using information from the whole interval.
The smoothed posterior is what we ultimately report.
m_smooth, P_smooth_sqr = jit_smoother(prior, result)
5. Visualise the results¶
We turn each square-root covariance into a standard deviation and plot the smoothed mean with its 2-sigma band. Because the logistic equation has a closed-form solution, we can overlay the exact curve and check that the band contains the true solution on this example.
def value_and_sigma(m_seq, P_seq_sqr):
"""Posterior value (component 0) and its standard deviation per grid point."""
m_seq, P_seq_sqr = np.array(m_seq), np.array(P_seq_sqr)
# Reconstruct each covariance from its square root: P = P_sqr.T @ P_sqr.
P_seq = np.einsum("nij,nik->njk", P_seq_sqr, P_seq_sqr)
return m_seq[:, 0], np.sqrt(P_seq[:, 0, 0])
def logistic_exact(t):
return 1.0 / (1.0 + (1.0 / x0[0] - 1.0) * np.exp(-t))
ts = np.linspace(tspan[0], tspan[1], N + 1)
t_dense = np.linspace(tspan[0], tspan[1], 400)
mean, sigma = value_and_sigma(m_smooth, P_smooth_sqr)
The logistic equation has the closed-form solution
$$x(t) = \frac{1}{1 + (1/x_0 - 1)\,e^{-t}},$$
which we overlay as a dashed reference. A well-calibrated band should contain this curve almost everywhere.
plt.figure(figsize=(9, 4), dpi=120)
plt.plot(t_dense, logistic_exact(t_dense), "k--", label="exact solution")
plt.plot(ts, mean, "o-", label="probabilistic estimate")
plt.fill_between(
ts, mean - 2 * sigma, mean + 2 * sigma, alpha=0.3, label=r"2$\sigma$ band"
)
plt.xlabel("t")
plt.ylabel("x(t)")
plt.legend()
plt.show()
The shaded band is the 2-sigma credible interval — the uncertainty quantification a classical solver does not give you. It contains the true curve everywhere and is widest through the middle, where the solution changes fastest and the coarse steps leave the solver least sure.
This band is calibrated: it reflects genuine confidence, not a fixed guess. The clearest way to see that is to vary the grid resolution.
def smoothed_solution(n_steps):
res = jit_filter(mu_0, P0_sqr, prior, measure, tspan, n_steps)
m_s, P_s = jit_smoother(prior, res)
grid = np.linspace(tspan[0], tspan[1], n_steps + 1)
m, s = value_and_sigma(m_s, P_s)
return grid, m, s
fig, axes = plt.subplots(1, 2, figsize=(11, 4), dpi=120, sharey=True)
for ax, n_steps in zip(axes, [6, 40]):
grid, m, s = smoothed_solution(n_steps)
ax.plot(t_dense, logistic_exact(t_dense), "k--", label="exact")
ax.plot(grid, m, label="estimate")
ax.fill_between(grid, m - 2 * s, m + 2 * s, alpha=0.3, label=r"2$\sigma$ band")
ax.set_title(f"N = {n_steps} steps")
ax.set_xlabel("t")
ax.legend()
axes[0].set_ylabel("x(t)")
fig.tight_layout()
plt.show()
With only 6 steps the band is wide: the solver reports that big jumps leave it uncertain between grid points. Refine to 40 steps and the band collapses to a thin line — more evaluations, more confidence — while still containing the truth at every resolution. That step-dependent uncertainty, tracking the discretization error, is what a probabilistic solver gives you on top of a point estimate.
That's the entire recipe: pick a prior, state the ODE as a measurement, filter forward, smooth backward — and read off a solution with its uncertainty. The other notebooks apply these same four ingredients to richer systems.