CSPTH：基于天河二号的晶体结构预测软件框架

CSPTH: A Crystal Structure Prediction Framework on Tianhe-2 Supercomputer

摘要: 晶体结构是深入理解材料的物理及化学性质的重要信息，发展可以从理论上预测晶体结构的方法具有重要意义。通过高性能计算集群甚至利用超级计算机来加速晶体结构预测已逐步成为趋势。本文中，我们基于天河二号超级计算机开发了一套开源的晶体结构预测软件框架，命名为CSPTH。在算法层面，我们基于当前效率最高的遗传算法进行晶体结构预测，并采用了多种技术提升结构预测效率，包括并行化生成种群结构；引入空间群限定，减少自由度搜索，提升结构多样性；引入晶体指纹进行相似性算法，排除相似结构干扰，避免“基因漂变”的问题。特别地，我们针对晶体结构预测算法的应用特点以及天河二号的系统环境，从任务以及数据管理两个方面做了优化。在任务管理上，我们设计了多层任务调度管理模块，根据计算任务的规模大小的分发细粒度作业(节点内)以及粗粒度作业(跨节点)，提升计算资源的高效使用；在数据管理上，我们将每个计算任务的数据都临时储存于计算节点的RAMDISK，提取有效信息后再存储于MongoDB数据库，避免大量小文件存储于公共存储。CSPTH已在15种已知一元、二元以及三元体系上进行了结构预测，实验结果表明CSPTH能根据给定的组分及外部压力条件下全部预测出相应的稳定结构。

Abstract: Crystal structure is the critical information for understanding the physical and chemical properties of materials. Therefore, theoretical prediction of crystal structures only with chemical composition and external conditions is significant. It has become a trend to design new materials through high-performance computing clusters or even using supercomputers. In this paper, we developed an open source framework for crystal structure prediction (CSP) based on Tianhe-2 supercomputer, named CSPTH. When designing the algorithm in CSP, we chose the most efficient genetic algorithm in our framework and adopted numerous technologies to improve prediction accuracy. Specifically, we used a multi-process parallel method to generate the trying structures. The space group restriction is introduced to reduce the searching space and improve the structural diversity in population. We utilized a crystal fingerprint to eliminate the similar structures, which can avoid the problem of “gene drift”. In particular, considering the characteristics of crystal structure algorithm and the system environment of Tianhe-2, we optimize the crystal structure prediction algorithm from two aspects: task management and data management. In terms of task management, we designed a multi-layer task scheduling management module to distribute fine-grained tasks (within a node) and coarse-grained tasks (multi-nodes) according to the system size of tasks to improve the efficiency of employing resources. In our data management module, the data of each computing task is temporarily cached in the RAMDISK of the computing node, and the useful information is extracted and later stored in the MongoDB database, which can avoid a large number of small files stored in the public storage. CSPTH has been used to predict the structures of 15 known element, binary and ternary systems. Experimental results show that CSPTH can predict all the correspond-ing stable structures with the only known of chemical composition and external pressure.

1. 引言

在过去的几十年中，利用计算机进行晶体结构预测的方法，使得系统地进行新材料设计已成为了可能 [1] [2] [3] [4] [5]。结构预测的本质就是在只给定化学组分和外界条件下，确定全局能量最低结构的原子排列方式。具体地说就是在势能面上寻找全局能量最低点。尽管原理上很简单，然而从理论上确定物质结构，是物理、化学和材料研究领域的长期难题。其根本原因在于物质势能面的高度复杂性，随着单胞内原子数目的增加，体系可能的结构数目成指数增加；并且精确地评估结构的能量也需要大量的计算，理论结构预测就是在如此复杂的势能面中搜索全局能量最低的结构，具有巨大挑战性。在材料的结构搜索领域，已报道的具代表性的方法有模拟退火法 [6] [7]、能谷跳跃 [8]、极小值跳跃 [9]、随机法 [10] [11]、遗传算法 [12] [13] [14] [15]、元动力学法 [16] 和粒子群优化算法 [17] [18]。表1列举了当前主流的结构搜索计算软件，其中USPEX [14] [19] [20] 和CALYPSO [21] [22] [23] 是目前报道的用户以及论文引用最多的软件。

Table 1. Summary of structure search software

表1. 结构搜索软件情况

相关晶体结构预测的算法是典型的计算密集性和数据密集型应用。首先，基于进化算法的晶体结构方法需要使用准确的第一性原理的方法去计算种群中的每个结构的能量，而此第一性原理方法非常耗时，使得晶体结构预测的效率比较低；其次，相关算法都需要产生大量的种群结构数，而使用第一性原理的方法评价每一个结构将产生批量的作业，并且需要迭代进行多轮计算，需要消耗大量的计算资源；最后，由于调用的是第三方程序对种群中的每个结构进行优化，在不更改第三方程序源码的情况下，难以对每个结构的输出进行统一管理，将会产生大量的临时文件数据。近年来，随着高性能计算领域的快速发展以及材料设计领域需求的提升，通过高性能计算集群甚至利用超级计算机来加速晶体结构预测已逐步成为趋势。然而相关的软件并没有针对高性能集群进行并行或者高通量管理的设计，不能充分地发挥高性能集群的性能。基于此，我们基于天河二号超级计算机开发了一套开源的晶体结构预测软件框架，命名为CSPTH。在任务管理上，我们设计了多层任务调度管理模块，根据计算任务的规模大小的调整细粒度作业(节点内)以及粗粒度作业(跨节点)，提升资源的利用率；在数据管理上，我们将每个计算任务的数据都临时储存于计算节点的RAMDISK，提取有效信息后再存储于MongoDB数据库，避免大量小文件存储于公共存储。在算法层面，我们基于当前效率最高的遗传算法进行晶体结构预测，并采用了多种技术提升结构预测效率，包括并行化生成种群结构；引入空间群限定，减少自由度搜索，提升结构多样性；引入晶体指纹进行相似性算法，排除相似结构干扰，避免“基因漂变”(genetric drift) [35] 的问题。目前这一软件框架已在多种已知体系上进行了实验，实验结果表明CSPTH能根据给定的组分及外部压力准确找到相应的稳定结构。

2. 晶体结构预测算法

2.1. 遗传算法

我们基于遗传算法设计了晶体结构预测算法。如图1所示，该算法主要包含四个步骤：1) 初始化种群结构：基于对称性限定随机生成结构；2) 结构优化：通过第一性原理计算对结构进行局域优化；3) 生成新结构：通过结构交叉、晶格突变、原子交换这三种变异操作及随机方法生成候选结构。图2为三种变异操作的示意图；4) 挑选下一代种群结构：基于晶体结构的能量构造适应度函数，挑选低能量结构进入下一代种群。重复寻优过程，直到找到稳定的结构。

Figure 1. Flowchart of genetic algorithm

图1. 遗传算法流程图

Figure 2. Diagram of variation operations

图2. 变异操作示意图

2.2. 基于对称性限定的结构生成优化

晶体结构的对称性由230种空间群决定。根据空间群的不同，可以将结构划分为7大晶系，不同空间群的晶体结构差异较大。然而，纯随机生成的晶体结构通常不具备对称性或仅有较低的对称性(空间群为P1)，这可能导致整个遗传算法对于化学空间的探索局限在空间群为P1的结构集合内。由于自然界中低能量结构通常具有一定的对称性，纯随机生成的空间群为P1的结构几乎都是无效结构。为了避免生成大量无效的低对称性结构，我们采用基于对称性限定的随机方法生成结构。首先，在230个空间群中随机选择一个。为了尽可能地探索化学空间，每一次挑选空间群时，都尽可能采用未使用过的空间群。随后，依据空间群构造晶格及原子坐标：根据空间群对应的布拉菲晶格确定部分晶格常数，其余晶格常数随机生成；通过随机组合空间群对应的维科夫(Wyckoff)位置生成原子坐标。

2.3. 结构生成并行化优化

在遗传算法中，种群中结构的数量影响着算法整体对化学空间的探索程度，使用较大的种群数可以快速采样化学空间，有助于算法尽快收敛到全局最优解。在现有的晶体结构预测软件中，结构生成是串行的。由于结构生成本身需要大量的尝试才能生成有效结构，尤其是在加入了多种条件限制后，在种群数较大的情况下，结构生成会成为整个算法的效率瓶颈。因此，我们针对结构生成(包含基于对称性限定的随机结构生成、结构交叉、晶格突变、原子交换)部分进行并行化优化。针对每一个结构生成操作提交相应的任务，通过HTCondor分发到天河节点上。为了充分利用天河节点24核的计算资源，提高计算效率，采用单节点多任务模式，在一个节点上并行多个结构生成任务。

2.4. 基于晶体指纹排除相似结构

相似的结构经过局域结构优化后，通常会落入势能面的同一个能谷中，即收敛为同一个低能量结构。因此，种群中的相似结构对于寻找全局最低能量的结构并没有帮助，相反地，会造成计算资源的浪费。我们采用局部径向分布函数 [36] 计算晶体指纹，通过比较晶体指纹的相似度消除种群中的相似结构。

局部径向分布函数 $g_{\alpha \beta }\left(r\right)$ 表示与类型为 $\alpha$ 的原子相距r的薄壳区域(薄壳厚度为dr)内存在 $\beta$ 类原子的概率，其计算公式如下：

$g_{\alpha \beta }\left(r\right)=\frac{N_{\beta }}{V}d{n}_{\alpha \beta }\left(r\right)V_{shell}=\frac{N_{\beta }}{V}d{n}_{\alpha \beta }\left(r\right)4\pi r^{2}dr\left(V_{shell}=\frac{4}{3}\pi {\left(r+dr\right)}^3-\frac{4}{3}\pi r^3\simeq 4\pi r^{2}dr\right)$

其中， $N_{\beta }$ 为 $\beta$ 类原子的数量，V、 $V_{shell}$ 分别为整个结构的体积、薄壳区域的体积， $d{n}_{\alpha \beta }\left(r\right)$ 表示与类型为 $\alpha$ 的原子相距r的薄壳区域(薄壳厚度为dr)内存在 $\beta$ 类原子的数量。

基于局部径向分布函数，可以通过下式计算得到每一个结构的晶体指纹fp：

$fp\left(i,j,r\right)=g_{a_i,a_j}\left(r\right),i=1,\cdots ,S;j=1,\cdots ,S;r=1,2,\cdots$

其中，S为原子种类数。

通过下式计算结构i和结构j的晶体指纹之间的余弦距离，如果余弦距离小于设定的阈值，则认为两个结构相似，消除两者中能量较高的结构。

$d_{ij}=\frac{1}{2}\left(1-\frac{f{p}_{i}f{p}_j}{\Vert f{p}_i\Vert \Vert f{p}_j\Vert }\right)$

3. CSPTH框架介绍

晶体结构预测算法是典型的计算密集性和数据密集型应用。晶体结构预测算法依赖准确的第一性原理的方法提供每个结构的能量信息，而第一性原理计算非常耗时。如果缺乏针对高性能集群设计的并行管理，多次批量地进行第一性原理方法计算会降低晶体结构预测算法的效率，也不能充分地发挥高性能集群的性能。在调用基于第一性原理方法的第三方程序进行结构优化时，会产生大量的临时文件数据，数据存储、数据移动以及数据分析将产生瓶颈，造成共享文件系统负担过重。基于此，我们设计了针对天河二号集群的晶体结构预测框架CSPTH，为晶体结构预测算法提供任务管理与数据管理。

CSPTH框架如图3所示，包含晶体结构预测算法、局域结构优化、任务管理以及数据管理四个子模块。下面将对局域结构优化、任务管理以及数据管理三个子模块进行介绍。

Figure 3. Framework of CSPTH

图3. CSPTH框架图

3.1. 局域结构优化

CSPTH目前支持第一性原理软件包或者基于力场的分子动力学程序作为接口来实现晶体结构的几何优化。这些软件通过调整晶体结构中各个原子的位置，可以得到一系列的结构，这些结构所具有的能量组成了一个势能面，几何优化计算的目标就是找到能量的最小值点对应的晶体结构。目前的结构优化

技术常用算法包括共轭梯度算法、最速下降、线性最小化和拟牛顿法等。这些极小值点将会为产生下一代结构提供更物理的结构信息。因此局域优化对于结构预测的成功具有非常重要的作用。其中第一性原理方法是进行结构优化最准确的方法，但同时这一方法非常耗时。

3.2. 任务管理

CSPTH和同类型的进化算法软件类似，通过调用第三方的软件对当前种群的每个结构进行结构评价，这些程序中既有并行效率高的程序(如LAMMPS [37]、VASP [38] )，也有只能在单节点内多线程并行的程序(如Gaussian [39] )，并且具体任务所用资源的大小也和计算体系的规模有关，因此需要支持灵活多样的任务分发与管理。然后大多数传统的大规模高性能集群支持的作业队列管理系统如SLURM [40]，PBS [41] 和SGE [42] 等，都配置为粗粒度的计算资源。当作业数扩展到一千以上时，容易出现作业失败的问题。这一问题主要来自两个方面，一个是输入文件的参数问题，另一方面是计算集群的节点故障(如：节点内存溢出、输入输出阻塞、网络延迟等)，如果手动处理这些作业，将会影响预测晶体结构的效率。

因此，在任务管理模块中，我们针对天河二号系统的环境，设计了多级任务调度和作业容错两个子模块。在天河二号集群上，采用的是SLURM作业管理系统，每个节点配置24个CPU物理核。在多级任务调度子模块中，我们通过一级的SLURM作业管理系统申请一定数量的资源，然后通过二级的HTCondor [43] 对资源进一步细化，并进行任务的分发与管理。通过这种方式解决不同任务对于调度粒度的需求，同时保证了对计算资源的统一管理。对于容错，我们主要考虑作业失败后自动处理的能力。对于输入文件参数导致的问题，我们针对应用软件的常见错误，建议一一对应的处理方法；对于系统节点故障问题，我们首先通过作业管理系统排除指定节点，然后将作业重新提交。

3.3. 数据管理

由于调用的是第三方程序进行结构的评估，在不更改源程序的前提下，难以对每个结构的输出进行统一管理，当计算任务扩展到成千上万以后，将会产生大量的临时文件数据，数据存储、数据移动以及数据分析将产生瓶颈。如果直接访问和存储文件将产生大量的输入输出(Input/output, IO)操作，容易导致共享文件系统负担过重。数据库为存储和数据查询提供了一个有吸引力的解决方案。在CSPTH的框架中，我们通过使用MongoDB [44] 和本地RAMDISK来避免使用共享文件系统(如图4所示)。在计算过程中产生的输出文件临时地存储在计算节点的RAMDISK中，然后将有效的信息提取出来直接存储到MongoDB中。MongoDB作为数据存储引擎，它是一个高性能、高可用、高可扩展的、开源的非结构化数据库。这种架构适用于稀疏和类文档的数据存储。通过使用MongoDB“索引”，可以轻松地对分子进行查询和排序。此外，MongoDB使用“分片”(一种在多台机器上分布数据的方法)来支持以高吞吐量方式部署大型数据集，随着数据库的增长平衡查询负载来提高计算性能。最后，MongDB接受高达16 MB的大数据，足以存储常规的输出文件。第三方应用软件涉及处理大量的纯文本文件。在不修改应用程序源代码的情况下，必须通过在共享文件磁盘上移动来处理大量的临时文件。因此，我们充分利用计算节点中的RAMDISK用于临时存储应用程序所需的IO文件(如图4所示)。RAMDISK提供高速、低延迟的IO操作来处理大量的小文件，同时将MongDB部署在共享存储上。通过在MongDB和RAMDISK之间移动数据，有效缓解了共享文件存储的IO压力。

Figure 4. Diagram of dataflow

图4. 数据流示意图

对于种群中结构的评价，需要根据每个结构的自由能进行排序，挑选出适合的结构进行下一代结构的生成。考虑到共享文件存储的高性能计算系统，需要避免大量小文件导致的IO过载问题。因此，分析共享存储磁盘上的输出文件是不明智的。在RAMDISK中完成计算后，将分析输出文件，并根据自由能以及相应的结构存储到数据库中进行排序。当小文件数量急剧增加时，此方法可最大限度地减少IO压力。

4. 实验结果与分析

我们参考了 USPEX [14] 以及CALYPSO [17] 测试过的体系进行CSPTH的性能测试。如表2所示，CSPTH在单质，二元以及三元体系中的共15个一定压强下的结构全部被成功预测出来，成功率为100%。其中单质测试案例包括Li、C、Mg以及Si，CSPTH在第一代就预测出了空间群为P6₃/mmc的单质Li材料，对于单质Mg也只在第二代预测出来。对于双元体系材料，我们对SiO₂、TiH₂以及Al₂O₃三种结构进行预测，对于SiO₂在不同压强下的三种结构均在第一代预测出来；TiH₂的两个已知结构也在第一和第四代中预测出来。CSPTH对于所含原子数较小的单质、二元化合物所在的化学空间表现出了较强的搜索能力，这主要归功于基于对称性限定的结构生成方法。引入对称性的限定，使得遗传算法在探索中能同时关注不同空间群下的多样化的结构，迅速锁定可能的低能量结构所在的势能面的能谷。对于较复杂的三元体系MgSiO₃以及MgAlO₄也分别在第十三代和第十五代预测出来。实验结果表明，CSPTH不仅具有较高的预测准确度，同时也具有较高的预测效率，均在前15代内预测出了结构。同时，我们使用CSPTH预测出了一些新颖的结构，如图5所示。此外，我们还在天河二号系统上进行了压力测试，同时申请了500节点进行所有结构的计算，最大的作业并发数为260，总共产生了43,980个临时文件。我们在三天的时间内完成了所有的计算和数据分析。

Figure 5. Structures predicted by CSPTH

图5. CSPTH成功预测的结构

Table 2. CSPTH tested on known structural systems

表2. CSPTH预测的已知结构体系

详细案例

算例硅的配置文件如图6所示。基本设置为原子类型为硅(Si)，对应的硅原子数量设置为8，外部压力为0 Gpa，种群大小为每代30个结构。作业相关的设置为该算例最多使用45个节点，每个vasp计算使用2个节点。根据配置文件，CSPTH首先会通过SLURM作业管理系统申请45个计算节点。在CSPTH运行过程中，由二级的HTCondor对这45个计算节点进行任务的分发与管理。在遗传算法中，主要会产生两类计算任务：结构生成任务及结构局域优化任务。对于结构生成任务，HTCondor会为其分配4个核的计算资源。对于结构局域优化任务，HTCondor会根据该任务的计算精度为其分配16~64个核的计算资源。结构局域优化的计算精度由其K点网格密度决定，我们设置粗、中等、细这三种网格密度，分别对应16、32、64个核的计算资源。在这两类任务执行过程中会产生大量的临时文件，这些文件都会临时存储在计算节点的RAMDISK。当结构局域优化任务完成时，CSPTH会从这些临时文件中读取优化后的结构信息以及结构对应的能量，并将它们写入MongDB数据库。在每一代种群完成结构生成、结构局域优化、结构挑选后，CSPTH会从MongDB数据库中获取该代所有结构的详细信息，写入到文本文件gen_x_info中(其中x表示种群代数)。图7展示了Si算例中第11代种群的信息汇总文件。文件中记录了三部分内容：1) 种群结构性质：每一个结构的信息，包括能量、空间群、该结构最早产生的代数；2) 最低能量列表：从第1代到第11代种群的最低能量列表；3) 结构来源：每一个结构的生成方式，如果该结构来自于变异操作，会记录父本结构。

Figure 6. Input file of Si case

图6. Si算例中的输入文件

Figure 7. The detail information of the 11th population in Si case is recorded named gen_11_info

图7. Si算例中第11代种群的信息汇总文件gen_11_info

根据种群的信息汇总文件，可以还原出遗传算法的搜寻稳定Si结构的过程，如图8所示。随着种群的演化，每一代中的最优结构的能量逐步下降，意味着CSPTH找到的结构逐渐趋于稳定。在第6代和11代分别对应较低能力与最低能量出现的时刻，这两代很可能首次发现了不同的稳定Si结构。随后，通过比较晶体指纹相似性，确定已在第6代和第11代分别找到已知的稳定结构。

Figure 8. The lowest enthalpy against generation in Si case

图8. Si算例中最低能量随种群代数变化

5. 结语

在本文中，我们介绍了CSPTH框架的实现，该框架可用于在给定化学成分和外部压力的情况下预测晶体结构。CSPTH实现了结构生成、结构评估和结构选择的自动化，使实验科学家也能够进行晶体结构预测。特别地，我们针对天河二号超级计算机的软硬件环境进行优化，包括任务管理以及数据管理，提高软件的效率、鲁棒性以及可扩展性，我们初步测试了最高260个并发任务，调用500个节点进行高通量计算，在三天的时间内完成了9个材料体系共15个算例的结构预测，预测出全部已知实验观察的结构，成功率为100%。CSPTH的高成功率和高性能证明了其作为晶体结构预测工具的可靠性和应用前景。

基金项目

广东省重点领域研发计划(2019B010940001)；广州市科技计划项目(201604016005)。

NOTES

^*通讯作者。