论文标题
FlowPM:FASTPM宇宙学N体求解器的分布式张量集实现
FlowPM: Distributed TensorFlow Implementation of the FastPM Cosmological N-body Solver
论文作者
论文摘要
我们提出FlowPM,这是一种在网格浓度流中实现的粒子网(PM)宇宙n体代码,用于GPU加速,分布式和可区分的模拟。我们基于多分辨率金字塔来实施和验证新型多网络方案的准确性,以在分布式平台上有效地计算大规模的力量。我们探讨了大规模超级计算机上模拟的缩放,并将其与相应的基于Python的PM代码进行了比较,从壁炉时间的平均速度提高了10倍。我们还展示了该新颖的工具如何用于有效地解决大规模宇宙论问题,特别是在具有混合PM和神经网络前进模型的远期模型贝叶斯框架中重建宇宙学领域。我们为这些示例提供骨架代码,整个代码可在https://github.com/modichirag/flowpm上公开获得。
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.