Cython оптимизировать критическую часть массива NumPy суммирования

Question

Пусть Ь список L = [A_1, A_2, ..., A_n], и каждый из A_i являются numpy.int32 массивы длины 1024.

(Большую часть времени 1000 < < п 4000).

После некоторого профилирования, я видел, что одна часть времени трудоемкой операцией суммирования:

def summation(): 
    # L is a global variable, modified outside of this function 
    b = numpy.zeros(1024, numpy.int32) 
    for a in L: 
     b += a 
    return b 

PS: Я не думаю, что я могу определить 2D массив размером 1024 x n, потому что n не неподвижная : некоторые элементы удаляются/добавляются в L динамически, поэтому len(L) = n может варьироваться от 1000 до 4000 в течение времени работы.

Могу ли я получить значительное улучшение с помощью Cython? Если да, то как я должен Cython-перекодировать эту маленькую функцию (я не должен добавить cdef печатать?)

Или вы можете увидеть некоторые другие возможные улучшения?

Answer 1

Вот код Cython, убедитесь, что каждый массив в L является C_CONTIGUOUS:

import cython 
import numpy as np 
cimport numpy as np 

@cython.boundscheck(False) 
@cython.wraparound(False) 
def sum_list(list a): 
    cdef int* x 
    cdef int* b 
    cdef int i, j 
    cdef int count 
    count = len(a[0]) 
    res = np.zeros_like(a[0]) 
    b = <int *>((<np.ndarray>res).data) 
    for j in range(len(a)): 
     x = <int *>((<np.ndarray>a[j]).data) 
     for i in range(count): 
      b[i] += x[i] 
    return res 

Один мой PC это примерно 4 раза быстрее.