ENH: Add PiecewiseLinearNorm and tests

phobson · jklymak · commit 512f2d45f470 · 2018-10-06T20:51:27.000-07: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/doc/users/whats_new.rst b/doc/users/whats_new.rst
@@ -10,7 +10,6 @@ revision, see the :ref:`github-stats`.
 .. contents:: Table of Contents
    :depth: 4
 
-
 ..
    For a release, add a new section after this, then comment out the include
    and toctree below by indenting them. Uncomment them after the release.
diff --git a/lib/matplotlib/colorbar.py b/lib/matplotlib/colorbar.py
@@ -772,8 +772,8 @@ def _process_values(self, b=None):
                                       + self._boundaries[1:])
                 if isinstance(self.norm, colors.NoNorm):
                     self._values = (self._values + 0.00001).astype(np.int16)
-                return
-            self._values = np.array(self.values)
+            else:
+                self._values = np.array(self.values)
             return
         if self.values is not None:
             self._values = np.array(self.values)
diff --git a/lib/matplotlib/colors.py b/lib/matplotlib/colors.py
@@ -992,6 +992,164 @@ def scaled(self):
         return self.vmin is not None and self.vmax is not None
 
 
+class PiecewiseLinearNorm(Normalize):
+    """
+    Normalizes data into the ``[0.0, 1.0]`` interval over linear segments.
+    """
+    def __init__(self, data_points=None, norm_points=None):
+        """Normalize data linearly between the defined stop points.
+
+        Parameters
+        ----------
+        data_points : tuple of floats
+            Floats spanning the data to be mapped between 0-1
+
+        norm_points : tuple of floats or None (default)
+            Floats spanning [0, 1] that the data points will map to.  If
+            *None*, the data points will be mapped equally into [0, 1].
+
+        Examples
+        --------
+        Note this example is equivalent to the `.DivergingNorm` example.
+        >>> import matplotlib.colors as mcolors
+        >>> offset = mcolors.PiecewiseLinearNorm([-2, 0, 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._data_points = (np.asarray(data_points)
+                                if data_points is not None
+                                else None)
+        self._norm_points = (np.asarray(norm_points)
+                                if norm_points is not None
+                                else None)
+        # self._norm_points = np.asarray(norm_points)
+        self.vmin = data_points[0]
+        self.vmax = data_points[-1]
+
+    def __call__(self, value, clip=None):
+        """
+        Map value to the interval [0, 1]. The clip argument is unused.
+        """
+
+        result, is_scalar = self.process_value(value)
+        self.autoscale_None(result)
+        vmin, vmax = self._data_points[0], self._data_points[-1]
+
+        # its possible some of the data points are less than vmin
+        # or vmax.
+        ind = np.where((self._data_points >= vmin) &
+                       (self._data_points <= vmax))[0]
+        result = np.ma.masked_array(np.interp(result,
+                                              self._data_points[ind],
+                                              self._norm_points[ind]),
+                                    mask=np.ma.getmask(result))
+        if is_scalar:
+            result = np.atleast_1d(result)[0]
+        return result
+
+    def autoscale_None(self, A):
+        """
+        Parameters:
+        -----------
+
+        A : tuple
+            data used for autoscaling, if appropriate.
+
+        If ``norm.vmin` is None, or ``norm.vmax`` are none, use
+        the min and max of ``A`` for the endpoints of the linear mapping.
+        """
+
+        # allow autoscaling with the data if the user specifies
+        # vmin or vmax is None.  Note we never set vmax/vmin in the class
+        # except at init.
+        vmin = self.vmin
+        if vmin is None:
+            vmin = np.ma.min(A)
+        vmax = self.vmax
+        if vmax is None:
+            vmax = np.ma.max(A)
+
+        if vmin > vmax:
+            raise ValueError('vmin must be less than or equal to vmax')
+
+        if self._data_points is None:
+            self._data_points = np.asarray([vmin, vmax])
+        if self._norm_points is None:
+            N = len(self._data_points)
+            self._norm_points = np.linspace(0, 1, N)
+
+        self._data_points[0] = vmin
+        self._data_points[-1] = vmax
+
+
+class DivergingNorm(PiecewiseLinearNorm):
+    def __init__(self, vmin=None, vcenter=None, vmax=None):
+        """
+        Normalize data with a set center.
+
+        Useful when mapping data with an unequal rates of change around a
+        conceptual center, e.g., data that range from -2 to 4, with 0 as
+        the midpoint.
+
+        Parameters
+        ----------
+        vmin : float, optional
+            The data value that defines ``0.0`` in the normalization.
+            Defaults to the min value of the dataset.
+        vcenter : float, optional
+            The data value that defines ``0.5`` in the normalization.
+            If not defined, the normalization just linearly maps between
+            *vmin* and *vmax*.
+        vmax : float, optional
+            The data value that defines ``1.0`` in the normalization.
+            Defaults to the the max value of the dataset.
+
+        Examples
+        --------
+        This maps data value -4000 to 0., 0 to 0.5, and +10000 to 1.0; data
+        between is linearly interpolated::
+
+            >>> import matplotlib.colors as mcolors
+            >>> offset = mcolors.DivergingNorm(vmin=-4000.,
+                                               vcenter=0., vmax=10000)
+            >>> data = [-4000., -2000., 0., 2500., 5000., 7500., 10000.]
+            >>> offset(data)
+            array([0., 0.25, 0.5, 0.625, 0.75, 0.875, 1.0])
+        """
+
+        # vmin/vmax will set the first and last value, so they don't
+        # matter very much.
+        if vcenter is None:
+            stops = [0, 1]
+        else:
+            stops = [vcenter * 0.75, vcenter, vcenter * 1.5]
+        super(DivergingNorm, self).__init__(stops)
+        # if these have been specified we need to set them here so
+        # DivergingNorm knows they are user-set limits.  Otherwise it
+        # will autoscale to data.
+        self.vmin = vmin
+        self.vmax = vmax
+        if vcenter is not None and vmax is not None and vcenter > vmax:
+            raise ValueError('vmin, vcenter, and vmax must be in '
+                             'ascending order')
+        if vcenter is not None and vmin is not None and vcenter < vmin:
+            raise ValueError('vmin, vcenter, and vmax must be in '
+                             'ascending order')
+
+        def __call__(value, clip=None):
+            """
+            Map value to the interval [0, 1]. The clip argument is unused.
+            """
+            try:
+                super().__call(value, clip=clip)
+            except ValueError:
+                # improve the more general error message.
+                raise ValueError('vmin, vcenter, and vmax must '
+                                 'increment monotonically for DivergingNorm '
+                                 'to work.')
+
+
 class LogNorm(Normalize):
     """Normalize a given value to the 0-1 range on a log scale."""
 
diff --git a/lib/matplotlib/tests/baseline_images/test_colors/test_offset_norm.png b/lib/matplotlib/tests/baseline_images/test_colors/test_offset_norm.png
diff --git a/lib/matplotlib/tests/test_colors.py b/lib/matplotlib/tests/test_colors.py
@@ -1,6 +1,5 @@
 import copy
 import itertools
-
 import numpy as np
 import pytest
 
@@ -221,6 +220,208 @@ def test_Normalize():
     assert 0 < norm(1 + 50 * eps) < 1
 
 
+class BaseNormMixin(object):
+    def test_call(self):
+        normed_vals = self.norm(self.vals)
+        np.assert_array_almost_equal(normed_vals, self.expected)
+
+    def test_inverse(self):
+        if self.test_inverse:
+            _inverse_tester(self.norm, self.vals)
+        else:
+            pass
+
+    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])
+        np.assert_equal(norm.vmin, 10.)
+        np.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])
+        np.assert_equal(norm.vmin, 0.)
+        np.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])
+        np.assert_equal(norm.vmin, 1.)
+        np.assert_equal(norm.vmax, 10.)
+
+    def test_scale(self):
+        norm = self.normclass()
+        np.assert_false(norm.scaled())
+
+        norm([1, 2, 3, 4])
+        np.assert_true(norm.scaled())
+
+    def test_process_value_scalar(self):
+        res, is_scalar = mcolors.Normalize.process_value(5)
+        np.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])
+        np.assert_false(is_scalar)
+        np.assert_array_equal(res, np.array([5., 10.]))
+
+    def test_process_value_tuple(self):
+        res, is_scalar = mcolors.Normalize.process_value((5, 10))
+        np.assert_false(is_scalar)
+        np.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]))
+        np.assert_false(is_scalar)
+        np.assert_array_equal(res, np.array([5., 10.]))
+
+
+class BaseDivergingNorm(BaseNormMixin):
+    normclass = mcolors.DivergingNorm
+    test_inverse = False
+
+
+class test_DivergingNorm_Even(BaseDivergingNorm):
+    def setup(self):
+        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_DivergingNorm_Odd(BaseDivergingNorm):
+    def setup(self):
+        self.normclass = mcolors.DivergingNorm
+        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_DivergingNorm_AllNegative(BaseDivergingNorm):
+    def setup(self):
+        self.normclass = mcolors.DivergingNorm
+        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_DivergingNorm_AllPositive(BaseDivergingNorm):
+    def setup(self):
+        self.normclass = mcolors.DivergingNorm
+        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_DivergingNorm_NoVs(BaseDivergingNorm):
+    def setup(self):
+        self.normclass = mcolors.DivergingNorm
+        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):
+        np.assert_true(self.norm.vmin is None)
+        self.norm(self.vals)
+        np.assert_equal(self.norm.vmin, self.expected_vmin)
+
+    def test_vcenter(self):
+        np.assert_true(self.norm.vcenter is None)
+        self.norm(self.vals)
+        np.assert_equal(self.norm.vcenter, self.expected_vcenter)
+
+    def test_vmax(self):
+        np.assert_true(self.norm.vmax is None)
+        self.norm(self.vals)
+        np.assert_equal(self.norm.vmax, self.expected_vmax)
+
+
+class test_DivergingNorm_VminEqualsVcenter(BaseDivergingNorm):
+    def setup(self):
+        self.normclass = mcolors.DivergingNorm
+        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_DivergingNorm_VmaxEqualsVcenter(BaseDivergingNorm):
+    def setup(self):
+        self.normclass = mcolors.DivergingNorm
+        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_DivergingNorm_VsAllEqual(BaseDivergingNorm):
+    def setup(self):
+        self.v = 10
+        self.normclass = mcolors.DivergingNorm
+        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_DivergingNorm_Errors(object):
+    def setup(self):
+        self.vals = np.arange(50)
+
+    def test_VminGTVcenter(self):
+        with pytest.raises(ValueError):
+            norm = mcolors.DivergingNorm(vmin=10, vcenter=0, vmax=20)
+            norm(self.vals)
+
+    def test_VminGTVmax(self):
+        with pytest.raises(ValueError):
+            norm = mcolors.DivergingNorm(vmin=10, vcenter=0, vmax=5)
+            norm(self.vals)
+
+    def test_VcenterGTVmax(self):
+        with pytest.raises(ValueError):
+            norm = mcolors.DivergingNorm(vmin=10, vcenter=25, vmax=20)
+            norm(self.vals)
+
+    def test_premature_scaling(self):
+        with pytest.raises(ValueError):
+            norm = mcolors.DivergingNorm()
+            norm.inverse(np.array([0.1, 0.5, 0.9]))
+
+
+@image_comparison(baseline_images=['test_offset_norm'], extensions=['png'],
+                  style='mpl20')
+def test_offset_norm_img():
+    x = np.linspace(-2, 7)
+    y = np.linspace(-1*np.pi, np.pi)
+    X, Y = np.meshgrid(x, y)
+    Z = x * np.sin(Y)**2
+
+    fig, (ax1, ax2) = plt.subplots(ncols=2)
+    cmap = plt.cm.coolwarm
+    norm = mcolors.DivergingNorm(vmin=-2, vcenter=0, vmax=7)
+
+    img1 = ax1.pcolormesh(Z, cmap=cmap, norm=None)
+    cbar1 = fig.colorbar(img1, ax=ax1)
+
+    img2 = ax2.pcolormesh(Z, cmap=cmap, norm=norm)
+    cbar2 = fig.colorbar(img2, ax=ax2)
+
+
 def test_SymLogNorm():
     """
     Test SymLogNorm behavior
