2
Implementation of a plut using a custom overlay and tool
5
from __future__ import with_statement
9
from enthought.traits.api import HasTraits, Instance, Enum
10
from enthought.traits.ui.api import View, Item
11
from enthought.enable.api import ComponentEditor
12
from enthought.chaco.api import Plot, ArrayPlotData, AbstractOverlay
13
from enthought.enable.api import BaseTool
14
from enthought.enable.markers import DOT_MARKER, DotMarker
16
class BoxSelectTool(BaseTool):
17
""" Tool for selecting all points within a box
19
There are 2 states for this tool, normal and selecting. While the
20
left mouse button is down the metadata on the datasources will be
21
updated with the current selected bounds.
23
Note that the tool does not actually store the selected point, but the
27
event_state = Enum("normal", "selecting")
29
def normal_left_down(self, event):
30
self.event_state = "selecting"
31
self.selecting_mouse_move(event)
33
def selecting_left_up(self, event):
34
self.event_state = "normal"
36
def selecting_mouse_move(self, event):
37
x1, y1 = self.map_to_data(event.x-25, event.y-25)
38
x2, y2 = self.map_to_data(event.x+25, event.y+25)
40
index_datasource = self.component.index
41
index_datasource.metadata['selections'] = (x1, x2)
43
value_datasource = self.component.value
44
value_datasource.metadata['selections'] = (y1, y2)
46
self.component.request_redraw()
48
def map_to_data(self, x, y):
49
""" Returns the data space coordinates of the given x and y.
51
Takes into account orientation of the plot and the axis setting.
55
if plot.orientation == "h":
56
index = plot.x_mapper.map_data(x)
57
value = plot.y_mapper.map_data(y)
59
index = plot.y_mapper.map_data(y)
60
value = plot.x_mapper.map_data(x)
65
class XRayOverlay(AbstractOverlay):
66
""" Overlay which draws scatter markers on top of plot data points.
68
This overlay should be combined with a tool which updates the
69
datasources metadata with selection bounds.
74
def overlay(self, component, gc, view_bounds=None, mode='normal'):
75
x_range = self._get_selection_index_screen_range()
76
y_range = self._get_selection_value_screen_range()
86
gc.set_fill_color((1.0,1.0,1.0))
87
gc.rect(x1, y1, x2-x1, y2-y1)
90
pts = self._get_selected_points()
93
screen_pts = self.component.map_screen(pts)
94
if hasattr(gc, 'draw_marker_at_points'):
95
gc.draw_marker_at_points(screen_pts, 3, DOT_MARKER)
98
for sx,sy in screen_pts:
99
gc.translate_ctm(sx, sy)
101
self.marker.add_to_path(gc, 3)
102
gc.draw_path(self.marker.draw_mode)
103
gc.translate_ctm(-sx, -sy)
106
def _get_selected_points(self):
107
""" gets all the points within the bounds defined in the datasources
110
index_datasource = self.component.index
111
index_selection = index_datasource.metadata['selections']
112
index = index_datasource.get_data()
114
value_datasource = self.component.value
115
value_selection = value_datasource.metadata['selections']
116
value = value_datasource.get_data()
118
x_indices = numpy.where((index > index_selection[0]) & (index < index_selection[-1]))
119
y_indices = numpy.where((value > value_selection[0]) & (value < value_selection[-1]))
121
indices = list(set(x_indices[0]) & set(y_indices[0]))
123
sel_index = index[indices]
124
sel_value = value[indices]
126
return zip(sel_index, sel_value)
128
def _get_selection_index_screen_range(self):
129
""" maps the selected bounds which were set by the tool into screen
130
space. The screen space points can be used for drawing the overlay
132
index_datasource = self.component.index
133
index_mapper = self.component.index_mapper
134
index_selection = index_datasource.metadata['selections']
135
return tuple(index_mapper.map_screen(numpy.array(index_selection)))
137
def _get_selection_value_screen_range(self):
138
""" maps the selected bounds which were set by the tool into screen
139
space. The screen space points can be used for drawing the overlay
141
value_datasource = self.component.value
142
value_mapper = self.component.value_mapper
143
value_selection = value_datasource.metadata['selections']
144
return tuple(value_mapper.map_screen(numpy.array(value_selection)))
146
class PlotExample(HasTraits):
148
plot = Instance(Plot)
150
traits_view = View(Item('plot', editor=ComponentEditor()),
151
width=600, height=600)
153
def __init__(self, index, value, *args, **kw):
154
super(PlotExample, self).__init__(*args, **kw)
156
plot_data = ArrayPlotData(index=index)
157
plot_data.set_data('value', value)
159
self.plot = Plot(plot_data)
160
line = self.plot.plot(('index', 'value'))[0]
162
line.overlays.append(XRayOverlay(line))
163
line.tools.append(BoxSelectTool(line))
165
index = numpy.arange(0, 25, 0.25)
166
value = numpy.sin(index) + numpy.arange(0, 10, 0.1)
168
example = PlotExample(index, value)
169
example.configure_traits()