Merge pull request matplotlib#766 from ajdawson/colorbar-extensions

NF - control the length of colorbar extensions

Author: WeatherGod

doc/api/api_changes.rst

@@ -51,6 +51,9 @@ Changes in 1.2.x
   matplotlib axes by providing a ``_as_mpl_axes`` method. See
   :ref:`adding-new-scales` for more detail.
 
+* A new keyword *extendfrac* in :meth:`~matplotlib.pyplot.colorbar` and
+  :class:`~matplotlib.colorbar.ColorbarBase` allows one to control the size of
+  the triangular minimum and maximum extensions on colorbars.
 
 Changes in 1.1.x
 ================

doc/users/whats_new.rst

@@ -26,6 +26,37 @@ Damon McDougall added a new plotting method for the
 
 .. plot:: mpl_examples/mplot3d/trisurf3d_demo.py
 
+Control the lengths of colorbar extensions
+------------------------------------------
+
+Andrew Dawson added a new keyword argument *extendfrac* to
+:meth:`~matplotlib.pyplot.colorbar` to control the length of
+minimum and maximum colorbar extensions.
+
+.. plot::
+
+    import matplotlib.pyplot as plt
+    import numpy as np
+
+    x = y = np.linspace(0., 2*np.pi, 100)
+    X, Y = np.meshgrid(x, y)
+    Z = np.cos(X) * np.sin(0.5*Y)
+
+    clevs = [-.75, -.5, -.25, 0., .25, .5, .75]
+    cmap = plt.cm.get_cmap(name='jet', lut=8)
+
+    ax1 = plt.subplot(211)
+    cs1 = plt.contourf(x, y, Z, clevs, cmap=cmap, extend='both')
+    cb1 = plt.colorbar(orientation='horizontal', extendfrac=None)
+    cb1.set_label('Default length colorbar extensions')
+
+    ax2 = plt.subplot(212)
+    cs2 = plt.contourf(x, y, Z, clevs, cmap=cmap, extend='both')
+    cb2 = plt.colorbar(orientation='horizontal', extendfrac='auto')
+    cb2.set_label('Custom length colorbar extensions')
+
+    plt.show()
+
 .. _whats-new-1-1:
 
 new in matplotlib-1.1

examples/api/colorbar_only.py

@@ -6,8 +6,9 @@
 
 # Make a figure and axes with dimensions as desired.
 fig = pyplot.figure(figsize=(8,3))
-ax1 = fig.add_axes([0.05, 0.65, 0.9, 0.15])
-ax2 = fig.add_axes([0.05, 0.25, 0.9, 0.15])
+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 norm to correspond to the data for which
 # the colorbar will be used.
@@ -47,5 +48,27 @@
                                      orientation='horizontal')
 cb2.set_label('Discrete intervals, some other units')
 
+# The third example illustrates the use of custom length colorbar
+# extensions, used on a colorbar with discrete intervals.
+cmap = mpl.colors.ListedColormap([[0., .4, 1.], [0., .8, 1.],
+    [1., .8, 0.], [1., .4, 0.]])
+cmap.set_over((1., 0., 0.))
+cmap.set_under((0., 0., 1.))
+
+bounds = [-1., -.5, 0., .5, 1.]
+norm = mpl.colors.BoundaryNorm(bounds, cmap.N)
+cb3 = mpl.colorbar.ColorbarBase(ax3, cmap=cmap,
+                                     norm=norm,
+                                     boundaries=[-10]+bounds+[10],
+                                     extend='both',
+                                     # Make the length of each extension
+                                     # the same as the length of the
+                                     # interior colors:
+                                     extendfrac='auto',
+                                     ticks=bounds,
+                                     spacing='uniform',
+                                     orientation='horizontal')
+cb3.set_label('Custom extension lengths, some other units')
+
 pyplot.show()
 

lib/matplotlib/__init__.py

@@ -992,7 +992,8 @@ def tk_window_focus():
     'matplotlib.tests.test_text',
     'matplotlib.tests.test_tightlayout',
     'matplotlib.tests.test_delaunay',
-    'matplotlib.tests.test_legend'
+    'matplotlib.tests.test_legend',
+    'matplotlib.tests.test_colorbar',
     ]
 
 def test(verbosity=1):

lib/matplotlib/colorbar.py

@@ -59,13 +59,29 @@
 
 colormap_kw_doc = '''
 
-    ===========   ====================================================
+    ============  ====================================================
     Property      Description
-    ===========   ====================================================
+    ============  ====================================================
     *extend*      [ 'neither' | 'both' | 'min' | 'max' ]
                   If not 'neither', make pointed end(s) for out-of-
                   range values.  These are set for a given colormap
                   using the colormap set_under and set_over methods.
+    *extendfrac*  [ *None* | 'auto' | length | lengths ]
+                  If set to *None*, both the minimum and maximum
+                  triangular colorbar extensions with have a length of
+                  5% of the interior colorbar length (this is the
+                  default setting). If set to 'auto', makes the
+                  triangular colorbar extensions the same lengths as
+                  the interior boxes (when *spacing* is set to
+                  'uniform') or the same lengths as the respective
+                  adjacent interior boxes (when *spacing* is set to
+                  'proportional'). If a scalar, indicates the length
+                  of both the minimum and maximum triangular colorbar
+                  extensions as a fraction of the interior colorbar
+                  length. A two-element sequence of fractions may also
+                  be given, indicating the lengths of the minimum and
+                  maximum colorbar extensions respectively as a
+                  fraction of the interior colorbar length.
     *spacing*     [ 'uniform' | 'proportional' ]
                   Uniform spacing gives each discrete color the same
                   space; proportional makes the space proportional to
@@ -82,7 +98,7 @@
                   given instead.
     *drawedges*   [ False | True ] If true, draw lines at color
                   boundaries.
-    ===========   ====================================================
+    ============  ====================================================
 
     The following will probably be useful only in the context of
     indexed colors (that is, when the mappable has norm=NoNorm()),
@@ -221,6 +237,7 @@ def __init__(self, ax, cmap=None,
                            format=None,
                            drawedges=False,
                            filled=True,
+                           extendfrac=None,
                            ):
         self.ax = ax
         self._patch_ax()
@@ -236,6 +253,7 @@ def __init__(self, ax, cmap=None,
         self.orientation = orientation
         self.drawedges = drawedges
         self.filled = filled
+        self.extendfrac = extendfrac
         self.solids = None
         self.lines = None
         self.outline = None
@@ -616,6 +634,35 @@ def _extended_N(self):
             N += 1
         return N
 
+    def _get_extension_lengths(self, frac, automin, automax, default=0.05):
+        '''
+        Get the lengths of colorbar extensions.
+
+        A helper method for _uniform_y and _proportional_y.
+        '''
+        # Set the default value.
+        extendlength = np.array([default, default])
+        if isinstance(frac, str):
+            if frac.lower() == 'auto':
+                # Use the provided values when 'auto' is required.
+                extendlength[0] = automin
+                extendlength[1] = automax
+            else:
+                # Any other string is invalid.
+                raise ValueError('invalid value for extendfrac')
+        elif frac is not None:
+            try:
+                # Try to set min and max extension fractions directly.
+                extendlength[:] = frac
+                # If frac is a sequence contaning None then NaN may
+                # be encountered. This is an error.
+                if np.isnan(extendlength).any():
+                    raise ValueError()
+            except (TypeError, ValueError):
+                # Raise an error on encountering an invalid value for frac.
+                raise ValueError('invalid value for extendfrac')
+        return extendlength
+
     def _uniform_y(self, N):
         '''
         Return colorbar data coordinates for *N* uniformly
@@ -624,16 +671,19 @@ def _uniform_y(self, N):
         if self.extend == 'neither':
             y = np.linspace(0, 1, N)
         else:
+            automin = automax = 1. / (N - 1.)
+            extendlength = self._get_extension_lengths(self.extendfrac,
+                    automin, automax, default=0.05)
             if self.extend == 'both':
                 y = np.zeros(N + 2, 'd')
-                y[0] = -0.05
-                y[-1] = 1.05
+                y[0] = 0. - extendlength[0]
+                y[-1] = 1. + extendlength[1]
             elif self.extend == 'min':
                 y = np.zeros(N + 1, 'd')
-                y[0] = -0.05
+                y[0] = 0. - extendlength[0]
             else:
                 y = np.zeros(N + 1, 'd')
-                y[-1] = 1.05
+                y[-1] = 1. + extendlength[1]
             y[self._inside] = np.linspace(0, 1, N)
         return y
 
@@ -648,10 +698,27 @@ def _proportional_y(self):
             y = y / (self._boundaries[-1] - self._boundaries[0])
         else:
             y = self.norm(self._boundaries.copy())
-        if self._extend_lower():
-            y[0] = -0.05
-        if self._extend_upper():
-            y[-1] = 1.05
+        if self.extend == 'min':
+            # Exclude leftmost interval of y.
+            clen = y[-1] - y[1]
+            automin = (y[2] - y[1]) / clen
+            automax = (y[-1] - y[-2]) / clen
+        elif self.extend == 'max':
+            # Exclude rightmost interval in y.
+            clen = y[-2] - y[0]
+            automin = (y[1] - y[0]) / clen
+            automax = (y[-2] - y[-3]) / clen
+        else:
+            # Exclude leftmost and rightmost intervals in y.
+            clen = y[-2] - y[1]
+            automin = (y[2] - y[1]) / clen
+            automax = (y[-2] - y[-3]) / clen
+        extendlength = self._get_extension_lengths(self.extendfrac,
+                automin, automax, default=0.05)
+        if self.extend in ('both', 'min'):
+            y[0] = 0. - extendlength[0]
+        if self.extend in ('both', 'max'):
+            y[-1] = 1. + extendlength[1]
         yi = y[self._inside]
         norm = colors.Normalize(yi[0], yi[-1])
         y[self._inside] = norm(yi)

lib/matplotlib/tests/baseline_images/test_colorbar/colorbar_extensions_proportional.png

@@ -0,0 +1,61 @@
+from matplotlib import rcParams, rcParamsDefault
+from matplotlib.testing.decorators import image_comparison
+import matplotlib.pyplot as plt
+from matplotlib.colors import BoundaryNorm
+from matplotlib.cm import get_cmap
+from matplotlib.colorbar import ColorbarBase
+
+
+def _colorbar_extensions(spacing):
+
+    # Create a color map and specify the levels it represents.
+    cmap = get_cmap("RdBu", lut=5)
+    clevs = [-5., -2.5, -.5, .5, 1.5, 3.5]
+
+    # Define norms for the color maps.
+    norms = dict()
+    norms['neither'] = BoundaryNorm(clevs, len(clevs)-1)
+    norms['min'] = BoundaryNorm([-10]+clevs[1:], len(clevs)-1)
+    norms['max'] = BoundaryNorm(clevs[:-1]+[10], len(clevs)-1)
+    norms['both'] = BoundaryNorm([-10]+clevs[1:-1]+[10], len(clevs)-1)
+
+    # Create a figure and adjust whitespace for subplots.
+    fig = plt.figure()
+    fig.subplots_adjust(hspace=.6)
+    
+    for i, extension_type in enumerate(('neither', 'min', 'max', 'both')):
+        # Get the appropriate norm and use it to get colorbar boundaries.
+        norm = norms[extension_type]
+        boundaries = values = norm.boundaries
+        for j, extendfrac in enumerate((None, 'auto', 0.1)):
+            # Create a subplot.
+            cax = fig.add_subplot(12, 1, i*3+j+1)
+            # Turn off text and ticks.
+            for item in cax.get_xticklabels() + cax.get_yticklabels() +\
+                    cax.get_xticklines() + cax.get_yticklines():
+                item.set_visible(False)
+            # Generate the colorbar.
+            cb = ColorbarBase(cax, cmap=cmap, norm=norm,
+                    boundaries=boundaries, values=values,
+                    extend=extension_type, extendfrac=extendfrac,
+                    orientation='horizontal', spacing=spacing)
+
+    # Return the figure to the caller.
+    return fig
+
+
+@image_comparison(
+        baseline_images=['colorbar_extensions_uniform', 'colorbar_extensions_proportional'],
+        extensions=['png'])
+def test_colorbar_extensions():
+    # Use default params so .matplotlibrc doesn't cause the test to fail.
+    rcParams.update(rcParamsDefault)
+    # Create figures for uniform and proportionally spaced colorbars.
+    fig1 = _colorbar_extensions('uniform')
+    fig2 = _colorbar_extensions('proportional')
+
+
+if __name__ == '__main__':
+    import nose
+    nose.runmodule(argv=['-s', '--with-doctest'], exit=False)
+