Add rectangular selection tool

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

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]]
+        ])

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)

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()