FlowPM: Distributed TensorFlow Implementation of the FastPM Cosmological N-body Solver
Berkeley Center for Cosmological Physics, Department of Physics, University of California, Berkeley, CA 94720, USA
arXiv:2010.11847 [astro-ph.CO], (22 Oct 2020)
@misc{modi2020flowpm,
title={FlowPM: Distributed TensorFlow Implementation of the FastPM Cosmological N-body Solver},
author={Chirag Modi and Francois Lanusse and Uros Seljak},
year={2020},
eprint={2010.11847},
archivePrefix={arXiv},
primaryClass={astro-ph.CO}
}
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.
October 25, 2020 by hgpu
