分子模拟

查看源文件

MindSpore SPONGE采取了一种独一无二的 “类AI”分子模拟程序架构:

MindSpore SPONGE

该架构的特点在于其将分子模拟变成了一种特殊的AI训练过程,这使得MindSpore SPONGE可以像使用MindSpore对训练网络模型那样进行分子动力学模拟。

首先,MindSpore和MindSpore SPONGE都需要设置三个基础单元(Cell)或功能模块,而且MindSpore和MindSpore SPONGE的这三个基础单元之间有着一一对应的关系:

MindSpore

MindSpore SPONGE

网络模型

network

模拟体系

system

损失函数

loss_fn

势能函数

potential

优化器

optimizer

更新器

updater

在设置好三个功能模块后,还需要将它们进行封装,这样才能执行训练或者模拟。

例如,一个典型的MindSpore训练网络的代码为:

from mindspore.nn import WithLossCell
from mindspore.train import Model
from mindspore.train import LossMonitor
from mindspore.train import ModelCheckpoint, CheckpointConfig
# 组合方式一:
model = Model(network, loss_fn, optimizer)
# 组合方式二:
loss_net = WithLossCell(network, loss_fn)
model = Model(loss_net, optimizer=optimizer)
# 设置回调函数:
monitor_cb = LossMonitor(16)
ckpoint_cb = ModelCheckpoint('network', config=CheckpointConfig(32))
# 执行训练过程:
model.train(1024, dataset, callbacks=[monitor_cb, ckpoint_cb])

而使用MindSpore SPONGE进行分子动力学模拟的代码则同上面的代码非常类似:

from mindsponge import SimulationCell
from mindsponge import Sponge
from mindsponge.callback import RunInfo
from mindsponge.callback import WriteH5MD
# 组合方式一:
md = Sponge(system, potential, updater)
# 组合方式二:
sim = SimulationCell(system, potential, cutoff=1.0)
md = Sponge(sim, optimizer= updater)
# 设置回调函数:
run_info = RunInfo(10)
cb_h5md = WriteH5MD(system, 'traj.h5md', save_freq=10)
# 执行模拟过程:
md.run(1000, callbacks=[run_info, cb_h5md])

可以看到,在MindSpore有两种封装方式,一种是直接使用MindSpore的顶层模块Model将网络模型network、损失函数loss_fn以及优化器optimizer三者进行封装,这种情况适用于最简单的监督式学习的训练。而另外一种封装方式,则是先使用WithLossCell模块将网络模型network和损失函数loss_fn封装起来,然后再用Model模块将封装好的WithLossCell模块和优化器optimizer封装起来。这种情况适用于较为复杂的训练,例如使用MindSpore进行生成对抗网络GANs(generative adversarial networks)的训练便可以通过重载WithLossCell的方式实现。

而MindSpore SPONGE用于对三大功能模块进行封装的顶层模块是Sponge,同样也有两种方式对三大基础单元进行封装。第一种是直接用Sponge将模拟体系system、势能函数potential以及更新器updater进行封装,这种情况同样也是适用于最普通的模拟。而另外一种则是先使用SimulationCell将模拟体系system和势能函数potential进行封装,然后在将封装好的SimulationCell同更新器updater一起封装在Sponge中。这种情况可以通过调整SimulationCell的参数实现一些比较复杂的模拟功能,如调整阶段(cutoff)半径以及设置偏向势(bias potential)等。

在封装完三大基础单元后,MindSpore通过调用Model模块的train()函数执行训练过程,而MindSpore SPONGE则通过调用Sponge模块的run()函数执行模拟过程。在MindSpore可以通过使用不同的回调(callback)函数来处理训练过程中的信息(不涉及改变计算图),例如使用LossMonitor()函数监视训练过程中的参数变化,以及使用ModelCheckpoint()函数保存网络参数等。而MindSpore SPONGE同样也可以使用回调函数来处理模拟过程中的信息,例如使用RunInfo()函数打印训练过程的参数信息,以及使用WriteH5MD()将模拟过程中体系坐标的变化记录为模拟轨迹(trajectory)。

体系(分子)模块:System(Molecule)

体系(分子)模块用于描述分子体系的化学性质,例如化学成分、拓扑信息、空间坐标等。体系模拟的基类(父类)是mindsponge.system.Molecule,其包含的参数主要有:

  • 原子名(atom_name):体系中各个原子的原子名(atom name),用于区分不同原子

  • 原子类型(atom_type):体系中各个原子的原子类型(atom type),用于设定力场参数

  • 原子质量(atom_mass):体系中各个原子的质量

  • 原子电荷(atom_charge):体系中各个原子的点电荷

  • 坐标(coordinate):体系原子的空间坐标

  • 周期性盒子(pbc_box):周期性边界条件“原胞”盒子的尺寸

  • 键(bond):原子间的键连接关系(不区分单双键)

  • 残基(residue):用于区分大分子中的分子片段或体系中的不同分子(注:Molecule中residue的概念跟PDB文件中的residue类似,每个residue不仅用于表示单个氨基酸残基,也用于表示单个小分子,如水分子和无机盐离子等)

Molecule类可以通过人为传入上述的参数进行初始化:

from mindsponge import Molecule

system = Molecule(
    atom_name=['O', 'H1', 'H2'],
    atom_type=['OW', 'HW', 'HW'],
    atom_mass=[16, 1.008, 1.008],
    coordinate=[
        [0, 0, 0],
        [0.1, 0, 0],
        [-0.0333, 0.0943, 0]
        ],
    bond=[[[0, 1], [0, 2]]],
)

而对于蛋白质分子,则可以使用Molecule类的子类Protein类,通过读取PDB文件进行初始化:

    from mindsponge import Protein
    system = Protein(pdb='protein.pdb')

势能(力场)模块:Potential (Force Field)

势能模块用于描述模拟过程中所使用的势能函数,势能模块的基类(父类)是mindsponge.potential.PotentialCell。而经典分子模拟一般使用分子力场(force field)作为势能函数,在MindSpore SPONGE则为力场模块。力场模块的基类是mindsponge.potential.ForceFieldBase,其为PotentialCell的子类。经典分子力场一般包含不同的能量项,常见的能量项有键长(bond length)、键角(bond angles)、二面角(dihedral angles)、库伦(Coulomb)相互作用、范德华(van der Waals)相互作用等。在mindsponge.potential.energy中内置了一些常用的能量项,可以将这些能量项以列表(list)的形式传入ForceFieldBase以对其进行初始化:

from mindsponge import ForceFieldBase
from mindsponge.potential.energy import *
bond_energy = BondEnergy(bond_index, force_constant=rk_init, bond_length=req_init)
angle_energy = AngleEnergy(angle_index, force_constant=tk_init, bond_angle=teq_init)
dihedral_energy = DihedralEnergy(dihedral_index, force_constant=pk_init, periodicity=pn_init, phase=phase_init)
ele_energy = CoulombEnergy(atom_charge=system.atom_charge)
vdw_energy = LennardJonesEnergy(sigma=sigma, epsilon=epsilon)
potential = ForceFieldBase(
    energy=[bond_energy,
            angle_energy,
            dihedral_energy,
            ele_energy,
            vdw_energy],
)

而对于成套的分子力场,也可以使用ForceFieldBase的子类ForceField类,通过传入体系模块和力场参数文件的方式进行初始化,例如建立一个AMBER FF14SB力场的势能函数:

    from mindsponge import ForceField
    potential = ForceField(system, 'AMBER.FF14SB')

这里的字符串’AMBER.FF14SB’本质上是一个YAML格式力场参数文件的文件名。MindSpore SPONGE读取力场参数的规则是,首先在当前目录下寻找是否存在一个同名文件,如果有则读取该文件,如果没有则会继续在mindsponge.data.forcefield目录中寻找并读取同名文件。

注意,如果使用ForceField类建立势能函数,那么其初始化过程中会对传入的体系(分子)模块system的原子类型(atom_type)等力场相关的参数进行重构。这是因为不同分子力场的原子类型(atom type)往往并不相同,而MindSpore SPONGE在使用体系模块进行建模时,并没有像绝大多数传统的MD模拟软件那样要求在建模前就指定分子力场。所以在很多情况下,体系(分子)模块在初始化的时候并没有设定原子类型(atom_type),所以在使用ForceField建立势能函数时,会根据体系(分子)模块中的原子名(atom_name)来自动设置原子类型(atom_type)、电荷(atom_charge)、质量(atom_mass)。

更新器(优化器)模块:Updater (Optimizer)

在MindSpore SPONGE中,更新器Updater用于在模拟过程中更新模拟体系的原子坐标。更新器Updater本质上其实是一个优化器(optimizer),其基类mindsponge.optimizer.Updater也是mindspore.nn.Optimizer的子类。实际上MindSpore的优化器,如Adam等可以直接作为MindSpore SPONGE的更新器使用,此时的模拟过程就相当于对模拟体系进行能量极小化(energy minimization)。

而MindSpore SPONGE的更新器(Updater)可以通过不同的“控制器”(controller)对模拟过程进行细致的调控。控制器(Controller)的基类(父类)是mindsponge.control.Controller,可以用来调节和更新模拟过程中体系的坐标(coordinate)、速度(velocity)、力(force)、能量(energy)、动能(kinetics)、维里(virial)、周期性盒子(PBC box)七个变量。常见的控制器(Controller)包括积分器(Integrator)、控温器(Thermostat)、控压器(Barostat)以及约束(Constraint)等。基类的更新器(Updater)可以通过接受一个控制器类的列表(list)进行初始化,而在模拟过程中会按照列表中控制器的顺序更新相应的参数:

from mindsponge import Updater
from mindsponge.control import VelocityVerlet
from mindsponge.control import Langevin
from mindsponge.control import BerendsenBarostat
from mindsponge.control import Lincs

integrator = VelocityVerlet(system, thermostat=Langevin(system, 300))

opt = Updater(
    system,
    controller=[
    integrator,
        Lincs(system, 'h-bonds’),
        BerendsenBarostat(system, 1)
    ],
velocity=velocity,
time_step=1e-3,
)

而对于一般的分子动力学模拟来说,则可以使用更新器(Updater)的子类mindsponge.control.MolecularDynamics来作为更新,直接设置积分器(integrator)、控温器(thermostat)、控压器(barostat)以及约束(constraint)来进行初始化:

from mindsponge import MolecularDynamics
from mindsponge.control import VelocityVerlet
from mindsponge.control import Langevin
from mindsponge.control import BerendsenBarostat
from mindsponge.control import Lincs

opt = MolecularDynamics (
    system,
    integrator=VelocityVerlet(system),
    thermostat=Langevin(system, 300),
    barostat=BerendsenBarostat(system, 1),
    constraint=Lincs(system, 'h-bonds'),
    velocity=velocity,
    time_step=1e-3,
)

核心模拟单元:SimulationCell

前面介绍过,除了直接用MindSpore SPONGE的顶层模块Sponge将三大功能模块,也就是体系(分子)模块、势能(力场)模块以及更新器(优化器)模块进行封装外,也可以先使用SimulationCell将体系(分子)模块、势能(力场)模块进行封装。除了封装两大模块外,SimulationCell在初始化时还可以设置近邻表(neighbour list)、偏向势(bias potential)以及势能后处理方式(wrapper)等:

from mindsponge import SimulationCell
from mindsponge.partition import NeighbourList
from mindsponge.potential.bias import MetaD
from mindsponge.core.wrapper import ReplicaExchange

neighbour_list = NeighbourList(system, cutoff=1, use_grids=True)

simulation = SimulationCell(
    system=system,
    potential=potential,
    neighbour_list=neighbour_list,
    bias=MetaDynamics(colvar=colvar),
    wrapper=ReplicaExchange(steps=100),
    )

回调函数:Callback

在AI框架中,回调函数(callback)用来进行一些不改变计算图的操作,例如监视和记录训练过程中损失函数的变化、验证训练过程中的参数是否达标、根据验证结果中止训练过程等。而在MindSpore SPONGE中,同样可以使用Callback来实现那些不影响模拟过程的操作,例如监视重要的模拟参数、记录模拟轨迹、计算模拟过程中的物理量和CVs(collective variables)等。

目前MindSpore SPONGE内置了两个Callback函数:

  • RunInfo:用于打印模拟过程的中间量,如能量、温度、压强等。

  • WriteH5MD:输出H5MD格式的模拟轨迹。

模拟轨迹文件:H5MD

MindSpore SPONGE采用H5MD作为记录模拟轨迹的默认文件格式。H5MD (HDF5 Molecular Data)是比利时布鲁塞尔自由大学的Pierre de Buyl博士等人于2014年提出的一种基于HDF5(Hierarchical Data Format 5)格式的MD模拟轨迹文件格式[1]。

德国斯图加特大学的Jonas Landsgesell和Sascha Ehrhardt和开发了一款VMD插件,安装后可以用VMD查看H5MD格式的模拟轨迹文件。但此插件存在Bug,且自2019以后就不再更新。我们Fork了原版仓库,修正了bug,并在原版程序的基础上进行了小幅修改,增加了坐标的单位转换等功能,并将默认的文件扩展名由.h5改为.h5md。此外,也可以使用MDAnalysis来读取H5MD文件的模拟轨迹信息。

得益于HDF5的多层级数据结构,H5MD除了可以记录模拟轨迹之外,本身也具有非常强大的可扩展性,可以记录模拟过程中的相关数据。MindSpore SPONGE将模拟过程中的势能(potential energy)、动能(kinetic energy)、温度(temperature)、压强(pressure)等信息记录在了H5MD文件的observables目录中。这些信息可以使用一些HDF5阅读器来进行查看,例如Silx Viewer

silx

使用教程

使用MindSpore SPONGE进行分子动力学模拟的教程可在MindScience仓库进行查看。

参考文献

[1] de Buyl, P.; Colberg, P. H.; Höfling, F. H5MD: A Structured, Efficient, and Portable File Format for Molecular Data [J]. Comput Phys Commun 2014, 185(6): 1546-1553.