use _find_nearest_contour

rcomer · anntzer · rcomer · commit 521c4d9305fa · 2023-10-14T19:33:05.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
@@ -1338,18 +1338,22 @@ def _process_linestyles(self, linestyles):
                 raise ValueError("Unrecognized type for linestyles kwarg")
         return tlinestyles
 
-    def _find_nearest_contour(self, xy, indices=None):
+    def _find_nearest_contour(self, xy, indices=None, pixel=True):
         """
-        Find the point in the unfilled contour plot that is closest (in screen
+        Find the point in the unfilled contour plot that is closest (in axes or screen
         space) to point *xy*.
 
         Parameters
         ----------
         xy : tuple[float, float]
-            The reference point (in screen space).
+            The reference point.
         indices : list of int or None, default: None
             Indices of contour levels to consider.  If None (the default), all levels
             are considered.
+        pixel : bool, default: True
+            If *True*, measure distance in pixel (screen) space, which is
+            useful for manual contour labeling; else, measure distance in axes
+            space.
 
         Returns
         -------
@@ -1378,7 +1382,9 @@ def _find_nearest_contour(self, xy, indices=None):
             path = self._paths[idx_level]
             if not len(path.vertices):
                 continue
-            lc = self.get_transform().transform(path.vertices)
+            lc = path.vertices
+            if pixel:
+                lc = self.get_transform().transform(lc)
             d2, proj, leg = _find_closest_point_on_path(lc, xy)
             if d2 < d2min:
                 d2min = d2
@@ -1422,53 +1428,19 @@ def find_nearest_contour(self, x, y, indices=None, pixel=True):
         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.
+        i_level, i_vtx, (xmin, ymin) = self._find_nearest_contour(
+            (x, y), indices, pixel)
 
-        if self.filled:
-            raise ValueError("Method does not support filled contours.")
-
-        paths = self.get_paths()
-        if indices is None:
-            indices = range(len(paths))
+        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])
-        trans = self.get_transform()
-
-        for icon in indices:
-            path = paths[icon]
-
-            for segNum, linepath in enumerate(path._iter_connected_components()):
-                lc = linepath.vertices
-                if lc.size == 0:
-                    continue
-                # 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/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)
 
 
