WIP: Add offset norm

phobson · phobson · commit 37ff5a5d67e5 · 2014-11-27T13:31:57.000-08:00
diff --git a/lib/matplotlib/colors.py b/lib/matplotlib/colors.py
@@ -969,6 +969,80 @@ def scaled(self):
         return (self.vmin is not None and self.vmax is not None)
 
 
+class OffsetNorm(Normalize):
+    def __init__(self, vmin=None, vcenter=None, vmax=None, clip=False):
+        self.vmin = vmin
+        self.vcenter = vcenter
+        self.vmax = vmax
+        self.clip = clip
+
+    def __call__(self, value, clip=False):
+        if clip is None:
+            clip = self.clip
+
+        result, is_scalar = self.process_value(value)
+
+        self.autoscale_None(result)
+        vmin, vcenter, vmax = self.vmin, self.vcenter, self.vmax
+        if vmin == vmax == vcenter:
+            result.fill(0)
+        elif not vmin <= vcenter <= vmax:
+            raise ValueError("minvalue must be less than or equal to "
+                             "centervalue which must be less than or "
+                             "equal to maxvalue")
+        else:
+            vmin = float(vmin)
+            vcenter = float(vcenter)
+            vmax = float(vmax)
+            if clip:
+                mask = ma.getmask(result)
+                result = ma.array(np.clip(result.filled(vmax), vmin, vmax),
+                                  mask=mask)
+
+            resdat = result.data
+            lower = resdat[resdat < vcenter]
+            upper = resdat[resdat >= vcenter]
+
+            lower -= vmin
+            lower /= ((vcenter-vmin) * 2)
+
+            upper -= vcenter
+            upper /= ((vmax-vcenter) * 2)
+            upper += 0.5
+
+            resdat = np.hstack([lower, upper])
+
+            result = np.ma.array(resdat, mask=result.mask, copy=False)
+
+        if is_scalar:
+            result = result[0]
+
+        return result
+
+    def inverse(self, value):
+        midpoint = 0.5
+        if not self.scaled():
+            raise ValueError("Not invertible until scaled")
+
+        vmin = float(self.vmin)
+        vcenter = float(self.vcenter)
+        vmax = float(self.vmax)
+
+        if cbook.iterable(value):
+            val = ma.asarray(value)
+        else:
+            val = ma.asarray([value])
+
+        lower = vmin + val[val <= midpoint] * (2 * (vcenter - vmin))
+        upper = vcenter + val[val > midpoint] * (2 * (vmax - vcenter)) - vmax
+
+        inverted = np.hstack([lower, upper])
+        if cbook.iterable(value):
+            return inverted
+        else:
+            return inverted[0]
+
+
 class LogNorm(Normalize):
     """
     Normalize a given value to the 0-1 range on a log scale
diff --git a/lib/matplotlib/tests/test_colors.py b/lib/matplotlib/tests/test_colors.py
@@ -80,6 +80,34 @@ def test_Normalize():
     _mask_tester(norm, vals)
 
 
+class _base_NormMixin(object):
+    def test_call(self):
+        normed_vals = self.norm(self.vals)
+        assert_array_almost_equal(normed_vals, self.expected)
+
+    def test_inverse(self):
+        _inverse_tester(self.norm, self.vals)
+
+    def test_scalar(self):
+        _scalar_tester(self.norm, self.vals)
+
+    def test_mask(self):
+        _mask_tester(self.norm, self.vals)
+
+
+class test_OffsetNorm_Even(_base_NormMixin):
+    def setup(self):
+        self.norm = mcolors.OffsetNorm(vmin=-1, vcenter=0, vmax=4)
+        self.vals = np.array([-1, -0.5, 0.0, 1.0, 2.0, 3.0, 4.0])
+        self.expected = np.array([0.0, 0.25, 0.5, 0.625, 0.75, 0.875, 1.0])
+
+class test_OffsetNorm_Even(_base_NormMixin):
+    def setup(self):
+        self.norm = mcolors.OffsetNorm(vmin=-2, vcenter=0, vmax=5)
+        self.vals = np.array([-2.0, -1.0, 0.0, 1.0, 2.0, 3.0, 4.0, 5.0])
+        self.expected = np.array([0.0, 0.25, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0])
+
+
 def test_SymLogNorm():
     """
     Test SymLogNorm behavior
