Add LassoSelector widget with demo.

Note: I put the demo in the "widgets" directory even though the `Lasso` demo is in "event_handling".
matplotlib · efiring · Mar 17, 2012 · Feb 27, 2012 · Feb 27, 2012 · Feb 27, 2012
commit 8345d4c6afcd8afa9d646d554440de27fcc7d49c
diff --git a/examples/widgets/lasso_selector_demo.py b/examples/widgets/lasso_selector_demo.py
@@ -0,0 +1,81 @@
+import numpy as np
+
+from matplotlib.widgets import LassoSelector
+from matplotlib.nxutils import points_inside_poly
+
+
+class SelectFromCollection(object):
+    """Select indices from a matplotlib collection using :class:`LassoSelector`.
+
+    Selected indices are saved in the `ind` attribute. This tool highlights
+    selected points by fading them out (i.e., reducing their alpha values).
+    If your collection has alpha < 1, this tool will permanently alter them.
+
+    Note that this tool selects collection objects based on their *origins*
+    (i.e., `offsets`).
+
+    Parameters
+    ----------
+    ax : :class:`Axes`
+        Axes to interact with.
+
+    collection : :class:`Collection`
+        Collection you want to select from.
+
+    alpha_other : 0 <= float <= 1
+        To highlight a selection, this tool sets all selected points to an
+        alpha value of 1 and non-selected points to `alpha_other`.
+    """
+    def __init__(self, ax, collection, alpha_other=0.3):
+        self.canvas = ax.figure.canvas
+        self.collection = collection
+        self.alpha_other = alpha_other
+
+        self.xys = collection.get_offsets()
+        self.Npts = len(self.xys)
+
+        # Ensure that we have separate colors for each object
+        self.fc = collection.get_facecolors()
+        if len(self.fc) == 0:
+            raise ValueError('Collection must have a facecolor')
+        elif len(self.fc) == 1:
+            self.fc = np.tile(self.fc, self.Npts).reshape(self.Npts, -1)
+
+        self.lasso = LassoSelector(ax, onselect=self.onselect)
+        self.ind = []
+
+    def onselect(self, verts):
+        self.ind = np.nonzero(points_inside_poly(self.xys, verts))[0]
+        self.fc[:, -1] = self.alpha_other
+        self.fc[self.ind, -1] = 1
+        self.collection.set_facecolors(self.fc)
+        self.canvas.draw_idle()
+
+    def disconnect(self):
+        self.lasso.disconnect_events()
+        self.fc[:, -1] = 1
+        self.collection.set_facecolors(self.fc)
+        self.canvas.draw_idle()
+
+
+if __name__ == '__main__':
+    import matplotlib.pyplot as plt
+
+    plt.ion()
+    data = np.random.rand(100, 2)
+
+    subplot_kw = dict(xlim=(0,1), ylim=(0,1), autoscale_on=False)
+    fig, ax = plt.subplots(subplot_kw=subplot_kw)
+
+    pts = ax.scatter(data[:, 0], data[:, 1], s=80)
+    selector = SelectFromCollection(ax, pts)
+
+    plt.draw()
+    raw_input('Press any key to accept selected points')
+    print "Selected points:"
+    print selector.xys[selector.ind]
+    selector.disconnect()
+
+    # Block end of script so you can check that lasso is disconnected.
+    raw_input('Press any key to quit')
+
diff --git a/lib/matplotlib/widgets.py b/lib/matplotlib/widgets.py
@@ -1409,6 +1409,71 @@ def get_active(self):
         """ Get status of active mode (boolean variable)"""
         return self.active
 
+
+class LassoSelector(AxesWidget):
+    def __init__(self, ax, onselect=None, useblit=True, lineprops=None):
+        AxesWidget.__init__(self, ax)
+
+        self.useblit = useblit
+        self.onselect = onselect
+        self.verts = None
+
+        if lineprops is None:
+            lineprops = dict()
+        self.line = Line2D([], [], **lineprops)
+        self.line.set_visible(False)
+        self.ax.add_line(self.line)
+
+        self.connect_event('button_press_event', self.onpress)
+        self.connect_event('button_release_event', self.onrelease)
+        self.connect_event('motion_notify_event', self.onmove)
+        self.connect_event('draw_event', self.update_background)
+
+    def ignore(self, event):
+        wrong_button = hasattr(event, 'button') and event.button != 1
+        return not self.active or wrong_button
+
+    def onpress(self, event):
+        if self.ignore(event) or event.inaxes != self.ax:
+            return
+        self.verts = [(event.xdata, event.ydata)]
+        self.line.set_visible(True)
+
+    def onrelease(self, event):
+        if self.ignore(event):
+            return
+        if self.verts is not None:
+            if event.inaxes == self.ax:
+                self.verts.append((event.xdata, event.ydata))
+            self.onselect(self.verts)
+        self.line.set_data([[], []])
+        self.line.set_visible(False)
+        self.verts = None
+
+    def onmove(self, event):
+        if self.ignore(event) or event.inaxes != self.ax:
+            return
+        if self.verts is None: return
+        if event.inaxes != self.ax: return
+        if event.button!=1: return
+        self.verts.append((event.xdata, event.ydata))
+
+        self.line.set_data(zip(*self.verts))
+
+        if self.useblit:
+            self.canvas.restore_region(self.background)
+            self.ax.draw_artist(self.line)
+            self.canvas.blit(self.ax.bbox)
+        else:
+            self.canvas.draw_idle()
+
+    def update_background(self, event):
+        if self.ignore(event):
+            return
+        if self.useblit:
+            self.background = self.canvas.copy_from_bbox(self.ax.bbox)
+
+
 class Lasso(AxesWidget):
     def __init__(self, ax, xy, callback=None, useblit=True):
         AxesWidget.__init__(self, ax)
@@ -1452,3 +1517,4 @@ def onmove(self, event):
             self.canvas.blit(self.ax.bbox)
         else:
             self.canvas.draw_idle()
+