DP力场与OpenMM的碰撞:OpenMM Plugin for DeepMD-kit
如同LAMMPS一样,OpenMM也是一款高性能的分子动力学软件。与LAMMPS不一样的是,OpenMM软件架构设计更加模块化,也易于扩展。同时在底层的核心程序都使用C++编写,用户交互及力场参数的处理则都是采用了Python编写。既保证了计算性能的同时也使得使用上更加便捷易懂。目前,OpenMM大多用于生物体系的模拟当中,同时也支持CHARMM, Amber, Drude 等大多数常用力场。随着DeepMD-kit的使用范围越来越广,将OpenMM跟DeepMD-kit结合起来也越来越成为必须。
目前针对DeepMD-kit的OpenMM Plugin已经开发完毕,并在Github上面公开内测,欢迎大家使用并提出issue和Pull Request.
https://github.com/JingHuangLab/openmm_deepmd_plugin.git
目前的OpenMM Plugin只支持DP2.0的模型,采用DP1.2/3训练的模型可以先通过dp convert
转成2.0模型再使用。
欢迎加入我们的DP+OpenMM交流群,如二维码过期,请在公众号后台回复“DP+OpenMM”获取最新二维码。
架构设计
OpenMM DeepMD Plugin (以下简称Plugin)的设计主要参照了OpenMM所提供的ExamplePlugin1和 TensorflowPlugin2.
Plugin首先会load一个已经预先训练好的graph model,并生成一个DeepPot Class 的实例。然后,在MD过程中的每一步,Plugin都会从context中提取出每一个粒子的坐标与box信息,并把这些信息提供给DeepPot,得到由graph预测的forces和energy。之后,Plugin就会将forces和energy返回给OpenMM的context,由context执行之后的积分和控温控压等操作。
安装Plugin
Plugin的安装需要依靠OpenMM,libdeepmd和libtensorflow,而这三个软件都可以通过conda的命令迅速安装好。
通过conda准备安装环境依赖OpenMM,libdeepmd, libtensorflow.
conda install -c deepmodeling libdeepmd=2.0.0=1_cuda10.1_gpu
conda install openmm
将Plugin的源码从Github上clone下来,并创建build文件夹。
git clone https://github.com/JingHuangLab/openmm_deepmd_plugin.git
cd openmm_deepmd_plugin && mkdir build && cd build.
通过cmake指定三个软件的文件夹。
cmake .. -DOPENMM_DIR={OPENMM_INSTALLED_DIR} -DDEEPMD_DIR={DEEPMD_INSTALLED_DIR} -DTENSORFLOW_DIR=${TENSORFLOW_DIR}
通过conda install安装的OpenMM,所指定的三个目录地址都应该是conda创建的环境文件夹地址,例如/home/dingye/anaconda3/envs/dp_plugin_test_build
编译并安装Plugin的C++ Interface。
make && make install
单元测试并安装Plugin的Python Interface。
make test && make PythonInstall
使用Plugin
Plugin的使用可以参照Github Repo下的tests3文件夹里面的各种.py文件。Plugin的使用流程一般如下:
传入graph的路径,初始化 DeepmdForce
.
from OpenMMDeepmdPlugin import *
dp_force = DeepmdForce(model_file, "", "", used4Alchemical)
DeepmdForce
的初始化需要四个参数,第一个是训练好的graph路径(str),后面两个参数则是在alchemical模拟中需要设置。然后是布尔类型的参数used4Alchemical
,指示这个DeepmdForce
是否做alchemical simulation.
然后我们需要给 DeepmdForce
添加原子类型,也就是类似于Deepmd-kit训练input.json中的type_map
参数。
dp_force.addType(0, element.oxygen.symbol)
dp_force.addType(1, element.hydrogen.symbol)
注意这里的第一个参数则是代表了type_id,需要与训练数据中的type.raw文件保持一致。第二个参数则是代表了这个这个type的类型(str),需要跟之后addParticle中的参数保持一致。
向 DeepmdForce
添加原子。
# Add particles into force.
for atom in topology.atoms():
if atom.element == element.oxygen:
dp_system.addParticle(element.oxygen.mass)
dp_force.addParticle(atom.index, element.oxygen.symbol)
for at in atom.residue.atoms():
topology.addBond(at, atom)
if at.index != atom.index:
dp_force.addBond(atom.index, at.index)
elif atom.element == element.hydrogen:
dp_system.addParticle(element.hydrogen.mass)
dp_force.addParticle(atom.index, element.hydrogen.symbol)
我们在这里通过遍历topology,向系统中添加了一个个的原子,同时dp_force
也需要addParticle
。值得一提的是dp_force
还调用了addBond
函数,这个函数只会在系统中构建一个BondList,不会影响最终的MD结果。因为OpenMM写入轨迹dcd文件的时候会根据原子所属的group做平移以保持分子都在同一个盒子内,而OpenMM判断原子是否属于一个group所依据的就是各个force下的BondList。
最后,我们设置
dp_force
的单位转换参数系数并将dp_force
加入system中。# Set the units transformation coefficients from openmm to graph input tensors.
# First is the coordinates coefficient, which used for transformation from nanometers to graph needed coordinate unit.
# Second number is force coefficient, which used for transformation graph output force unit to openmm used unit (kJ/(mol * nm))
# Third number is energy coefficient, which used for transformation graph output energy unit to openmm used unit (kJ/mol)
dp_force.setUnitTransformCoefficients(10.0, 964.8792534459, 96.48792534459)
# Add force in dp_system
dp_system.addForce(dp_force)
因为OpenMM做动力学过程中的使用的单位分别是. 所以我们需要在这里设置OpenMM跟graph model之间的单位转换系数。然后我们将dp_force
添加进system中,便完成了dp_force
的所有设置。
Plugin的测试
我们首先比较了动力学过程中,使用同一个graph,OpenMM给出的力与能量与DP的C++接口给出的力与能量的差别有多大。这个测试可以在安装会有make test
的单元测试来做。通过比较,这两者的差距小于.
同时,我们也对Plugin做了其他的一系列测试,包括NVE的守恒,NVT系综的效果,水中氧原子之间的RDF和水的水和自由能(Hydration Free Energy)。
NVE守恒
我们使用VerletIntegrator,0.2fs的时间步,200ps,256个水分子的NVE模拟轨迹来分析。在200ps的模拟中,总能量的方差是0.1343 kJ/mol. 对于平均到每一个自由度的总能量波动,也远小于0.001 kcal/mol.
RDF of Oxygen-Oxygen
我们比较了在同一个graph下,OpenMM和LAMMPS,Gromacs使用DP做NVT模拟下的RDF结果。这里我们使用了0.5fs的时间步,Nose-Hoover Thermostat,298 K, 256 个水分子的4 ns NVT模拟轨迹。最终的结果也显示这三个软件的RDF完全一致。
NVT系综
我们对NVT系综的测试基于Shirts的文章. 如果我们在统计在两个不同的温度下NVT模拟的系统势能概率分布,理论上,他们应该满足下面这样的线性分布。
其中.则是亥姆霍兹自由能。
我们利用Plugin做了5个NVT模拟,温度分别是300,305,310,315和320K。可以看出,利用Plugin的DP下的NVT模拟也显示出了非常好的线性关系。
水的水和自由能(Hydration Free Energy)
我们在Plugin当中还实现了基于DP的alchemical protocol。利用这种alchemical模拟策略计算了水的水和自由能,并跟实验值做了比较。
Alchemical 模拟是基于下图的protocol。
对于每一个中间态,系统的哈密顿量可以写为
所以在这个Protocol中,从1变成0就是系统的alchemical解耦过程。原则上,我们需要三个graph完成这个protocol,但是因为水分子体系的各向均一性,而且DP训练的graph泛化性能也非常不错,我们使用了一个graph来代替三个graph来计算水的水和自由能。
在Alchemical模拟的使用中,与之前介绍的使用有一小部分的不同。首先是DeepmdForce
的初始化需要三个model_file,分别代表了, used4Alchemical也应该设置为True。
dp_force = DeepmdForce(model_file, model_file, model_file, used4Alchemical)
同时,我们也应该设置对于graph 的粒子索引和。
for atom in topology.atoms():
if int(atom.residue.id) == alchemical_resid:
graph2_particles.append(atom.index)
else:
graph1_particles.append(atom.index)
if atom.element == element.oxygen:
dp_system.addParticle(element.oxygen.mass)
dp_force.addParticle(atom.index, element.oxygen.symbol)
for at in atom.residue.atoms():
topology.addBond(at, atom)
if at.index != atom.index:
dp_force.addBond(atom.index, at.index)
elif atom.element == element.hydrogen:
dp_system.addParticle(element.hydrogen.mass)
dp_force.addParticle(atom.index, element.hydrogen.symbol)
dp_force.setAtomsIndex4Graph1(graph1_particles)
dp_force.setAtomsIndex4Graph2(graph2_particles)
dp_force.setLambda(Lambda)
在最终的Alchemical模拟中,我们设置了10个,每一个都在300K,0.2fs时间步,3 ns NVT模拟参数中运行。MBAR用于对最后的结果进行了计算。
可以看到最后的结果中,. 两个值还是比较接近的。
小结一下
在这里我们实现了给Deepmd-kit的OpenMM Plugin,以后大家也可以使用OpenMM去做DP的模拟了。而且你甚至都可以做基于DP的Alchemical模拟。尽管目前的性能上Plugin的模拟会比LAMMPS要慢一些(256个水分子,1.5 ns/day, A40),但是后续我们会将neighbor list作为输入给DeepPot,以及其他各种性能优化手段提升Plugin的性能。未来我们还将使Plugin支持NNP/MM的混合力场下的模拟,希望能使得DP模型在生物体系的模拟当中带来更多的应用。
我们非常期待大家一起帮忙测试、使用、提issue,共同改善目前的Plugin。
感谢黄晶老师的大力支持!感谢王涵老师,张林峰提供的graph与帮助。同时也非常感谢刘歆子建、王应泽还有张与之的讨论与帮助!
补充资料
1.https://github.com/openmm/openmmexampleplugin
2.https://github.com/openmm/openmm-tensorflow
3.https://github.com/JingHuangLab/openmm_deepmd_plugin/tree/master/tests