PFJAX: Particle Filtering in JAX

What is PFJAX?

PFJAX is a collection of tools for estimating the parameters of state-space models using particle filtering methods, with JAX as the backend for JIT-compiling models and automatic differentiation.

---

Installation

This will clone the repo into a subfolder pfjax, from where you (i) issue the git clone command and (ii) install the package from source.

git clone https://github.com/mlysy/pfjax
cd pfjax
pip install .

Quickstart

A brief introduction to PFJAX.

Documentation

This is a work in progress! Current modules include:

• The quickstart guide.

• A comparison of gradient and hessian algorithms based on particle filters, which in turn are used for conducting parameter inference.

• An example of parameter inference using stochastic optimization.

• An example of parameter inference using Markov chain Monte Carlo.

• The API reference documentation.

Developers

The instructions below assume that uv is being used for dependency management. Instructions for installing uv are available here.

Testing

From within the pfjax folder:

uv run --group test pytest

Building Documentation

The documentation is build using Quarto + MkDocs-Material. The latter comes as a Python package installed by uv, but the former must be installed separately. Once Quarto is installed, from within the pfjax folder:

uv run --group docs quarto render docs/
uv run --group docs mkdocs build