Advanced features: model knowledge and data¶
Beyond a bare ODE, ode_filters lets you fold in extra model knowledge
(conservation laws) and external data (noisy observations). This notebook
covers both — and neither requires leaving the public gaussian_filter API.
The two ideas share one mechanism. A probabilistic solver works by treating the
ODE itself as a measurement: at each step it asks "how badly does my current
guess violate dx/dt = f(x, t)?" and corrects the Gaussian posterior
accordingly. Anything else you know about the solution — that three quantities
must sum to a constant, or that a sensor reported a noisy value at some time —
is just another measurement folded into the same Kalman update.
Two acronyms up front: an ODE is an ordinary differential equation, and an IVP is the initial value problem of solving one forward from a known start. We use the extended Kalman filter (EKF) under the hood and the Rauch–Tung–Striebel smoother (RTS) to refine the trajectory after a forward pass.
We import the fixed-grid filter, the RTS smoother, the IWP prior, and the
Taylor-mode initializer from the top level. The measurement building blocks
— ODEInformation (the always-on ODE model), Conservation, Measurement,
and prepare_observations — live in the ode_filters.measurement submodule.
The final line sets the shared documentation plot style.
import jax
import jax.numpy as np
import jax.random as jrandom
import matplotlib.pyplot as plt
from ode_filters import (
IWP,
ODEInformation,
gaussian_filter,
rts_smoother,
taylor_mode_initialization,
)
from ode_filters.measurement import (
Conservation,
Measurement,
prepare_observations,
)
jax.config.update("jax_enable_x64", True)
plt.style.use("ode_filters.mplstyle")
Conservation laws¶
Many systems obey an exact algebraic invariant: a quantity that must stay
constant no matter what the trajectory does. The classic example is an SIR
epidemic model, where the population splits into Susceptible, Infected,
and Recovered fractions. People only ever move S → I → R, so the total
S + I + R is conserved — it should equal 1 for all time.
A plain ODE solve does not know this. Each step introduces a tiny numerical error, and those errors accumulate, so the total slowly drifts away from 1. We can pin it down by handing the solver the invariant as a noiseless pseudo-observation.
Here is the SIR vector field. Susceptibles become infected at rate
beta * S * I; infected recover at rate gamma * I. We start with a 1%
infected fraction.
beta_sir = 0.5
gamma_sir = 0.1
def vf_sir(x, *, t):
"""SIR epidemic model with state x = [S, I, R]."""
S, I, R = x # noqa: E741
return np.array(
[
-beta_sir * S * I,
beta_sir * S * I - gamma_sir * I,
gamma_sir * I,
]
)
x0_sir = np.array([0.99, 0.01, 0.0])
tspan_sir = (0.0, 100.0)
N_sir = 100
ts_sir = np.linspace(*tspan_sir, N_sir + 1)
prior_sir = IWP(q=2, d=3, Xi=1.0 * np.eye(3))
mu_0_sir, P0_sqr_sir = taylor_mode_initialization(vf_sir, x0_sir, q=2)
A Conservation(A, p) object enforces the linear invariant A @ x = p
as a noiseless pseudo-observation: at every step the filter additionally
corrects toward states that satisfy it exactly. For the population sum we want
[1, 1, 1] @ [S, I, R] = 1, so A is a single row of ones and p is [1.0].
We attach it through the constraints= argument of ODEInformation. Because a
conservation law is always active and has a fixed shape, it rides along inside
the ODE-information model — unlike time-gated data, which we handle separately
later.
A_sum = np.array([[1.0, 1.0, 1.0]])
p_sum = np.array([1.0])
conservation = Conservation(A_sum, p_sum)
# Two measurement models: the plain ODE, and the ODE plus the invariant.
measure_free = ODEInformation(vf_sir, prior_sir.E0, prior_sir.E1)
measure_cons = ODEInformation(
vf_sir, prior_sir.E0, prior_sir.E1, constraints=[conservation]
)
We solve the IVP both ways — once without the constraint, once with it —
and smooth each forward pass with the RTS smoother. Everything uses the default
("dynamic") calibration; here we focus on the mean trajectory and its
invariant rather than the band.
def solve(measure):
result = gaussian_filter(
mu_0_sir, P0_sqr_sir, prior_sir, measure, tspan_sir, N_sir
)
m_smooth, _ = rts_smoother(prior_sir, result)
return np.array(m_smooth)
m_free = solve(measure_free)
m_cons = solve(measure_cons)
# The conserved quantity along each trajectory, and its drift from 1.
total_free = m_free[:, :3].sum(axis=1)
total_cons = m_cons[:, :3].sum(axis=1)
resid_free = np.abs(total_free - 1.0)
resid_cons = np.abs(total_cons - 1.0)
The left panel shows the SIR curves themselves (from the constrained
solve — they are visually identical either way). The right panel is the point:
on a logarithmic axis we plot how far S + I + R strays from 1. Without the
constraint the residual drifts up to roughly 1e-4; with it, the invariant is
held at float64 machine precision (~1e-16) for the entire run.
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(13, 4))
ax1.plot(ts_sir, m_cons[:, 0], label="S (susceptible)")
ax1.plot(ts_sir, m_cons[:, 1], label="I (infected)")
ax1.plot(ts_sir, m_cons[:, 2], label="R (recovered)")
ax1.set_xlabel("time (days)")
ax1.set_ylabel("population fraction")
ax1.set_title("SIR trajectory")
ax1.legend()
ax2.semilogy(ts_sir, resid_free, label="without constraint")
ax2.semilogy(ts_sir, resid_cons, label="with Conservation")
ax2.set_xlabel("time (days)")
ax2.set_ylabel("|S + I + R - 1|")
ax2.set_title("Drift of the conserved quantity")
ax2.legend()
plt.tight_layout()
plt.show()
print(f"max drift without constraint: {resid_free.max():.2e}")
print(f"max drift with constraint: {resid_cons.max():.2e}")
max drift without constraint: 6.78e-05 max drift with constraint: 4.44e-16
Assimilating data (observations)¶
The second kind of extra knowledge is data: noisy measurements of the true solution recorded at specific times. Folding them in is called data assimilation — the trajectory is pulled toward the observations wherever they land, and floats on the ODE model in between.
To keep this purely about assimilation (and not parameter fitting), we use a system whose dynamics are fully known: exponential decay
$$\frac{dx}{dt} = -k\,x, \qquad x(0) = x_0,$$
with a fixed, known rate k. Nothing about the model is unknown — we are
only injecting external measurements to sharpen an already-correct solve.
Why are data observations handled differently from the ODE measurement
and the conservation law? The ODE model is always on: it fires at every step
and has a fixed shape. Data observations are time-gated — they exist only at
the few instants a sensor reported — and matching "which observation belongs to
which step" is a Python-side decision. To keep the whole solve inside
jax.lax.scan (and therefore jit-safe and fast), we pre-bake the observations
against the time grid once, with prepare_observations, and pass the result
as obs_model=. The ODE model and the data thus stay cleanly decoupled.
First the known system and a probabilistic solve with no data, to use as a
reference. We also record the true analytic solution x0 * exp(-k t).
k_decay = 0.8 # fixed, known decay rate
def vf_decay(x, *, t):
"""Exponential decay dx/dt = -k x with a known rate k."""
return -k_decay * x
x0_decay = np.array([2.0])
tspan_decay = (0.0, 6.0)
N_decay = 60
ts_decay = np.linspace(*tspan_decay, N_decay + 1)
prior_decay = IWP(q=2, d=1, Xi=0.5 * np.eye(1))
mu_0_decay, P0_sqr_decay = taylor_mode_initialization(vf_decay, x0_decay, q=2)
measure_decay = ODEInformation(vf_decay, prior_decay.E0, prior_decay.E1)
x_true = float(x0_decay[0]) * np.exp(-k_decay * ts_decay)
Now we manufacture some synthetic data: take the true solution at a handful of times and add Gaussian noise. This stands in for a real sensor. The observations are sparse (every eighth grid point) and noisy.
obs_idx = np.arange(8, N_decay, 8)
z_t = ts_decay[obs_idx]
noise_std = 0.1
key = jrandom.PRNGKey(0)
z_clean = float(x0_decay[0]) * np.exp(-k_decay * z_t)
z = z_clean[:, None] + noise_std * jrandom.normal(key, (len(z_t), 1))
A Measurement(A, z, z_t, noise) describes a linear observation
A @ x = z recorded at times z_t, corrupted by Gaussian noise of the given
variance. We observe the state directly, so A = [[1.0]]. Then
prepare_observations aligns these measurements to the solver's time grid and
returns an ObsModel — a fixed-shape container the scan can consume. Its
boolean mask records which grid steps carry an observation.
A_obs = np.array([[1.0]])
measurement = Measurement(A_obs, z, z_t, noise=noise_std**2)
obs_model = prepare_observations([measurement], prior_decay.E0, ts_decay)
print(f"grid steps: {N_decay}")
print(f"steps with data: {int(obs_model.mask.sum())}")
grid steps: 60 steps with data: 7
When an obs_model is supplied, each filter step becomes a two-stage
Kalman update:
- an ODE update — correct the prediction toward
dx/dt = -k x; - an observation update — correct again toward any data at this step.
The order matters: the observation update sees the tighter, ODE-informed
posterior rather than the diffuse prediction. The two contributions are scored
separately: result.log_likelihood is the (post-calibration) marginal
likelihood of the ODE residuals, and result.log_likelihood_obs is the marginal
likelihood of the data. For data-driven model comparison or parameter fitting use
log_likelihood_obs; the ODE-residual log_likelihood is rescaled by the dynamic
calibration each step and so is not comparable across calibration modes.
result = gaussian_filter(
mu_0_decay,
P0_sqr_decay,
prior_decay,
measure_decay,
tspan_decay,
N_decay,
obs_model=obs_model,
)
m_smooth, P_smooth_sqr = rts_smoother(prior_decay, result)
m_smooth = np.array(m_smooth)
P_smooth_sqr = np.array(P_smooth_sqr)
# Recover the marginal standard deviation of x from the square-root covariance.
P_smooth = np.einsum("nij,nik->njk", P_smooth_sqr, P_smooth_sqr)
x_std = np.sqrt(P_smooth[:, 0, 0])
print(f"ODE log-likelihood: {float(result.log_likelihood):.3f}")
print(f"data log-likelihood: {float(result.log_likelihood_obs):.3f}")
ODE log-likelihood: 337.630 data log-likelihood: 6.909
The smoothed mean tracks the true decay curve, the two-sigma band narrows wherever an observation pins the trajectory down, and the noisy data points scatter around the truth as expected.
fig, ax = plt.subplots(figsize=(10, 4))
ax.plot(ts_decay, x_true, "k--", label="true solution")
ax.plot(ts_decay, m_smooth[:, 0], label="smoothed estimate")
ax.fill_between(
ts_decay,
m_smooth[:, 0] - 2 * x_std,
m_smooth[:, 0] + 2 * x_std,
alpha=0.3,
label="2-sigma band",
)
ax.scatter(z_t, z[:, 0], color="orange", zorder=3, label="noisy observations")
ax.set_xlabel("t")
ax.set_ylabel("x(t)")
ax.set_title("Exponential decay with assimilated data")
ax.legend()
plt.show()
Summary¶
Both extra-knowledge mechanisms reuse the same Kalman machinery as the bare ODE
solve, and both stay inside the public gaussian_filter API:
- Conservation laws. A
Conservation(A, p)object enforces the linear invariantA @ x = pas a noiseless pseudo-observation. It is always active, so it rides along inODEInformation(..., constraints=[...]). With the SIR population sum constrained, the invariant held at float64 machine precision (~1e-16) instead of drifting to~1e-4. - Data assimilation. A
Measurement(A, z, z_t, noise)describes a noisy linear observation at given times. Because data is time-gated, it is decoupled from the always-on ODE model: pre-bake it withprepare_observationsand pass the resultingObsModelasobs_model=. Each step then runs a two-stage update (ODE, then data), and the ODE and data fits are scored separately asresult.log_likelihoodandresult.log_likelihood_obs.
For more exotic needs, the library also offers nonlinear observation models
through BlackBoxMeasurement and TransformedMeasurement, and Gaussian-process
priors via MaternPrior. To pick a correction scheme (EKF order, iterated
updates) see the corrections guide, and the how-to-choose guide for matching a
prior and step size to your problem.