Узкое место в производительности индексации с массивами NumPy (или создании кортежей)

Question

Представьте, что у нас есть следующая функция:

def return_slice(self, k):
    return (self.A[self.C[k]:self.C[k+1]], self.B[self.C[k]:self.C[k+1]])

, которая является частью класса с массивами A, B и C, которые содержат тонну целых чисел (свыше 10 ^ 5). Хотя вызов этой функции несколько раз выполняется достаточно быстро, я заметил, что ~ 2 миллиона вызовов этой функции занимают очень много времени (мои последние несколько опытов показывали ~ 12 секунд). Мне удалось сделать немного лучше с этим:

def return_slice(self, k):
    pos = slice(self.C[k], self.C[k + 1])
    return (self.A[pos], self.B[pos])

, что снижает это до ~ 6 секунд. Это все еще немного неприемлемо для меня ... Я чувствую, что должен изменить всю структуру моих массивов, но я задаю вам этот вопрос, потому что может быть что-то, что я упускаю из-за того, почему это так медленно .

Имейте в виду, что «структура» для значений k не может быть принята, просто предположите, что она случайна для каждого выполнения.

Я также думаю, что создание кортежа перед возвратом Может быть, здесь проблема, но для ее устранения потребуется тонна работы - я бы предпочел изучить другие альтернативы.

Редактировать: A и B имеют одинаковый размер, но не одинаковый тип данных .

Answer 1

Давайте рассмотрим некоторые вариации срезов:

In [447]: A = np.ones(10000)                                                                           
In [448]: timeit A[24:500]                                                                             

285 ns ± 0.807 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)
In [449]:                                                                                              
In [449]: C = np.array([24, 500])                                                                      
In [450]: timeit A[C[0]:C[1]]                                                                          
632 ns ± 3.99 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)
In [451]: def foo(k): 
     ...:     pos = slice(C[k],C[k+1]) 
     ...:     return A[pos] 
     ...:                                                                                              
In [452]: timeit foo(0)                                                                                
989 ns ± 4.33 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)
In [453]: def foo(k): 
     ...:     pos = slice(C[k],C[k+1]) 
     ...:     return A[pos], A[pos] 
     ...:      
     ...:                                                                                              
In [454]: timeit foo(0)                                                                                
1.31 µs ± 30 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)

Таким образом, вызов foo 2 миллиона раз займет более 2 секунд.

Обычно, когда я делаю тестирование по времени, время под микросекунда выглядит хорошо, если только это не тривиальная операция. Ключ к скорости в numpy заключается в сокращении количества вызовов, а не в ускорении отдельных. «Векторизация» пытается исключить множество вызовов / итераций с помощью операций с целым массивом - один вызов с использованием скомпилированных numpy методов. Это может дать 10 раз или лучше.

numba может переместить нас в скомпилированном направлении. С таким же простым приложением:

In [456]: @numba.njit 
     ...: def foo(k): 
     ...:     pos = slice(C[k],C[k+1]) 
     ...:     return A[pos], A[pos] 

In [459]: timeit foo(0)                                                                                
555 ns ± 4.7 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)

и сбором миллиона таких звонков:

In [473]: timeit [foo(0) for _ in range(1000000)]                                                      
1.09 s ± 38.2 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

Answer 2

Как насчет?

 self.D = np.vstack(self.A, self.B)

 def return_slice(self, k):
     pos = slice(self.C[k], self.C[k + 1])
     return tuple(self.D[:, pos])