Question

我正在用Python进行氩气液体的分子动力学模拟。我有一个稳定版本运行，但它运行缓慢超过100个原子。我发现瓶颈是以下嵌套的for循环。这是一个强制计算，放在从我的main.py脚本调用的函数中：

def computeForce(currentPositions):
    potentialEnergy = 0
    force = zeros((NUMBER_PARTICLES,3))
    for iParticle in range(0,NUMBER_PARTICLES-1):
        for jParticle in range(iParticle + 1, NUMBER_PARTICLES):
            distance = currentPositions[iParticle] - currentPositions[jParticle]
            distance = distance - BOX_LENGTH * (distance/BOX_LENGTH).round()
            #note: this is so much faster than scipy.dot()
            distanceSquared = distance[0]*distance[0] + distance[1]*distance[1] + distance[2]*distance[2]            
            if distanceSquared < CUT_OFF_RADIUS_SQUARED:
                r2i = 1. / distanceSquared
                r6i = r2i*r2i*r2i
                lennardJones = 48. * r2i * r6i * (r6i - 0.5)
                force[iParticle] += lennardJones*distance
                force[jParticle] -= lennardJones*distance
                potentialEnergy += 4.* r6i * (r6i - 1.) - CUT_OFF_ENERGY
return(force,potentialEnergy)

CAPITAL字母中的变量是常量，在config.py文件中定义。 “currentPositions”是一个3个按粒子数矩阵。

我已经使用scipy.weave实现了嵌套for循环的一个版本，该版本的灵感来自这个网站：http://www.scipy.org/PerformancePython。

但是，我不喜欢失去灵活性。我对这个for循环“向量化”很感兴趣。我只是不知道它是如何工作的。任何人都可以给我一个线索或教导这个的好教程吗？

Answer 1

在纯python中编写像MD引擎这样的东西会很慢。我会看看Numba（http://numba.pydata.org/）或Cython（http://cython.org/）。在Cython方面，我使用cython编写了一个简单的Langevin Dynamics引擎，可以作为一个例子来帮助你入门：

https://bitbucket.org/joshua.adelman/pylangevin-integrator

我非常喜欢的另一个选择是使用OpenMM。有一个python包装器，允许您将MD引擎的所有部分组合在一起，实现自定义力等。它还具有针对GPU设备的能力：

https://simtk.org/home/openmm

总的来说，有很多高度调整的MD代码，除非你出于某种一般教育目的这样做，否则从头开始编写自己的代码是没有意义的。一些主要代码，仅举几例：

NAMD - http://www.ks.uiuc.edu/Research/namd/（免费）
Gromacs - http://www.gromacs.org/（免费）
琥珀 - http://ambermd.org/
LAMMPS - http://lammps.sandia.gov/（免费）
HOOMD - http://codeblue.umich.edu/hoomd-blue/（免费）
OpenMM - https://simtk.org/home/openmm（免费）
Charmm - http://www.charmm.org/

Answer 2

以下是我的代码的矢量化版本。对于1000点的数据集，我的代码大约比原始代码快50倍：

In [89]: xyz = 30 * np.random.uniform(size=(1000, 3))

In [90]: %timeit a0, b0 = computeForce(xyz)
1 loops, best of 3: 7.61 s per loop

In [91]: %timeit a, b = computeForceVector(xyz)
10 loops, best of 3: 139 ms per loop

代码：

from numpy import zeros

NUMBER_PARTICLES = 1000
BOX_LENGTH = 100
CUT_OFF_ENERGY = 1
CUT_OFF_RADIUS_SQUARED = 100

def computeForceVector(currentPositions):
    potentialEnergy = 0
    force = zeros((NUMBER_PARTICLES, 3))
    for iParticle in range(0, NUMBER_PARTICLES - 1):
        positionsJ =  currentPositions[iParticle + 1:, :]
        distance = currentPositions[iParticle, :] - positionsJ
        distance = distance - BOX_LENGTH * (distance / BOX_LENGTH).round()
        distanceSquared = (distance**2).sum(axis=1)
        ind = distanceSquared < CUT_OFF_RADIUS_SQUARED

        if ind.any():
            r2i = 1. / distanceSquared[ind]
            r6i = r2i * r2i * r2i
            lennardJones = 48. * r2i * r6i * (r6i - 0.5)
            ljdist = lennardJones[:, None] * distance[ind, :]
            force[iParticle, :] += (ljdist).sum(axis=0)
            force[iParticle+1:, :][ind, :] -= ljdist
            potentialEnergy += (4.* r6i * (r6i - 1.) - CUT_OFF_ENERGY).sum()
    return (force, potentialEnergy)

我还检查过代码会产生相同的结果

Answer 3

为了使这篇文章完整，我在C代码中编译了我的实现。请注意，您需要导入编织和转换器才能工作。而且，weave现在只适用于python 2.7。再次感谢所有的帮助！这比矢量化版本快10倍。

from scipy import weave
from scipy.weave import converters
def computeForceC(currentPositions):        
    code = """
    using namespace blitz;
    Array<double,1> distance(3);
    double distanceSquared, r2i, r6i, lennardJones;
    double potentialEnergy = 0.;

    for( int iParticle = 0; iParticle < (NUMBER_PARTICLES - 1); iParticle++){
        for( int jParticle = iParticle + 1; jParticle < NUMBER_PARTICLES; jParticle++){
            distance(0) = currentPositions(iParticle,0)-currentPositions(jParticle,0);
            distance(0) = distance(0) - BOX_LENGTH * round(distance(0)/BOX_LENGTH);
            distance(1) = currentPositions(iParticle,1)-currentPositions(jParticle,1);
            distance(1) = distance(1) - BOX_LENGTH * round(distance(1)/BOX_LENGTH);
            distance(2) = currentPositions(iParticle,2)-currentPositions(jParticle,2);
            distance(2) = distance(2) - BOX_LENGTH * round(distance(2)/BOX_LENGTH);
            distanceSquared = distance(0)*distance(0) + distance(1)*distance(1) + distance(2)*distance(2);
            if(distanceSquared < CUT_OFF_RADIUS_SQUARED){
                r2i = 1./distanceSquared;
                r6i = r2i * r2i * r2i;
                lennardJones = 48. * r2i * r6i * (r6i - 0.5);
                force(iParticle,0) += lennardJones*distance(0);
                force(iParticle,1) += lennardJones*distance(1);
                force(iParticle,2) += lennardJones*distance(2);
                force(jParticle,0) -= lennardJones*distance(0);
                force(jParticle,1) -= lennardJones*distance(1);
                force(jParticle,2) -= lennardJones*distance(2);
                potentialEnergy += 4.* r6i * (r6i - 1.)-CUT_OFF_ENERGY;

                }

            }//end inner for loop
    }//end outer for loop
    return_val = potentialEnergy;

    """
    #args that are passed into weave.inline and created inside computeForce
    #potentialEnergy = 0.
    force = zeros((NUMBER_PARTICLES,3))

    #all args
    arguments = ['currentPositions','force','NUMBER_PARTICLES','CUT_OFF_RADIUS_SQUARED','BOX_LENGTH','CUT_OFF_ENERGY']
    #evaluate stuff in code
    potentialEnergy = weave.inline(code,arguments,type_converters = converters.blitz,compiler = 'gcc')    

    return force, potentialEnergy

改善嵌套循环性能

3 个答案: