使用索引循环的最快方法答案

【问题标题】：Fastest way to loop with index使用索引循环的最快方法
【发布时间】：2017-08-19 07:53:34
【问题描述】：

如果我需要知道索引，循环 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

【问题讨论】：

方法没有太大区别，如果你必须循环的话。最好的方法是根本不循环。我会使用A[i,j,k] 索引，如果不是速度的话，为了清楚起见。你能为i,j,k'一次'的所有组合构造index吗？
我有点怀疑你的index(i, j, k) 是另一种存储元素的布局 - Fortran 顺序与 C 顺序/行主要与列主要等......如果是这样的话，有比遍历所有数组更好的方法。
我不认为我可以一次构建索引。我编辑了帖子，它现在有了实际的循环。
你可以作为一个单一的班轮做到这一点：[do_something for i in range(N) for j in range(N) for k in range(N)]，也许这会更快。
我不确定这是否是禁食，但nditer 可能对你有帮助。

标签： python performance numpy

【解决方案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+ 加速。

【讨论】：