MindSponge分子动力学模拟——多路径分子模拟（2024.05）
随着硬件算力的发展，以及AI技术的日益增进，我们不仅可以借助深度学习框架来加速分子动力学模拟，以及降低分子模拟开发的门槛。还可以实现高通量模拟，使得用最小的开销并行的运行多个分子模拟成为可能。
技术背景
在前面的MindSponge教程系列博客中，我们已经介绍过MindSponge分子动力学模拟框架的基础功能使用方法，例如MindSponge的安装与使用、定义分子系统、计算单点能和迭代器等等。这些模块和功能，更多的是凭借MindSpore深度学习框架的自动微分、GPU加速和Python语言的灵活性，而本文要介绍的是深度学习框架可以带来的另一个重大优势：多路径的高通量分子模拟。
使用逻辑
在MindSponge分子动力学模拟框架中，原生的支持了多batch的计算模式，这个场景跟深度学习中处理Dataset时用到的batchwise是非常类似的。在分子模拟中，我们可以给定同一个体系的多个初始位置，可以相同也可以不同，然后使用同一个力场和迭代器进行优化、更新，最终得到不同的轨迹和不同的构象。在MindSponge框架下，单路径和多路径这两种模式的切换，只需要在定义模拟系统system时，先定义单路径的system，再按照不同的初始构象重构一次system即可：
mol_file = 'alad_water_opt.pdb'
mol = Protein(mol_file, template=['protein0.yaml', 'water.spce.yaml'])
mol.num_walker = 2
mol.build_system()
# 对coordinate的batch维度进行扩维，并以wrap形式填充
mol.build_space(msnp.pad(mol.coordinate, ((0, 1), (0, 0), (0, 0)), mode='wrap'))
其他的配置可以直接复用单路径分子模拟的代码。
真实案例
这里我们使用到了一个丙氨酸二肽+一个水分子的简单测试体系，其pdb文件内容为：
MODEL 1
ATOM 1 H1 ACE A 1 2.012 1.007 0.000 1.0 0.0 H
ATOM 2 CH3 ACE A 1 1.996 2.097 0.002 1.0 0.0 C
ATOM 3 H2 ACE A 1 1.484 2.455 0.891 1.0 0.0 H
ATOM 4 H3 ACE A 1 1.487 2.455 -0.890 1.0 0.0 H
ATOM 5 C ACE A 1 3.409 2.623 0.011 1.0 0.0 C
ATOM 6 O ACE A 1 4.349 1.840 0.032 1.0 0.0 O
ATOM 7 N ALA A 2 3.559 3.946 -0.023 1.0 0.0 N
ATOM 8 H ALA A 2 2.741 4.544 -0.004 1.0 0.0 H
ATOM 9 CA ALA A 2 4.862 4.614 -0.000 1.0 0.0 C
ATOM 10 HA ALA A 2 5.411 4.315 0.894 1.0 0.0 H
ATOM 11 CB ALA A 2 5.665 4.220 -1.246 1.0 0.0 C
ATOM 12 HB1 ALA A 2 5.123 4.520 -2.141 1.0 0.0 H
ATOM 13 HB2 ALA A 2 6.638 4.707 -1.220 1.0 0.0 H
ATOM 14 HB3 ALA A 2 5.834 3.147 -1.264 1.0 0.0 H
ATOM 15 C ALA A 2 4.706 6.138 0.016 1.0 0.0 C
ATOM 16 O ALA A 2 3.588 6.651 0.002 1.0 0.0 O
ATOM 17 N NME A 3 5.838 6.851 0.016 1.0 0.0 N
ATOM 18 H NME A 3 6.722 6.359 0.001 1.0 0.0 H
ATOM 19 CH3 NME A 3 5.854 8.301 0.002 1.0 0.0 C
ATOM 20 HH31 NME A 3 4.827 8.674 -0.000 1.0 0.0 H
ATOM 21 HH32 NME A 3 6.363 8.660 0.894 1.0 0.0 H
ATOM 22 HH33 NME A 3 6.363 8.656 -0.894 1.0 0.0 H
ATOM 23 O WAT A 4 -0.314 8.560 2.297 1.0 0.0 O
ATOM 24 H1 WAT A 4 0.499 9.142 2.295 1.0 0.0 H
ATOM 25 H2 WAT A 4 -1.119 9.148 2.298 1.0 0.0 H
TER
ENDMDL
END
其构象大概长这个样子：
在安装好MindSponge之后，我们可以直接使用如下代码来构建一个多路径分子模拟的示例：
import mindspore as ms
from mindspore import numpy as msnp
# 设置MindSpore后台计算使用GPU静态图模式，速度会更快
ms.set_context(mode=ms.GRAPH_MODE, device_target='GPU', device_id=0)
# 引入必要的MindSponge模块
from sponge import Protein, ForceField, WithEnergyCell, UpdaterMD, RunOneStepCell, Sponge
from sponge.control import Langevin, LeapFrog
from sponge.callback import RunInfo
# 导入分子构象文件，构建模拟系统
mol_file = './alad_water_opt.pdb'
mol = Protein(mol_file, template=['protein0.yaml', 'water.spce.yaml'])
# 设置多路径的数量为2
mol.num_walker = 2
# 先重构分子系统的拓扑信息
mol.build_system()
# 再构建分子系统的空间信息，此时输入的坐标batch维度要跟多路径的数量对齐
mol.build_space(msnp.pad(mol.coordinate, ((0, 1), (0, 0), (0, 0)), mode='wrap'))
# 查看多路径数量是否更新成功
print (mol.num_walker)
# 打印每条轨迹下前三个原子的空间坐标
print (mol.coordinate.asnumpy()[:, :3])
# 打印分子系统的原子类型信息
print (mol.atom_type)
# 力场构建，与单路径构建方法一致
potential = ForceField(mol, parameters=['amber.ff14sb', 'spce'], use_pme=False)
with_energy = WithEnergyCell(mol, potential)
# 计算单点能
print (with_energy.calc_energies().sum(-1))
# 构建朗之万热浴积分器，与单路径下的使用方法一致
temp = 300
thermostat = Langevin(mol, temp)
updater = UpdaterMD(mol, time_step=2e-3, integrator=LeapFrog(mol), thermostat=thermostat)
one_step = RunOneStepCell(energy=with_energy, optimizer=updater)
md = Sponge(network=one_step)
# 定义回调函数，用于打印模拟过程中的势能等信息
run_info = RunInfo(1)
md.run(10, callbacks=[run_info])
# 打印迭代之后的体系空间信息和能量信息
print (mol.coordinate.asnumpy()[:, :3])
print (with_energy.calc_energies().sum(-1))
模拟的输出结果为：
2 # 两条模拟路径
[[[2.0120001e-01 1.0070000e-01 0.0000000e+00]
[1.9960001e-01 2.0969999e-01 2.0000001e-04]
[1.4839999e-01 2.4550000e-01 8.9099996e-02]]
[[2.0120001e-01 1.0070000e-01 0.0000000e+00]
[1.9960001e-01 2.0969999e-01 2.0000001e-04]
[1.4839999e-01 2.4550000e-01 8.9099996e-02]]] # 初始构象一致
[['H1' 'CH3' 'H2' 'H3' 'C' 'O' 'N' 'H' 'CA' 'HA' 'CB' 'HB1' 'HB2' 'HB3'
'C' 'O' 'N' 'H' 'CH3' 'HH31' 'HH32' 'HH33' 'OW' 'HW' 'HW']] # 多路径模拟共享同一个模拟体系
[-67.05737 -67.05737] # 初始能量一致
[MindSPONGE] Started simulation at 2024-05-27 16:43:31
[MindSPONGE] Step: 1, E_pot: [-67.05737 -67.05737], E_kin: [0.93526113 1.1376405 ], E_tot: [-66.12211 -65.91973], Temperature: [3.1246128 3.8007417]
[MindSPONGE] Step: 2, E_pot: [-67.76422 -67.67892], E_kin: [2.87101 2.2547605], E_tot: [-64.89321 -65.42416], Temperature: [9.591754 7.5329266]
[MindSPONGE] Step: 3, E_pot: [-68.390305 -68.228424], E_kin: [4.2298875 3.9742615], E_tot: [-64.160416 -64.254166], Temperature: [14.131627 13.277606]
[MindSPONGE] Step: 4, E_pot: [-68.723 -68.80056], E_kin: [5.0383153 5.081285 ], E_tot: [-63.684685 -63.719276], Temperature: [16.8325 16.976059]
[MindSPONGE] Step: 5, E_pot: [-69.14583 -69.199196], E_kin: [6.2894073 5.8883085], E_tot: [-62.856422 -63.310886], Temperature: [21.012274 19.672243]
[MindSPONGE] Step: 6, E_pot: [-69.68199 -69.531006], E_kin: [7.1465344 7.5664873], E_tot: [-62.535458 -61.96452 ], Temperature: [23.875849 25.27887 ]
[MindSPONGE] Step: 7, E_pot: [-70.71329 -70.14986], E_kin: [9.099633 8.052189], E_tot: [-61.613655 -62.097668], Temperature: [30.400955 26.901548]
[MindSPONGE] Step: 8, E_pot: [-71.7045 -70.37026], E_kin: [9.534681 9.304852], E_tot: [-62.169815 -61.06541 ], Temperature: [31.854406 31.086567]
[MindSPONGE] Step: 9, E_pot: [-71.351135 -70.86687 ], E_kin: [ 9.35821 10.023544], E_tot: [-61.992928 -60.843323], Temperature: [31.264833 33.487644]
[MindSPONGE] Step: 10, E_pot: [-71.8908 -71.19918], E_kin: [10.796071 11.098853], E_tot: [-61.09473 -60.100327], Temperature: [36.068584 37.080147]
[MindSPONGE] Finished simulation at 2024-05-27 16:43:36
[MindSPONGE] Simulation time: 5.89 seconds.
--------------------------------------------------------------------------------
[[[ 0.20221242 0.10218664 -0.001829 ]
[ 0.1975184 0.210064 -0.00026266]
[ 0.14476728 0.24293597 0.08909092]]
[[ 0.20099933 0.10118634 -0.00139746]
[ 0.1981248 0.20958814 0.00327332]
[ 0.14694569 0.2449529 0.09299506]]] # 最终构象出现差异
[-71.71762 -71.29018] # 最终能量出现差异
由于这里使用的朗之万积分器带有随机项，因此即时是具备相同的初始构象，在多步的模拟之后，构象和能量都与初始构象发生偏离，而且两条路径下的构象也是有所差异的。当然，类似的，我们可以模拟不同的初始构象在没有随机性的迭代器下的迭代过程，感兴趣的童鞋可以自行尝试一下。
总结概要
随着硬件算力的发展，以及AI技术的日益增进，我们不仅可以借助深度学习框架来加速分子动力学模拟，以及降低分子模拟开发的门槛。还可以实现高通量模拟，使得用最小的开销并行的运行多个分子模拟成为可能。
版权声明
本文首发链接为：https://www.cnblogs.com/dechinphy/p/multi-md.html
作者ID：DechinPhy
更多原著文章：https://www.cnblogs.com/dechinphy/
请博主喝咖啡：https://www.cnblogs.com/dechinphy/gallery/image/379634.html
