按索引复制numpy数组的单个轴

Question

我正在寻找一种优雅的方法，通过索引提取numpy数组的单个轴的值。例如：

x = np.arange(16).reshape((4,4))
a = x[0]
b = x[:, 0] 

是我通常所做的，但我正在寻找类似的东西：

^{pr2}$

也许有什么特别的功能可以做到这一点？在

Answer 1

您可以直接使用__getitem__magic方法在功能上获得相同的结果，但是更容易支持动态参数。下面是一个使用itertools包的示例：

def get(matrix, axis, index):
    return a.__getitem__(tuple(chain(repeat(slice(None), axis), (index,))))

这将创建一个具有slice对象的元组，该对象表示以a[:]重复axis次的冒号，最后是index。我认为元组一代可以被清理干净，但我目前还想不出更干净的方法。在

示例用法如下：

^{pr2}$

Answer 2

您可以使用^{}将感兴趣的轴移到前面，然后按正常方式索引到它：

def get(x, axis=0, index=0):
    return np.rollaxis(x, axis, 0)[index]

x = np.arange(27).reshape(3, 3, 3)

assert np.all(get(x, 1, 2) == x[:, 2, :])

正如Joe正确指出的，这将返回x上的视图。为了强制制作副本，可以使用.copy()方法：

^{pr2}$