Fixed bug in find_nearest_contour; Added test

matplotlib · Apr 8, 2022 · 6eea6e6 · 6eea6e6
1 parent 6c3412b
commit 6eea6e6
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 1 deletion.
diff --git a/lib/matplotlib/contour.py b/lib/matplotlib/contour.py
@@ -1371,7 +1371,7 @@ def find_nearest_contour(self, x, y, indices=None, pixel=True):
         # Nonetheless, improvements could probably be made.
 
         if indices is None:
-            indices = range(len(self.levels))
+            indices = range(len(self.collections))
 
         d2min = np.inf
         conmin = None

diff --git a/lib/matplotlib/tests/test_contour.py b/lib/matplotlib/tests/test_contour.py
@@ -470,6 +470,30 @@ def test_contour_line_start_on_corner_edge():
     cbar.add_lines(lines)
 
 
+def test_find_nearest_contour():
+    xy = np.indices((15, 15))
+    img = np.exp(-np.pi * (np.sum((xy - 5)**2, 0)/5.**2))
+    cs = plt.contour(img, 10)
+
+    nearest_contour = cs.find_nearest_contour(1, 1, pixel=False)
+    expected_nearest = (1, 0, 33, 1.965966, 1.965966, 1.866183)
+    assert_array_almost_equal(nearest_contour, expected_nearest)
+
+    nearest_contour = cs.find_nearest_contour(8, 1, pixel=False)
+    expected_nearest = (1, 0, 5, 7.550173, 1.587542, 0.547550)
+    assert_array_almost_equal(nearest_contour, expected_nearest)
+
+    nearest_contour = cs.find_nearest_contour(2, 5, pixel=False)
+    expected_nearest = (3, 0, 21, 1.884384, 5.023335, 0.013911)
+    assert_array_almost_equal(nearest_contour, expected_nearest)
+
+    nearest_contour = cs.find_nearest_contour(2, 5,
+                                              indices=(5, 7),
+                                              pixel=False)
+    expected_nearest = (5, 0, 16, 2.628202, 5.0, 0.394638)
+    assert_array_almost_equal(nearest_contour, expected_nearest)
+
+
 @mpl.style.context("default")
 def test_contour_autolabel_beyond_powerlimits():
     ax = plt.figure().add_subplot()