Numpy Subarray per numero arbitrario di dimensione

Question

In Numpy, supponiamo che tu abbia un array a, puoi tagliare l'ultima dimensione facendo A[...,0] oppure puoi tagliare la prima dimensione facendo A[0]. Voglio generalizzare questa operazione per tutta la dimensione (non solo prima o ultima) e lo voglio per ND-ARRAY di N arbitrario in modo che se A3 è un 3D-array e A4 è un array 4D, func(A3, dim = 1, slice = 0) Dammi A3[ : , 0 , :] e func(A4, dim = 1, slice = 0) Dammi A4[ : , 0 , : , : ].

Ho cercato questo per un po 'di tempo e alla fine ho capito come farlo senza fare un hack orribile (come lo scambio di dimensione fino a quando quello di interesse non è all'ultimo posto). Quindi il codice che sto postando qui funziona per il mio bisogno ma

1) Sono sempre alla ricerca di consigli per migliorare me stesso

2) Come ho detto, lo cerco da tempo e non ho mai trovato nulla in modo che possa essere utile per gli altri.

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

Quindi questa piccola funzione crea solo una tupla affettata concatenando slice(None,None,1) Nella posizione della dimensione non vogliamo tagliare e slice(fslice, fslice+1, 1) nella dimensione di interesse. Che restituisce un subarray. Gestisce gli indici negativi.

C'è una leggera differenza con questo e indicizzazione diretta: se A3 è 3x4x5, A3[:,0,:] sarà 3x5 mentre fancy_subarray(A3, fdim = 1, fslice = 0) sarà 3x1x5. Anche la gestione della funzione fuori dalla dimensione vincolata e indici "naturalmente". Se fdim >= farray.ndim La funzione restituisce solo l'intero array perché la condizione if all'interno del loop non è mai vera e se fslice >= farray.shape[fdim] Il subarray restituito avrà una dimensione di 0 in dimensione FDIM.

Questo può essere facilmente esteso a qualcosa di più della semplice selezione di un elemento in una dimensione ovviamente.

Grazie!

Answer 1

Penso che tu stia complicando eccessivamente le cose, specialmente durante la gestione fslice. Se hai semplicemente fatto:

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]

Quindi non solo rendi il tuo codice più compatto, ma la stessa funzione può assumere singoli indici, fette o persino elenchi di indici:

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