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

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.

Paper

References (11)

05XRISM Science Team2020 · arXiv e-prints
062025, arXiv e-prints
07Demonstrated forward model performance and initial-field reconstruction in strongly nonlinear settings (turbulence, instabilities, and self-gravitating flows)
08Reworked the code for scalable multi-device parallelism, with demonstrations up to 1024 3
09Implemented a differentiable multigrid Poisson solver for gravity and related elliptic fluxes/forces
10Implemented techniques to reduce memory and accelerator-time costs during backpropagation, including adjoint methods and frozen gradients
11Introduced new physics modules, including (non-equilibrium) heating/cooling, turbulence, conduction, and multigrid Poisson gravity

Similar papers

© 2026 NYSGPT2525 LLC