2
Renders a colormapped image of a scalar value field, and a cross section
3
chosen by a line interactor.
6
# Standard library imports
7
from optparse import OptionParser
10
# Major library imports
11
from numpy import array, linspace, meshgrid, nanmin, nanmax, pi, zeros
13
# Enthought library imports
14
from chaco.api import ArrayDataSource, ArrayPlotData, ColorBar, ContourLinePlot, \
15
ColormappedScatterPlot, CMapImagePlot, \
16
ContourPolyPlot, DataRange1D, VPlotContainer, \
17
DataRange2D, GridMapper, GridDataSource, \
18
HPlotContainer, ImageData, LinearMapper, \
19
LinePlot, OverlayPlotContainer, Plot, PlotAxis
20
from chaco.default_colormaps import *
21
from enable.component_editor import ComponentEditor
22
from chaco.tools.api import LineInspector, PanTool, RangeSelection, \
23
RangeSelectionOverlay, ZoomTool
24
from enable.api import Window
25
from traits.api import Any, Array, Callable, CFloat, CInt, Enum, Event, Float, HasTraits, \
26
Int, Instance, Str, Trait, on_trait_change
27
from traitsui.api import Group, Handler, HGroup, Item, View
28
from traitsui.menu import Action, CloseAction, Menu, \
29
MenuBar, NoButtons, Separator
32
class Model(HasTraits):
34
#Traits view definitions:
36
Group(Item('function'),
37
HGroup(Item('npts_x', label="Number X Points"),
38
Item('npts_y', label="Number Y Points")),
39
HGroup(Item('min_x', label="Min X value"),
40
Item('max_x', label="Max X value")),
41
HGroup(Item('min_y', label="Min Y value"),
42
Item('max_y', label="Max Y value"))),
43
buttons=["OK", "Cancel"])
45
function = Str("tanh(x**2+y)*cos(y)*jn(0,x+y*2)")
52
min_y = CFloat(-1.5*pi)
53
max_y = CFloat(1.5*pi)
64
def __init__(self, *args, **kwargs):
65
super(Model, self).__init__(*args, **kwargs)
68
def compute_model(self):
69
# The xs and ys used for the image plot range need to be the
71
self.xs = linspace(self.min_x, self.max_x, self.npts_x+1)
72
self.ys = linspace(self.min_y, self.max_y, self.npts_y+1)
74
# The grid of points at which we will evaluate the 2D function
75
# is located at cell centers, so use halfsteps from the
76
# min/max values (which are edges)
77
xstep = (self.max_x - self.min_x) / self.npts_x
78
ystep = (self.max_y - self.min_y) / self.npts_y
79
gridx = linspace(self.min_x+xstep/2, self.max_x-xstep/2, self.npts_x)
80
gridy = linspace(self.min_y+xstep/2, self.max_y-xstep/2, self.npts_y)
81
x, y = meshgrid(gridx, gridy)
84
exec "from scipy import *" in d
85
exec "from scipy.special import *" in d
86
self.zs = eval(self.function, d)
87
self.minz = nanmin(self.zs)
88
self.maxz = nanmax(self.zs)
89
self.model_changed = True
90
self._function = self.function
92
self.set(function = self._function, trait_change_notify=False)
94
def _anytrait_changed(self, name, value):
95
if name in ['function', 'npts_x', 'npts_y',
96
'min_x', 'max_x', 'min_y', 'max_y']:
100
class PlotUI(HasTraits):
102
#Traits view definitions:
104
Group(Item('container',
105
editor=ComponentEditor(size=(800,600)),
110
plot_edit_view = View(
111
Group(Item('num_levels'),
113
buttons=["OK","Cancel"])
117
colormap = Enum(color_map_name_dict.keys())
119
#---------------------------------------------------------------------------
121
#---------------------------------------------------------------------------
123
_image_index = Instance(GridDataSource)
124
_image_value = Instance(ImageData)
126
_cmap = Trait(jet, Callable)
128
#---------------------------------------------------------------------------
129
# Public View interface
130
#---------------------------------------------------------------------------
132
def __init__(self, *args, **kwargs):
133
super(PlotUI, self).__init__(*args, **kwargs)
136
def create_plot(self):
138
# Create the mapper, etc
139
self._image_index = GridDataSource(array([]),
141
sort_order=("ascending","ascending"))
142
image_index_range = DataRange2D(self._image_index)
143
self._image_index.on_trait_change(self._metadata_changed,
146
self._image_value = ImageData(data=array([]), value_depth=1)
147
image_value_range = DataRange1D(self._image_value)
151
# Create the contour plots
152
self.polyplot = ContourPolyPlot(index=self._image_index,
153
value=self._image_value,
154
index_mapper=GridMapper(range=
157
self._cmap(image_value_range),
158
levels=self.num_levels)
160
self.lineplot = ContourLinePlot(index=self._image_index,
161
value=self._image_value,
162
index_mapper=GridMapper(range=
163
self.polyplot.index_mapper.range),
164
levels=self.num_levels)
167
# Add a left axis to the plot
168
left = PlotAxis(orientation='left',
170
mapper=self.polyplot.index_mapper._ymapper,
171
component=self.polyplot)
172
self.polyplot.overlays.append(left)
174
# Add a bottom axis to the plot
175
bottom = PlotAxis(orientation='bottom',
177
mapper=self.polyplot.index_mapper._xmapper,
178
component=self.polyplot)
179
self.polyplot.overlays.append(bottom)
182
# Add some tools to the plot
183
self.polyplot.tools.append(PanTool(self.polyplot,
184
constrain_key="shift"))
185
self.polyplot.overlays.append(ZoomTool(component=self.polyplot,
186
tool_mode="box", always_on=False))
187
self.polyplot.overlays.append(LineInspector(component=self.polyplot,
189
inspect_mode="indexed",
193
self.polyplot.overlays.append(LineInspector(component=self.polyplot,
195
inspect_mode="indexed",
200
# Add these two plots to one container
201
contour_container = OverlayPlotContainer(padding=20,
204
contour_container.add(self.polyplot)
205
contour_container.add(self.lineplot)
209
cbar_index_mapper = LinearMapper(range=image_value_range)
210
self.colorbar = ColorBar(index_mapper=cbar_index_mapper,
212
padding_top=self.polyplot.padding_top,
213
padding_bottom=self.polyplot.padding_bottom,
218
self.pd = ArrayPlotData(line_index = array([]),
219
line_value = array([]),
220
scatter_index = array([]),
221
scatter_value = array([]),
222
scatter_color = array([]))
224
self.cross_plot = Plot(self.pd, resizable="h")
225
self.cross_plot.height = 100
226
self.cross_plot.padding = 20
227
self.cross_plot.plot(("line_index", "line_value"),
229
self.cross_plot.plot(("scatter_index","scatter_value","scatter_color"),
232
color_mapper=self._cmap(image_value_range),
236
self.cross_plot.index_range = self.polyplot.index_range.x_range
238
self.pd.set_data("line_index2", array([]))
239
self.pd.set_data("line_value2", array([]))
240
self.pd.set_data("scatter_index2", array([]))
241
self.pd.set_data("scatter_value2", array([]))
242
self.pd.set_data("scatter_color2", array([]))
244
self.cross_plot2 = Plot(self.pd, width = 140, orientation="v", resizable="v", padding=20, padding_bottom=160)
245
self.cross_plot2.plot(("line_index2", "line_value2"),
247
self.cross_plot2.plot(("scatter_index2","scatter_value2","scatter_color2"),
250
color_mapper=self._cmap(image_value_range),
254
self.cross_plot2.index_range = self.polyplot.index_range.y_range
258
# Create a container and add components
259
self.container = HPlotContainer(padding=40, fill_padding=True,
260
bgcolor = "white", use_backbuffer=False)
261
inner_cont = VPlotContainer(padding=0, use_backbuffer=True)
262
inner_cont.add(self.cross_plot)
263
inner_cont.add(contour_container)
264
self.container.add(self.colorbar)
265
self.container.add(inner_cont)
266
self.container.add(self.cross_plot2)
269
def update(self, model):
270
self.minz = model.minz
271
self.maxz = model.maxz
272
self.colorbar.index_mapper.range.low = self.minz
273
self.colorbar.index_mapper.range.high = self.maxz
274
self._image_index.set_data(model.xs, model.ys)
275
self._image_value.data = model.zs
276
self.pd.set_data("line_index", model.xs)
277
self.pd.set_data("line_index2", model.ys)
278
self.container.invalidate_draw()
279
self.container.request_redraw()
282
#---------------------------------------------------------------------------
284
#---------------------------------------------------------------------------
286
def _metadata_changed(self, old, new):
287
""" This function takes out a cross section from the image data, based
288
on the line inspector selections, and updates the line and scatter
291
self.cross_plot.value_range.low = self.minz
292
self.cross_plot.value_range.high = self.maxz
293
self.cross_plot2.value_range.low = self.minz
294
self.cross_plot2.value_range.high = self.maxz
295
if self._image_index.metadata.has_key("selections"):
296
x_ndx, y_ndx = self._image_index.metadata["selections"]
298
self.pd.set_data("line_value",
299
self._image_value.data[y_ndx,:])
300
self.pd.set_data("line_value2",
301
self._image_value.data[:,x_ndx])
302
xdata, ydata = self._image_index.get_data()
303
xdata, ydata = xdata.get_data(), ydata.get_data()
304
self.pd.set_data("scatter_index", array([xdata[x_ndx]]))
305
self.pd.set_data("scatter_index2", array([ydata[y_ndx]]))
306
self.pd.set_data("scatter_value",
307
array([self._image_value.data[y_ndx, x_ndx]]))
308
self.pd.set_data("scatter_value2",
309
array([self._image_value.data[y_ndx, x_ndx]]))
310
self.pd.set_data("scatter_color",
311
array([self._image_value.data[y_ndx, x_ndx]]))
312
self.pd.set_data("scatter_color2",
313
array([self._image_value.data[y_ndx, x_ndx]]))
315
self.pd.set_data("scatter_value", array([]))
316
self.pd.set_data("scatter_value2", array([]))
317
self.pd.set_data("line_value", array([]))
318
self.pd.set_data("line_value2", array([]))
320
def _colormap_changed(self):
321
self._cmap = color_map_name_dict[self.colormap]
322
if hasattr(self, "polyplot"):
323
value_range = self.polyplot.color_mapper.range
324
self.polyplot.color_mapper = self._cmap(value_range)
325
value_range = self.cross_plot.color_mapper.range
326
self.cross_plot.color_mapper = self._cmap(value_range)
327
# FIXME: change when we decide how best to update plots using
328
# the shared colormap in plot object
329
self.cross_plot.plots["dot"][0].color_mapper = self._cmap(value_range)
330
self.cross_plot2.plots["dot"][0].color_mapper = self._cmap(value_range)
331
self.container.request_redraw()
333
def _num_levels_changed(self):
334
if self.num_levels > 3:
335
self.polyplot.levels = self.num_levels
336
self.lineplot.levels = self.num_levels
340
class Controller(Handler):
342
#---------------------------------------------------------------------------
344
#---------------------------------------------------------------------------
346
model = Instance(Model)
347
view = Instance(PlotUI)
349
#---------------------------------------------------------------------------
351
#---------------------------------------------------------------------------
353
def init(self, info):
354
self.model = info.object.model
355
self.view = info.object.view
356
self.model.on_trait_change(self._model_changed, "model_changed")
359
#---------------------------------------------------------------------------
360
# Public Controller interface
361
#---------------------------------------------------------------------------
363
def edit_model(self, ui_info):
364
self.model.configure_traits()
366
def edit_plot(self, ui_info):
367
self.view.configure_traits(view="plot_edit_view")
370
#---------------------------------------------------------------------------
371
# Private Controller interface
372
#---------------------------------------------------------------------------
374
def _model_changed(self):
375
if self.view is not None:
376
self.view.update(self.model)
378
class ModelView(HasTraits):
380
model = Instance(Model)
381
view = Instance(PlotUI)
382
traits_view = View(Item('@view',
384
menubar=MenuBar(Menu(Action(name="Edit Model",
385
action="edit_model"),
386
Action(name="Edit Plot",
390
handler = Controller,
391
title = "Function Inspector",
394
@on_trait_change('model, view')
395
def update_view(self):
396
if self.model is not None and self.view is not None:
397
self.view.update(self.model)
399
options_dict = {'colormap' : "jet",
401
'function' : "tanh(x**2+y)*cos(y)*jn(0,x+y*2)"}
402
model=Model(**options_dict)
403
view=PlotUI(**options_dict)
404
popup = ModelView(model=model, view=view)
406
def show_plot(**kwargs):
407
model = Model(**kwargs)
408
view = PlotUI(**kwargs)
409
modelview=ModelView(model=model, view=view)
410
modelview.configure_traits()
417
usage = "usage: %prog [options]"
418
parser = OptionParser(usage=usage, version="%prog 1.0")
420
parser.add_option("-c", "--colormap",
421
action="store", type="string", dest="colormap", default="jet",
422
metavar="CMAP", help="choose a default colormapper")
424
parser.add_option("-n", "--nlevels",
425
action="store", type="int", dest="num_levels", default=15,
426
help="number countour levels to plot [default: %default]")
428
parser.add_option("-f", "--function",
429
action="store", type="string", dest="function",
430
default="tanh(x**2+y)*cos(y)*jn(0,x+y*2)",
431
help="function of x and y [default: %default]")
433
opts, args = parser.parse_args(argv[1:])
436
parser.error("Incorrect number of arguments")
438
show_plot(colormap=opts.colormap, num_levels=opts.num_levels,
439
function=opts.function)
441
if __name__ == "__main__":
added
removed
examples/demo/advanced/scalar_image_function_inspector_old.py
