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

TL;DR

diffhydro is a scalable, differentiable hydrodynamics code that integrates physics and machine learning for astrophysical simulations, enabling efficient inference, error correction, and complex initial condition reconstructions.

Contribution

This work extends diffhydro with new physics modules, a neural corrector, and multi-device scaling, advancing PDE-constrained inference in astrophysics.

Findings

Good agreement with Athena++ on standard tests

Successful gradient-based reconstruction of complex flows

Simulations on up to 1024^3 elements on GPU clusters

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.