ENH: Adding callbacks to Norms for update signals

greglucas · greglucas · commit 9c82e3d04793 · 2021-05-28T07:50:33.000-06:00
This adds a callback registry to Norm instances that can be
connected to by other objects to be notified when the Norm
is updated. This is particularly relevant for ScalarMappables
to be notified when the vmin/vmax are changed on the Norm.

Quadcontourset overrides ScalarMappable's `changed()` function,
which meant that autoscaling would get called with the wrong data
too early. Therefore, we force an autoscale with the proper data
earlier in the `Quadcontourset.changed()` function.

The Quadcontourset.changed() method assumes some attributes to
be there. If we called changed from a parent class, the object
may not have been initialized with those attributes yet, so
skip that portion of the update.

Remove the ScalarMappable callback from axes_grid as the
state isn't fully cleared when updating the axes.
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,7 @@
+A callback registry has been added to Normalize objects
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+`.colors.Normalize` objects now have a callback registry, ``callbacksNorm``,
+that can be connected to by other objects to be notified when the norm is
+updated. `.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
@@ -257,7 +257,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.
@@ -415,6 +415,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
@@ -424,7 +426,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):
         """
@@ -450,6 +451,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.callbacksNorm.disconnect(self._id_norm)
+        self._norm = norm
+        self._id_norm = self.norm.callbacksNorm.connect('changed',
+                                                        self.changed)
+        if not in_init:
+            self.changed()
+
     def set_norm(self, norm):
         """
         Set the normalization instance.
@@ -464,13 +489,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):
         """
@@ -479,8 +498,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):
         """
@@ -489,8 +509,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
@@ -961,7 +961,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
@@ -1149,10 +1149,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._vmin = _sanitize_extrema(vmin)
+        self._vmax = _sanitize_extrema(vmax)
+        self._clip = clip
         self._scale = scale.LinearScale(axis=None)
+        self.callbacksNorm = 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.callbacksNorm.process('changed')
 
     @staticmethod
     def process_value(value):
@@ -1298,16 +1338,25 @@ def __init__(self, vcenter, vmin=None, vmax=None):
             array([0., 0.25, 0.5, 0.625, 0.75, 0.875, 1.0])
         """
 
-        self.vcenter = vcenter
-        self.vmin = vmin
-        self.vmax = vmax
+        super().__init__(vmin=vmin, vmax=vmax)
+        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
@@ -1369,6 +1418,7 @@ def __init__(self, vcenter=0, halfrange=None, clip=False):
             >>> norm(data)
             array([0.25, 0.5 , 1.  ])
         """
+        super().__init__()
         self._vcenter = vcenter
         self.vmin = None
         self.vmax = None
@@ -1404,7 +1454,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
@@ -1712,9 +1764,7 @@ def __init__(self, boundaries, ncolors, clip=False, *, extend='neither'):
         """
         if clip and extend != 'neither':
             raise ValueError("'clip=True' is not compatible with 'extend'")
-        self.clip = clip
-        self.vmin = boundaries[0]
-        self.vmax = boundaries[-1]
+        super().__init__(vmin=boundaries[0], vmax=boundaries[-1], clip=clip)
         self.boundaries = np.asarray(boundaries)
         self.N = len(self.boundaries)
         if self.N < 2:
diff --git a/lib/matplotlib/contour.py b/lib/matplotlib/contour.py
@@ -1092,6 +1092,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/tests/test_colors.py b/lib/matplotlib/tests/test_colors.py
@@ -17,7 +17,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', [
@@ -1409,3 +1409,73 @@ def test_norm_deepcopy():
     assert isinstance(norm2._scale, mscale.LinearScale)
     assert norm2.vmin == norm.vmin
     assert norm2._scale is not norm._scale
+
+
+def test_norm_callback():
+    norm = mcolors.Normalize()
+    x = 0
+
+    def increment():
+        nonlocal x
+        x += 1
+
+    norm.callbacksNorm.connect('changed', increment)
+    # Haven't updated anything, so x should be 0
+    assert x == 0
+
+    # Now change vmin and vmax to test callbacks
+    norm.vmin = 1
+    assert x == 1
+    norm.vmax = 5
+    assert x == 2
+    # callback shouldn't be called if setting to the same value
+    norm.vmin = 1
+    assert x == 2
+    norm.vmax = 5
+    assert x == 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",
diff --git a/lib/mpl_toolkits/axes_grid1/axes_grid.py b/lib/mpl_toolkits/axes_grid1/axes_grid.py
@@ -32,6 +32,10 @@ def colorbar(self, mappable, *, ticks=None, **kwargs):
         kwargs['userax'] = False
         cb = mpl.colorbar.Colorbar(
             self, mappable, orientation=orientation, ticks=ticks, **kwargs)
+        # If this gets called multiple times on the same axes, some state can
+        # be left on the axes between calls, so for now we will just remove
+        # the callback for all axes_grid colorbars
+        mappable.callbacksSM.disconnect(mappable.colorbar_cid)
         self._config_axes()
         return cb
 
diff --git a/lib/mpl_toolkits/tests/test_axes_grid.py b/lib/mpl_toolkits/tests/test_axes_grid.py
@@ -44,7 +44,7 @@ def test_imagegrid_cbar_mode_edge():
         # "second" ones.  To achieve this, clear out the axes first.
         for ax in grid:
             ax.cax.cla()
-            cb = ax.cax.colorbar(ax.images[0])
+            ax.cax.colorbar(ax.images[0])
 
 
 def test_imagegrid():
