在TensorFlow里模拟宇宙！提升模拟计算速度，成本降低1000倍

原标题：在TensorFlow里模拟宇宙！提升模拟计算速度，成本降低1000倍

大数据文摘出品

来源：tensorflowblog

编译：睡不着的iris、刘俊寰

去年6月，俄罗斯天文学家兼码农小哥开发了一款名叫Space Engine的宇宙模拟游戏，目前，该游戏已经收到2639份测评，绝大部分都是好评。

不过，Space Engine还只是在游戏层面对宇宙的一种呈现，试想以下，如果用数值模拟的基础工具，从观测的海量宇宙数据进行分析……

停！看到这儿也大概能猜到，这样的模拟计算肯定代价不菲，而且通常需要在大规模的超级算计算机上离线运行。

但是，若是提升模拟计算速度，将其与机器学习组件整合到统一框架内，会有什么不同呢？

这便是新的N体宇宙模拟代码FlowPM在做的。

没错，文摘菌今天给大家展示的就是，如何使用TensorFlow模拟自己的小宇宙，并说明为什么它对宇宙学有着激动人心的前景展望，搓搓小手跟进了~

图1：（蓝色）通过2dFGRS测量的宇宙结构，（红色）使用N体模拟生成的千禧年宇宙结构。

N体宇宙模拟

简单来说，N体模拟的方法是为了解释宇宙演化中的非线性演化问题，科学家们提出的一套研究宇宙学的数值模拟方法。

在宇宙学领域，N体模拟还原了宇宙是如何从诞生之初的大爆炸演变至今的。

如果我们按照物质初始分布，将大量的粒子放置在一个盒子中，由于引力的牵引，粒子会不断移动。而最终的粒子状态，就是对浩瀚宇宙结构的再现。例如，从这个分布中，你可以看到星系和星系群是如何分布的。

这些结构的形成和演化取决于一些与宇宙相关的基础问题，诸如宇宙中有多少物质？宇宙膨胀有多快？导致宇宙膨胀的黑物质和暗能量有哪些特性？

现代的宇宙学研究用强大的望远镜覆盖了辽阔的观测区域，绘制出了这些结构，时间甚至可以追溯到几十亿年前。通过将N体模拟的预测与观测数据进行比较，可以帮助我们解答前述基础问题，从而提高对宇宙的诞生和演化的理解。

FlowPM：一个TensorFlow的宇宙N体求解程序

随着几十年研究的更迭，N体模拟已是宇宙学界的主要研究工具。

得益于机器学习和统计推理领域的最新进展，一个完全基于TensorFlow实现的N体宇宙模拟程序FlowPM应运而生了。

在这里，我们使用Google Colab notebook对其进行实验，同时，以下主要提供TF1的编写代码。

首先需要设置好边长、网格大小和执行代码片段等模拟参数，随后的执行模拟代码片段就非常简单了，完整的代码如下所示，如此就能生成下图所示的大型结构。

# 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真有这么香？文摘菌把话撂这，还这真有这么香，TensorFlow的优势主要体现在以下两个方面。

首先是分析和推理，TensorFlow模拟为宇宙学家提供了前所未有的独特功能，也就是我们所说的可微性（differentiability）。这个功能为科学家提供一种新的分析工具，例如开发基于模拟的有效推理技术。

不仅如此，这个功能还支持量化各类输入参数的调整对最终观察的影响。在伯克利宇宙中心，研究者们对时光倒流和重构初始状态十分感兴趣，这也有助于更好地了解现在的宇宙结构，但这是有前提的，必须先解决一个百万维度的高度非线性优化问题，即模拟过程得可微，而这，FlowPM就能做到，重构后的图如下所示。

从最终的黑物质的大型宇宙结构（右边）如何重建宇宙初始状态（左边）

另一个就是混合物理/深度学习模拟，在N体模拟运用深度学习组件时，你会需要开发一个混合前向模型。但是，利用当前的计算资源，不可能高度准确地同时模拟所有的组件，只能依据科研目标，在不同元素之间进行取舍，例如宇宙可观察的深度、星系可见的范围和形成星系的多维物理过程等。

然而现在，我们可以用深度学习替代模型，将这些元素自然地整合进N体模拟中。在BCCP，科学家们开发了一个这样的混合模拟，如下图所示。假如我们要模拟宇宙中的气体动量密度，目前，这类观察采用混合模拟实现代价非常大，但使用一个端到端的FlowPM模型，成本至少可以降低1000倍。

一个混合模拟的示例，用两层的非线性变换（网络）补充PM模拟的暗物质输出（左），以模拟宇宙（中心）的气体动量密度。将使用昂贵1000倍的混合模拟所模拟与真相进行比较（右）

采用Mesh TensorFlow分布式框架仿真

虽然原理看上去很简单，但实际的挑战还是很大的，毕竟宇宙真的很庞大！

为了准确匹配当前和未来的观测，我们要同时模拟数以亿计的粒子的变化。

要实现这个，N体宇宙模拟主要有以下两个难点：

• 连续演化数十亿个粒子，计算成本非常高。要估算所有粒子之间的重力，需要计算所有的粒子对，当计算规模扩展至N2，无法得出结果。但幸运的是，可以在FlowPM中使用一种叫做particle-mesh（PM）的方法，即把一个Ng规格的网格离散空间，使用高度优化的3D快速傅里叶变换计算整个空间的力，这将计算成本从N2g降到了Nglog(Ng)。
• 尽管算法进行了优化，但模拟数以万计的粒子也非常占用内存。结果是一个有效尺寸为1024x1024x1024的最小的模拟网格，都无法再单个GPU上显示。因此，我们需要一个模型并行框架来支撑大规模模拟，这也正是Mesh TensorFlow的用武之地。

Mesh TensorFlow框架可以轻松地使用分布式张量来表示模拟，并持续追踪分布式梯度和设备之间的存储通信，通过Mesh TensorFlow编写N体求解程序，可以将巨大的仿真量跨多个超级计算机进行分布。

在这种模拟中，每一步操作和网格组件在一个时间步长可以演化出不同的空间区域。在云端TPU部署相同的模拟代码，可以模拟128个128x128x128网格的独立宇宙，或者部署在国家计算机（NERSC）的64个GPU上模拟一个1024x1024x1024的宇宙。

除了支持大规模仿真，模型并行框架允许通过将计算任务分配给多个处理器来加速中间尺寸的模拟，如下图所示，FlowPM的平均模拟速度比当前的可微python模拟器FastPM快了40倍不止。

使用FastPM（在Cori Haswell内核上运行基于CPU的python代码）和FlowPM（在Cori GPU上运行基于GPU的Mesh TensorFlow代码）模拟2563网格PM，比较1个步长内时间缩放与处理器数量

展望

过去的三十多年，数值模拟是大型宇宙的研究基石。现在，有了FlowPM，研究即将迈出新的一步，将模拟器集成在包含深度学习组件的统一框架内，同时保持对潜在现象的准确物理理解。

在宇宙学中，这种结合不仅促生了新颖的分析工具，在研究体系中加入深度模型攻克棘手的历史难题。这些活跃的研究领域，加速催生下一代宇宙研究，在未来的十年之际可以观察宇宙中成千上万的个体。

多亏了Mesh TensorFlow模型并行框架，使得物理建模和机器学习的融合得以实现，希望成分分析和FlowPM计算工具的发展，让宇宙学以外的领域也得以受益。

相关报道：

https://blog.tensorflow.org/2020/03/simulating-universe-in-tensorflow.html