Numpy 向量化和加速

Question

我发现了一个小代码片段，它曾经是一个双 for 循环，我设法通过矢量化将它带到一个单 for 循环中。这样做会大大缩短时间，所以我想知道是否可以通过矢量化摆脱第二个 for 循环，以及它是否会提高性能。import numpy as npfrom timeit import default_timer as timernlin, npix = 478, 480bb = np.random.rand(nlin,npix)slope = -8fac = 4offset= 0barray = 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带有 option 的装饰器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一个版本快。希望这可以帮助。