始终返回数组的numpy切片方法

Question

给定一个 numpy 数组和一个类似于 __getitem__ 的索引，有没有一种简单的方法可以获取这个数组的切片，确保总是返回一个数组，而不是单个数值?

有效的索引示例包括：一个 int（整数）、一个 slice（切片）、省略号或者上述的元组。

假设我有一个这样的数组：

a = np.array([[1,2],[3,4]])

我想要一个操作，这个操作在所有情况下都等同于 a[whatever]，除了当 a[whatever] 返回一个单个数值的时候（例如，a[1,1]）。在这些情况下，我希望这个替代操作返回一个只有一个元素的数组。

Answer 1

除了 np.array(a[whatever]) 这个写法，似乎没有比这个更简单或者更符合常规的方法了。

Answer 2

这里有一个稍微复杂一点的版本，它总是返回原始数组的一个视图（当然前提是你没有进行任何复杂的索引操作；这应该通过你对有效索引的规定来保证）：

def get(a, item):
    if not isinstance(item, tuple):
        item = (item,)
    if len(item) == a.ndim and all(isinstance(x, int) for x in item):
        return a[item + (None,)]
    else:
        return a[item]

Answer 3

如果你只是想在本来会返回一个单一数值的情况下，返回一个只有一个元素的数组，那为什么不直接对切片的结果使用 numpy.atleast_1d 呢？

比如：

import numpy as np
x = np.arange(100).reshape(10,10)
print x[0,0]
print np.atleast_1d(x[0,0])
print np.atleast_1d(x[:,:3])