1
""" Defines the LassoSelection controller class.
3
# Major library imports
5
from numpy import array, empty, sometrue, transpose, vstack, zeros
7
# Enthought library imports
8
from enthought.traits.api import Any, Array, Enum, Event, Bool, Instance, \
10
from enthought.kiva.agg import points_in_polygon
13
from enthought.chaco.api import AbstractController, AbstractDataSource, \
14
BaseXYPlot, Base2DPlot
17
class LassoSelection(AbstractController):
18
""" A controller that represents the interaction of "lassoing" a set of
21
"Lassoing" means drawing an arbitrary selection region around the points
22
by dragging the mouse along the outline of the region.
24
# An Nx2 array of points in data space representing all selected points.
25
dataspace_points = Property(Array)
27
# A list of all the selection polygons.
28
disjoint_selections = Property(List)
30
# Fires whenever **dataspace_points** changes, necessitating a redraw of the
34
# Fires when the selection mask changes.
35
selection_changed = Event
37
# Fires when the user release the mouse button and finalizes the selection.
38
selection_completed = Event
40
# If True, the selection mask is updated as the mouse moves, rather
41
# than only at the beginning and end of the selection operation.
42
incremental_select = Bool(False)
44
# The selection mode of the lasso pointer: "include", "exclude" or
45
# "invert" points from the selection. The "include" and "exclude"
46
# settings essentially invert the selection mask. The "invert" setting
47
# differs from "exclude" in that "invert" inverses the selection of all
48
# points the the lasso'ed polygon, while "exclude" operates only on
49
# points included in a previous selection.
50
selection_mode = Enum("include", "exclude", "invert")
52
# The data source that the mask of selected points is attached to. Note
53
# that the indices in this data source must match the indices of the data
55
selection_datasource = Instance(AbstractDataSource)
57
# Mapping from screen space to data space. By default, it is just
61
# The possible event states of this selection tool (overrides
65
# Nothing has been selected, and the user is not dragging the mouse.
67
# The user is dragging the mouse and is actively changing the
69
event_state = Enum('normal', 'selecting')
71
#----------------------------------------------------------------------
73
#----------------------------------------------------------------------
75
# The PlotComponent associated with this tool.
76
_plot = Trait(None, Any)
78
# To support multiple selections, a list of cached selections and the
79
# active selection are maintained. A single list is not used because the
80
# active selection is re-created every time a new point is added via
81
# the vstack function.
82
_active_selection = Array
83
_previous_selections = List(Array)
85
#----------------------------------------------------------------------
87
#----------------------------------------------------------------------
89
def _get_dataspace_points(self):
90
""" Returns a complete list of all selected points.
92
This property exists for backwards compatibility, as the
93
disjoint_selections property is almost always the preferred
94
method of accessingselected points
96
composite = empty((0,2))
97
for region in self.disjoint_selections:
99
composite = vstack((composite, region))
103
def _get_disjoint_selections(self):
104
""" Returns a list of all disjoint selections composed of
105
the previous selections and the active selection
107
if len(self._active_selection) == 0:
108
return self._previous_selections
110
return self._previous_selections + [self._active_selection]
112
#----------------------------------------------------------------------
114
#----------------------------------------------------------------------
116
def normal_left_down(self, event):
117
""" Handles the left mouse button being pressed while the tool is
118
in the 'normal' state.
120
Puts the tool into 'selecting' mode, and starts defining the selection.
122
# We may want to generalize this for the n-dimensional case...
124
self._active_selection = empty((0,2))
126
if self.selection_datasource is not None:
127
self.selection_datasource.metadata['selection'] = zeros(len(self.selection_datasource.get_data()))
128
self.selection_mode = "include"
129
self.event_state = 'selecting'
130
self.selecting_mouse_move(event)
132
if (not event.shift_down) and (not event.control_down):
133
self._previous_selections = []
135
if event.control_down:
136
self.selection_mode = "exclude"
138
self.selection_mode = "include"
139
self.trait_property_changed("disjoint_selections", [], self.disjoint_selections)
142
def selecting_left_up(self, event):
143
""" Handles the left mouse coming up in the 'selecting' state.
145
Completes the selection and switches to the 'normal' state.
147
self.event_state = 'normal'
148
self.selection_completed = True
149
self._update_selection()
151
self._previous_selections.append(self._active_selection)
152
self._active_selection = empty((0,2))
155
def selecting_mouse_move(self, event):
156
""" Handles the mouse moving when the tool is in the 'selecting' state.
158
The selection is extended to the current mouse position.
160
# Translate the event's location to be relative to this container
161
xform = self.component.get_event_transform(event)
162
event.push_transform(xform, caller=self)
163
new_point = self._map_data(array((event.x, event.y)))
164
self._active_selection = vstack((self._active_selection, array((new_point,))))
166
if self.incremental_select:
167
self._update_selection()
168
# Report None for the previous selections
169
self.trait_property_changed("disjoint_selections", None)
172
def selecting_mouse_leave(self, event):
173
""" Handles the mouse leaving the plot when the tool is in the
176
Ends the selection operation.
178
self.selecting_left_up(event)
181
def normal_key_pressed(self, event):
182
""" Handles the user pressing a key in the 'normal' state.
184
If the user presses the Escape key, the tool is reset.
186
if event.character == "Esc":
188
elif event.character == 'a' and event.control_down:
191
elif event.character == 'i' and event.control_down:
192
self.selecting_left_up(None)
193
self.selection_mode = 'invert'
197
#----------------------------------------------------------------------
199
#----------------------------------------------------------------------
201
def _dataspace_points_default(self):
205
""" Resets the selection
207
self.event_state='normal'
208
self._active_selection = empty((0,2))
209
self._previous_selections = []
210
self._update_selection()
212
def _select_all(self):
213
""" Selects all points in the plot. This is done by making a rectangle
214
using the corners of the plot, which is simple but effective. A
215
much cooler, but more time-intensive solution would be to make
216
a selection polygon representing the convex hull.
218
points = [self._map_data(array((self.plot.x, self.plot.y2))),
219
self._map_data(array((self.plot.x2, self.plot.y2))),
220
self._map_data(array((self.plot.x2, self.plot.y))),
221
self._map_data(array((self.plot.x, self.plot.y)))]
223
self._active_selection = numpy.array(points)
224
self._update_selection()
227
def _update_selection(self):
228
""" Sets the selection datasource's 'selection' metadata element
229
to a mask of all the points selected
231
if self.selection_datasource is None:
234
selected_mask = zeros(self.selection_datasource._data.shape, dtype=numpy.int32)
235
data = self._get_data()
237
# Compose the selection mask from the cached selections first, then
238
# the active selection, taking into account the selection mode only
239
# for the active selection
241
for selection in self._previous_selections:
242
selected_mask |= (points_in_polygon(data, selection, False))
244
if self.selection_mode == 'exclude':
245
selected_mask |= (points_in_polygon(data, self._active_selection, False))
246
selected_mask = 1 - selected_mask
248
elif self.selection_mode == 'invert':
249
selected_mask = -1 * (selected_mask -points_in_polygon(data, self._active_selection, False))
251
selected_mask |= (points_in_polygon(data, self._active_selection, False))
253
if sometrue(selected_mask != self.selection_datasource.metadata['selection']):
254
self.selection_datasource.metadata['selection'] = selected_mask
255
self.selection_changed = True
258
def _map_screen(self, points):
259
""" Maps a point in data space to a point in screen space on the plot.
261
Normally this method is a pass-through, but it may do more in
264
return self.plot.map_screen(points)[:,:2]
266
def _map_data(self, point):
267
""" Maps a point in screen space to data space.
269
Normally this method is a pass-through, but for plots that have more
270
data than just (x,y), proper transformations need to happen here.
272
if isinstance(self.plot, Base2DPlot):
273
# Base2DPlot.map_data takes an array of points, for some reason
274
return self.plot.map_data([point])[0]
275
elif isinstance(self.plot, BaseXYPlot):
276
return self.plot.map_data(point, all_values=True)[:2]
278
raise RuntimeError("LassoSelection only supports BaseXY and Base2D plots")
281
""" Returns the datapoints in the plot, as an Nx2 array of (x,y).
283
return transpose(array((self.plot.index.get_data(), self.plot.value.get_data())))
286
#------------------------------------------------------------------------
287
# Property getter/setters
288
#------------------------------------------------------------------------
291
if self._plot is not None:
294
return self.component
296
def _set_plot(self, val):
added
removed
enthought/chaco/tools/lasso_selection.py
