python中numpy多维数组的非相邻切片

Question

我有一个多维数组 a：

a = np.random.uniform(1,10,(2,4,2,3,10,10))

对于维度 4-6，我有 3 个列表，其中包含用于切片数组“a”的该维度的索引

dim4 = [0,2]
dim5 = [3,5,9]
dim6 = [1,2,7,8]

如何切出数组“a”以便我得到：

b = a[0,:,0,dim4,dim5,dim6]

因此 b 应该是一个形状为 (4,2,3,4) 的数组，并且包含 a 相应维度的元素。当我尝试上面的代码时，我收到一条错误消息，指出轴 4-6 无法一起广播不同的形状，但如果我这样做：

b = a[0,:,0:2,0:3,0:4]

那么它确实有效，即使切片列表都有不同的长度。那么如何对具有非相邻索引的多维数组进行切片呢？

Answer 1

您可以使用该numpy.ix_函数来构建像这样的复杂索引。它需要一系列array_like，并从中创建一个“开放网格”。文档字符串中的示例非常清楚：

使用ix_一个可以快速构造索引数组来索引叉积。a[np.ix_([1,3],[2,5])]返回数组 [[a[1,2] a[1,5]], [a[3,2] a[3,5]]]。

因此，对于您的数据，您会这样做：

>>> indices = np.ix_((0,), np.arange(a.shape[1]), (0,), dim4, dim5, dim6)
>>> a[indices].shape
(1, 4, 1, 2, 3, 4)

使用以下命令消除 size-1 尺寸np.squeeze：

>>> np.squeeze(a[indices]).shape
(4, 2, 3, 4)