mardi 28 avril 2020

2D version of numpy random choice with weighting

This relates to this earlier post: Numpy random choice of tuples

I have a 2D numpy array and want to choose from it using a 2D probability array. The only way I could think to do this was to flatten and then use the modulo and remainder to convert the result back to a 2D index

import numpy as np
# dummy data
x=np.arange(100).reshape(10,10)

# dummy probability array
p=np.zeros([10,10])
p[4:7,1:4]=1.0/9

xy=np.random.choice(x.flatten(),1,p=p.flatten())
index=[int(xy/10),(xy%10)[0]] # convert back to index
print(index)

which gives

[5, 2]

but is there a cleaner way that avoids flattening and the modulo? i.e. I could pass a list of coordinate tuples as x, but how can I then handle the weights?




Aucun commentaire:

Enregistrer un commentaire