Vectorizing a comparison in numpy

Question

How can I vectorize this loop in NumPy? It uses sampling from NumPy's binomial() function to estimate the probability that out of 55 events exactly m of a particular type occur, where the probability of m occuring is 5%; ie it estimates 55Cm.(0.05)^m.(0.95)^(55-m). where 55Cm = 55!/(m!.(55-m)!)

import numpy as np
M = 7
m = np.arange(M+1)
ntrials = 1000000
p = np.empty(M+1)
for r in m:
    p[r] = np.sum(np.random.binomial(55, 0.05, ntrials)==r)/ntrials

Answer 1

Here is the equivalent code:

p = np.zeros(M+1)
print p

I imagine you didn't intend for your output to always be all zero, but it is! So the first thing to do is add a dtype=float argument to your np.sum() call. With that out of the way, we can vectorize the whole thing like this:

samples = np.random.binomial(55, 0.05, (ntrials, M+1))
p = np.sum(samples == m, dtype=float, axis=0) / ntrials

This produces an equivalent, though not identical, result. The reason is that the random number generation is done in a different sequence, so you will get an answer which is "correct" but not identical to the old code. If you want the identical result to before, you can get that by changing the first line to this:

samples = p.random.binomial(55, 0.05, (M+1, ntrials)).T

Then you draw in the same order as before, with no real performance penalty.