Merge pull request #4055 from tacaswell/power_norm_fix

tacaswell · tacaswell · commit c14240769aa5 · 2015-02-01T17:24:10.000-05:00
BUG : fix PowerNorm with scalars Merging own PR due to prior review by @jenshnielsen
diff --git a/lib/matplotlib/colors.py b/lib/matplotlib/colors.py
@@ -1168,16 +1168,18 @@ def __call__(self, value, clip=None):
         elif vmin == vmax:
             result.fill(0)
         else:
+            res_mask = result.data < 0
             if clip:
                 mask = ma.getmask(result)
-                val = ma.array(np.clip(result.filled(vmax), vmin, vmax),
-                                mask=mask)
+                result = ma.array(np.clip(result.filled(vmax), vmin, vmax),
+                                  mask=mask)
             resdat = result.data
             resdat -= vmin
             np.power(resdat, gamma, resdat)
             resdat /= (vmax - vmin) ** gamma
+
             result = np.ma.array(resdat, mask=result.mask, copy=False)
-            result[value < 0] = 0
+            result[res_mask] = 0
         if is_scalar:
             result = result[0]
         return result
diff --git a/lib/matplotlib/tests/test_colors.py b/lib/matplotlib/tests/test_colors.py
@@ -3,7 +3,7 @@
 
 import six
 
-from nose.tools import assert_raises
+from nose.tools import assert_raises, assert_equal
 
 import numpy as np
 from numpy.testing.utils import assert_array_equal, assert_array_almost_equal
@@ -62,11 +62,29 @@ def test_PowerNorm():
     assert_array_almost_equal(norm(a), pnorm(a))
 
     a = np.array([-0.5, 0, 2, 4, 8], dtype=np.float)
-    expected = [0, 0, 1./16, 1./4, 1]
+    expected = [0, 0, 1/16, 1/4, 1]
     pnorm = mcolors.PowerNorm(2, vmin=0, vmax=8)
     assert_array_almost_equal(pnorm(a), expected)
+    assert_equal(pnorm(a[0]), expected[0])
+    assert_equal(pnorm(a[2]), expected[2])
     assert_array_almost_equal(a[1:], pnorm.inverse(pnorm(a))[1:])
 
+    # Clip = True
+    a = np.array([-0.5, 0, 1, 8, 16], dtype=np.float)
+    expected = [0, 0, 0, 1, 1]
+    pnorm = mcolors.PowerNorm(2, vmin=2, vmax=8, clip=True)
+    assert_array_almost_equal(pnorm(a), expected)
+    assert_equal(pnorm(a[0]), expected[0])
+    assert_equal(pnorm(a[-1]), expected[-1])
+
+    # Clip = True at call time
+    a = np.array([-0.5, 0, 1, 8, 16], dtype=np.float)
+    expected = [0, 0, 0, 1, 1]
+    pnorm = mcolors.PowerNorm(2, vmin=2, vmax=8, clip=False)
+    assert_array_almost_equal(pnorm(a, clip=True), expected)
+    assert_equal(pnorm(a[0], clip=True), expected[0])
+    assert_equal(pnorm(a[-1], clip=True), expected[-1])
+
 
 def test_Normalize():
     norm = mcolors.Normalize()
