import numpy as np from scipy.special import gammaln class MultinomialDistribution(object): def __init__(self, p, rso=np.random): """Initialize the multinomial random variable. Parameters ---------- p: numpy array of length `k` The outcome probabilities rso: numpy RandomState object (default: np.random) The random number generator """ # Check that the probabilities sum to 1. If they don't, then # something is wrong! We use `np.isclose` rather than checking # for exact equality because in many cases, we won't have # exact equality due to floating-point error. if not np.isclose(np.sum(p), 1.0): raise ValueError("outcome probabilities do not sum to 1") # Store the parameters that were passed in self.p = p self.rso = rso # Precompute log probabilities, for use by the log-PMF, for # each element of `self.p` (the function `np.log` operates # elementwise over NumPy arrays, as well as on scalars.) self.logp = np.log(self.p) def sample(self, n): """Samples draws of `n` events from a multinomial distribution with outcome probabilities `self.p`. Parameters ---------- n: integer The number of total events Returns ------- numpy array of length `k` The sampled number of occurrences for each outcome """ x = self.rso.multinomial(n, self.p) return x def log_pmf(self, x): """Evaluates the log-probability mass function (log-PMF) of a multinomial with outcome probabilities `self.p` for a draw `x`. Parameters ---------- x: numpy array of length `k` The number of occurrences of each outcome Returns ------- The evaluated log-PMF for draw `x` """ # Get the total number of events n = np.sum(x) # equivalent to log(n!) log_n_factorial = gammaln(n + 1) # equivalent to log(x1! * ... * xk!) sum_log_xi_factorial = np.sum(gammaln(x + 1)) # If one of the values of self.p is 0, then the corresponding # value of self.logp will be -inf. If the corresponding value # of x is 0, then multiplying them together will give nan, but # we want it to just be 0. log_pi_xi = self.logp * x log_pi_xi[x == 0] = 0 # equivalent to log(p1^x1 * ... * pk^xk) sum_log_pi_xi = np.sum(log_pi_xi) # Put it all together log_pmf = log_n_factorial - sum_log_xi_factorial + sum_log_pi_xi return log_pmf def pmf(self, x): """Evaluates the probability mass function (PMF) of a multinomial with outcome probabilities `self.p` for a draw `x`. Parameters ---------- x: numpy array of length `k` The number of occurrences of each outcome Returns ------- The evaluated PMF for draw `x` """ pmf = np.exp(self.log_pmf(x)) return pmf