numpy array: efficient way to compute argmax within a fixed window at a set of rows and columns given as input

user_650885 0
Viewed 352

Lets assume we have a 2D array arr, a set of positions defined by rows and cols, and a window of shape (5, 1). For (i, j), we need the index of the max value within arr[i-2:i+2, j]. We want to repeat for all input (i, j) pairs.

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)]

I was curious if there's a faster numpy way of doing this instead of using list comprehension.

It has some semblance to "morphological dilation" but here it's on a subset of the array cells and we want the indices.

numpy arrays python
1 Answers

We can leverage np.lib.stride_tricks.as_strided based scikit-image's view_as_windows to get sliding windowed views into a minimum-value-padded version of the input (to account for boundary cases) and then get argmax values in those and offset against the given rows.

Hence, the implementation would be -

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]

Sample run -

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]])
user_32938810
Related