Déclinaison des tableaux Numpy

Question

Dans une grande base de code, j'utilise np.broadcast_to pour diffuser des tableaux (nous n'utilisons ici que des exemples simples) :

In [1]: x = np.array([1,2,3])

In [2]: y = np.broadcast_to(x, (2,1,3))

In [3]: y.shape
Out[3]: (2, 1, 3)

Ailleurs dans le code, j'utilise des fonctions tierces qui peuvent opérer de manière vectorielle sur des tableaux Numpy mais qui ne sont pas des ufuncs. Ces fonctions ne comprennent pas la diffusion, ce qui signifie qu'appeler une telle fonction sur des tableaux tels que y est inefficace. Des solutions comme celle de Numpy vectorize ne sont pas bonnes non plus car, bien qu'elles comprennent la radiodiffusion, elles introduisent une for boucle sur les éléments du tableau, ce qui est alors très inefficace.

Idéalement, j'aimerais pouvoir disposer d'une fonction que l'on pourrait appeler par exemple unbroadcast qui renvoie un tableau avec une forme minimale qui peut être retransmise à la taille complète si nécessaire. Ainsi, par exemple

In [4]: z = unbroadcast(y)

In [5]: z.shape
Out[5]: (1, 1, 3)

Je peux alors exécuter les fonctions tierces sur z puis diffuse le résultat à l'adresse suivante y.shape .

Existe-t-il un moyen de mettre en œuvre unbroadcast qui s'appuie sur l'API publique de Numpy ? Sinon, existe-t-il des astuces qui permettraient d'obtenir le résultat souhaité ?

Answer 1

J'ai une solution possible, je vais donc la poster ici (mais si quelqu'un en a une meilleure, n'hésitez pas à répondre aussi !) Une solution consiste à vérifier le strides argument des tableaux, qui seront 0 le long des dimensions diffusées :

def unbroadcast(array):
    slices = []
    for i in range(array.ndim):
        if array.strides[i] == 0:
            slices.append(slice(0, 1))
        else:
            slices.append(slice(None))
    return array[slices]

Cela donne :

In [14]: unbroadcast(y).shape
Out[14]: (1, 1, 3)

Answer 2

C'est probablement équivalent à votre propre solution, mais un peu plus intégré. Elle utilise as_strided en numpy.lib.stride_tricks :

import numpy as np
from numpy.lib.stride_tricks import as_strided

x = np.arange(16).reshape(2,1,8,1)  # shape (2,1,8,1)
y = np.broadcast_to(x,(2,3,8,5))    # shape (2,3,8,5) broadcast

def unbroadcast(arr):
    #determine unbroadcast shape
    newshape = np.where(np.array(arr.strides) == 0,1,arr.shape) # [2,1,8,1], thanks to @Divakar
    return as_strided(arr,shape=newshape)    # strides are automatically set here

z = unbroadcast(x)
np.all(z==x)  # is True

---

Notez que dans ma réponse originale je n'ai pas défini de fonction, et le résultat z Le tableau a (64,0,8,0) comme strides alors que l'entrée a (64,64,8,8) . Dans la version actuelle, le retour z Le tableau a des stridences identiques à x Je suppose que passer et retourner le tableau force la création d'une copie. Quoi qu'il en soit, nous pourrions toujours définir les strides manuellement dans le fichier as_strided pour obtenir des tableaux identiques en toutes circonstances, mais cela ne semble pas nécessaire dans la configuration ci-dessus.