массив numpy: эффективный способ вычисления argmax в фиксированном окне для набора строк и столбцов, заданных в качестве входных данных

Question

Предположим, что у нас есть двумерный массив arr, набор позиций, заданных rows и cols, и окно формы (5, 1). Для (i, j) нам нужен индекс максимального значения в пределах arr[i-2:i+2, j]. Мы хотим повторить для всех входных (i, j) пар.

import numpy as np

def get_max_index(arr, i, j, stride):
    nr, nc = arr.shape
    i_low = max(0, i - stride)
    i_up = min(i + stride, nr)

    idx = np.argmax(arr[i_low:i_up, j])

    # transform back to original index.
    return i_low + idx, j

# Given numpy array
arr = np.array([
    [1,2,3,6,7],
    [4,5,6,15,8],
    [7,8,9,24,9],
    [1,1,1,3,10],
    [2,2,2,6,11],
    [3,3,3,9,12],
    [4,4,4,4,42]
    ])

# Rows and columns at which the windows will be centered.
rows = np.array([2, 4, 6, 6])
cols = np.array([1, 1, 3, 4])

# Measure corresponding to window of size 5
stride = 2

# Apply the function on the input rows and cols.
res = [get_max_index(arr, i, j, stride) for i, j in zip(rows, cols)]

assert res == [(2, 1), (2, 1), (5, 3), (6, 4)]

Мне было интересно, есть ли более быстрый numpy способ сделать это вместо использования списка.

У него есть некоторыепохоже на « морфологическое расширение », но здесь оно находится на подмножестве ячеек массива, и нам нужны индексы.

Answer 1

Мы можем использовать np.lib.stride_tricks.as_strided на основе scikit-image's view_as_windows, чтобы получить скользящие оконные представления в версии ввода с минимальным значением (для учета граничных случаев), а затемполучить значения argmax в них и сместить относительно заданных строк.

Следовательно, реализация будет -

from skimage.util.shape import view_as_windows

def windowed_argmax(arr, rows, cols, stride):
    # Get full window extent 
    W = 2*stride+1

    #  Pad with minimum value, so that on boundaries we will skip those
    a = np.pad(arr,((stride,stride),(0,0)),'constant',constant_values=arr.min()-1)

    # Get sliding windows
    w = view_as_windows(a,(W,1))[...,0]

    # Index into those specific rows, cols positions; get argmax, offset back
    return np.c_[rows+w[rows,cols].argmax(1)-stride,cols]

Пример выполнения -

In [75]: arr
Out[75]: 
array([[ 1,  2,  3,  6,  7],
       [ 4,  5,  6, 15,  8],
       [ 7,  8,  9, 24,  9],
       [ 1,  1,  1,  3, 10],
       [ 2,  2,  2,  6, 11],
       [ 3,  3,  3,  9, 12],
       [ 4,  4,  4,  4, 42]])

In [76]: rows
Out[76]: array([2, 4, 6, 6])

In [77]: cols
Out[77]: array([1, 1, 3, 4])

In [78]: windowed_argmax(arr, rows, cols, stride=2)
Out[78]: 
array([[2, 1],
       [2, 1],
       [5, 3],
       [6, 4]])