ENH: Adding zoom and pan to colorbar

greglucas · greglucas · commit 5a6c99b8852a · 2021-05-26T17:23:56.000-06:00
The zoom and pan funcitons change the vmin/vmax of the norm
attached to the colorbar. The colorbar is rendered as an inset
axis, but the event handler is implemented on the parent axis.
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,6 @@
+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.
diff --git a/lib/matplotlib/colorbar.py b/lib/matplotlib/colorbar.py
@@ -31,6 +31,7 @@
 import copy
 import logging
 import textwrap
+import types
 
 import numpy as np
 
@@ -1199,6 +1200,11 @@ def __init__(self, ax, mappable, **kwargs):
         mappable.colorbar_cid = mappable.callbacksSM.connect(
             'changed', self.update_normal)
 
+        # Handle interactiveness on the outer axis
+        for x in ["_get_view", "_set_view", "_set_view_from_bbox",
+                  "start_pan", "end_pan", "drag_pan"]:
+            setattr(self.ax.outer_ax, x, getattr(self, x))
+
     @_api.deprecated("3.3", alternative="update_normal")
     def on_mappable_changed(self, mappable):
         """
@@ -1326,6 +1332,242 @@ def remove(self):
             # use_gridspec was True
             ax.set_subplotspec(subplotspec)
 
+    def _get_view(self):
+        """
+        Save information required to reproduce the current view.
+
+        The colorbar is controlled by the norm's vmin and vmax
+        """
+        return self.norm.vmin, self.norm.vmax
+
+    def _set_view(self, view):
+        """
+        Apply a previously saved view.
+
+        This sets the vmin and vmax of the norm associated with
+        the scalar mappable. We then update the normal to redraw
+        the content that was updated.
+        """
+        self.mappable.norm.vmin, self.mappable.norm.vmax = view
+        self.update_normal(self.mappable)
+
+    def _set_view_from_bbox(self, bbox, direction='in',
+                            mode=None, twinx=False, twiny=False):
+        """
+        Update view from a selection bbox.
+
+        The view of a colorbar is controlled by the vmin and vmax of the
+        norm. Updating the view from a bbox will change the vmin and vmax
+        to the extent of the selection when zooming in and scale the values
+        away from the selection when zooming out.
+
+        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.
+        """
+        ax = self.ax
+        if len(bbox) == 3:
+            Xmin, Xmax = ax.get_xlim()
+            Ymin, Ymax = ax.get_ylim()
+
+            xp, yp, scl = bbox  # Zooming code
+
+            if scl == 0:  # Should not happen
+                scl = 1.
+
+            if scl > 1:
+                direction = 'in'
+            else:
+                direction = 'out'
+                scl = 1/scl
+
+            # get the limits of the axes
+            tranD2C = ax.transData.transform
+            xmin, ymin = tranD2C((Xmin, Ymin))
+            xmax, ymax = tranD2C((Xmax, Ymax))
+
+            # set the range
+            xwidth = xmax - xmin
+            ywidth = ymax - ymin
+            xcen = (xmax + xmin)*.5
+            ycen = (ymax + ymin)*.5
+            xzc = (xp*(scl - 1) + xcen)/scl
+            yzc = (yp*(scl - 1) + ycen)/scl
+
+            bbox = [xzc - xwidth/2./scl, yzc - ywidth/2./scl,
+                    xzc + xwidth/2./scl, yzc + ywidth/2./scl]
+        elif len(bbox) != 4:
+            # should be len 3 or 4 but nothing else
+            _api.warn_external(
+                "Warning in _set_view_from_bbox: bounding box is not a tuple "
+                "of length 3 or 4. Ignoring the view change.")
+            return
+
+        # Original limits.
+        xmin0, xmax0 = ax.get_xbound()
+        ymin0, ymax0 = ax.get_ybound()
+        # The zoom box in screen coords.
+        startx, starty, stopx, stopy = bbox
+
+        # Convert to data coords.
+        (startx, starty), (stopx, stopy) = ax.transData.inverted().transform(
+            [(startx, starty), (stopx, stopy)])
+        # Clip to axes limits.
+        xmin, xmax = np.clip(sorted([startx, stopx]), xmin0, xmax0)
+        ymin, ymax = np.clip(sorted([starty, stopy]), ymin0, ymax0)
+        # Don't double-zoom twinned axes or if zooming only the other axis.
+        if twinx or mode == "y":
+            xmin, xmax = xmin0, xmax0
+        if twiny or mode == "x":
+            ymin, ymax = ymin0, ymax0
+
+        if direction == "in":
+            new_xbound = xmin, xmax
+            new_ybound = ymin, ymax
+
+        elif direction == "out":
+            x_trf = ax.xaxis.get_transform()
+            sxmin0, sxmax0, sxmin, sxmax = x_trf.transform(
+                [xmin0, xmax0, xmin, xmax])  # To screen space.
+            factor = (sxmax0 - sxmin0) / (sxmax - sxmin)  # Unzoom factor.
+            # Move original bounds away by
+            # (factor) x (distance between unzoom box and axes bbox).
+            sxmin1 = sxmin0 - factor * (sxmin - sxmin0)
+            sxmax1 = sxmax0 + factor * (sxmax0 - sxmax)
+            # And back to data space.
+            new_xbound = x_trf.inverted().transform([sxmin1, sxmax1])
+
+            y_trf = ax.yaxis.get_transform()
+            symin0, symax0, symin, symax = y_trf.transform(
+                [ymin0, ymax0, ymin, ymax])
+            factor = (symax0 - symin0) / (symax - symin)
+            symin1 = symin0 - factor * (symin - symin0)
+            symax1 = symax0 + factor * (symax0 - symax)
+            new_ybound = y_trf.inverted().transform([symin1, symax1])
+
+        norm = self.mappable.norm
+        if self.orientation == 'horizontal' and not twinx and mode != "y":
+            norm.vmin, norm.vmax = new_xbound
+        if self.orientation == 'vertical' and not twiny and mode != "x":
+            norm.vmin, norm.vmax = new_ybound
+        self.update_normal(self.mappable)
+
+    def start_pan(self, x, y, button):
+        """
+        Called when a pan operation has started.
+
+        Parameters
+        ----------
+        x, y : float
+            The mouse coordinates in display coords.
+        button : `.MouseButton`
+            The pressed mouse button.
+        """
+        self._pan_start = types.SimpleNamespace(
+            lim=self.ax.viewLim.frozen(),
+            trans=self.ax.transData.frozen(),
+            trans_inverse=self.ax.transData.inverted().frozen(),
+            bbox=self.ax.bbox.frozen(),
+            x=x,
+            y=y)
+
+    def end_pan(self):
+        """
+        Called when a pan operation completes (when the mouse button is up.)
+        """
+        del self._pan_start
+
+    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.
+        """
+        def format_deltas(key, dx, dy):
+            if key == 'control':
+                if abs(dx) > abs(dy):
+                    dy = dx
+                else:
+                    dx = dy
+            elif key == 'x':
+                dy = 0
+            elif key == 'y':
+                dx = 0
+            elif key == 'shift':
+                if 2 * abs(dx) < abs(dy):
+                    dx = 0
+                elif 2 * abs(dy) < abs(dx):
+                    dy = 0
+                elif abs(dx) > abs(dy):
+                    dy = dy / abs(dy) * abs(dx)
+                else:
+                    dx = dx / abs(dx) * abs(dy)
+            return dx, dy
+
+        p = self._pan_start
+        dx = x - p.x
+        dy = y - p.y
+        if dx == dy == 0:
+            return
+        if button == 1:
+            dx, dy = format_deltas(key, dx, dy)
+            result = p.bbox.translated(-dx, -dy).transformed(p.trans_inverse)
+        elif button == 3:
+            try:
+                dx = -dx / self.ax.bbox.width
+                dy = -dy / self.ax.bbox.height
+                dx, dy = format_deltas(key, dx, dy)
+                if self.ax.get_aspect() != 'auto':
+                    dx = dy = 0.5 * (dx + dy)
+                alpha = np.power(10.0, (dx, dy))
+                start = np.array([p.x, p.y])
+                oldpoints = p.lim.transformed(p.trans)
+                newpoints = start + alpha * (oldpoints - start)
+                result = (mtransforms.Bbox(newpoints)
+                          .transformed(p.trans_inverse))
+            except OverflowError:
+                _api._warn_external('Overflow while panning')
+                return
+        else:
+            return
+
+        valid = np.isfinite(result.transformed(p.trans))
+        points = result.get_points().astype(object)
+        # Just ignore invalid limits (typically, underflow in log-scale).
+        points[~valid] = None
+
+        norm = self.mappable.norm
+        if self.orientation == 'horizontal':
+            norm.vmin, norm.vmax = points[:, 0]
+        if self.orientation == 'vertical':
+            norm.vmin, norm.vmax = points[:, 1]
+        self.update_normal(self.mappable)
+
 
 def _normalize_location_orientation(location, orientation):
     if location is None:
