MindSponge分子动力学模拟——使用迭代器进行系统演化（2023.09）
在经过前面几篇博客的介绍之后，我们可以定义一些目标的分子体系，并且计算其单点能。而分子模拟的精髓就在于快速的迭代和演化，也就是本文所要介绍的迭代器相关的内容。在具备了分子系统、单点能和迭代器这三者之后，就可以正式开始进行分子动力学模拟。常见的模拟过程有：能量极小化、NVT恒温恒容过程、NPT恒温恒压过程以及NVE微正则系综，本文所涉及的主要是能量极小化以及NVT恒温恒容过程，更多的模拟方法有待大家一起研究探讨。
技术背景
在前面几篇博客中，我们已经介绍过使用MindSponge去定义一个系统以及使用MindSponge计算一个分子系统的单点能。这篇文章我们将介绍一下在MindSponge中定义迭代器Updater，并使用Sponge对系统进行演化，最后使用CallBack对输出结果进行追踪和保存。
分子动力学迭代器与深度学习优化器
首先我们回顾一下这张MindSponge的软件架构图：
在这里Updater从WithForceCell接收一个力（这个力可以是直接从外界输入的，也可以是利用MindSpore的自动微分功能，对WithEnergyCell进行自动微分得到的力），并将其作用在Molecule系统上，得到一个Molecule新的状态。这个流程可以对照深度学习里面的优化器，比如梯度下降法，我们也是从损失函数中计算一个梯度（依赖于自动微分或者差分），然后将这个梯度根据不同的优化算法计算一个迭代值，最后作用在网络参数上，得到一组新的网络参数。所以，分子动力学中的迭代器跟深度学习中的优化器，本质上可以是相同的，换句话说，我们可以直接使用深度学习中的一些优化算法（如Adam等）来作为分子动力学模拟中的迭代器。比如说，我们可以参考如下方案做一个能量极小化。
能量极小化
其实做能量极小化的思路是简单的，我们假设单点能\(E(R)\)（维度为[B,1]）是一个关于原子坐标\(R\)（维度为[B,A,D]）的一个函数。那么所谓的能量极小化，就是找到这样的一个最低能量值：
\[E_{min}=min_{R}\{E(R)\}
\]
假定我们就使用MindSpore中内置的Adam算法，那么相应的代码实现可以参考如下案例。首先我们有一个用于模拟的pdb案例文件：
REMARK Generated By Xponge (Molecule)
ATOM 1 N ALA 1 -0.095 -11.436 -0.780
ATOM 2 CA ALA 1 -0.171 -10.015 -0.507
ATOM 3 CB ALA 1 1.201 -9.359 -0.628
ATOM 4 C ALA 1 -1.107 -9.319 -1.485
ATOM 5 O ALA 1 -1.682 -9.960 -2.362
ATOM 6 N ARG 2 -1.303 -8.037 -1.397
ATOM 7 CA ARG 2 -2.194 -7.375 -2.328
ATOM 8 CB ARG 2 -3.606 -7.943 -2.235
ATOM 9 CG ARG 2 -4.510 -7.221 -3.228
ATOM 10 CD ARG 2 -5.923 -7.789 -3.136
ATOM 11 NE ARG 2 -6.831 -7.111 -4.087
ATOM 12 CZ ARG 2 -8.119 -7.421 -4.205
ATOM 13 NH1 ARG 2 -8.686 -8.371 -3.468
ATOM 14 NH2 ARG 2 -8.844 -6.747 -5.093
ATOM 15 C ARG 2 -2.273 -5.882 -2.042
ATOM 16 O ARG 2 -1.630 -5.388 -1.119
ATOM 17 N ALA 3 -3.027 -5.119 -2.777
ATOM 18 CA ALA 3 -3.103 -3.697 -2.505
ATOM 19 CB ALA 3 -1.731 -3.041 -2.625
ATOM 20 C ALA 3 -4.039 -3.001 -3.483
ATOM 21 O ALA 3 -4.614 -3.643 -4.359
ATOM 22 N ALA 4 -4.235 -1.719 -3.394
ATOM 23 CA ALA 4 -5.126 -1.057 -4.325
ATOM 24 CB ALA 4 -6.538 -1.625 -4.233
ATOM 25 C ALA 4 -5.205 0.436 -4.039
ATOM 26 O ALA 4 -4.561 0.930 -3.116
ATOM 27 OXT ALA 4 -5.915 1.166 -4.728
TER
然后可以用MindSponge实现一个简单的用Adam进行能量极小化的代码：
from mindspore import nn, context
from sponge import ForceField, Sponge, set_global_units, Protein
from sponge.callback import RunInfo, WriteH5MD
# 配置MindSpore的执行环境
context.set_context(mode=context.GRAPH_MODE, device_target='GPU', device_id=1)
# 配置全局单位
set_global_units('A', 'kcal/mol')
# 定义一个基于case1.pdb的分子系统
system = Protein('case1.pdb', template=['protein0.yaml'], rebuild_hydrogen=True)
# 定义一个amber.ff99sb的力场
energy = ForceField(system, parameters=['AMBER.FF99SB'])
# 定义一个学习率为1e-03的Adam优化器
min_opt = nn.Adam(system.trainable_params(), 1e-03)
# 定义一个用于执行分子模拟的Sponge实例
md = Sponge(system, potential=energy, optimizer=min_opt)
# RunInfo这个回调函数可以在屏幕上根据指定频次输出能量参数
run_info = RunInfo(200)
# WriteH5MD回调函数，可以将轨迹、能量、力和速度等参数保留到一个hdf5文件中，文件后缀为h5md
cb_h5md = WriteH5MD(system, 'test.h5md', save_freq=10, write_image=False)
# 开始执行分子动力学模拟，运行2000次迭代
md.run(2000, callbacks=[run_info, cb_h5md])
上述代码的运行结果如下：
[MindSPONGE] Adding 45 hydrogen atoms for the protein molecule in 0.003 seconds.
[MindSPONGE] Started simulation at 2023-09-04 15:14:29
[MindSPONGE] Step: 0, E_pot: 1200.4639
[MindSPONGE] Step: 200, E_pot: 7.763489
[MindSPONGE] Step: 400, E_pot: -70.34643
[MindSPONGE] Step: 600, E_pot: -96.88522
[MindSPONGE] Step: 800, E_pot: -109.98717
[MindSPONGE] Step: 1000, E_pot: -117.33747
[MindSPONGE] Step: 1200, E_pot: -121.95378
[MindSPONGE] Step: 1400, E_pot: -125.20764
[MindSPONGE] Step: 1600, E_pot: -127.72044
[MindSPONGE] Step: 1800, E_pot: -129.79828
[MindSPONGE] Finished simulation at 2023-09-04 15:15:16
[MindSPONGE] Simulation time: 46.79 seconds.
--------------------------------------------------------------------------------
此时我们可以看到整个分子能量是一直在下降的，同时在该路径下生成了一个test.h5md的轨迹文件。打开这个轨迹文件的方式有两种，一种是使用silx view test.h5md（可以使用python3 -m pip install silx来进行安装）来查看文件的具体内容，比如可以看到这样的一个结果：
既可以使用曲线图的方式来进行浏览，也可以使用表格的方式来进行浏览。而如果参考这个README.md文件的指示安装一个VMD插件的话，就可以在本地直接用VMD来可视化分子模拟的轨迹，输出为gif动态图如下所示：
分子动力学模拟
在上一个章节中，我们演示了一下使用MindSpore中自带的优化器做了一个能量极小化，得到了一个稳定的构象。需要说明的是，在上一步的过程中，如果我们想保留最后一帧的结果，既可以使用VMD导出，也可以在WriteH5MD中进行相应的参数配置，比如在上面的案例中将对应代码修改为：
cb_h5md = WriteH5MD(system, 'test.h5md', save_freq=10, write_image=False, save_last_pdb='case1_min.pdb')
就可以在运行结束之后生成一个case1_min.pdb文件。因为在上一步的运行过程中我们已经对氢原子进行了重构，因此这里如果我们重新执行一个动力学模拟的任务的话，可以不需要重构氢原子，对应的代码应调整为：
system = Protein('case1_min.pdb', template=['protein0.yaml'], rebuild_hydrogen=False)
那么最终整体的代码如下所示：
from mindspore import context
from sponge import ForceField, Sponge, set_global_units, Protein, UpdaterMD, WithEnergyCell, RunOneStepCell
from sponge.function import VelocityGenerator
from sponge.callback import RunInfo, WriteH5MD
context.set_context(mode=context.GRAPH_MODE, device_target='GPU', device_id=1)
set_global_units('A', 'kcal/mol')
# 这里设置rebuild_hydrogen为False，意为不对氢原子进行重构
system = Protein('case1_min.pdb', template=['protein0.yaml'], rebuild_hydrogen=False)
energy = ForceField(system, parameters=['AMBER.FF99SB'])
# 定义一个速度生成器，使得生成的随机速度可以让系统处在300K的温度下
vgen = VelocityGenerator(300)
# 根据速度生成器生成相应的原子速度
velocity = vgen(system.shape, system.atom_mass)
# UpdaterMD迭代器，这里给定了temperature和thermostat，是一个NVT过程，积分器使用的是velocity_verlet算法
opt = UpdaterMD(system=system,
time_step=1e-3,
velocity=velocity,
integrator='velocity_verlet',
temperature=300,
thermostat='langevin')
# 定义Sponge可以有很多种方法，这里采用的是RunOneStep来定义
sim = WithEnergyCell(system, energy)
one_step = RunOneStepCell(energy=sim, optimizer=opt)
md = Sponge(one_step)
run_info = RunInfo(200)
cb_h5md = WriteH5MD(system, 'test.h5md', save_freq=10, write_image=False)
md.run(2000, callbacks=[run_info, cb_h5md])
除了前面提到的两处代码修改之外，其他的调整见代码中的注释。上述代码的运行结果为：
[MindSPONGE] Started simulation at 2023-09-04 17:06:26
[MindSPONGE] Step: 0, E_pot: -131.60876, E_kin: 52.25413, E_tot: -79.35463, Temperature: 313.0393
[MindSPONGE] Step: 200, E_pot: -112.29764, E_kin: 39.081543, E_tot: -73.216095, Temperature: 234.12614
[MindSPONGE] Step: 400, E_pot: -125.55388, E_kin: 66.579636, E_tot: -58.974243, Temperature: 398.85922
[MindSPONGE] Step: 600, E_pot: -127.8087, E_kin: 59.09694, E_tot: -68.71176, Temperature: 354.03256
[MindSPONGE] Step: 800, E_pot: -150.88245, E_kin: 69.84598, E_tot: -81.03647, Temperature: 418.42694
[MindSPONGE] Step: 1000, E_pot: -163.9544, E_kin: 62.79237, E_tot: -101.16203, Temperature: 376.1708
[MindSPONGE] Step: 1200, E_pot: -159.15376, E_kin: 55.752754, E_tot: -103.40101, Temperature: 333.99854
[MindSPONGE] Step: 1400, E_pot: -159.43027, E_kin: 57.967392, E_tot: -101.462875, Temperature: 347.26575
[MindSPONGE] Step: 1600, E_pot: -163.72491, E_kin: 57.850487, E_tot: -105.87443, Temperature: 346.56543
[MindSPONGE] Step: 1800, E_pot: -168.06078, E_kin: 54.730705, E_tot: -113.33007, Temperature: 327.87573
[MindSPONGE] Finished simulation at 2023-09-04 17:07:13
[MindSPONGE] Simulation time: 47.41 seconds.
--------------------------------------------------------------------------------
因为上面这个案例我们运行的是一个NVT恒温过程，因此我们可以用silx view看到，在结果中所保存的温度最终会逐渐趋近于300K附近：
同样的，我们可以用VMD的插件来可视化这个分子运动的轨迹：
迭代器切换
之所以采用Python这一编程语言来实现，很大程度上就考虑到了各种方法实现的便捷性。比如上述章节中定义好一个Sponge实例之后，我们需要切换其中的优化器——这其实是一个比较常用的方法，例如我们运行完了一个能量极小化的过程，我们甚至都不需要退出程序，直接用演化好的system可以继续执行NVT过程，然后执行NPT过程。而这一系列的操作只需要用到一个函数：change_optimizer。
from mindspore import nn, context
from sponge import ForceField, Sponge, set_global_units, Protein, UpdaterMD
from sponge.function import VelocityGenerator
from sponge.callback import RunInfo, WriteH5MD
context.set_context(mode=context.GRAPH_MODE, device_target='GPU', device_id=1)
set_global_units('A', 'kcal/mol')
system = Protein('case1.pdb', template=['protein0.yaml'], rebuild_hydrogen=True)
energy = ForceField(system, parameters=['AMBER.FF99SB'])
min_opt = nn.Adam(system.trainable_params(), 1e-03)
md = Sponge(system, potential=energy, optimizer=min_opt)
run_info = RunInfo(200)
cb_h5md = WriteH5MD(system, 'test_min.h5md', save_freq=10, write_image=False)
md.run(2000, callbacks=[run_info, cb_h5md])
# 定义一个新的迭代器，并用change_optimizer完成切换
vgen = VelocityGenerator(300)
velocity = vgen(system.shape, system.atom_mass)
opt = UpdaterMD(system=system,
time_step=1e-3,
velocity=velocity,
integrator='velocity_verlet',
temperature=300,
thermostat='langevin')
md.change_optimizer(opt)
nvt_h5md = WriteH5MD(system, 'test_nvt.h5md', save_freq=10, write_image=False)
md.run(2000, callbacks=[run_info, nvt_h5md])
上述代码的运行结果如下：
[MindSPONGE] Adding 45 hydrogen atoms for the protein molecule in 0.003 seconds.
[MindSPONGE] Started simulation at 2023-09-04 17:29:29
[MindSPONGE] Step: 0, E_pot: 1200.4639
[MindSPONGE] Step: 200, E_pot: 7.7634506
[MindSPONGE] Step: 400, E_pot: -70.34645
[MindSPONGE] Step: 600, E_pot: -96.885216
[MindSPONGE] Step: 800, E_pot: -109.987114
[MindSPONGE] Step: 1000, E_pot: -117.33747
[MindSPONGE] Step: 1200, E_pot: -121.95381
[MindSPONGE] Step: 1400, E_pot: -125.20767
[MindSPONGE] Step: 1600, E_pot: -127.72041
[MindSPONGE] Step: 1800, E_pot: -129.79825
[MindSPONGE] Finished simulation at 2023-09-04 17:30:13
[MindSPONGE] Simulation time: 43.96 seconds.
--------------------------------------------------------------------------------
[MindSPONGE] Started simulation at 2023-09-04 17:30:14
[MindSPONGE] Step: 0, E_pot: -131.61736, E_kin: 58.91965, E_tot: -72.69771, Temperature: 352.9705
[MindSPONGE] Step: 200, E_pot: -114.52379, E_kin: 45.360672, E_tot: -69.16312, Temperature: 271.74258
[MindSPONGE] Step: 400, E_pot: -120.40167, E_kin: 48.916946, E_tot: -71.484726, Temperature: 293.04718
[MindSPONGE] Step: 600, E_pot: -127.64978, E_kin: 60.41433, E_tot: -67.23545, Temperature: 361.92465
[MindSPONGE] Step: 800, E_pot: -152.59546, E_kin: 72.86487, E_tot: -79.73059, Temperature: 436.5122
[MindSPONGE] Step: 1000, E_pot: -152.563, E_kin: 71.374565, E_tot: -81.18844, Temperature: 427.58426
[MindSPONGE] Step: 1200, E_pot: -160.1132, E_kin: 68.66432, E_tot: -91.44888, Temperature: 411.34796
[MindSPONGE] Step: 1400, E_pot: -152.07262, E_kin: 58.346176, E_tot: -93.72644, Temperature: 349.53497
[MindSPONGE] Step: 1600, E_pot: -152.71495, E_kin: 50.630737, E_tot: -102.08421, Temperature: 303.314
[MindSPONGE] Step: 1800, E_pot: -148.00838, E_kin: 52.72198, E_tot: -95.28639, Temperature: 315.84204
[MindSPONGE] Finished simulation at 2023-09-04 17:31:00
[MindSPONGE] Simulation time: 46.71 seconds.
--------------------------------------------------------------------------------
总结概要
在经过前面几篇博客的介绍之后，我们可以定义一些目标的分子体系，并且计算其单点能。而分子模拟的精髓就在于快速的迭代和演化，也就是本文所要介绍的迭代器相关的内容。在具备了分子系统、单点能和迭代器这三者之后，就可以正式开始进行分子动力学模拟。常见的模拟过程有：能量极小化、NVT恒温恒容过程、NPT恒温恒压过程以及NVE微正则系综，本文所涉及的主要是能量极小化以及NVT恒温恒容过程，更多的模拟方法有待大家一起研究探讨。
版权声明
本文首发链接为：https://www.cnblogs.com/dechinphy/p/updater-md.html
作者ID：DechinPhy
更多原著文章：https://www.cnblogs.com/dechinphy/
请博主喝咖啡：https://www.cnblogs.com/dechinphy/gallery/image/379634.html
