5

我正在 Python 中对 Argon 液体进行分子动力学模拟。我有一个稳定版本正在运行，但是它运行缓慢，超过 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)

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

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

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

4

3 回答 3

4

以下是我的代码的矢量化版本。对于具有 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)

我还检查了代码是否产生相同的结果

于 2013-02-13T15:35:40.590 回答
3

用纯 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 代码可用，除非您出于某种一般教育目的这样做，否则从头开始编写自己的代码是没有意义的。一些主要代码，仅举几例：

于 2013-02-13T14:34:03.393 回答
1

只是为了使这篇文章完整，我在 C 代码中添加了我的实现。请注意，您需要导入 weave 和转换器才能使其正常工作。此外，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
于 2013-02-15T08:13:08.833 回答