[英]Fastest way to loop with index
如果我需要知道indeces,那么在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
如果我們可以訪問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)
時間:
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
在這些輸入上看到120x+
加速。
聲明:本站的技術帖子網頁,遵循CC BY-SA 4.0協議,如果您需要轉載,請注明本站網址或者原文地址。任何問題請咨詢:yoyou2525@163.com.