Numpy: randomly choose 1 sample from each of the N discrete probability distributions

user_12109772 0
Viewed 70

I have a 2-D Numpy array P where each P[ii,:] is a discrete probability distribution over 0,1,...,M-1. I would like to sample from each of the P[ii,:]'s once. I've tried to use the numpy.random.choice function:

for ii in range(N):
    samples[ii] = np.random.choice(np.arange(M), p=P[ii,:])

How can I do the same thing without for loop?

Thank you.

numpy python
1 Answers

I think you probably want np.apply_along_axis

import numpy as np

P = np.array([[0,1/2,1/2],[1,0,0],[1/3,1/3,1/3]])

def rand_choice(c):
  return np.random.choice(c.size, p=c)

samples = np.apply_along_axis(rand_choice, 1, P)
user_107348800
Related