一文带你认识 MindSpore 新一代分子模拟库 SPONGE

关注

发布于: 2021 年 05 月 25 日

摘要：基于 MindSpore 自动并行、图算融合等特性，SPONGE 可高效地完成传统分子模拟过程，利用 MindSpore 自动微分的特性，可以将神经网络等 AI 方法与传统分子模拟进行结合。

本文分享自华为云社区《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 argparsefrom 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 Simulationfrom 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