Skip to content

Commit 5281d4e

Browse files
jhamrickjhamrick
committed
Add multinomial tests
1 parent d675599 commit 5281d4e

2 files changed

Lines changed: 146 additions & 1 deletion

File tree

‎sampler/multinomial.py‎

Lines changed: 16 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,11 @@ def __init__(self, p, rso=None):
2727
else:
2828
self._sample_func = np.random.multinomial
2929

30+
# check that the probabilities sum to 1 -- if they don't, then
31+
# something is wrong
32+
if not np.isclose(np.sum(self.p), 1.0):
33+
raise ValueError("event probabilities do not sum to 1")
34+
3035
def logpmf(self, x):
3136
"""Evaluates the log-probability mass function (log-PMF) of a
3237
multinomial with event probabilities `self.p` for a draw `x`.
@@ -43,12 +48,22 @@ def logpmf(self, x):
4348
"""
4449
# get the total number of events
4550
n = np.sum(x)
51+
4652
# equivalent to log(n!)
4753
numerator = gammaln(n + 1)
4854
# equivalent to log(x1! * ... * xk!)
4955
denominator = np.sum(gammaln(x + 1))
56+
57+
# if one of the values of self.p is 0, then the corresponding
58+
# value of self.logp will be -inf. If the corresponding value
59+
# of x is 0, then multiplying them together will give nan, but
60+
# we want it to just be 0.
61+
all_weights = self.logp * x
62+
all_weights[x == 0] = 0
5063
# equivalent to log(p1^x1 * ... * pk^xk)
51-
weights = np.sum(self.logp * x)
64+
weights = np.sum(all_weights)
65+
66+
# put it all together
5267
log_pmf = numerator - denominator + weights
5368
return log_pmf
5469

‎sampler/test_multinomial.py‎

Lines changed: 130 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,130 @@
1+
import numpy as np
2+
import pytest
3+
from multinomial import MultinomialDistribution
4+
5+
6+
def test_init_without_rso():
7+
"""Initialize without rso"""
8+
p = np.array([0.1, 0.5, 0.3, 0.1])
9+
dist = MultinomialDistribution(p)
10+
assert (dist.p == p).all()
11+
assert (dist.logp == np.log(p)).all()
12+
assert dist.rso is None
13+
assert dist._sample_func == np.random.multinomial
14+
15+
16+
def test_init_with_rso():
17+
"""Initialize with rso"""
18+
p = np.array([0.1, 0.5, 0.3, 0.1])
19+
rso = np.random.RandomState(29348)
20+
dist = MultinomialDistribution(p, rso=rso)
21+
assert (dist.p == p).all()
22+
assert (dist.logp == np.log(p)).all()
23+
assert dist.rso == rso
24+
assert dist._sample_func == rso.multinomial
25+
26+
27+
def test_init_bad_probabilities():
28+
"""Initialize with probabilities that don't sum to 1"""
29+
p = np.array([0.1, 0.5, 0.3, 0.0])
30+
rso = np.random.RandomState(29348)
31+
with pytest.raises(ValueError):
32+
MultinomialDistribution(p, rso=rso)
33+
34+
35+
def test_pmf_1():
36+
"""Test PMF with only one possible event"""
37+
p = np.array([1.0])
38+
rso = np.random.RandomState(29348)
39+
dist = MultinomialDistribution(p, rso=rso)
40+
assert dist.pmf(np.array([1])) == 1.0
41+
assert dist.pmf(np.array([2])) == 1.0
42+
assert dist.pmf(np.array([10])) == 1.0
43+
44+
45+
def test_pmf_2():
46+
"""Test PMF with two possible events, one with zero probability"""
47+
p = np.array([1.0, 0.0])
48+
rso = np.random.RandomState(29348)
49+
dist = MultinomialDistribution(p, rso=rso)
50+
assert dist.pmf(np.array([1, 0])) == 1.0
51+
assert dist.pmf(np.array([0, 1])) == 0.0
52+
assert dist.pmf(np.array([2, 0])) == 1.0
53+
assert dist.pmf(np.array([2, 2])) == 0.0
54+
assert dist.pmf(np.array([10, 0])) == 1.0
55+
assert dist.pmf(np.array([10, 3])) == 0.0
56+
57+
p = np.array([0.0, 1.0])
58+
rso = np.random.RandomState(29348)
59+
dist = MultinomialDistribution(p, rso=rso)
60+
assert dist.pmf(np.array([0, 1])) == 1.0
61+
assert dist.pmf(np.array([1, 0])) == 0.0
62+
assert dist.pmf(np.array([0, 2])) == 1.0
63+
assert dist.pmf(np.array([2, 2])) == 0.0
64+
assert dist.pmf(np.array([0, 10])) == 1.0
65+
assert dist.pmf(np.array([3, 10])) == 0.0
66+
67+
68+
def test_pmf_3():
69+
"""Test PMF with two possible events, both with nonzero probability"""
70+
p = np.array([0.5, 0.5])
71+
rso = np.random.RandomState(29348)
72+
dist = MultinomialDistribution(p, rso=rso)
73+
assert dist.pmf(np.array([1, 0])) == 0.5
74+
assert dist.pmf(np.array([0, 1])) == 0.5
75+
assert dist.pmf(np.array([2, 0])) == 0.25
76+
assert dist.pmf(np.array([0, 2])) == 0.25
77+
assert dist.pmf(np.array([1, 1])) == 0.5
78+
79+
80+
def test_sample_1():
81+
"""Test sampling with only one possible event"""
82+
p = np.array([1.0])
83+
rso = np.random.RandomState(29348)
84+
dist = MultinomialDistribution(p, rso=rso)
85+
samples = np.array([dist.sample(1) for i in xrange(100)])
86+
assert samples.shape == (100, 1)
87+
assert (samples == 1).all()
88+
samples = np.array([dist.sample(3) for i in xrange(100)])
89+
assert samples.shape == (100, 1)
90+
assert (samples == 3).all()
91+
92+
93+
def test_sample_2():
94+
"""Test sampling with two possible events, one with zero probability"""
95+
p = np.array([1.0, 0.0])
96+
rso = np.random.RandomState(29348)
97+
dist = MultinomialDistribution(p, rso=rso)
98+
samples = np.array([dist.sample(1) for i in xrange(100)])
99+
assert samples.shape == (100, 2)
100+
assert (samples == np.array([1, 0])).all()
101+
samples = np.array([dist.sample(3) for i in xrange(100)])
102+
assert samples.shape == (100, 2)
103+
assert (samples == np.array([3, 0])).all()
104+
105+
p = np.array([0.0, 1.0])
106+
rso = np.random.RandomState(29348)
107+
dist = MultinomialDistribution(p, rso=rso)
108+
samples = np.array([dist.sample(1) for i in xrange(100)])
109+
assert samples.shape == (100, 2)
110+
assert (samples == np.array([0, 1])).all()
111+
samples = np.array([dist.sample(3) for i in xrange(100)])
112+
assert samples.shape == (100, 2)
113+
assert (samples == np.array([0, 3])).all()
114+
115+
116+
def test_sample_3():
117+
"""Test sampling with two possible events, both with nonzero probability"""
118+
p = np.array([0.5, 0.5])
119+
rso = np.random.RandomState(29348)
120+
dist = MultinomialDistribution(p, rso=rso)
121+
samples = np.array([dist.sample(1) for i in xrange(100)])
122+
assert samples.shape == (100, 2)
123+
assert ((samples == np.array([1, 0])) |
124+
(samples == np.array([0, 1]))).all()
125+
samples = np.array([dist.sample(3) for i in xrange(100)])
126+
assert samples.shape == (100, 2)
127+
assert ((samples == np.array([3, 0])) |
128+
(samples == np.array([2, 1])) |
129+
(samples == np.array([1, 2])) |
130+
(samples == np.array([0, 3]))).all()

0 commit comments

Comments
 (0)