MindSponge分子动力学模拟——使用迭代器进行系统演化（2023.09）

技术背景

在前面几篇博客中，我们已经介绍过使用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恒温恒容过程，更多的模拟方法有待大家一起研究探讨。