沿任意维度切片 numpy 数组答案

【问题标题】：slicing numpy array along an arbitrary dimension沿任意维度切片 numpy 数组
【发布时间】：2013-05-24 14:49:07
【问题描述】：

假设我有一个 (40,20,30) numpy 数组，并且我有一个函数，经过一些工作后，它将沿选定的输入轴返回输入数组的一半。有自动的方法吗？我想避免这样丑陋的代码：

def my_function(array,axis=0):

    ...

    if axis == 0:
        return array[:array.shape[0]/2,:,:] --> (20,20,30) array
    elif axis = 1:
        return array[:,:array.shape[1]/2,:] --> (40,10,30) array
    elif axis = 2: 
        return array[:,:,:array.shape[2]/2] --> (40,20,15) array

感谢您的帮助

埃里克

【问题讨论】：

标签： python numpy

【解决方案1】：

我认为您可以将np.split 用于此[docs]，并简单地获取返回的第一个或第二个元素，具体取决于您想要的元素。例如：

>>> a = np.random.random((40,20,30))
>>> np.split(a, 2, axis=0)[0].shape
(20, 20, 30)
>>> np.split(a, 2, axis=1)[0].shape
(40, 10, 30)
>>> np.split(a, 2, axis=2)[0].shape
(40, 20, 15)
>>> (np.split(a, 2, axis=0)[0] == a[:a.shape[0]/2, :,:]).all()
True

【讨论】：

仅供参考：split() 还接受一个指定任意分割点的元组。

【解决方案2】：

感谢您的帮助，帝斯曼。我会用你的方法。

与此同时，我发现了一个（肮脏的？）黑客：

>>> a = np.random.random((40,20,30))
>>> s = [slice(None),]*a.ndim
>>> s[axis] = slice(f,l,s)
>>> a1 = a[s]

可能比 np.split 更通用一点，但不那么优雅！

【讨论】：

这同样优雅。这只是删除了一些语法糖：: 是 slice(None)，a:b:c 是 slice(a, b, c)，等等。

【解决方案3】：

numpy.rollaxis 是一个很好的工具：

def my_func(array, axis=0):
    array = np.rollaxis(array, axis)
    out = array[:array.shape[0] // 2]
    # Do stuff with array and out knowing that the axis of interest is now 0
    ...

    # If you need to restore the order of the axes
    if axis == -1:
        axis = out.shape[0] - 1
    out = np.rollaxis(out, 0, axis + 1)

【讨论】：