python-在NumPy中向量化嵌套嵌套循环的循环

我具有以下函数,该函数在四面体上生成一系列网格点.

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中是否有简单的方法可以对此向量化？

解决方法:

您可以做的第一件事就是使用广播将三个计算转换为一个：

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,而是在返回之前立即完成所有操作：

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

现在,您需要做的就是找出如何根据p以向量化的方式创建i,j,k的方法.

作为一般说明,我发现从“由内而外”解决这些问题最简单,查看最内层的循环中的代码,并尽可能将其推入尽可能多的循环之外.一遍又一遍地执行此操作,最终将没有循环.

标签：vectorization,python,numpy
来源： https://codeday.me/bug/20191111/2019249.html