Differentiable Cosmological Hydrodynamics for Field-Level Inference and High Dimensional Parameter Constraints

TL;DR

diffHydro presents a fully differentiable cosmological hydrodynamics pipeline that integrates a TVD Euler solver with a PM gravity solver and differentiable subgrid processes via Gumbel-Softmax, enabling gradient-based and Bayesian inference in high-dimensional parameter spaces. The approach is implemented in JAX, supports GPU acceleration, and is validated through Sedov blast tests and synthetic field-level inferences. Key contributions include differentiating through stochastic discrete events, enabling joint cosmology–subgrid parameter inference from baryon statistics and reconstructing initial conditions from noisy data. As a proof-of-principle, the work demonstrates the feasibility of end-to-end differentiable forward modeling for large-scale structure, with clear paths toward scalability and richer physics on future hardware.

Abstract

Hydrodynamical simulations are the most accurate way to model structure formation in the universe, but they often involve a large number of astrophysical parameters modeling subgrid physics, in addition to cosmological parameters. This results in a high-dimensional space that is difficult to jointly constrain using traditional statistical methods due to prohibitive computational costs. To address this, we present a fully differentiable approach for cosmological hydrodynamical simulations and a proof-of-concept implementation, diffhydro. By back-propagating through an upwind finite volume scheme for solving the Euler Equations jointly with a dark matter particle-mesh method for Poisson equation, we are able to efficiently evaluate derivatives of the output baryonic fields with respect to input density and model parameters. Importantly, we demonstrate how to differentiate through stochastically sampled discrete random variables, which frequently appear in subgrid models. We use this framework to rapidly sample sub-grid physics and cosmological parameters as well as perform field level inference of initial conditions using high dimensional optimization techniques. Our code is implemented in JAX (python), allowing easy code development and GPU acceleration.

Figures (9)

• Figure 1: Central slide of a Sedov-Taylor blast-wave test conducted in a box with $128^3$ cells after 30 time steps. Initial energy to density ratio tuned for the analytical radius to be $r_{sh}(t_f)=20.0.$
• Figure 3: Example simulation in a $256^3$ resolution box evolved to $z=2.5$. Shown is the dark matter density field, baryon density field, temperature field, and the tracked star particles.
• Figure 4: Convergence test of diffHydro (without sub-grid physics) for evolved baryon density field for an 8 $h^{-1}$ Mpc box at $z=2.5$, changing box size while keeping the same initial conditions. We find excellent agreement even at relatively coarse resolutions up to $k\sim 7.0$$h$ Mpc$^{-1}$.
• Figure 5: Demonstration on how baryon density, internal energy, and stellar field can change with varying the subgrid physics parameters $\delta_c \in [0.5,1.5]$ and $E_0 \in [0.01,0.05]$ with fixed initial conditions at $z=2.5$. Slices are projected over a 100 $h^{-1}$ kpc thick slab.
• Figure 6: Comparison of our automatically calculated power spectra derivatives vs. those found via finite differences at $z=2.97$. We perform our comparison at a fixed realization in order to show the accuracy of our derivative calculation. While methods are susceptible to numerical noise, downstream likelihood calculations aren't significantly effected. Note that at these cosmological values, $dP_k/d\Omega_m$ and $dP_k/d\sigma_8$ have very similar shapes, but their ratio does have a k-dependence.