加速python中的嵌套for循环/遍历numpy数组

Question

假设我有 4 个 numpy 数组 A,B,C,D ，每个数组的大小为 (256,256,1792)。 我想遍历这些数组的每个元素并对它做一些事情，但是我需要以 256x256x256 立方体的块来做。

我的代码如下所示：

for l in range(7): 
    x, y, z, t = 0,0,0,0
    for m in range(a.shape[0]):
        for n in range(a.shape[1]):
            for o in range(256*l,256*(l+1)):
                t += D[m,n,o] * constant
                x += A[m,n,o] * D[m,n,o] * constant
                y += B[m,n,o] * D[m,n,o] * constant
                z += C[m,n,o] * D[m,n,o] * constant
    final = (x+y+z)/t
    doOutput(final)

代码可以正常工作并输出我想要的内容，但速度非常慢。我在网上读到，在 python 中应该避免那种嵌套的 for 循环。什么是最干净的解决方案？ （现在我正在尝试在 C 中完成这部分代码，并以某种方式通过 Cython 或其他工具导入它，但我喜欢纯 python 解决方案）

谢谢

添加

Willem Van Onsem 对第一部分的解决方案似乎工作得很好，我想我理解它。但现在我想在对它们求和之前修改我的值。好像

（在外 l 循环内）

for m in range(a.shape[0]):
    for n in range(a.shape[1]):
        for o in range(256*l,256*(l+1)):
            R += (D[m,n,o] * constant * (A[m,n,o]**2 
            + B[m,n,o]**2 + C[m,n,o]**2)/t - final**2)
doOutput(R)

我显然不能只求和 x = (A[:a.shape[0],:a.shape[1],256*l:256*(l+1)]*Dsub).sum()**2*constant 的平方，因为 (A²+B²) != (A+B)² 我怎样才能重做这个最后的 for 循环？

Answer 1

由于您使用m in range(a.shape[0])、n in range(a.shape[1]) 和o in range(256*l,256*(l+1)) 的每个元素更新t，因此您可以替换：

for m in range(a.shape[0]):
    for n in range(a.shape[1]):
        for o in range(256*l,256*(l+1)):
            t += D[m,n,o]

与：

t += D[:a.shape[0],:a.shape[1],256*l:256*(l+1)].sum()

其他任务也一样。因此，您可以将代码重写为：

for l in range(7): 
    Dsub = D[:a.shape[0],:a.shape[1],256*l:256*(l+1)]
    x = (A[:a.shape[0],:a.shape[1],256*l:256*(l+1)]*Dsub).sum()*constant
    y = (B[:a.shape[0],:a.shape[1],256*l:256*(l+1)]*Dsub).sum()*constant
    z = (C[:a.shape[0],:a.shape[1],256*l:256*(l+1)]*Dsub).sum()*constant
    t = Dsub.sum()*constant
   final = (x+y+z)/t
   doOutput(final)

请注意，numpy 中的* 是逐元素 乘法，不是 矩阵乘积。您可以在求和之前进行乘法运算，但由于乘法与常数的总和等于该常数与总和的乘积，因此我认为在循环外执行此操作更有效。

如果a.shape[0] 等于D.shape[0] 等。您可以使用: 代替:a.shape[0]。根据您的问题，似乎是这样。所以：

# only when `a.shape[0] == D.shape[0], a.shape[1] == D.shape[1] (and so for A, B and C)`
for l in range(7): 
    Dsub = D[:,:,256*l:256*(l+1)]
    x = (A[:,:,256*l:256*(l+1)]*Dsub).sum()*constant
    y = (B[:,:,256*l:256*(l+1)]*Dsub).sum()*constant
    z = (C[:,:,256*l:256*(l+1)]*Dsub).sum()*constant
    t = Dsub.sum()*constant
    final = (x+y+z)/t
    doOutput(final)

在numpy 级别处理.sum() 将提高性能，因为您不会来回转换值，而使用.sum()，您使用紧密 循环。

编辑：

您更新后的问题没有太大变化。您可以简单地使用：

m,n,_* = a.shape
lo,hi = 256*l,256*(l+1)
R = (D[:m,:n,lo:hi]*constant*(A[:m,:n,lo:hi]**2+B[:m,:n,lo:hi]**2+D[:m,:n,lo:hi]**2)/t-final**2)).sum()
doOutput(R)