ENH: Add PiecewiseLinearNorm and tests

Borrows heavily from @Tillsen's solution found on
StackOverflow here: http://goo.gl/RPXMYB

Used with his permission dicussesd on Github here:
matplotlib#3858

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.

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)

lib/matplotlib/colors.py

@@ -62,6 +62,7 @@
 
 from collections.abc import Sized
 import itertools
+import operator
 import re
 
 import numpy as np
@@ -992,6 +993,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."""