我如何加快python嵌套循环的速度？

Question

I'm trying to calculate the gravity effect of a buried object by calculating the effect on each side of the body then summing up the contributions to get one measurement at one station, an repeating for a number of stations. the code is as follows( the body is a square and the code calculates clockwise around it, that's why it goes from -x back to -x coordinates)

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)

Any help in speeding this loop up would be appreciated Thanks.

Answer 1

Your xcorn and zcorn values repeat, so consider caching the result of some of the computations.

Take a look at the timeit and profile modules to get more information about what is taking the most computational time.

Answer 2

It is very inefficient to access individual elements of a numpy array in a Python loop. For example, this Python loop:

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

would be much slower than:

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

You could use a better algorithm (less time complexity) or use vector operations on numpy arrays as in the example above. But the quicker way might be just to compile the code using 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

It is not necessary to define all types but if you need a significant speed up compared to Python you should define at least types of arrays and loop indexes.

You could use cProfile (Cython supports it) instead of manual calls to time.clock() .

To call 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)