Merge pull request #25848 from anntzer/ld

timhoffm · web-flow · commit 8faa835c81d9 · 2023-05-11T05:28:49.000+02:00
Simplify lasso_demo example.
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()
