任意维数的Numpy子阵列

时间:2013-11-07 17:18:58

标签: numpy multidimensional-array slice matrix-indexing

在Numpy中,假设你有一个Nd阵列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函数只返回完整数组,因为for循环中的if条件永远不是True,而fslice >= farray.shape[fdim]返回的子数组在维度fdim中的大小为0。

这很容易扩展到不仅仅是在一个维度中选择一个元素。

谢谢!

1 个答案:

答案 0 :(得分: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]]])