Merge pull request #13234 from jklymak/fix-colorbar-norm-changed

NelleV · web-flow · commit ce33de7afedb · 2019-02-02T11:56:19.000-08:00
FIX: allow colorbar mappable norm to change and do right thing
diff --git a/doc/api/next_api_changes/2019-01-27-JMK.rst b/doc/api/next_api_changes/2019-01-27-JMK.rst
@@ -0,0 +1,12 @@
+`matplotlib.colorbar.ColorbarBase` is no longer a subclass of `.ScalarMappable`
+-------------------------------------------------------------------------------
+
+This inheritance lead to a confusing situation where the
+`ScalarMappable` passed to `matplotlib.colorbar.Colorbar` (`~.Figure.colorbar`)
+had a ``set_norm`` method, as did the colorbar.  The colorbar is now purely a
+slave to the `ScalarMappable` norm and colormap, and the old inherited methods
+`~matplotlib.colorbar.ColorbarBase.set_norm`,
+`~matplotlib.colorbar.ColorbarBase.set_cmap`,
+`~matplotlib.colorbar.ColorbarBase.set_clim` are deprecated, as are the
+getter versions of those calls.  To set the norm associated with a colorbar do
+``colorbar.mappable.set_norm()`` etc.
diff --git a/lib/matplotlib/cm.py b/lib/matplotlib/cm.py
@@ -349,6 +349,13 @@ def set_norm(self, norm):
         Parameters
         ----------
         norm : `.Normalize`
+
+        Notes
+        -----
+        If there are any colorbars using the mappable for this norm, setting
+        the norm of the mappable will reset the norm, locator, and formatters
+        on the colorbar to default.
+
         """
         if norm is None:
             norm = colors.Normalize()
diff --git a/lib/matplotlib/colorbar.py b/lib/matplotlib/colorbar.py
@@ -304,7 +304,51 @@ def tick_values(self, vmin, vmax):
         return ticks
 
 
-class ColorbarBase(cm.ScalarMappable):
+class _ColorbarMappableDummy(object):
+    """
+    Private class to hold deprecated ColorbarBase methods that used to be
+    inhereted from ScalarMappable.
+    """
+    @cbook.deprecated("3.1", alternative="ScalarMappable.set_norm")
+    def set_norm(self, norm):
+        """
+        `.colorbar.Colorbar.set_norm` does nothing; set the norm on
+        the mappable associated with this colorbar.
+        """
+        pass
+
+    @cbook.deprecated("3.1", alternative="ScalarMappable.set_cmap")
+    def set_cmap(self, cmap):
+        """
+        `.colorbar.Colorbar.set_cmap` does nothing; set the norm on
+        the mappable associated with this colorbar.
+        """
+        pass
+
+    @cbook.deprecated("3.1", alternative="ScalarMappable.set_clim")
+    def set_clim(self, cmap):
+        """
+        `.colorbar.Colorbar.set_clim` does nothing; set the limits on
+        the mappable associated with this colorbar.
+        """
+        pass
+
+    @cbook.deprecated("3.1", alternative="ScalarMappable.get_cmap")
+    def get_cmap(self):
+        """
+        return the colormap
+        """
+        return self.cmap
+
+    @cbook.deprecated("3.1", alternative="ScalarMappable.get_clim")
+    def get_clim(self):
+        """
+        return the min, max of the color limits for image scaling
+        """
+        return self.norm.vmin, self.norm.vmax
+
+
+class ColorbarBase(_ColorbarMappableDummy):
     '''
     Draw a colorbar in an existing axes.
 
@@ -371,7 +415,8 @@ def __init__(self, ax, cmap=None,
         if norm is None:
             norm = colors.Normalize()
         self.alpha = alpha
-        cm.ScalarMappable.__init__(self, cmap=cmap, norm=norm)
+        self.cmap = cmap
+        self.norm = norm
         self.values = values
         self.boundaries = boundaries
         self.extend = extend
@@ -387,30 +432,26 @@ def __init__(self, ax, cmap=None,
         self.outline = None
         self.patch = None
         self.dividers = None
+        self.locator = None
+        self.formatter = None
         self._manual_tick_data_values = None
 
         if ticklocation == 'auto':
             ticklocation = 'bottom' if orientation == 'horizontal' else 'right'
         self.ticklocation = ticklocation
 
         self.set_label(label)
+        self._reset_locator_formatter_scale()
+
         if np.iterable(ticks):
             self.locator = ticker.FixedLocator(ticks, nbins=len(ticks))
         else:
             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)
-            else:
-                self.formatter = ticker.ScalarFormatter()
-        elif isinstance(format, str):
+
+        if isinstance(format, str):
             self.formatter = ticker.FormatStrFormatter(format)
         else:
-            self.formatter = format  # Assume it is a Formatter
-        # The rest is in a method so we can recalculate when clim changes.
+            self.formatter = format  # Assume it is a Formatter or None
         self.draw_all()
 
     def _extend_lower(self):
@@ -432,7 +473,6 @@ def draw_all(self):
         Calculate any free parameters based on the current cmap and norm,
         and do all the drawing.
         '''
-
         # sets self._boundaries and self._values in real data units.
         # takes into account extend values:
         self._process_values()
@@ -451,12 +491,6 @@ 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':
             long_axis, short_axis = ax.yaxis, ax.xaxis
@@ -504,6 +538,20 @@ def _get_ticker_locator_formatter(self):
             else:
                 b = self._boundaries[self._inside]
                 locator = ticker.FixedLocator(b, nbins=10)
+
+        if formatter is None:
+            if isinstance(self.norm, colors.LogNorm):
+                formatter = ticker.LogFormatterSciNotation()
+            elif isinstance(self.norm, colors.SymLogNorm):
+                formatter = ticker.LogFormatterSciNotation(
+                                        linthresh=self.norm.linthresh)
+            else:
+                formatter = ticker.ScalarFormatter()
+        else:
+            formatter = self.formatter
+
+        self.locator = locator
+        self.formatter = formatter
         _log.debug('locator: %r', locator)
         return locator, formatter
 
@@ -517,6 +565,24 @@ def _use_auto_colorbar_locator(self):
                 and ((type(self.norm) == colors.Normalize)
                     or (type(self.norm) == colors.LogNorm)))
 
+    def _reset_locator_formatter_scale(self):
+        """
+        Reset the locator et al to defaults.  Any user-hardcoded changes
+        need to be re-entered if this gets called (either at init, or when
+        the mappable normal gets changed: Colorbar.update_normal)
+        """
+        self.locator = None
+        self.formatter = None
+        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.
+            self.ax.set_xscale('log')
+            self.ax.set_yscale('log')
+        else:
+            self.ax.set_xscale('linear')
+            self.ax.set_yscale('linear')
+
     def update_ticks(self):
         """
         Force the update of the ticks and ticklabels. This must be
@@ -526,7 +592,6 @@ def update_ticks(self):
         # 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:
@@ -1102,7 +1167,6 @@ def __init__(self, ax, mappable, **kw):
             kw['boundaries'] = CS._levels
             kw['values'] = CS.cvalues
             kw['extend'] = CS.extend
-            #kw['ticks'] = CS._levels
             kw.setdefault('ticks', ticker.FixedLocator(CS.levels, nbins=10))
             kw['filled'] = CS.filled
             ColorbarBase.__init__(self, ax, **kw)
@@ -1125,8 +1189,7 @@ def on_mappable_changed(self, mappable):
         by :func:`colorbar_factory` and should not be called manually.
 
         """
-        self.set_cmap(mappable.get_cmap())
-        self.set_clim(mappable.get_clim())
+        _log.debug('colorbar mappable changed')
         self.update_normal(mappable)
 
     def add_lines(self, CS, erase=True):
@@ -1156,9 +1219,24 @@ def update_normal(self, mappable):
         Update solid patches, lines, etc.
 
         Unlike `.update_bruteforce`, this does not clear the axes.  This is
-        meant to be called when the image or contour plot to which this
-        colorbar belongs changes.
+        meant to be called when the norm of the image or contour plot to which
+        this colorbar belongs changes.
+
+        If the norm on the mappable is different than before, this resets the
+        locator and formatter for the axis, so if these have been customized,
+        they will need to be customized again.  However, if the norm only
+        changes values of *vmin*, *vmax* or *cmap* then the old formatter
+        and locator will be preserved.
         """
+
+        _log.debug('colorbar update normal %r %r', mappable.norm, self.norm)
+        self.mappable = mappable
+        self.set_alpha(mappable.get_alpha())
+        self.cmap = mappable.cmap
+        if mappable.norm != self.norm:
+            self.norm = mappable.norm
+            self._reset_locator_formatter_scale()
+
         self.draw_all()
         if isinstance(self.mappable, contour.ContourSet):
             CS = self.mappable
@@ -1180,15 +1258,16 @@ def update_bruteforce(self, mappable):
         # properties have been changed by methods other than the
         # colorbar methods, those changes will be lost.
         self.ax.cla()
+        self.locator = None
+        self.formatter = None
+
         # clearing the axes will delete outline, patch, solids, and lines:
         self.outline = None
         self.patch = None
         self.solids = None
         self.lines = list()
         self.dividers = None
-        self.set_alpha(mappable.get_alpha())
-        self.cmap = mappable.cmap
-        self.norm = mappable.norm
+        self.update_normal(mappable)
         self.draw_all()
         if isinstance(self.mappable, contour.ContourSet):
             CS = self.mappable
diff --git a/lib/matplotlib/tests/test_colorbar.py b/lib/matplotlib/tests/test_colorbar.py
@@ -4,10 +4,10 @@
 from matplotlib import rc_context
 from matplotlib.testing.decorators import image_comparison
 import matplotlib.pyplot as plt
-from matplotlib.colors import BoundaryNorm, LogNorm, PowerNorm
+from matplotlib.colors import BoundaryNorm, LogNorm, PowerNorm, Normalize
 from matplotlib.cm import get_cmap
-from matplotlib.colorbar import ColorbarBase
-from matplotlib.ticker import LogLocator, LogFormatter
+from matplotlib.colorbar import ColorbarBase, _ColorbarLogLocator
+from matplotlib.ticker import LogLocator, LogFormatter, FixedLocator
 
 
 def _get_cmap_norms():
@@ -442,23 +442,69 @@ def test_colorbar_renorm():
     fig, ax = plt.subplots()
     im = ax.imshow(z)
     cbar = fig.colorbar(im)
+    assert np.allclose(cbar.ax.yaxis.get_majorticklocs(),
+                       np.arange(0, 120000.1, 15000))
+
+    cbar.set_ticks([1, 2, 3])
+    assert isinstance(cbar.locator, FixedLocator)
 
     norm = LogNorm(z.min(), z.max())
     im.set_norm(norm)
-    cbar.set_norm(norm)
-    cbar.locator = LogLocator()
-    cbar.formatter = LogFormatter()
-    cbar.update_normal(im)
+    assert isinstance(cbar.locator, _ColorbarLogLocator)
+    assert np.allclose(cbar.ax.yaxis.get_majorticklocs(),
+                       np.logspace(-8, 5, 14))
+    # note that set_norm removes the FixedLocator...
     assert np.isclose(cbar.vmin, z.min())
+    cbar.set_ticks([1, 2, 3])
+    assert isinstance(cbar.locator, FixedLocator)
+    assert np.allclose(cbar.ax.yaxis.get_majorticklocs(),
+                       [1.0, 2.0, 3.0])
 
     norm = LogNorm(z.min() * 1000, z.max() * 1000)
     im.set_norm(norm)
-    cbar.set_norm(norm)
-    cbar.update_normal(im)
     assert np.isclose(cbar.vmin, z.min() * 1000)
     assert np.isclose(cbar.vmax, z.max() * 1000)
 
 
+def test_colorbar_format():
+    # make sure that format is passed properly
+    x, y = np.ogrid[-4:4:31j, -4:4:31j]
+    z = 120000*np.exp(-x**2 - y**2)
+
+    fig, ax = plt.subplots()
+    im = ax.imshow(z)
+    cbar = fig.colorbar(im, format='%4.2e')
+    fig.canvas.draw()
+    assert cbar.ax.yaxis.get_ticklabels()[4].get_text() == '6.00e+04'
+
+    # make sure that if we change the clim of the mappable that the
+    # formatting is *not* lost:
+    im.set_clim([4, 200])
+    fig.canvas.draw()
+    assert cbar.ax.yaxis.get_ticklabels()[4].get_text() == '8.00e+01'
+
+    # but if we change the norm:
+    im.set_norm(LogNorm(vmin=0.1, vmax=10))
+    fig.canvas.draw()
+    assert (cbar.ax.yaxis.get_ticklabels()[0].get_text() ==
+            r'$\mathdefault{10^{-1}}$')
+
+
+def test_colorbar_scale_reset():
+    x, y = np.ogrid[-4:4:31j, -4:4:31j]
+    z = 120000*np.exp(-x**2 - y**2)
+
+    fig, ax = plt.subplots()
+    pcm = ax.pcolormesh(z, cmap='RdBu_r', rasterized=True)
+    cbar = fig.colorbar(pcm, ax=ax)
+    assert cbar.ax.yaxis.get_scale() == 'linear'
+
+    pcm.set_norm(LogNorm(vmin=1, vmax=100))
+    assert cbar.ax.yaxis.get_scale() == 'log'
+    pcm.set_norm(Normalize(vmin=-20, vmax=20))
+    assert cbar.ax.yaxis.get_scale() == 'linear'
+
+
 def test_colorbar_get_ticks():
     with rc_context({'_internal.classic_mode': False}):
 
diff --git a/tutorials/colors/colorbar_only.py b/tutorials/colors/colorbar_only.py
@@ -8,12 +8,10 @@
 Customized Colorbars
 ====================
 
-:class:`~matplotlib.colorbar.ColorbarBase` derives from
-:mod:`~matplotlib.cm.ScalarMappable` and puts a colorbar in a specified axes,
-so it has everything needed for a standalone colorbar. It can be used as-is to
-make a colorbar for a given colormap; it does not need a mappable object like
-an image. In this tutorial we will explore what can be done with standalone
-colorbar.
+`~matplotlib.colorbar.ColorbarBase` puts a colorbar in a specified axes,
+and can make a colorbar for a given colormap; it does not need a mappable
+object like an image. In this tutorial we will explore what can be done with
+standalone colorbar.
 
 Basic continuous colorbar
 -------------------------
