来自伯克利宇宙物理中心、CosmoStat 实验室和劳伦斯伯克利国家实验室的 Chirag Modi、François Lanusse、Mustafa Mustafa、Uroŝ Seljak 的客座文章

引言

数值 模拟1 宇宙大尺度结构是宇宙学家用来理解宇宙学调查收集的大量数据的基本工具。这些模拟通常非常计算密集，通常在大型超级计算机上离线运行。但是，如果我们可以使这些模拟非常快，并将它们与单个统一框架中的机器学习组件集成在一起呢？这就是一个新的 N 体宇宙学模拟代码 FlowPM2 的目标。在这篇博文中，我们将向您展示如何在 TensorFlow 中模拟自己的微型宇宙，并解释为什么这对宇宙学家来说是一个令人兴奋的前景。

图 1：(蓝色) 在 2dFGRS 调查中观察到的宇宙结构。(红色) 在 Millenium N 体模拟中生成的对应结构

N 体宇宙学模拟

宇宙学中的 N 体模拟从宇宙在大爆炸诞生时开始演化到今天。它们通过根据物质的初始分布在一个盒子里分布大量粒子来实现，然后根据重力作用在时间内移动这些粒子。由此产生的最终粒子分布再现了我们宇宙的大尺度结构。例如，它反映了星系和星系团在今天的分布方式。

这些结构的形成和演化取决于我们宇宙的一些基本问题，例如宇宙中有多少物质？宇宙膨胀的速度有多快？驱动这种膨胀的暗物质和暗能量的性质是什么？

现代宇宙学调查使用最强大的望远镜绘制了这些结构，覆盖了最大的天空区域，并追溯到数十亿年前。将 N 体模拟的预测与这些调查的观测结果相匹配，有助于我们回答上述基本问题，从而提高我们对宇宙诞生和演化的理解。

FlowPM，一个 TensorFlow 宇宙学 N 体求解器

N 体模拟几十年来一直是宇宙学界的核心工具。为了从机器学习和统计推断领域的最新进展中获益，我们引入了 FlowPM，一个纯 TensorFlow 实现的宇宙学 N 体模拟。我们提供了一个 Google Colab 笔记本 来试验这些模拟。注意：我们的代码目前使用 TF1 编写（我们将在未来探索更新到 TF2）。

在设置模拟参数（例如箱体长度和网格大小）之后，执行模拟的代码段非常简单，并包含在下面。它生成观察到的如第 2 图所示的大尺度结构。

# Generate Gaussian initial conditions for the matter distribution 
initial_conditions = flowpm.linear_field(N,  L, ipklin, batch_size=batch)   
# Sample particles, i.e. generate the initial displacement and velocity
state = flowpm.lpt_init(initial_conditions, a0=a0)   
# Evolve particles from initial state down to present time with N-Body simulations
final_state = flowpm.nbody(state, stages, N)         
# Visualize final density field i.e interpolate the particles to a grid
Final_field = flowpm.cic_paint(tf.zeros_like(initial_conditions), final_state[0])
#Execute the graph!
with tf.Session() as sess:
    ic, istate, fstate, sim = sess.run([initial_conditions, state, final_state, final_field])

图 2：(左) 宇宙中物质的初始分布，在 N 体模拟开始时。(右) 物质的最终分布，在模拟的最终快照中。可以清楚地看到大尺度结构，包括坍缩的晕、长丝和空洞。

那么，在 TensorFlow 中拥有这些模拟有什么好处呢？这种框架的优势主要可以分为两大类
1. 分析和推断：TensorFlow 中的模拟具有宇宙学家以前没有的独特能力 - 可微性。这种能力为科学家打开了新的分析工具的大门，例如开发有效的基于模拟的推断技术。它还允许我们量化最终观测结果相对于各种输入参数的响应。在我们伯克利宇宙物理中心 (BCCP，加州大学伯克利分校) 的小组中，我们有兴趣追溯时间并重建我们宇宙的 初始条件4 从今天的大尺度结构的观测结果来看。这涉及解决数百万维度的非线性优化问题，这只有使用 FlowPM 等可微模拟才可行。下面 GIF 1 显示了这种重建的示例。

GIF 1：gif 显示从最终暗物质宇宙场（右）中的大尺度结构重建宇宙的初始条件（左）。

2. 混合物理/深度学习模拟：它允许人们开发混合正向模型，我们在其中使用深度学习组件作为 N 体模拟的一部分。凭借当前的计算资源，不可能以非常高的精度同时模拟所有组件。例如，根据他们的科学目标，当前的模拟经常在不同元素之间进行权衡：例如宇宙中观察到的广泛长度尺度、观察到的星系质量范围以及形成这些星系所涉及的各种物理过程。

但是现在，根据科学要求，我们可以使用深度学习代理模型来改进这种权衡，以一种与底层 N 体模拟自然交互的方式包含这些元素。图 3 显示了我们在 BCCP 小组中开发的这种混合模拟 3 的示例。假设我们想模拟宇宙中的气体动量密度。这些可观测值目前使用昂贵的流体动力学模拟进行模拟，但现在可以使用 FlowPM 端到端以 1000 倍的成本生成。

图 3：混合模拟的示例，其中我们用 2 层非线性变换（网络）补充 PM 模拟（左）的暗物质输出，以模拟宇宙中的气体动量密度（中）。将其与通过 1000 倍昂贵的流体动力学模拟模拟的真实情况进行比较（右）

分布式模拟 使用 Mesh TensorFlow5,6

所有这些在原则上可能看起来很简单，但还有另一个挑战；我们的宇宙非常大！为了准确匹配当前和未来调查的观测结果，我们需要同时演化数十亿个粒子。这使得 N 体宇宙学模拟在两个方面都具有挑战性

1. 连续演化数十亿个粒子在计算上非常昂贵。 此外，为了估计所有粒子的重力，我们将需要计算模拟中的所有粒子对。这按 N2 缩放，这使得这种计算变得过高。幸运的是，有一些近似方案可以使它变得易于处理。我们在这里在 FlowPM 中采用的一个方案是粒子网格 (PM) 方法。在 PM 方法中，为了估计重力，我们将在大小为 Ng 的规则网格上离散化空间，然后使用高度优化的 3D 快速傅里叶变换在整个空间上计算力。这将计算成本从 N2g 降低到 Nglog(Ng)。

2. 尽管有这些算法优化，数十亿个粒子也使得这些模拟非常占内存。 结果，需要至少 1024x1024x1024 网格的有用大小的模拟不适合单个 GPU。因此，我们需要一个模型并行框架来开发大规模模拟，这就是 Mesh TensorFlow 的作用。

Mesh TensorFlow 框架允许我们轻松地根据分布式张量描述我们的模拟，并在幕后跟踪分布式梯度和设备之间内存通信。通过在 Mesh TensorFlow 中编写我们的 N 体求解器，我们可以将这些庞大的模拟体积分布在超级计算机上的多个设备上。在这种模拟中，每个进程和网格组件在每个时间步长演化不同的空间区域。使用相同的模拟代码，我们可以同时演化 128 个独立的 128x128x128 网格的宇宙，或在国家计算设施（如 NERSC）上的 64 个 GPU 上演化 1024x1024x1024 网格大小的宇宙。除了能够进行大型模拟之外，模型并行框架还允许我们通过将计算分散到多个处理器上来加快中型模拟的速度。这在下面的图 4 中得到了证明，我们展示了平均而言，FlowPM 模拟比当前可微分的 Python 模拟 FastPM 快 40 倍。

图 4：我们比较了 2563 网格 PM 模拟中 1 步的时间缩放与处理器数量的关系，FastPM（基于 CPU 的 Python 代码，在 Cori Haswell 内核上运行）和 FlowPM（基于 GPU 的 Mesh TensorFlow 代码，在 Cori GPU 上运行）

展望

超过三十年来，对我们宇宙的数值模拟一直是大尺度结构宇宙学的基础。通过 FlowPM，我们正在迈出第一步，将这些模拟与深度学习组件整合到一个统一的框架中，同时保持对底层现象的精确物理理解。在宇宙学中，这种结合为开发新颖的分析工具打开了大门，并将建模推向了以前难以企及的领域。这些是活跃的研究领域，随着下一代宇宙学调查的到来，这一趋势变得更加紧迫，这些调查将在十年之交开始观测宇宙中数千万个天体。物理建模和机器学习的这种融合在很大程度上得益于 Mesh TensorFlow 的模型并行框架，我们希望 FlowPM 开发的组件分析和计算工具也将惠及宇宙学以外其他学科的大规模科学应用。

我们衷心感谢 NERSC 的同事 Wahid Bhimji、Steve Farrell、Peter Harrington、Prabhat 以及 Google 的同事 Niki Parmar、Thiru Palanisamy、Noam Shazeer、Youlong Cheng、Zak Stone 以及其他为我们指明相关资源、积极讨论优化和改进这些模拟方法以及提供有用反馈的人。

参考资料

1. FastPM (FlowPM 的底层 PM 方案) - https://arxiv.org/abs/1603.00476
   
2. FlowPM 代码在 TF 中: https://github.com/modichirag/flowpm
   
3. Dai 等人: https://drive.google.com/open?id=0B7_TnnOHCrvBcWxHR2tVUkR2N0xDbHo3TUxyN2hZemtZSUJn
   
4. 使用神经网络重建初始条件 - https://arxiv.org/abs/1805.02247
   
5. Mesh Tensorflow: https://github.com/tensorflow/mesh
   https://arxiv.org/abs/1811.02084
   
6. 使用 MeshTF 的并行 FlowPM 代码: https://github.com/modichirag/flowpm/tree/mesh