Python NumPy中具有循环依赖项的循环的向量化嵌套
我有下面的函数,它在四面体上生成一系列网格点Python NumPy中具有循环依赖项的循环的向量化嵌套,python,numpy,vectorization,Python,Numpy,Vectorization,我有下面的函数,它在四面体上生成一系列网格点 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
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)
扩展到(nsize,n,3)
,因为i,j,k,l
不依赖于xv
的哪一列被乘以
我们可以跳过循环中的计算l
,而是在返回之前立即执行所有操作:
xg[:,3] = n - xg[:,0:3].sum(1)
现在,您需要做的就是找出如何根据p
以矢量化方式创建i,j,k
一般来说,我发现从“内部到外部”解决这些问题最容易,在最内部的循环中查看代码,并尽可能多地将代码推到尽可能多的循环之外。反复这样做,最终将不会有循环。您可以摆脱依赖的嵌套循环,但要以内存浪费为代价(在矢量化问题中比通常情况下更为如此)。您正在为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
将numpy导入为np
定义测试向量(n):
xv=np.array([
[-1.,-1.,-1.],
[ 1.,-1.,-1.],
[-1., 1.,-1.],
[-1.,-1., 1.],
])
#根据范围(0,n+1)生成三维网格的阵列
ii,jj,kk=np.ogrid[:n+1,:n+1,:n+1]
#位于原始for循环内的三元组的索引
inds=(jj
关键步骤是生成跨越三维网格的ii,jj,kk
索引。这一步是内存高效的,因为np.ogrid
创建了可用于广播操作的跨数组,以引用整个网格。在下一步中,我们只生成一个完整的(n+1)^3
大小的数组:inds
布尔数组选择3d框中位于感兴趣区域内的点(它通过使用数组广播来实现)。在下面的步骤中,np.where(inds)
选择这个大数组中的真实元素,我们最终得到的是数量较少的nsize
元素。因此,单个内存浪费步骤就是创建inds
剩下的很简单:我们需要为数组生成一个额外的列,该列包含每行的[i,j,k]
索引,这可以通过对数组的列求和来完成(同样是一个向量化操作)。一旦我们有了(nsize,4)
形状的辅助数组,该数组在其行中包含每个(i,j,k,l)
:我们需要执行该对象与xv
的矩阵乘法,我们就完成了
使用较小的n
大小进行测试表明,上述函数产生的结果与您的相同。n=100的计时:1.15秒原始,19毫秒矢量化。矢量化这些依赖项很有趣@的确是迪瓦卡!我仍然在等待你的神奇解决方案,它不会浪费5/6的必要内存:)看起来你已经在做我想做的事情了。我们只需要等待内存优化的更好的硬件。