Merge pull request #14257 from greglucas/colormap_nan_handling

tacaswell · web-flow · commit 89105f188c39 · 2019-05-20T19:07:47.000-04:00
FIX: Changing cmap(np.nan) to 'bad' value rather than 'under' value
diff --git a/lib/matplotlib/colors.py b/lib/matplotlib/colors.py
@@ -517,15 +517,17 @@ def __call__(self, X, alpha=None, bytes=False):
         if not self._isinit:
             self._init()
         mask_bad = None
-        if not np.iterable(X):
-            vtype = 'scalar'
-            xa = np.array([X])
-        else:
-            vtype = 'array'
-            xma = np.ma.array(X, copy=True)  # Copy here to avoid side effects.
-            mask_bad = xma.mask              # Mask will be used below.
-            xa = xma.filled()                # Fill to avoid infs, etc.
-            del xma
+        if np.ma.is_masked(X):
+            mask_bad = X.mask
+        elif np.any(np.isnan(X)):
+            # mask nan's
+            mask_bad = np.isnan(X)
+
+        xa = np.array(X, copy=True)
+        # Fill bad values to avoid warnings
+        # in the boolean comparisons below.
+        if mask_bad is not None:
+            xa[mask_bad] = 0.
 
         # Calculations with native byteorder are faster, and avoid a
         # bug that otherwise can occur with putmask when the last
@@ -548,10 +550,8 @@ def __call__(self, X, alpha=None, bytes=False):
         xa[xa > self.N - 1] = self._i_over
         xa[xa < 0] = self._i_under
         if mask_bad is not None:
-            if mask_bad.shape == xa.shape:
-                np.copyto(xa, self._i_bad, where=mask_bad)
-            elif mask_bad:
-                xa.fill(self._i_bad)
+            xa[mask_bad] = self._i_bad
+
         if bytes:
             lut = (self._lut * 255).astype(np.uint8)
         else:
@@ -572,8 +572,9 @@ def __call__(self, X, alpha=None, bytes=False):
                 # override its alpha just as for any other value.
 
         rgba = lut.take(xa, axis=0, mode='clip')
-        if vtype == 'scalar':
-            rgba = tuple(rgba[0, :])
+        if not np.iterable(X):
+            # Return a tuple if the input was a scalar
+            rgba = tuple(rgba)
         return rgba
 
     def __copy__(self):
diff --git a/lib/matplotlib/tests/test_colors.py b/lib/matplotlib/tests/test_colors.py
@@ -75,6 +75,60 @@ def test_colormap_endian():
         assert_array_equal(cmap(anative), cmap(aforeign))
 
 
+def test_colormap_invalid():
+    """
+    Github issue #9892: Handling of nan's were getting mapped to under
+    rather than bad. This tests to make sure all invalid values
+    (-inf, nan, inf) are mapped respectively to (under, bad, over).
+    """
+    cmap = cm.get_cmap("plasma")
+    x = np.array([-np.inf, -1, 0, np.nan, .7, 2, np.inf])
+
+    expected = np.array([[0.050383, 0.029803, 0.527975, 1.],
+                         [0.050383, 0.029803, 0.527975, 1.],
+                         [0.050383, 0.029803, 0.527975, 1.],
+                         [0.,       0.,       0.,       0.],
+                         [0.949217, 0.517763, 0.295662, 1.],
+                         [0.940015, 0.975158, 0.131326, 1.],
+                         [0.940015, 0.975158, 0.131326, 1.]])
+    assert_array_equal(cmap(x), expected)
+
+    # Test masked representation (-inf, inf) are now masked
+    expected = np.array([[0.,       0.,       0.,       0.],
+                         [0.050383, 0.029803, 0.527975, 1.],
+                         [0.050383, 0.029803, 0.527975, 1.],
+                         [0.,       0.,       0.,       0.],
+                         [0.949217, 0.517763, 0.295662, 1.],
+                         [0.940015, 0.975158, 0.131326, 1.],
+                         [0.,       0.,       0.,       0.]])
+    assert_array_equal(cmap(np.ma.masked_invalid(x)), expected)
+
+    # Test scalar representations
+    assert_array_equal(cmap(-np.inf), cmap(0))
+    assert_array_equal(cmap(np.inf), cmap(1.0))
+    assert_array_equal(cmap(np.nan), np.array([0., 0., 0., 0.]))
+
+
+def test_colormap_return_types():
+    """
+    Make sure that tuples are returned for scalar input and
+    that the proper shapes are returned for ndarrays.
+    """
+    cmap = cm.get_cmap("plasma")
+    # Test return types and shapes
+    # scalar input needs to return a tuple of length 4
+    assert isinstance(cmap(0.5), tuple)
+    assert len(cmap(0.5)) == 4
+
+    # input array returns an ndarray of shape x.shape + (4,)
+    x = np.ones(4)
+    assert cmap(x).shape == x.shape + (4,)
+
+    # multi-dimensional array input
+    x2d = np.zeros((2, 2))
+    assert cmap(x2d).shape == x2d.shape + (4,)
+
+
 def test_BoundaryNorm():
     """
     Github issue #1258: interpolation was failing with numpy
