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

あなたが行うことができます最初に簡単な事は一つに3つの計算をオンにする放送使用です：

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

次はnによる除算は最後の最後に移動することができることに注意することです。

return xg/n 

次に、4つの乗算を分割して結果を別々に格納し、最後にそれらを組み合わせることができます。今、私たちは持っている：

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)を放送することである - i,j,k,lがxvの列がで乗算されているに依存しないので、我々は(nsize,n,3)に(nsize,n)xgを展開します。

我々は、ループ内lを計算スキップし、代わりに右return前に一度それをすべて行うことができます。

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

今

すべてを行う必要がに従ってベクトル化の方法でi,j,kを作成する方法を見つけ出すれますp。

一般的な注意点として、これらの問題を最も内側のループのコードを見て、可能な限り多くのループの外側にできるだけ多く押し込むことが、最も簡単な方法です。これを何度も繰り返し、最終的にはループはありません。

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 

キーステップは、3Dグリッドにまたがるii,jj,kkインデックスを生成しています。 np.ogridは、ブロードキャスト操作でフルグリッドを参照するために使用できるスパニングアレイを作成するので、この手順はメモリ効率が良い方法です。次のステップでは、完全な(n+1)^3サイズの配列のみを生成します。inds boolean配列は、興味のある領域の内側にある3dボックス内の点を選択します（配列ブロードキャストを使用して行います）。次のステップでは、np.where(inds)がこの大きな配列のtruthy要素を選び出し、より小さな数のnsize要素で終わります。したがって、単一のメモリ消費ステップは、indsの作成です。

残りの部分は簡単です：各行に[i,j,k]のインデックスを含む配列の追加の列を生成する必要があります。これは、配列の列を合計してベクトル化することによって行うことができます。(i,j,k,l)が行内にある(nsize,4)型補助配列を取得したら、このオブジェクトの行列乗算をxvで実行する必要があります。

サイズが小さいとテストすると、上記の機能があなたと同じ結果をもたらすことが示唆されます。 n = 100のタイミング：1.15秒元、19msベクトル化。