使用NumPy向量化具有循环依赖的嵌套for循环

Question

我有一个函数，用于在四面体上生成一系列网格点。

def tet_grid(n):

    xv = np.array([
        [-1.,-1.,-1.],
        [ 1.,-1.,-1.],
        [-1., 1.,-1.],
        [-1.,-1., 1.],
        ])

    nsize = int((n+1)*(n+2)*(n+3)/6);
    xg = np.zeros((nsize,3))
    p = 0

    for i in range ( 0, n + 1 ):
        for j in range ( 0, n + 1 - i ):
            for k in range ( 0, n + 1 - i - j ):
                l = n - i - j - k
                xg[p,0]=(i * xv[0,0] + j * xv[1,0] + k * xv[2,0] + l * xv[3,0])/n 
                xg[p,1]=(i * xv[0,1] + j * xv[1,1] + k * xv[2,1] + l * xv[3,1])/n 
                xg[p,2]=(i * xv[0,2] + j * xv[1,2] + k * xv[2,2] + l * xv[3,2])/n 
                p = p + 1

    return xg

有没有在NumPy中向量化此操作的简便方法？

- Tasam Farkie

Answer 1

你可以使用广播将三个计算合并为一个，这是最简单的方法：

xg[p]=(i * xv[0] + j * xv[1] + k * xv[2] + l * xv[3])/n

下一步需要注意的是，除以n可以移动到最后：

return xg / n

然后，我们可以将四个乘法拆分并分别存储结果，最后再进行合并。现在我们有：

xg = np.empty((nsize,4)) # n.b. zeros not required
p = 0

for i in range ( 0, n + 1 ):
    for j in range ( 0, n + 1 - i ):
        for k in range ( 0, n + 1 - i - j ):
            l = n - i - j - k
            xg[p,0] = i
            xg[p,1] = j
            xg[p,2] = k
            xg[p,3] = l
            p = p + 1

return (xg[:,:,None] * xv).sum(1) / n

下面的技巧是通过广播 (nsize,n) * (n,3) 实现的，其中xg[:,:,None] 的关键在于我们将 (nsize, n) 的 xg 扩展为 (nsize, n, 3)，因为 i,j,k,l 不依赖于与 xv 的哪一列相乘。

在循环中可以跳过计算 l ，而是在 return 之前一次性完成所有操作。

xg[:,3] = n - xg[:,0:3].sum(1)

现在你需要做的就是找出如何按照 p 的向量化方式创建 i,j,k。

一般来说，我会从“内部向外”解决这些问题，先看最内部循环中的代码，然后尽可能地将其推到尽可能多的循环之外。不断重复此过程，最终就没有循环了。

- John Zwinck

Answer 2

你可以摆脱你的依赖嵌套循环，但代价是浪费内存（比向量化问题通常更多）。在你的for循环中，你正在处理3d盒子的一个小子集。如果你不介意生成(n+1)^3个项目，仅使用(n+1)(n+2)(n+3)/6个项目，并且可以适应内存，则向量化版本确实可能更快。

我的建议如下解释：

import numpy as np

def tet_vect(n):

    xv = np.array([
        [-1.,-1.,-1.],
        [ 1.,-1.,-1.],
        [-1., 1.,-1.],
        [-1.,-1., 1.],
        ])

    # spanning arrays of a 3d grid according to range(0,n+1)
    ii,jj,kk = np.ogrid[:n+1,:n+1,:n+1]
    # indices of the triples which fall inside the original for loop
    inds = (jj < n+1-ii) & (kk < n+1-ii-jj)
    # the [i,j,k] indices of the points that fall inside the for loop, in the same order
    combs = np.vstack(np.where(inds)).T
    # combs is now an (nsize,3)-shaped array

    # compute "l" column too
    lcol = n - combs.sum(axis=1)
    combs = np.hstack((combs,lcol[:,None]))
    # combs is now an (nsize,4)-shaped array

    # all we need to do now is to take the matrix product of combs and xv, divide by n in the end
    xg = np.matmul(combs,xv)/n

    return xg

关键步骤是生成ii，jj，kk索引，这些索引跨越3D网格。这一步是内存高效的，因为np.ogrid创建跨度数组，可用于广播操作以引用您的完整网格。我们仅在下一步生成完整的(n+1)^3大小的数组：inds布尔数组通过利用数组广播选择位于您感兴趣的区域内的3D框中的点。在以下步骤中，np.where(inds)挑选出这个大数组的真值元素，最终获得较少数量的nsize元素。因此，单一内存浪费步骤是创建inds。

其余步骤很简单：我们需要为包含每行中的[i，j，k]索引的数组生成额外列，可以通过对数组的列求和来完成（再次进行向量化操作）。一旦我们有了包含每个(i，j，k，l)的辅助数组的(nsize，4)形状对象：我们需要对这个对象进行矩阵乘法，然后就完成了。

测试小的n大小表明，以上函数产生与您的相同结果。当n = 100时的时间：1.15秒原始，19毫秒向量化。

- Andras Deak -- Слава Україні

---

那些有依赖关系的部分很有趣，可以进行向量化！ - Divakar

@Divakar 确实如此！我仍在等待您的神奇解决方案，它不会浪费5/6必要的内存 :) - Andras Deak -- Слава Україні

看起来你已经在做我会做的事情了。我们只需要等待更好的内存优化硬件。 - Divakar