Abstract

Contributed Talk - Splinter CompQuantum   (MW-2050)

LUNA: A Performance Portable, Scalable, Differentiable, Collisionless N-Body + Hydrodynamics Codefor Astrophysical Fluid Dynamics

Ulrich Philipp Steinwandel
Max Planck Institute for Astrophysics

We present LUNA, a performance-portable, differentiable astrophysical fluid-dynamics code written entirely in JAX that couples a classical finite-volume hydrodynamics scheme to self-gravity, a collisionless particle–mesh sector, and radiative cooling within a single end-to-end differentiable framework. The compressible Euler equations are solved in one, two, and three dimensions on uniform Cartesian grids using the unsplit MUSCL–Hancock method with piecewise-linear reconstruction, TVD slope limiting, and the HLLC approximate Riemann solver (with HLL fallback). Self-gravity is provided by a geometric full-multigrid Poisson solver coupled to the particle–mesh sector, and the code further supports super-comoving cosmological integration and optically-thin cooling. Because every operation—including the Riemann solver, slope limiter, multigrid solve, and cooling sub-cycle—is expressed as composable JAX transformations, the entire solver is differentiable: gradients of any scalar output with respect to initial conditions, boundary data, or model parameters are obtained at machine precision by reverse- and forward-mode automatic differentiation, with a custom implicit-function VJP for the Poisson solve, without hand-written adjoint code. Domain decomposition across GPUs uses JAX's `shard_map` with explicit halo exchange; we demonstrate weak and strong scaling to 1024³ zones on up to 256 AMD MI250X accelerators, with parallel scaling behaviour matching the established AthenaK code. A suite of standard tests—Sod, Sedov–Taylor, double Mach reflection, Kelvin–Helmholtz, and self-gravitating problems—verifies second-order convergence and correct shock capturing, while single-GPU throughput is benchmarked against AthenaK. We illustrate the framework on protostellar collapse, driven turbulence, a galactic wind, cosmological dark matter, and gradient-based inverse problems, opening a path toward differentiable simulation in astrophysics.