python优化循环_python中的双循环优化

我正在尝试优化以下循环：def numpy(nx, nz, c, rho):

for ix in range(2, nx-3):

for iz in range(2, nz-3):

a[ix, iz] = sum(c*rho[ix-1:ix+3, iz])

b[ix, iz] = sum(c*rho[ix-2:ix+2, iz])

return a, b

我尝试了不同的解决方案，发现使用numba计算产品总和可以获得更好的性能：

^{pr2}$

这导致Time numpy : 4.1595

Time numba1 : 0.6993

Time numba2 : 1.0135

使用numba版本的sum函数(sum\u opt)性能非常好。但是我想知道为什么numba版本的双循环函数(numba2)会导致执行速度变慢。我试图使用jit而不是autojit，指定参数类型，但情况更糟。在

我还注意到，先在最小的循环上循环比先在最大的循环上循环慢。有什么解释吗？在

不管怎样，我确信这个双循环函数可以改进很多向量化问题(比如this)或使用其他方法(map？)但我对这些方法有点困惑。在

在我代码的其他部分，我使用numba和numpy切片方法来替换所有显式循环，但是在这个特殊的例子中，我不知道如何设置它。在

有什么想法吗？在

编辑

谢谢你的评论。我在这个问题上做了一些工作：import numba as nb

import numpy as np

from scipy import signal

import time

@nb.jit(['float64(float64[:], float64[:])'], nopython=True)

def sum_opt(arr1, arr2):

s = arr1[0]*arr2[0]

for i in xrange(1, len(arr1)):

s+=arr1[i]*arr2[i]

return s

@nb.autojit

def numba1(nx, nz, c, rho, a, b):

for ix in range(2, nx-3):

for iz in range(2, nz-3):

a[ix, iz] = sum_opt(c, rho[ix-1:ix+3, iz])

b[ix, iz] = sum_opt(c, rho[ix-2:ix+2, iz])

return a, b

@nb.jit(nopython=True)

def numba2(nx, nz, c, rho, a, b):

for ix in range(2, nx-3):

for iz in range(2, nz-3):

a[ix, iz] = sum_opt(c, rho[ix-1:ix+3, iz])

b[ix, iz] = sum_opt(c, rho[ix-2:ix+2, iz])

return a, b

@nb.jit(['float64[:,:](int16, int16, float64[:], float64[:,:], float64[:,:])'], nopython=True)

def numba3a(nx, nz, c, rho, a):

for ix in range(2, nx-3):

for iz in range(2, nz-3):

a[ix, iz] = sum_opt(c, rho[ix-1:ix+3, iz])

return a

@nb.jit(['float64[:,:](int16, int16, float64[:], float64[:,:], float64[:,:])'], nopython=True)

def numba3b(nx, nz, c, rho, b):

for ix in range(2, nx-3):

for iz in range(2, nz-3):

b[ix, iz] = sum_opt(c, rho[ix-2:ix+2, iz])

return b

def convol(nx, nz, c, aa, bb):

s1 = rho[1:nx-1,2:nz-3]

s2 = rho[0:nx-2,2:nz-3]

kernel = c[:,None][::-1]

aa[2:nx-3,2:nz-3] = signal.convolve2d(s1, kernel, boundary='symm', mode='valid')

bb[2:nx-3,2:nz-3] = signal.convolve2d(s2, kernel, boundary='symm', mode='valid')

return aa, bb

nx = 1024

nz = 256

rho = np.random.rand(nx, nz)

c = np.random.rand(4)

a = np.zeros((nx, nz))

b = np.zeros((nx, nz))

ti = time.clock()

for i in range(1000):

a, b = numba1(nx, nz, c, rho, a, b)

print 'Time numba1 : ' + `round(time.clock() - ti, 4)`

ti = time.clock()

for i in range(1000):

a, b = numba2(nx, nz, c, rho, a, b)

print 'Time numba2 : ' + `round(time.clock() - ti, 4)`

ti = time.clock()

for i in range(1000):

a = numba3a(nx, nz, c, rho, a)

b = numba3b(nx, nz, c, rho, b)

print 'Time numba3 : ' + `round(time.clock() - ti, 4)`

ti = time.clock()

for i in range(1000):

a, b = convol(nx, nz, c, a, b)

print 'Time convol : ' + `round(time.clock() - ti, 4)`

您的解决方案非常优雅，但我必须在代码中大量使用此函数。因此，对于1000次迭代，这将导致Time numba1 : 3.2487

Time numba2 : 3.7012

Time numba3 : 3.2088

Time convol : 22.7696

autojit和{}非常接近。

但是，在使用jit时，指定所有参数类型似乎很重要。在

当函数有多个输出时，我不知道是否有方法在jit修饰符中指定参数类型。有人吗？在

现在我没有找到其他的解决办法，除了使用numba。欢迎有新想法！在