CODEWITHSUNDEEP
pythonnumpyrandom
ClimateUnboxed
ClimateUnboxed

Reputation: 8067

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?

Upvotes: 2

Views: 1901

Answers (2)

yatu
yatu

Reputation: 88226

I don't think it's possible to directly specify a 2D shaped array of probabilities. So raveling should be fine. However to get the corresponding 2D shaped indices from the flat index you can use np.unravel_index

index= np.unravel_index(xy.item(), x.shape)
# (4, 2)

For multiple indices, you can just stack the result:

xy=np.random.choice(x.flatten(),3,p=p.flatten())

indices = np.unravel_index(xy, x.shape)
# (array([4, 4, 5], dtype=int64), array([1, 2, 3], dtype=int64))
np.c_[indices]
array([[4, 1],
       [4, 2],
       [5, 3]], dtype=int64)

where np.c_ stacks along the right hand axis and gives the same result as 

np.column_stack(indices)

Upvotes: 3

Paddy Harrison
Paddy Harrison

Reputation: 1992

You could use numpy.random.randint to generate an index, for example:

# assumes p is a square array
ij = np.random.randint(p.shape[0], size=p.ndim) # size p.ndim = 2 generates 2 coords

# need to convert to tuple to index correctly
p[tuple(i for i in ij))]
>>> 0.0

You can also index multiple random values at once:

ij = np.random.randint(p.shape[0], size=(p.ndim, 5)) # get 5 values
p[tuple(i for i in ij))]
>>> array([0.        , 0.        , 0.        , 0.11111111, 0.        ])

Upvotes: 0

Related Questions