Simplify lasso_demo example.

anntzer · anntzer · commit 99b5f0b00c82 · 2023-05-10T21:48:47.000+02:00
- No need for a separate Datum data model, which feels a bit
  overengineered.
- Showing how to directly use the collection's array to store the
  relevant info and colormapping for display purposes seems more
  interesting.  This also allows directly using the return value of
  contains_point() to update the collection's array.
- In general it is a good idea to not store a reference to the canvas
  (which can be swapped underneath you).
diff --git a/galleries/examples/event_handling/lasso_demo.py b/galleries/examples/event_handling/lasso_demo.py
@@ -3,12 +3,8 @@
 Lasso Demo
 ==========
 
-Show how to use a lasso to select a set of points and get the indices
-of the selected points.  A callback is used to change the color of the
-selected points
-
-This is currently a proof-of-concept implementation (though it is
-usable as is).  There will be some refinement of the API.
+Use a lasso to select a set of points and get the indices of the selected points.
+A callback is used to change the color of the selected points.
 
 .. note::
     This example exercises the interactive capabilities of Matplotlib, and this
@@ -28,79 +24,43 @@
 from matplotlib.widgets import Lasso
 
 
-class Datum:
-    colorin = mcolors.to_rgba("red")
-    colorout = mcolors.to_rgba("blue")
-
-    def __init__(self, x, y, include=False):
-        self.x = x
-        self.y = y
-        if include:
-            self.color = self.colorin
-        else:
-            self.color = self.colorout
-
-
 class LassoManager:
     def __init__(self, ax, data):
-        self.axes = ax
-        self.canvas = ax.figure.canvas
-        self.data = data
-
-        self.Nxy = len(data)
-
-        facecolors = [d.color for d in data]
-        self.xys = [(d.x, d.y) for d in data]
+        # The information of whether a point has been selected or not is stored in the
+        # collection's array (0 = out, 1 = in), which then gets colormapped to blue
+        # (out) and red (in).
         self.collection = RegularPolyCollection(
-            6, sizes=(100,),
-            facecolors=facecolors,
-            offsets=self.xys,
-            offset_transform=ax.transData)
-
+            6, sizes=(100,), offset_transform=ax.transData,
+            offsets=data, array=np.zeros(len(data)),
+            clim=(0, 1), cmap=mcolors.ListedColormap(["tab:blue", "tab:red"]))
         ax.add_collection(self.collection)
-
-        self.cid_press = self.canvas.mpl_connect('button_press_event',
-                                                 self.on_press)
-        self.cid_release = self.canvas.mpl_connect('button_release_event',
-                                                   self.on_release)
+        canvas = ax.figure.canvas
+        canvas.mpl_connect('button_press_event', self.on_press)
+        canvas.mpl_connect('button_release_event', self.on_release)
 
     def callback(self, verts):
-        facecolors = self.collection.get_facecolors()
-        p = path.Path(verts)
-        ind = p.contains_points(self.xys)
-        for i in range(len(self.xys)):
-            if ind[i]:
-                facecolors[i] = Datum.colorin
-            else:
-                facecolors[i] = Datum.colorout
-
-        self.canvas.draw_idle()
+        data = self.collection.get_offsets()
+        self.collection.set_array(path.Path(verts).contains_points(data))
+        canvas = self.collection.figure.canvas
+        canvas.draw_idle()
         del self.lasso
 
     def on_press(self, event):
-        if self.canvas.widgetlock.locked():
+        canvas = self.collection.figure.canvas
+        if event.inaxes is not self.collection.axes or canvas.widgetlock.locked():
             return
-        if event.inaxes is None:
-            return
-        self.lasso = Lasso(event.inaxes,
-                           (event.xdata, event.ydata),
-                           self.callback)
-        # acquire a lock on the widget drawing
-        self.canvas.widgetlock(self.lasso)
+        self.lasso = Lasso(event.inaxes, (event.xdata, event.ydata), self.callback)
+        canvas.widgetlock(self.lasso)  # acquire a lock on the widget drawing
 
     def on_release(self, event):
-        if hasattr(self, 'lasso') and self.canvas.widgetlock.isowner(self.lasso):
-            self.canvas.widgetlock.release(self.lasso)
+        canvas = self.collection.figure.canvas
+        if hasattr(self, 'lasso') and canvas.widgetlock.isowner(self.lasso):
+            canvas.widgetlock.release(self.lasso)
 
 
 if __name__ == '__main__':
-
     np.random.seed(19680801)
-
-    data = [Datum(*xy) for xy in np.random.rand(100, 2)]
-    ax = plt.axes(xlim=(0, 1), ylim=(0, 1), autoscale_on=False)
-    ax.set_title('Lasso points using left mouse button')
-
-    lman = LassoManager(ax, data)
-
+    ax = plt.figure().add_subplot(
+        xlim=(0, 1), ylim=(0, 1), title='Lasso points using left mouse button')
+    manager = LassoManager(ax, np.random.rand(100, 2))
     plt.show()
