Merge pull request #4824 from fmaussion/maintenance-colors

efiring · efiring · commit 5d9245d23ded · 2015-08-08T02:24:48.000-11:00
Two bugs in colors.BoundaryNorm
diff --git a/lib/matplotlib/colors.py b/lib/matplotlib/colors.py
@@ -1252,9 +1252,12 @@ 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
+        Out-of-range values are mapped
+        to -1 if low and ncolors if high; these are converted
+        to valid indices by
         :meth:`Colormap.__call__` .
+        If clip == True, out-of-range values
+        are mapped to 0 if low and ncolors-1 if high.
         """
         self.clip = clip
         self.vmin = boundaries[0]
@@ -1267,25 +1270,29 @@ def __init__(self, boundaries, ncolors, clip=False):
         else:
             self._interp = True
 
-    def __call__(self, x, clip=None):
+    def __call__(self, value, clip=None):
         if clip is None:
             clip = self.clip
-        x = ma.asarray(x)
-        mask = ma.getmaskarray(x)
-        xx = x.filled(self.vmax + 1)
+
+        xx, is_scalar = self.process_value(value)
+        mask = ma.getmaskarray(xx)
+        xx = np.atleast_1d(xx.filled(self.vmax + 1))
         if clip:
-            np.clip(xx, self.vmin, self.vmax)
-        iret = np.zeros(x.shape, dtype=np.int16)
+            np.clip(xx, self.vmin, self.vmax, out=xx)
+            max_col = self.Ncmap - 1
+        else:
+            max_col = self.Ncmap
+        iret = np.zeros(xx.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
+        iret[xx >= self.vmax] = max_col
         ret = ma.array(iret, mask=mask)
-        if ret.shape == () and not mask:
-            ret = int(ret)  # assume python scalar
+        if is_scalar:
+            ret = int(ret[0])  # assume python scalar
         return ret
 
     def inverse(self, value):
diff --git a/lib/matplotlib/tests/test_colors.py b/lib/matplotlib/tests/test_colors.py
@@ -5,7 +5,7 @@
 import itertools
 from distutils.version import LooseVersion as V
 
-from nose.tools import assert_raises, assert_equal
+from nose.tools import assert_raises, assert_equal, assert_true
 
 import numpy as np
 from numpy.testing.utils import assert_array_equal, assert_array_almost_equal
@@ -39,15 +39,80 @@ def test_BoundaryNorm():
     Github issue #1258: interpolation was failing with numpy
     1.7 pre-release.
     """
-    # TODO: expand this into a more general test of BoundaryNorm.
+
     boundaries = [0, 1.1, 2.2]
-    vals = [-1, 0, 2, 2.2, 4]
-    expected = [-1, 0, 2, 3, 3]
+    vals = [-1, 0, 1, 2, 2.2, 4]
+
+    # Without interpolation
+    expected = [-1, 0, 0, 1, 2, 2]
+    ncolors = len(boundaries) - 1
+    bn = mcolors.BoundaryNorm(boundaries, ncolors)
+    assert_array_equal(bn(vals), expected)
+
     # ncolors != len(boundaries) - 1 triggers interpolation
+    expected = [-1, 0, 0, 2, 3, 3]
     ncolors = len(boundaries)
     bn = mcolors.BoundaryNorm(boundaries, ncolors)
     assert_array_equal(bn(vals), expected)
 
+    # more boundaries for a third color
+    boundaries = [0, 1, 2, 3]
+    vals = [-1, 0.1, 1.1, 2.2, 4]
+    ncolors = 5
+    expected = [-1, 0, 2, 4, 5]
+    bn = mcolors.BoundaryNorm(boundaries, ncolors)
+    assert_array_equal(bn(vals), expected)
+
+    # a scalar as input should not trigger an error and should return a scalar
+    boundaries = [0, 1, 2]
+    vals = [-1, 0.1, 1.1, 2.2]
+    bn = mcolors.BoundaryNorm(boundaries, 2)
+    expected = [-1, 0, 1, 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)
+
+    # same with interp
+    bn = mcolors.BoundaryNorm(boundaries, 3)
+    expected = [-1, 0, 2, 3]
+    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)
+
+    # Clipping
+    bn = mcolors.BoundaryNorm(boundaries, 3, clip=True)
+    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 = np.ma.masked_invalid([-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.ma.masked_invalid([np.Inf, np.NaN])
+    assert_true(np.all(bn(vals).mask))
+    vals = np.ma.masked_invalid([np.Inf])
+    assert_true(np.all(bn(vals).mask))
+
 
 def test_LogNorm():
     """
