Merge pull request matplotlib#9903 from jklymak/enh-colorbar-ticks

ENH: adjustable colorbar ticks

Author: dopplershift

doc/api/api_changes/2017-12-01-JMK.rst

@@ -0,0 +1,8 @@
+The ticks for colorbar now adjust for the size of the colorbar
+--------------------------------------------------------------
+
+Colorbar ticks now adjust for the size of the colorbar if the
+colorbar is made from a mappable that is not a contour or
+doesn't have a BoundaryNorm, or boundaries are not specified.
+If boundaries, etc are specified, the colorbar maintains the
+original behaviour.

doc/users/next_whats_new/colorbarticks.rst

@@ -0,0 +1,7 @@
+Colorbar ticks can now be automatic
+-----------------------------------
+
+The number of ticks on colorbars was appropriate for a large colorbar, but
+looked bad if the colorbar was made smaller (i.e. via the ``shrink`` kwarg).
+This has been changed so that the number of ticks is now responsive to how
+large the colorbar is.  

lib/matplotlib/colorbar.py

@@ -19,6 +19,7 @@
 
 '''
 
+import logging
 import warnings
 
 import numpy as np
@@ -39,6 +40,8 @@
 import matplotlib._constrained_layout as constrained_layout
 from matplotlib import docstring
 
+_log = logging.getLogger(__name__)
+
 make_axes_kw_doc = '''
 
     ============= ====================================================
@@ -212,6 +215,63 @@ def _set_ticks_on_axis_warn(*args, **kw):
     warnings.warn("Use the colorbar set_ticks() method instead.")
 
 
+class _ColorbarAutoLocator(ticker.MaxNLocator):
+    """
+    AutoLocator for Colorbar
+
+    This locator is just a `.MaxNLocator` except the min and max are
+    clipped by the norm's min and max (i.e. vmin/vmax from the
+    image/pcolor/contour object).  This is necessary so ticks don't
+    extrude into the "extend regions".
+    """
+
+    def __init__(self, colorbar):
+        """
+        This ticker needs to know the *colorbar* so that it can access
+        its *vmin* and *vmax*.  Otherwise it is the same as
+        `~.ticker.AutoLocator`.
+        """
+
+        self._colorbar = colorbar
+        nbins = 'auto'
+        steps = [1, 2, 2.5, 5, 10]
+        ticker.MaxNLocator.__init__(self, nbins=nbins, steps=steps)
+
+    def tick_values(self, vmin, vmax):
+        vmin = max(vmin, self._colorbar.norm.vmin)
+        vmax = min(vmax, self._colorbar.norm.vmax)
+        return ticker.MaxNLocator.tick_values(self, vmin, vmax)
+
+
+class _ColorbarLogLocator(ticker.LogLocator):
+    """
+    LogLocator for Colorbarbar
+
+    This locator is just a `.LogLocator` except the min and max are
+    clipped by the norm's min and max (i.e. vmin/vmax from the
+    image/pcolor/contour object).  This is necessary so ticks don't
+    extrude into the "extend regions".
+
+    """
+    def __init__(self, colorbar, *args, **kwargs):
+        """
+        _ColorbarLogLocator(colorbar, *args, **kwargs)
+
+        This ticker needs to know the *colorbar* so that it can access
+        its *vmin* and *vmax*.  Otherwise it is the same as
+        `~.ticker.LogLocator`.  The ``*args`` and ``**kwargs`` are the
+        same as `~.ticker.LogLocator`.
+        """
+        self._colorbar = colorbar
+        ticker.LogLocator.__init__(self, *args, **kwargs)
+
+    def tick_values(self, vmin, vmax):
+        vmin = self._colorbar.norm.vmin
+        vmax = self._colorbar.norm.vmax
+        ticks = ticker.LogLocator.tick_values(self, vmin, vmax)
+        return ticks[(ticks >= vmin) & (ticks <= vmax)]
+
+
 class ColorbarBase(cm.ScalarMappable):
     '''
     Draw a colorbar in an existing axes.
@@ -341,8 +401,15 @@ def draw_all(self):
         and do all the drawing.
         '''
 
+        # sets self._boundaries and self._values in real data units.
+        # takes into account extend values:
         self._process_values()
+        # sets self.vmin and vmax in data units, but just for
+        # the part of the colorbar that is not part of the extend
+        # patch:
         self._find_range()
+        # returns the X and Y mesh, *but* this was/is in normalized
+        # units:
         X, Y = self._mesh()
         C = self._values[:, np.newaxis]
         self._config_axes(X, Y)
@@ -351,35 +418,105 @@ def draw_all(self):
 
     def config_axis(self):
         ax = self.ax
+        if (isinstance(self.norm, colors.LogNorm)
+                and self._use_auto_colorbar_locator()):
+            # *both* axes are made log so that determining the
+            # mid point is easier.
+            ax.set_xscale('log')
+            ax.set_yscale('log')
+
         if self.orientation == 'vertical':
-            ax.xaxis.set_ticks([])
-            # location is either one of 'bottom' or 'top'
-            ax.yaxis.set_label_position(self.ticklocation)
-            ax.yaxis.set_ticks_position(self.ticklocation)
+            long_axis, short_axis = ax.yaxis, ax.xaxis
         else:
-            ax.yaxis.set_ticks([])
-            # location is either one of 'left' or 'right'
-            ax.xaxis.set_label_position(self.ticklocation)
-            ax.xaxis.set_ticks_position(self.ticklocation)
+            long_axis, short_axis = ax.xaxis, ax.yaxis
+
+        long_axis.set_label_position(self.ticklocation)
+        long_axis.set_ticks_position(self.ticklocation)
+        short_axis.set_ticks([])
+        short_axis.set_ticks([], minor=True)
 
         self._set_label()
 
+    def _get_ticker_locator_formatter(self):
+        """
+        This code looks at the norm being used by the colorbar
+        and decides what locator and formatter to use.  If ``locator`` has
+        already been set by hand, it just returns
+        ``self.locator, self.formatter``.
+        """
+        locator = self.locator
+        formatter = self.formatter
+        if locator is None:
+            if self.boundaries is None:
+                if isinstance(self.norm, colors.NoNorm):
+                    nv = len(self._values)
+                    base = 1 + int(nv / 10)
+                    locator = ticker.IndexLocator(base=base, offset=0)
+                elif isinstance(self.norm, colors.BoundaryNorm):
+                    b = self.norm.boundaries
+                    locator = ticker.FixedLocator(b, nbins=10)
+                elif isinstance(self.norm, colors.LogNorm):
+                    locator = _ColorbarLogLocator(self)
+                elif isinstance(self.norm, colors.SymLogNorm):
+                    # The subs setting here should be replaced
+                    # by logic in the locator.
+                    locator = ticker.SymmetricalLogLocator(
+                                      subs=np.arange(1, 10),
+                                      linthresh=self.norm.linthresh,
+                                      base=10)
+                else:
+                    if mpl.rcParams['_internal.classic_mode']:
+                        locator = ticker.MaxNLocator()
+                    else:
+                        locator = _ColorbarAutoLocator(self)
+            else:
+                b = self._boundaries[self._inside]
+                locator = ticker.FixedLocator(b, nbins=10)
+        _log.debug('locator: %r', locator)
+        return locator, formatter
+
+    def _use_auto_colorbar_locator(self):
+        """
+        Return if we should use an adjustable tick locator or a fixed
+        one.  (check is used twice so factored out here...)
+        """
+        return (self.boundaries is None
+                and self.values is None
+                and ((type(self.norm) == colors.Normalize)
+                    or (type(self.norm) == colors.LogNorm)))
+
     def update_ticks(self):
         """
         Force the update of the ticks and ticklabels. This must be
         called whenever the tick locator and/or tick formatter changes.
         """
         ax = self.ax
-        ticks, ticklabels, offset_string = self._ticker()
-        if self.orientation == 'vertical':
-            ax.yaxis.set_ticks(ticks)
-            ax.set_yticklabels(ticklabels)
-            ax.yaxis.get_major_formatter().set_offset_string(offset_string)
+        # get the locator and formatter.  Defaults to
+        # self.locator if not None..
+        locator, formatter = self._get_ticker_locator_formatter()
 
+        if self.orientation == 'vertical':
+            long_axis, short_axis = ax.yaxis, ax.xaxis
         else:
-            ax.xaxis.set_ticks(ticks)
-            ax.set_xticklabels(ticklabels)
-            ax.xaxis.get_major_formatter().set_offset_string(offset_string)
+            long_axis, short_axis = ax.xaxis, ax.yaxis
+
+        if self._use_auto_colorbar_locator():
+            _log.debug('Using auto colorbar locator on colorbar')
+            _log.debug('locator: %r', locator)
+            long_axis.set_major_locator(locator)
+            long_axis.set_major_formatter(formatter)
+            if type(self.norm) == colors.LogNorm:
+                long_axis.set_minor_locator(_ColorbarLogLocator(self,
+                            base=10., subs='auto'))
+                long_axis.set_minor_formatter(
+                    ticker.LogFormatter()
+                )
+        else:
+            _log.debug('Using fixed locator on colorbar')
+            ticks, ticklabels, offset_string = self._ticker(locator, formatter)
+            long_axis.set_ticks(ticks)
+            long_axis.set_ticklabels(ticklabels)
+            long_axis.get_major_formatter().set_offset_string(offset_string)
 
     def set_ticks(self, ticks, update_ticks=True):
         """
@@ -515,6 +652,7 @@ def _add_solids(self, X, Y, C):
         # since the axes object should already have hold set.
         _hold = self.ax._hold
         self.ax._hold = True
+        _log.debug('Setting pcolormesh')
         col = self.ax.pcolormesh(*args, **kw)
         self.ax._hold = _hold
         #self.add_observer(col) # We should observe, not be observed...
@@ -568,39 +706,11 @@ def add_lines(self, levels, colors, linewidths, erase=True):
         self.ax.add_collection(col)
         self.stale = True
 
-    def _ticker(self):
+    def _ticker(self, locator, formatter):
         '''
         Return the sequence of ticks (colorbar data locations),
         ticklabels (strings), and the corresponding offset string.
         '''
-        locator = self.locator
-        formatter = self.formatter
-        if locator is None:
-            if self.boundaries is None:
-                if isinstance(self.norm, colors.NoNorm):
-                    nv = len(self._values)
-                    base = 1 + int(nv / 10)
-                    locator = ticker.IndexLocator(base=base, offset=0)
-                elif isinstance(self.norm, colors.BoundaryNorm):
-                    b = self.norm.boundaries
-                    locator = ticker.FixedLocator(b, nbins=10)
-                elif isinstance(self.norm, colors.LogNorm):
-                    locator = ticker.LogLocator(subs='all')
-                elif isinstance(self.norm, colors.SymLogNorm):
-                    # The subs setting here should be replaced
-                    # by logic in the locator.
-                    locator = ticker.SymmetricalLogLocator(
-                                      subs=np.arange(1, 10),
-                                      linthresh=self.norm.linthresh,
-                                      base=10)
-                else:
-                    if mpl.rcParams['_internal.classic_mode']:
-                        locator = ticker.MaxNLocator()
-                    else:
-                        locator = ticker.AutoLocator()
-            else:
-                b = self._boundaries[self._inside]
-                locator = ticker.FixedLocator(b, nbins=10)
         if isinstance(self.norm, colors.NoNorm) and self.boundaries is None:
             intv = self._values[0], self._values[-1]
         else:
@@ -840,17 +950,29 @@ def _mesh(self):
         transposition for a horizontal colorbar are done outside
         this function.
         '''
+        # if boundaries and values are None, then we can go ahead and
+        # scale this up for Auto tick location.  Otherwise we
+        # want to keep normalized between 0 and 1 and use manual tick
+        # locations.
+
         x = np.array([0.0, 1.0])
         if self.spacing == 'uniform':
             y = self._uniform_y(self._central_N())
         else:
             y = self._proportional_y()
+        if self._use_auto_colorbar_locator():
+            y = self.norm.inverse(y)
+            x = self.norm.inverse(x)
         self._y = y
         X, Y = np.meshgrid(x, y)
+        if self._use_auto_colorbar_locator():
+            xmid = self.norm.inverse(0.5)
+        else:
+            xmid = 0.5
         if self._extend_lower() and not self.extendrect:
-            X[0, :] = 0.5
+            X[0, :] = xmid
         if self._extend_upper() and not self.extendrect:
-            X[-1, :] = 0.5
+            X[-1, :] = xmid
         return X, Y
 
     def _locate(self, x):