Get nearest graphics (#519)

kushalkolar · web-flow · commit a07636c0b334 · 2024-06-18T13:55:51.000-04:00
* bugfix

* center works, need to figure out what's wrong with edge

* allow passing graphic collection to subset

* allow passing graphic collection to subset

* move get_nearest_graphics to utils

* move stuff

* test for get nearest

* black

* more black
diff --git a/fastplotlib/graphics/line_collection.py b/fastplotlib/graphics/line_collection.py
@@ -81,7 +81,7 @@ def cmap(self, args):
         if isinstance(args, str):
             name = args
             transform, alpha = None, 1.0
-        if len(args) == 1:
+        elif len(args) == 1:
             name = args[0]
             transform, alpha = None, None
 
diff --git a/fastplotlib/utils/__init__.py b/fastplotlib/utils/__init__.py
@@ -3,6 +3,7 @@
 
 from .functions import *
 from .gpu import enumerate_adapters, select_adapter, print_wgpu_report
+from ._plot_helpers import *
 
 
 @dataclass
diff --git a/fastplotlib/utils/_plot_helpers.py b/fastplotlib/utils/_plot_helpers.py
@@ -0,0 +1,53 @@
+from typing import Sequence
+
+import numpy as np
+
+from ..graphics._base import Graphic
+from ..graphics._collection_base import GraphicCollection
+
+
+def get_nearest_graphics(
+    pos: tuple[float, float] | tuple[float, float, float],
+    graphics: Sequence[Graphic] | GraphicCollection,
+) -> np.ndarray[Graphic]:
+    """
+    Returns the nearest ``graphics`` to the passed position ``pos`` in world space.
+    Uses the distance between ``pos`` and the center of the bounding sphere for each graphic.
+
+    Parameters
+    ----------
+    pos: (x, y) | (x, y, z)
+        position in world space, z-axis is ignored when calculating L2 norms if ``pos`` is 2D
+
+    graphics: Sequence, i.e. array, list, tuple, etc. of Graphic | GraphicCollection
+        the graphics from which to return a sorted array of graphics in order of closest
+        to furthest graphic
+
+    Returns
+    -------
+    tuple[Graphic]
+        nearest graphics to ``pos`` in order
+
+    """
+
+    if isinstance(graphics, GraphicCollection):
+        graphics = graphics.graphics
+
+    if not all(isinstance(g, Graphic) for g in graphics):
+        raise TypeError("all elements of `graphics` must be Graphic objects")
+
+    pos = np.asarray(pos)
+
+    if pos.shape != (2,) or not pos.shape != (3,):
+        raise TypeError
+
+    # get centers
+    centers = np.empty(shape=(len(graphics), len(pos)))
+    for i in range(centers.shape[0]):
+        centers[i] = graphics[i].world_object.get_world_bounding_sphere()[: len(pos)]
+
+    # l2
+    distances = np.linalg.norm(centers[:, : len(pos)] - pos, ord=2, axis=1)
+
+    sort_indices = np.argsort(distances)
+    return np.asarray(graphics)[sort_indices]
diff --git a/tests/test_plot_helpers.py b/tests/test_plot_helpers.py
@@ -0,0 +1,33 @@
+import numpy as np
+import fastplotlib as fpl
+
+
+def make_circle(center, radius: float, n_points: int = 75) -> np.ndarray:
+    theta = np.linspace(0, 2 * np.pi, n_points)
+    xs = radius * np.sin(theta)
+    ys = radius * np.cos(theta)
+
+    return np.column_stack([xs, ys]) + center
+
+
+def test_get_nearest_graphics():
+    circles = list()
+
+    centers = [[0, 0], [0, 20], [20, 0], [20, 20]]
+
+    for center in centers:
+        circles.append(make_circle(center, 5, n_points=75))
+
+    fig = fpl.Figure()
+
+    lines = fig[0, 0].add_line_collection(circles, cmap="jet", thickness=5)
+
+    fig[0, 0].add_scatter(np.array([[0, 12, 0]]))
+
+    # check distances
+    nearest = fpl.utils.get_nearest_graphics((0, 12), lines)
+    assert nearest[0] is lines[1]  # closest
+    assert nearest[1] is lines[0]
+    assert nearest[2] is lines[3]
+    assert nearest[3] is lines[2]  # furthest
+    assert nearest[-1] is lines[2]
