ENH: Add PiecewiseLinearNorm and tests

phobson · jklymak · commit 7fe02bcebaaa · 2018-10-05T21:14:46.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/colors.py b/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):
+    """
+    A subclass of matplotlib.colors.Normalize.
+
+    Normalizes data into the ``[0.0, 1.0]`` interval.
+    """
+    # TODO rewrite the internals of this class once we support OrderedDicts
+    # i.e. after we drop support for python 2.6
+    def __init__(self, stops=None):
+        """Normalize data linearly between the defined stop points.
+        Use this as more generic form of ``DivergingNorm``
+
+        Parameters
+        ----------
+        stops : dict-like, optional
+            Accepts a dictionary or anything that can get converted to a
+            dictionary which maps the space [0.0, 1.0] to data point, i.e. key
+            value pairs.
+
+        Examples
+        --------
+        Note this example is equivalent to the DivergingNorm example.
+        >>> import matplotlib.colors as mcolors
+        >>> offset = mcolors.PiecewiseLinearNorm({0.: -2., 0.5: 0., 1.=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._set_stops(stops)
+
+    @property
+    def vmin(self):
+        try:
+            if self._stops[0][0] == 0:
+                return self._stops[0][1]
+        except IndexError:
+            return None
+
+    @vmin.setter
+    def vmin(self, vmin):
+        try:
+            if self._stops[0][0] == 0:
+                self._stops[0] = (self._stops[0][0], vmin)
+                return
+        except IndexError:
+            pass
+        self.append_stop(0., vmin)
+
+    @property
+    def vmax(self):
+        try:
+            if self._stops[-1][0] == 1:
+                return self._stops[-1][1]
+        except IndexError:
+            return None
+
+    @vmax.setter
+    def vmax(self, vmax):
+        try:
+            if self._stops[-1][0] == 1:
+                self._stops[-1] = (self._stops[-1][0], vmax)
+                return
+        except IndexError:
+            pass
+        self.append_stop(1., vmax)
+
+    # TODO Change this to a property when we drop 2.6 and use Ordered Dicts
+    def _set_stops(self, stops):
+        if not stops:
+            self._stops = []
+            return
+
+        stops = dict(stops)
+        self._stops = sorted(stops.items(), key=operator.itemgetter(0))
+        map_points, data_points = zip(*self._stops)
+        if not np.all(np.diff(data_points) > 0):
+            raise ValueError("stops must increase monotonically")
+
+    def append_stop(self, cmap_fraction, data_value):
+        i = -1
+        for i, (map_point, data_point) in enumerate(self._stops):
+            if map_point >= cmap_fraction:
+                d1 = data_point  # the current index
+                break
+        else:
+            i += 1  # the index to insert before
+            d1 = np.inf
+
+        if i > 0:
+            d0 = self._stops[i-1][1]
+        else:
+            d0 = -np.inf
+
+        if not (d0 < data_value < d1):
+            raise ValueError(('Stops must increase monotonically, due to the '
+                             + 'stops already set, the cmap_fraction specified'
+                             + ' (%f) means that the data_value must lie '
+                             + 'between %f and %f, but %f given') %
+                             (cmap_fraction, d0, d1, data_value))
+
+        self._stops.insert(i, (cmap_fraction, data_value))
+
+    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)
+
+        map_points, data_points = zip(*self._stops)
+        result = np.ma.masked_array(np.interp(result, data_points, map_points),
+                                 mask=np.ma.getmask(result))
+        if is_scalar:
+            result = np.atleast_1d(result)[0]
+        return result
+
+    def autoscale_None(self, A):
+        """Ensures we have the upper and lower bounds set, using the data A"""
+        if len(self._stops) == 0 or self._stops[0][0] != 0:
+            self.append_stop(0., ma.min(A))
+        if self._stops[-1][0] != 1:
+            self.append_stop(1., ma.max(A))
+
+
+class DivergingNorm(PiecewiseLinearNorm):
+    def __init__(self, vmin=None, vcenter=None, vmax=None):
+        """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 : float, optional
+            The data value that defines ``0.0`` in the normalized data.
+            Defaults to the min value of the dataset.
+        vcenter : float, optional
+            The data value that defines ``0.5`` in the normalized data.
+            Defaults to halfway between *vmin* and *vmax*.
+        vmax : float, optional
+            The data value that defines ``1.0`` in the normalized data.
+            Defaults to the the max value of the dataset.
+        Examples
+        --------
+        >>> import matplotlib.colors as mcolors
+        >>> offset = mcolors.PiecewiseLinearNorm(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])
+        """
+        stops = {}
+        if vmin is not None:
+            stops[0.] = vmin
+        if vcenter is not None:
+            stops[0.5] = vcenter
+        if vmax is not None:
+            stops[1.] = vmax
+        super(DivergingNorm, self).__init__(stops)
+
+
 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)
+        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])
+        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 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):
+        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_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
