forked from aosabook/500lines
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathmultinomial.py
More file actions
103 lines (80 loc) · 3.07 KB
/
Copy pathmultinomial.py
File metadata and controls
103 lines (80 loc) · 3.07 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
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