NumPy提取沿一个轴任意子数组

Question

例如，我有一个3维数组：

[[[ 0  1  2  3]
  [ 4  5  6  7]
  [ 8  9 10 11]]

 [[12 13 14 15]
  [16 17 18 19]
  [20 21 22 23]]

 [[24 25 26 27]
  [28 29 30 31]
  [32 33 34 35]]]

我想要的最终数组：

[[[ 0  1]
  [ 4  5]]

 [[18 19]
  [22 23]]

 [[26 27]
  [30 31]]]

有没有一种更有效的方法来获取数组，而不使用for循环？

提出这个问题的原因是，如果我们想要获取单个轴上的任意元素，例如：

[[ 0  1  2  3]
 [ 4  5  6  7]
 [ 8  9 10 11]]

我可以使用 a[np.arange(a.shape[0]), [2, 3, 1]] 获得数组 [ 2 7 9]。那么当元素变成子数组时，是否有类似的方法呢？

- Jeff Dong

---

这不是切片。在切片中，您会在整个数组中切割。在这里，您将数组的不同部分粘合在一起。 - zegkljan

@zegkljan：谢谢您指出。我已经更改了我的表达方式。 - Jeff Dong

Answer 1

一个简单的想法可能是 a[[0,1,2],[0:2,1:3,0:2],[0:2,2:4,2:4]] ，但它并没有被实现。

可以通过np.lib.stride_tricks.as_strided获得解决方法。只需定义：

ast=np.lib.stride_tricks.as_strided(a,a.shape*2,a.strides*2)
#ast.shape is (3, 3, 4, 3, 3, 4).

然后您可以分别定义块的起始位置和大小:

In [4]: ast[[0,1,2],[0,1,0],[0,2,2],0,:2,:2]
Out[4]: 
array([[[ 0,  1],
        [ 4,  5]],

       [[18, 19],
        [22, 23]],

       [[26, 27],
        [30, 31]]])

一些解释：

• 原始数组：

你想找到以元素0,18,26开头的块。

它们在重塑后的数组中的索引可以通过以下方式找到：

In [316]: np.unravel_index([0,18,26],a.shape)
Out[316]: 
(array([0, 1, 2], dtype=int64),
 array([0, 1, 0], dtype=int64),
 array([0, 2, 2], dtype=int64))

ast [[0,1,2]，[0,1,0]，[0,2,2]]是一个(3,3,3,4)的数组。每个(3,3,4)的数组都以一个选定的元素开头。

array([[[[          0,           1,           2,           3],
         [          4,           5,           6,           7],
         [          8,           9,          10,          11]],

        [[         12,          13,          14,          15],
         [         16,          17,          18,          19],
         [         20,          21,          22,          23]],

        [[         24,          25,          26,          27],
         [         28,          29,          30,          31],
         [         32,          33,          34,          35]]],


       [[[         18,          19,          20,          21],
         [         22,          23,          24,          25],
         [         26,          27,          28,          29]],

        [[         30,          31,          32,          33],
         [         34,          35,    23592960,       18335],
         [  697780028, -2147480064,   540876865,  1630433390]],

        [[ 2036429426,   538970664,   538976288,  1532698656],
         [  741355058,   808334368,   775168044,   874523696],
         [  744304686,   538976266,   538976288,   811278368]]],


       [[[         26,          27,          28,          29],
         [         30,          31,          32,          33],
         [         34,          35,    23592960,       18335]],

        [[  697780028, -2147480064,   540876865,  1630433390],
         [ 2036429426,   538970664,   538976288,  1532698656],
         [  741355058,   808334368,   775168044,   874523696]],

        [[  744304686,   538976266,   538976288,   811278368],
         [  539766830,   741355058,   808333600,   775036972],
         [  170679600,   538976288,   538976288,   774920992]]]])

如文档所述，as_strided是一种危险的技巧，必须谨慎使用，因为如果使用不当，它可以访问不在数组中的元素。下一步将确保选择有效的元素。

• size：

有趣的元素是每个第一个块左上角的四个元素。因此，ast[[0,1,2],[0,1,0],[0,2,2],0,:2,:2]会选择它们。

您还可以像这样定义大小的块：bloc122=ast[...,0,:2,:2]（bloc122.shape为(3, 3, 4, 2, 2)）：

In [8]: bloc122[[0,1,2],[0,1,0],[0,2,2]]=0

In [9]: a
Out[9]: 
array([[[ 0,  0,  2,  3],
        [ 0,  0,  6,  7],
        [ 8,  9, 10, 11]],

       [[12, 13, 14, 15],
        [16, 17,  0,  0],
        [20, 21,  0,  0]],

       [[24, 25,  0,  0],
        [28, 29,  0,  0],
        [32, 33, 34, 35]]])

- B. M.

---

这个 as_strided 能用，但很难理解，看起来很危险。显示整个 ast 数组会显示很多原始 a.data 缓冲区之外的垃圾值。 - hpaulj

是的，我添加了一些解释。就像ix_和as_strided一样，需要一些经验才能理解它们 ;)。 - B. M.

通过将as_strided语句包装起来，使其始终提供块的有效视图，可以通过将形状的一半设置为块大小，另一半设置为形状减去块大小来解决这个问题。 - Eelco Hoogendoorn

你使用as_strided和我之间的一个关键区别是，我只扩展一个维度，即具有重叠索引的维度。在两个或更多维度上扩展会导致不规则或对角边界。 - hpaulj

Answer 2

在第二维上重叠的块需要使用类似于as_strided的东西。B.M.先前用as_strided解决了这个问题，但我发现很难理解。它看起来也很危险，因为显示了越界数据。我正在分步骤接近任务。

对于“cols”的最后一维进行选择[0,1]和[2,3]相对容易。通过变形使其更容易。

In [27]: A=np.arange(36).reshape(3,3,4)
In [28]: A1=A.reshape(3,3,2,2)
In [29]: A2=A1[[0,1,2],:,[0,1,1],:]

In [30]: A2
Out[30]: 
array([[[ 0,  1],
        [ 4,  5],
        [ 8,  9]],

       [[14, 15],
        [18, 19],
        [22, 23]],

       [[26, 27],
        [30, 31],
        [34, 35]]])

针对一个子数组，我发现我可以将其视为2个重叠的数组：

In [59]: as_strided(A2[0],shape=(2,2,2),strides=(8,8,4))
Out[59]: 
array([[[0, 1],
        [4, 5]],

       [[4, 5],
        [8, 9]]])

np.lib.stride_tricks.as_strided是一个难以使用的函数。我看到它主要用于移动窗口的应用程序。

应用于整个数组：

In [65]: A3=as_strided(A2,shape=(3,2,2,2),strides=(24,8,8,4))

In [66]: A3
Out[66]: 
array([[[[ 0,  1],
         [ 4,  5]],
         ...
        [[30, 31],
         [34, 35]]]])

并且目标:

In [71]: A3[[0,1,2],[0,1,0]]
Out[71]: 
array([[[ 0,  1],
        [ 4,  5]],

       [[18, 19],
        [22, 23]],

       [[26, 27],
        [30, 31]]])

这可以通过一种方式组合在一起，允许赋值（形状和步幅从A1的值适应）。

In [105]: A1 = A.reshape(3,3,2,2)
In [106]: A1s = as_strided(A1, shape=(3,2,2,2,2), strides=(48,16,16,8,4))

In [107]: A1s[[0,1,2],[0,1,0],:,[0,1,1],:]
Out[107]: 
array([[[ 0,  1],
        [ 4,  5]],

       [[18, 19],
        [22, 23]],

       [[26, 27],
        [30, 31]]])

分配测试：

In [108]: A1s[[0,1,2],[0,1,0],:,[0,1,1],:] = 99

In [109]: A
Out[109]: 
array([[[99, 99,  2,  3],
        [99, 99,  6,  7],
        [ 8,  9, 10, 11]],

       [[12, 13, 14, 15],
        [16, 17, 99, 99],
        [20, 21, 99, 99]],

       [[24, 25, 99, 99],
        [28, 29, 99, 99],
        [32, 33, 34, 35]]])

---

如果您不需要分配任务，那么在没有 striding 的情况下理解逻辑会更容易：

np.array([A2[0,:-1], A2[1,1:], A2[2,:-1]])

slices=(slice(-1),slice(1,None))    
np.array([A2[i,slices[j]] for i,j in zip([0,1,2],[0,1,0])])

===========================

之前在寻找答案时遇到的问题：

我认为可以通过高级索引来提取它。

Ind = np.ix_([0,0,1,1,2,2], [0,1,1,2,0,1], [0,1,2,3,2,3])
a[ind]

这些值可能需要微调，因为我无法在这里测试。

思路是枚举每个维度中需要哪些行和列，并使用ix_（我认为这是正确的函数）添加newaxis以便它们一起广播。

您对2d情况进行泛化的想法是正确的。诀窍是弄清楚何时需要 np.array([0,1,2])，以及何时需要使用ix_或[:,None]等旋转它。最终我会尝试各种想法。

它可能需要进行调整。

  Ind = np.ix_([0,1,2], [0,1,1,2,0,1], [0,1,2,3,2,3])

这不完全正确。它将生成一个3x6x6的数组; 您需要的是3x2x2的数组。

将其重新调整为3x3x2x2可能会更容易地索引最后一维。

另一个答案提到的步幅技巧可能有助于将第二个重叠选择转换为类似的2x2块。 但我需要尝试一下。

我想象中的索引元组形式为([1,2,3],[?,?,?],:,[?,?,?],:)

- hpaulj

Answer 3

你可以使用索引来单独获取预期的输出：

>>> a = np.array([[[ 0, 1, 2, 3],
...   [ 4, 5, 6, 7],
...   [ 8, 9,10,11]],
... 
...  [[12,13,14,15],
...   [16,17,18,19],
...   [20,21,22,23]],
... 
...  [[24,25,26,27],
...   [28,29,30,31],
...   [32,33,34,35]]])

>>> a[0,:2,:2]
array([[0, 1],
       [4, 5]])
>>> a[1,1:,2:]
array([[18, 19],
       [22, 23]])
>>> a[2,:2,2:]
array([[26, 27],
       [30, 31]])
>>> 

- Mazdak

---

这是正确的，但只是一个for循环。我要找的是一个单一的表达式a[something]，这样当我想要为这些元素分配值时，我只需使用a[something] = b。 - Jeff Dong

Answer 4

当我提出这个问题时，我没有想到会使用像 hpaulj 和 B. M. 提到的 ix_ 和 as_strided 这样棘手的方法，对于像我这样的初学者来说很难理解。

我想出了一种更易懂但不太高效的方法，受到了 B. M. 的启发。

结合 高级索引，我们只需要将a[[0,1,2],[0:2,1:3,0:2],[0:2,2:4,2:4]]转换为a的高级索引。在我们的情况下是：

a[[0 0 0 0 1 1 1 1 2 2 2 2],
  [0 0 1 1 1 1 2 2 0 0 1 1],
  [0 1 0 1 2 3 2 3 2 3 2 3]]

一个简单的翻译函数可以像这样：

（代码中的笛卡尔函数来自于pv.对this问题的回答）

# a[[0,1,2],[0:2,1:3,0:2],[0:2,2:4,2:4]]
def translate_idx(idx1=[0, 1, 2], idx2=[(0,2),(1,3),(0,2)], idx3=[(0,2),(2,4),(2,4)]):
    # first we need to get the combinations for correponding intervals
    # e.g for (1,3) and (2,4) we get [[1, 2], [1, 3], [2, 2], [2, 3]]
    idx23 = []
    for i in range(len(idx2)):
        # the function cartesian here just generates all conbinations from some arrays
        idx23.append(cartesian((np.arange(*idx2[i]), np.arange(*idx3[i]))))
    idx23 = np.concatenate(idx23).T
    # now we get index for 2nd and 3rd axis
    # [[0 0 1 1 1 1 2 2 0 0 1 1]
    #  [0 1 0 1 2 3 2 3 2 3 2 3]]
    # we can repeat 4 times for idx1 and append idx23
    step = (idx2[0][1] - idx2[0][0]) * (idx3[0][1] - idx3[0][0])
    idx123 = np.append(np.array(idx1).repeat(step).reshape((1, -1)), idx23, axis=0)
    return idx123

然后我可以使用

idx = translate_idx()
a[idx[0], idx[1], idx[2]]

为了得到我想要的。

虽然这种方法并不是很高效，但我认为它揭示了这类问题与高级索引之间的一些关联。在实践中，as_strided 绝对是一个更好的选择。

非常感谢大家提供的详细答案 :)

- Jeff Dong