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

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.