什么是为numpy ndarray寻找最大值及其移动子数组的索引的pythonic方法？

Question

我有可能是三维或四维空间。我希望在具有指定步幅的移动子数组窗口中找到最大值及其索引。

例如，假设我有一个4x42d数组，并且为了简单起见，我的移动子数组窗口为2x2，带2步长：

[[ 1, 2, 3, 4], 
 [ 5, 6, 7, 8], 
 [ 9,10,11,12], 
 [13,14,15,16]].

我想找到

[[ 6  8],
 [14 16]]

对于最大值和

[(1,1), (3,1),
 (3,1), (3,3)]

作为输出的指数。

对于不使用循环的ndarray，是否有一个简洁、高效的实现？

Answer 1

这里有一个使用stride_tricks的解决方案

def make_panes(arr, window):
    arr = np.asarray(arr)
    r,c = arr.shape
    s_r, s_c = arr.strides
    w_r, w_c = window

    if c % w_c != 0 or r % w_r != 0:
        raise ValueError("Window doesn't fit array.")

    shape = (r / w_r, c / w_c, w_r, w_c)
    strides = (w_r*s_r, w_c*s_c, s_r, s_c)

    return np.lib.stride_tricks.as_strided(arr, shape, strides)

def max_in_panes(arr, window):
    w_r, w_c = window
    r, c = arr.shape
    panes = make_panes(arr, window)
    v = panes.reshape((-1, w_r * w_c))

    ix = np.argmax(v, axis=1)
    max_vals = v[np.arange(r/w_r * c/w_c), ix]

    i = np.repeat(np.arange(0,r,w_r), c/w_c)
    j = np.tile(np.arange(0, c, w_c), r/w_r)
    rel_i, rel_j = np.unravel_index(ix, window)
    max_ix = i + rel_i, j + rel_j

    return max_vals, max_ix

演示：

>>> vals, ix = max_in_panes(x, (2,2))
>>> print vals
[[ 6  8]
 [14 16]]
>>> print ix
(array([1, 1, 3, 3]), array([1, 3, 1, 3]))

请注意，这是相当未经测试的，并设计用于处理2d数组。我将把泛化问题留给读者.