python - numpy argmin 向量化

Question

我正在尝试迭代 numpy 行，并将包含最低值的每个簇的 3 个元素的索引放入另一行中。这应该是在左、中、右的背景下；左右边缘只查看两个值(“左和中”或“中和右”)，但中间的所有内容都应该查看所有 3 个值。

For 循环可以简单地完成此操作，但速度非常慢。某种 numpy 矢量化可能会加快速度。

例如:

 [1 18 3 6 2]
 # should give the indices...
 [0 0 2 4 4] # matching values 1 1 3 2 2

实现的缓慢 for 循环:

for y in range(height):
    for x in range(width):
        i = 0 if x == 0 else x - 1
        other_array[y,x] = np.argmin(array[y,i:x+2]) + i

Answer 1

注意:请参阅下面的更新，了解没有 for 循环的解决方案。

这适用于任意维数的数组:

def window_argmin(arr):
    padded = np.pad(
        arr,
        [(0,)] * (arr.ndim-1) + [(1,)],
        'constant',
        constant_values=np.max(arr)+1,
    )
    slices = np.concatenate(
        [
            padded[..., np.newaxis, i:i+3]
            for i in range(arr.shape[-1])
        ],
        axis=-2,
    )
    return (
        np.argmin(slices, axis=-1) + 
        np.arange(-1, arr.shape[-1]-1)
    )

代码使用 np.pad 在数组的最后一个维度上向左填充一个额外的数字，向右填充一个额外的数字，因此我们始终可以使用 3 个元素的窗口作为 argmin。它将额外的元素设置为 max+1，因此它们永远不会被 argmin 选择。

然后，它使用切片列表的 np.concatenate 为每个 3 元素窗口添加一个新维度。这是我们唯一使用 for 循环的地方，并且我们只在最后一个维度上循环一次，以创建单独的三元素窗口。 (请参阅下面的更新，了解删除此 for 循环的解决方案。)

最后，我们在每个窗口上调用 np.argmin。

我们需要调整它们，我们可以通过添加窗口第一个元素的偏移量来完成(对于第一个窗口来说，实际上是 -1，因为它是一个填充元素。)我们可以用一个简单的方法来进行调整arange 数组的总和，适用于广播。

这是对示例数组的测试:

>>> x = np.array([1, 18,  3,  6,  2])

>>> window_argmin(x)

array([0, 0, 2, 4, 4])

还有一个 3d 示例:

>>> z

array([[[ 1, 18,  3,  6,  2],
        [ 1,  2,  3,  4,  5],
        [ 3,  6, 19, 19,  7]],

       [[ 1, 18,  3,  6,  2],
        [99,  4,  4, 67,  2],
        [ 9,  8,  7,  6,  3]]])

>>> window_argmin(z)

array([[[0, 0, 2, 4, 4],
        [0, 0, 1, 2, 3],
        [0, 0, 1, 4, 4]],

       [[0, 0, 2, 4, 4],
        [1, 1, 1, 4, 4],
        [1, 2, 3, 4, 4]]])

<小时/>

更新:这是使用 stride_tricks 的版本不使用任何 for 循环:

def window_argmin(arr):
    padded = np.pad(
        arr,
        [(0,)] * (arr.ndim-1) + [(1,)],
        'constant',
        constant_values=np.max(arr)+1,
    )
    slices = np.lib.stride_tricks.as_strided(
        padded,
        shape=arr.shape + (3,),
        strides=padded.strides + (padded.strides[-1],),
    )
    return (
        np.argmin(slices, axis=-1) + 
        np.arange(-1, arr.shape[-1]-1)
    )

帮助我想出跨步技巧解决方案的是 this numpy issue要求添加滑动窗口功能，链接到example implementation of it ，所以我只是针对这个具体情况进行了调整。这对我来说仍然很神奇，但它确实有效。 😁

针对不同维数的数组进行了测试并按预期工作。