FIX colorbars for Norms that do not have a scale.

dstansby · jklymak · commit e175ffcd4086 · 2020-02-02T19:24:55.000-08:00
diff --git a/lib/matplotlib/colorbar.py b/lib/matplotlib/colorbar.py
@@ -452,6 +452,7 @@ def __init__(self, ax, cmap=None,
         self.locator = None
         self.formatter = None
         self._manual_tick_data_values = None
+        self._scale = None  # linear, log10 for now.  Hopefully more?
 
         if ticklocation == 'auto':
             ticklocation = 'bottom' if orientation == 'horizontal' else 'right'
@@ -585,8 +586,8 @@ def _use_auto_colorbar_locator(self):
         one.  (check is used twice so factored out here...)
         """
         contouring = self.boundaries is not None and self.spacing == 'uniform'
-        return (type(self.norm) in [colors.Normalize, colors.LogNorm]
-                and not contouring)
+        return (type(self.norm) in [colors.Normalize, colors.LogNorm] and
+                not contouring)
 
     def _reset_locator_formatter_scale(self):
         """
@@ -602,9 +603,16 @@ def _reset_locator_formatter_scale(self):
             self.ax.set_xscale('log')
             self.ax.set_yscale('log')
             self.minorticks_on()
+            self._scale = 'log'
         else:
             self.ax.set_xscale('linear')
             self.ax.set_yscale('linear')
+            if (isinstance(self.norm, colors.Normalize) and
+                not issubclass(type(self.norm), colors.Normalize)):
+                self._scale = 'linear'
+            else:
+                self._scale = 'manual'
+
 
     def update_ticks(self):
         """
@@ -1119,13 +1127,13 @@ def _mesh(self):
         else:
             y = self._proportional_y()
         xmid = np.array([0.5])
-        try:
+        if not self._scale == 'manual':
             y = norm.inverse(y)
             x = norm.inverse(x)
             xmid = norm.inverse(xmid)
-        except ValueError:
-            # occurs for norms that don't have an inverse, in
-            # which case manually scale:
+        else:
+            # if a norm doesn't have a named scale, or
+            # we are not using a norm
             dv = self.vmax - self.vmin
             x = x * dv + self.vmin
             y = y * dv + self.vmin
diff --git a/lib/matplotlib/tests/baseline_images/test_colorbar/colorbar_norms.png b/lib/matplotlib/tests/baseline_images/test_colorbar/colorbar_norms.png
diff --git a/lib/matplotlib/tests/test_colorbar.py b/lib/matplotlib/tests/test_colorbar.py
@@ -373,7 +373,6 @@ def test_colorbar_get_ticks():
     data = np.arange(1200).reshape(30, 40)
     levels = [0, 200, 400, 600, 800, 1000, 1200]
 
-    plt.subplot()
     plt.contourf(data, levels=levels)
 
     # testing getter for user set ticks
@@ -581,3 +580,18 @@ def test_colorbar_int(clim):
     im = ax.imshow([[*map(np.int16, clim)]])
     fig.colorbar(im)
     assert (im.norm.vmin, im.norm.vmax) == clim
+
+
+@image_comparison(
+    baseline_images=['colorbar_norms'], extensions=['png'])
+def test_cbar_norms():
+    # Test colormaps with different norms
+    norms = [mcolors.Normalize(vmin=-1, vmax=2),
+             mcolors.BoundaryNorm(boundaries=[-1, 0, 1], ncolors=4),
+             mcolors.LogNorm(vmin=1, vmax=1e4),
+             mcolors.PowerNorm(gamma=0.2, vmin=0.1, vmax=0.6),
+             mcolors.SymLogNorm(1, vmin=-10, vmax=10),
+             mcolors.TwoSlopeNorm(1, vmin=0, vmax=2)]
+    fig, axs = plt.subplots(ncols=len(norms), constrained_layout=True)
+    for ax, norm in zip(axs, norms):
+        fig.colorbar(cm.ScalarMappable(norm=norm, cmap='viridis'), cax=ax, extend='both')
