jaxion

Differentiable Astrophysics for Fuzzy Dark Matter

Jaxion is a scientific research software for conducting astrophysical and cosmological simulations of Fuzzy Dark Matter, implemented in JAX and utilizing automatic differentiation and multi-GPU acceleration.

Differentiable Physics

Solvers in Jaxion are fully differentiable, enabling gradient-based inference, parameter optimization, and coupling to ML models directly through the simulation.

Scalable Performance

Built with JAX, Jaxion seamlessly scales across CPUs, GPUs, and TPUs, allowing large-scale simulations to run efficiently on modern accelerators.

Modular & Composable

Jaxion’s physics modules and processing tools are composable, making it easy to extend or embed into larger workflows.

Installation

Installation

Quickstart

Quickstart

Examples

Examples

Getting Started

• Installation
• Quickstart

Tutorials & Examples

• Examples
• Tutorial 1: Basic Simulation
• Tutorial 2: Analysis
• Tutorial 3: Sponge Boundary Conditions
• Tutorial 4: Custom Callbacks
• Tutorial 5: Custom Fields
• Tutorial 6: Multiple GPUs

References

• Parameters
• Equations
• Jaxion API
• Developing
• Citing
• About