MindSpore新一代分子模拟库：SPONGE

作者：于璠，MindSpore算法科学家

MindSpore新一代分子模拟库：SPONGE，由北大和深圳湾实验室高毅勤课题组与华为MindSpore团队联合开发，具有高性能、模块化等特性，是一个完全自主研发的分子模拟软件库。基于MindSpore自动并行、图算融合等特性，SPONGE可高效地完成传统分子模拟过程，利用MindSpore自动微分的特性，可以将神经网络等AI方法与传统分子模拟进行结合。

背景介绍

分子模拟是指利用计算机以原子水平的分子模型来模拟分子结构与行为，进而模拟分子体系的各种物理、化学性质的方法。它是在实验基础上，通过基本原理，构筑起一套模型和算法，从而计算出合理的分子结构与分子行为。近年来，分子模拟技术发展迅速并且在多个学科领域得到了广泛的应用。在药物设计领域，可用于研究病毒、药物的作用机理等；在生物科学领域，可用于表征蛋白质的多级结构与性质；在材料学领域，可用于研究结构与力学性能、材料的优化设计等；在化学领域，可用于研究表面催化及机理；在石油化工领域，可用于分子筛催化剂结构表征、合成设计、吸附扩散，可构建和表征高分子链以及晶态或非晶态本体聚合物的结构，预测包括共混行为、机械性质、扩散、内聚等重要性质。

由于模拟的时空限制，传统分子动力学仿真软件的应用受到很大的限制，科研工作者需要不断地开发新的力场、抽样方法、结合新的技术(如AI算法)来拓展分子动力学仿真的场景。因此， SPONGE应运而生，具有完全自主的知识产权。SPONGE使用模块化的设计特性，支持科学家进行高效且便捷地搭建分子动力学模拟中所需要的相关计算模块。同时， SPONGE也具有传统模拟所需要的高效性。除此之外，SPONGE也天然地支持与人工智能算法的自然融合，并且能运用MindSPore框架自身的高性能计算特性。

相比于之前在传统分子模拟软件上结合SITS方法进行生物分子增强抽样，SPONGE原生支持SITS并对计算流程进行优化使得其使用SITS方法模拟生物体系更加高效。针对极化体系，传统分子模拟采用结合量化计算等方式来解决电荷浮动等问题。即使采用机器学习降低计算量也会浪费大量时间在程序数据传送的问题上。而SPONGE利用模块化的特点可支持内存上直接与机器学习程序通信大大降低了整体计算时间。

图1：结合SITS等方法对显性溶剂中的丙氨酸二肽进行增强抽样

随MindSpore1.2版本开源的SPONGE具备以下优势：

1. 全模块化分子模拟。模块化构建分子模拟算法，易于领域研发人员进行理论和算法的快速实现，并为外部开发人员贡献子模块提供友好的开源社区环境。

2. 传统分子模拟与MindSpore结合的人工智能算法的全流程实现。在MindSpore中，研发人员能便利的将AI方法作用于分子模拟中。全算子化的SPONGE将与MindSpore进一步结合成为新一代端到端可微的分子模拟软件，实现人工智能与分子模拟的自然融合。

案例介绍

下面，这里将简单介绍一下MindSpore上的SPONGE的一个简单案例，该案例使用SPONGE模拟了丙氨酸三肽水溶液体系。

实践前，确保已经正确安装MindSpore。如果没有，可以通过MindSpore安装页面(MindSpore官网)安装MindSpore。

1. 输入文件准备

本教程模拟体系中需要加载三个输入文件，分别是：

· 属性文件（后缀为.in的文件），声明模拟的基本条件，对整个模拟过程进行参数控制。

· 拓扑文件（后缀为.param7的文件），拓扑文件描述的是体系内部分子的拓扑关系及各种参数。

· 坐标文件（后缀为.rst7的文件），坐标文件描述的是每个原子在体系中的初始时刻的坐标。

拓扑文件和坐标文件可以通过建模过程由AmberTools中自带的tleap工具建模完成，下载地址(Download Amber MD)。

通过tleap构建了所需要的拓扑文件和坐标文件后，需要通过属性文件声明模拟的基本条件，对整个模拟过程进行参数控制。以本教程中的属性文件为例，其文件内容如下：

NVT 290k
 mode = 1, # 分子动力学（MD）模式，1 表示模拟采用 NVT 系综
 dt= 0.001, # 模拟步长
 step_limit = 1, # 模拟总步数
 thermostat=1, # 控温方法，1 表示采用的是 Liujian-Langevin 方法
 langevin_gamma=1.0, # 控温器中的 Gamma_ln 参数
 target_temperature=290, # 目标温度
 write_information_interval=1000, # 输出频率
 amber_irest=0, # 输入方式，0 表示读入amber格式的输入坐标文件，其中不包含速度
 cut=10.0,  # 非键相互作用的距离

案例的输入文件完成后，分别命名为 http://NVT_290_10ns.in 、WATER_ALA.parm7和WATER_ALA_350_cool_290.rst7，这三个文件可以存放在本地工作区的自定义路径下。

2. 加载数据

从三个输入文件中，读取模拟体系所需要的参数，用于最后体系的计算，其加载代码如下：

import argparse
from mindspore import context
 
parser = argparse.ArgumentParser(description='Sponge Controller')
parser.add_argument('--i', type=str, default=None, help='input file')
parser.add_argument('--amber_parm', type=str, default=None, help='paramter file in AMBER type')
parser.add_argument('--c', type=str, default=None, help='initial coordinates file')
parser.add_argument('--r', type=str, default="restrt", help='')
parser.add_argument('--x', type=str, default="mdcrd", help='')
parser.add_argument('--o', type=str, default="mdout", help="")
parser.add_argument('--box', type=str, default="mdbox", help='')
parser.add_argument('--device_id', type=int, default=0, help='')
args_opt = parser.parse_args()
 
context.set_context(mode=context.GRAPH_MODE, device_target="GPU", device_id=args_opt.device_id, save_graphs=False)

3. 构建模拟流程

使用SPONGE中定义的计算力模块和计算能量模块，通过多次迭代进行分子动力学过程演化，使得体系达到我们所需要的平衡态，并记录每一个模拟步骤中得到的能量等数据。其模拟流程构建代码如下：

from src.simulation_initial import Simulation
from mindspore import Tensor
 
if __name__ == "__main__":
 simulation = Simulation(args_opt)
 save_path = args_opt.o
 for steps in range(simulation.md_info.step_limit):
 print_step = steps % simulation.ntwx
 if steps == simulation.md_info.step_limit - 1:
 print_step = 0
 temperature, total_potential_energy, sigma_of_bond_ene, sigma_of_angle_ene, sigma_of_dihedral_ene, \
 nb14_lj_energy_sum, nb14_cf_energy_sum, LJ_energy_sum, ee_ene, _ = simulation(Tensor(steps), Tensor(print_step))
 # compute energy and temperature

4. 运行脚本

python main.py --i /path/NVT_290_10ns.in \
 --amber_parm /path/WATER_ALA.parm7 \
 --c /path/WATER_ALA_350_cool_290.rst7 \
 --o /path/ala_NVT_290_10ns.out

其中，--i 为MD模拟的属性文件，控制模拟过程，--amber_parm 为MD模拟体系的拓扑文件，--c 为我们输入的初始坐标文件，--o 为我们模拟输出的记录文件，其记录了输出每步的能量等信息，--path 为输入文件所在的路径，在本教程中为 sponge_in 文件夹。

使用输入文件，通过在指定温度下进行模拟，计算力和能量，进行分子动力学过程演化。

5. 运行结果

运行得到的结果在 .out 文件中，体系的能量变化都被记录在该文件中，可以查看模拟体系的热力学信息。在 .out 文件中记录了体系的如下信息：

_steps_ _TEMP_ _TOT_POT_ENE_ _BOND_ENE_ _ANGLE_ENE_ _DIHEDRAL_ENE_ _14LJ_ENE_ _14CF_ENE_ _LJ_ENE_ _CF_PME_ENE_

其中记录了模拟过程中输出的各类能量， 分别是迭代次数（_steps_），温度（_TEMP_），总能量（_TOT_POT_E_），键长（_BOND_ENE_），键角（_ANGLE_ENE_），二面角相互作用（_DIHEDRAL_ENE_），非键相互作用，其包含静电力及Leonard-Jones相互作用。

教程文档：https://gitee.com/mindspore/docs/blob/master/tutorials/training/source_zh_cn/advanced_use/hpc_sponge.md

展望

在未来的版本中，会加入更多实用的分子动力学模拟模块，支持更多应用。之后，SPONGE各模块会逐步支持自动微分和自动并行，对于衔接机器学习方案提供更友好的支持。欢迎广大的分子动力学爱好者和研究者加入我们，共同拓展和维护SPONGE。