如何在numpy中循环回到数组的开头以获取越界索引？

Question

我有一个 2D numpy 数组，我想从中提取一个子矩阵。 我通过如下切片数组来获得子矩阵。 在这里，我想要一个 3*3 子矩阵围绕索引为 (2,3) 的项目。

>>> import numpy as np
>>> a = np.array([[0, 1, 2, 3],
...        [4, 5, 6, 7],
...        [8, 9, 0, 1],
...        [2, 3, 4, 5]])
>>> a[1:4, 2:5]
array([[6, 7],
   [0, 1],
   [4, 5]])

但我想要的是，对于超出范围的索引，它会回到数组的开头并从那里继续。这是我想要的结果：

array([[6, 7, 4],
   [0, 1, 8],
   [4, 5, 2]])

我知道我可以做一些事情，比如将索引的 mod 获取到数组的宽度；但我正在寻找一个可以做到这一点的 numpy 函数。 而且对于一维数组，这将导致索引超出范围错误，这并不是真正有用...

Answer 1

这是使用带有环绕模式的np.pad 的一种方式。

>>> a = np.array([[0, 1, 2, 3],
                  [4, 5, 6, 7],
                  [8, 9, 0, 1],
                  [2, 3, 4, 5]])

>>> pad_width = 1
>>> i, j = 2, 3

>>> startrow, endrow = i-1+pad_width, i+2+pad_width # for 3 x 3 submatrix
>>> startcol, endcol = j-1+pad_width, j+2+pad_width

>>> np.pad(a, (pad_width, pad_width), 'wrap')[startrow:endrow, startcol:endcol]
array([[6, 7, 4],
       [0, 1, 8],
       [4, 5, 2]])

根据补丁的形状（例如 5 x 5 而不是 3 x 3），您可以增加 pad_width 并相应地增加开始和结束行和列的索引。

Answer 2

np.take 确实有一个 mode 参数，它可以环绕超出范围的索引。但是将np.take 用于多维数组有点麻烦，因为axis 必须是标量。

但是，在您的特定情况下，您可以这样做：

a = np.array([[0, 1, 2, 3],
              [4, 5, 6, 7],
              [8, 9, 0, 1],
              [2, 3, 4, 5]])

np.take(a, np.r_[2:5], axis=1, mode='wrap')[1:4]

输出：

array([[6, 7, 4],
       [0, 1, 8],
       [4, 5, 2]])

编辑

这个函数可能就是你要找的（？）

def select3x3(a, idx):
    x,y = idx
    return np.take(np.take(a, np.r_[x-1:x+2], axis=0, mode='wrap'), np.r_[y-1:y+2], axis=1, mode='wrap')

但回想起来，我建议对这种操作使用模数和花式索引（这基本上是 mode='wrap' 在内部所做的）：

def select3x3(a, idx):
    x,y = idx
    return a[np.r_[x-1:x+2][:,None] % a.shape[0], np.r_[y-1:y+2][None,:] % a.shape[1]]

上述解决方案也适用于a 上的任何二维形状。