Merge branch 'colorfix-#4354' of https://github.com/QiCuiHub/matplotlib into scatter_color

efiring · efiring · commit 0fb837ab24c8 · 2018-10-06T09:26:51.000-10:00
diff --git a/lib/matplotlib/axes/_axes.py b/lib/matplotlib/axes/_axes.py
@@ -4010,7 +4010,7 @@ def dopatch(xs, ys, **kwargs):
                       label_namer="y")
     def scatter(self, x, y, s=None, c=None, marker=None, cmap=None, norm=None,
                 vmin=None, vmax=None, alpha=None, linewidths=None,
-                verts=None, edgecolors=None,
+                verts=None, edgecolors=None, plotinvalid=False,
                 **kwargs):
         """
         A scatter plot of *y* vs *x* with varying marker size and/or color.
@@ -4087,6 +4087,10 @@ def scatter(self, x, y, s=None, c=None, marker=None, cmap=None, norm=None,
             For non-filled markers, the *edgecolors* kwarg is ignored and
             forced to 'face' internally.
 
+        plotinvalid : boolean, optional, default: False
+            Set to plot valid points with invalid color, in conjunction with
+            `~matplotlib.colors.Colormap.set_bad`.
+
         Returns
         -------
         paths : `~matplotlib.collections.PathCollection`
@@ -4236,11 +4240,12 @@ def scatter(self, x, y, s=None, c=None, marker=None, cmap=None, norm=None,
         else:
             colors = None  # use cmap, norm after collection is created
 
-        # `delete_masked_points` only modifies arguments of the same length as
-        # `x`.
-        x, y, s, c, colors, edgecolors, linewidths =\
-            cbook.delete_masked_points(
-                x, y, s, c, colors, edgecolors, linewidths)
+        if plotinvalid is False:
+            # `delete_masked_points` only modifies arguments of the same length
+            #  as `x`.
+            x, y, s, c, colors, edgecolors, linewidths =\
+               cbook.delete_masked_points(
+                   x, y, s, c, colors, edgecolors, linewidths)
 
         scales = s   # Renamed for readability below.
 
@@ -4284,7 +4289,12 @@ def scatter(self, x, y, s=None, c=None, marker=None, cmap=None, norm=None,
             if norm is not None and not isinstance(norm, mcolors.Normalize):
                 raise ValueError(
                     "'norm' must be an instance of 'mcolors.Normalize'")
-            collection.set_array(np.asarray(c))
+
+            if plotinvalid is False:
+                collection.set_array(c)
+            else:
+                collection.set_array(ma.masked_invalid(c))
+
             collection.set_cmap(cmap)
             collection.set_norm(norm)
 
diff --git a/lib/matplotlib/tests/baseline_images/test_axes/scatter_invalid_color.png b/lib/matplotlib/tests/baseline_images/test_axes/scatter_invalid_color.png
diff --git a/lib/matplotlib/tests/test_axes.py b/lib/matplotlib/tests/test_axes.py
@@ -1685,6 +1685,18 @@ def test_hist2d_transpose():
     ax.hist2d(x, y, bins=10, rasterized=True)
 
 
+@image_comparison(baseline_images=['scatter_invalid_color'], remove_text=True,
+                  extensions=['png'])
+def test_scatter_invalid_color():
+    fig, ax = plt.subplots()
+    data = {"x": [0, 1, 2, 3], "c": [1, np.nan, 9, np.nan],
+            "s": 1e4}
+    cmap = plt.get_cmap('jet', 16)
+    cmap.set_bad("k", 1)
+    ax.scatter(data["x"], data["x"], c=data["c"], s=data["s"],
+               cmap=cmap, plotinvalid=True)
+
+
 class TestScatter(object):
     @image_comparison(baseline_images=['scatter', 'scatter'])
     def test_scatter_plot(self):
