matplotlib · fmaussion · May 2, 2017 · Sep 9, 2015 · Nov 10, 2015 · May 2, 2017
diff --git a/doc/users/whats_new/extend_kwarg_to_BoundaryNorm.rst b/doc/users/whats_new/extend_kwarg_to_BoundaryNorm.rst
@@ -0,0 +1,46 @@
+New "extend" keyword to colors.BoundaryNorm
+-------------------------------------------
+
+:func:`~matplotlib.colors.BoundaryNorm` now has an ``extend`` kwarg. This is
+useful when creating a discrete colorbar from a continuous colormap: when
+setting ``extend`` to ``'both'``, ``'min'`` or ``'max'``, the colors are
+interpolated so that the extensions have a different color than the inner
+cells.
+
+Example
+```````
+::
+
+    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
+    # my levels are chosen so that the color bar should be extended
+    levels = [-0.8, -0.5, -0.2, 0.2, 0.5, 0.8]
+    cmap = plt.get_cmap('PiYG')
+
+    # Before this change
+    plt.subplot(2, 1, 1)
+    norm = BoundaryNorm(levels, ncolors=cmap.N)
+    im = plt.pcolormesh(x, y, z, cmap=cmap, norm=norm)
+    plt.colorbar(extend='both')
+    plt.axis([x.min(), x.max(), y.min(), y.max()])
+    plt.title('pcolormesh with extended colorbar')
+
+    # With the new keyword
+    norm = BoundaryNorm(levels, ncolors=cmap.N, extend='both')
+    plt.subplot(2, 1, 2)
+    im = plt.pcolormesh(x, y, z, cmap=cmap, norm=norm)
+    plt.colorbar()  # note that the colorbar is updated accordingly
+    plt.axis([x.min(), x.max(), y.min(), y.max()])
+    plt.title('pcolormesh with extended BoundaryNorm')
+
+    plt.show()
diff --git a/lib/matplotlib/colorbar.py b/lib/matplotlib/colorbar.py
@@ -348,7 +348,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,
@@ -365,6 +365,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
         cm.ScalarMappable.__init__(self, cmap=cmap, norm=norm)
         self.values = values

diff --git a/lib/matplotlib/colors.py b/lib/matplotlib/colors.py
@@ -1246,7 +1246,7 @@ class BoundaryNorm(Normalize):
     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
         ----------
@@ -1263,25 +1263,41 @@ 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 :meth:`Colormap.__call__`.
+        extend : {'neither', 'both', 'min', 'max'}, optional
+            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
+            :class:`~matplotlib.colorbar.Colorbar` will be drawn with
+            the triangle extension on the left side.
 
         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, doesn't equal
+        *ncolors*, the color is chosen by linear interpolation of the
+        bin number onto color numbers.
         """
+        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
 
     def __call__(self, value, clip=None):
         if clip is None:
@@ -1295,11 +1311,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
@@ -144,6 +144,86 @@ def test_BoundaryNorm():
     vals = np.ma.masked_invalid([np.Inf])
     assert np.all(bn(vals).mask)
 
+    # Testing extend keyword
+    bounds = [1, 2, 3]
+    cmap = cm.get_cmap('jet')
+
+    refnorm = mcolors.BoundaryNorm([0] + bounds + [4], cmap.N)
+    mynorm = mcolors.BoundaryNorm(bounds, cmap.N, extend='both')
+    x = np.random.random(100) + 1.5
+    np.testing.assert_array_equal(refnorm(x), mynorm(x))
+
+    # Min and max
+    cmref = mcolors.ListedColormap(['blue', 'red'])
+    cmref.set_over('black')
+    cmref.set_under('white')
+
+    cmshould = mcolors.ListedColormap(['white', 'blue', 'red', 'black'])
+    cmshould.set_over(cmshould(cmshould.N))
+    cmshould.set_under(cmshould(0))
+
+    refnorm = mcolors.BoundaryNorm(bounds, cmref.N)
+    mynorm = mcolors.BoundaryNorm(bounds, cmshould.N, extend='both')
+    np.testing.assert_array_equal(refnorm.vmin, mynorm.vmin)
+    np.testing.assert_array_equal(refnorm.vmax, mynorm.vmax)
+    x = [-1, 1.2, 2.3, 9.6]
+    np.testing.assert_array_equal(cmshould([0, 1, 2, 3]), cmshould(mynorm(x)))
+    x = np.random.randn(100) * 10 + 2
+    np.testing.assert_array_equal(cmref(refnorm(x)), cmshould(mynorm(x)))
+
+    np.testing.assert_array_equal(-1, mynorm(-1))
+    np.testing.assert_array_equal(1, mynorm(1.1))
+    np.testing.assert_array_equal(4, mynorm(12))
+
+    # Test raises
+    with pytest.raises(ValueError):
+        mcolors.BoundaryNorm(bounds, cmref.N, extend='both', clip=True)
+
+    # Just min
+    cmref = mcolors.ListedColormap(['blue', 'red'])
+    cmref.set_under('white')
+    cmshould = mcolors.ListedColormap(['white', 'blue', 'red'])
+    cmshould.set_under(cmshould(0))
+
+    np.testing.assert_array_equal(2, cmref.N)
+    np.testing.assert_array_equal(3, cmshould.N)
+    refnorm = mcolors.BoundaryNorm(bounds, cmref.N)
+    mynorm = mcolors.BoundaryNorm(bounds, cmshould.N, extend='min')
+    np.testing.assert_array_equal(refnorm.vmin, mynorm.vmin)
+    np.testing.assert_array_equal(refnorm.vmax, mynorm.vmax)
+    x = [-1, 1.2, 2.3]
+    np.testing.assert_array_equal(cmshould([0, 1, 2]), cmshould(mynorm(x)))
+    x = np.random.randn(100) * 10 + 2
+    np.testing.assert_array_equal(cmref(refnorm(x)), cmshould(mynorm(x)))
+
+    # Just max
+    cmref = mcolors.ListedColormap(['blue', 'red'])
+    cmref.set_over('black')
+    cmshould = mcolors.ListedColormap(['blue', 'red', 'black'])
+    cmshould.set_over(cmshould(2))
+
+    np.testing.assert_array_equal(2, cmref.N)
+    np.testing.assert_array_equal(3, cmshould.N)
+    refnorm = mcolors.BoundaryNorm(bounds, cmref.N)
+    mynorm = mcolors.BoundaryNorm(bounds, cmshould.N, extend='max')
+    np.testing.assert_array_equal(refnorm.vmin, mynorm.vmin)
+    np.testing.assert_array_equal(refnorm.vmax, mynorm.vmax)
+    x = [1.2, 2.3, 4]
+    np.testing.assert_array_equal(cmshould([0, 1, 2]), cmshould(mynorm(x)))
+    x = np.random.randn(100) * 10 + 2
+    np.testing.assert_array_equal(cmref(refnorm(x)), cmshould(mynorm(x)))
+
+    # General case
+    bounds = [1, 2, 3, 4]
+    cmap = cm.get_cmap('jet')
+    mynorm = mcolors.BoundaryNorm(bounds, cmap.N, extend='both')
+    refnorm = mcolors.BoundaryNorm([-100] + bounds + [100], cmap.N)
+    x = np.random.randn(100) * 10 - 5
+    ref = refnorm(x)
+    ref = np.where(ref == 0, -1, ref)
+    ref = np.where(ref == cmap.N-1, cmap.N, ref)
+    np.testing.assert_array_equal(ref, mynorm(x))
+
 
 def test_LogNorm():
     """
@@ -304,6 +384,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