如何加速Python嵌套循环？

Question

我正在尝试计算一个埋在地下物体的引力影响。具体做法是先计算物体每一侧的引力效果，然后把这些效果加起来，得到一个测量值，接着在多个测量点重复这个过程。下面是我的代码（这个物体是个正方形，代码是顺时针方向计算的，所以坐标会从-x回到-x）。

grav = []
x = si.arange(-30.0, 30.0, 0.5)

# -9.79742526     9.78716693    22.32153704    27.07382349  2138.27146193
xcorn = (-9.79742526,  9.78716693,  9.78716693, -9.79742526, -9.79742526)
zcorn = (22.32153704, 22.32153704, 27.07382349, 27.07382349, 22.32153704)
gamma = (6.672*(10**-11)) # 'N m^2 / Kg^2'
rho = 2138.27146193 # 'Kg / m^3'
grav = []
iter_time = []
def procedure():
    for i in si.arange(len(x)): # Cycles position
        t0 = time.clock()
        sum_lines = 0.0

        for n in si.arange(len(xcorn)-1): # Cycles corners
            x1 = xcorn[n]   - x[i]
            x2 = xcorn[n+1] - x[i]
            z1 = zcorn[n]   -0.0  # Just depth to corner since all observations are on the surface.
            z2 = zcorn[n+1] -0.0
            r1 = ((z1**2) + (x1**2))**0.5
            r2 = ((z2**2) + (x2**2))**0.5
            O1 = si.arctan2(z1, x1)
            O2 = si.arctan2(z2, x2)
            denom = z2-z1
            if denom == 0.0:
                denom = 1.0e-6

            alpha = (x2-x1)/denom

            beta = ((x1*z2)-(x2*z1))/denom
            factor = (beta/(1.0 + (alpha**2)))
            term1 = si.log(r2/r1) # Log base 10
            term2 = alpha*(O2-O1)
            sum_lines = sum_lines + (factor*(term1-term2))
        sum_lines = sum_lines*2*gamma*rho
        grav.append(sum_lines)
        t1 = time.clock()
        dt = t1-t0
        iter_time.append(dt)

我该如何加快这个循环的速度呢？

Answer 1

在Python的循环中，逐个访问NumPy数组里的元素是非常低效的。例如，下面这个Python循环：

for i in xrange(0, len(a), 2):
    a[i] = i

执行起来会比下面这个慢很多：

a[::2] = np.arange(0, len(a), 2)

你可以使用更好的算法（也就是时间复杂度更低的算法），或者像上面那样在NumPy数组上使用向量运算。但更简单的办法可能就是用Cython来编译代码：

#cython: boundscheck=False, wraparound=False
#procedure_module.pyx
import numpy as np
cimport numpy as np

ctypedef np.float64_t dtype_t

def procedure(np.ndarray[dtype_t,ndim=1] x,
              np.ndarray[dtype_t,ndim=1] xcorn):
    cdef:
        Py_ssize_t i, j
        dtype_t x1, x2, z1, z2, r1, r2, O1, O2
        np.ndarray[dtype_t,ndim=1] grav = np.empty_like(x)
    for i in range(x.shape[0]):
        for j in range(xcorn.shape[0]-1):
            x1 = xcorn[j]-x[i]
            x2 = xcorn[j+1]-x[i]
            ...
        grav[i] = ...
    return grav

虽然不需要定义所有的数据类型，但如果你想要比Python快很多，至少应该定义数组的类型和循环的索引。

你可以使用cProfile（Cython是支持这个的），这样就不用手动调用time.clock()来计时了。

要调用procedure()：

#!/usr/bin/env python
import pyximport; pyximport.install() # pip install cython
import numpy as np

from procedure_module import procedure

x = np.arange(-30.0, 30.0, 0.5)
xcorn = np.array((-9.79742526, 9.78716693, 9.78716693, -9.79742526, -9.79742526))
grav = procedure(x, xcorn)

Answer 2

你的xcorn和zcorn的值是重复的，所以可以考虑把一些计算的结果缓存起来，这样可以提高效率。

可以看看timeit和profile这两个模块，它们可以帮助你了解哪些操作花费的计算时间最多。