diff --git a/chaco/tools/api.py b/chaco/tools/api.py index e02b816cb..e3b147295 100644 --- a/chaco/tools/api.py +++ b/chaco/tools/api.py @@ -19,6 +19,7 @@ from .range_selection import RangeSelection from .range_selection_2d import RangeSelection2D from .range_selection_overlay import RangeSelectionOverlay +from .rectangular_selection import RectangularSelection from .regression_lasso import RegressionLasso, RegressionOverlay from .save_tool import SaveTool from .scatter_inspector import ScatterInspector diff --git a/chaco/tools/rectangular_selection.py b/chaco/tools/rectangular_selection.py new file mode 100644 index 000000000..fdd9dbc5e --- /dev/null +++ b/chaco/tools/rectangular_selection.py @@ -0,0 +1,51 @@ +""" Defines the RectangularSelection controller class. +""" +import numpy as np + +from chaco.tools.lasso_selection import LassoSelection +from traits.api import ArrayOrNone + + +class RectangularSelection(LassoSelection): + """ A lasso selection tool whose selection shape is rectangular + """ + + #: The first click. This represents a corner of the rectangle. + first_corner = ArrayOrNone(shape=(2,)) + + def selecting_mouse_move(self, event): + """ This function is the same as the super except that it injects + `_make_rectangle` as the `_active_selection` assignment. + """ + # Translate the event's location to be relative to this container + xform = self.component.get_event_transform(event) + event.push_transform(xform, caller=self) + new_point = self._map_data(np.array((event.x, event.y))) + if self.first_corner is None: + self.first_corner = new_point + self._active_selection = self._make_rectangle( + self.first_corner, new_point) + self.updated = True + if self.incremental_select: + self._update_selection() + # Report None for the previous selections + self.trait_property_changed("disjoint_selections", None) + + def selecting_mouse_up(self, event): + super(RectangularSelection, self).selecting_mouse_up(event) + # Clear the first click + self.first_corner = None + + def _make_rectangle(self, p1, p2): + """ Makes an array that represents that path that follows the + corner points of the rectangle with two corners p1 and p2: + *-----p2 + | | + p1----* + """ + return np.array([ + p1, + [p1[0], p2[1]], + p2, + [p2[0], p1[1]] + ]) diff --git a/chaco/tools/tests/test_rectangular_selection_tool.py b/chaco/tools/tests/test_rectangular_selection_tool.py new file mode 100644 index 000000000..954913724 --- /dev/null +++ b/chaco/tools/tests/test_rectangular_selection_tool.py @@ -0,0 +1,53 @@ +import unittest + +import numpy as np + +from chaco.array_plot_data import ArrayPlotData +from chaco.plot import Plot +from chaco.tools.rectangular_selection import RectangularSelection +from enable.testing import EnableTestAssistant + + +class RectangularSelectionTestCase(EnableTestAssistant, unittest.TestCase): + + def test_selection_mask(self): + + plot_data = ArrayPlotData() + plot = Plot(plot_data) + arr = np.array([-2, -1, 1, 2]) + plot_data.set_data("x", arr) + plot_data.set_data("y", arr) + splot = plot.plot(('x', 'y'), type='scatter')[0] + tool = RectangularSelection( + component=splot, + selection_datasource=splot.index, + metadata_name='selections', + ) + splot.tools.append(tool) + + # Set the cursor start and stop positions to be such + # that the middle two points of the four possible are selected. + cursor_start = splot.map_screen([-1.5, -1.5])[0] + cursor_stop = splot.map_screen([1.5, 1.5])[0] + + self.mouse_down( + interactor=tool, + x=cursor_start[0], + y=cursor_start[1] + ) + + self.mouse_move( + interactor=tool, + x=cursor_stop[0], + y=cursor_stop[1] + ) + + self.mouse_up( + interactor=tool, + x=cursor_stop[0], + y=cursor_stop[1] + ) + + expected_mask = [False, True, True, False] + selection_mask = list(splot.index.metadata['selections']) + self.assertEqual(expected_mask, selection_mask) diff --git a/examples/demo/basic/scatter_rect_select.py b/examples/demo/basic/scatter_rect_select.py new file mode 100644 index 000000000..cf6d2e845 --- /dev/null +++ b/examples/demo/basic/scatter_rect_select.py @@ -0,0 +1,135 @@ +""" +Rectangular selection of data points + +Draws a simple scatterplot of random data. Drag the mouse to use the +selector, which allows you to select points via a bounding box. + +Upon completion of the selection operation, the indices of the selected +points are printed to the console and highlighted visually. +""" + +import sys + +# Major library imports +from numpy import sort, compress, arange +from numpy.random import random + +# Enthought library imports +from enable.api import Component, ComponentEditor +from traits.api import HasTraits, Instance +from traitsui.api import Item, Group, View + +# Chaco imports +from chaco.api import ( + ArrayPlotData, Plot, LassoOverlay, ScatterInspectorOverlay) +from chaco.tools.api import RectangularSelection, ScatterInspector + + + +# =============================================================================== +# # Create the Chaco plot. +# =============================================================================== +def _create_plot_component(): + # Create some data + npts = 200 + x = sort(random(npts)) + y = random(npts) + + # Create a plot data obect and give it this data + pd = ArrayPlotData() + pd.set_data("index", x) + pd.set_data("value", y) + + # Create the plot + plot = Plot(pd) + plot.plot(("index", "value"), + type="scatter", + name="my_plot", + marker="circle", + index_sort="ascending", + color="red", + marker_size=4, + bgcolor="white") + + # Tweak some of the plot properties + plot.title = "Scatter Plot With Rectangular Selection" + plot.line_width = 1 + plot.padding = 50 + + # Right now, some of the tools are a little invasive, and we need the + # actual ScatterPlot object to give to them + my_plot = plot.plots["my_plot"][0] + + # Attach some tools to the plot + rect_selection = RectangularSelection( + component=my_plot, + selection_datasource=my_plot.index, + drag_button="left", + metadata_name='selections', + ) + my_plot.tools.append(rect_selection) + my_plot.tools.append(ScatterInspector(my_plot, selection_mode='toggle')) + my_plot.active_tool = rect_selection + + lasso_overlay = LassoOverlay(lasso_selection=rect_selection, + component=my_plot) + my_plot.overlays.append(lasso_overlay) + + scatter_overlay = ScatterInspectorOverlay( + component=my_plot, + selection_color='cornflowerblue', + selection_marker_size=int(my_plot.marker_size)+3, + selection_marker='circle' + ) + my_plot.overlays.append(scatter_overlay) + + return plot + + +# =============================================================================== +# Attributes to use for the plot view. +size = (650, 650) +title = "Scatter plot with selection" +bg_color = "lightgray" + + +# =============================================================================== +# # Demo class that is used by the demo.py application. +# =============================================================================== +class Demo(HasTraits): + plot = Instance(Component) + + traits_view = View( + Group( + Item('plot', editor=ComponentEditor(size=size), + show_label=False), + orientation="vertical"), + resizable=True, title=title + ) + + def _selection_changed(self): + mask = self.index_datasource.metadata['selections'] + print("New selection: ") + print(compress(mask, arange(len(mask)))) + # Ensure that the points are printed immediately: + sys.stdout.flush() + + def _plot_default(self): + plot = _create_plot_component() + + # Retrieve the plot hooked to the RectangularSelection tool. + my_plot = plot.plots["my_plot"][0] + rect_selection = my_plot.active_tool + + # Set up the trait handler for the selection + self.index_datasource = my_plot.index + rect_selection.on_trait_change(self._selection_changed, + 'selection_changed') + + return plot + + +demo = Demo() + +if __name__ == "__main__": + demo.configure_traits()