Merge pull request #17534 from efiring/fmaussion-extended-BoundaryNorm

QuLogic · web-flow · commit 4356b677a334 · 2020-06-02T16:45:31.000-04:00
Fmaussion extended boundary norm
diff --git a/doc/users/next_whats_new/2020-05-28-extend_kwarg_to_BoundaryNorm.rst b/doc/users/next_whats_new/2020-05-28-extend_kwarg_to_BoundaryNorm.rst
@@ -0,0 +1,49 @@
+New "extend" keyword to colors.BoundaryNorm
+-------------------------------------------
+
+`~.colors.BoundaryNorm` now has an ``extend`` keyword argument, analogous to
+``extend`` in `~.axes.Axes.contourf`. When set to 'both', 'min', or 'max',
+it maps the corresponding out-of-range values to `~.colors.Colormap`
+lookup-table indices near the appropriate ends of their range so that the
+colors for out-of range values are adjacent to, but distinct from, their
+in-range neighbors.  The colorbar inherits the ``extend`` argument from the
+norm, so with ``extend='both'``, for example, the colorbar will have triangular
+extensions for out-of-range values with colors that differ from adjacent in-range
+colors.
+
+  .. plot::
+
+    import matplotlib.pyplot as plt
+    from matplotlib.colors import BoundaryNorm
+    import numpy as np
+
+    # Make the data
+    dx, dy = 0.05, 0.05
+    y, x = np.mgrid[slice(1, 5 + dy, dy),
+                    slice(1, 5 + dx, dx)]
+    z = np.sin(x) ** 10 + np.cos(10 + y * x) * np.cos(x)
+    z = z[:-1, :-1]
+
+    # Z roughly varies between -1 and +1.
+    # Color boundary levels range from -0.8 to 0.8, so there are out-of-bounds
+    # areas.
+    levels = [-0.8, -0.5, -0.2, 0.2, 0.5, 0.8]
+    cmap = plt.get_cmap('PiYG')
+
+    fig, axs = plt.subplots(nrows=2, constrained_layout=True, sharex=True)
+
+    # Before this change:
+    norm = BoundaryNorm(levels, ncolors=cmap.N)
+    im = axs[0].pcolormesh(x, y, z, cmap=cmap, norm=norm)
+    fig.colorbar(im, ax=axs[0], extend='both')
+    axs[0].axis([x.min(), x.max(), y.min(), y.max()])
+    axs[0].set_title("Colorbar with extend='both'")
+
+    # With the new keyword:
+    norm = BoundaryNorm(levels, ncolors=cmap.N, extend='both')
+    im = axs[1].pcolormesh(x, y, z, cmap=cmap, norm=norm)
+    fig.colorbar(im, ax=axs[1])  # note that the colorbar is updated accordingly
+    axs[1].axis([x.min(), x.max(), y.min(), y.max()])
+    axs[1].set_title("BoundaryNorm with extend='both'")
+
+    plt.show()
diff --git a/lib/matplotlib/colorbar.py b/lib/matplotlib/colorbar.py
@@ -402,7 +402,7 @@ def __init__(self, ax, cmap=None,
                  boundaries=None,
                  orientation='vertical',
                  ticklocation='auto',
-                 extend='neither',
+                 extend=None,
                  spacing='uniform',  # uniform or proportional
                  ticks=None,
                  format=None,
@@ -430,6 +430,11 @@ def __init__(self, ax, cmap=None,
             cmap = cm.get_cmap()
         if norm is None:
             norm = colors.Normalize()
+        if extend is None:
+            if hasattr(norm, 'extend'):
+                extend = norm.extend
+            else:
+                extend = 'neither'
         self.alpha = alpha
         self.cmap = cmap
         self.norm = norm
diff --git a/lib/matplotlib/colors.py b/lib/matplotlib/colors.py
@@ -530,7 +530,7 @@ def __call__(self, X, alpha=None, bytes=False):
         """
         Parameters
         ----------
-        X : float, ndarray
+        X : float or int, ndarray or scalar
             The data value(s) to convert to RGBA.
             For floats, X should be in the interval ``[0.0, 1.0]`` to
             return the RGBA values ``X*100`` percent along the Colormap line.
@@ -1410,7 +1410,7 @@ class BoundaryNorm(Normalize):
     interpolation, but using integers seems simpler, and reduces the number of
     conversions back and forth between integer and floating point.
     """
-    def __init__(self, boundaries, ncolors, clip=False):
+    def __init__(self, boundaries, ncolors, clip=False, *, extend='neither'):
         """
         Parameters
         ----------
@@ -1427,25 +1427,50 @@ def __init__(self, boundaries, ncolors, clip=False):
             they are below ``boundaries[0]`` or mapped to *ncolors* if they are
             above ``boundaries[-1]``. These are then converted to valid indices
             by `Colormap.__call__`.
+        extend : {'neither', 'both', 'min', 'max'}, default: 'neither'
+            Extend the number of bins to include one or both of the
+            regions beyond the boundaries.  For example, if ``extend``
+            is 'min', then the color to which the region between the first
+            pair of boundaries is mapped will be distinct from the first
+            color in the colormap, and by default a
+            `~matplotlib.colorbar.Colorbar` will be drawn with
+            the triangle extension on the left or lower end.
+
+        Returns
+        -------
+        int16 scalar or array
 
         Notes
         -----
         *boundaries* defines the edges of bins, and data falling within a bin
         is mapped to the color with the same index.
 
-        If the number of bins doesn't equal *ncolors*, the color is chosen
-        by linear interpolation of the bin number onto color numbers.
+        If the number of bins, including any extensions, is less than
+        *ncolors*, the color index is chosen by linear interpolation, mapping
+        the ``[0, nbins - 1]`` range onto the ``[0, ncolors - 1]`` range.
         """
+        if clip and extend != 'neither':
+            raise ValueError("'clip=True' is not compatible with 'extend'")
         self.clip = clip
         self.vmin = boundaries[0]
         self.vmax = boundaries[-1]
         self.boundaries = np.asarray(boundaries)
         self.N = len(self.boundaries)
         self.Ncmap = ncolors
-        if self.N - 1 == self.Ncmap:
-            self._interp = False
-        else:
-            self._interp = True
+        self.extend = extend
+
+        self._N = self.N - 1  # number of colors needed
+        self._offset = 0
+        if extend in ('min', 'both'):
+            self._N += 1
+            self._offset = 1
+        if extend in ('max', 'both'):
+            self._N += 1
+        if self._N > self.Ncmap:
+            raise ValueError(f"There are {self._N} color bins including "
+                             f"extensions, but ncolors = {ncolors}; "
+                             "ncolors must equal or exceed the number of "
+                             "bins")
 
     def __call__(self, value, clip=None):
         if clip is None:
@@ -1459,11 +1484,9 @@ def __call__(self, value, clip=None):
             max_col = self.Ncmap - 1
         else:
             max_col = self.Ncmap
-        iret = np.zeros(xx.shape, dtype=np.int16)
-        for i, b in enumerate(self.boundaries):
-            iret[xx >= b] = i
-        if self._interp:
-            scalefac = (self.Ncmap - 1) / (self.N - 2)
+        iret = np.digitize(xx, self.boundaries) - 1 + self._offset
+        if self.Ncmap > self._N:
+            scalefac = (self.Ncmap - 1) / (self._N - 1)
             iret = (iret * scalefac).astype(np.int16)
         iret[xx < self.vmin] = -1
         iret[xx >= self.vmax] = max_col
diff --git a/lib/matplotlib/tests/baseline_images/test_colors/boundarynorm_and_colorbar.png b/lib/matplotlib/tests/baseline_images/test_colors/boundarynorm_and_colorbar.png
diff --git a/lib/matplotlib/tests/test_colors.py b/lib/matplotlib/tests/test_colors.py
@@ -265,6 +265,84 @@ def test_BoundaryNorm():
     vals = np.ma.masked_invalid([np.Inf])
     assert np.all(bn(vals).mask)
 
+    # Incompatible extend and clip
+    with pytest.raises(ValueError, match="not compatible"):
+        mcolors.BoundaryNorm(np.arange(4), 5, extend='both', clip=True)
+
+    # Too small ncolors argument
+    with pytest.raises(ValueError, match="ncolors must equal or exceed"):
+        mcolors.BoundaryNorm(np.arange(4), 2)
+
+    with pytest.raises(ValueError, match="ncolors must equal or exceed"):
+        mcolors.BoundaryNorm(np.arange(4), 3, extend='min')
+
+    with pytest.raises(ValueError, match="ncolors must equal or exceed"):
+        mcolors.BoundaryNorm(np.arange(4), 4, extend='both')
+
+    # Testing extend keyword, with interpolation (large cmap)
+    bounds = [1, 2, 3]
+    cmap = cm.get_cmap('viridis')
+    mynorm = mcolors.BoundaryNorm(bounds, cmap.N, extend='both')
+    refnorm = mcolors.BoundaryNorm([0] + bounds + [4], cmap.N)
+    x = np.random.randn(100) * 10 + 2
+    ref = refnorm(x)
+    ref[ref == 0] = -1
+    ref[ref == cmap.N - 1] = cmap.N
+    assert_array_equal(mynorm(x), ref)
+
+    # Without interpolation
+    cmref = mcolors.ListedColormap(['blue', 'red'])
+    cmref.set_over('black')
+    cmref.set_under('white')
+    cmshould = mcolors.ListedColormap(['white', 'blue', 'red', 'black'])
+
+    refnorm = mcolors.BoundaryNorm(bounds, cmref.N)
+    mynorm = mcolors.BoundaryNorm(bounds, cmshould.N, extend='both')
+    assert mynorm.vmin == refnorm.vmin
+    assert mynorm.vmax == refnorm.vmax
+
+    assert mynorm(bounds[0] - 0.1) == -1  # under
+    assert mynorm(bounds[0] + 0.1) == 1   # first bin -> second color
+    assert mynorm(bounds[-1] - 0.1) == cmshould.N - 2  # next-to-last color
+    assert mynorm(bounds[-1] + 0.1) == cmshould.N  # over
+
+    x = [-1, 1.2, 2.3, 9.6]
+    assert_array_equal(cmshould(mynorm(x)), cmshould([0, 1, 2, 3]))
+    x = np.random.randn(100) * 10 + 2
+    assert_array_equal(cmshould(mynorm(x)), cmref(refnorm(x)))
+
+    # Just min
+    cmref = mcolors.ListedColormap(['blue', 'red'])
+    cmref.set_under('white')
+    cmshould = mcolors.ListedColormap(['white', 'blue', 'red'])
+
+    assert cmref.N == 2
+    assert cmshould.N == 3
+    refnorm = mcolors.BoundaryNorm(bounds, cmref.N)
+    mynorm = mcolors.BoundaryNorm(bounds, cmshould.N, extend='min')
+    assert mynorm.vmin == refnorm.vmin
+    assert mynorm.vmax == refnorm.vmax
+    x = [-1, 1.2, 2.3]
+    assert_array_equal(cmshould(mynorm(x)), cmshould([0, 1, 2]))
+    x = np.random.randn(100) * 10 + 2
+    assert_array_equal(cmshould(mynorm(x)), cmref(refnorm(x)))
+
+    # Just max
+    cmref = mcolors.ListedColormap(['blue', 'red'])
+    cmref.set_over('black')
+    cmshould = mcolors.ListedColormap(['blue', 'red', 'black'])
+
+    assert cmref.N == 2
+    assert cmshould.N == 3
+    refnorm = mcolors.BoundaryNorm(bounds, cmref.N)
+    mynorm = mcolors.BoundaryNorm(bounds, cmshould.N, extend='max')
+    assert mynorm.vmin == refnorm.vmin
+    assert mynorm.vmax == refnorm.vmax
+    x = [1.2, 2.3, 4]
+    assert_array_equal(cmshould(mynorm(x)), cmshould([0, 1, 2]))
+    x = np.random.randn(100) * 10 + 2
+    assert_array_equal(cmshould(mynorm(x)), cmref(refnorm(x)))
+
 
 @pytest.mark.parametrize("vmin,vmax", [[-1, 2], [3, 1]])
 def test_lognorm_invalid(vmin, vmax):
@@ -537,6 +615,35 @@ def test_cmap_and_norm_from_levels_and_colors():
     ax.tick_params(labelleft=False, labelbottom=False)
 
 
+@image_comparison(baseline_images=['boundarynorm_and_colorbar'],
+                  extensions=['png'])
+def test_boundarynorm_and_colorbarbase():
+
+    # Make a figure and axes with dimensions as desired.
+    fig = plt.figure()
+    ax1 = fig.add_axes([0.05, 0.80, 0.9, 0.15])
+    ax2 = fig.add_axes([0.05, 0.475, 0.9, 0.15])
+    ax3 = fig.add_axes([0.05, 0.15, 0.9, 0.15])
+
+    # Set the colormap and bounds
+    bounds = [-1, 2, 5, 7, 12, 15]
+    cmap = cm.get_cmap('viridis')
+
+    # Default behavior
+    norm = mcolors.BoundaryNorm(bounds, cmap.N)
+    cb1 = mcolorbar.ColorbarBase(ax1, cmap=cmap, norm=norm, extend='both',
+                                 orientation='horizontal')
+    # New behavior
+    norm = mcolors.BoundaryNorm(bounds, cmap.N, extend='both')
+    cb2 = mcolorbar.ColorbarBase(ax2, cmap=cmap, norm=norm,
+                                 orientation='horizontal')
+
+    # User can still force to any extend='' if really needed
+    norm = mcolors.BoundaryNorm(bounds, cmap.N, extend='both')
+    cb3 = mcolorbar.ColorbarBase(ax3, cmap=cmap, norm=norm,
+                                 extend='neither', orientation='horizontal')
+
+
 def test_cmap_and_norm_from_levels_and_colors2():
     levels = [-1, 2, 2.5, 3]
     colors = ['red', (0, 1, 0), 'blue', (0.5, 0.5, 0.5), (0.0, 0.0, 0.0, 1.0)]
diff --git a/tutorials/colors/colorbar_only.py b/tutorials/colors/colorbar_only.py
@@ -38,11 +38,31 @@
 fig.colorbar(mpl.cm.ScalarMappable(norm=norm, cmap=cmap),
              cax=ax, orientation='horizontal', label='Some Units')
 
+
+###############################################################################
+# Extended colorbar with continuous colorscale
+# --------------------------------------------
+#
+# The second example shows how to make a discrete colorbar based on a
+# continuous cmap. With the "extend" keyword argument the appropriate colors
+# are chosen to fill the colorspace, including the extensions:
+fig, ax = plt.subplots(figsize=(6, 1))
+fig.subplots_adjust(bottom=0.5)
+
+cmap = mpl.cm.viridis
+bounds = [-1, 2, 5, 7, 12, 15]
+norm = mpl.colors.BoundaryNorm(bounds, cmap.N, extend='both')
+cb2 = mpl.colorbar.ColorbarBase(ax, cmap=cmap,
+                                norm=norm,
+                                orientation='horizontal')
+cb2.set_label("Discrete intervals with extend='both' keyword")
+fig.show()
+
 ###############################################################################
 # Discrete intervals colorbar
 # ---------------------------
 #
-# The second example illustrates the use of a
+# The third example illustrates the use of a
 # :class:`~matplotlib.colors.ListedColormap` which generates a colormap from a
 # set of listed colors, `.colors.BoundaryNorm` which generates a colormap
 # index based on discrete intervals and extended ends to show the "over" and
@@ -54,11 +74,15 @@
 # bounds array must be one greater than the length of the color list. The
 # bounds must be monotonically increasing.
 #
-# This time we pass some more arguments in addition to previous arguments to
-# `~.Figure.colorbar`. For the out-of-range values to
-# display on the colorbar, we have to use the *extend* keyword argument. To use
-# *extend*, you must specify two extra boundaries. Finally spacing argument
-# ensures that intervals are shown on colorbar proportionally.
+# This time we pass additional arguments to
+# `~.Figure.colorbar`. For the out-of-range values to display on the colorbar
+# without using the *extend* keyword with
+# `.colors.BoundaryNorm`, we have to use the *extend* keyword argument directly
+# in the colorbar call, and supply an additional boundary on each end of the
+# range.  Here we also
+# use the spacing argument to make
+# the length of each colorbar segment proportional to its corresponding
+# interval.
 
 fig, ax = plt.subplots(figsize=(6, 1))
 fig.subplots_adjust(bottom=0.5)
@@ -72,7 +96,7 @@
 fig.colorbar(
     mpl.cm.ScalarMappable(cmap=cmap, norm=norm),
     cax=ax,
-    boundaries=[0] + bounds + [13],
+    boundaries=[0] + bounds + [13],  # Adding values for extensions.
     extend='both',
     ticks=bounds,
     spacing='proportional',
@@ -84,7 +108,7 @@
 # Colorbar with custom extension lengths
 # --------------------------------------
 #
-# Here we illustrate the use of custom length colorbar extensions, used on a
+# Here we illustrate the use of custom length colorbar extensions, on a
 # colorbar with discrete intervals. To make the length of each extension the
 # same as the length of the interior colors, use ``extendfrac='auto'``.
 
diff --git a/tutorials/colors/colormapnorms.py b/tutorials/colors/colormapnorms.py
