我可以从 numpy 数组的最后一个维度 select 任意 windows 吗？

Question

我想编写一个 numpy 函数，它接受一个 MxN 数组 A、一个 window 长度 L 和一个起始索引的 MxP 数组 idxs进入 A 的 M 行，从 A 的 M 行中的每一行中选择 P 个长度为 L 的任意切片。除了，我希望它能在 A 的最后一个维度上工作，而不一定关心 A 有多少维度，所以 A 和 idxs 的所有维度都匹配除了最后一个。示例：

如果 A 只是一维：

A = np.array([1, 2, 3, 4, 5, 6])

window_len = 3

idxs = np.array([1, 3])

result = magical_routine(A, idxs, window_len)

其中 result 是一个 2x3 数组，因为我选择了 len 3 的 2 个切片：

np.array([[ 2, 3, 4],
          [ 4, 5, 6]])

如果 A 是二维的：

A = np.array([[ 1, 2, 3, 4, 5, 6],
              [ 7, 8, 9,10,11,12],
              [13,14,15,16,17,18]])

window_len = 3
idxs = np.array([[1, 3],
                 [0, 1],
                 [2, 2]])

result = magical_routine(A, idxs, window_len)

其中 result 是一个 3x2x3 数组，因为有 3 行 A，我从每行中选择了 2 片 len 3:

np.array([[[ 2, 3, 4], [ 4, 5, 6]],
          [[ 7, 8, 9], [ 8, 9,10]],
          [[15,16,17], [15,16,17]]])

以此类推

我发现了很多执行此操作的低效方法，以及适用于 A 特定维数的方法。对于 2D，以下内容非常整洁：

col_idxs = np.add.outer(idxs, np.arange(window_len))
np.take_along_axis(A[:, np.newaxis], col_idxs, axis=-1)

不过我看不出有什么好的方法可以将其推广到 1D 和其他 D...

有没有人知道一种推广到任意数量暗淡的有效方法？

Answer 1

对于您的第一个案例

In [271]: A=np.arange(1,7)
In [272]: idxs = np.array([1,3])

使用这个问题通常得到的迭代类型：

In [273]: np.vstack([A[i:i+3] for i in idxs])
Out[273]: 
array([[2, 3, 4],
       [4, 5, 6]])

或者生成所有索引和一个索引。 linspace 对此很方便（虽然它不是唯一的选择）：

In [278]: j = np.linspace(idxs,idxs+3,3,endpoint=False)
In [279]: j
Out[279]: 
array([[1., 3.],
       [2., 4.],
       [3., 5.]])
In [282]: A[j.T.astype(int)]
Out[282]: 
array([[2, 3, 4],
       [4, 5, 6]])

为 2d

In [284]: B
Out[284]: 
array([[ 1,  2,  3,  4,  5,  6],
       [ 7,  8,  9, 10, 11, 12],
       [13, 14, 15, 16, 17, 18]])

In [285]: idxs = np.array([[1, 3],
     ...:                  [0, 1],
     ...:                  [2, 2]])

In [286]: j = np.linspace(idxs,idxs+3,3,endpoint=False)

In [287]: j
Out[287]: 
array([[[1., 3.],
        [0., 1.],
        [2., 2.]],

       [[2., 4.],
        [1., 2.],
        [3., 3.]],

       [[3., 5.],
        [2., 3.],
        [4., 4.]]])

经过一些反复试验，配对指数得到：

In [292]: B[np.arange(3)[:,None,None],j.astype(int).transpose(1,2,0)]
Out[292]: 
array([[[ 2,  3,  4],
        [ 4,  5,  6]],

       [[ 7,  8,  9],
        [ 8,  9, 10]],

       [[15, 16, 17],
        [15, 16, 17]]])

或者像第一种情况一样迭代，但是多了一层：

In [294]: np.array([[B[j,i:i+3] for i in idxs[j]] for j in range(3)])
Out[294]: 
array([[[ 2,  3,  4],
        [ 4,  5,  6]],

       [[ 7,  8,  9],
        [ 8,  9, 10]],

       [[15, 16, 17],
        [15, 16, 17]]])

滑动windows:

In [295]: aa = np.lib.stride_tricks.sliding_window_view(A,3)

In [296]: aa.shape
Out[296]: (4, 3)

In [297]: aa
Out[297]: 
array([[1, 2, 3],
       [2, 3, 4],
       [3, 4, 5],
       [4, 5, 6]])

In [298]: aa[[1,3]]
Out[298]: 
array([[2, 3, 4],
       [4, 5, 6]])

和

In [300]: bb = np.lib.stride_tricks.sliding_window_view(B,(1,3))

In [301]: bb.shape
Out[301]: (3, 4, 1, 3)

In [302]: bb[np.arange(3)[:,None],idxs,0,:]
Out[302]: 
array([[[ 2,  3,  4],
        [ 4,  5,  6]],

       [[ 7,  8,  9],
        [ 8,  9, 10]],

       [[15, 16, 17],
        [15, 16, 17]]])

Answer 2

我明白了！我快到了：

def magical_routine(A, idxs, window_len=2000):
    col_idxs = np.add.outer(idxs, np.arange(window_len))
    return np.take_along_axis(A[..., np.newaxis, :], col_idxs, axis=-1)

我只需要始终将新轴添加到 A 的倒数第二个暗淡的位置，然后不理会其余轴。

我可以从 numpy 数组的最后一个维度 select 任意 windows 吗？

Can I select arbitrary windows from the last dimension of a numpy array?

numpy