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带有 option 的装饰器nopython=True，它可以为您提供：import numpy as npfrom numba import jitnlin, npix = 478, 480bb = np.random.rand(nlin,npix)slope = -8fac = 4offset= 0def 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))>>> Truenp.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一个版本快。希望这可以帮助。