PySAGES结合CUDA SPONGE增强采样
本文探索并梳理了一下CUDA SPONGE高性能分子模拟采样软件，和PySAGES高性能增强采样软件，这两者强强联合的MD模拟新范式。
技术背景
在前面的一篇博客中，我们介绍过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')
特别注意
在使用PySAGES时有一点特别需要注意的是MetaDynamics方法更新的interval，原始更新步长是这样设定的：
ncalls = state.ncalls + 1
in_deposition_step = (ncalls > 1) & (ncalls % stride == 1)
当in_deposition_step=True时才会计算bias force。虽然这么算也没问题，但在有些软件里面使用的更新步数设定是取余数为0，那么可以对应的修改为：
ncalls = state.ncalls + 1
in_deposition_step = ncalls % stride == 0
可以根据自己的场景和需求来进行修改。
总结概要
本文探索并梳理了一下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
