分子动力学借助计算机来模拟分子、原子体系的运动，是模拟微观世界的基本方法，在物理、化学、生物、材料等领域有着广泛的应用。分子动力学将体系中的每一个原子视为遵守牛顿第二定律的基本粒子，在给定原子初始状态和时间步长的情况下，通过对原子进行受力分析(图1红色框内的部分)，便可得到原子的运动轨迹。
分子动力学模拟的核心在于分子的受力计算。传统方法使用经验力场计算分子的受力，在特定体系上，可以模拟百亿原子规模的系统，但是精度受限且难以应对复杂体系的模拟；基于密度泛函理论的第一性原理分子动力学，计算精度高，但是计算量大，模拟尺度通常限于数百个原子，100皮秒以内。深度学习的兴起，为原子的受力分析提供了一种新的选择，DeePMD-kit便是其中的代表，兼具经验立场的计算效率和第一性原理的精度。
使用DeePMD-kit进行受力分析
使用DeePMD-kit进行受力计算的网络结构如图2所示，从前往后分别是：输入构型Coord，Environment matrix，Embedding matrix，以及构建满足对称性的Descriptor，体系总能量E以及反向受力计算。需要强调的是，Environment matrix的生成和受力计算需要以Tensorflow Custom Operator的方式实现，其余部分均使用Tensorflow原生的Operator。下面是对各个部分功能的简单介绍，详细介绍参考DeepModeling Hackathon。
1. 输入构型
以坐标Coord为例，典型的MD(Molecular Dynamics) 输入构型除了原子坐标外，还包含了原子的类型等信息，详情参考表1。
2. Environment matrix
Environment matrix通常依据一定的规则进行构建。对于一个大的系统来说，需要先将系统栅格化成若干个较小的系统(图1中橙色的方块)，在小系统内根据截断半径(图1中的圆的半径)等参数，以中心原子为视点，选取邻居原子，构建系统的Environment matrix。从公式1和公式2可以看出，Environment matrix本身满足平移不变性。
3. Embedding matrix
Embedding matrix是最原始的特征，embedding net(图2)包含可训练参数，可以为不同类型的中心原子与邻居关系构建适合的表示，而这对于神经网络的训练非常重要。
4. 构建Descriptor
这一步主要是给构造的神经网络模型添加旋转和置换不变性，这可以提高模型的适应性，使得在原子发生旋转和置换的情况下，模型依然能对能量进行准确预测，同时在一定程度上减少对训练集的依赖。
5. 体系能量E拟合
这一步通过将Descriptor输入基于MLP构建的神经网络，计算整个系统的能量。在训练中，整个系统的能量通常使用第一性原理等高精度方法计算得到。
6. 受力分析
通过计算总能量E对位置的导数，可以得到每个原子的受力，因此受力分析部分包含一次神经网络的反向，这会导致内存占用较高。同时，Environment matrix对位置的求导使用Tensorflow Custom Operator实现。
使用IPU进行水分子模拟
在IPU上进行水分子模拟时，我们使用SavedModel格式的模型。为了在IPU上运行模型，我们做了以下几件事：
- 在IPU上实现并优化了上文提到的Custom Operator
- 设置模型的device属性为IPU
1. Custom Operator
EnvMat，即上述Enviornment matrix，实际处理时，需要对每个中心原子的领居原子计算相似度S(rij），然后进行排序，选出每个中心原子的topK个邻居原子，再代入公式计算，因此这个Op的瓶颈并不在计算部分，而在排序部分。实际的实现中，通过QuickSort+InsertSort的组合方式，最终实现高效的排序。
Force和Virial Force, 求力和维里力，根据网络的前向结果，对x求偏导。这两个Op涉及到全局内存访问的问题，考虑到IPU的限制，无法将所有原子的数据放在一个tile上，因此我们需要将全局数据拆分成N份，最后再对N份数据进行reduce_sum，最后得到一份数完整的数据。这么做虽然可以让IPU的内存不再受限， 但是，显然限制了IPU算力的发挥，需要寻找更好的解决方案。
2. 设置模型device属性
Tensorflow SavedModel格式的模型，每个操作对应一个Node，通过为不同的Node设置不同的device属性，我们可以设置该操作的执行device，我们通过编辑Node的device属性为IPU，将操作放在IPU上执行。
3. 在IPU上模拟水分子
水分子模型的详细输入输出信息如表1所示，除了Coord以外，还包括中心原子的邻居信息，包括坐标、原子类型等。我们根据firstneigh、typeneigh、posineigh信息构建每个中心原子的Environment matrix。
| Input/Output name | shape | meaning |
| box(i) | (9,) | 体系的三个顶点 |
| coord (i) | (20394,) | 所有原子的三维坐标 |
| default_mesh (i) | (16,) | |
| firstneigh (i) | (1536, 228) | 中心原子的邻居列表 |
| ilist (i) | (1536,) | 中心原子的数量 |
| nall (i) | (6798,) | 体系中原子的数量 |
| natoms_vec (i) | (4,) | |
| nloc (i) | (1536,) | 中心原子的数量 |
| numneigh (i) | (1536,) | 中心原子的数量 |
| type (i) | (6798,) | 所有院子的类型 |
| typeneigh (i) | (1536, 228) | 邻居原子的类型 |
| posineigh (i) | (1536, 684) | 邻居原子的坐标 |
| atom_ener(o) | (1, 1536) | 中心原子的能量 |
| energy (o) | (1,) | 体系的能量 |
| force (o) | (1, 20394) | 原子在三维空间中的受力 |
| virial(o) | (1, 61182) | 原子在三维空间中的维里力 |
在完成Custom Operator和设置device属性后，最终在1536个中心原子的体系上我们可以得到和V100相近的性能。
4. 内存优化
我们对水分子模拟进行了深入分析，因为模型本身存在一次反向，我们需要存储一些前向计算的结果，总大小约为300MB，而这部分内存开销可以通过re-compute来缓解，从而可以运行更大的系统。
总结
目前我们成功在IPU上进行了水分子模拟，在小体系上性能可以基本合V100持平，我们还需要做更多的工作使得IPU可以对大体系进行模拟。