A simple JAX-powered simulation library for numerical experiments of fuzzy dark matter, stars, gas + more!
Author: Philip Mocz (@pmocz)
Jaxion is built for multi-GPU scalability and is fully differentiable. It is a high-performance JAX-based simulation library for modeling fuzzy dark matter alongside stars, gas, and cosmological dynamics. Being differentiable, Jaxion can seamlessly integrate with pipelines for inverse-problems, inference, optimization, and coupling to ML models.
Jaxion is the simpler companion project to differentiable astrophysics code Adirondax
Install with:
pip install jaxionor see the docs for how to build from source.
Check out the examples/ directory for demonstrations of using Jaxion.
Launch a live demo in Google Colab:
Jaxion is scalable on multiple GPUs!
Jaxion welcomes community contributions of all kinds. Open an issue or fork the code and submit a pull request. Please check out the Contributing Guidelines
- cosmological initial condition generator
- add cosmological factors to gas evolution
- Code repository on GitHub (this page).
- Documentation for up-to-date information about installing and using Jaxion.
If you use this software, please cite the software
@software{Mocz_Jaxion_2025,
author = {Mocz, Philip},
doi = {10.5281/zenodo.17438467},
month = apr,
title = {{Jaxion}},
url = {https://github.com/JaxionProject/jaxion},
version = {0.0.11},
year = {2026}
}and the published JOSS paper:
@article{Mocz_Jaxion_JOSS_2026,
author = {Mocz, Philip},
doi = {10.21105/joss.09542},
journal = {Journal of Open Source Software},
month = jul,
number = {123},
pages = {9542},
title = {{Jaxion: A JAX package for Fuzzy Dark Matter}},
url = {https://joss.theoj.org/papers/10.21105/joss.09542},
volume = {11},
year = {2026}
}