Question

我有一个数组（N = 10 ^ 4），我需要找到每两个条目之间的差异（计算给定原子坐标的电位）这是我在纯python中编写的代码，但它真的没有效果，有人能告诉我如何加快速度吗？（使用numpy或编织）。这里x，y是原子坐标数组（只是简单的1D数组）

def potential(r):
   U = 4.*(np.power(r,-12) - np.power(r,-6))
   return U
def total_energy(x):
   E = 0.
   #need to speed up this part
   for i in range(N-1):
     for j in range(i):
         E += potential(np.sqrt((x[i]-x[j])**2))  
return E

Answer 1

首先你可以使用数组算术：

def potential(r):
    return 4.*(r**(-12) - r**(-6))

def total_energy(x):
    E = 0.
    for i in range(N-1):
        E += potential(np.sqrt((x[i]-x[:i])**2)).sum()
    return E

或者您可以测试完全矢量化的版本：

def total_energy(x):
    b=np.diag(x).cumsum(1)-x
    return potential(abs(b[np.triu_indices_from(b,1)])).sum()

Answer 2

我建议调查scipy.spatial.distance。特别是使用pdist计算数组的所有成对距离。

我假设你有一个形状（Nx3）的数组，因此我们需要稍微改变你的代码：

def potential(r):
       U = 4.*(np.power(r,-12) - np.power(r,-6))
       return U
def total_energy(x):
   E = 0.
   #need to speed up this part
   for i in range(N):                                    #To N here
     for j in range(i):
         E += potential(np.sqrt(np.sum((x[i]-x[j])**2))) #Add sum here
   return E

现在让我们使用spatial：

重写它

import scipy.spatial.distance as sd


def scipy_LJ(arr, sigma=None):
    """
    Computes the Lennard-Jones potential for an array (M x N) of M points
    in N dimensional space. Usage of a sigma parameter is optional.
    """

    if len(arr.shape)==1:
        arr = arr[:,None]

    r = sd.pdist(arr)

    if sigma==None:
        np.power(r, -6, out=r)
        return np.sum(r**2 - r)*4

    else:
       r *= sigma
       np.power(r, -6, out=r)
       return np.sum(r**2 - r)*4

让我们进行一些测试：

N = 1000
points = np.random.rand(N,3)+0.1

np.allclose(total_energy(points), scipy_LJ(points))
Out[43]: True

%timeit total_energy(points)
1 loops, best of 3: 13.6 s per loop

%timeit scipy_LJ(points)
10 loops, best of 3: 24.3 ms per loop

现在它快〜500倍！

N = 10000
points = np.random.rand(N,3)+0.1

%timeit scipy_LJ(points)
1 loops, best of 3: 3.05 s per loop

这使用了~2GB的ram。

Answer 3

以下是一些时间的最终答案

0）普通版（真的很慢）

In [16]: %timeit total_energy(points)
1 loops, best of 3: 14.9 s per loop

1）SciPy版

In [9]: %timeit scipy_LJ(points)
10 loops, best of 3: 44 ms per loop

1-2）Numpy版

 %timeit sum( potential(np.sqrt((x[i]-x[:i])**2 + (y[i]-y[:i])**2 + (z[i] - z[:i])**2)).sum() for i in range(N-1))
10 loops, best of 3: 126 ms per loop

2）疯狂快速的Fortran版本（！ - 意味着评论）

    subroutine EnergyForces(Pos, PEnergy, Dim, NAtom)
    implicit none
    integer, intent(in) :: Dim, NAtom
    real(8), intent(in), dimension(0:NAtom-1, 0:Dim-1) :: Pos
!    real(8), intent(in) :: L
    real(8), intent(out) :: PEnergy
    real(8), dimension(Dim) :: rij, Posi
    real(8) :: d2, id2, id6, id12
    real(8) :: rc2, Shift
    integer :: i, j
    PEnergy = 0.
    do i = 0, NAtom - 1
        !store Pos(i,:) in a temporary array for faster access in j loop
        Posi = Pos(i,:)
        do j = i + 1, NAtom - 1
            rij = Pos(j,:) - Posi
!            rij = rij - L * dnint(rij / L)
            !compute only the squared distance and compare to squared cut
            d2 = sum(rij * rij)
            id2 = 1. / d2            !inverse squared distance
            id6 = id2 * id2 * id2    !inverse sixth distance
            id12 = id6 * id6         !inverse twelvth distance
            PEnergy = PEnergy + 4. * (id12 - id6)
      enddo
    enddo
end subroutine

之后

In [14]: %timeit ljlib.energyforces(points.transpose(), 3, N)
10000 loops, best of 3: 61 us per loop

3）结论Fortran比scipy快1000倍，比numpy快3000倍，比纯python快数百万倍。这是因为Scipy版本创建了一个差异矩阵然后对其进行分析，而Fortran版本可以即时执行所有操作。

Answer 4

感谢您的帮助。这是我发现的。

最短版本

return sum( potential(np.sqrt((x[i]-x[:i])**2)).sum() for i in range(N-1))

scipy版本也不错。
可以考虑的最快版本是使用f2py程序，即在纯Fortran中编写瓶颈慢的部分（这是非常快的），编译它然后将它作为库插入你的python代码中例如，我有： program_lj.f90

$ gfortran -c program_lj.f90

如果在fortran程序中明确定义了所有类型，我们就可以了。

$ f2py -c -m program_lj program_lj.f90

编译之后，唯一剩下的就是从python调用程序。在python程序中：

import program_lj
result = program_lj.subroutine_in_program(parameters)

如果您需要更一般的参考，请参阅精彩的webpage。

如何以最有效的方式写这个

4 个答案: