jaxion 0.0.5

A differentiable simulation library for fuzzy dark matter in JAX

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

Getting started

Install with:

pip install jaxion

or, for GPU support, use:

pip install jaxion[cuda12]

See the docs for more info on how to build from source.

Examples

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

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

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 = oct,
    title = {{Jaxion}},
      url = {https://github.com/JaxionProject/jaxion},
  version = {0.0.4},
     year = {2025}
}