jaxion

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 Jaxion

Install with:

pip install jaxion

or see the docs for how to build from source.

Examples

Check out the examples/ directory for demonstrations of using Jaxion.

Try it out!

Launch a live demo in Google Colab:

High-Performance

Jaxion is scalable on multiple GPUs!

Contributing

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

TODO/Wishlist

• cosmological initial condition generator
• add cosmological factors to gas evolution

Links

• Code repository on GitHub (this page).
• Documentation for up-to-date information about installing and using Jaxion.

Cite this repository

If you use this software, please cite it as below.

@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}
}