任意尺寸数量的Numpy子阵列

Question

在numpy中，假设您有一个nd-array a，您可以通过执行执行 A[...,0] 或者，您可以通过执行第一个维度切片 A[0]. 。我想对所有维度（不仅是第一或最后）概括此操作，并且我希望为任意n的ND阵列进行此操作，以便如果A3是3D阵列，并且A4是4D阵列，则 func(A3, dim = 1, slice = 0) 给我吗 A3[ : , 0 , :] 和 func(A4, dim = 1, slice = 0) 给我吗 A4[ : , 0 , : , : ].

我已经寻找了一段时间，终于弄清楚了如何在不进行可怕的黑客的情况下做到这一点（例如，交换维度直到感兴趣的是最后一个地方）。因此，我在此处发布的代码适合我的需要，但是

1）我一直在寻找建议自己的建议

2）正如我所说，我已经寻找了一段时间，从未找到任何东西，因此对其他人可能有用。

def fancy_subarray(farray, fdim, fslice):
    # Return the farray slice at position fslice in dimension fdim

    ndim = farray.ndim

    # Handle negative dimension and slice indexing
    if fdim < 0:
        fdim += ndim
    if fslice < 0:
        fslice += v.shape[fdim]

    # Initilize slicing tuple
    obj = ()

    for i in range(ndim):
        if i == fdim:
            # Only element fslice in that dimension
            obj += (slice(fslice, fslice+1, 1),)
        else:
            # "Full" dimension
            obj += (slice(None,None,1),)

    return farray[obj].copy()

因此，这个小功能只是通过连接来建立切片元组 slice(None,None,1) 在尺寸的位置，我们不想切片 slice(fslice, fslice+1, 1) 在感兴趣的方面。它比返回子阵列。它处理负索引。

与此索引有很小的区别：如果A3为3x4x5，则 A3[:,0,:] 会是3x5的 fancy_subarray(A3, fdim = 1, fslice = 0) 将是3x1x5。还可以“自然”函数处理界限和索引。如果 fdim >= farray.ndim 该功能只是返回完整的数组 fslice >= farray.shape[fdim] 返回的子阵列的大小为0尺寸FDIM。

当然，这可以轻松地扩展到不仅仅是在一个维度中选择一个元素。

谢谢！

Answer 1

我认为您过于复杂的事情，尤其是在处理时 fslice. 。如果您只是这样做：

def fancy_subarray(farray, fdim, fslice):
    fdim += farray.ndim if fdim < 0 else 0
    index = ((slice(None),) * fdim + (fslice,) +
             (slice(None),) * (farray.ndim - fdim - 1))
    return farray[index]

那么，您不仅使代码更加紧凑，而且相同的功能可以采用单个索引，切片甚至列表：

>>> fancy_subarray(a, 1, 2)
array([[10, 11, 12, 13, 14],
       [30, 31, 32, 33, 34],
       [50, 51, 52, 53, 54]])
>>> fancy_subarray(a, 1, slice(2,3))
array([[[10, 11, 12, 13, 14]],

       [[30, 31, 32, 33, 34]],

       [[50, 51, 52, 53, 54]]])
>>> fancy_subarray(a, 1, [2, 3])
array([[[10, 11, 12, 13, 14],
        [15, 16, 17, 18, 19]],

       [[30, 31, 32, 33, 34],
        [35, 36, 37, 38, 39]],

       [[50, 51, 52, 53, 54],
        [55, 56, 57, 58, 59]]])