ENH: Add OffsetNorm and tests

phobson · OceanWolf · commit 05c1b76af7a1 · 2015-09-14T00:42:59.000+02:00
Borrows heavily from @Tillsen's solution found on StackOverflow here: http://goo.gl/RPXMYB Used with his permission dicussesd on Github here: https://github.com/matplotlib/matplotlib/pull/3858`
diff --git a/lib/matplotlib/colors.py b/lib/matplotlib/colors.py
@@ -225,6 +225,7 @@ def rgb2hex(rgb):
     a = '#%02x%02x%02x' % tuple([int(np.round(val * 255)) for val in rgb[:3]])
     return a
 
+
 hexColorPattern = re.compile("\A#[a-fA-F0-9]{6}\Z")
 
 
@@ -963,6 +964,127 @@ def scaled(self):
         return (self.vmin is not None and self.vmax is not None)
 
 
+class OffsetNorm(Normalize):
+    """
+    A subclass of matplotlib.colors.Normalize.
+
+    Normalizes data into the ``[0.0, 1.0]`` interval.
+    """
+    def __init__(self, vmin=None, vcenter=None, vmax=None, clip=False):
+        """Normalize data with an offset midpoint
+
+        Useful when mapping data unequally centered around a conceptual
+        center, e.g., data that range from -2 to 4, with 0 as the midpoint.
+
+        Parameters
+        ----------
+        vmin : optional float
+            The data value that defines ``0.0`` in the normalized data.
+            Defaults to the min value of the dataset.
+
+        vcenter : optional float
+            The data value that defines ``0.5`` in the normalized data.
+            Defaults to halfway between *vmin* and *vmax*.
+
+        vmax : option float
+            The data value that defines ``1.0`` in the normalized data.
+            Defaults to the the max value of the dataset.
+
+        clip : optional bool (default is False)
+            If *clip* is True, values beyond *vmin* and *vmax* will be set
+            to ``0.0`` or ``1.0``, respectively. Otherwise, values outside
+            the ``[0.0, 1.0]`` will be returned.
+
+        Examples
+        --------
+        >>> import matplotlib.colors as mcolors
+        >>> offset = mcolors.OffsetNorm(vmin=-2., vcenter=0., vmax=4.)
+        >>> data = [-2., -1., 0., 1., 2., 3., 4.]
+        >>> offset(data)
+        array([0., 0.25, 0.5, 0.625, 0.75, 0.875, 1.0])
+
+        """
+
+        self.vmin = vmin
+        self.vcenter = vcenter
+        self.vmax = vmax
+        self.clip = clip
+
+    def __call__(self, value, clip=None):
+        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)
+
+            # ma division is very slow; we can take a shortcut
+            resdat = result.data
+
+            #First scale to -1 to 1 range, than to from 0 to 1.
+            resdat -= vcenter
+            resdat[resdat > 0] /= abs(vmax - vcenter)
+            resdat[resdat < 0] /= abs(vmin - vcenter)
+
+            resdat /= 2.
+            resdat += 0.5
+            result = np.ma.array(resdat, mask=result.mask, copy=False)
+
+        if is_scalar:
+            result = result[0]
+
+        return result
+
+    def inverse(self, value):
+        if not self.scaled():
+            raise ValueError("Not invertible until scaled")
+
+        vmin, vcenter, vmax = self.vmin, self.vcenter, self.vmax
+        vmin = float(self.vmin)
+        vcenter = float(self.vcenter)
+        vmax = float(self.vmax)
+
+        if cbook.iterable(value):
+            val = ma.asarray(value)
+            val = 2 * (val - 0.5)
+            val[val > 0] *= abs(vmax - vcenter)
+            val[val < 0] *= abs(vmin - vcenter)
+            val += vcenter
+            return val
+        else:
+            val = 2 * (val - 0.5)
+            if val < 0:
+                return val * abs(vmin - vcenter) + vcenter
+            else:
+                return val * abs(vmax - vcenter) + vcenter
+
+    def autoscale_None(self, A):
+        ' autoscale only None-valued vmin or vmax'
+        if self.vmin is None and np.size(A) > 0:
+            self.vmin = ma.min(A)
+
+        if self.vmax is None and np.size(A) > 0:
+            self.vmax = ma.max(A)
+
+        if self.vcenter is None:
+            self.vcenter = (self.vmax + self.vmin) * 0.5
+
+
 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
@@ -5,7 +5,8 @@
 import itertools
 from distutils.version import LooseVersion as V
 
-from nose.tools import assert_raises, assert_equal, assert_true
+from nose.tools import (assert_raises, assert_equal, assert_true, assert_false
+                        raises)
 
 import numpy as np
 from numpy.testing.utils import assert_array_equal, assert_array_almost_equal
@@ -163,6 +164,182 @@ 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)
+
+    def test_autoscale(self):
+        norm = self.normclass()
+        norm.autoscale([10, 20, 30, 40])
+        assert_equal(norm.vmin, 10.)
+        assert_equal(norm.vmax, 40.)
+
+    def test_autoscale_None_vmin(self):
+        norm = self.normclass(vmin=0, vmax=None)
+        norm.autoscale_None([1, 2, 3, 4, 5])
+        assert_equal(norm.vmin, 0.)
+        assert_equal(norm.vmax, 5.)
+
+    def test_autoscale_None_vmax(self):
+        norm = self.normclass(vmin=None, vmax=10)
+        norm.autoscale_None([1, 2, 3, 4, 5])
+        assert_equal(norm.vmin, 1.)
+        assert_equal(norm.vmax, 10.)
+
+    def test_scale(self):
+        norm = self.normclass()
+        assert_false(norm.scaled())
+
+        norm([1, 2, 3, 4])
+        assert_true(norm.scaled())
+
+    def test_process_value_scalar(self):
+        res, is_scalar = mcolors.Normalize.process_value(5)
+        assert_true(is_scalar)
+        assert_array_equal(res, np.array([5.]))
+
+    def test_process_value_list(self):
+        res, is_scalar = mcolors.Normalize.process_value([5, 10])
+        assert_false(is_scalar)
+        assert_array_equal(res, np.array([5., 10.]))
+
+    def test_process_value_tuple(self):
+        res, is_scalar = mcolors.Normalize.process_value((5, 10))
+        assert_false(is_scalar)
+        assert_array_equal(res, np.array([5., 10.]))
+
+    def test_process_value_array(self):
+        res, is_scalar = mcolors.Normalize.process_value(np.array([5, 10]))
+        assert_false(is_scalar)
+        assert_array_equal(res, np.array([5., 10.]))
+
+
+class test_OffsetNorm_Even(_base_NormMixin):
+    def setup(self):
+        self.normclass = mcolors.OffsetNorm
+        self.norm = self.normclass(vmin=-1, vcenter=0, vmax=4)
+        self.vals = np.array([-1.0, -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_Odd(_base_NormMixin):
+    def setup(self):
+        self.normclass = mcolors.OffsetNorm
+        self.norm = self.normclass(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])
+
+
+class test_OffsetNorm_AllNegative(_base_NormMixin):
+    def setup(self):
+        self.normclass = mcolors.OffsetNorm
+        self.norm = self.normclass(vmin=-10, vcenter=-8, vmax=-2)
+        self.vals = np.array([-10., -9., -8., -6., -4., -2.])
+        self.expected = np.array([0.0, 0.25, 0.5, 0.666667, 0.833333, 1.0])
+
+
+class test_OffsetNorm_AllPositive(_base_NormMixin):
+    def setup(self):
+        self.normclass = mcolors.OffsetNorm
+        self.norm = self.normclass(vmin=0, vcenter=3, vmax=9)
+        self.vals = np.array([0., 1.5, 3., 4.5, 6.0, 7.5, 9.])
+        self.expected = np.array([0.0, 0.25, 0.5, 0.625, 0.75, 0.875, 1.0])
+
+
+class test_OffsetNorm_NoVs(_base_NormMixin):
+    def setup(self):
+        self.normclass = mcolors.OffsetNorm
+        self.norm = self.normclass(vmin=None, vcenter=None, vmax=None)
+        self.vals = np.array([-2.0, -1.0, 0.0, 1.0, 2.0, 3.0, 4.0])
+        self.expected = np.array([0., 0.16666667, 0.33333333,
+                                  0.5, 0.66666667, 0.83333333, 1.0])
+        self.expected_vmin = -2
+        self.expected_vcenter = 1
+        self.expected_vmax = 4
+
+    def test_vmin(self):
+        assert_true(self.norm.vmin is None)
+        self.norm(self.vals)
+        assert_equal(self.norm.vmin, self.expected_vmin)
+
+    def test_vcenter(self):
+        assert_true(self.norm.vcenter is None)
+        self.norm(self.vals)
+        assert_equal(self.norm.vcenter, self.expected_vcenter)
+
+    def test_vmax(self):
+        assert_true(self.norm.vmax is None)
+        self.norm(self.vals)
+        assert_equal(self.norm.vmax, self.expected_vmax)
+
+
+class test_OffsetNorm_VminEqualsVcenter(_base_NormMixin):
+    def setup(self):
+        self.normclass = mcolors.OffsetNorm
+        self.norm = self.normclass(vmin=-2, vcenter=-2, vmax=2)
+        self.vals = np.array([-2.0, -1.0, 0.0, 1.0, 2.0])
+        self.expected = np.array([0.5, 0.625, 0.75, 0.875, 1.0])
+
+
+class test_OffsetNorm_VmaxEqualsVcenter(_base_NormMixin):
+    def setup(self):
+        self.normclass = mcolors.OffsetNorm
+        self.norm = self.normclass(vmin=-2, vcenter=2, vmax=2)
+        self.vals = np.array([-2.0, -1.0, 0.0, 1.0, 2.0])
+        self.expected = np.array([0.0, 0.125, 0.25, 0.375, 0.5])
+
+
+class test_OffsetNorm_VsAllEqual(_base_NormMixin):
+    def setup(self):
+        self.v = 10
+        self.normclass = mcolors.OffsetNorm
+        self.norm = self.normclass(vmin=self.v, vcenter=self.v, vmax=self.v)
+        self.vals = np.array([-2.0, -1.0, 0.0, 1.0, 2.0])
+        self.expected = np.array([0.0, 0.0, 0.0, 0.0, 0.0])
+        self.expected_inv = self.expected + self.v
+
+    def test_inverse(self):
+        assert_array_almost_equal(
+            self.norm.inverse(self.norm(self.vals)),
+            self.expected_inv
+        )
+
+
+class test_OffsetNorm_Errors(object):
+    def setup(self):
+        self.vals = np.arange(50)
+
+    @raises(ValueError)
+    def test_VminGTVcenter(self):
+        norm = mcolors.OffsetNorm(vmin=10, vcenter=0, vmax=20)
+        norm(self.vals)
+
+    @raises(ValueError)
+    def test_VminGTVmax(self):
+        norm = mcolors.OffsetNorm(vmin=10, vcenter=0, vmax=5)
+        norm(self.vals)
+
+    @raises(ValueError)
+    def test_VcenterGTVmax(self):
+        norm = mcolors.OffsetNorm(vmin=10, vcenter=25, vmax=20)
+        norm(self.vals)
+
+    @raises(ValueError)
+    def test_premature_scaling(self):
+        norm = mcolors.OffsetNorm()
+        norm.inverse(np.array([0.1, 0.5, 0.9]))
+
+
 def test_SymLogNorm():
     """
     Test SymLogNorm behavior
@@ -281,7 +458,12 @@ def test_cmap_and_norm_from_levels_and_colors2():
                                'Wih extend={0!r} and data '
                                'value={1!r}'.format(extend, d_val))
 
-    assert_raises(ValueError, mcolors.from_levels_and_colors, levels, colors)
+    assert_raises(
+        ValueError,
+        mcolors.from_levels_and_colors,
+        levels,
+        colors
+    )
 
 
 def test_rgb_hsv_round_trip():
