CODEWITHSUNDEEP
pythonprobabilitysampling
Sakurai.JJ
Sakurai.JJ

Reputation: 684

Sampling from joint probability mass function in python

I have a non-negative normalized vector p. I would like to sample an index from the index set of the vector. The probability getting sample k is p[k]. Using np.random.choice function, I can write the following code.

p = [0.2, 0.3, 0.1, 0.3, 0.1]
indices = np.arange(len(p))
k = np.random.choice(indices, p=p)

My question is, how can I generalize this code for multi-dimensional arrays? For example, given three dimensional non-negative normalized IxJxK tensor p = np.random.rand(I,J,K) how can I sample the index (i,j,k) with the probability p[i,j,k]?

Upvotes: 1

Views: 38

Answers (1)

Frank Yellin
Frank Yellin

Reputation: 11297

Suppose that x is a probability matrix:

x = np.random.random((3, 4, 5))
x /= np.sum(x)

You can use flatten(x) and np.random.choice() to get values in the range 0..3x4x5 with the associated probabilities:

flat = x.flatten()
temp = np.random.choice(np.arange(len(flat)), 10, p=flat)
print(temp)

You can now convert the indices in temp into indices in the original array:

np.unravel_index(temp, x.shape))

Note that this will return a tuple of three numpy arrays, with the first array being the first index, the second element being the second index, and so on.

Upvotes: 2

Related Questions