MindSponge分子动力学模拟——多路径分子模拟（2024.05）

技术背景

在前面的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