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

TL;DR

diffhydro extends differentiable hydrodynamics to include radiative heating/cooling, OU-driven turbulence, and self-gravity via a multigrid Poisson solver, enabling end-to-end gradient-based inference and solver-in-the-loop learning on large-scale astrophysical flows. Built in JAX with accelerator-native execution, it provides a modular finite-volume solver with custom adjoints, supports multi-device scaling, and validates against Athena++ across standard tests while performing gradient-based reconstructions of complex initial conditions. The work demonstrates practical PDE-constrained optimization and differentiable ML augmentation within a rigorous astrophysical context, including solver-in-the-loop neural correctors and high-resolution simulations up to $1024^3$ cells. These capabilities offer a pathway to data-driven, physically constrained inference and model calibration directly inside the forward model, with potential to connect simulations to multi-wavelength observations and statistical cosmology through differentiable pipelines.

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.

Figures (19)

• Figure 1: Schematic overview of forward and inverse workflows enabled by diffhydro. (a) Traditional simulation: an initial condition (IC) specified by parameters $\boldsymbol{\phi}_0$ is mapped to the conservative state $\mathbf{U}_0$ and advanced forward through repeated applications of the hydrodynamic update operator $\mathcal{H}$, yielding a final state $\mathbf{U}_f$ from which observables $\mathcal{O}$ are computed. (b) Initial--condition/parameter reconstruction: the same forward model is paired with a likelihood $\mathcal{L}(\mathcal{O};\mathrm{data})$; automatic differentiation backpropagates gradients through the observable postprocessing map and each timestep ($d\mathrm{Obs}$ and $d\mathcal{H}$) to obtain $\partial \mathcal{L}/\partial \boldsymbol{\phi}_0$, enabling gradient-based updates of ICs or physical parameters. (c) Solver--in--the--loop optimization: a reference simulation provides target states or observables at multiple times, defining a summed loss $\mathcal{L}=\sum_i \mathcal{L}_i$. A differentiable forward run $\hat{\mathbf{U}}$ is controlled by trainable parameters $\boldsymbol{\theta}$ (e.g., closure/source-term models $\mathcal{B}_\theta$); gradients $\partial \mathcal{L}/\partial \boldsymbol{\theta}$ are obtained by backpropagating through both $\mathcal{H}$ and $\mathcal{B}_\theta$, allowing optimization of embedded models while preserving the underlying finite-volume numerics.
• Figure 2: Forward model of Turbulence. The same turbulent initial conditions in a $384^3$ box are evolved with and without driving with diffhydro. The driving force injects kinetic energy, which decays dynamically into thermal energy.
• Figure 3: Example output from the cooling/heating test problem in diffhydro. Shown are mid-plane slices through an $384^3$ ISM-like box evolved from turbulent initial conditions with Ornstein–Uhlenbeck forcing and the cooling/heating module enabled. Top: logarithmic gas temperature, spanning the cold to warm neutral phases, with cooling producing sharp temperature contrasts in compressed regions. Bottom: logarithmic mass density in the same slice, exhibiting a connected network of overdense filaments and shells generated by the driven turbulence. The close spatial correspondence between cold structures and density enhancements, and the warm/hot volume-filling background, illustrates the emergence of a multiphase medium and demonstrates that the implemented cooling/heating routines couple self-consistently to the turbulent flow.
• Figure 4: Thermodynamic phase structure of the turbulent two-phase ISM. Joint probability distribution of gas temperature and overdensity from the driven, cooling–heating hydrodynamic simulation shown in the accompanying ISM slice (Figure \ref{['fig:ism']}) showcasing the bi-phase structure. Colors indicate the logarithm of the mass-weighted probability density. The dashed curve shows the thermal equilibrium heating–cooling balance for the adopted microphysics. At low temperatures, the gas is roughly at thermal equilibrum and therefore follows the heating/cooling curve. At high temperature, the dense ridge follows an approximately isobaric locus, reflecting turbulent mixing and shock heating rather than local thermal equilibrium. Cold and warm gas accumulate near the stable branches of the equilibrium curve, while hot gas occupies a broad, off-equilibrium tail produced by intermittent turbulent dissipation and long cooling times. Marginalized distributions in temperature (top) and overdensity (right) highlight the bimodal phase structure of the medium.
• Figure 5: Volume rendering of forward model of supernova remnant. Three-dimensional rendering of the same radiatively–cooling supernova remnant shown in Figure \ref{['fig:snr']}, coloured by internal energy. The remnant forms a roughly spherical hot bubble whose surface is strongly corrugated by expansion into a turbulent, inhomogeneous ambient medium. Bright filaments and knots trace dense, rapidly–cooling material in the swept-up shell and in mixed cloud fragments, while the diffuse interior shows a web of turbulent structures generated by shock–cloud interactions and deceleration-driven instabilities. The rendering emphasizes the multi-phase character of the remnant: a hot, volume-filling cavity bounded by a rippled, cooling shell and punctuated by small-scale clumps produced by nonlinear RT/KH growth.