A JAX-based gravitational-wave population inference toolkit for parametric models

Installation | Documentation | Tutorials | Analysis on 🤗 | FAQs | Citing GWKokab

Overview

GWKokab is a high-performance, flexible, and easy-to-use toolkit for gravitational-wave population inference. Built on top of JAX, it enables efficient Bayesian inference for a wide range of parametric population models while remaining fully compatible with modern GPU/TPU-accelerated workflows.

The framework is designed to support scalable hierarchical inference and rapid experimentation with astrophysical population models, including mass, spin, redshift, and eccentricity distributions of compact binary mergers.

Contributing

We welcome contributions from the community. If you would like to contribute to GWKokab, please see the contributing guidelines.

Citing GWKokab

If you use GWKokab in your research, please cite the following works:

@ARTICLE{2026PhRvD.113j3003Q,
  author          = {{Qazalbash}, M. and {Zeeshan}, M. and {O'Shaughnessy}, R.},
  title           = "{Implementation to identify the properties of multiple
                  populations of gravitational wave sources}",
  journal         = {\prd},
  keywords        = {Astrophysics and astroparticle physics, General Relativity
                  and Quantum Cosmology, High Energy Astrophysical Phenomena,
                  Instrumentation and Methods for Astrophysics},
  year            = 2026,
  month           = may,
  volume          = 113,
  number          = 10,
  eid             = 103003,
  pages           = 103003,
  doi             = {10.1103/krnm-3vrf},
  archivePrefix   = {arXiv},
  eprint          = {2509.13638},
  primaryClass    = {gr-qc},
  adsurl          = {https://ui.adsabs.harvard.edu/abs/2026PhRvD.113j3003Q},
  adsnote         = {Provided by the SAO/NASA Astrophysics Data System}
}

@Misc{gwkokab2024github,
  author          = {{Qazalbash}, Meesum and {Zeeshan}, Muhammad and
                  {O'Shaughnessy}, Richard},
  title           = {{GWKokab}: A JAX-based gravitational-wave population
                  inference toolkit for parametric models},
  url             = {https://github.com/kokabsc/gwkokab},
  year            = 2024
}