python - numpy 中带有数组索引的嵌套循环

Question

我想知道如何执行以下操作:

def GetFlux(self, time):
    bx = self.GetField("bx", time) * self.wpewce
    by = self.GetField("by", time) * self.wpewce
    bz = self.GetField("bz", time) * self.wpewce              

    flux  = np.zeros((self.ncells[0]+1,self.ncells[1]+1),"float32", order='FORTRAN')
    flux2  = np.zeros((self.ncells[0]+1,self.ncells[1]+1),"float32", order='FORTRAN')

    dx = self.dl[0]
    dz = self.dl[1]

    nx = self.ncells[0]
    nz = self.ncells[1]

    j = 0

    for i in np.arange(1, nx):
        flux2[i,0] = flux2[i-1,0] + bz[i-1,0]*dx
    flux[1:,0] = flux[0,0] + np.cumsum(bz[:-1,0]*dx)

    for j in np.arange(1,nz):
        flux2[0,j] = flux2[0,j-1] - bx[0,j-1]*dz
    flux[0,1:] = flux[0,0] - np.cumsum(bx[0,:-1]*dz)

    for i in np.arange(1,nx):
        for j in np.arange(1,nz):
            flux2[i,j] = 0.5*(flux2[i-1,j] + bz[i-1,j]*dx) + 0.5*(flux2[i,j-1] - bx[i,j-1]*dz)

    return flux2 

但是没有两个嵌套循环，这需要很长时间。 Bx、Bz 和 flux 是相同大小的数组。

我已设法用数组索引和 cumsum 替换前两个单循环，但我不知道如何替换嵌套循环。

有什么想法吗？

谢谢

Answer 1

有可能(ab)使用 scipy.ndimage.convolve 来解决此类问题。也许在 scipy 中使用一些过滤方法也可能有效并且更好，因为它不依赖于 scipy.ndimage.convolve 就地工作(我可以想象这种情况在遥远的将来会发生变化)。 (编辑:首先写了 scipy.signal.convolve ，它像 numpy.convolve 一样，不能做到这一点)

诀窍是这个卷积函数可以就地使用，因此双 for 循环:

for i in xrange(1, flux.shape[0]):
    for j in xrange(1, flux.shape[1]):
        flux[i,j] = 0.5*(flux[i-1,j] + bz[i-1,j]*dx) + 0.5*(flux[i,j-1] - bx[i,j-1]*dz)

可以替换为(抱歉，需要这么多临时数组...):

from scipy.ndimage import convolve
_flux = np.zeros((flux.shape[0]+1, flux.shape[1]+1), dtype=flux.dtype)
temp_bx = np.zeros((bx.shape[0]+1, bx.shape[1]+1), dtype=bx.dtype)
temp_bz = np.zeros((bz.shape[0]+1, bz.shape[1]+1), dtype=bz.dtype)

_flux[:-1,:-1] = flux
convolve(_flux[:-1,:-1], [[0, 0.5], [0.5, 0]], _flux[1:,1:])

temp_bz[1:,1:-1] = bz[:,1:]*dx
temp_bx[1:-1,1:] = bx[1:,:]*dz

conv_b = np.array([[0.0, 0.5], [0.5, 0.5]])
convolve(temp_bz[:-1,:-1], [[0.5, 0.5], [0.5, 0.]], temp_bz[1:,1:])
convolve(temp_bx[:-1,:-1], [[-0.5, 0.5], [0.5, 0.]], temp_bx[1:,1:])

flux = _flux[:-1,:-1] + temp_by[:-1,:-1] + temp_bx[:-1,:-1]

不幸的是，这意味着我们需要调整 bx、bz 如何进入最终结果，但这种方法可以避免创建大的 2 的幂，并且应该比之前的答案明显更快。

(请注意，numpy 卷积函数不允许这种就地使用。)