Skip to content

ODE Filters

PyPI Python CI Docs Coverage

A JAX-based implementation of probabilistic ODE solvers using Gaussian filtering and smoothing. This package provides tools for solving ordinary differential equations while quantifying uncertainty through Bayesian inference.

Features

  • Pure JAX implementation - Fully differentiable and JIT-compilable
  • Square-root filtering - Numerically stable EKF and RTS smoothing
  • Flexible priors - Integrated Wiener Process (IWP), Matern, and joint priors
  • First and second-order ODEs - Native support for both ODE types
  • Constraint handling - Conservation laws and time-varying measurements
  • State-parameter estimation - Joint inference with hidden states
  • Black-box measurements - Custom observation models with autodiff Jacobians
  • Transformed measurements - Nonlinear state transformations with chain-rule Jacobians
  • Diffusion calibration - Per-step quasi-MLE, running aggregate, and post-hoc MLE estimators for the prior diffusion scale (see Calibration)
  • Adaptive step-size control - Tolerance-based accept/reject with a PI controller, built on the calibrated local error estimate (see Adaptive steps)

Installation

Install the latest release from PyPI:

pip install ode-filters

Or install from source with development dependencies:

git clone https://github.com/paufisch/ode_filters.git
cd ode_filters
pip install -e ".[dev]"

Quick Example

import jax.numpy as np
from ode_filters.filters import ekf1_sqr_loop, rts_sqr_smoother_loop
from ode_filters.measurement import ODEInformation
from ode_filters.priors import IWP, taylor_mode_initialization

# Define ODE: dx/dt = -x (exponential decay)
def vf(x, *, t):
    return -x

x0 = np.array([1.0])
tspan = [0, 5]

# Setup prior and measurement model
prior = IWP(q=2, d=1, Xi=0.5 * np.eye(1))
mu_0, Sigma_0_sqr = taylor_mode_initialization(vf, x0, q=2)
measure = ODEInformation(vf, prior.E0, prior.E1)

# Run filter and smoother
m_seq, P_sqr, *_, G, d, P_back, _, _ = ekf1_sqr_loop(
    mu_0, Sigma_0_sqr, prior, measure, tspan, N=50
)
m_smooth, P_smooth_sqr = rts_sqr_smoother_loop(
    m_seq[-1], P_sqr[-1], G, d, P_back, N=50
)

Package Structure

ode_filters/
├── filters/          # EKF and RTS smoothing loops
├── inference/        # Square-root Gaussian algebra
├── measurement/      # ODE and observation models
└── priors/           # Gaussian Markov process priors

Next Steps