Merge pull request #19515 from greglucas/colorbar-axes-zoom

jklymak · web-flow · commit b9193bae0499 · 2021-09-21T10:49:34.000+02:00
Colorbar axis zoom and pan
diff --git a/doc/api/next_api_changes/behavior/19515-GL.rst b/doc/api/next_api_changes/behavior/19515-GL.rst
@@ -0,0 +1,10 @@
+Colorbars now have pan and zoom functionality
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+Interactive plots with colorbars can now be zoomed and panned on
+the colorbar axis. This adjusts the *vmin* and *vmax* of the
+``ScalarMappable`` associated with the colorbar. This is currently
+only enabled for continuous norms. Norms used with contourf and
+categoricals, such as ``BoundaryNorm`` and ``NoNorm``, have the
+interactive capability disabled by default. ``cb.ax.set_navigate()``
+can be used to set whether a colorbar axes is interactive or not.
diff --git a/lib/matplotlib/axes/_base.py b/lib/matplotlib/axes/_base.py
@@ -4255,41 +4255,14 @@ def _set_view(self, view):
         self.set_xlim((xmin, xmax))
         self.set_ylim((ymin, ymax))
 
-    def _set_view_from_bbox(self, bbox, direction='in',
-                            mode=None, twinx=False, twiny=False):
+    def _prepare_view_from_bbox(self, bbox, direction='in',
+                                mode=None, twinx=False, twiny=False):
         """
-        Update view from a selection bbox.
-
-        .. note::
-
-            Intended to be overridden by new projection types, but if not, the
-            default implementation sets the view limits to the bbox directly.
-
-        Parameters
-        ----------
-        bbox : 4-tuple or 3 tuple
-            * If bbox is a 4 tuple, it is the selected bounding box limits,
-              in *display* coordinates.
-            * If bbox is a 3 tuple, it is an (xp, yp, scl) triple, where
-              (xp, yp) is the center of zooming and scl the scale factor to
-              zoom by.
+        Helper function to prepare the new bounds from a bbox.
 
-        direction : str
-            The direction to apply the bounding box.
-                * `'in'` - The bounding box describes the view directly, i.e.,
-                           it zooms in.
-                * `'out'` - The bounding box describes the size to make the
-                            existing view, i.e., it zooms out.
-
-        mode : str or None
-            The selection mode, whether to apply the bounding box in only the
-            `'x'` direction, `'y'` direction or both (`None`).
-
-        twinx : bool
-            Whether this axis is twinned in the *x*-direction.
-
-        twiny : bool
-            Whether this axis is twinned in the *y*-direction.
+        This helper function returns the new x and y bounds from the zoom
+        bbox. This a convenience method to abstract the bbox logic
+        out of the base setter.
         """
         if len(bbox) == 3:
             xp, yp, scl = bbox  # Zooming code
@@ -4360,6 +4333,46 @@ def _set_view_from_bbox(self, bbox, direction='in',
             symax1 = symax0 + factor * (symax0 - symax)
             new_ybound = y_trf.inverted().transform([symin1, symax1])
 
+        return new_xbound, new_ybound
+
+    def _set_view_from_bbox(self, bbox, direction='in',
+                            mode=None, twinx=False, twiny=False):
+        """
+        Update view from a selection bbox.
+
+        .. note::
+
+            Intended to be overridden by new projection types, but if not, the
+            default implementation sets the view limits to the bbox directly.
+
+        Parameters
+        ----------
+        bbox : 4-tuple or 3 tuple
+            * If bbox is a 4 tuple, it is the selected bounding box limits,
+              in *display* coordinates.
+            * If bbox is a 3 tuple, it is an (xp, yp, scl) triple, where
+              (xp, yp) is the center of zooming and scl the scale factor to
+              zoom by.
+
+        direction : str
+            The direction to apply the bounding box.
+                * `'in'` - The bounding box describes the view directly, i.e.,
+                           it zooms in.
+                * `'out'` - The bounding box describes the size to make the
+                            existing view, i.e., it zooms out.
+
+        mode : str or None
+            The selection mode, whether to apply the bounding box in only the
+            `'x'` direction, `'y'` direction or both (`None`).
+
+        twinx : bool
+            Whether this axis is twinned in the *x*-direction.
+
+        twiny : bool
+            Whether this axis is twinned in the *y*-direction.
+        """
+        new_xbound, new_ybound = self._prepare_view_from_bbox(
+            bbox, direction=direction, mode=mode, twinx=twinx, twiny=twiny)
         if not twinx and mode != "y":
             self.set_xbound(new_xbound)
             self.set_autoscalex_on(False)
@@ -4400,22 +4413,13 @@ def end_pan(self):
         """
         del self._pan_start
 
-    def drag_pan(self, button, key, x, y):
+    def _get_pan_points(self, button, key, x, y):
         """
-        Called when the mouse moves during a pan operation.
+        Helper function to return the new points after a pan.
 
-        Parameters
-        ----------
-        button : `.MouseButton`
-            The pressed mouse button.
-        key : str or None
-            The pressed key, if any.
-        x, y : float
-            The mouse coordinates in display coords.
-
-        Notes
-        -----
-        This is intended to be overridden by new projection types.
+        This helper function returns the points on the axis after a pan has
+        occurred. This is a convenience method to abstract the pan logic
+        out of the base setter.
         """
         def format_deltas(key, dx, dy):
             if key == 'control':
@@ -4469,8 +4473,29 @@ def format_deltas(key, dx, dy):
         points = result.get_points().astype(object)
         # Just ignore invalid limits (typically, underflow in log-scale).
         points[~valid] = None
-        self.set_xlim(points[:, 0])
-        self.set_ylim(points[:, 1])
+        return points
+
+    def drag_pan(self, button, key, x, y):
+        """
+        Called when the mouse moves during a pan operation.
+
+        Parameters
+        ----------
+        button : `.MouseButton`
+            The pressed mouse button.
+        key : str or None
+            The pressed key, if any.
+        x, y : float
+            The mouse coordinates in display coords.
+
+        Notes
+        -----
+        This is intended to be overridden by new projection types.
+        """
+        points = self._get_pan_points(button, key, x, y)
+        if points is not None:
+            self.set_xlim(points[:, 0])
+            self.set_ylim(points[:, 1])
 
     def get_children(self):
         # docstring inherited.
diff --git a/lib/matplotlib/backend_bases.py b/lib/matplotlib/backend_bases.py
@@ -3136,7 +3136,7 @@ def zoom(self, *args):
             a.set_navigate_mode(self.mode._navigate_mode)
         self.set_message(self.mode)
 
-    _ZoomInfo = namedtuple("_ZoomInfo", "direction start_xy axes cid")
+    _ZoomInfo = namedtuple("_ZoomInfo", "direction start_xy axes cid cbar")
 
     def press_zoom(self, event):
         """Callback for mouse button press in zoom to rect mode."""
@@ -3151,20 +3151,34 @@ def press_zoom(self, event):
             self.push_current()  # set the home button to this view
         id_zoom = self.canvas.mpl_connect(
             "motion_notify_event", self.drag_zoom)
+        # A colorbar is one-dimensional, so we extend the zoom rectangle out
+        # to the edge of the axes bbox in the other dimension. To do that we
+        # store the orientation of the colorbar for later.
+        if hasattr(axes[0], "_colorbar"):
+            cbar = axes[0]._colorbar.orientation
+        else:
+            cbar = None
         self._zoom_info = self._ZoomInfo(
             direction="in" if event.button == 1 else "out",
-            start_xy=(event.x, event.y), axes=axes, cid=id_zoom)
+            start_xy=(event.x, event.y), axes=axes, cid=id_zoom, cbar=cbar)
 
     def drag_zoom(self, event):
         """Callback for dragging in zoom mode."""
         start_xy = self._zoom_info.start_xy
         ax = self._zoom_info.axes[0]
         (x1, y1), (x2, y2) = np.clip(
             [start_xy, [event.x, event.y]], ax.bbox.min, ax.bbox.max)
-        if event.key == "x":
+        key = event.key
+        # Force the key on colorbars to extend the short-axis bbox
+        if self._zoom_info.cbar == "horizontal":
+            key = "x"
+        elif self._zoom_info.cbar == "vertical":
+            key = "y"
+        if key == "x":
             y1, y2 = ax.bbox.intervaly
-        elif event.key == "y":
+        elif key == "y":
             x1, x2 = ax.bbox.intervalx
+
         self.draw_rubberband(event, x1, y1, x2, y2)
 
     def release_zoom(self, event):
@@ -3178,10 +3192,17 @@ def release_zoom(self, event):
         self.remove_rubberband()
 
         start_x, start_y = self._zoom_info.start_xy
+        key = event.key
+        # Force the key on colorbars to ignore the zoom-cancel on the
+        # short-axis side
+        if self._zoom_info.cbar == "horizontal":
+            key = "x"
+        elif self._zoom_info.cbar == "vertical":
+            key = "y"
         # Ignore single clicks: 5 pixels is a threshold that allows the user to
         # "cancel" a zoom action by zooming by less than 5 pixels.
-        if ((abs(event.x - start_x) < 5 and event.key != "y")
-                or (abs(event.y - start_y) < 5 and event.key != "x")):
+        if ((abs(event.x - start_x) < 5 and key != "y") or
+                (abs(event.y - start_y) < 5 and key != "x")):
             self.canvas.draw_idle()
             self._zoom_info = None
             return
@@ -3195,7 +3216,7 @@ def release_zoom(self, event):
                         for prev in self._zoom_info.axes[:i])
             ax._set_view_from_bbox(
                 (start_x, start_y, event.x, event.y),
-                self._zoom_info.direction, event.key, twinx, twiny)
+                self._zoom_info.direction, key, twinx, twiny)
 
         self.canvas.draw_idle()
         self._zoom_info = None
diff --git a/lib/matplotlib/colorbar.py b/lib/matplotlib/colorbar.py
@@ -422,7 +422,6 @@ def __init__(self, ax, mappable=None, *, cmap=None,
 
         self.ax = ax
         self.ax._axes_locator = _ColorbarAxesLocator(self)
-        ax.set(navigate=False)
 
         if extend is None:
             if (not isinstance(mappable, contour.ContourSet)
@@ -496,6 +495,29 @@ def __init__(self, ax, mappable=None, *, cmap=None,
         if isinstance(mappable, contour.ContourSet) and not mappable.filled:
             self.add_lines(mappable)
 
+        # Link the Axes and Colorbar for interactive use
+        self.ax._colorbar = self
+        # Don't navigate on any of these types of mappables
+        if (isinstance(self.norm, (colors.BoundaryNorm, colors.NoNorm)) or
+                isinstance(self.mappable, contour.ContourSet)):
+            self.ax.set_navigate(False)
+
+        # These are the functions that set up interactivity on this colorbar
+        self._interactive_funcs = ["_get_view", "_set_view",
+                                   "_set_view_from_bbox", "drag_pan"]
+        for x in self._interactive_funcs:
+            setattr(self.ax, x, getattr(self, x))
+        # Set the cla function to the cbar's method to override it
+        self.ax.cla = self._cbar_cla
+
+    def _cbar_cla(self):
+        """Function to clear the interactive colorbar state."""
+        for x in self._interactive_funcs:
+            delattr(self.ax, x)
+        # We now restore the old cla() back and can call it directly
+        del self.ax.cla
+        self.ax.cla()
+
     # Also remove ._patch after deprecation elapses.
     patch = _api.deprecate_privatize_attribute("3.5", alternative="ax")
 
@@ -1276,6 +1298,36 @@ def _short_axis(self):
             return self.ax.xaxis
         return self.ax.yaxis
 
+    def _get_view(self):
+        # docstring inherited
+        # An interactive view for a colorbar is the norm's vmin/vmax
+        return self.norm.vmin, self.norm.vmax
+
+    def _set_view(self, view):
+        # docstring inherited
+        # An interactive view for a colorbar is the norm's vmin/vmax
+        self.norm.vmin, self.norm.vmax = view
+
+    def _set_view_from_bbox(self, bbox, direction='in',
+                            mode=None, twinx=False, twiny=False):
+        # docstring inherited
+        # For colorbars, we use the zoom bbox to scale the norm's vmin/vmax
+        new_xbound, new_ybound = self.ax._prepare_view_from_bbox(
+            bbox, direction=direction, mode=mode, twinx=twinx, twiny=twiny)
+        if self.orientation == 'horizontal':
+            self.norm.vmin, self.norm.vmax = new_xbound
+        elif self.orientation == 'vertical':
+            self.norm.vmin, self.norm.vmax = new_ybound
+
+    def drag_pan(self, button, key, x, y):
+        # docstring inherited
+        points = self.ax._get_pan_points(button, key, x, y)
+        if points is not None:
+            if self.orientation == 'horizontal':
+                self.norm.vmin, self.norm.vmax = points[:, 0]
+            elif self.orientation == 'vertical':
+                self.norm.vmin, self.norm.vmax = points[:, 1]
+
 
 ColorbarBase = Colorbar  # Backcompat API
 
diff --git a/lib/matplotlib/tests/test_backend_bases.py b/lib/matplotlib/tests/test_backend_bases.py
@@ -181,6 +181,66 @@ def test_interactive_zoom():
     assert not ax.get_autoscalex_on() and not ax.get_autoscaley_on()
 
 
+@pytest.mark.parametrize("plot_func", ["imshow", "contourf"])
+@pytest.mark.parametrize("orientation", ["vertical", "horizontal"])
+@pytest.mark.parametrize("tool,button,expected",
+                         [("zoom", MouseButton.LEFT, (4, 6)),  # zoom in
+                          ("zoom", MouseButton.RIGHT, (-20, 30)),  # zoom out
+                          ("pan", MouseButton.LEFT, (-2, 8))])
+def test_interactive_colorbar(plot_func, orientation, tool, button, expected):
+    fig, ax = plt.subplots()
+    data = np.arange(12).reshape((4, 3))
+    vmin0, vmax0 = 0, 10
+    coll = getattr(ax, plot_func)(data, vmin=vmin0, vmax=vmax0)
+
+    cb = fig.colorbar(coll, ax=ax, orientation=orientation)
+    if plot_func == "contourf":
+        # Just determine we can't navigate and exit out of the test
+        assert not cb.ax.get_navigate()
+        return
+
+    assert cb.ax.get_navigate()
+
+    # Mouse from 4 to 6 (data coordinates, "d").
+    vmin, vmax = 4, 6
+    # The y coordinate doesn't matter, it just needs to be between 0 and 1
+    # However, we will set d0/d1 to the same y coordinate to test that small
+    # pixel changes in that coordinate doesn't cancel the zoom like a normal
+    # axes would.
+    d0 = (vmin, 0.5)
+    d1 = (vmax, 0.5)
+    # Swap them if the orientation is vertical
+    if orientation == "vertical":
+        d0 = d0[::-1]
+        d1 = d1[::-1]
+    # Convert to screen coordinates ("s").  Events are defined only with pixel
+    # precision, so round the pixel values, and below, check against the
+    # corresponding xdata/ydata, which are close but not equal to d0/d1.
+    s0 = cb.ax.transData.transform(d0).astype(int)
+    s1 = cb.ax.transData.transform(d1).astype(int)
+
+    # Set up the mouse movements
+    start_event = MouseEvent(
+        "button_press_event", fig.canvas, *s0, button)
+    stop_event = MouseEvent(
+        "button_release_event", fig.canvas, *s1, button)
+
+    tb = NavigationToolbar2(fig.canvas)
+    if tool == "zoom":
+        tb.zoom()
+        tb.press_zoom(start_event)
+        tb.drag_zoom(stop_event)
+        tb.release_zoom(stop_event)
+    else:
+        tb.pan()
+        tb.press_pan(start_event)
+        tb.drag_pan(stop_event)
+        tb.release_pan(stop_event)
+
+    # Should be close, but won't be exact due to screen integer resolution
+    assert (cb.vmin, cb.vmax) == pytest.approx(expected, abs=0.15)
+
+
 def test_toolbar_zoompan():
     expected_warning_regex = (
         r"Treat the new Tool classes introduced in "
