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.
Cite
@article{arxiv.2010.11847,
title = {FlowPM: Distributed TensorFlow Implementation of the FastPM Cosmological N-body Solver},
author = {Chirag Modi and Francois Lanusse and Uros Seljak},
journal= {arXiv preprint arXiv:2010.11847},
year = {2020}
}
Comments
14 pages, 17 figures. Code provided at https://github.com/modichirag/flowpm