Update find_nearest_contour to not use collections attribute

rcomer · anntzer · rcomer · commit 25ec2c0e4306 · 2023-10-26T19:10:17.000+01:00
Co-authored-by: Antony Lee &lt;anntzer.lee@gmail.com&gt;
diff --git a/lib/matplotlib/contour.py b/lib/matplotlib/contour.py
@@ -2,6 +2,7 @@
 Classes to support contour plotting and labelling for the Axes class.
 """
 
+from contextlib import ExitStack
 import functools
 import math
 from numbers import Integral
@@ -1409,64 +1410,39 @@ def find_nearest_contour(self, x, y, indices=None, pixel=True):
 
         Returns
         -------
-        contour : `.Collection`
-            The contour that is closest to ``(x, y)``.
-        segment : int
-            The index of the `.Path` in *contour* that is closest to
-            ``(x, y)``.
+        path : int
+            The index of the path that is closest to ``(x, y)``.  Each path corresponds
+            to one contour level.
+        subpath : int
+            The index within that closest path of the subpath that is closest to
+            ``(x, y)``.  Each subpath corresponds to one unbroken contour line.
         index : int
-            The index of the path segment in *segment* that is closest to
+            The index of the vertices within that subpath that are closest to
             ``(x, y)``.
         xmin, ymin : float
             The point in the contour plot that is closest to ``(x, y)``.
         d2 : float
             The squared distance from ``(xmin, ymin)`` to ``(x, y)``.
         """
+        segment = index = d2 = None
 
-        # This function uses a method that is probably quite
-        # inefficient based on converting each contour segment to
-        # pixel coordinates and then comparing the given point to
-        # those coordinates for each contour.  This will probably be
-        # quite slow for complex contours, but for normal use it works
-        # sufficiently well that the time is not noticeable.
-        # Nonetheless, improvements could probably be made.
+        with ExitStack() as stack:
+            if not pixel:
+                # _find_nearest_contour works in pixel space. We want axes space, so
+                # effectively disable the transformation here by setting to identity.
+                stack.enter_context(self._cm_set(
+                    transform=mtransforms.IdentityTransform()))
 
-        if self.filled:
-            raise ValueError("Method does not support filled contours.")
+            i_level, i_vtx, (xmin, ymin) = self._find_nearest_contour((x, y), indices)
 
-        if indices is None:
-            indices = range(len(self.collections))
+        if i_level is not None:
+            cc_cumlens = np.cumsum(
+                [*map(len, self._paths[i_level]._iter_connected_components())])
+            segment = cc_cumlens.searchsorted(i_vtx, "right")
+            index = i_vtx if segment == 0 else i_vtx - cc_cumlens[segment - 1]
+            d2 = (xmin-x)**2 + (ymin-y)**2
 
-        d2min = np.inf
-        conmin = None
-        segmin = None
-        imin = None
-        xmin = None
-        ymin = None
-
-        point = np.array([x, y])
-
-        for icon in indices:
-            con = self.collections[icon]
-            trans = con.get_transform()
-            paths = con.get_paths()
-
-            for segNum, linepath in enumerate(paths):
-                lc = linepath.vertices
-                # transfer all data points to screen coordinates if desired
-                if pixel:
-                    lc = trans.transform(lc)
-
-                d2, xc, leg = _find_closest_point_on_path(lc, point)
-                if d2 < d2min:
-                    d2min = d2
-                    conmin = icon
-                    segmin = segNum
-                    imin = leg[1]
-                    xmin = xc[0]
-                    ymin = xc[1]
-
-        return (conmin, segmin, imin, xmin, ymin, d2min)
+        return (i_level, segment, index, xmin, ymin, d2)
 
     def draw(self, renderer):
         paths = self._paths
diff --git a/lib/matplotlib/contour.pyi b/lib/matplotlib/contour.pyi
@@ -159,6 +159,6 @@ class ContourSet(ContourLabeler, Collection):
     ) -> tuple[list[Artist], list[str]]: ...
     def find_nearest_contour(
         self, x: float, y: float, indices: Iterable[int] | None = ..., pixel: bool = ...
-    ) -> tuple[Collection, int, int, float, float, float]: ...
+    ) -> tuple[int, int, int, float, float, float]: ...
 
 class QuadContourSet(ContourSet): ...
diff --git a/lib/matplotlib/tests/test_contour.py b/lib/matplotlib/tests/test_contour.py
@@ -557,15 +557,15 @@ def test_find_nearest_contour_no_filled():
     cs = plt.contourf(img, 10)
 
     with pytest.warns(mpl._api.MatplotlibDeprecationWarning), \
-         pytest.raises(ValueError, match="Method does not support filled contours."):
+         pytest.raises(ValueError, match="Method does not support filled contours"):
         cs.find_nearest_contour(1, 1, pixel=False)
 
     with pytest.warns(mpl._api.MatplotlibDeprecationWarning), \
-         pytest.raises(ValueError, match="Method does not support filled contours."):
+         pytest.raises(ValueError, match="Method does not support filled contours"):
         cs.find_nearest_contour(1, 10, indices=(5, 7), pixel=False)
 
     with pytest.warns(mpl._api.MatplotlibDeprecationWarning), \
-         pytest.raises(ValueError, match="Method does not support filled contours."):
+         pytest.raises(ValueError, match="Method does not support filled contours"):
         cs.find_nearest_contour(2, 5, indices=(2, 7), pixel=True)
 
 
