技术背景

在传统的分子动力学模拟软件中，对于分子体系的控制，例如控制体系的相对位置亦或是绝对位置，通常都是通过施加一些约束算法来实现的。例如用于限制化学键的LINCS算法，又比如水分子体系非常常用的SETTLE约束算法，这两种算法都属于Constraint（硬约束）。除此之外还有很多Restraint（软约束），例如施加谐振势等等，在MindSponge中有相关的实现：sponge.potential.OscillatorBias。

得益于Python的灵活性，使用MindSponge框架可以自己继承control里面的父类，实现一些自己的体系控制算法。

体系建模

首先我们要构建一个用于模拟的体系，这里只是为了演示，我们用一个简单的体系：丙氨酸二肽加上一个水分子。那么如果只有一个丙氨酸二肽的pdb文件，怎么加水分子呢？其实用MindSponge就可以直接完成这种简单的建模，原始的pdb文件是这样的：

txt 复制代码

ATOM      1  H1  ACE     1       2.000   1.000  -0.000  1.00  0.00
ATOM      2  CH3 ACE     1       2.000   2.090   0.000  1.00  0.00
ATOM      3  H2  ACE     1       1.486   2.454   0.890  1.00  0.00
ATOM      4  H3  ACE     1       1.486   2.454  -0.890  1.00  0.00
ATOM      5  C   ACE     1       3.427   2.641  -0.000  1.00  0.00
ATOM      6  O   ACE     1       4.391   1.877  -0.000  1.00  0.00
ATOM      7  N   ALA     2       3.555   3.970  -0.000  1.00  0.00
ATOM      8  H   ALA     2       2.733   4.556  -0.000  1.00  0.00
ATOM      9  CA  ALA     2       4.853   4.614  -0.000  1.00  0.00
ATOM     10  HA  ALA     2       5.408   4.316   0.890  1.00  0.00
ATOM     11  CB  ALA     2       5.661   4.221  -1.232  1.00  0.00
ATOM     12  HB1 ALA     2       5.123   4.521  -2.131  1.00  0.00
ATOM     13  HB2 ALA     2       6.630   4.719  -1.206  1.00  0.00
ATOM     14  HB3 ALA     2       5.809   3.141  -1.241  1.00  0.00
ATOM     15  C   ALA     2       4.713   6.129   0.000  1.00  0.00
ATOM     16  O   ALA     2       3.601   6.653   0.000  1.00  0.00
ATOM     17  N   NME     3       5.846   6.835   0.000  1.00  0.00
ATOM     18  H   NME     3       6.737   6.359  -0.000  1.00  0.00
ATOM     19  CH3 NME     3       5.846   8.284   0.000  1.00  0.00
ATOM     20 HH31 NME     3       4.819   8.648   0.000  1.00  0.00
ATOM     21 HH32 NME     3       6.360   8.648   0.890  1.00  0.00
ATOM     22 HH33 NME     3       6.360   8.648  -0.890  1.00  0.00
TER   
END

然后我们用MindSponge框架给这个构象加水、优化初始构象：

python 复制代码

from sponge import Protein, ForceField, Sponge
from sponge.optimizer import SteepestDescent
from sponge.callback import RunInfo, SaveLastPdb
# 加载丙氨酸二肽的多肽构象文件
mol_file = '../pdb/alad.pdb'
# 体系建模
mol = Protein(mol_file)
# 只加一个水分子
mol.fill_water(edge = 0.2, num_water=1, template='water.spce.yaml')
# 力场建模
potential = ForceField(mol, parameters=['amber.ff14sb', 'spce'], use_pme=False)
# 定义梯度下降优化算法
opt = SteepestDescent(mol.trainable_params(), 1e-6)
# 体系封装
mini = Sponge(mol, potential, opt)
# 两个回调函数，一个用于打印信息，一个用于存储pdb构象文件
run_info = RunInfo(1)
save_pdb = SaveLastPdb(mol, save_freq=1, pdb_name='alad_water_opt.pdb')
# 运行10个step
mini.run(10, callbacks=[run_info, save_pdb])

运行结束之后，会在当前路径下生成一个alad_water_opt.pdb文件：

txt 复制代码

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

可以用vmd看一下这个MindSponge加水并且优化之后的体系构象：

体系控制

这里我们以展示功能为主，看看MindSponge框架现在可以对分子体系进行何等程度的控制。

基于力的控制

像前面提到的LINCS约束算法或者是SETTLE约束算法，虽然最终实现的是对分子体系坐标的约束效果，但是其本质上是对体系的作用力进行约束，使得体系在约束后的作用力之下，遵循设计好的路线去运动到指定的坐标位置。在MindSponge中要添加LINCS约束算法或者SETTLE算法是非常容易的：

python 复制代码

from sponge import UpdaterMD
from sponge.control import SETTLE, Lincs
...
updater = UpdaterMD(system, time_step=2e-3, integrator=LeapFrog(system), constraint=[Lincs(system), SETTLE(system)])
...

直接把这两者加到优化器中即可。

另外还有一种直接对作用力进行操作的方法，我们可以在MindSponge中定义一个WithForceCell，然后向其中传入一个"作用力修改器"，参数名为modifier，即可完成对作用力的动态调整：

python 复制代码

import mindspore as ms
from mindspore import Tensor
from sponge.core.simulation.force import WithForceCell
from sponge.sampling import MaskedDriven
...
# 这里mask的作用是把所有原子标记为0，最后3个原子标记为1
mask = np.zeros((system.coordinate.shape[0], system.coordinate.shape[1]))
mask[:, -3:] += 1
mask = Tensor(mask.astype(np.int32), ms.int32)
# 把mask乘到Force里面去，直接修改作用力
modifier = MaskedDriven(mask=mask, update_pace=1)
with_force = WithForceCell(system, modifier=modifier)

在上述案例中，经过modifier的修饰，只有最后3个原子有作用力，其他原子的作用力都是0，这就是基于力的控制。

基于坐标的控制

MindSponge提供了Constraint的接口，我们可以直接基于这个接口去改变体系中的原子坐标。像前面提到的SETTLE和LINCS约束算法，本质上都是基于Constraint基类实现的约束算法。那么除了像上一个章节的内容那样去控制原子作用力，我们还可以直接操纵原子坐标。例如我们可以基于Constraint实现一个基于mask直接控制原子坐标的模块：

python 复制代码

class MaskedConstraint(Constraint):
    def __init__(self, system, mask, **kwargs):
        super().__init__(system, **kwargs)
        # 传入mask
        if not isinstance(mask, Tensor):
            self.mask = Tensor(mask, ms.int32)
        else:
            self.mask = mask
        # 反向mask
        self.inv_mask = 1 - self.mask
        # 记录原始坐标
        self.masked_crd = system.coordinate

    def construct(self, coordinate, velocity, force, energy, virial = None, pbc_box = None, step = 0, **kwargs):
        # 只对指定mask的原子更新坐标
        coordinate += self.inv_mask[..., None] * (self.masked_crd - coordinate)
        # mask之外的原子速度设置为0
        velocity *= self.mask[..., None]
        return {'coordinate': coordinate,
                'velocity': velocity,
                'force': force,
                'energy': energy,
                'virial': virial,
                'pbc_box': pbc_box,
                }
...
updater = UpdaterMD(system, time_step=2e-3, integrator=LeapFrog(system), constraint=[onsite_constraint, SETTLE(system)])
...

在这个我们自己实现的MaskedConstraint模块中，我们同时控制了坐标和速度，而最终的目标也是确保被mask的原子保持不动。这里需要说明一下，同样是保持原子位置不变，那么基于力的控制和基于坐标的控制有什么区别呢？其实可以思考一下朗之万动力学，其速度不仅受到加速度（也就是作用力）的影响，还受到一个与温度有关的随机作用力、以及跟速度相关的粘滞作用力的影响。因此，单纯的控制作用力，是无法达到控制被mask原子保持原位不动的操作的。所以，如果切实需要实现这样的一种场景，还是需要从Constraint硬约束的角度着手。

基于反应坐标的控制

当我们对更加广义的反应坐标进行操控时，就已经进入到了增强采样的范畴。通常来说，增强采样改变的是势能面，进而改变作用力，最后影响到真实的原子坐标上，本质上也是一种软约束。但是跟普通的约束算法不同的是，常用的约束条件都是在缩小采样子空间，而增强采样是在扩大采样子空间（个人见解）。

因为增强采样要基于反应坐标（Collective Variable, CV），所以我们需要先定义好反应坐标。我们在MindSponge中已经实现了一些常用的反应坐标，例如距离、角度、二面角和质心等等，这里我们演示一个案例是基于质心坐标的三维增强采样。因为这里使用到的MetaDynamics元动力学增强采样算法，只能接收一维的输入，因此我们需要手动实现一个CV把质心坐标的CV分成三个轴，分别构建独立的增强采样：

python 复制代码

class IndexedCenter(Center):
    def __init__(self, 
                 atoms, 
                 mass, 
                 index,
                 batched = False, 
                 keep_in_box = False, 
                 keepdims = False, 
                 axis = -2, 
                 name = 'atoms_center'):
        super().__init__(atoms, mass, batched, keep_in_box, keepdims, axis, name)
        # 质心轴的索引
        self.index = index
        self._set_shape((1, ))
    
    def construct(self, coordinate: Tensor, pbc_box: Tensor = None):
        res = super().construct(coordinate, pbc_box)
        # 输出格式要keepdim
        return res[..., [self.index]]

自己实现的方法很简单，我们只需要继承质心的计算模块，然后加一个Index索引，分别去索引xyz三个轴的坐标值就可以了：

python 复制代码

...
# 指定编号为[22, 23, 24]的原子计算质心，质量分别为[16., 1.01, 1.01]
cv_x = IndexedCenter([22, 23, 24], [16., 1.01, 1.01], 0, axis=-1)
cv_y = IndexedCenter([22, 23, 24], [16., 1.01, 1.01], 1, axis=-1)
cv_z = IndexedCenter([22, 23, 24], [16., 1.01, 1.01], 2, axis=-1)
# 配置Meta参数
pace = 10
height = 2.
sigma = 0.05
bins = 50
# 定义三维的MetaDynamics
metad = Metadynamics(
    colvar=[cv_x, cv_y, cv_z],
    update_pace=pace,
    height=height,
    sigma=sigma,
    grid_min=[-1.2, -1, -1.4],
    grid_max=[2.0, 2.1, 1.5],
    grid_bin=bins,
    temperature=temp,
)
# 定义Meta的边界墙约束
lwall = LowerWall([cv_x, cv_y, cv_z], [-0.9, -0.7, -1.1])
uwall = UpperWall([cv_x, cv_y, cv_z], [1.7, 1.8, 1.2])
# 构建分子模拟的操作对象
with_energy = WithEnergyCell(system, potential, bias=[metad, lwall, uwall])
one_step = RunOneStepCell(energy=with_energy, optimizer=updater)
...

经过这个增强采样，对应的质心坐标会在给定的范围内不断的波动，我们也可以通过reweight来重现真实的自由能面。

其他控制

在统计力学中，恒温恒压可以对应于不同的统计系综，对应于分子动力学模拟中的温度控制模块和压强控制模块。虽然在MindSponge的实现中我们把控温控压模块分别作为参数thermostat和barostat传入到迭代器中，但实际上这两种控制器也可以认为是对系统的控制对象，其在MindSponge中的定义和使用也是比较简单的：

python 复制代码

temp = 300
thermostat = BerendsenThermostat(system, temp)
updater = UpdaterMD(system, time_step=2e-3, integrator=LeapFrog(mol), thermostat=thermostat, 
                    constraint=[onsite_constraint, SETTLE(mol)])

完整示例

针对于前面提到的一些控制模块，我们做一个完整的示例，用来演示各种控制器的实现方法：

python 复制代码

import numpy as np
import mindspore as ms
from mindspore import Tensor
ms.set_context(mode=ms.GRAPH_MODE, device_target='GPU')

from sponge import Protein, ForceField, Sponge, UpdaterMD, WithEnergyCell, RunOneStepCell
from sponge.core.simulation.force import WithForceCell
from sponge.control import LeapFrog, BerendsenThermostat, Constraint, SETTLE
from sponge.colvar import Center
from sponge.callback import RunInfo, WriteH5MD
from sponge.potential import LowerWall, UpperWall
from sponge.sampling import Metadynamics, MaskedDriven

# 定义CV
class IndexedCenter(Center):
    def __init__(self, 
                 atoms, 
                 mass, 
                 index,
                 batched = False, 
                 keep_in_box = False, 
                 keepdims = False, 
                 axis = -2, 
                 name = 'atoms_center'):
        super().__init__(atoms, mass, batched, keep_in_box, keepdims, axis, name)
        self.index = index
        self._set_shape((1, ))
    
    @ms.jit
    def construct(self, coordinate: Tensor, pbc_box: Tensor = None):
        res = super().construct(coordinate, pbc_box)
        return res[..., [self.index]]

# 定义坐标控制器
class MaskedConstraint(Constraint):
    def __init__(self, system, mask, **kwargs):
        super().__init__(system, **kwargs)
        if not isinstance(mask, Tensor):
            self.mask = Tensor(mask, ms.int32)
        else:
            self.mask = mask
        self.inv_mask = 1 - self.mask
        self.masked_crd = system.coordinate

    def construct(self, coordinate, velocity, force, energy, virial = None, pbc_box = None, step = 0, **kwargs):
        coordinate += self.inv_mask[..., None] * (self.masked_crd - coordinate)
        velocity *= self.mask[..., None]
        return {'coordinate': coordinate,
                'velocity': velocity,
                'force': force,
                'energy': energy,
                'virial': virial,
                'pbc_box': pbc_box,
                }

# 输入分子文件
mol_file = './alad_water_opt.pdb'
mol = Protein(mol_file, template=['protein0.yaml', 'water.spce.yaml'])
potential = ForceField(mol, parameters=['amber.ff14sb', 'spce'], use_pme=False)
# 最后三个原子是一个水分子，我们要控制的CV就是这个水分子的质心
mask = np.zeros((mol.coordinate.shape[0], mol.coordinate.shape[1]))
mask[:, -3:] += 1
mask = Tensor(mask.astype(np.int32), ms.int32)

temp = 300
thermostat = BerendsenThermostat(mol, temp)
onsite_constraint = MaskedConstraint(mol, mask)
updater = UpdaterMD(mol, time_step=2e-3, integrator=LeapFrog(mol), thermostat=thermostat, 
                    constraint=[onsite_constraint, SETTLE(mol)])

run_info = RunInfo(10)
cb_h5md = WriteH5MD(mol, 'alad_water.h5md', save_freq=10, write_image=False)

modifier = MaskedDriven(mask=mask, update_pace=1)

cv_x = IndexedCenter([22, 23, 24], [16., 1.01, 1.01], 0, axis=-1)
cv_y = IndexedCenter([22, 23, 24], [16., 1.01, 1.01], 1, axis=-1)
cv_z = IndexedCenter([22, 23, 24], [16., 1.01, 1.01], 2, axis=-1)

pace = 10
height = 2.
sigma = 0.05
bins = 50

metad = Metadynamics(
    colvar=[cv_x, cv_y, cv_z],
    update_pace=pace,
    height=height,
    sigma=sigma,
    grid_min=[-1.2, -1, -1.4],
    grid_max=[2.0, 2.1, 1.5],
    grid_bin=bins,
    temperature=temp,
)

lwall = LowerWall([cv_x, cv_y, cv_z], [-0.9, -0.7, -1.1])
uwall = UpperWall([cv_x, cv_y, cv_z], [1.7, 1.8, 1.2])

with_energy = WithEnergyCell(mol, potential, bias=[metad, lwall, uwall])
with_force = WithForceCell(mol, modifier=modifier)
one_step = RunOneStepCell(energy=with_energy, force=with_force, optimizer=updater)
md = Sponge(network=one_step, metrics={'x': cv_x, 'y': cv_y, 'z': cv_z})
md.run(100, callbacks=[run_info, cb_h5md])

运行结果如下所示：

bash 复制代码

[MindSPONGE] The settle constraint is used for the molecule system.
[MindSPONGE] Started simulation at 2024-05-24 14:55:27
[MindSPONGE] Step: 10, E_pot: -120.77967, E_kin: 25.428715, E_tot: -95.35095, Temperature: 88.648445, x: 1.0543859, y: 1.178435, z: 0.931545
[MindSPONGE] Step: 20, E_pot: -125.460464, E_kin: 23.364426, E_tot: -102.09604, Temperature: 81.45201, x: 1.0473968, y: 1.1388566, z: 0.9233233
[MindSPONGE] Step: 30, E_pot: -105.13234, E_kin: 6.2417407, E_tot: -98.8906, Temperature: 21.759678, x: 1.0330578, y: 1.0696822, z: 1.0379511
[MindSPONGE] Step: 40, E_pot: -125.4366, E_kin: 22.289291, E_tot: -103.14731, Temperature: 77.70393, x: 1.0292426, y: 0.9710549, z: 0.918953
[MindSPONGE] Step: 50, E_pot: -119.97059, E_kin: 16.77571, E_tot: -103.19488, Temperature: 58.482727, x: 1.0150183, y: 1.0174756, z: 0.9012259
[MindSPONGE] Step: 60, E_pot: -117.3958, E_kin: 17.552721, E_tot: -99.84308, Temperature: 61.191513, x: 0.9964286, y: 1.12876, z: 0.9246655
[MindSPONGE] Step: 70, E_pot: -117.40674, E_kin: 19.596474, E_tot: -97.810265, Temperature: 68.31635, x: 0.9824194, y: 1.1459445, z: 0.9315006
[MindSPONGE] Step: 80, E_pot: -126.24636, E_kin: 26.27534, E_tot: -99.97102, Temperature: 91.59991, x: 0.9752382, y: 1.0937127, z: 0.89053756
[MindSPONGE] Step: 90, E_pot: -124.537186, E_kin: 19.668875, E_tot: -104.86831, Temperature: 68.56875, x: 0.9713231, y: 1.0286477, z: 0.80201
[MindSPONGE] Step: 100, E_pot: -120.2195, E_kin: 17.595142, E_tot: -102.62436, Temperature: 61.339397, x: 0.95973897, y: 1.0473468, z: 0.80481935
[MindSPONGE] Finished simulation at 2024-05-24 14:55:47
[MindSPONGE] Simulation time: 20.34 seconds.
--------------------------------------------------------------------------------

运行结束之后可以在当前路径下生成一个h5md轨迹文件，可以用VMD进行可视化。可以看到，在原位Constraint的作用下，丙氨酸二肽一致保持原位不动。而水分子在全原子力场和SETTLE约束算法的双重作用下，在给定的范围内不停的旋转和平移：