如何在Python中最佳优化在NxM网格上迭代的计算

时间:2019-03-30 22:04:36

标签: python performance numpy optimization iteration

在Python中工作时,我正在NxM值网格上进行一些物理计算,其中N从1到3108,M从1到2304(这相当于一张大图片)。我需要在该空间的每个点计算一个值,总计约700万次计算。我当前的方法非常缓慢,我很想知道是否有一种方法可以完成此任务,并且不需要花费数小时...

我的第一种方法只是使用嵌套的for循环,但这似乎是解决我的问题的最不有效的方法。我尝试使用NumPy的nditer并逐个遍历每个轴,但是我读到它实际上并没有加快我的计算速度。我没有尝试遍历每个轴,而是尝试制作3D数组并遍历外轴,如Brian How can I, in python, iterate over multiple 2d lists at once, cleanly?的答案所示。这是我的代码的当前状态:

import numpy as np
x,y = np.linspace(1,3108,num=3108),np.linspace(1,2304,num=2304) # x&y dimensions of image
X,Y = np.meshgrid(x,y,indexing='ij')
all_coords = np.dstack((X,Y)) # moves to 3-D
all_coords = all_coords.astype(int) # sets coords to int

作为参考,all_coords看起来像这样:

array([[[1.000e+00, 1.000e+00],
        [1.000e+00, 2.000e+00],
        [1.000e+00, 3.000e+00],
        ...,
        [1.000e+00, 2.302e+03],
        [1.000e+00, 2.303e+03],
        [1.000e+00, 2.304e+03]],

       [[2.000e+00, 1.000e+00],
        [2.000e+00, 2.000e+00],
        [2.000e+00, 3.000e+00],
        ...,
        [2.000e+00, 2.302e+03],
        [2.000e+00, 2.303e+03],
        [2.000e+00, 2.304e+03]],

,依此类推。回到我的代码...

'''
- below is a function that does a calculation on the full grid using the distance between x0,y0 and each point on the grid.
- the function takes x0,y0 and returns the calculated values across the grid
'''
def do_calc(x0,y0):
    del_x, del_y = X-x0, Y-y0
    np.seterr(divide='ignore', invalid='ignore')
    dmx_ij = (del_x/((del_x**2)+(del_y**2))) # x component
    dmy_ij = (del_y/((del_x**2)+(del_y**2))) # y component
    return dmx_ij,dmy_ij

# now the actual loop

def do_loop():
    dmx,dmy = 0,0
    for pair in all_coords:
        for xi,yi in pair:
            DM = do_calc(xi,yi)
            dmx,dmy = dmx+DM[0],dmy+DM[1]
    return dmx,dmy

您可能会看到,此代码需要花费非常长的时间才能运行...如果有任何方法可以修改我的代码,而无需花费数小时即可完成,那么我将非常有兴趣了解如何做那。先谢谢您的帮助。

1 个答案:

答案 0 :(得分:1)

这是一种在N=310, M=230下提供10,000倍加速的方法。由于该方法的扩展性好于原始代码,因此我预计在整个问题范围内,该方法的影响因子将超过一百万。

该方法利用了问题的转移不变性。例如,del_x**2在每次调用do_calc时基本上是相同的,因此我们只计算一次。

如果在求和之前对do_calc的输出进行加权,则问题不再是完全平移不变的,并且此方法不再起作用。但是,结果可以用线性卷积表示。在N=310, M=230,这仍然使我们的速度提高了1000倍以上。而且,在整个问题规模上,这将更多

原始问题的代码

import numpy as np

#N, M = 3108, 2304
N, M = 310, 230

### OP's code

x,y = np.linspace(1,N,num=N),np.linspace(1,M,num=M) # x&y dimensions of image
X,Y = np.meshgrid(x,y,indexing='ij')
all_coords = np.dstack((X,Y)) # moves to 3-D
all_coords = all_coords.astype(int) # sets coords to int

'''
- below is a function that does a calculation on the full grid using the distance between x0,y0 and each point on the grid.
- the function takes x0,y0 and returns the calculated values across the grid
'''
def do_calc(x0,y0):
    del_x, del_y = X-x0, Y-y0
    np.seterr(divide='ignore', invalid='ignore')
    dmx_ij = (del_x/((del_x**2)+(del_y**2))) # x component
    dmy_ij = (del_y/((del_x**2)+(del_y**2))) # y component
    return np.nan_to_num(dmx_ij), np.nan_to_num(dmy_ij)

# now the actual loop

def do_loop():
    dmx,dmy = 0,0
    for pair in all_coords:
        for xi,yi in pair:
            DM = do_calc(xi,yi)
            dmx,dmy = dmx+DM[0],dmy+DM[1]
    return dmx,dmy

from time import time

t = [time()]

### pp's code

x, y = np.ogrid[-N+1:N-1:2j*N - 1j, -M+1:M-1:2j*M - 1J]
den = x*x + y*y
den[N-1, M-1] = 1
xx = x / den
yy = y / den
for zz in xx, yy:
    zz[N:] -= zz[:N-1]
    zz[:, M:] -= zz[:, :M-1]
XX = xx.cumsum(0)[N-1:].cumsum(1)[:, M-1:]
YY = yy.cumsum(0)[N-1:].cumsum(1)[:, M-1:]
t.append(time())

### call OP's code for reference

X_OP, Y_OP = do_loop()
t.append(time())

# make sure results are equal

assert np.allclose(XX, X_OP)
assert np.allclose(YY, Y_OP)
print('pp {}\nOP {}'.format(*np.diff(t)))

样品运行:

pp 0.015251636505126953
OP 149.1642508506775

加权问题的代码:

import numpy as np

#N, M = 3108, 2304
N, M = 310, 230

values = np.random.random((N, M))
x,y = np.linspace(1,N,num=N),np.linspace(1,M,num=M) # x&y dimensions of image
X,Y = np.meshgrid(x,y,indexing='ij')
all_coords = np.dstack((X,Y)) # moves to 3-D
all_coords = all_coords.astype(int) # sets coords to int

'''
- below is a function that does a calculation on the full grid using the distance between x0,y0 and each point on the grid.
- the function takes x0,y0 and returns the calculated values across the grid
'''
def do_calc(x0,y0, v):
    del_x, del_y = X-x0, Y-y0
    np.seterr(divide='ignore', invalid='ignore')
    dmx_ij = (del_x/((del_x**2)+(del_y**2))) # x component
    dmy_ij = (del_y/((del_x**2)+(del_y**2))) # y component
    return v*np.nan_to_num(dmx_ij), v*np.nan_to_num(dmy_ij)

# now the actual loop

def do_loop():
    dmx,dmy = 0,0
    for pair, vv in zip(all_coords, values):
        for (xi,yi), v in zip(pair, vv):
            DM = do_calc(xi,yi, v)
            dmx,dmy = dmx+DM[0],dmy+DM[1]
    return dmx,dmy

from time import time
from scipy import signal

t = [time()]
x, y = np.ogrid[-N+1:N-1:2j*N - 1j, -M+1:M-1:2j*M - 1J]
den = x*x + y*y
den[N-1, M-1] = 1
xx = x / den
yy = y / den
XX, YY = (signal.fftconvolve(zz, values, 'valid') for zz in (xx, yy))

t.append(time())
X_OP, Y_OP = do_loop()
t.append(time())
assert np.allclose(XX, X_OP)
assert np.allclose(YY, Y_OP)
print('pp {}\nOP {}'.format(*np.diff(t)))

样品运行:

pp 0.12683939933776855
OP 158.35225439071655