jaxion 0.0.4

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.

Links

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

Testing

Jaxion is tested with pytest. Tests are included in the tests/ folder.

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

Cite this repository

TODO XXX