python - Numpy 矢量化和加速

Question

我发现是一个小的代码片段，它曾经是一个双 for 循环，我设法通过矢量化将它变成了一个单一的 for 循环。完成后这导致了显着的时间改进，所以我想知道是否有可能通过矢量化来摆脱这里的第二个 for 循环，以及它是否会提高性能。

import numpy as np
from timeit import default_timer as timer
nlin, npix = 478, 480
bb = np.random.rand(nlin,npix)
slope = -8
fac = 4
offset= 0
barray = np.zeros([2,2259]);

timex = timer()
for y in range(nlin):
    for x in range(npix):
        ling=(np.ceil((x-y/slope)*fac)+1-offset).astype(np.int);
        barray[0,ling] +=1;
        barray[1,ling] +=bb[y,x];
newVar = np.copy(barray)
print(timer() - timex)

因此可以通过创建以下矩阵将 ling 从循环中取出

lingMat = (np.ceil((np.vstack(npixrange)-nlinrange/slope)*fac)+1-offset).astype(np.int);

满足 lingMat[x,y] = "ling in the for loop at x and y"。这给出了矢量化的第一步。

Answer 1

在矢量化方面，您可能会使用基于 np.add.at 的东西。 :

def yaco_addat(bb,slope,fac,offset):
    barray = np.zeros((2,2259),dtype=np.float64)
    nlin_range = np.arange(nlin)
    npix_range = np.arange(npix)
    ling_mat = (np.ceil((npix_range-nlin_range[:,None]/slope)*fac)+1-offset).astype(np.int)  
    np.add.at(barray[0,:],ling_mat,1)
    np.add.at(barray[1,:],ling_mat,bb) 
    return barray

但是，我建议您直接使用 numba 进行优化。 , 使用 @jit带选项的装饰器 nopython=True ，这给你:

import numpy as np
from numba import jit

nlin, npix = 478, 480
bb = np.random.rand(nlin,npix)
slope = -8
fac = 4
offset= 0

def yaco_plain(bb,slope,fac,offset):
    barray = np.zeros((2,2259),dtype=np.float64)
    for y in range(nlin):
        for x in range(npix):
            ling=(np.ceil((x-y/slope)*fac)+1-offset).astype(np.int)
            barray[0,ling] += 1
            barray[1,ling] += bb[y,x]
    return barray

@jit(nopython=True)
def yaco_numba(bb,slope,fac,offset):
    barray = np.zeros((2,2259),dtype=np.float64)
    for y in range(nlin):
        for x in range(npix):
            ling = int((np.ceil((x-y/slope)*fac)+1-offset))
            barray[0,ling] += 1
            barray[1,ling] += bb[y,x]    
    return barray

让我们检查输出

np.allclose(yaco_plain(bb,slope,fac,offset),yaco_addat(bb,slope,fac,offset))
>>> True
np.allclose(yaco_plain(bb,slope,fac,offset),yaco_jit(bb,slope,fac,offset))
>>> True

现在时间这些

%timeit yaco_plain(bb,slope,fac,offset)
>>> 648 ms ± 4.14 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit yaco_addat(bb,slope,fac,offset)
>>> 27.2 ms ± 92.3 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

%timeit yaco_jit(bb,slope,fac,offset)
>>> 505 µs ± 995 ns per loop (mean ± std. dev. of 7 runs, 1000 loops each)

导致优化的函数比最初的 2 个循环版本和 53x 快得多比 np.add.at 快一。希望这可以帮助。