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, for GPU support, use:
pip install jaxion[cuda12]See the docs for more info on how to build from source.
Check out the examples/ directory for demonstrations of using Jaxion.
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
- Code repository on GitHub (this page).
- Documentation for up-to-date information about installing and using Jaxion.
If you use this software, please cite it as below.
@software{Mocz_Jaxion_2025,
author = {Mocz, Philip},
doi = {10.5281/zenodo.17438467},
month = oct,
title = {{Jaxion}},
url = {https://github.com/JaxionProject/jaxion},
version = {0.0.4},
year = {2025}
}