改善嵌套循環性能

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) 

大寫字母的變量是不變的，在config.py文件中定義。 「currentPositions」是一個3乘以顆粒數矩陣。

我已經使用scipy.weave實現了嵌套for循環版本，該版本受此網站的啓發：http://www.scipy.org/PerformancePython。

但是，我不喜歡靈活性的喪失。我對「向量化」這個循環感興趣。我真的不明白這是如何工作的。任何人都可以給我一個線索或一個教這個教程的好教程嗎？

Answer 1

以下是我的代碼向量化版本。對於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 2

在純Python中寫入類似於MD引擎的東西會變得很慢。我會看看Numba（http://numba.pydata.org/）或Cython（http://cython.org/）。在用Cython的一面，我已經寫了使用用Cython簡單的朗之萬動力學引擎，可以作爲一個例子，讓你開始：

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

另一種選擇，我想了不少，就是用OpenMM 。有一個Python包裝，可以讓你把所有在一起的MD發動機件，實現自定義的力量等，也有針對GPU設備的能力：

https://simtk.org/home/openmm

雖然在一般，有有很多高度調整的MD代碼可供使用，除非您爲了某種通用教育目的而這樣做，否則從頭開始編寫自己的代碼是沒有意義的。一些主要的代碼，只是僅舉幾例：

• 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 3

爲了使這篇文章完整，我在C代碼中添加了我的實現。請注意，您需要導入編織和轉換器才能使用。此外，編織現在只適用於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