Comment aplatir une tranche numpy ?

Question

Je suis en train d'implémenter une sous-classe de ndarray de numpy et je dois modifier __getitem__ pour récupérer les éléments d'une représentation aplatie du tableau. Le problème est que __getitem__ peut être appelé soit avec un indice entier, soit avec une tranche multidimensionnelle.

Quelqu'un sait-il comment convertir une tranche multidimensionnelle en une liste d'indices (ou une tranche unidimensionnelle) sur le tableau aplati ?

Answer 1

Il se peut qu'il ne soit pas possible de convertir une tranche multidimensionnelle en une tranche plate, par exemple :

>>> a = np.arange(16).reshape(4, 4)
>>> a
array([[ 0,  1,  2,  3],
       [ 4,  5,  6,  7],
       [ 8,  9, 10, 11],
       [12, 13, 14, 15]])
>>> a[::3, 1::2]
array([[ 1,  3],
       [13, 15]])

Et vous ne pouvez pas accéder au sous-réseau [ 1, 3, 13, 15] avec un start:stop:step notation. Mais vous pouvez construire une liste d'indices plats à partir d'indices multidimensionnels, en procédant de la manière suivante :

>>> row_idx = np.arange(4)[::3]
>>> col_idx = np.arange(4)[1::2]
>>> row_idx = np.repeat(row_idx, 2)
>>> col_idx = np.tile(col_idx, 2)
>>> np.ravel_multi_index((row_idx, col_idx), dims=(4,4))
array([ 1,  3, 13, 15], dtype=int64)

---

Dans un cadre plus général, une fois que vous avez un tableau d'indices pour chaque dimension, vous devez calculer le produit cartésien de tous les tableaux d'indices, donc itertools.product est probablement la meilleure solution. Par exemple :

>>> indices = [np.array([0, 4, 8]), np.array([1,7]), np.array([3, 5, 9])]
>>> indices = zip(*itertools.product(*indices))
>>> indices
[(0, 0, 0, 0, 0, 0, 4, 4, 4, 4, 4, 4, 8, 8, 8, 8, 8, 8),
 (1, 1, 1, 7, 7, 7, 1, 1, 1, 7, 7, 7, 1, 1, 1, 7, 7, 7),
 (3, 5, 9, 3, 5, 9, 3, 5, 9, 3, 5, 9, 3, 5, 9, 3, 5, 9)]
>>> np.ravel_multi_index(indices, dims=(10, 11, 12))
array([  15,   17,   21,   87,   89,   93,  543,  545,  549,  615,  617,
        621, 1071, 1073, 1077, 1143, 1145, 1149], dtype=int64)