Differentiable Cosmological Simulation with Adjoint Method

Abstract

Rapid advances in deep learning have brought not only myriad powerful neural networks, but also breakthroughs that benefit established scientific research. In particular, automatic differentiation (AD) tools and computational accelerators like GPUs have facilitated forward modeling of the Universe with differentiable simulations. Based on analytic or automatic backpropagation, current differentiable cosmological simulations are limited by memory, and thus are subject to a trade-off between time and space/mass resolution, usually sacrificing both. We present a new approach free of such constraints, using the adjoint method and reverse time integration. It enables larger and more accurate forward modeling at the field level, and will improve gradient based optimization and inference. We implement it in an open-source particle-mesh (PM) $N$-body library pmwd (particle-mesh with derivatives). Based on the powerful AD system JAX, pmwd is fully differentiable, and is highly performant on GPUs.

arXiv comment5 figures + 2 tables; repo at https://github.com/eelregit/pmwd ; v2 matches published version with better typesetting

References in corpus (20)

• SciPy 1.0--Fundamental Algorithms for Scientific Computing in Python
• Array Programming with NumPy
• The Cosmic Linear Anisotropy Solving System (CLASS) II: Approximation schemes
• Baryonic Features in the Matter Transfer Function
• The frontier of simulation-based inference
• Bayesian physical reconstruction of initial conditions from large scale structure surveys
• HACC: Simulating Sky Surveys on State-of-the-Art Supercomputing Architectures
• Fast likelihood-free cosmology with neural density estimators and active learning
• FastPM: a new scheme for fast simulations of dark matter and halos
• Large-scale dark matter simulations
• ELUCID - Exploring the Local Universe with reConstructed Initial Density field I: Hamiltonian Markov Chain Monte Carlo Method with Particle Mesh Dynamics
• Physical Bayesian modelling of the non-linear matter distribution: new insights into the Nearby Universe
• The $\text{Abacus}$ Cosmological $N$-body Code
• Towards optimal extraction of cosmological information from nonlinear data
• Accurate initial conditions for cosmological N-body simulations: Minimizing truncation and discreteness errors
• FlowPM: Distributed TensorFlow Implementation of the FastPM Cosmological N-body Solver
• Learning effective physical laws for generating cosmological hydrodynamics with Lagrangian Deep Learning
• pmwd: A Differentiable Cosmological Particle-Mesh $N$-body Library
• Hybrid Physical-Neural ODEs for Fast N-body Simulations
• Emulating cosmological growth functions with B-Splines

Cited by in corpus (10)

• Bayesian Inference of Initial Conditions from Non-Linear Cosmic Structures using Field-Level Emulators
• Differentiable and accelerated spherical harmonic and Wigner transforms
• BullFrog: Multi-step perturbation theory as a time integrator for cosmological simulations
• A Differentiable, End-to-End Forward Model for 21 cm Cosmology: Estimating the Foreground, Instrument, and Signal Joint Posterior
• Differentiable Cosmological Hydrodynamics for Field-Level Inference and High Dimensional Parameter Constraints
• Learning the Universe: Learning to Optimize Cosmic Initial Conditions with Non-Differentiable Structure Formation Models
• Introducing cosmosTNG: simulating galaxy formation with constrained realizations of the COSMOS field
• Analytic auto-differentiable $\Lambda$CDM cosmography
• A Hamiltonian, post-Born, three-dimensional, on-the-fly ray tracing algorithm for gravitational lensing
• Differentiable Fuzzy Cosmic-Web for Field Level Inference