calculate binomial distribution probability matrix with python

Question

Give the N and P, I want to get a 2D binomial distribution probability matrix M,

for i in range(1, N+1):
   for j in range(i+1):
      M[i,j] = choose(i, j) * p**j * (1-p)**(i-j)
other value = 0

I want to know is there any fast way to get this matrix, instead of the for loop. the N may be bigger than 100,000

Answer 1

I believe that scipy.stats.binom can take advantage of broadcasting in the way that you're looking for.

# Binomial PMF: Pr(X=k) = choose(n, k) * p**k * (1-p)**(n-k)
# Probability of getting exactly k successes in n trials

>>> from scipy.stats import binom

>>> n = np.arange(1, N+1, dtype=np.int64)
>>> dist = binom(p=0.25, n=n)
>>> M = dist.pmf(k=np.arange(N+1, dtype=np.int64)[:, None])

>>> M.round(2)
array([[0.75, 0.56, 0.42, 0.32, 0.24, 0.18, 0.13, 0.1 , 0.08, 0.06],
       [0.25, 0.38, 0.42, 0.42, 0.4 , 0.36, 0.31, 0.27, 0.23, 0.19],
       [0.  , 0.06, 0.14, 0.21, 0.26, 0.3 , 0.31, 0.31, 0.3 , 0.28],
       [0.  , 0.  , 0.02, 0.05, 0.09, 0.13, 0.17, 0.21, 0.23, 0.25],
       [0.  , 0.  , 0.  , 0.  , 0.01, 0.03, 0.06, 0.09, 0.12, 0.15],
       [0.  , 0.  , 0.  , 0.  , 0.  , 0.  , 0.01, 0.02, 0.04, 0.06],
       [0.  , 0.  , 0.  , 0.  , 0.  , 0.  , 0.  , 0.  , 0.01, 0.02],
       [0.  , 0.  , 0.  , 0.  , 0.  , 0.  , 0.  , 0.  , 0.  , 0.  ],
       [0.  , 0.  , 0.  , 0.  , 0.  , 0.  , 0.  , 0.  , 0.  , 0.  ],
       [0.  , 0.  , 0.  , 0.  , 0.  , 0.  , 0.  , 0.  , 0.  , 0.  ],
       [0.  , 0.  , 0.  , 0.  , 0.  , 0.  , 0.  , 0.  , 0.  , 0.  ]])

Here, the rows are k (0-indexed) and the columns are n (1-indexed):

>>> from math import factorial as fac
>>> def manual_pmf(p, n, k):
...     return fac(n) / (fac(k) * fac(n - k)) * p**k * (1-p)**(n-k)

>>> manual_pmf(p=0.25, n=3, k=2)
0.140625  # (2, 2) in M because M's columns are effectively 1-indexed 

You could also start n at zero to get an array that is 0-indexed on both rows and columns:

>>> n = np.arange(N+1, dtype=np.int64)  # M.shape is (11, 11)