インデックスを使用してループする最速の方法

Question

インデックスを知る必要がある場合、3d numpy配列をループする最も速い方法は何ですか。

... some sort of loop ... 
    do something with each element that requires a knowledge of i,j,k. 

など。

for i in range(N): 
    for j in range(N): 
     for k in range(N): 
      index = # something that depends on i,j,k 
      B[index] = A[i][j][k]**2 

実際のループは次のようになります。問題がある場合、

for z in range(Ngrid):                               
    kz = 2*pi/LMAX*(z - Ngrid/2)             
    for y in range(Ngrid):              
     ky = 2*pi/LMAX*(y - Ngrid/2)            
     for x in range(Ngrid):             
      kx = 2*pi/LMAX*(x - Ngrid/2)           
      kk = sqrt(kx**2 + ky**2 + kz**2)          
      bind = int((kk - kmin)/dk)           
      if bind >= Nk:              
       continue               
      delk = delta_k[x][y][z]            
      Pk[bind] += (delk.real**2 + delk.imag**2)/2       
      Numk[bind] += 1 

Answer 1

numpyのツールに我々がアクセス権を持っていることを与えられた問題を解決するための最速の方法は全くしないループです並列化/ベクトル化可能。手元にある問題については、それをベクトル化できると思われます。この問題は以前のQ&Aによく似ています。だから、私たちはそのポストからほんの少しのものを借りて、それは主にbroadcastingを使うことを中心にしています。

したがって、我々はそうのようなソリューション、希望 -

KXYZ = 2*np.pi/LMAX*(np.arange(Ngrid) - Ngrid/2) 
KK = np.sqrt(KXYZ[:,None,None]**2 + KXYZ[:,None]**2 + KXYZ**2) 
BIND = ((KK - kmin)/dk).astype(int) 

valid_mask = BIND<Nk 
IDs = BIND[valid_mask] 
vals = (delta_k.real[valid_mask]**2 + delta_k.imag[valid_mask]**2)/2 

Pk += np.bincount(IDs, vals, minlength=len(Pk)) 
Numk += np.bincount(IDs, minlength=len(Numk)) 

ランタイムテスト

アプローチ -

def loopy_app(Ngrid, LMAX, kmin, dk, Nk, delta_k): 
    Pk = np.zeros(Nk) 
    Numk = np.zeros(Nk) 
    for z in range(Ngrid):     
     kz = 2*np.pi/LMAX*(z - Ngrid/2)   
     for y in range(Ngrid):  
      ky = 2*np.pi/LMAX*(y - Ngrid/2)  
      for x in range(Ngrid):     
       kx = 2*np.pi/LMAX*(x - Ngrid/2)   
       kk = np.sqrt(kx**2 + ky**2 + kz**2)   
       bind = int((kk - kmin)/dk)  
       if bind >= Nk:  
        continue        
       delk = delta_k[x,y,z]        
       Pk[bind] += (delk.real**2 + delk.imag**2)/2 
       Numk[bind] += 1 
    return Pk, Numk   

def vectorized_app(Ngrid, LMAX, kmin, dk, Nk, delta_k): 
    Pk = np.zeros(Nk) 
    Numk = np.zeros(Nk) 

    KXYZ = 2*np.pi/LMAX*(np.arange(Ngrid) - Ngrid/2) 
    KK = np.sqrt(KXYZ[:,None,None]**2 + KXYZ[:,None]**2 + KXYZ**2) 
    BIND = ((KK - kmin)/dk).astype(int) 

    valid_mask = BIND<Nk 
    IDs = BIND[valid_mask] 
    vals = (delta_k.real[valid_mask]**2 + delta_k.imag[valid_mask]**2)/2 

    Pk += np.bincount(IDs, vals, minlength=len(Pk)) 
    Numk += np.bincount(IDs, minlength=len(Numk)) 
    return Pk, Numk 

入力セットアップ：

# Setup inputs with random numbers 
Ngrid = 100 
LMAX = 3.45 
kmin = 0.345 
dk = 1.56 
Nk = 80 
delta_k = np.random.rand(Ngrid,Ngrid,Ngrid) + 1j * \ 
      np.random.rand(Ngrid,Ngrid,Ngrid) 

これらの入力に存在し 120x+スピードアップを見て

In [186]: app1_out1, app1_out2 = loopy_app(Ngrid, LMAX, kmin, dk, Nk, delta_k) 
    ...: app2_out1, app2_out2 = vectorized_app(Ngrid, LMAX, kmin, dk, Nk, delta_k) 
    ...: print np.allclose(app1_out1, app2_out1) 
    ...: print np.allclose(app1_out2, app2_out2) 
    ...: 
True 
True 

In [187]: %timeit loopy_app(Ngrid, LMAX, kmin, dk, Nk, delta_k) 
    ...: %timeit vectorized_app(Ngrid, LMAX, kmin, dk, Nk, delta_k) 
    ...: 
1 loops, best of 3: 2.61 s per loop 
10 loops, best of 3: 20.7 ms per loop 

In [188]: 2610/20.7 
Out[188]: 126.08695652173914 

：の

タイミング。