Made colorbar.py accept numpy array input, compatible with output from subplots

jklymak · jklymak · commit 0e339a16a998 · 2017-06-10T15:28:16.000-07:00
diff --git a/lib/matplotlib/colorbar.py b/lib/matplotlib/colorbar.py
@@ -39,7 +39,7 @@
 import matplotlib.patches as mpatches
 import matplotlib.path as mpath
 import matplotlib.ticker as ticker
-import matplotlib.transforms as mtransforms
+import matplotlib.transforms as mtrans
 
 from matplotlib import docstring
 
@@ -256,8 +256,6 @@ class ColorbarBase(cm.ScalarMappable):
                    'min': slice(1, None),
                    'max': slice(0, -1)}
 
-    n_rasterize = 50  # rasterize solids if number of colors >= n_rasterize
-
     def __init__(self, ax, cmap=None,
                  norm=None,
                  alpha=None,
@@ -311,13 +309,10 @@ def __init__(self, ax, cmap=None,
             self.locator = ticks    # Handle default in _ticker()
         if format is None:
             if isinstance(self.norm, colors.LogNorm):
-                self.formatter = ticker.LogFormatterSciNotation()
-            elif isinstance(self.norm, colors.SymLogNorm):
-                self.formatter = ticker.LogFormatterSciNotation(
-                                        linthresh=self.norm.linthresh)
+                self.formatter = ticker.LogFormatterMathtext()
             else:
                 self.formatter = ticker.ScalarFormatter()
-        elif isinstance(format, six.string_types):
+        elif cbook.is_string_like(format):
             self.formatter = ticker.FormatStrFormatter(format)
         else:
             self.formatter = format  # Assume it is a Formatter
@@ -344,7 +339,6 @@ def draw_all(self):
         Calculate any free parameters based on the current cmap and norm,
         and do all the drawing.
         '''
-
         self._process_values()
         self._find_range()
         X, Y = self._mesh()
@@ -400,10 +394,6 @@ def set_ticks(self, ticks, update_ticks=True):
             self.update_ticks()
         self.stale = True
 
-    def get_ticks(self, minor=False):
-        """Return the x ticks as a list of locations"""
-        return self._tick_data_values
-
     def set_ticklabels(self, ticklabels, update_ticks=True):
         """
         set tick labels. Tick labels are updated immediately unless
@@ -509,10 +499,10 @@ def _add_solids(self, X, Y, C):
         # Save, set, and restore hold state to keep pcolor from
         # clearing the axes. Ordinarily this will not be needed,
         # since the axes object should already have hold set.
-        _hold = self.ax._hold
-        self.ax._hold = True
+        _hold = self.ax.ishold()
+        self.ax.hold(True)
         col = self.ax.pcolormesh(*args, **kw)
-        self.ax._hold = _hold
+        self.ax.hold(_hold)
         #self.add_observer(col) # We should observe, not be observed...
 
         if self.solids is not None:
@@ -527,8 +517,6 @@ def _add_solids(self, X, Y, C):
                                     colors=(mpl.rcParams['axes.edgecolor'],),
                                     linewidths=linewidths)
             self.ax.add_collection(self.dividers)
-        elif len(self._y) >= self.n_rasterize:
-            self.solids.set_rasterized(True)
 
     def add_lines(self, levels, colors, linewidths, erase=True):
         '''
@@ -582,23 +570,13 @@ def _ticker(self):
                     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)
+                    locator = ticker.LogLocator()
                 else:
-                    if mpl.rcParams['_internal.classic_mode']:
-                        locator = ticker.MaxNLocator()
-                    else:
-                        locator = ticker.AutoLocator()
+                    locator = ticker.MaxNLocator()
             else:
                 b = self._boundaries[self._inside]
                 locator = ticker.FixedLocator(b, nbins=10)
-        if isinstance(self.norm, colors.NoNorm) and self.boundaries is None:
+        if isinstance(self.norm, colors.NoNorm):
             intv = self._values[0], self._values[-1]
         else:
             intv = self.vmin, self.vmax
@@ -616,7 +594,6 @@ def _ticker(self):
         else:
             eps = (intv[1] - intv[0]) * 1e-10
             b = b[(b <= intv[1] + eps) & (b >= intv[0] - eps)]
-        self._tick_data_values = b
         ticks = self._locate(b)
         formatter.set_locs(b)
         ticklabels = [formatter(t, i) for i, t in enumerate(b)]
@@ -687,24 +664,15 @@ def _process_values(self, b=None):
                 self.norm.vmin = 0
                 self.norm.vmax = 1
 
-            self.norm.vmin, self.norm.vmax = mtransforms.nonsingular(
-                self.norm.vmin,
-                self.norm.vmax,
-                expander=0.1)
+            self.norm.vmin, self.norm.vmax = mtrans.nonsingular(self.norm.vmin,
+                                                                self.norm.vmax,
+                                                                expander=0.1)
 
             b = self.norm.inverse(self._uniform_y(self.cmap.N + 1))
-
-            if isinstance(self.norm, colors.LogNorm):
-                # If using a lognorm, ensure extensions don't go negative
-                if self._extend_lower():
-                    b[0] = 0.9 * b[0]
-                if self._extend_upper():
-                    b[-1] = 1.1 * b[-1]
-            else:
-                if self._extend_lower():
-                    b[0] = b[0] - 1
-                if self._extend_upper():
-                    b[-1] = b[-1] + 1
+            if self._extend_lower():
+                b[0] = b[0] - 1
+            if self._extend_upper():
+                b[-1] = b[-1] + 1
         self._process_values(b)
 
     def _find_range(self):
@@ -1050,7 +1018,9 @@ def make_axes(parents, location=None, orientation=None, fraction=0.15,
               shrink=1.0, aspect=20, **kw):
     '''
     Resize and reposition parent axes, and return a child
-    axes suitable for a colorbar.
+    axes suitable for a colorbar::
+
+        cax, kw = make_axes(parent, **kw)
 
     Keyword arguments may include the following (with defaults):
 
@@ -1117,18 +1087,18 @@ def make_axes(parents, location=None, orientation=None, fraction=0.15,
     parent_anchor = kw.pop('panchor', loc_settings['panchor'])
     pad = kw.pop('pad', loc_settings['pad'])
 
-    # turn parents into a list if it is not already
-    if not isinstance(parents, (list, tuple)):
-        parents = [parents]
+    # turn parents into a list if it is not already.  Using numpy because
+    # `fig.subplots` returns an ndarray...
+    parents = np.atleast_1d(parents).ravel().tolist()
 
     fig = parents[0].get_figure()
     if not all(fig is ax.get_figure() for ax in parents):
         raise ValueError('Unable to create a colorbar axes as not all '
                          'parents share the same figure.')
 
     # take a bounding box around all of the given axes
-    parents_bbox = mtransforms.Bbox.union(
-        [ax.get_position(original=True).frozen() for ax in parents])
+    parents_bbox = mtrans.Bbox.union([ax.get_position(original=True).frozen()
+                                      for ax in parents])
 
     pb = parents_bbox
     if location in ('left', 'right'):
@@ -1149,12 +1119,12 @@ def make_axes(parents, location=None, orientation=None, fraction=0.15,
 
     # define a transform which takes us from old axes coordinates to
     # new axes coordinates
-    shrinking_trans = mtransforms.BboxTransform(parents_bbox, pb1)
+    shrinking_trans = mtrans.BboxTransform(parents_bbox, pb1)
 
     # transform each of the axes in parents using the new transform
     for ax in parents:
         new_posn = shrinking_trans.transform(ax.get_position())
-        new_posn = mtransforms.Bbox(new_posn)
+        new_posn = mtrans.Bbox(new_posn)
         ax.set_position(new_posn)
         if parent_anchor is not False:
             ax.set_anchor(parent_anchor)
@@ -1184,7 +1154,9 @@ def make_axes_gridspec(parent, **kw):
         of the parent with a new one.
 
     While this function is meant to be compatible with *make_axes*,
-    there could be some minor differences.
+    there could be some minor differences.::
+
+        cax, kw = make_axes_gridspec(parent, **kw)
 
     Keyword arguments may include the following (with defaults):
 
@@ -1287,8 +1259,8 @@ def _add_solids(self, X, Y, C):
         # Save, set, and restore hold state to keep pcolor from
         # clearing the axes. Ordinarily this will not be needed,
         # since the axes object should already have hold set.
-        _hold = self.ax._hold
-        self.ax._hold = True
+        _hold = self.ax.ishold()
+        self.ax.hold(True)
 
         kw = {'alpha': self.alpha, }
 
@@ -1334,7 +1306,7 @@ def _add_solids(self, X, Y, C):
                     linewidths=(0.5 * mpl.rcParams['axes.linewidth'],))
             self.ax.add_collection(self.dividers)
 
-        self.ax._hold = _hold
+        self.ax.hold(_hold)
 
 
 def colorbar_factory(cax, mappable, **kwargs):
