CudaSPONGE之Python接口

技术背景

上一篇博客中我们介绍了CudaSPONGE的基本安装和使用方法。为了性能考虑,CudaSPONGE是基于纯CUDA C开发的,但是现在很多轮子都是Python开发的。为兼容更多的框架和平台,CudaSPONGE也提供了相应的Python API,方便Python开发者调用与二次开发。

接口逻辑

虽然安装和操作的过程并不复杂,但是这里面的交互逻辑还是得大概梳理一下。CudaSPONGE本身支持从plugin中调用几个固定的接口函数,如Calculate_Force()用于更新作用力,还有Mdout_Print()打印输出回调函数等等。调用的方式是通过动态链接库加载,也就是说,plugin的开发逻辑是先有一个python文件或者C语言文件,其中的API要跟CudaSPONGE对齐,然后编译成so动态链接库,供CudaSPONGE模拟的过程去调用,这是一个CudaSPONGE plugin开发的逻辑链条。

此外还有另外一个形式的plugin开发,可以参考本文的参考链接1中的内容,CudaSPONGE官方提供了一个prips插件,这个Python插件的逻辑是两头调用,本质上是对上述动态链接库接口的进一步封装。这就使得我们可以直接从Python文件中调用相应的接口函数,而不需要再编译成一个动态链接库文件,大大简化了Python Plugin开发的工作量。两种模式的差异如下图所示(非官方,个人理解):

prips安装与测试

prips插件支持pip直接安装:

bash 复制代码
$ python3 -m pip install prips
Looking in indexes: http://mirrors.aliyun.com/pypi/simple/
Collecting prips
  Downloading http://mirrors.aliyun.com/pypi/packages/d1/c0/35e829fb82fd6d4bcb5debd0a0fa7cfeec85325f9d015a2babb68123a3ee/prips-1.4.tar.gz (78 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 78.4/78.4 kB 1.9 MB/s eta 0:00:00
  Installing build dependencies ... done
  Getting requirements to build wheel ... done
  Installing backend dependencies ... done
  Preparing metadata (pyproject.toml) ... done
Building wheels for collected packages: prips
  Building wheel for prips (pyproject.toml) ... done
  Created wheel for prips: filename=prips-1.4-cp37-cp37m-linux_x86_64.whl size=4182966 sha256=dec62c0a31359dfda10b67f1cac94a4648e800284bc575014ad24f5984acb393
  Stored in directory: /root/.cache/pip/wheels/66/cd/8d/0fe470330380020b3d2395589216cb37387027eb687a203672
Successfully built prips
Installing collected packages: prips
Successfully installed prips-1.4

测试安装可以直接在命令行中执行:

bash 复制代码
$ python3 -c "import prips"

  PRIPS: Python Runtime Interface Plugin of SPONGE

Version: 1.4
Path: /usr/local/python-3.7.5/lib/python3.7/site-packages/prips/_prips.so

Error: 
    PRIPS replies on the python package "cupy".
    Please install cupy

这里发现少装了一个cupy,那就用pip装一个跟本地CUDA驱动匹配的cupy版本:

bash 复制代码
$ python3 -m pip install cupy-cuda11x
Looking in indexes: http://mirrors.aliyun.com/pypi/simple/
Collecting cupy-cuda11x
  Downloading http://mirrors.aliyun.com/pypi/packages/31/36/38a34d8bf2bcf9ac44be99c072e6a97bf882892ac506fa69cb70925a845f/cupy_cuda11x-11.6.0-cp37-cp37m-manylinux1_x86_64.whl (90.2 MB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 90.2/90.2 MB 3.8 MB/s eta 0:00:00
Requirement already satisfied: numpy<1.27,>=1.20 in /usr/local/python-3.7.5/lib/python3.7/site-packages (from cupy-cuda11x) (1.21.6)
Collecting fastrlock>=0.5 (from cupy-cuda11x)
  Using cached http://mirrors.aliyun.com/pypi/packages/42/4e/8bff5aa98ba1406c23a7dded13fea0bf2f536b4f8f7096fcbea0303e9cf5/fastrlock-0.8.2-cp37-cp37m-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_24_x86_64.whl (47 kB)
Installing collected packages: fastrlock, cupy-cuda11x
Successfully installed cupy-cuda11x-11.6.0 fastrlock-0.8.2

再次运行测试:

bash 复制代码
$ python3 -c "import prips"

  PRIPS: Python Runtime Interface Plugin of SPONGE

Version: 1.4
Path: /usr/local/python-3.7.5/lib/python3.7/site-packages/prips/_prips.so

Usage:
    1. Copy the path printed above
    2. Paste it to the value of the command "plugin" of SPONGE

这表示安装成功了,并且给出了具体的动态链接库地址,方便我们直接把这个动态链接库地址拷贝到CudaSPONGE运行的mdin配置文件中。

CudaSPONGE-pyplugin测试

如果上述章节显示安装成功,并且本地已经配置好CudaSPONGE的环境,那就可以开始测试简单案例了。例如使用上一篇博客中的能量极小化的案例,来一个基础的CudaSPONGE版本的hello world

txt 复制代码
case1 Minimization

mode = Minimization 
minimization_dynamic_dt 1
default_in_file_prefix = protein/case1

pbc=0 
cutoff=999

dt = 1e-02
step_limit = 500
write_information_interval = 50

rst = restart

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

其中test.py的文件内容为:

python 复制代码
import Sponge

def Mdout_Print():
    print("Hellow SPONGE World!")

直接使用$ ../SPONGE -mdin minimize.txt命令行运行结果输出为:

txt 复制代码
------------------------------------------------------------------------------------------------------------
Hellow SPONGE World!
           step =              50,            time =           0.000,     temperature =            0.00, 
      potential =       120784.66,              LJ =          191.83,         Coulomb =         -134.70, 
        nb14_LJ =       119114.87,         nb14_EE =          -20.69,            bond =            4.54, 
          angle =         1535.60,        dihedral =           93.22, 
------------------------------------------------------------------------------------------------------------
Hellow SPONGE World!
           step =             100,            time =           0.000,     temperature =            0.00, 
      potential =        47867.34,              LJ =          192.20,         Coulomb =         -134.83, 
        nb14_LJ =        46200.27,         nb14_EE =          -20.75,            bond =            8.92, 
          angle =         1528.36,        dihedral =           93.18, 
------------------------------------------------------------------------------------------------------------
Hellow SPONGE World!
           step =             150,            time =           0.000,     temperature =            0.01, 
      potential =        21746.57,              LJ =          192.48,         Coulomb =         -134.96, 
        nb14_LJ =        20079.19,         nb14_EE =          -20.75,            bond =           15.94, 
          angle =         1521.55,        dihedral =           93.14, 
------------------------------------------------------------------------------------------------------------
Hellow SPONGE World!
           step =             200,            time =           0.000,     temperature =            0.02, 
      potential =        11373.98,              LJ =          192.42,         Coulomb =         -135.10, 
        nb14_LJ =         9703.54,         nb14_EE =          -20.71,            bond =           26.00, 
          angle =         1514.73,        dihedral =           93.10, 
------------------------------------------------------------------------------------------------------------
Hellow SPONGE World!
           step =             250,            time =           0.000,     temperature =            0.05, 
      potential =         7180.14,              LJ =          191.38,         Coulomb =         -135.23, 
        nb14_LJ =         5504.79,         nb14_EE =          -20.65,            bond =           39.29, 
          angle =         1507.51,        dihedral =           93.06, 
------------------------------------------------------------------------------------------------------------
Hellow SPONGE World!
           step =             300,            time =           0.000,     temperature =            0.13, 
      potential =         5464.48,              LJ =          187.58,         Coulomb =         -135.35, 
        nb14_LJ =         3785.66,         nb14_EE =          -20.68,            bond =           55.14, 
          angle =         1499.09,        dihedral =           93.04, 
------------------------------------------------------------------------------------------------------------
Hellow SPONGE World!
           step =             350,            time =           0.000,     temperature =            0.32, 
      potential =         4587.00,              LJ =          177.15,         Coulomb =         -135.43, 
        nb14_LJ =         2914.02,         nb14_EE =          -20.95,            bond =           71.93, 
          angle =         1487.24,        dihedral =           93.03, 
------------------------------------------------------------------------------------------------------------
Hellow SPONGE World!
           step =             400,            time =           0.000,     temperature =            0.63, 
      potential =         3817.87,              LJ =          154.56,         Coulomb =         -135.67, 
        nb14_LJ =         2169.29,         nb14_EE =          -21.40,            bond =           92.39, 
          angle =         1465.65,        dihedral =           93.05, 
------------------------------------------------------------------------------------------------------------
Hellow SPONGE World!
           step =             450,            time =           0.000,     temperature =            0.96, 
      potential =         2990.47,              LJ =          119.14,         Coulomb =         -136.70, 
        nb14_LJ =         1380.29,         nb14_EE =          -21.51,            bond =          136.01, 
          angle =         1420.15,        dihedral =           93.09, 
------------------------------------------------------------------------------------------------------------
Hellow SPONGE World!
           step =             500,            time =           0.000,     temperature =            1.11, 
      potential =         2276.00,              LJ =           82.45,         Coulomb =         -139.29, 
        nb14_LJ =          725.04,         nb14_EE =          -20.48,            bond =          215.85, 
          angle =         1319.28,        dihedral =           93.16, 
------------------------------------------------------------------------------------------------------------

可以看到,在每一个打印环节都会调用py文件中的打印内容。如果需要看帮助文档,可以将test.py的内容修改为:

python 复制代码
import Sponge
help(Sponge.controller)
help(Sponge.cv_controller)
help(Sponge.md_info)

但是这里我建议要看文档还是直接进Gitee仓库直接看。关于这里面可以访问的force的数据类型,我们也可以打印出来看一下:

python 复制代码
import Sponge

Sponge.controller.Step_Print_Initial("Force_TYPE", "%s")

def Mdout_Print():
    Sponge.controller.Step_Print("Force_TYPE", type(Sponge.md_info.frc))

输出内容为:

txt 复制代码
------------------------------------------------------------------------------------------------------------
           step =             500,            time =           0.000,     temperature =            1.11, 
      potential =         2276.00,      Force_TYPE = <class 'cupy.ndarray'>,              LJ =           82.45, 
        Coulomb =         -139.29,         nb14_LJ =          725.04,         nb14_EE =          -20.48, 
           bond =          215.85,           angle =         1319.28,        dihedral =           93.16, 
------------------------------------------------------------------------------------------------------------

可以看到是一个封装好的cupy的数组类型。那么我们也可以查看相应参量的Shape:

python 复制代码
import Sponge

Sponge.controller.Step_Print_Initial("Force_Shape_0", "%d")
Sponge.controller.Step_Print_Initial("Force_Shape_1", "%d")

def Mdout_Print():
    Sponge.controller.Step_Print("Force_Shape_0", Sponge.md_info.frc.shape[0])
    Sponge.controller.Step_Print("Force_Shape_1", Sponge.md_info.frc.shape[1])

输出内容大概是这样的:

txt 复制代码
------------------------------------------------------------------------------------------------------------
           step =              50,            time =           0.000,     temperature =            0.00, 
      potential =       120784.66,   Force_Shape_0 =              57,   Force_Shape_1 =               3, 
             LJ =          191.83,         Coulomb =         -134.70,         nb14_LJ =       119114.87, 
        nb14_EE =          -20.69,            bond =            4.54,           angle =         1535.60, 
       dihedral =           93.22, 
------------------------------------------------------------------------------------------------------------

这个Shape也就是我们输入的分子体系的Shape了。

为了方便查看结果,我们把mdin改成单步的优化:

txt 复制代码
case1 Minimization

mode = Minimization 
minimization_dynamic_dt 1
default_in_file_prefix = protein/case1

pbc=0 
cutoff=999

dt = 1e-02
step_limit = 1
write_information_interval = 1

rst = restart

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

然后输出一个force求和的数据结果:

python 复制代码
import Sponge

Sponge.controller.Step_Print_Initial("Force_SUM", "%2f")

def Mdout_Print():
    Sponge.controller.Step_Print("Force_SUM", Sponge.md_info.frc.sum())

输出结果为:

txt 复制代码
------------------------------------------------------------------------------------------------------------
           step =               1,            time =           0.000,     temperature =            0.00, 
      potential =       424228.34,       Force_SUM =       -0.258057,              LJ =          191.39, 
        Coulomb =         -134.55,         nb14_LJ =       422551.28,         nb14_EE =          -20.52, 
           bond =            2.78,           angle =         1544.69,        dihedral =           93.28, 
------------------------------------------------------------------------------------------------------------

可以看到正常大概是在0.x这个数量级,如果我们对这个force进行操作,使用Calculate_Force()函数将其放大100倍,再看看结果:

python 复制代码
import Sponge

Sponge.controller.Step_Print_Initial("Force_SUM", "%2f")

def Calculate_Force():
    Sponge.md_info.frc *= 100

def Mdout_Print():
    Sponge.controller.Step_Print("Force_SUM", Sponge.md_info.frc.sum())

输出内容为:

txt 复制代码
------------------------------------------------------------------------------------------------------------
           step =               1,            time =           0.000,     temperature =            0.00, 
      potential =       424228.34,       Force_SUM =      -30.750000,              LJ =          191.39, 
        Coulomb =         -134.55,         nb14_LJ =       422551.28,         nb14_EE =          -20.52, 
           bond =            2.78,           angle =         1544.69,        dihedral =           93.28, 
------------------------------------------------------------------------------------------------------------

可以看到这是一个被放大100倍之后的结果,数量级已经不一样了。关于CudaSPONGE-python调用的案例就先介绍这么多,在这个基础上去扩展一些Force Wrapper的应用应该是很容易的。

总结概要

本文介绍了高性能GPU分子动力学模拟软件CudaSPONGE的Python API接口,通过官方开发的prips插件,使得我们可以在Python框架下很方便的开发一些分子动力学模拟的Force Wrapper,例如Meta Dynamics中就有很多可以外界的工具,非常方便开发者的二次开发,同时又能够兼顾到性能。

版权声明

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

作者ID:DechinPhy

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

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

参考链接

  1. https://gitee.com/gao_hyp_xyj_admin/prips