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序列，并从它们生成一个“开放网格”。docstring中的示例非常清楚：

使用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消除尺寸-1维

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