-
-
Notifications
You must be signed in to change notification settings - Fork 1.9k
/
test_plots.py
81 lines (63 loc) · 2.35 KB
/
test_plots.py
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
import matplotlib
matplotlib.use('Agg', warn=False)
import numpy as np
import pymc3 as pm
from .checks import close_to
from .models import multidimensional_model
from ..plots import traceplot, forestplot, autocorrplot, plot_posterior, make_2d
from ..step_methods import Slice, Metropolis
from ..sampling import sample
from ..tuning.scaling import find_hessian
from .test_examples import build_disaster_model
from pymc3.examples import arbitrary_stochastic as asmod
def test_plots():
# Test single trace
with asmod.build_model() as model:
start = model.test_point
h = find_hessian(start)
step = Metropolis(model.vars, h)
trace = sample(3000, step, start)
traceplot(trace)
forestplot(trace)
plot_posterior(trace)
autocorrplot(trace)
def test_plots_multidimensional():
# Test single trace
start, model, _ = multidimensional_model()
with model as model:
h = np.diag(find_hessian(start))
step = Metropolis(model.vars, h)
trace = sample(3000, step, start)
traceplot(trace)
plot_posterior(trace)
def test_multichain_plots():
model = build_disaster_model()
with model:
# Run sampler
step1 = Slice([model.early_mean_log_, model.late_mean_log_])
step2 = Metropolis([model.switchpoint])
start = {'early_mean': 2., 'late_mean': 3., 'switchpoint': 50}
ptrace = sample(1000, [step1, step2], start, njobs=2)
forestplot(ptrace, varnames=['early_mean', 'late_mean'])
autocorrplot(ptrace, varnames=['switchpoint'])
plot_posterior(ptrace)
def test_make_2d():
a = np.arange(4)
close_to(make_2d(a), a[:, None], 0)
n = 7
a = np.arange(n * 4 * 5).reshape((n, 4, 5))
res = make_2d(a)
assert res.shape == (n, 20)
close_to(a[:, 0, 0], res[:, 0], 0)
close_to(a[:, 3, 2], res[:, 2 * 4 + 3], 0)
def test_plots_transformed():
with pm.Model() as model:
pm.Uniform('x', 0, 1)
step = pm.Metropolis()
trace = pm.sample(100, step=step)
assert traceplot(trace).shape == (1, 2)
assert traceplot(trace, plot_transformed=True).shape == (2, 2)
assert autocorrplot(trace).shape == (1, 1)
assert autocorrplot(trace, plot_transformed=True).shape == (2, 1)
assert plot_posterior(trace).shape == (1, )
assert plot_posterior(trace, plot_transformed=True).shape == (2, )