Numpy索引异常:如何从多维数组中进行子选择并保持所有轴

Question

我有一个多维数组，并且有两个整数列表L_i和L_j，对应于我想要保留的轴I和轴j上的元素。我还想满足以下条件:

1. 保持数组的原始维度，即使是L_i或L_j只包含一个元素(换句话说，我不想要一个单例)要折叠的轴)
2. 保持轴的顺序

最干净的方法是什么?

下面是一个可复制的例子，显示了我得到的一些意想不到的行为:

import numpy as np
aa = np.arange(120).reshape(5,4,3,2)
aa.shape
### (5,4,3,2) as expected
aa[:,:,:,[0,1]].shape
### (5, 4, 3, 2) as expected
aa[:,:,:,[0]].shape
### (5,4,3,1) as desired. Notice that even though the [0] is one element, 
### that last axis is preserved, which is what I want
aa[:,[1,3],:,[0]].shape
### (2, 5, 3) NOT WHAT I EXPECTED!!
### I was expecting (5, 2, 3, 1)

很好奇为什么numpy会折叠和重新排序轴，也是正确进行子集的最佳方法。

Answer 1

关于你的问题的答案…

• 为什么numpy在折叠坐标轴?

由于高级索引[1,3]和[0]一起广播形成形状为(2,)的子空间，取代了它们索引的子空间(即大小分别为4和2的轴)。

• 为什么numpy要重新排序坐标轴?

因为高级索引被切片分隔，所以没有明确的位置可以放置新的形状(2,)子空间。结果，numpy把它放在数组的前面，切片的尺寸放在后面(形状为(5, 3))。

…这样你就得到了一个形状为(2, 5, 3)的数组。

有关更多信息，请参阅numpy指南中关于结合基本索引和高级索引的部分。

---

PS:仍然可以使用单个索引调用来获得所需的形状，但是您必须与切片分开，而不是定义广播到形状(5, 2, 3, 1)的索引，例如使用np.ix_:

>>> aa[ np.ix_([0, 1, 2, 3, 4], [1, 3], [0, 1, 2], [0]) ].shape
(5, 2, 3, 1)

Answer 2

一次减少一个轴:

aa[:, [1, 3], :, :][..., [0]].shape
(5, 2, 3, 1)