Second-order systems¶
Many physical models are written as second-order ordinary differential equations (ODEs) of the form
$$x''(t) = f\big(x(t),\, x'(t),\, t\big),$$
where the acceleration $x''$ depends on the position $x$ and the velocity $x'$. The textbook trick is to reduce such a system to first order by stacking $[x, x']$ into a new state and solving a larger first-order ODE. It works, but it throws away structure: the solver no longer knows that the second component is the derivative of the first.
This notebook solves a second-order ODE directly. The idea is simple and fits the probabilistic solver naturally:
- We use an integrated-Wiener-process prior
IWP(q, d=1). A second-order ODE constrains the acceleration $x''$, so the prior must model the state at least up to $x''$ — that is, $q \ge 2$. We takeq=3(state $[x, x', x'', x''']$): modelling one derivative beyond the constraint gives a smoother, higher-order prior that tracks the oscillation accurately over many periods. As in the quickstart, higherqmeans higher-order accuracy —q=2would visibly lag the true solution over this horizon. - The information operator
SecondOrderODEInformationthen asserts the dynamics directly on the acceleration: it forces $x''(t) = f(x, x', t)$ at every step. The matricesE0,E1,E2extract $x$, $x'$, $x''$ from the state, and the constraint is applied throughE2.
Because of this, the vector field you write has the second-order signature directly:
def vf(x, dx, *, t):
return ddx # the acceleration
As with any probabilistic solve, the output is not a single trajectory but a Gaussian posterior over the solution: a mean plus a calibrated uncertainty band. Our test case is the damped harmonic oscillator, which has a known closed-form solution we can check against.
We import the public pieces we need: the IWP prior, the
SecondOrderODEInformation measurement model, gaussian_filter for the
forward pass, rts_smoother (RTS = Rauch-Tung-Striebel) for the backward
smoothing pass, and taylor_mode_initialization to set up the initial
state. The last line loads the shared plot style used across the docs.
import jax
import jax.numpy as np
import matplotlib.pyplot as plt
from ode_filters import (
gaussian_filter,
rts_smoother,
taylor_mode_initialization,
)
from ode_filters.measurement import SecondOrderODEInformation
from ode_filters.priors import IWP
jax.config.update("jax_enable_x64", True)
plt.style.use("ode_filters.mplstyle")
The damped harmonic oscillator¶
A mass on a spring with friction obeys
$$x'' = -\omega^2\, x - \gamma\, x',$$
where $\omega$ is the natural frequency and $\gamma = 2\zeta\omega$ sets the damping (with damping ratio $\zeta$). The restoring force pulls the mass back toward zero and the friction term bleeds off energy, so the oscillation decays over time.
Notice how the vector field below reads off the second-order equation
verbatim: it takes the position x and velocity dx and returns the
acceleration. We start the oscillator displaced to $x=1$ and at rest
($x'=0$), and integrate over $t \in [0, 30]$ on a grid of 100 steps.
omega, zeta = 1.0, 0.1
gamma = 2 * zeta * omega
omega_d = np.sqrt(1 - zeta**2) * omega # damped frequency
tspan, N = (0, 30), 100
ts = np.linspace(tspan[0], tspan[1], N + 1)
def vf(x, dx, *, t):
return -(omega**2) * x - gamma * dx
x0, dx0 = np.array([1.0]), np.array([0.0])
Prior, measurement model, and initial state¶
We choose an IWP(q=3, d=1) prior. The order q=3 means the modeled
state at each time is $[x, x', x'', x''']$, and d=1 because the oscillator
is scalar (one spatial dimension). A second-order ODE only requires q=2
(so that $x''$ is in the state and can be constrained), but a slightly higher
order buys markedly better accuracy on an oscillation that runs for many
periods — q=2 here would noticeably lag the true curve. Xi is the prior's
diffusion scale matrix (here 1×1); the identity is a fine default, and the
per-step diffusion calibration sets the overall scale on top of it.
SecondOrderODEInformation wraps the vector field together with the
extraction matrices E0, E1, E2 (which pick out $x$, $x'$, $x''$).
It is what enforces $x''(t) = f(x, x', t)$ during filtering.
Finally, taylor_mode_initialization builds a consistent initial mean (and a
zero/Dirac initial covariance) from the initial position and velocity. We
pass order=2 (the ODE is second order, so we supply $x(0)$ and $x'(0)$);
it then differentiates the vector field to fill in the remaining initial
derivatives, matching the q=3 state.
prior = IWP(q=3, d=1, Xi=np.eye(1))
measure = SecondOrderODEInformation(vf, prior.E0, prior.E1, prior.E2)
mu_0, Sigma_0_sqr = taylor_mode_initialization(vf, (x0, dx0), q=3, order=2)
Solve: filter, then smooth¶
gaussian_filter runs the forward extended Kalman filter (EKF) pass,
sweeping left to right and enforcing the ODE at each step. We leave
calibration at its default ("dynamic"), so the solver rescales the
diffusion per step from the ODE residuals (an approximate quasi-MLE) and the
uncertainty band adapts to the observed misfit instead of using a fixed scale.
rts_smoother then runs the backward pass, refining every state estimate
using information from later times. We jit both passes for speed; the
prior, measurement model, time span, and step count are static arguments.
jit_filter = jax.jit(gaussian_filter, static_argnums=(2, 3, 4, 5))
jit_smoother = jax.jit(rts_smoother, static_argnums=(0,))
result = jit_filter(mu_0, Sigma_0_sqr, prior, measure, tspan, N)
m_smooth, P_smooth_sqr = jit_smoother(prior, result)
m_smooth, P_smooth_sqr = np.array(m_smooth), np.array(P_smooth_sqr)
Compare against the closed-form solution¶
An underdamped oscillator ($\zeta < 1$) has the standard closed-form solution
$$x(t) = e^{-\zeta\omega t}\Big(\cos(\omega_d t) + \tfrac{\zeta\omega}{\omega_d}\sin(\omega_d t)\Big),$$
with damped frequency $\omega_d = \sqrt{1-\zeta^2}\,\omega$; this is the textbook result for our initial conditions $x(0)=1$, $x'(0)=0$.
We recover the marginal variance of the position from the square-root covariance via $P = P_{\text{sqr}}^\top P_{\text{sqr}}$, take its first diagonal entry, and plot a 2-sigma band around the smoothed mean. The band hugs the analytic curve and contains the true solution on this example, illustrating the step-dependent uncertainty the probabilistic solver provides on top of a point estimate.
x_true = np.exp(-zeta * omega * ts) * (
np.cos(omega_d * ts) + (zeta * omega / omega_d) * np.sin(omega_d * ts)
)
P_smooth = np.einsum("nij,nik->njk", P_smooth_sqr, P_smooth_sqr)
margin = 2 * np.sqrt(P_smooth[:, 0, 0])
plt.figure(figsize=(10, 3))
plt.plot(ts, m_smooth[:, 0], label="smoothed")
plt.plot(ts, x_true, "k--", label="true")
plt.fill_between(
ts, m_smooth[:, 0] - margin, m_smooth[:, 0] + margin, alpha=0.3
)
plt.xlabel("t")
plt.ylabel("x(t)")
plt.legend()
plt.show()
When to use second-order form¶
If your model is naturally second order — oscillators, $N$-body and
celestial mechanics, structural dynamics, anything Newtonian where forces
set accelerations — reach for IWP(q=2) and SecondOrderODEInformation.
Solving in second-order form keeps the position-velocity-acceleration
relationship baked into the prior instead of asking the solver to
rediscover it from a flattened first-order system. The code also stays
closer to the physics: you write the acceleration law vf(x, dx, *, t)
directly. For genuinely first-order problems, the ordinary first-order
interface (see the basic-usage tutorial) remains the right tool.