使用索引列表切片 3D numpy 数组

Slicing 3D numpy array using list of index

objective 是使用索引列表对 3D 数组进行切片。

这里,数组的形状是2,5,5。为简单起见,假设索引 0 到 4 标签为 A,B,C,D,E.

假设我们有如下的 3d 数组

array([[[44, 47, 64, 67, 67],
        [ 9, 83, 21, 36, 87],
        [70, 88, 88, 12, 58],
        [65, 39, 87, 46, 88],
        [81, 37, 25, 77, 72]],

       [[ 9, 20, 80, 69, 79],
        [47, 64, 82, 99, 88],
        [49, 29, 19, 19, 14],
        [39, 32, 65,  9, 57],
        [32, 31, 74, 23, 35]]], dtype=int64)

感兴趣的指数是[1,3,4]。同样,我们将其标记为 B,D,E`。根据索引对 3D 数组进行切片时的预期输出如下

array([[[83, 36, 87],
        [39, 46, 88],
        [37, 77, 72]],

       [[64, 99, 88],
        [32,  9, 57],
        [31, 23, 35]]], dtype=int64)

但是,如下所示对数组进行切片

import numpy as np
np.random.seed(0)
arr = np.random.randint(0, 100, size=(2, 5, 5))
k=arr[:,(1,3,4),(1,3,4)]

没有产生预期的输出。

在实际用例中,要切片的元素数量> 3个元素(> B,D,E)。很抱歉没有使用正确的术语

问题

高级索引需要明确索引所有维度。您在这里所做的是沿轴 0.

在每个数组中获取坐标 (1, 1), (3, 3), (4, 4) 处的元素

解决方法

你需要做的是:

idx = (1, 3, 4)  # the indices of interest
arr[np.ix_((0, 1), idx, idx)]

其中 (0, 1) 对应于轴 0 的前两个数组。

输出:

array([[[83, 36, 87],
        [39, 46, 88],
        [37, 77, 72]],

       [[64, 99, 88],
        [32,  9, 57],
        [31, 23, 35]]], dtype=int64)

如上所示,np.ix_((0, 1), idx, idx)) 生成可用于高级索引的对象。 (0, 1) 意味着您显式地 selecting 数组 arr[0]arr[1] 中的元素。如果你有一个更通用的形状 (n, m, q) 的 3D 数组,并且想从轴 0 的每个数组中获取相同的子数组,你可以使用

np.ix_(np.arange(arr.shape[0]), idx, idx))

作为你的指数。请注意,idx 在这里重复,因为您需要那些特定的索引,但通常它们不需要匹配。

泛化

更一般地说,您可以随意切片和切块,如下所示:

In [1]: arrays_to_select = (0, 1)

In [2]: rows_to_select = (1, 3, 4)

In [3]: cols_to_select = (1, 3, 4)

In [4]: indices = np.ix_(arrays_to_select, rows_to_select, cols_to_select)

In [5]: arr[indices]
Out[5]:
array([[[83, 36, 87],
        [39, 46, 88],
        [37, 77, 72]],

       [[64, 99, 88],
        [32,  9, 57],
        [31, 23, 35]]], dtype=int64)

让我们考虑一些其他形状:

In [4]: x = np.random.randint(0, 9, (4, 3, 5))

In [5]: x
Out[5]:
array([[[1, 0, 2, 1, 0],
        [3, 5, 1, 4, 3],
        [1, 8, 1, 4, 2]],

       [[1, 6, 8, 2, 8],
        [0, 0, 4, 2, 3],
        [8, 5, 6, 2, 5]],

       [[4, 4, 8, 6, 0],
        [3, 0, 1, 2, 8],
        [0, 8, 2, 4, 3]],

       [[7, 8, 8, 1, 4],
        [5, 7, 4, 8, 5],
        [7, 5, 5, 3, 4]]])

In [6]: rows = (0, 2)

In [7]: cols = (0, 2, 3, 4)

通过使用 rowscols,您将仅从第 0 行和第 2 行获取由第 0 列到第 4 列的所有元素组成的子数组。让我们验证一下沿轴 0 的第一个数组:

In [8]: arrs = (0,)  # A 1-tuple which will give us only the first array along axis 0

In [9]: x[np.ix_(arrs, rows, cols)]
Out[9]:
array([[[1, 2, 1, 0],
        [1, 1, 4, 2]]])

现在假设您想要由 rowscols 生成的仅沿轴 0 的第一个和最后一个数组的子数组。您可以明确地 select (0, -1):

In [10]: arrs = (0, -1)

In [11]: x[np.ix_(arrs, rows, cols)]
Out[11]:
array([[[1, 2, 1, 0],
        [1, 1, 4, 2]],

       [[7, 8, 1, 4],
        [7, 5, 3, 4]]])

相反,如果您想要来自 all 沿轴 0 的数组的相同子数组:

In [12]: arrs = np.arange(x.shape[0])

In [13]: arrs
Out[13]: array([0, 1, 2, 3])

In [14]: x[np.ix_(arrs, rows, cols)]
Out[14]:
array([[[1, 2, 1, 0],
        [1, 1, 4, 2]],

       [[1, 8, 2, 8],
        [8, 6, 2, 5]],

       [[4, 8, 6, 0],
        [0, 2, 4, 3]],

       [[7, 8, 1, 4],
        [7, 5, 3, 4]]])

试试这个,它与您的 arr[:,idx,idx] 结构相似,但使用的是 np.ix_()。请阅读 documentation for np.ix().-

idx = [1,3,4]
ixgrid = np.ix_(idx,idx)
arr[:,ixgrid[0],ixgrid[1]]
array([[[83, 36, 87],
        [39, 46, 88],
        [37, 77, 72]],

       [[64, 99, 88],
        [32,  9, 57],
        [31, 23, 35]]])

说明

您想要做的是从数组的最后两个轴中提取网格。但是您正在做的是从 2 个轴中的每一个中提取精确的索引。

  • 当您使用 arr[:,(1,3,4),(1,3,4)] 时,您实际上是在从两个矩阵 arr[0]arr[1] 中请求 (1,1)(3,3)(4,4) ]

  • 您需要的是提取网格。这可以通过 np.ix_ 和广播的魔力来实现。

如果你要求...

[[1],
 [3],   and   [1,3,4]
 [4]]

... 这就是 np.ix_ 构造的内容,您广播索引并要求它们之间的叉积,即 (1,1), (1,3), (1,4), (3,1), (3,3)...

希望这能说明您为什么会得到现在的结果,以及您如何才能真正得到您需要的东西。