Sous-réseau Numpy pour un nombre arbitraire de dimension

Question

Dans Numpy, disons que vous avez un nd-array a, vous pouvez trancher la dernière dimension en faisant A[...,0] ou vous pouvez couper la première dimension en faisant A[0]. Je veux généraliser cette opération pour toutes les dimensions (pas seulement le premier ou le dernier) et je veux cela pour ND-Array of Arbitrary N afin que si A3 soit un array 3D et A4 est un array 4D, func(A3, dim = 1, slice = 0) donne-moi A3[ : , 0 , :] et func(A4, dim = 1, slice = 0) donne-moi A4[ : , 0 , : , : ].

Je le recherche depuis un certain temps et j'ai finalement compris comment le faire sans faire du hack horrible (comme l'échange de dimension jusqu'à ce que celui qui intéresse est en dernière place). Donc le code que je publie ici fonctionne pour mon besoin mais

1) Je suis toujours à la recherche de conseils pour m'améliorer

2) Comme je l'ai dit, je le recherche depuis un certain temps et je n'ai jamais rien trouvé, ce qui pourrait être utile pour les autres.

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

Donc, cette petite fonction construit simplement un tuple de tranchage en concaténant slice(None,None,1) En position de la dimension, nous ne voulons pas trancher et slice(fslice, fslice+1, 1) dans la dimension d'intérêt. Il renvoie un sous-réseau. Il gère les indices négatifs.

Il y a une légère différence avec cette indexation directe: si A3 est 3x4x5, A3[:,0,:] va être 3x5 pendant fancy_subarray(A3, fdim = 1, fslice = 0) va être 3x1x5. Aussi la poignée de fonction hors de la dimension et les index "naturellement". Si fdim >= farray.ndim La fonction renvoie simplement le tableau complet car la condition IF à l'intérieur de la boucle pour fslice >= farray.shape[fdim] Le sous-réseau retourné aura une taille de 0 dans la dimension fdim.

Cela peut facilement être étendu à plus que la simple sélection d'un élément dans une dimension bien sûr.

Merci!

Answer 1

Je pense que vous compliquez des choses, surtout lors de la manipulation fslice. Si vous l'avez simplement fait:

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]

Ensuite, non seulement vous rendez votre code plus compact, mais la même fonction peut prendre des indices, des tranches ou même des listes d'indices:

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