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)

使用 np.squeeze 摆脱 size-1 维度：

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