Differentiable Cosmological Hydrodynamics for Field-Level Inference and High Dimensional Parameter Constraints
arXiv:2502.02294 · doi:10.1093/mnras/staf1785
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.
14 pages, 10 figures
References in corpus (60)
- E pur si muove: Galiliean-invariant cosmological hydrodynamical simulations on a moving mesh
- Cosmological SPH simulations: A hybrid multi-phase model for star formation
- Simulating Galaxy Formation with the IllustrisTNG Model
- GADGET: A code for collisionless and gasdynamical cosmological simulations
- Simulating galaxy formation with black hole driven thermal and kinetic feedback
- Introducing the Illustris Project: the evolution of galaxy populations across cosmic time
- A model for cosmological simulations of galaxy formation physics
- Athena: A New Code for Astrophysical MHD
- HARM: A Numerical Scheme for General Relativistic Magnetohydrodynamics
- The Athena++ Adaptive Mesh Refinement Framework: Design and Magnetohydrodynamic Solvers
- The Cosmic Linear Anisotropy Solving System (CLASS) I: Overview
- The Arepo public code release
- Cusp-core transformations in dwarf galaxies: observational predictions
- Neural General Circulation Models for Weather and Climate
- The Coyote Universe I: Precision Determination of the Nonlinear Matter Power Spectrum
- Bayesian physical reconstruction of initial conditions from large scale structure surveys
- The CAMELS project: Cosmology and Astrophysics with MachinE Learning Simulations
- Molecular Hydrogen and Global Star Formation Relations in Galaxies
- HACC: Simulating Sky Surveys on State-of-the-Art Supercomputing Architectures
- FastPM: a new scheme for fast simulations of dark matter and halos
- New Constraints on IGM Thermal Evolution from the Lyα Forest Power Spectrum
- The MillenniumTNG Project: The hydrodynamical full physics simulation and a first look at its galaxy clusters
- Solver-in-the-Loop: Learning from Differentiable Physics to Interact with Iterative PDE-Solvers
- Towards optimal extraction of cosmological information from nonlinear data
- TreePM: A code for Cosmological N-Body Simulations
- Cholla : A New Massively-Parallel Hydrodynamics Code For Astrophysical Simulation
- Emulation of baryonic effects on the matter power spectrum and constraints from galaxy cluster data
- Cosmological Reconstruction From Galaxy Light: Neural Network Based Light-Matter Connection
- Cosmological Adaptive Mesh Refinement
- A Primer on Eulerian Computational Fluid Dynamics for Astrophysics
- JAX-COSMO: An End-to-End Differentiable and GPU Accelerated Cosmology Library
- How much information can be extracted from galaxy clustering at the field level?
- Constraining the Cosmic Baryon Distribution with Fast Radio Burst Foreground Mapping
- A Fast Algorithm for Solving the Poisson Equation on a Nested Grid
- TARDIS Paper I: A Constrained Reconstruction Approach to Modeling the z~2.5 Cosmic Web Probed by Lyman-alpha Forest Tomography
- MUSE: Marginal Unbiased Score Expansion and Application to CMB Lensing
- FLIMFLAM DR1: The First Constraints on the Cosmic Baryon Distribution from 8 FRB sightlines
- A gradient based method for modeling baryons and matter in halos of fast simulations
- Predicted future fate of COSMOS galaxy protoclusters over 11 Gyr with constrained simulations
- Osaka Feedback Model II: Modeling Supernova Feedback Based on High-Resolution Simulations
- Towards a self-consistent numerical model of late-type galaxies: Calibrating the effects of sub-grid physics on galactic models
- XLB: A differentiable massively parallel lattice Boltzmann library in Python
- Efficient Optimal Reconstruction of Linear Fields and Band-powers from Cosmological Data
- The Prime Focus Spectrograph Galaxy Evolution Survey
- TARDIS Paper II: Synergistic Density Reconstruction from Lyman-alpha Forest and Spectroscopic Galaxy Surveys with Applications to Protoclusters and the Cosmic Web
- Differentiable Cosmological Simulation with Adjoint Method
- Observational Evidence for Large-Scale Gas Heating in a Galaxy Protocluster at z=2.30
- The THESAN project: Lyman-alpha emitter luminosity function calibration
- Fast, high-fidelity Lyman forests with convolutional neural networks
- pmwd: A Differentiable Cosmological Particle-Mesh -body Library
- HyPhy: Deep Generative Conditional Posterior Mapping of Hydrodynamical Physics
- DISCO-DJ I: a differentiable Einstein-Boltzmann solver for cosmology
- Forecasting the power of Higher Order Weak Lensing Statistics with automatically differentiable simulations
- Constraining the Fluctuating Gunn-Peterson Approximation Using Lyman- Forest Tomography at
- A field-level emulator for modeling baryonic effects across hydrodynamic simulations
- Hybrid Physical-Neural ODEs for Fast N-body Simulations
- Introducing cosmosTNG: simulating galaxy formation with constrained realizations of the COSMOS field
- TensorFlow Hydrodynamics Analysis for Ly- Simulations
- A Gigaparsec-Scale Hydrodynamic Volume Reconstructed with Deep Learning
- Learning an Effective Evolution Equation for Particle-Mesh Simulations Across Cosmologies