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
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
|
# -*- encoding: utf-8 -*-
# traits
from traits.api import HasTraits, Instance, Button
# ui
from traitsui.view import View
from traitsui.item import Item
# general
from chaco.plot import Plot
from chaco.errorbar_plot import ErrorBarPlot
from chaco.plot_containers import GridPlotContainer as PCont
from chaco.array_data_source import ArrayDataSource
from chaco.default_colormaps import jet
# tools
from chaco.tools.pan_tool import PanTool as PTool
from chaco.tools.better_selecting_zoom import BetterSelectingZoom as ZTool
from chaco.tools.save_tool import SaveTool
from chaco.tools.line_inspector import LineInspector
# enable
from enable.component_editor import ComponentEditor
from enable.base_tool import BaseTool
# other
from numpy import linspace, sin
from math import ceil, sqrt
class FocusTool(BaseTool):
def __init__(self, mclayout, desc):
super(FocusTool, self).__init__()
self.mclayout = mclayout
self.desc = desc
def normal_left_dclick(self, *args):
if len(self.mclayout.focus) > 1:
# focus on chosen plot
self.mclayout.focus = [self.desc]
# reinit layout
self.mclayout.reinit()
def normal_key_pressed(self, event):
if event.character == 'l':
self.desc.log = not self.desc.log
self.mclayout.reinit()
class Snapper(LineInspector):
def __init__(self, mclayout, plot, **kwargs):
super(Snapper, self).__init__(plot, axis='index_x', **kwargs)
self.layout = mclayout
self.plot = plot
self.pos = None
self.enabled = False
self.title = plot.title
def normal_key_pressed(self, event):
if event.character == 't':
self.enabled = not self.enabled
if not self.enabled:
self.pos = None
self.plot.title = self.title
self.plot.request_redraw()
def normal_mouse_move(self, event):
if not self.enabled:
return
mx = self.plot.index_mapper.map_data(event.x)
px = self.plot.data[self.plot.x_axis.title][0]
for i, x in enumerate(self.plot.data[self.plot.x_axis.title]):
if x > mx:
break
px = x
if (px - mx) ** 2 < (x - mx)**2:
i -= 1
mx = px
else:
mx = x
my = self.plot.data[self.plot.y_axis.title][i]
err = self.plot.data['value_error'][i]
x = self.plot.index_mapper.map_screen(mx)
y = self.plot.value_mapper.map_screen(my)
self.pos = (x, y)
self.plot.title = '%s (%.6f, %.6f ± %.6f)' % (self.title, mx, my, err)
self.plot.request_redraw()
def draw(self, gc, view_bounds=None):
if self.pos is None or not self.enabled:
self._last_position = None
else:
x, y = self.pos
super(Snapper, self)._draw_vertical_line(gc, x)
super(Snapper, self)._draw_horizontal_line(gc, y)
super(Snapper, self).draw(gc, view_bounds)
class ZoomWithAspectRatio(ZTool):
def __init__(self, plot, *args, **kwargs):
super(ZoomWithAspectRatio, self).__init__(plot, *args, **kwargs)
self.plot = plot
def selecting_right_up(self, event):
super(ZoomWithAspectRatio, self).selecting_right_up(event)
fx, tx, fy, ty = self._history[-1].next
w, h = (tx - fx, ty - fy)
print(w, h)
self.plot.aspect_ratio = float(w) / h
class PlotDesc(object):
def __init__(self, x, y, data, title, log=False, **kwargs):
self.x = x
self.y = y
self.data = data
self.title = title
self.log = log
self.params = dict(type='line', color='blue')
self.params.update(kwargs)
class McPlot(Plot):
def __init__(self, *args, **kwargs):
super(McPlot, self).__init__(*args, **kwargs)
class McLayout(HasTraits):
layout = Instance(PCont)
reset = Button('Reset')
traits_view = View(
Item('reset', show_label=False),
Item('layout',editor=ComponentEditor(), show_label=False),
width=500, height=500, resizable=True, title='mcplot')
def __init__(self, descs, focus=None, window_title=None):
super(McLayout, self).__init__()
self.descs = descs
self.pans = []
self.zooms = []
self.plots = []
self.focus = descs if focus is None else focus
if window_title:
self.trait_view().title = window_title
if descs:
self.reinit()
def reinit(self):
# clean up plots
for p in self.plots:
map(p.remove_trait, p.__dict__.keys())
self.pans = []
self.zooms = []
num = len(self.focus)
w = int(ceil(sqrt(num)))
h = int(ceil(num / float(w)))
self.layout = PCont(shape=(w, h))
self.plots = map(self.plot_desc, self.focus)
map(self.layout.add, self.plots)
def plot_desc(self, desc):
plot = McPlot(desc.data)
plot.title = desc.title
plot.x_axis.title = desc.x
plot.y_axis.title = desc.y
log = desc.log
ztool = ZTool
if None not in (desc.data['value_low'], desc.data['value_high']):
# 1d
value_scale = 'log' if log else 'linear'
plot.plot((desc.x, desc.y), name='data',
value_scale=value_scale, **desc.params)
# extract and prepare data for error bars
ylow = ArrayDataSource(desc.data['value_low'])
yhigh = ArrayDataSource(desc.data['value_high'])
# create error bars object
x = ArrayDataSource(desc.data[desc.x])
scp = plot.plots['data'][0]
ebp = ErrorBarPlot(index=x,
value_low=ylow,
value_high=yhigh,
index_mapper=scp.index_mapper,
value_mapper=scp.value_mapper,
origin=scp.origin)
# add error bars to plot
plot.add(ebp)
# set aspect ratio
plot.aspect_ratio = 1.0
elif desc.data['imagedata'] is not None:
# 2d
data_name = 'imagedata_log' if log else 'imagedata'
plot.img_plot(data_name, colormap=jet)
# TODO: contour plots: plot.contour_plot('imagedata')
# set apsect ratio
img = desc.data[data_name]
w, h = len(img), len(img[0])
plot.aspect_ratio = float(h) / w
ztool = ZoomWithAspectRatio
# pan
pan = PTool(plot, restrict_to_data=True)
plot.tools.append(pan)
self.pans.append(pan)
# zoom
zoom = ztool(plot, tool_mode='box', always_on=True, drag_button='right')
plot.overlays.append(zoom)
self.zooms.append(zoom)
# line inspector
plot.overlays.append(Snapper(self.layout, plot))
# save
save = SaveTool(self.layout, always_on=True, filename="plot.pdf")
plot.tools.append(save)
# focus
focus = FocusTool(self, desc)
plot.tools.append(focus)
return plot
def _reset_changed(self):
self.focus = self.descs
self.reinit()
|
