PySAGES结合CUDA SPONGE增强采样

技术背景

在前面的一篇博客中，我们介绍过PySAGES这个增强采样软件的基本安装和使用方法。该软件类似于Plumed是一个外挂增强采样软件，但是PySAGES是基于Python语言和Jax框架来实现的，在性能上有一定的优势。这里我们结合PySAGES的易开发特性，和CUDA SPONGE的高性能特性，做一个简单的扩展将二者联合起来进行分子动力学模拟与增强采样。

耦合框架

在前面的文章中我们介绍过SPONGE的Python接口以及调用模式：

这里再总结一个PySAGES的基本调用模式：

目前PySAGES已经集成了一些MD Backend，例如前面介绍过的OpenMM，以及基于Jax框架开发的Jax-MD等等。这些框架都可以通过Extension模块跟PySAGES进行对接，其中Jax-MD的Extension是在PySAGES中直接定义的，因为两者都基于Jax，底层兼容性较好。而OpenMM的Extension单独出来一个叫openmm_dlext的扩展包，这是因为其底层基于cupy开发，需要通过dlpack这个标准化工具在GPU上进行免拷贝操作。比较遗憾的是，目前MindSpore暂不支持dlpack。

PySAGES大概的模拟流程是这样的，先在MD Backend中定义好力场和积分器，将输入坐标传递给Extension再到PySAGES中构建SnapShot、SnapShot Method和Helper。然后在PySAGES中定义增强采样Method，例如MetaDynamics Method，传入构建好的SnapShot和Helper，就可以得到一个用于更新的函数Update Function和一个State增强采样状态参量。到这里PySAGES的初始化就结束了。然后在MD Backend中计算好Force，把相应的Force传给PySAGES进行更新，PySAGES的Update Function接收一个State和一个Snap就可以得到新的State，这个State中有一个bias变量，就是偏置势对应的Bias Force。把Bias Force直接加到MD Backend中传过来的Force中，就可以得到一个全新的Force用于积分器的迭代。

PySAGES与CUDA SPONGE

接上一个章节内容，这里需要特别说明一下CUDA SPONGE提供的Python接口与PySAGES的运用方法。因为CUDA SPONGE的调用形式是通过CUDA去调用Python的函数内容，而PySAGES的调用形式是从Python去控制特定Backend的CUDA函数的句柄，来实现二者的结合。一个以CUDA为主，一个以Python为主。那么要结合的话，使用CUDA Sponge自身的调用形式会相对容易一些，因为我们将PySAGES中仅当作一个用于计算Bias Force的模块，集成到SPONGE的运行流程中。当然，为了结合PySAGES，我们还是需要手动去定义一些SPONGE的SnapShot和Helper，首先是SnapShot：

from pysages.backends.snapshot import Snapshot

def build_sponge_snapshot(num_atoms):
    crd = jnp.array(np.random.random((num_atoms, 3)), jnp.float32)
    ids = jnp.arange(num_atoms)
    forces = jnp.zeros_like(crd)
    return Snapshot(crd, None, forces, ids, None, None, None)

这里因为是初始化，所以可以直接用一个jax的随机array来构造。其实也可以直接传一个backend的crd进来，但是为了区分initialization和update的功能，这里还是直接随机生成了一个。这里的SnapShot是一个namedtuple格式，存储了坐标和力等系统信息，这里为了最简化，我们只传入三个最基本的参数：坐标、力、原子索引。然后是SnapShotMethod和Helper，这两个标准化接口用于实时获取SnapShot中的信息：

from pysages.backends.snapshot import SnapshotMethods, HelperMethods

def build_sponge_snapshot_methods():
    def positions(snapshot):
        return snapshot.positions
    def indices(snapshot):
        return snapshot.ids
    return SnapshotMethods(positions, indices, None, None)

def build_sponge_helper(dims=3):
    def get_dims():
        return dims
    return HelperMethods(build_data_querier(build_sponge_snapshot_methods(), {"positions", "indices"}), get_dims)

这样我们就完成了一个基本的SPONGE的Extension的构建。

完整示例

这里是一个把CUDA SPONGE作为MD Backend，然后用PySAGES来进行增强采样的一个简单示例：

"""
SPONGE Usage:
    $ ../SPONGE -mdin nvt.txt
"""
import Sponge

Sponge.controller.Step_Print_Initial("Phi", "%2f")
Sponge.controller.Step_Print_Initial("Psi", "%2f")

import pysages
from pysages.colvars import DihedralAngle
from numpy import pi
from pysages.methods import Metadynamics
from pysages.backends.snapshot import Snapshot, SnapshotMethods, HelperMethods, build_data_querier

import numpy as np
from jax import numpy as jnp
from jax.dlpack import to_dlpack, from_dlpack
from cupy import fromDlpack as cufd

kB = 0.00831446261815324
T = 300

def build_sponge_snapshot_methods():
    def positions(snapshot):
        return snapshot.positions
    def indices(snapshot):
        return snapshot.ids
    return SnapshotMethods(positions, indices, None, None)

def build_sponge_snapshot(num_atoms):
    crd = jnp.array(np.random.random((num_atoms, 3)), jnp.float32)
    ids = jnp.arange(num_atoms)
    forces = jnp.zeros_like(crd)
    return Snapshot(crd, None, forces, ids, None, None, None)

def build_sponge_helper(dims=3):
    def get_dims():
        return dims
    return HelperMethods(build_data_querier(build_sponge_snapshot_methods(), {"positions", "indices"}), get_dims)

# 定义增强采样方法
def phi_psi():
    cvs = [DihedralAngle([4, 6, 8, 14]), DihedralAngle([6, 8, 14, 16])]
    height = 5.0  # kJ/mol
    sigma = [0.4, 0.4]  # radians
    stride = 3
    ngauss = 500
    grid = pysages.Grid(lower=(-pi, -pi), upper=(pi, pi), shape=(50, 50), periodic=True)
    method = Metadynamics(cvs, height, sigma, stride, ngauss, grid=grid)
    return method

initialized = False
state = None
snap = None
helper = None
method = None
update_func = None

# SPONGE用于更新Force的标准化调用
def Calculate_Force():
    global initialized
    global state
    global snap
    global helper
    global method
    global update_func

    if not initialized:
        initialized = True
        num_atoms = Sponge.md_info.frc.shape[0]
        snap = build_sponge_snapshot(num_atoms)
        helper = build_sponge_helper()
        method = phi_psi()
        res = method.build(snap, helper)
        state = res[1]()
        update_func = res[2]

    snap = snap._replace(positions=from_dlpack(Sponge.md_info.crd.toDlpack()))
    state = update_func(snap, state)
    # 加Meta
    Sponge.md_info.frc += cufd(to_dlpack(state.bias))
    # 不加Meta
    # Sponge.md_info.frc += 0

# 手动记录CV值
import os
record_name = "cv_record.txt"
if os.path.exists(record_name):
    os.remove(record_name)

# SPONGE的标准化打印输出接口
def Mdout_Print():
    global state
    global record_name
    Sponge.controller.Step_Print("Phi", state.xi[0][0])
    Sponge.controller.Step_Print("Psi", state.xi[0][1])
    with open(record_name, 'a+') as file:
        file.write("{},{}\n".format(state.xi[0][0], state.xi[0][1]))

其中用到的SPONGE配置文件为：

case1 MD simulation

mode = NVT
default_in_file_prefix = protein/alad

pbc=0
cutoff=999

dt = 1e-3
step_limit = 5000
write_information_interval = 10

thermostat = middle_langevin
middle_langevin_gamma = 10

rst = nvt_restart

coordinate_in_file = protein/alad_coordinate.txt
plugin = /usr/local/python-3.7.5/lib/python3.7/site-packages/prips/_prips.so
py = pysages_test.py

原始的pdb文件为：

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

使用Xponge根据pdb文件进行建模的方法，可以参考这篇文章中的流程。需要注意的是，生成一个坐标文件之后，需要手动把坐标文件里面的Box信息（最后一行的前三列）修改为999（不加PBC Box的情况下），改完大概是这样的一个txt文件：

22
3.514000 3.000000 5.288678
3.514000 4.090000 5.288679
3.000000 4.454000 6.021000
3.000000 4.454000 4.241000
4.941000 4.641000 5.288676
5.905000 3.877000 5.288673
5.069000 5.970000 5.288679
4.247000 6.556000 5.288679
6.367000 6.614000 5.288679
6.922000 6.316000 6.021000
7.175000 6.221000 3.899000
6.637000 6.521000 3.000000
8.144000 6.719000 3.925000
7.323000 5.141000 3.890000
6.227000 8.129000 5.288675
5.115000 8.653000 5.288673
7.360000 8.835000 5.288687
8.251000 8.359000 5.288693
7.360000 10.284000 5.288682
6.333000 10.648000 5.288674
7.874000 10.648000 6.021000
7.874000 10.648000 4.241000
999 999 999 90.000000 90.000000 90.000000

在确保SPONGE和PySAGES两者都正常安装的情况下，如果不加Meta，跑出来的CV轨迹是这样的：

加上Meta之后，CV轨迹是这样的：

可以看到，我们加上的Meta明显提升了MD的采样空间。

其中作图脚本如下：

import numpy as np
import matplotlib.pyplot as plt

def gaussian2(x1, x2, sigma1=1.0, sigma2=1.0, A=0.5):
    return np.sum(A*np.exp(-0.5*(x1**2/sigma1**2+x2**2/sigma2**2))/np.pi/sigma1/sigma2, axis=-1)

def potential_energy(position, psi, phi, sigma1, sigma2):
    # (A, )
    psi_, phi_ = position[:, 0], position[:, 1]
    # (A, R)
    delta_psi = psi_[:, None] - psi[None]
    delta_phi = phi_[:, None] - phi[None]
    # (A, )
    Z = -np.log(gaussian2(delta_psi, delta_phi, sigma1=sigma1, sigma2=sigma2, A=2.0)+1)
    return Z

data = np.genfromtxt('./cv_record.txt', delimiter=',')
phi = data[:, 0]
psi = data[:, 1]

num_grids = 100
num_levels = 10
psi_grids = np.linspace(-np.pi, np.pi, num_grids)
phi_grids = np.linspace(-np.pi, np.pi, num_grids)
grids = np.array(np.meshgrid(psi_grids, phi_grids)).T.reshape((-1, 2))

Z = potential_energy(grids, phi, psi, 1.0, 1.0).reshape((psi_grids.shape[0], phi_grids.shape[0])).T
X,Y = np.meshgrid(psi_grids, phi_grids)
levels = np.linspace(np.min(Z), np.max(Z), num_levels)

plt.figure()
plt.title("Biased MD Traj")
plt.xlabel(r'$\phi$')
plt.ylabel(r'$\psi$')
fc = plt.contourf(X, Y, Z, cmap='Greens', levels=levels)
plt.colorbar(fc)

plt.xlim(-np.pi, np.pi)
plt.ylim(-np.pi, np.pi)
plt.plot(phi, psi, 'o', alpha=0.4, color='red')
plt.savefig('meta.png')

总结概要

本文探索并梳理了一下CUDA SPONGE高性能分子模拟采样软件，和PySAGES高性能增强采样软件，这两者强强联合的MD模拟新范式。

版权声明

本文首发链接为：https://www.cnblogs.com/dechinphy/p/pysages-sponge.html

作者ID：DechinPhy

更多原著文章：https://www.cnblogs.com/dechinphy/

请博主喝咖啡：https://www.cnblogs.com/dechinphy/gallery/image/379634.html