added tests for masked array and clipping

fmaussion · fmaussion · commit 7fab12947646 · 2015-08-01T15:56:04.000+02:00
diff --git a/lib/matplotlib/colors.py b/lib/matplotlib/colors.py
@@ -1252,8 +1252,9 @@ def __init__(self, boundaries, ncolors, clip=False):
         as i varies from 0 to len(boundaries)-2,
         j goes from 0 to ncolors-1.
 
-        Out-of-range values are mapped to -1 if low and ncolors
-        if high; these are converted to valid indices by
+        If clip == False, Out-of-range values are mapped
+        to -1 if low and ncolors if high; these are converted
+        to valid indices by
         :meth:`Colormap.__call__` .
         """
         self.clip = clip
@@ -1270,21 +1271,24 @@ def __init__(self, boundaries, ncolors, clip=False):
     def __call__(self, x, clip=None):
         if clip is None:
             clip = self.clip
-        _x = ma.asarray(x)
-        mask = ma.getmaskarray(_x)
-        xx = np.atleast_1d(_x.filled(self.vmax + 1))
+        x = ma.masked_invalid(x, copy=False)
+        mask = ma.getmaskarray(x)
+        xx = np.atleast_1d(x.filled(self.vmax + 1))
         if clip:
             np.clip(xx, self.vmin, self.vmax, out=xx)
-        iret = np.atleast_1d(np.zeros(_x.shape, dtype=np.int16))
+            max_col = self.Ncmap - 1
+        else:
+            max_col = self.Ncmap
+        iret = np.atleast_1d(np.zeros(x.shape, dtype=np.int16))
         for i, b in enumerate(self.boundaries):
             iret[xx >= b] = i
         if self._interp:
             scalefac = float(self.Ncmap - 1) / (self.N - 2)
             iret = (iret * scalefac).astype(np.int16)
         iret[xx < self.vmin] = -1
-        iret[xx >= self.vmax] = self.Ncmap  # ISSUE here: clip is ignored
+        iret[xx >= self.vmax] = max_col
         ret = ma.array(iret, mask=mask)
-        if _x.shape == () and not mask:
+        if x.shape == ():
             ret = int(ret[0])  # assume python scalar
         return ret
 
diff --git a/lib/matplotlib/tests/test_colors.py b/lib/matplotlib/tests/test_colors.py
@@ -84,15 +84,37 @@ def test_BoundaryNorm():
         assert_array_equal(bn([v]), ex)
 
     # Clipping
-    # For the "3" I do not know: isn't clipping also supposed to clip the max?
     bn = mcolors.BoundaryNorm(boundaries, 3, clip=True)
-    expected = [0, 0, 2, 3]
+    expected = [0, 0, 2, 2]
     for v, ex in zip(vals, expected):
         ret = bn(v)
         assert_true(isinstance(ret, six.integer_types))
         assert_array_equal(ret, ex)
         assert_array_equal(bn([v]), ex)
 
+    # Masked arrays
+    boundaries = [0, 1.1, 2.2]
+    vals = [-1., np.NaN, 0, 1.4, 9]
+
+    # Without interpolation
+    ncolors = len(boundaries) - 1
+    bn = mcolors.BoundaryNorm(boundaries, ncolors)
+    expected = np.ma.masked_array([-1, -99, 0, 1, 2], mask=[0, 1, 0, 0, 0])
+    assert_array_equal(bn(vals), expected)
+
+    # With interpolation
+    bn = mcolors.BoundaryNorm(boundaries, len(boundaries))
+    expected = np.ma.masked_array([-1, -99, 0, 2, 3], mask=[0, 1, 0, 0, 0])
+    assert_array_equal(bn(vals), expected)
+
+    # Non-trivial masked arrays
+    vals = [np.Inf, np.NaN]
+    assert_true(np.all(bn(vals).mask))
+    vals = [np.Inf]
+    assert_true(np.all(bn(vals).mask))
+
+    # Invalid scalar raises an exception
+    assert_raises(Exception, bn, np.NaN)
 
 def test_LogNorm():
     """
