-
Notifications
You must be signed in to change notification settings - Fork 1
/
plot.py
147 lines (123 loc) · 5.64 KB
/
plot.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
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
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
"""Classes for generating plot data."""
from __future__ import division
import numpy as np
import sympy as sy
import emcee
from sympy.parsing.sympy_parser import parse_expr
class Plot:
"""Generic plot class."""
def __init__(self, eq, nums):
"""Initialization requires a string the represents the equation for
the plot, as well as the plot boudaries, which are xmin and xmax in
one dimension and ymin and ymax in two dimensions. The boundaries
are passed as a dictionary."""
self.eq = eq
for key in nums:
setattr(self, key, nums[key])
#try to parse the given equation
try:
self.exp = parse_expr(str(sy.sympify(self.eq)))
# str b/c sympy does not like unicode
#sympify so that ^ is an acceptable exponentiation operator
except:
raise RuntimeError("Equation parsing failed.")
#validate ranges
if self.xmin >= self.xmax:
raise RuntimeError("Invalid range.")
if hasattr(self, 'ymin'):
if self.ymin >= self.ymax:
raise RuntimeError("Invalid range.")
class OneDeePlot(Plot):
"""One dimensional plot class."""
def evaluate(self, steps=1000):
"""Evaluate the given function at each of `step` points
in the x domain."""
stepsize = (self.xmax - self.xmin)/steps
try:
self.points = [(i, float(self.exp.evalf(subs={sy.Symbol('x'): i})))
for i in np.arange(self.xmin,
self.xmax,
stepsize)]
except:
raise RuntimeError("Evaluation failed.")
self.ymin = min(zip(*self.points)[1])
self.ymax = max(zip(*self.points)[1])
def __lnprobfn(self, pos):
"""Return the ln of the posterior probability of being in a
position. The posterior probability is just 1/error^2, where
error is the difference between the y-value at this position
and the minimum y-value."""
if pos[0] > self.xmax or pos[0] < self.xmin:
# We are outside the acceptable range, probability = 0
return -np.inf
res = float(self.exp.evalf(subs={sy.Symbol('x'): float(pos[0])}))
return np.log(1/((res - self.ymin)**2))
def sample(self, nsamples=1000, vardiv=10):
"""Generate a Markov chain of length `nsamples`
vardim controls the variance of the proposal distribution, which
will be equal to (domain width)/(vardiv). So making this big will
cause small jumps and vice versa."""
var = (self.xmax - self.xmin) / vardiv
sampler = emcee.MHSampler([[var, var], [var, var]], # cov
2, # dim
self.__lnprobfn)
initval = float(np.random.uniform(self.xmin, self.xmax, 1))
sampler.run_mcmc([initval, initval], nsamples)
self.chain = zip(*sampler.chain)[0]
def toclient(self):
"""Return a dictionary of everything that needs to be passed
to the client in order for plotting to take place."""
return {"points": self.points,
"xmin": self.xmin,
"xmax": self.xmax,
"chain": self.chain}
class TwoDeePlot(Plot):
"""Subclass for a two dimensional plot."""
def evaluate(self, xsteps=100, ysteps=100):
"""Evaluate the plot's equation at a set of points determined
by the number of x and y steps passed in."""
stepsizex = (self.xmax - self.xmin)/xsteps
stepsizey = (self.ymax - self.ymin)/ysteps
# get list of x and y points to eval at
self.xs = list(np.arange(self.xmin, self.xmax, stepsizex))
self.ys = list(np.arange(self.ymin, self.ymax, stepsizey))
self.zs = [[float(self.exp.evalf(subs={sy.Symbol('x'): i,
sy.Symbol('y'): j}))
for j in self.ys]
for i in self.xs]
# get min and max z value
self.zmin = min([min(row) for row in self.zs])
self.zmax = max([max(row) for row in self.zs])
def __lnprobfn(self, pos):
"""Returns theln of the posterior probability of any point
in the two dimensional space. This is just 1/error^2, where
error is the difference between f(pos) and the minimum z value."""
if (pos[0] > self.xmax or pos[0] < self.xmin or
pos[1] > self.ymax or pos[1] < self.ymin):
return -np.inf
res = float(self.exp.evalf(subs={sy.Symbol('x'): float(pos[0]),
sy.Symbol('y'): float(pos[1])}))
return np.log(1/((res-self.zmin)**2))
def sample(self, nsamples=1000, vardivx=10, vardivy=10):
"""Generate a Markov chain of length `nsamples`"""
xvar = (self.xmax-self.xmin)/vardivx
yvar = (self.ymax-self.ymin)/vardivy
cov = np.array([[xvar, 0],
[0, yvar]])
sampler = emcee.MHSampler(cov, 2, self.__lnprobfn)
initstate = np.concatenate((
np.random.uniform(self.xmin, self.xmax, 1),
np.random.uniform(self.ymin, self.ymax, 1)
))
sampler.run_mcmc(initstate, nsamples)
self.chain = sampler.chain.tolist()
def toclient(self):
"""Package up everything we need to send to the client in a dict."""
return {"zs": self.zs,
"xmin": self.xmin,
"xmax": self.xmax,
"ymin": self.ymin,
"ymax": self.ymax,
"zmin": self.zmin,
"zmax": self.zmax,
"chain": self.chain}