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.

💡 Research Summary

This paper introduces an advanced extension of “diffhydro,” a differentiable hydrodynamics framework designed for scalable PDE-constrained inference and integrated hybrid physics-ML modeling in astrophysical contexts. The core innovation lies in leveraging JAX-based differentiable programming to enable end-to-end gradient computation through complex, long-term fluid simulations. This capability is transformative, as it allows researchers to solve “inverse problems”—reconstructing complex initial physical conditions from observed final states—using gradient-based optimization.

The authors have significantly expanded the physics capabilities of the code, incorporating essential astrophysical processes such as radiative heating and cooling, OU-driven turbulence, and self-gravity implemented via a multigrid Poisson solver. To ensure scientific rigor, the framework was validated against the industry-standard Athena++ code, demonstrating high fidelity in fundamental tests including Sedov-Taylor blast waves, Kelvin-Helmholtz instabilities, and both driven and decaying turbulence simulations.

A major technical contribution is the introduction of a “solver-in-the-loop neural corrector.” This hybrid component utilizes a neural network to learn and mitigate discretization errors inherent in coarse-grid numerical integration, effectively bridging the gap between computational efficiency and physical accuracy without compromising numerical stability. Furthermore, the implementation utilizes XLA (Accelerated Linear Algebra) to compile modular finite-volume components into fused accelerator kernels. This optimization allows for high-throughput forward simulations and tractable differentiation even during long-duration integrations.

The scalability of the framework is demonstrated through massive simulations involving up to $1024^3$ elements, executed on distributed GPU systems. The paper concludes by showcasing the practical utility of diffhydro through the gradient-based reconstruction of complex initial conditions within turbulent, self-gravitating, and radiatively cooling flows. By combining the precision of classical computational fluid dynamics (CFD) with the learning power of machine learning, diffhydro establishes a powerful new paradigm for high-fidelity astrophysical modeling and parameter estimation.