numpy: векторизованная накопленная сумма по двум осям

Question

Я хотел бы реализовать функциональность numpy.cumsum одновременно по двум осям в векторизованном виде. Таким образом, значение любой ячейки в выходных данных должно быть суммой всех ячеек выше и слева включительно соответствующей ячейки во входных данных.

Это эквивалентно применению cumsum последовательно по первой оси, а затем по другой:

>>> X = np.arange(25).reshape(5,5)
>>> X
array([[ 0,  1,  2,  3,  4],
       [ 5,  6,  7,  8,  9],
       [10, 11, 12, 13, 14],
       [15, 16, 17, 18, 19],
       [20, 21, 22, 23, 24]])
>>> np.cumsum(np.cumsum(X, axis=0),axis=1)
array([[  0,   1,   3,   6,  10],
       [  5,  12,  21,  32,  45],
       [ 15,  33,  54,  78, 105],
       [ 30,  64, 102, 144, 190],
       [ 50, 105, 165, 230, 300]])

Однако я думаю, что это относительно неэффективный способ сделать это.

Возможна ли векторизация в numpy / Python так, что массив нужно суммировать только один раз?

Answer 1

Возможно, есть лучшее решение, но это возможно с numba.stencil:

import numba
# Define some input array
x = np.arange(25).reshape(5,5)
# Pad the array to one cell width with 0
y = np.pad(x, 1, mode='constant', constant_values=0)

# Instantiate the output array
ii = np.zeros_like(y)

@numba.stencil
def integral_image(A, B):
    # A is the input array, B is the output
    return A[0,0] + B[-1,0] + B[0,-1] - B[-1,-1]

# Slice the output to get rid of padded cells
ii = integral_image(y, ii, out=ii)[1:-1,1:-1]

Создает функцию трафарета, которая использует значения из входного изображения вместе с предварительно вычисленными значениями из выходных данных. Выходные данные должны быть созданы в виде массива нулей перед выполнением вычисления и должны быть переданы в функцию трафарета в качестве выходного массива, чтобы функция имела доступ к ранее вычисленным значениям.

Моя основная проблема заключается в том, что требуется скопировать входной массив для заполнения его для использования с numba.stencil. Кажется, что numba разработчики думают о добавлении добавления различных режимов для обработки границ в какой-то момент, но в настоящее время необходимо дополнить ввод, чтобы трафарет был применен к каждой ячейке.

Я также не уверен, насколько это быстрее по сравнению с x.cumsum(0).cumsum(1). Когда я подойду к тестированию производительности с последующими вызовами на np.cumsum, я постараюсь опубликовать спецификации.