Merge pull request #19553 from greglucas/norm-updater

tacaswell · web-flow · commit 5707b3095217 · 2021-09-01T16:04:06.000-04:00
ENH: Adding callbacks to Norms for update signals
diff --git a/doc/users/next_whats_new/callbacks_on_norms.rst b/doc/users/next_whats_new/callbacks_on_norms.rst
@@ -0,0 +1,8 @@
+A callback registry has been added to Normalize objects
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+`.colors.Normalize` objects now have a callback registry, ``callbacks``,
+that can be connected to by other objects to be notified when the norm is
+updated. The callback emits the key ``changed`` when the norm is modified.
+`.cm.ScalarMappable` is now a listener and will register a change
+when the norm's vmin, vmax or other attributes are changed.
diff --git a/lib/matplotlib/cm.py b/lib/matplotlib/cm.py
@@ -337,7 +337,7 @@ def __init__(self, norm=None, cmap=None):
             The colormap used to map normalized data values to RGBA colors.
         """
         self._A = None
-        self.norm = None  # So that the setter knows we're initializing.
+        self._norm = None  # So that the setter knows we're initializing.
         self.set_norm(norm)  # The Normalize instance of this ScalarMappable.
         self.cmap = None  # So that the setter knows we're initializing.
         self.set_cmap(cmap)  # The Colormap instance of this ScalarMappable.
@@ -496,6 +496,8 @@ def set_clim(self, vmin=None, vmax=None):
 
              .. ACCEPTS: (vmin: float, vmax: float)
         """
+        # If the norm's limits are updated self.changed() will be called
+        # through the callbacks attached to the norm
         if vmax is None:
             try:
                 vmin, vmax = vmin
@@ -505,7 +507,6 @@ def set_clim(self, vmin=None, vmax=None):
             self.norm.vmin = colors._sanitize_extrema(vmin)
         if vmax is not None:
             self.norm.vmax = colors._sanitize_extrema(vmax)
-        self.changed()
 
     def get_alpha(self):
         """
@@ -531,6 +532,30 @@ def set_cmap(self, cmap):
         if not in_init:
             self.changed()  # Things are not set up properly yet.
 
+    @property
+    def norm(self):
+        return self._norm
+
+    @norm.setter
+    def norm(self, norm):
+        _api.check_isinstance((colors.Normalize, None), norm=norm)
+        if norm is None:
+            norm = colors.Normalize()
+
+        if norm is self.norm:
+            # We aren't updating anything
+            return
+
+        in_init = self.norm is None
+        # Remove the current callback and connect to the new one
+        if not in_init:
+            self.norm.callbacks.disconnect(self._id_norm)
+        self._norm = norm
+        self._id_norm = self.norm.callbacks.connect('changed',
+                                                    self.changed)
+        if not in_init:
+            self.changed()
+
     def set_norm(self, norm):
         """
         Set the normalization instance.
@@ -545,13 +570,7 @@ def set_norm(self, norm):
         the norm of the mappable will reset the norm, locator, and formatters
         on the colorbar to default.
         """
-        _api.check_isinstance((colors.Normalize, None), norm=norm)
-        in_init = self.norm is None
-        if norm is None:
-            norm = colors.Normalize()
         self.norm = norm
-        if not in_init:
-            self.changed()  # Things are not set up properly yet.
 
     def autoscale(self):
         """
@@ -560,8 +579,9 @@ def autoscale(self):
         """
         if self._A is None:
             raise TypeError('You must first set_array for mappable')
+        # If the norm's limits are updated self.changed() will be called
+        # through the callbacks attached to the norm
         self.norm.autoscale(self._A)
-        self.changed()
 
     def autoscale_None(self):
         """
@@ -570,8 +590,9 @@ def autoscale_None(self):
         """
         if self._A is None:
             raise TypeError('You must first set_array for mappable')
+        # If the norm's limits are updated self.changed() will be called
+        # through the callbacks attached to the norm
         self.norm.autoscale_None(self._A)
-        self.changed()
 
     def changed(self):
         """
diff --git a/lib/matplotlib/colorbar.py b/lib/matplotlib/colorbar.py
@@ -471,6 +471,7 @@ def __init__(self, ax, mappable=None, *, cmap=None,
         self.ax.add_collection(self.dividers)
 
         self.locator = None
+        self.minorlocator = None
         self.formatter = None
         self.__scale = None  # linear, log10 for now.  Hopefully more?
 
@@ -1096,7 +1097,7 @@ def _mesh(self):
         # vmax of the colorbar, not the norm.  This allows the situation
         # where the colormap has a narrower range than the colorbar, to
         # accommodate extra contours:
-        norm = copy.copy(self.norm)
+        norm = copy.deepcopy(self.norm)
         norm.vmin = self.vmin
         norm.vmax = self.vmax
         x = np.array([0.0, 1.0])
diff --git a/lib/matplotlib/colors.py b/lib/matplotlib/colors.py
@@ -1123,10 +1123,50 @@ def __init__(self, vmin=None, vmax=None, clip=False):
         -----
         Returns 0 if ``vmin == vmax``.
         """
-        self.vmin = _sanitize_extrema(vmin)
-        self.vmax = _sanitize_extrema(vmax)
-        self.clip = clip
-        self._scale = None  # will default to LinearScale for colorbar
+        self._vmin = _sanitize_extrema(vmin)
+        self._vmax = _sanitize_extrema(vmax)
+        self._clip = clip
+        self._scale = None
+        self.callbacks = cbook.CallbackRegistry()
+
+    @property
+    def vmin(self):
+        return self._vmin
+
+    @vmin.setter
+    def vmin(self, value):
+        value = _sanitize_extrema(value)
+        if value != self._vmin:
+            self._vmin = value
+            self._changed()
+
+    @property
+    def vmax(self):
+        return self._vmax
+
+    @vmax.setter
+    def vmax(self, value):
+        value = _sanitize_extrema(value)
+        if value != self._vmax:
+            self._vmax = value
+            self._changed()
+
+    @property
+    def clip(self):
+        return self._clip
+
+    @clip.setter
+    def clip(self, value):
+        if value != self._clip:
+            self._clip = value
+            self._changed()
+
+    def _changed(self):
+        """
+        Call this whenever the norm is changed to notify all the
+        callback listeners to the 'changed' signal.
+        """
+        self.callbacks.process('changed')
 
     @staticmethod
     def process_value(value):
@@ -1273,14 +1313,24 @@ def __init__(self, vcenter, vmin=None, vmax=None):
         """
 
         super().__init__(vmin=vmin, vmax=vmax)
-        self.vcenter = vcenter
+        self._vcenter = vcenter
         if vcenter is not None and vmax is not None and vcenter >= vmax:
             raise ValueError('vmin, vcenter, and vmax must be in '
                              'ascending order')
         if vcenter is not None and vmin is not None and vcenter <= vmin:
             raise ValueError('vmin, vcenter, and vmax must be in '
                              'ascending order')
 
+    @property
+    def vcenter(self):
+        return self._vcenter
+
+    @vcenter.setter
+    def vcenter(self, value):
+        if value != self._vcenter:
+            self._vcenter = value
+            self._changed()
+
     def autoscale_None(self, A):
         """
         Get vmin and vmax, and then clip at vcenter
@@ -1387,7 +1437,9 @@ def vcenter(self):
 
     @vcenter.setter
     def vcenter(self, vcenter):
-        self._vcenter = vcenter
+        if vcenter != self._vcenter:
+            self._vcenter = vcenter
+            self._changed()
         if self.vmax is not None:
             # recompute halfrange assuming vmin and vmax represent
             # min and max of data
diff --git a/lib/matplotlib/contour.py b/lib/matplotlib/contour.py
@@ -1090,6 +1090,15 @@ def _make_paths(self, segs, kinds):
                     in zip(segs, kinds)]
 
     def changed(self):
+        if not hasattr(self, "cvalues"):
+            # Just return after calling the super() changed function
+            cm.ScalarMappable.changed(self)
+            return
+        # Force an autoscale immediately because self.to_rgba() calls
+        # autoscale_None() internally with the data passed to it,
+        # so if vmin/vmax are not set yet, this would override them with
+        # content from *cvalues* rather than levels like we want
+        self.norm.autoscale_None(self.levels)
         tcolors = [(tuple(rgba),)
                    for rgba in self.to_rgba(self.cvalues, alpha=self.alpha)]
         self.tcolors = tcolors
diff --git a/lib/matplotlib/image.py b/lib/matplotlib/image.py
@@ -537,11 +537,14 @@ def _make_image(self, A, in_bbox, out_bbox, clip_bbox, magnification=1.0,
                 if isinstance(self.norm, mcolors.LogNorm) and s_vmin <= 0:
                     # Don't give 0 or negative values to LogNorm
                     s_vmin = np.finfo(scaled_dtype).eps
-                with cbook._setattr_cm(self.norm,
-                                       vmin=s_vmin,
-                                       vmax=s_vmax,
-                                       ):
-                    output = self.norm(resampled_masked)
+                # Block the norm from sending an update signal during the
+                # temporary vmin/vmax change
+                with self.norm.callbacks.blocked():
+                    with cbook._setattr_cm(self.norm,
+                                           vmin=s_vmin,
+                                           vmax=s_vmax,
+                                           ):
+                        output = self.norm(resampled_masked)
             else:
                 if A.ndim == 2:  # _interpolation_stage == 'rgba'
                     self.norm.autoscale_None(A)
diff --git a/lib/matplotlib/tests/test_colors.py b/lib/matplotlib/tests/test_colors.py
@@ -1,5 +1,6 @@
 import copy
 import itertools
+import unittest.mock
 
 from io import BytesIO
 import numpy as np
@@ -17,7 +18,7 @@
 import matplotlib.cbook as cbook
 import matplotlib.pyplot as plt
 import matplotlib.scale as mscale
-from matplotlib.testing.decorators import image_comparison
+from matplotlib.testing.decorators import image_comparison, check_figures_equal
 
 
 @pytest.mark.parametrize('N, result', [
@@ -1408,3 +1409,69 @@ def test_norm_deepcopy():
     norm2 = copy.deepcopy(norm)
     assert norm2._scale is None
     assert norm2.vmin == norm.vmin
+
+
+def test_norm_callback():
+    increment = unittest.mock.Mock(return_value=None)
+
+    norm = mcolors.Normalize()
+    norm.callbacks.connect('changed', increment)
+    # Haven't updated anything, so call count should be 0
+    assert increment.call_count == 0
+
+    # Now change vmin and vmax to test callbacks
+    norm.vmin = 1
+    assert increment.call_count == 1
+    norm.vmax = 5
+    assert increment.call_count == 2
+    # callback shouldn't be called if setting to the same value
+    norm.vmin = 1
+    assert increment.call_count == 2
+    norm.vmax = 5
+    assert increment.call_count == 2
+
+
+def test_scalarmappable_norm_update():
+    norm = mcolors.Normalize()
+    sm = matplotlib.cm.ScalarMappable(norm=norm, cmap='plasma')
+    # sm doesn't have a stale attribute at first, set it to False
+    sm.stale = False
+    # The mappable should be stale after updating vmin/vmax
+    norm.vmin = 5
+    assert sm.stale
+    sm.stale = False
+    norm.vmax = 5
+    assert sm.stale
+    sm.stale = False
+    norm.clip = True
+    assert sm.stale
+    # change to the CenteredNorm and TwoSlopeNorm to test those
+    # Also make sure that updating the norm directly and with
+    # set_norm both update the Norm callback
+    norm = mcolors.CenteredNorm()
+    sm.norm = norm
+    sm.stale = False
+    norm.vcenter = 1
+    assert sm.stale
+    norm = mcolors.TwoSlopeNorm(vcenter=0, vmin=-1, vmax=1)
+    sm.set_norm(norm)
+    sm.stale = False
+    norm.vcenter = 1
+    assert sm.stale
+
+
+@check_figures_equal()
+def test_norm_update_figs(fig_test, fig_ref):
+    ax_ref = fig_ref.add_subplot()
+    ax_test = fig_test.add_subplot()
+
+    z = np.arange(100).reshape((10, 10))
+    ax_ref.imshow(z, norm=mcolors.Normalize(10, 90))
+
+    # Create the norm beforehand with different limits and then update
+    # after adding to the plot
+    norm = mcolors.Normalize(0, 1)
+    ax_test.imshow(z, norm=norm)
+    # Force initial draw to make sure it isn't already stale
+    fig_test.canvas.draw()
+    norm.vmin, norm.vmax = 10, 90
diff --git a/lib/matplotlib/tests/test_image.py b/lib/matplotlib/tests/test_image.py
@@ -1017,8 +1017,8 @@ def test_imshow_bool():
 def test_full_invalid():
     fig, ax = plt.subplots()
     ax.imshow(np.full((10, 10), np.nan))
-    with pytest.warns(UserWarning):
-        fig.canvas.draw()
+
+    fig.canvas.draw()
 
 
 @pytest.mark.parametrize("fmt,counted",
