forked from aosabook/500lines
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathtest_multinomial.py
More file actions
128 lines (108 loc) · 4.42 KB
/
Copy pathtest_multinomial.py
File metadata and controls
128 lines (108 loc) · 4.42 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
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
import numpy as np
import pytest
from multinomial import MultinomialDistribution
def test_init_without_rso():
"""Initialize without rso"""
p = np.array([0.1, 0.5, 0.3, 0.1])
dist = MultinomialDistribution(p)
assert (dist.p == p).all()
assert (dist.logp == np.log(p)).all()
assert dist.rso is np.random
def test_init_with_rso():
"""Initialize with rso"""
p = np.array([0.1, 0.5, 0.3, 0.1])
rso = np.random.RandomState(29348)
dist = MultinomialDistribution(p, rso=rso)
assert (dist.p == p).all()
assert (dist.logp == np.log(p)).all()
assert dist.rso == rso
def test_init_bad_probabilities():
"""Initialize with probabilities that don't sum to 1"""
p = np.array([0.1, 0.5, 0.3, 0.0])
rso = np.random.RandomState(29348)
with pytest.raises(ValueError):
MultinomialDistribution(p, rso=rso)
def test_pmf_1():
"""Test PMF with only one possible event"""
p = np.array([1.0])
rso = np.random.RandomState(29348)
dist = MultinomialDistribution(p, rso=rso)
assert dist.pmf(np.array([1])) == 1.0
assert dist.pmf(np.array([2])) == 1.0
assert dist.pmf(np.array([10])) == 1.0
def test_pmf_2():
"""Test PMF with two possible events, one with zero probability"""
p = np.array([1.0, 0.0])
rso = np.random.RandomState(29348)
dist = MultinomialDistribution(p, rso=rso)
assert dist.pmf(np.array([1, 0])) == 1.0
assert dist.pmf(np.array([0, 1])) == 0.0
assert dist.pmf(np.array([2, 0])) == 1.0
assert dist.pmf(np.array([2, 2])) == 0.0
assert dist.pmf(np.array([10, 0])) == 1.0
assert dist.pmf(np.array([10, 3])) == 0.0
p = np.array([0.0, 1.0])
rso = np.random.RandomState(29348)
dist = MultinomialDistribution(p, rso=rso)
assert dist.pmf(np.array([0, 1])) == 1.0
assert dist.pmf(np.array([1, 0])) == 0.0
assert dist.pmf(np.array([0, 2])) == 1.0
assert dist.pmf(np.array([2, 2])) == 0.0
assert dist.pmf(np.array([0, 10])) == 1.0
assert dist.pmf(np.array([3, 10])) == 0.0
def test_pmf_3():
"""Test PMF with two possible events, both with nonzero probability"""
p = np.array([0.5, 0.5])
rso = np.random.RandomState(29348)
dist = MultinomialDistribution(p, rso=rso)
assert dist.pmf(np.array([1, 0])) == 0.5
assert dist.pmf(np.array([0, 1])) == 0.5
assert dist.pmf(np.array([2, 0])) == 0.25
assert dist.pmf(np.array([0, 2])) == 0.25
assert dist.pmf(np.array([1, 1])) == 0.5
def test_sample_1():
"""Test sampling with only one possible event"""
p = np.array([1.0])
rso = np.random.RandomState(29348)
dist = MultinomialDistribution(p, rso=rso)
samples = np.array([dist.sample(1) for i in xrange(100)])
assert samples.shape == (100, 1)
assert (samples == 1).all()
samples = np.array([dist.sample(3) for i in xrange(100)])
assert samples.shape == (100, 1)
assert (samples == 3).all()
def test_sample_2():
"""Test sampling with two possible events, one with zero probability"""
p = np.array([1.0, 0.0])
rso = np.random.RandomState(29348)
dist = MultinomialDistribution(p, rso=rso)
samples = np.array([dist.sample(1) for i in xrange(100)])
assert samples.shape == (100, 2)
assert (samples == np.array([1, 0])).all()
samples = np.array([dist.sample(3) for i in xrange(100)])
assert samples.shape == (100, 2)
assert (samples == np.array([3, 0])).all()
p = np.array([0.0, 1.0])
rso = np.random.RandomState(29348)
dist = MultinomialDistribution(p, rso=rso)
samples = np.array([dist.sample(1) for i in xrange(100)])
assert samples.shape == (100, 2)
assert (samples == np.array([0, 1])).all()
samples = np.array([dist.sample(3) for i in xrange(100)])
assert samples.shape == (100, 2)
assert (samples == np.array([0, 3])).all()
def test_sample_3():
"""Test sampling with two possible events, both with nonzero probability"""
p = np.array([0.5, 0.5])
rso = np.random.RandomState(29348)
dist = MultinomialDistribution(p, rso=rso)
samples = np.array([dist.sample(1) for i in xrange(100)])
assert samples.shape == (100, 2)
assert ((samples == np.array([1, 0])) |
(samples == np.array([0, 1]))).all()
samples = np.array([dist.sample(3) for i in xrange(100)])
assert samples.shape == (100, 2)
assert ((samples == np.array([3, 0])) |
(samples == np.array([2, 1])) |
(samples == np.array([1, 2])) |
(samples == np.array([0, 3]))).all()