Найти диагонали суммы в NumPy (быстрее)

Question

у меня есть некоторые board Numpy массивов как то:

array([[0, 0, 0, 1, 0, 0, 0, 0], 
     [1, 0, 0, 0, 0, 1, 0, 1], 
     [0, 0, 0, 0, 0, 0, 0, 1], 
     [0, 1, 0, 1, 0, 0, 0, 0], 
     [0, 0, 0, 0, 0, 0, 0, 1], 
     [0, 0, 0, 0, 1, 0, 0, 0], 
     [0, 0, 1, 0, 0, 0, 1, 0], 
     [1, 0, 0, 0, 0, 1, 0, 0]]) 

И я использую следующий код, чтобы найти сумму элементов по каждому й диагональным от -7 до 8 (и его зеркальной версии).

n = 8 
rate = [b.diagonal(i).sum() 
     for b in (board, board[::-1]) 
     for i in range(-n+1, n)] 

После некоторого профилирования, эта операция занимает около 2/3 общего времени работы, и это, как представляется, из-2 фактора:

• Метод .diagonal создает новый массив вместо того, чтобы в целях (выглядит как Numpy 1.7 будет иметь новый метод .diag решить, что)
• итерация сделано в питон внутри списка понимания

Итак, есть ли способы найти эти суммы быстрее (возможно, в слое C numpy)?

---

После еще нескольких испытаний, я мог бы уменьшить 7.5x общего времени за счет кэширования этой операции ... Может быть, я искал неправильное узкое место? более

---

одно:

Просто нашли .trace метод, который заменяет diagonal(i).sum() вещи и ... Там не было много улучшений в производительности (примерно от 2 до 4%).

Таким образом, проблема должна заключаться в понимании. Есть идеи?

Answer 1

Возможное решение: stride_tricks. Это частично основано на множестве информации, доступной в ответах на вопрос this question, но проблема, по-моему, не такая уж иная, как дубликат. Вот основная идея, примененная к квадратной матрице - см. Ниже функцию, реализующую более общее решение.

>>> cols = 8 
>>> a = numpy.arange(cols * cols).reshape((cols, cols)) 
>>> fill = numpy.zeros((cols - 1) * cols, dtype='i8').reshape((cols - 1, cols)) 
>>> stacked = numpy.vstack((a, fill, a)) 
>>> major_stride, minor_stride = stacked.strides 
>>> strides = major_stride, minor_stride * (cols + 1) 
>>> shape = (cols * 2 - 1, cols) 
>>> numpy.lib.stride_tricks.as_strided(stacked, shape, strides) 
array([[ 0, 9, 18, 27, 36, 45, 54, 63], 
     [ 8, 17, 26, 35, 44, 53, 62, 0], 
     [16, 25, 34, 43, 52, 61, 0, 0], 
     [24, 33, 42, 51, 60, 0, 0, 0], 
     [32, 41, 50, 59, 0, 0, 0, 0], 
     [40, 49, 58, 0, 0, 0, 0, 0], 
     [48, 57, 0, 0, 0, 0, 0, 0], 
     [56, 0, 0, 0, 0, 0, 0, 0], 
     [ 0, 0, 0, 0, 0, 0, 0, 7], 
     [ 0, 0, 0, 0, 0, 0, 6, 15], 
     [ 0, 0, 0, 0, 0, 5, 14, 23], 
     [ 0, 0, 0, 0, 4, 13, 22, 31], 
     [ 0, 0, 0, 3, 12, 21, 30, 39], 
     [ 0, 0, 2, 11, 20, 29, 38, 47], 
     [ 0, 1, 10, 19, 28, 37, 46, 55]]) 
>>> diags = numpy.lib.stride_tricks.as_strided(stacked, shape, strides) 
>>> diags.sum(axis=1) 
array([252, 245, 231, 210, 182, 147, 105, 56, 7, 21, 42, 70, 105, 
     147, 196]) 

Конечно, я не знаю, как быстро это будет на самом деле. Но я уверен, что это будет быстрее, чем понимание Python.

Для удобства здесь представлена ​​общая функция diagonals. Предполагается, что вы хотите переместить диагональ вдоль самой длинной оси:

def diagonals(a): 
    rows, cols = a.shape 
    if cols > rows: 
     a = a.T 
     rows, cols = a.shape 
    fill = numpy.zeros(((cols - 1), cols), dtype=a.dtype) 
    stacked = numpy.vstack((a, fill, a)) 
    major_stride, minor_stride = stacked.strides 
    strides = major_stride, minor_stride * (cols + 1) 
    shape = (rows + cols - 1, cols) 
    return numpy.lib.stride_tricks.as_strided(stacked, shape, strides) 

Answer 2

Как я писал в комментарии, я бы не стал входить в код C.

Попытайтесь пойти с PyPy. На самом деле это поддержка numpy очень хорошая (однако она не поддерживает напрямую array.diagonal). Я не проверял, есть ли для этого другой метод buidin. бессильных, я попытался следующий код:

try: 
    import numpypy # required by PyPy 
except ImportError: 
    pass 
import numpy 

board = numpy.array([[0, 0, 0, 1, 0, 0, 0, 0], 
     [1, 0, 0, 0, 0, 1, 0, 1], 
     [0, 0, 0, 0, 0, 0, 0, 1], 
     [0, 1, 0, 1, 0, 0, 0, 0], 
     [0, 0, 0, 0, 0, 0, 0, 1], 
     [0, 0, 0, 0, 1, 0, 0, 0], 
     [0, 0, 1, 0, 0, 0, 1, 0], 
     [1, 0, 0, 0, 0, 1, 0, 0]]) 

n=len(board) 
def diag_sum(i, b): 
    s = 0 
    if i>=0: 
     row = 0 
     end = n 
    else: 
     row = -i 
     end = n+i 
     i = 0 
    while i<end: 
     s += b[row, i] 
     i+=1 
     row+=1 
    return s 


import time 
t=time.time() 
for i in xrange(50000): 
    # rate = [b.diagonal(i).sum() 
    #   for b in (board, board[::-1]) 
    #   for i in range(-n+1, n)] 

    rate = [diag_sum(i,b) 
      for b in (board, board[::-1]) 
      for i in range(-n+1, n)] 

print time.time() - t 

Результаты являются:

• 0.64s PyPy с diag_sum
• 6.01s CPython версии с diag_sum
• 5.60s CPython версии с b.diagonal