Échantillonnage aléatoire pondéré à partir du tableau 2d numpy

Question

J'ai un tableau 2d numpy Z et je veux choisir aléatoirement un index de Z où la chance d'un index choisi est proportionnelle à la valeur de Z à cet index.

En ce moment, je fais ce qui suit:

yar = list(np.ndenumerate(Z)) 
x,y = yar[np.random.choice(len(yar), p=Z.ravel()/Z.sum())][0] 

Ce qui fait le travail, mais se sent hideux (et est extrêmement lent en plus). Y a-t-il un meilleur moyen?

Answer 1

Nous pouvons optimiser en évitant la création de yar. Nous obtiendrions simplement l'équivalent linéaire de np.random.choice, le convertissons en indices de dimension avec np.unravel_index pour nous donner x et y.

Ainsi, la mise en œuvre serait -

linear_idx = np.random.choice(Z.size, p=Z.ravel()/float(Z.sum())) 
x, y = np.unravel_index(linear_idx, Z.shape) 

juste pour donner un peu de contexte sur les chiffres dont la création de yar a été à l'origine du goulot d'étranglement dans cette configuration, voici un test de synchronisation de l'échantillon -

In [402]: Z = np.random.randint(0,9,(300,400)) 

In [403]: yar = list(np.ndenumerate(Z)) 

In [404]: %timeit list(np.ndenumerate(Z)) 
10 loops, best of 3: 46.3 ms per loop 

In [405]: %timeit yar[np.random.choice(len(yar), p=Z.ravel()/float(Z.sum()))][0] 
1000 loops, best of 3: 1.34 ms per loop 

In [406]: 46.3/(46.3+1.34) 
Out[406]: 0.971872376154492 

Donc, la création yar était manger jusqu'à97% de l'exécution là.