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 优化它，使用带有选项 nopython=True 的 @jit 装饰器，这会给您：

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 版本快。希望这有帮助。

Numpy 向量化和加速

Numpy Vectorization And Speedup

python

optimization

numpy

vectorization