Merge pull request #14313 from efiring/masked_to_rgba

dstansby · web-flow · commit 4a0d9c51a379 · 2019-05-25T11:46:19.000+02:00
Support masked array inputs for to_rgba and to_rgba_array.
diff --git a/lib/matplotlib/colors.py b/lib/matplotlib/colors.py
@@ -152,7 +152,7 @@ def to_rgba(c, alpha=None):
 
     Parameters
     ----------
-    c : Matplotlib color
+    c : Matplotlib color or ``np.ma.masked``
 
     alpha : scalar, optional
         If *alpha* is not ``None``, it forces the alpha value, except if *c* is
@@ -189,6 +189,8 @@ def _to_rgba_no_colorcycle(c, alpha=None):
     ``"none"`` (case-insensitive), which always maps to ``(0, 0, 0, 0)``.
     """
     orig_c = c
+    if c is np.ma.masked:
+        return (0., 0., 0., 0.)
     if isinstance(c, str):
         if c.lower() == "none":
             return (0., 0., 0., 0.)
@@ -260,19 +262,25 @@ def to_rgba_array(c, alpha=None):
 
     If *alpha* is not ``None``, it forces the alpha value.  If *c* is
     ``"none"`` (case-insensitive) or an empty list, an empty array is returned.
+    If *c* is a masked array, an ndarray is returned with a (0, 0, 0, 0)
+    row for each masked value or row in *c*.
     """
     # Special-case inputs that are already arrays, for performance.  (If the
     # array has the wrong kind or shape, raise the error during one-at-a-time
     # conversion.)
     if (isinstance(c, np.ndarray) and c.dtype.kind in "if"
             and c.ndim == 2 and c.shape[1] in [3, 4]):
+        mask = c.mask.any(axis=1) if np.ma.is_masked(c) else None
+        c = np.ma.getdata(c)
         if c.shape[1] == 3:
             result = np.column_stack([c, np.zeros(len(c))])
             result[:, -1] = alpha if alpha is not None else 1.
         elif c.shape[1] == 4:
             result = c.copy()
             if alpha is not None:
                 result[:, -1] = alpha
+        if mask is not None:
+            result[mask] = 0
         if np.any((result < 0) | (result > 1)):
             raise ValueError("RGBA values should be within 0-1 range")
         return result
diff --git a/lib/matplotlib/tests/test_colors.py b/lib/matplotlib/tests/test_colors.py
@@ -821,6 +821,16 @@ def test_conversions():
         hex_color
 
 
+def test_conversions_masked():
+    x1 = np.ma.array(['k', 'b'], mask=[True, False])
+    x2 = np.ma.array([[0, 0, 0, 1], [0, 0, 1, 1]])
+    x2[0] = np.ma.masked
+    assert mcolors.to_rgba(x1[0]) == (0, 0, 0, 0)
+    assert_array_equal(mcolors.to_rgba_array(x1),
+                       [[0, 0, 0, 0], [0, 0, 1, 1]])
+    assert_array_equal(mcolors.to_rgba_array(x2), mcolors.to_rgba_array(x1))
+
+
 def test_to_rgba_array_single_str():
     # single color name is valid
     assert_array_equal(mcolors.to_rgba_array("red"), [(1, 0, 0, 1)])
