ENH: Add PiecewiseLinearNorm and tests

phobson · jklymak · commit 9fe878e1598b · 2019-03-07T16:20:05.000-08: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` TST: add tests for DivergingNorm DOC: add to colors_api.rst DOC: add tutorial FIX: fix extend=both
diff --git a/doc/api/colors_api.rst b/doc/api/colors_api.rst
@@ -25,6 +25,7 @@ Classes
 
    BoundaryNorm
    Colormap
+   DivergingNorm
    LightSource
    LinearSegmentedColormap
    ListedColormap
diff --git a/lib/matplotlib/colorbar.py b/lib/matplotlib/colorbar.py
@@ -875,8 +875,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)
@@ -1114,19 +1114,17 @@ def _locate(self, x):
             xn = self.norm(x, clip=False).filled()
 
         # The rest is linear interpolation with extrapolation at ends.
-        ii = np.searchsorted(b, xn)
-        i0 = ii - 1
-        itop = (ii == len(b))
-        ibot = (ii == 0)
-        i0[itop] -= 1
-        ii[itop] -= 1
-        i0[ibot] += 1
-        ii[ibot] += 1
-
-        y = self._y
-        db = b[ii] - b[i0]
-        dy = y[ii] - y[i0]
-        z = y[i0] + (xn - b[i0]) * dy / db
+
+        bunique = b
+        yunique = self._y
+        if b[0] == b[1]:
+            bunique = bunique[1:]
+            yunique = yunique[1:]
+        if b[-1] == b[-2]:
+            bunique = bunique[:-1]
+            yunique = yunique[:-1]
+
+        z = np.interp(xn, bunique, yunique)
         return z
 
     def set_alpha(self, alpha):
diff --git a/lib/matplotlib/colors.py b/lib/matplotlib/colors.py
@@ -958,6 +958,76 @@ def scaled(self):
         return self.vmin is not None and self.vmax is not None
 
 
+class DivergingNorm(Normalize):
+    def __init__(self, vcenter, vmin=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
+        ----------
+        vcenter : float
+            The data value that defines ``0.5`` in the normalization.
+        vmin : float, optional
+            The data value that defines ``0.0`` in the normalization.
+            Defaults to the min value of the dataset.
+        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])
+        """
+
+        self.vcenter = vcenter
+        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 autoscale_None(self, A):
+        """
+        Get vmin and vmax, and then clip at vcenter
+        """
+        super().autoscale_None(A)
+        if self.vmin > self.vcenter:
+            self.vmin = self.vcenter
+        if self.vmax < self.vcenter:
+            self.vmax = self.vcenter
+
+    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)  # sets self.vmin, self.vmax if None
+
+        if not self.vmin <= self.vcenter <= self.vmax:
+            raise ValueError("vmin, vcenter, vmax must increase monotonically")
+        result = np.ma.masked_array(
+            np.interp(result, [self.vmin, self.vcenter, self.vmax],
+                      [0, 0.5, 1.]), mask=np.ma.getmask(result))
+        if is_scalar:
+            result = np.atleast_1d(result)[0]
+        return result
+
+
 class LogNorm(Normalize):
     """Normalize a given value to the 0-1 range on a log scale."""
 
diff --git a/lib/matplotlib/mpl-data/sample_data/exampleTopo.npz b/lib/matplotlib/mpl-data/sample_data/exampleTopo.npz
diff --git a/lib/matplotlib/tests/test_colorbar.py b/lib/matplotlib/tests/test_colorbar.py
@@ -4,7 +4,8 @@
 from matplotlib import rc_context
 from matplotlib.testing.decorators import image_comparison
 import matplotlib.pyplot as plt
-from matplotlib.colors import BoundaryNorm, LogNorm, PowerNorm, Normalize
+from matplotlib.colors import (BoundaryNorm, LogNorm, PowerNorm, Normalize,
+                               DivergingNorm)
 from matplotlib.cm import get_cmap
 from matplotlib.colorbar import ColorbarBase, _ColorbarLogLocator
 from matplotlib.ticker import LogLocator, LogFormatter, FixedLocator
@@ -539,3 +540,21 @@ def test_colorbar_inverted_ticks():
     cbar.ax.invert_yaxis()
     np.testing.assert_allclose(ticks, cbar.get_ticks())
     np.testing.assert_allclose(minorticks, cbar.get_ticks(minor=True))
+
+
+def test_extend_colorbar_customnorm():
+    # This was a funny error with DivergingNorm, maybe with other norms,
+    # when extend='both'
+    N = 100
+    X, Y = np.mgrid[-3:3:complex(0, N), -2:2:complex(0, N)]
+    Z1 = np.exp(-X**2 - Y**2)
+    Z2 = np.exp(-(X - 1)**2 - (Y - 1)**2)
+    Z = (Z1 - Z2) * 2
+
+    fig, ax = plt.subplots(2, 1)
+    pcm = ax[0].pcolormesh(X, Y, Z,
+                           norm=DivergingNorm(vcenter=0., vmin=-2, vmax=1),
+                           cmap='RdBu_r')
+    cb = fig.colorbar(pcm, ax=ax[0], extend='both')
+    np.testing.assert_allclose(cb.ax.get_position().extents,
+                               [0.78375, 0.536364, 0.796147, 0.9], rtol=1e-3)
diff --git a/lib/matplotlib/tests/test_colors.py b/lib/matplotlib/tests/test_colors.py
@@ -221,6 +221,96 @@ def test_Normalize():
     assert 0 < norm(1 + 50 * eps) < 1
 
 
+def test_DivergingNorm_autoscale():
+    norm = mcolors.DivergingNorm(vcenter=20)
+    norm.autoscale([10, 20, 30, 40])
+    assert norm.vmin == 10.
+    assert norm.vmax == 40.
+
+
+def test_DivergingNorm_autoscale_None_vmin():
+    norm = mcolors.DivergingNorm(2, vmin=0, vmax=None)
+    norm.autoscale_None([1, 2, 3, 4, 5])
+    assert norm(5) == 1
+    assert norm.vmax == 5
+
+
+def test_DivergingNorm_autoscale_None_vmax():
+    norm = mcolors.DivergingNorm(2, vmin=None, vmax=10)
+    norm.autoscale_None([1, 2, 3, 4, 5])
+    assert norm(1) == 0
+    assert norm.vmin == 1
+
+
+def test_DivergingNorm_scale():
+    norm = mcolors.DivergingNorm(2)
+    assert norm.scaled() is False
+    norm([1, 2, 3, 4])
+    assert norm.scaled() is True
+
+
+def test_DivergingNorm_scaleout_center():
+    # test the vmin never goes above vcenter
+    norm = mcolors.DivergingNorm(vcenter=0)
+    x = norm([1, 2, 3, 5])
+    assert norm.vmin == 0
+    assert norm.vmax == 5
+
+
+def test_DivergingNorm_scaleout_center_max():
+    # test the vmax never goes below vcenter
+    norm = mcolors.DivergingNorm(vcenter=0)
+    x = norm([-1, -2, -3, -5])
+    assert norm.vmax == 0
+    assert norm.vmin == -5
+
+
+def test_DivergingNorm_Even():
+    norm = mcolors.DivergingNorm(vmin=-1, vcenter=0, vmax=4)
+    vals = np.array([-1.0, -0.5, 0.0, 1.0, 2.0, 3.0, 4.0])
+    expected = np.array([0.0, 0.25, 0.5, 0.625, 0.75, 0.875, 1.0])
+    assert_array_equal(norm(vals), expected)
+
+
+def test_DivergingNorm_Odd():
+    norm = mcolors.DivergingNorm(vmin=-2, vcenter=0, vmax=5)
+    vals = np.array([-2.0, -1.0, 0.0, 1.0, 2.0, 3.0, 4.0, 5.0])
+    expected = np.array([0.0, 0.25, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0])
+    assert_array_equal(norm(vals), expected)
+
+
+def test_DivergingNorm_VminEqualsVcenter():
+    with pytest.raises(ValueError):
+        norm = mcolors.DivergingNorm(vmin=-2, vcenter=-2, vmax=2)
+
+
+def test_DivergingNorm_VmaxEqualsVcenter():
+    with pytest.raises(ValueError):
+        norm = mcolors.DivergingNorm(vmin=-2, vcenter=2, vmax=2)
+
+
+def test_DivergingNorm_VminGTVcenter():
+    with pytest.raises(ValueError):
+        norm = mcolors.DivergingNorm(vmin=10, vcenter=0, vmax=20)
+
+
+def test_DivergingNorm_DivergingNorm_VminGTVmax():
+    with pytest.raises(ValueError):
+        norm = mcolors.DivergingNorm(vmin=10, vcenter=0, vmax=5)
+
+
+def test_DivergingNorm_VcenterGTVmax():
+    vals = np.arange(50)
+    with pytest.raises(ValueError):
+        norm = mcolors.DivergingNorm(vmin=10, vcenter=25, vmax=20)
+
+
+def test_DivergingNorm_premature_scaling():
+    norm = mcolors.DivergingNorm(vcenter=2)
+    with pytest.raises(ValueError):
+        norm.inverse(np.array([0.1, 0.5, 0.9]))
+
+
 def test_SymLogNorm():
     """
     Test SymLogNorm behavior
diff --git a/lib/matplotlib/tests/test_contour.py b/lib/matplotlib/tests/test_contour.py
@@ -404,7 +404,10 @@ def test_contourf_log_extension():
 
 
 @image_comparison(baseline_images=['contour_addlines'],
-                  extensions=['png'], remove_text=True, style='mpl20')
+                  extensions=['png'], remove_text=True, style='mpl20',
+                  tol=0.03)
+# tolerance is because image changed minutely when tick finding on
+# colorbars was cleaned up...
 def test_contour_addlines():
     fig, ax = plt.subplots()
     np.random.seed(19680812)
diff --git a/tutorials/colors/colormapnorms.py b/tutorials/colors/colormapnorms.py
