在 NumPy 中为具有循环依赖关系的循环进行嵌套矢量化

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 中对此进行矢量化？

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

底部xg[:,:,None]的诀窍是广播(nsize,n) * (n,3)- 我们将(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。

作为一般说明，我发现从"由内而外"解决这些问题是最容易的，查看最内层循环中的代码，并尽可能多地推动尽可能多的循环。 一遍又一遍地这样做，最终不会有循环。

Answer 2

您可以摆脱依赖的嵌套循环，但代价是内存浪费(在矢量化问题中比平时更是如此)。您正在处理 for 循环中 3D 框的一小部分。如果您不介意生成仅使用(n+1)(n+2)(n+3)/6项的(n+1)^3项，并且可以放入内存中，那么矢量化版本确实可能会快得多。

我的建议，解释如下：

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

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

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

使用小n进行测试表明上述函数产生的结果与您的相同。n = 100时序：原始 1.15 秒，矢量化 19 毫秒。