在NumPy中使用循环依赖关系嵌套for循环

时间:2017-05-21 10:49:08

标签: 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 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中进行矢量化?

2 个答案:

答案 0 :(得分:2)

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

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)

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

总的来说,我发现最容易从“由内而外”处理这些问题,查看最内层循环中的代码并尽可能多地推出尽可能多的循环。一遍又一遍地这样做,最终没有循环。

答案 1 :(得分:1)

可以摆脱你的依赖嵌套循环,但代价是浪费在内存中(在矢量化问题中比平常更多)。您在for循环中处理3d框的一小部分。如果您不介意生成仅使用$sqlConnection = new-object system.data.SqlClient.SqlConnection("Data Source=.\SQLExpress;Integrated Security=SSPI;Initial Catalog=master") try { $sqlConnection.Open() $commandText = @" exec sp_msforeachdb 'IF ''?'' NOT IN (''master'', ''model'', ''msdb'', ''tempdb'') BEGIN drop database [?] END' "@ $sqlCommand = New-Object System.Data.SqlClient.SqlCommand $sqlCommand.CommandText = $commandText $sqlCommand.Connection = $sqlConnection $SQLCommand.CommandTimeout = 0 $sqlCommand.ExecuteNonQuery() } finally{ $sqlConnection.Close() } 项的(n+1)^3项,并且您可以适应内存,那么矢量化版本的确可能更快。

我的建议,解释如下:

(n+1)(n+2)(n+3)/6

关键步骤是生成跨越3d网格的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创建了可用于广播操作以引用完整网格的跨越数组。我们只在下一步中生成一个完整的np.ogrid大小的数组:(n+1)^3布尔数组选择位于您感兴趣的区域内的3d框中的那些点(并且它通过使用数组来实现广播)。在接下来的步骤inds中挑选出这个大数组的真实元素,最后得到的np.where(inds)元素数量较少。因此,单个内存浪费步骤是nsize的创建。

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

使用小xv尺寸进行测试表明,上述功能会产生与您相同的结果。 n的时间:原始的1.15秒,向量化的19毫秒。