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

TL;DR

This paper introduces a fully differentiable cosmological hydrodynamics simulation framework, enabling efficient high-dimensional parameter inference and joint constraint of cosmological and astrophysical parameters using gradient-based methods.

Contribution

The authors develop diffhydro, a differentiable simulation tool that allows backpropagation through hydrodynamical and subgrid models, including stochastic variables, for the first time.

Findings

Enables rapid sampling of parameters using gradients.

Facilitates field-level inference of initial conditions.

Demonstrates efficient joint constraint of cosmological and astrophysical parameters.

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.