С точки зрения векторизации вы могли бы потенциально использовать что-то на основе 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
. Надеюсь, это поможет.