Merge pull request #17546 from rlung/legend_elements

jkseppan · web-flow · commit ede8e566e5c7 · 2020-08-04T13:04:13.000+03:00
`func` argument in `legend_elements` with non-monotonically increasing functions
diff --git a/lib/matplotlib/collections.py b/lib/matplotlib/collections.py
@@ -1076,8 +1076,10 @@ def legend_elements(self, prop="colors", num="auto",
             cond = ((label_values >= func(arr).min()) &
                     (label_values <= func(arr).max()))
             label_values = label_values[cond]
-            xarr = np.linspace(arr.min(), arr.max(), 256)
-            values = np.interp(label_values, func(xarr), xarr)
+            yarr = np.linspace(arr.min(), arr.max(), 256)
+            xarr = func(yarr)
+            ix = np.argsort(xarr)
+            values = np.interp(label_values, xarr[ix], yarr[ix])
 
         kw = dict(markeredgewidth=self.get_linewidths()[0],
                   alpha=self.get_alpha())
diff --git a/lib/matplotlib/tests/test_collections.py b/lib/matplotlib/tests/test_collections.py
@@ -660,3 +660,30 @@ def test_quadmesh_set_array():
     coll.set_array(np.ones(9))
     fig.canvas.draw()
     assert np.array_equal(coll.get_array(), np.ones(9))
+
+
+def test_legend_inverse_size_label_relationship():
+    """
+    Ensure legend markers scale appropriately when label and size are
+    inversely related.
+    Here label = 5 / size
+    """
+
+    np.random.seed(19680801)
+    X = np.random.random(50)
+    Y = np.random.random(50)
+    C = 1 - np.random.random(50)
+    S = 5 / C
+
+    legend_sizes = [0.2, 0.4, 0.6, 0.8]
+    fig, ax = plt.subplots()
+    sc = ax.scatter(X, Y, s=S)
+    handles, labels = sc.legend_elements(
+      prop='sizes', num=legend_sizes, func=lambda s: 5 / s
+    )
+
+    # Convert markersize scale to 's' scale
+    handle_sizes = [x.get_markersize() for x in handles]
+    handle_sizes = [5 / x**2 for x in handle_sizes]
+
+    assert_array_almost_equal(handle_sizes, legend_sizes, decimal=1)
