JAX brings automatic differentiation and the XLA compiler together through a NumPy-like API for high performance machine learning research on accelerators like GPUs and TPUs.
This is a curated list of awesome JAX libraries, projects, and other resources specifically coming from the field of astronomy/astrophysics. Contributions are welcome!
For a broader list of JAX resources, check out the awesome-jax repository.
- flowMC - Normalizing-flow enhanced sampling package for probabilistic inference in Jax.
- Papers: arXiv:2211.06397, arXiv:2105.12603
- simple 2D fluid solver in JAX - solves the 2D Euler equations on distributed GPU systems in 300 lines of code
- ripple - Differentiable Gravitational Waveforms with JAX
- Paper: arXiv:2302.05329
- jim - Gravitational-wave data analysis tools in Jax
- Papers: arXiv:2302.05333, arXiv:2404.11397
- ringdown - Black hole ringdowns with JAX and numpyro
- Papers: arXiv:2107.05609
- fiesta - Fast inference of electromagnetic signals and transients with JAX
- jester - JAX-based EOS and TOV solver
- LINX - A fast, differentiable, and extensible public BBN code
- Papers: arXiv:2408.14538, arXiv:2408.14531
- candl - Differentiable Likelihood for CMB Analysis
- Papers: arXiv:2412.00826
- synax - A GPU-accelerated automatic differentiable Galactic synchrotron simulation package
- Papers: arXiv:2410.01136
- jaxoplanet - a
functional-programming-forward
implementation of many features from the
exoplanet and
starry packages built on top of
JAX
- squishyplanet - Tranists and phase curves of non-spherical exoplanets in JAX
- Paper: arXiv:2409.00167
- ExoJAX - differentiable spectrum model from exoplanet and substellar atmosphere
- Papers: arXiv:2105.14782, arXiv:2410.06900
- ∂Lux — differentiable physical optical models from Louis Desdoigts et al.
- jaxspec — differentiable likelihoods for X-ray spectral fitting
- Paper : arxiv:2409.05757
