ode-filters 0.4.0

Basic ODE filtering and smooting implementation.

ODE Filters

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

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

Documentation

Full documentation is available at paufisch.github.io/ode_filters.

Development

Run the test suite:

uv run pytest --cov=ode_filters --cov-report=term-missing

Build documentation locally:

uv run mkdocs serve

License

MIT License - see LICENSE for details.