JAX-SMfSB (jsmfsb)

SMfSB code in Python+JAX

Python code relating to the book Stochastic Modelling for Systems Biology, third edition.

There is a regular Python+Numpy package on PyPI, smfsb, which has complete coverage of the book. If you are new to the book and/or this codebase, that might be a simpler place to start.

This package covers all of the core simulation and inference algorithms from the book, including the parsing of SBML and SBML-shorthand models. These core algorithms will run very fast, using JAX. Computationally intensive algorithms will typically run between 50 and 150 times faster than they would using the regular smfsb package, even without a GPU (but YMMV). You must install JAX (which is system dependent), before attempting to install this package. See the JAX documentation for details, but for a CPU-only installation, it should be as simple as pip install jax.

Once you have JAX installed and working correctly, you can install this package with:

pip install jsmfsb

To upgrade already installed package:

pip install --upgrade jsmfsb

You can test that your installation is working by entering the following at a python prompt:

import jax
import jsmfsb

If these both return silently, you are probably good to go.

Documentation

Note that the book, and its associated github repo is the main source of documentation for this library. The code in the book is in R, but the code in this library is supposed to mirror the R code, but in Python.

For an introduction to this library, see the jax-smfsb tutorial.

Converting from the smfsb python package

The API for this package is very similar to that of the smfsb package. The main difference is that non-deterministic (random) functions have first argument corresponding to a JAX random number key rather than a numpy random number generator. Users are responsible for ensuring that keys are not re-used. See the relevant section of the JAX documentation for further information regarding random numbers in JAX code.

Further information

For further information, see the demo directory and the API documentation. Within the demos directory, see shbuild.py for an example of how to specify a (SEIR epidemic) model using SBML-shorthand and sbml-params.py for an example of how to modify the parameters of a model parsed from SBML (or SBML-shorthand). Also see step_cle_2df.py for a 2-d reaction-diffusion simulation. For parameter inference (from time course data), see abc-cal.py for ABC inference, abc_smc.py for ABC-SMC inference and pmmh.py for particle marginal Metropolis-Hastings MCMC-based inference. There are many other demos besides these.

You can view this package on GitHub or PyPI.

Citing

If you use this software in your research, please cite the paper:

Wilkinson, D. J. (2025). jax-smfsb: A Python library for stochastic systems biology modelling and inference. Journal of Open Source Software, 10(106), 7491.

Contributing

If you have problems with this software, please start an Issue or a Discussion. Pull requests containing bug fixes are welcome.

Copyright (C) 2024-2026 Darren J Wilkinson