A JAX-based gravitational-wave inference toolkit
Jim is a JAX-based toolkit for estimating parameters of gravitational-wave sources through Bayesian inference. At its core, Jim uses the normalizing-flow enhanced sampler flowMC to improve the convergence of gradient-based MCMC sampling.
Built on JAX, Jim leverages hardware acceleration to achieve significant speedups on GPUs. The toolkit also implements likelihood-heterodyning (Cornish et al. 2010, Cornish & Littenberg 2021) for efficient gravitational-wave likelihood computation.
See the accompanying paper, Wong, Isi, Edwards (2023), for more details.
Warning
Jim is under active development and the API may change. Use at your own risk! Consider forking a specific version if you need API stability. We aim to release a stable v1.0.0 version in 2026.
[Documentation and examples are a work in progress]
The simplest way to install Jim is through pip:
pip install jimGW
This will install the latest stable release and its dependencies. Jim is built on JAX and flowMC. By default, this installs the CPU version of JAX from PyPI. If you have a GPU and want to leverage hardware acceleration, install the CUDA-enabled version:
pip install jimGW[cuda]
If you want to install the latest version of Jim, you can clone this repo and install it locally:
git clone https://github.com/GW-JAX-Team/jim.git
cd jim
pip install -e .
Jim's performance varies with available hardware. Under optimal conditions with CUDA, parameter estimation for a binary neutron star can complete in ~1 minute on an NVIDIA A100 GPU (see paper for details).
If no GPU is available, JAX will automatically fall back to CPU execution, displaying:
No GPU/TPU found, falling back to CPU.
If you use Jim in your research, please cite the accompanying paper:
@article{Wong:2023lgb,
author = "Wong, Kaze W. K. and Isi, Maximiliano and Edwards, Thomas D. P.",
title = "{Fast Gravitational-wave Parameter Estimation without Compromises}",
eprint = "2302.05333",
archivePrefix = "arXiv",
primaryClass = "astro-ph.IM",
doi = "10.3847/1538-4357/acf5cd",
journal = "Astrophys. J.",
volume = "958",
number = "2",
pages = "129",
year = "2023"
}