CODEWITHSUNDEEP
pythonnumpyrandom
simonzack
simonzack

Reputation: 20938

numpy draw a few bernoulli distributions at once

Say I want to draw from a list of bernoulli distributions at once, e.g. with probabilities [0.1, 0.2, 0.3]. I could do this with a for loop but is there a better numpy way (use of scipy is also fine)?

ps = [0.1, 0.2, 0.3]
[np.random.choice(2, p=[1 - p, p]) for p in ps]

Upvotes: 7

Views: 12478

Answers (2)

user6601029
user6601029

Reputation:

I suggest to rather use (np.random.uniform(size=d) < probs) * 1, where probs is your vector of probabilities. This generates uniform random variables that are then thresholded to be 0 or 1 and the probability of getting 1 is exactly what you have in probs.

If you need only a few bernoulli variables and just once the solution with np.random.binomial should be fine. However, it seems to be several times slower than generating bernoulli variables using uniform distribution. For instance, if we have some random vector of probabilities probs = np.uniform(size=10000), then we get

%timeit np.random.binomial(n=1, p=probs, size=10000)
Result: 775 µs ± 5.35 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

%timeit (np.random.uniform(size=10000) < probs) * 1
Result: 103 µs ± 510 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

Upvotes: 4

user2285236
user2285236

Reputation:

You can draw n=1 from a binomial, which is equivalent to Bernoulli. Since binomial accepts an array as the probability parameter, you can use:

np.random.binomial(1, p=ps)

You can test if it's working by giving very small/large probabilities to some elements and call the function several times.

For example, let ps = [0.23, 0.48, 0.64, 0.98]

In [90]: np.sum([np.random.binomial(1,p=ps) for i in range(100000)], axis=0)
Out[90]: array([23000, 48115, 64128, 97957])

Upvotes: 12

Related Questions