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

We present the extension of the differentiable hydrodynamics code, diffhydro , enabling scalable partial differential equation-constrained inference and integrated hybrid physics–machine learning (ML) models for a wide range of astrophysical applications. New physics additions include radiative heating/cooling, Ornstein–Uhlenbeck–driven turbulence, and self-gravity via multigrid Poisson. We demonstrate good agreement with the Athena++ code on standard validation tests like 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 multidevice 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, and self-gravitating, radiatively cooling flows. The code is written in JAX, and the solver’s modular finite-volume components are compiled by Accelerated Linear Algebra into fused accelerator kernels, delivering high-throughput forward runs and tractable differentiation through long integrations.