bugs in colors.BoundaryNorm

fmaussion · fmaussion · commit 32e727b00a7d · 2015-07-30T12:24:37.000+02:00
diff --git a/lib/matplotlib/colors.py b/lib/matplotlib/colors.py
@@ -1270,22 +1270,22 @@ 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 = x.filled(self.vmax + 1)
+        _x = ma.asarray(x)
+        mask = ma.getmaskarray(_x)
+        xx = np.atleast_1d(_x.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)
+        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
+        iret[xx >= self.vmax] = self.Ncmap  # ISSUE here: clip is ignored
         ret = ma.array(iret, mask=mask)
-        if ret.shape == () and not mask:
-            ret = int(ret)  # assume python scalar
+        if _x.shape == () and not mask:
+            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,60 @@ 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
+    # 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]
+    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)
+
 
 def test_LogNorm():
     """
