diffhydro: Inverse Multiphysics Modeling and Embedded Machine Learning in Astrophysical Flows

Abstract

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 1024^3 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.