FlowPM: Distributed TensorFlow implementation of the FastPM cosmological N-body solver

• Published in: Astronomy and computing Vol. 37; no. C; p. 100505
• Main Authors: Modi, C., Lanusse, F., Seljak, U.
• Format: Journal Article
• Language: English
• Published: Netherlands Elsevier B.V 01.10.2021; Elsevier
• Subjects: Cosmology; Large-scale structure of universe, methods; N-body simulations; Cosmology; Large-scale structure of universe, methods; N-body simulations
• ISSN: 2213-1337; 2213-1345; 2213-1345
• DOI: 10.1016/j.ascom.2021.100505

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▪. •End-to-end differentiable cosmological N-Body simulations.•GPU based simulations with 10x speed gain over current CPU simulations.•First N-body simulation written in TensorFlow interfacing with ML and DL components.•Novel multi-grid force algorithm for distributed computing of large scale forces.•Support for large-scale distribution on supercomputers with Mesh TensorFlow.