使用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和{}是相同大小的数组。在

我已经成功地用数组索引和cumsum替换了前两个单循环，但是我无法找到如何替换嵌套循环。在

有什么想法吗？在

谢谢

Answer 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)

被替换为

^{pr2}$

我认为第一个卷积（在通量上）将使用模具[[0，a]，[b，0]]。但是，假设a、b、c和d是标量，那么bz和bx的另外两个模板是什么呢？在

Answer 2

有可能使用scipy.ndimage.卷积对于这种问题。也许在scipy中使用一些过滤方法也可以，而且更好，因为它不依赖于scipy.ndimage.卷积在适当的地方工作（我可以想象这在遥远的将来会发生变化）。（编辑：第一作者信号卷积就像纽比卷积，无法执行此操作）

诀窍在于，这个卷积函数可以在适当的地方使用，这样double 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)

可以替换为（抱歉，需要这么多临时数组…）：

^{pr2}$

不幸的是，这意味着我们需要分清bx，bz如何进入最终结果，但这种方法避免了产生2的大幂次，而且应该比之前的答案快得多。在

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

Answer 3

内环的矢量化相当简单。你有一个基本方程，如下所示：

X[n] = a * X[n-1] + b[n]

这个方程可以展开和重写，而不依赖于X[n-1]：

^{pr2}$

如果原始代码是这样的：

for i in np.arange(1,nx+1):
    for j in np.arange(1,nz+1):
        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)

你可以像这样去掉内部循环：

a = 0.5
aexp = np.arange(nz).reshape(nz, 1) - np.arange(nz).reshape(1, nz)
abcoeff = a**aexp
abcoeff[aexp<0] = 0
for i in np.arange(1,nx+1):
    b = 0.5*flux2[i-1, 1:] + 0.5*bz[i-1, 1:]*dx - 0.5*bx[i,:-1]*dz
    bvals = (abcoeff * b.reshape(1, nz)).sum(axis=1)
    n = np.arange(1, nz+1)
    x0 = flux2[i, 0]
    flux2[i, 1:] = a**n * x0 + bvals

由于浮点错误，这些值不会完全相同，但足够接近。 从理论上讲，您可以应用相同的过程来消除这两个循环，但这会变得相当复杂，并且根据阵列的形状，可能不会提供太多的性能优势。在