FlowPM: Distributed TensorFlow Implementation of the FastPM Cosmological N-body Solver

Abstract

We present FlowPM, a Particle-Mesh (PM) cosmological N-body code implemented in Mesh-TensorFlow for GPU-accelerated, distributed, and differentiable simulations. We implement and validate the accuracy of a novel multi-grid scheme based on multiresolution pyramids to compute large scale forces efficiently on distributed platforms. We explore the scaling of the simulation on large-scale supercomputers and compare it with corresponding python based PM code, finding on an average 10x speed-up in terms of wallclock time. We also demonstrate how this novel tool can be used for efficiently solving large scale cosmological inference problems, in particular reconstruction of cosmological fields in a forward model Bayesian framework with hybrid PM and neural network forward model. We provide skeleton code for these examples and the entire code is publicly available at https://github.com/modichirag/flowpm.

arXiv comment14 pages, 17 figures. Code provided at https://github.com/modichirag/flowpm

References in corpus (28)

• Adam: A Method for Stochastic Optimization
• Large scale bias and the peak background split
• Cosmological Hydrodynamics with Adaptive Mesh Refinement: a new high resolution code called RAMSES
• The DESI Experiment Part I: Science,Targeting, and Survey Design
• The frontier of simulation-based inference
• LSST Science Book, Version 2.0
• Improving Cosmological Distance Measurements by Reconstruction of the Baryon Acoustic Peak
• Simulating cosmic structure formation with the GADGET-4 code
• Solving Large Scale Structure in Ten Easy Steps with COLA
• Bayesian physical reconstruction of initial conditions from large scale structure surveys
• Nonlinear perturbation theory with halo bias and redshift-space distortions via the Lagrangian picture
• Learning to Predict the Cosmological Structure Formation
• FastPM: a new scheme for fast simulations of dark matter and halos
• Mock galaxy catalogs using the quick particle mesh method
• Modelling Baryon Acoustic Oscillations with Perturbation Theory and Stochastic Halo Biasing
• Massive optimal data compression and density estimation for scalable, likelihood-free inference in cosmology
• High Performance P3M N-body code: CUBEP3M
• Modeling Biased Tracers at the Field Level
• ICE-COLA: Towards fast and accurate synthetic galaxy catalogues optimizing a quasi $N$-body method
• The mass-Peak Patch algorithm for fast generation of deep all-sky dark matter halo catalogues and its N-Body validation
• Towards optimal parallel PM N-body codes: PMFAST
• Comparing approximate methods for mock catalogues and covariance matrices III: Bispectrum
• Super-resolution emulator of cosmological simulations using deep physical models
• Mesh-TensorFlow: Deep Learning for Supercomputers
• TARDIS Paper I: A Constrained Reconstruction Approach to Modeling the z~2.5 Cosmic Web Probed by Lyman-alpha Forest Tomography
• Cosmology Inference from Biased Tracers using the EFT-based Likelihood
• Intensity mapping with neutral hydrogen and the Hidden Valley simulations
• An Adaptive Multigrid Solver for High-Resolution Cosmological Simulations

Cited by in corpus (33)

• Large-scale dark matter simulations
• The DAWES review 10: The impact of deep learning for the analysis of galaxy surveys
• Field Level Neural Network Emulator for Cosmological N-body Simulations
• Second Data Release of the COSMOS Lyman-alpha Mapping and Tomographic Observation: The First 3D Maps of the Detailed Cosmic Web at 2.05<z<2.55
• Machine Learning for Observational Cosmology
• $\mathrm{S}\mathrm{IM}\mathrm{BIG}$: Mock Challenge for a Forward Modeling Approach to Galaxy Clustering
• Towards a Deep Learning Model for Hadronization
• The Power of Locality: Primordial Non-Gaussianity at the Map Level
• Reconstructing the cosmological density and velocity fields from redshifted galaxy distributions using V-net
• Bayesian Inference of Initial Conditions from Non-Linear Cosmic Structures using Field-Level Emulators
• Joint velocity and density reconstruction of the Universe with nonlinear differentiable forward modeling
• DSPS: Differentiable Stellar Population Synthesis
• Differentiable Cosmological Simulation with Adjoint Method
• Reconstructing the Universe with Variational self-Boosted Sampling
• Constrained cosmological simulations of the Local Group using Bayesian hierarchical field-level inference
• Forecasting the power of Higher Order Weak Lensing Statistics with automatically differentiable simulations
• Bayesian Control Variates for optimal covariance estimation with pairs of simulations and surrogates
• Painting baryons onto N-body simulations of galaxy clusters with image-to-image deep learning
• The HalfDome Multi-Survey Cosmological Simulations: N-body Simulations
• Interpretable Uncertainty Quantification in AI for HEP
• BullFrog: Multi-step perturbation theory as a time integrator for cosmological simulations
• Improved Lyman Alpha Tomography using Optimized Reconstruction with Constraints on Absorption (ORCA)
• Massive $\nu$s through the CNN lens: interpreting the field-level neutrino mass information in weak lensing
• Perturbation-theory informed integrators for cosmological simulations
• COmoving Computer Acceleration (COCA): $N$-body simulations in an emulated frame of reference
• 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
• Learning from Topology: Cosmological Parameter Estimation from the Large-scale Structure
• Diagnosing Systematic Effects Using the Inferred Initial Power Spectrum
• Benchmarking field-level cosmological inference from galaxy redshift surveys
• Differentiable Fuzzy Cosmic-Web for Field Level Inference
• Variational autoencoder for generating realistic $N$-body simulations for dark matter halos