Differentiable Astrophysics

Philip Mocz

Keywords: astrophysics, differentiable physics, inverse-problem

TL;DR: We introduce a JAX-powered differentiable astrophysics code (fuzzy dark matter, gas, stars) to solve inverse-problems

Abstract: We introduce Jaxion, a simple and extensible simulation library built on JAX for numerical experiments in astrophysics and scientific machine learning. Jaxion enables multiphysics simulations of fuzzy dark matter, stars, and gas, combining spectral methods, particle-mesh, and finite volume schemes. Leveraging JAX's automatic differentiation and hardware acceleration, Jaxion provides a flexible platform for rapid prototyping and differentiable simulations. This differentiability allows gradients to be computed through entire simulation pipelines, enabling seamless integration with optimization, inference, and machine learning workflows. As a result, users can efficiently run simulations on GPUs and treat them as black-box differentiable functions for inverse problems or hybrid physics-ML modeling.

Email Sharing: We authorize the sharing of all author emails with Program Chairs.

Data Release: We authorize the release of our submission and author names to the public in the event of acceptance.

Submission Number: 21