A jax based affine-invariant MCMC sampler that can leverage GPUs to speed up sampling for computationally intensive likelihoods. It implements the Goodman-Weare algorithm as described in dfm++ and is inspired by the popular emcee library. The just-in-time compilation together with vectorized likelihood evaluation for the walkers gives significant speed-up even on CPUs when compared to emcee
To install jaims, please clone this repository and then run python setup.py install inside it
You can also install this via pip using
pip install jaims
To run it on a GPU, you must have an installation of jaxlib compatible with your CUDA version. For more information, please refer to the official guidelines
The API for jaims is slightly different from emcee. This might change in the future.