如何以最有效的方式编写这段代码
我有一个数组（大小为10的4次方），我需要计算数组中每两个元素之间的差值（根据原子的坐标计算一个潜在值）。我正在用纯Python写代码，但效率真的不高，有人能告诉我怎么提高速度吗？（可以使用numpy或weave）。这里的x和y是原子的坐标数组（只是简单的一维数组）。
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
4 个回答
谢谢你的帮助。以下是我找到的一些内容。
最简单的版本
return sum( potential(np.sqrt((x[i]-x[:i])**2)).sum() for i in range(N-1))scipy版本也不错。
如果想要更快的版本，可以考虑使用f2py程序，也就是把性能较慢的部分用纯Fortran写出来（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)
如果你需要更一般的参考，可以查看这个很棒的网页。
这里是最终的答案和一些时间上的比较
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 版本则是直接实时处理的。
我建议你看看 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
现在我们用空间库来重写这个：
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的内存。
首先，你可以使用数组运算：
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()