2025/12/15 by Benjamin Horowitz, Zarija Lukić, Horowitz, Benjamin +3
Physics and Astronomy · Computer Science · #Galaxies: Formation, Evolution, Phenomena #Stochastic Gradient Optimization Techniques #Model Reduction and Neural Networks
paper · pdf · doi:10.48550/arxiv.2512.13403
We present the extension of the differentiable hydrodynamics code, diffhydro, enabling scalable PDE-constrained inference and integrated hybrid physics-ML models for a wide range of astrophysical applications. New physics additions include radiative heating/cooling, OU-driven turbulence, and self-gravity via multigrid Poisson. We demonstrate good agreement with the Athena++ code on standard validation tests such as Sedov-Taylor, Kelvin-Helmholtz, and driven/decaying turbulence. We further introduce a solver-in-the-loop neural corrector that reduces coarse-grid errors during time integration while preserving stability. The addition of custom adjoints facilitates efficient end-to-end gradients and multi-device scaling. We present simulations up to 10243 elements, run on distributed GPU systems, and we show gradient-based reconstructions of complex initial conditions in turbulent, self-gravitating, radiatively cooling flows. The code is written in JAX, and the solver's modular finite-volume components are compiled by XLA into fused accelerator kernels, delivering high-throughput forward runs and tractable differentiation through long integrations.