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 Snapshotdef 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, HelperMethodsdef 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 SpongeSponge.controller.Step_Print_Initial("Phi", "%2f")Sponge.controller.Step_Print_Initial("Psi", "%2f")import pysagesfrom pysages.colvars import DihedralAnglefrom numpy import pifrom pysages.methods import Metadynamicsfrom pysages.backends.snapshot import Snapshot, SnapshotMethods, HelperMethods, build_data_querierimport numpy as npfrom jax import numpy as jnpfrom jax.dlpack import to_dlpack, from_dlpackfrom cupy import fromDlpack as cufdkB = 0.00831446261815324T = 300def 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()] height = 5.0# kJ/mol sigma = # 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 methodinitialized = Falsestate = Nonesnap = Nonehelper = Nonemethod = Noneupdate_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 snap = build_sponge_snapshot(num_atoms) helper = build_sponge_helper() method = phi_psi() res = method.build(snap, helper) state = res() update_func = res 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 osrecord_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) Sponge.controller.Step_Print("Psi", state.xi) with open(record_name, 'a+') as file: file.write("{},{}\n".format(state.xi, state.xi))其中用到的SPONGE配置文件为：
case1 MD simulationmode = NVTdefault_in_file_prefix = protein/aladpbc=0 cutoff=999dt = 1e-3step_limit = 5000write_information_interval = 10thermostat = middle_langevinmiddle_langevin_gamma = 10rst = nvt_restartcoordinate_in_file = protein/alad_coordinate.txtplugin = /usr/local/python-3.7.5/lib/python3.7/site-packages/prips/_prips.sopy = pysages_test.py原始的pdb文件为：
ATOM 1H1ACE 1 2.000 1.000-0.0001.000.00ATOM 2CH3 ACE 1 2.000 2.090 0.0001.000.00ATOM 3H2ACE 1 1.486 2.454 0.8901.000.00ATOM 4H3ACE 1 1.486 2.454-0.8901.000.00ATOM 5C ACE 1 3.427 2.641-0.0001.000.00ATOM 6O ACE 1 4.391 1.877-0.0001.000.00ATOM 7N ALA 2 3.555 3.970-0.0001.000.00ATOM 8H ALA 2 2.733 4.556-0.0001.000.00ATOM 9CAALA 2 4.853 4.614-0.0001.000.00ATOM 10HAALA 2 5.408 4.316 0.8901.000.00ATOM 11CBALA 2 5.661 4.221-1.2321.000.00ATOM 12HB1 ALA 2 5.123 4.521-2.1311.000.00ATOM 13HB2 ALA 2 6.630 4.719-1.2061.000.00ATOM 14HB3 ALA 2 5.809 3.141-1.2411.000.00ATOM 15C ALA 2 4.713 6.129 0.0001.000.00ATOM 16O ALA 2 3.601 6.653 0.0001.000.00ATOM 17N NME 3 5.846 6.835 0.0001.000.00ATOM 18H NME 3 6.737 6.359-0.0001.000.00ATOM 19CH3 NME 3 5.846 8.284 0.0001.000.00ATOM 20 HH31 NME 3 4.819 8.648 0.0001.000.00ATOM 21 HH32 NME 3 6.360 8.648 0.8901.000.00ATOM 22 HH33 NME 3 6.360 8.648-0.8901.000.00TER END 使用Xponge根据pdb文件进行建模的方法，可以参考这篇文章中的流程。需要注意的是，生成一个坐标文件之后，需要手动把坐标文件里面的Box信息（最后一行的前三列）修改为999（不加PBC Box的情况下），改完大概是这样的一个txt文件：
223.514000 3.000000 5.2886783.514000 4.090000 5.2886793.000000 4.454000 6.0210003.000000 4.454000 4.2410004.941000 4.641000 5.2886765.905000 3.877000 5.2886735.069000 5.970000 5.2886794.247000 6.556000 5.2886796.367000 6.614000 5.2886796.922000 6.316000 6.0210007.175000 6.221000 3.8990006.637000 6.521000 3.0000008.144000 6.719000 3.9250007.323000 5.141000 3.8900006.227000 8.129000 5.2886755.115000 8.653000 5.2886737.360000 8.835000 5.2886878.251000 8.359000 5.2886937.360000 10.284000 5.2886826.333000 10.648000 5.2886747.874000 10.648000 6.0210007.874000 10.648000 4.241000999 999 999 90.000000 90.000000 90.000000在确保SPONGE和PySAGES两者都正常安装的情况下，如果不加Meta，跑出来的CV轨迹是这样的：
加上Meta之后，CV轨迹是这样的：
可以看到，我们加上的Meta明显提升了MD的采样空间。
其中作图脚本如下：
import numpy as npimport matplotlib.pyplot as pltdef 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 delta_phi = phi_[:, None] - phi # (A, ) Z = -np.log(gaussian2(delta_psi, delta_phi, sigma1=sigma1, sigma2=sigma2, A=2.0)+1) return Zdata = np.genfromtxt('./cv_record.txt', delimiter=',')phi = data[:, 0]psi = data[:, 1]num_grids = 100num_levels = 10psi_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, phi_grids.shape)).TX,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 + 1in_deposition_step = (ncalls > 1) & (ncalls % stride == 1)当in_deposition_step=True时才会计算bias force。虽然这么算也没问题，但在有些软件里面使用的更新步数设定是取余数为0，那么可以对应的修改为：
ncalls = state.ncalls + 1in_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
页:
[1]