Add add_polygon_selector method to ScatterGraphic and create an e…

…xample for it
lkeegan · Jul 21, 2025 · Jul 22, 2025 · Jul 23, 2025 · Jul 23, 2025 · Sep 29, 2025
commit e2d6b9f1536cda56e95ca6baf71a9e865c93e3e3
@@ -0,0 +1,49 @@
+"""
+Polygon Selectors with ScatterGraphic
+=====================================
+
+Example showing how to use a `PolygonSelector` with a scatter plot.
+"""
+
+# test_example = false
+# sphinx_gallery_pygfx_docs = 'screenshot'
+
+import numpy as np
+import fastplotlib as fpl
+
+# create a figure
+figure = fpl.Figure(
+    (1, 2),
+    size=(700, 560),
+    names=["scatter", "zoomed selection"],
+)
+
+xys = (100 * np.random.random_sample(size=(2000, 2))).astype(np.float32)
+
+# add image
+scatter = figure[0, 0].add_scatter(xys, cmap="jet", sizes=4)
+
+# add polygon selector to scatter graphic
+polygon_selector = scatter.add_polygon_selector()
+
+# add event handler to highlight selected indices and display selected data in zoomed plot
+@polygon_selector.add_event_handler("selection")
+def color_indices(ev):
+    figure[0, 1].clear()
+    scatter.cmap = "jet"
+    scatter.sizes = 4
+    ixs = ev.get_selected_indices()
+    if ixs.size == 0:
+        return
+    scatter.colors[ixs] = 'w'
+    scatter.sizes[ixs] = 8
+    figure[0, 1].add_scatter(ev.get_selected_data(), sizes=16)
+    figure[0, 1].auto_scale()
+
+figure.show()
+
+# NOTE: fpl.loop.run() should not be used for interactive sessions
+# See the "JupyterLab and IPython" section in the user guide
+if __name__ == "__main__":
+    print(__doc__)
+    fpl.loop.run()
@@ -4,6 +4,7 @@
 import pygfx
 
 from ._positions_base import PositionsGraphic
+from .selectors import PolygonSelector
 from .features import (
     PointsSizesFeature,
     UniformSize,
@@ -143,3 +144,46 @@ def sizes(self, value):
 
         elif isinstance(self._sizes, UniformSize):
             self._sizes.set_value(self, value)
+
+    def add_polygon_selector(
+        self,
+        selection: List[tuple[float, float]] = None,
+        **kwargs,
+    ) -> PolygonSelector:
+        """
+        Add a :class:`.PolygonSelector`.
+
+        Selectors are just ``Graphic`` objects, so you can manage, remove, or delete them from a
+        plot area just like any other ``Graphic``.
+
+        Parameters
+        ----------
+        selection: List of positions, optional
+            Initial points for the polygon. If not given or None, you'll start drawing the selection (clicking adds points to the polygon).
+        """
+
+        # remove any nans
+        data = self.data.value[~np.any(np.isnan(self.data.value), axis=1)]
+
+        x_axis_vals = data[:, 0]
+        y_axis_vals = data[:, 1]
+
+        ymin = np.floor(y_axis_vals.min()).astype(int)
+        ymax = np.ceil(y_axis_vals.max()).astype(int)
+        y25p = 0.25 * (ymax - ymin)
+        xmin = np.floor(x_axis_vals.min()).astype(int)
+        xmax = np.ceil(x_axis_vals.max()).astype(int)
+
+        # min/max limits include the data + 25% padding in the y-direction
+        limits = (xmin, xmax, ymin - y25p, ymax + y25p)
+
+        selector = PolygonSelector(
+            selection,
+            limits,
+            parent=self,
+            **kwargs,
+        )
+
+        self._plot_area.add_graphic(selector, center=False)
+
+        return selector
@@ -19,7 +19,7 @@ class MoveInfo:
     """Movement info specific to the polygon selector."""
 
     # The interaction mode: None, 'create', or 'drag'
-    mode: str
+    mode: str | None
 
     # The index of the point in the polygon that is currently being manipulated
     index: int
@@ -89,7 +89,7 @@ def __init__(
         self._resizable = bool(resizable)
 
         BaseSelector.__init__(self, name=name, parent=parent)
-        self._move_info = MoveInfo("none", -1, -1, None, None)
+        self._move_info = MoveInfo(None, -1, -1, None, None)
 
         # Initialize geometry with space for 8 points. The buffers are oversized, so we only need to create new buffers when the allocated space is full.
         # The points are 3D, even though the z-component is always 0. Indices represent the faces (i.e. the triangles).
@@ -207,6 +207,9 @@ def get_selected_data(
         if "Image" in source.__class__.__name__:
             return source.data[ixs[:, 1], ixs[:, 0]]
 
+        if "Scatter" in source.__class__.__name__:
+            return source.data[ixs]
+
         if mode not in ["full", "partial", "ignore"]:
             raise ValueError(
                 f"`mode` must be one of 'full', 'partial', or 'ignore', you have passed {mode}"
@@ -309,10 +312,11 @@ def get_selected_indices(
         Returns
         -------
         Union[np.ndarray, List[np.ndarray]]
-            data indicies of the selection
+            data indices of the selection
             | array of (x, y) indices if the graphic is an image
             | list of indices along the x-dimension for each line if graphic is a line collection
             | array of indices along the x-dimension if graphic is a line
+            | array of indices if graphic is a scatter
         """
         # get indices from source
         source = self._get_source(graphic)
@@ -324,7 +328,10 @@ def get_selected_indices(
         if len(polygon) == 0:
             if "Image" in source.__class__.__name__:
                 return np.zeros((0, 2), np.int32)
-            if "Line" in source.__class__.__name__:
+            if (
+                "Line" in source.__class__.__name__
+                or "Scatter" in graphic.__class__.__name__
+            ):
                 if isinstance(source, GraphicCollection):
                     return [np.zeros((0, 1), np.int32) for _ in source.graphics]
                 else:
@@ -348,7 +355,10 @@ def get_selected_indices(
                         indices.append(p)
             return np.array(indices, np.int32).reshape(-1, 2)
 
-        if "Line" in source.__class__.__name__:
+        if (
+            "Line" in source.__class__.__name__
+            or "Scatter" in source.__class__.__name__
+        ):
             if isinstance(source, GraphicCollection):
                 ixs = list()
                 for g in source.graphics:
@@ -366,7 +376,10 @@ def get_selected_indices(
                     ixs.append(g_ixs)
             else:
                 # map only this graphic
-                points = source.data.value[:2]
+                if "Scatter" in source.__class__.__name__:
+                    points = source.data.value[:, :]
+                else:
+                    points = source.data.value[:2]
                 ixs = np.where(
                     (points[:, 0] >= xmin)
                     & (points[:, 0] <= xmax)