Subarray Numpy para o número arbitrário de dimensão

Question

Em Numpy, digamos que você tenha um nd-Array A, você pode cortar a última dimensão fazendo A[...,0] ou você pode cortar a primeira dimensão fazendo A[0]. Quero generalizar esta operação para toda a dimensão (não apenas em primeiro ou último) e quero isso para Nd-Array de arbitrário n, para que se A3 seja um 3D-matray e A4 seja um 4D-matriz, func(A3, dim = 1, slice = 0) Me dê A3[ : , 0 , :] e func(A4, dim = 1, slice = 0) Me dê A4[ : , 0 , : , : ].

Procuro isso por algum tempo e finalmente descobri como fazê -lo sem fazer hacker horrível (como trocar a dimensão até que o interesse esteja em último lugar). Então, o código que estou postando aqui funciona para minha necessidade, mas

1) Estou sempre procurando conselhos para melhorar eu

2) Como eu disse, procuro isso há algum tempo e nunca encontrei nada, para que possa ser útil para os outros.

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()

Portanto, essa pequena função só cria uma tupla faticando concatenando slice(None,None,1) Na posição da dimensão, não queremos cortar e slice(fslice, fslice+1, 1) na dimensão do interesse. Do que retorna um subarray. Ele lida com índices negativos.

Há uma pequena diferença com esta e indexação direta: se A3 for 3x4x5, A3[:,0,:] vai ser 3x5 enquanto fancy_subarray(A3, fdim = 1, fslice = 0) vai ser 3x1x5. Além disso, a função manipula a dimensão limitada e os índices "naturalmente". Se fdim >= farray.ndim A função apenas retorna a matriz completa porque a condição se dentro do loop for nunca é verdadeira e se fslice >= farray.shape[fdim] O subarray retornado terá um tamanho de 0 na dimensão FDIM.

Isso pode ser facilmente estendido a mais do que apenas selecionar um elemento em uma dimensão, é claro.

Obrigado!

Answer 1

Eu acho que você está supercomplicando as coisas, especialmente ao manusear fslice. Se você simplesmente fez:

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]

Então você não apenas torna seu código mais compacto, mas a mesma função pode levar índices, fatias ou até listas de índices:

>>> 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]]])