Можно ли получить доступ к текущим индексам во время операции векторизованного вещания Numpy?

Question

Я хотел бы ускорить функцию в одном массиве в Numpy, используя причудливую индексацию, векторизацию и / или вещание.Для каждого значения в моем массиве мне нужно сделать расчет, который включает смежные значения.Поэтому в моей векторизованной операции мне нужно иметь доступ к текущему индексу, чтобы я мог получать индексы вокруг него.Рассмотрим следующую простую операцию с массивом:

x = np.arange(36).reshape(6, 6)
y = np.zeros((6, 6))
y[:] = x + 1

Я хотел бы использовать подобный синтаксис, но вместо простого приращения я хотел бы сделать что-то вроде добавления всех значений в соседних индексах к текущему значению.в векторизованном цикле.Например, если область вокруг индекса [i, j] == 7 выглядит как

3 2 5
2 7 6
5 5 5

, я бы хотел, чтобы вычисленное значение для [i, j] было 3 + 2 + 5 + 2 + 7 + 6 + 5 + 5 + 5, и я хочу сделать это для всех индексов [i, j].

Это прямой вложенный цикл (или единственный цикл for, использующий np.sum для каждого индекса) ... но я хочу использовать широковещательную и / или необычную индексацию, если это возможно.Это может быть слишком сложной проблемой для синтаксиса Numpy, но я чувствую, что это должно быть возможно.

По сути, это сводится к следующему: как я могу ссылаться на текущий индекс во время операции широковещания?

Answer 1

попробуйте это:

x = np.arange(36).reshape(6, 6)
y = np.zeros((6, 6))
for i in range(x.shape[0]):
    for j in range(x.shape[1]):
        if i>0 and i<x.shape[0]-1 and j>0 and j<x.shape[1]-1:
            y[i,j]=x[i,j]+x[i-1,j]+x[i,j-1]+x[i-1,j-1]+x[i+1,j]+x[i,j+1]+x[i+1,j+1]+x[i-1,j+1]+x[i+1,j-1]
        if j==0:
            if i==0:
                y[i,j]=x[i,j]+x[i,j+1]+x[i+1,j+1]+x[i+1,j]
            elif i==x.shape[0]-1:
                y[i,j]=x[i,j]+x[i,j+1]+x[i-1,j+1]+x[i-1,j]
            else:
                y[i,j]=x[i,j]+x[i,j+1]+x[i+1,j+1]+x[i+1,j]+x[i-1,j]+x[i-1,j+1]

        if j==x.shape[1]-1:
            if i==0:
                y[i,j]=x[i,j]+x[i,j-1]+x[i+1,j-1]+x[i+1,j]
            elif i==x.shape[0]-1:
                y[i,j]=x[i,j]+x[i,j-1]+x[i-1,j-1]+x[i-1,j] 
            else:
                y[i,j]=x[i,j]+x[i,j-1]+x[i-1,j-1]+x[i+1,j]+x[i-1,j]+x[i+1,j-1]
        if i==0 and j in range(1,x.shape[1]-1):
            y[i,j]=x[i,j]+x[i,j-1]+x[i+1,j-1]+x[i+1,j]+x[i+1,j+1]+x[i,j+1]
        if i==x.shape[0]-1 and j in range(1,x.shape[1]-1):
            y[i,j]=x[i,j]+x[i,j-1]+x[i-1,j-1]+x[i-1,j]+x[i-1,j+1]+x[i,j+1]
print(y)

Answer 2

Начните с 1D примера:

x = np.arange(10)

Вам нужно сделать выбор: отбрасывать края или нет, поскольку у них нет двух соседей?Если вы это сделаете, вы можете создать выходной массив по существу за один шаг:

result = x[:-2] + x[1:-1] + x[2:]

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

Если вы предпочитаете сохранять края, вы можете предварительно выделить буфер вывода и добавить непосредственно в него:

result = x.copy()
result[:-1] += x[1:]
result[1:] += x[:-1]

Фундаментальная идея в обоих случаях заключается в том, что для применения операции к всем соседним элементам вы просто сдвигаете массив на +/- 1.Вам не нужно знать какие-либо показатели или делать что-то необычное.Чем проще, тем лучше.

Надеюсь, вы увидите, как обобщить это в 2D-случай.Вместо одного индекса, сдвигающегося между -1, 0, 1, у вас есть два индекса в каждой возможной комбинации -1, 0, 1 между двумя из них.

Приложение

Вот обобщенный подход для результата no-egde:

from itertools import product
def sum_shifted(a):
    result = np.zeros(tuple(x - 2 for x in a.shape), dtype=a.dtype)
    for index in product([slice(0, -2), slice(1, -1), slice(2, None)], repeat=a.ndim):
        result += a[index]
    return result

Эта реализация несколько зачаточна, потому что она не проверяет входные данные без измерений или форм <2, но работает дляпроизвольное число измерений. </p>

Обратите внимание, что для одномерного случая цикл будет выполняться ровно три раза, для двумерного девяти раз и для ND 3N.Это один из случаев, когда я нахожу явный цикл for подходящим для numpy.Цикл очень маленький по сравнению с работой над большим массивом, достаточно быстрый для небольшого массива и, безусловно, лучше, чем выписывание всех 27 вариантов вручную для трехмерного случая.

Еще одна вещь, на которую следует обратить вниманиеЭто то, как генерируются последовательные индексы.В Python индекс с двоеточием, например x[1:2:3], преобразуется в относительно неизвестный slice объект: slice(1, 2, 3).Поскольку (почти) все с запятыми интерпретируется как кортеж, индекс, подобный выражению x[1:2, ::-1, :2], в точности эквивалентен (slice(1, 2), slice(None, None, -1), slice(None, 2)).Цикл генерирует именно такое выражение с одним элементом для каждого измерения.Таким образом, в результате получается простое индексирование по всем измерениям.

Аналогичный подход возможен, если вы хотите сохранить ребра.Единственное существенное отличие состоит в том, что вам нужно индексировать как входные, так и выходные массивы:

from itertools import product
def sum_shifted(a):
    result = np.zeros_like(a)
    for r_index, a_index in zip(product([slice(0, -1), slice(None), slice(1, None)], repeat=a.ndim),
                                product([slice(1, None), slice(None), slice(0, -1)], repeat=a.ndim)):
        result[r_index] += a[a_index]
    return result

Это работает, потому что itertools.product гарантирует порядок итерации, поэтому два сжатыхитераторы останутся в шаге.