ENH: add an inset_axes to the axes class

jklymak · jklymak · commit 0abd014b17e7 · 2018-07-18T13:17:35.000-07:00
diff --git a/.flake8 b/.flake8
@@ -253,6 +253,7 @@ per-file-ignores =
     examples/subplots_axes_and_figures/axes_zoom_effect.py: E402
     examples/subplots_axes_and_figures/demo_tight_layout.py: E402
     examples/subplots_axes_and_figures/two_scales.py: E402
+    examples/subplots_axes_and_figures/zoom_inset_axes.py: E402
     examples/tests/backend_driver_sgskip.py: E402, E501
     examples/text_labels_and_annotations/annotation_demo.py: E501
     examples/text_labels_and_annotations/custom_legends.py: E402
diff --git a/doc/api/axes_api.rst b/doc/api/axes_api.rst
@@ -181,6 +181,9 @@ Text and Annotations
    Axes.text
    Axes.table
    Axes.arrow
+   Axes.inset_axes_from_bounds
+   Axes.indicate_inset_bounds
+   Axes.indicate_inset_zoom
 
 
 Fields
diff --git a/examples/subplots_axes_and_figures/zoom_inset_axes.py b/examples/subplots_axes_and_figures/zoom_inset_axes.py
@@ -0,0 +1,61 @@
+"""
+======================
+Zoom region inset axes
+======================
+
+Example of an inset axes and a rectangle showing where the zoom is located.
+
+"""
+
+import matplotlib.pyplot as plt
+import numpy as np
+
+
+def get_demo_image():
+    from matplotlib.cbook import get_sample_data
+    import numpy as np
+    f = get_sample_data("axes_grid/bivariate_normal.npy", asfileobj=False)
+    z = np.load(f)
+    # z is a numpy array of 15x15
+    extent = (-3, 4, -4, 3)
+    return z, extent
+
+fig, ax = plt.subplots(figsize=[5, 4])
+
+# make data
+Z, extent = get_demo_image()
+Z2 = np.zeros([150, 150], dtype="d")
+ny, nx = Z.shape
+Z2[30:30 + ny, 30:30 + nx] = Z
+
+ax.imshow(Z2, extent=extent, interpolation="nearest",
+          origin="lower")
+
+# inset axes....
+axins = ax.inset_axes_from_bounds([0.5, 0.5, 0.47, 0.47])
+axins.imshow(Z2, extent=extent, interpolation="nearest",
+          origin="lower")
+# sub region of the original image
+x1, x2, y1, y2 = -1.5, -0.9, -2.5, -1.9
+axins.set_xlim(x1, x2)
+axins.set_ylim(y1, y2)
+axins.set_xticklabels('')
+axins.set_yticklabels('')
+
+ax.indicate_inset_zoom(axins)
+
+plt.show()
+
+#############################################################################
+#
+# ------------
+#
+# References
+# """"""""""
+#
+# The use of the following functions and methods is shown in this example:
+
+import matplotlib
+matplotlib.axes.Axes.inset_axes_from_bounds
+matplotlib.axes.Axes.indicate_inset_zoom
+matplotlib.axes.Axes.imshow
diff --git a/lib/matplotlib/axes/_axes.py b/lib/matplotlib/axes/_axes.py
@@ -84,9 +84,36 @@ def _plot_args_replacer(args, data):
                          "multiple plotting calls instead.")
 
 
+def _make_inset_locator(rect, trans, parent):
+    """
+    Helper function to locate inset axes, used in
+    `.Axes.inset_axes_from_bounds`.
+
+    A locator gets used in `Axes.set_aspect` to override the default
+    locations...  It is a function that takes an axes object and
+    a renderer and tells `set_aspect` where it is to be placed.
+
+    Here *rect* is a rectangle [l, b, w, h] that specifies the
+    location for the axes in the transform given by *trans* on the
+    *parent*.
+    """
+    _rect = mtransforms.Bbox.from_bounds(*rect)
+    _trans = trans
+    _parent = parent
+
+    def inset_locator(ax, renderer):
+        bb = mtransforms.TransformedBbox(_rect, _trans)
+        tr = _parent.figure.transFigure.inverted()
+        bb = mtransforms.TransformedBbox(bb, tr)
+        return bb
+
+    return inset_locator
+
+
 # The axes module contains all the wrappers to plotting functions.
 # All the other methods should go in the _AxesBase class.
 
+
 class Axes(_AxesBase):
     """
     The :class:`Axes` contains most of the figure elements:
@@ -390,6 +417,227 @@ def legend(self, *args, **kwargs):
     def _remove_legend(self, legend):
         self.legend_ = None
 
+    def inset_axes_from_bounds(self, bounds, *, transform=None, zorder=5,
+            **kwargs):
+        """
+        Add a child inset axes to this existing axes.
+
+        Warnings
+        --------
+
+        This method is experimental as of 3.0, and the API may change.
+
+        Parameters
+        ----------
+
+        bounds : [x0, y0, width, height]
+            Lower-left corner of inset axes, and its width and height.
+
+        transform : `.Transform`
+            Defaults to `ax.transAxes`, i.e. the units of *rect* are in
+            axes-relative coordinates.
+
+        zorder : number
+            Defaults to 5 (same as `.Axes.legend`).  Adjust higher or lower
+            to change whether it is above or below data plotted on the
+            parent axes.
+
+        **kwargs
+
+            Other *kwargs* are passed on to the `axes.Axes` child axes.
+
+        Returns
+        -------
+
+        Axes
+            The created `.axes.Axes` instance.
+
+        Examples
+        --------
+
+        This example makes two inset axes, the first is in axes-relative
+        coordinates, and the second in data-coordinates::
+
+            fig, ax = plt.suplots()
+            ax.plot(range(10))
+            axin1 = ax.inset_axes_from_bounds([0.8, 0.1, 0.15, 0.15])
+            axin2 = ax.inset_axes_from_bounds(
+                    [5, 7, 2.3, 2.3], transform=ax.transData)
+
+        """
+        if transform is None:
+            transform = self.transAxes
+        label = kwargs.pop('label', 'inset_axes_from_bounds')
+
+        # This puts the rectangle into figure-relative coordinates.
+        inset_locator = _make_inset_locator(bounds, transform, self)
+        bb = inset_locator(None, None)
+
+        inset_ax = Axes(self.figure, bb.bounds, zorder=zorder,
+                label=label, **kwargs)
+
+        # this locator lets the axes move if in data coordinates.
+        # it gets called in `ax.apply_aspect() (of all places)
+        inset_ax.set_axes_locator(inset_locator)
+
+        self.add_child_axes(inset_ax)
+
+        return inset_ax
+
+    def indicate_inset_bounds(self, bounds, inset_ax=None, *, transform=None,
+            facecolor='none', edgecolor='0.5', alpha=0.5,
+            zorder=4.99, **kwargs):
+        """
+        Add an inset indicator to the axes.  This is a rectangle on the plot
+        at the position indicated by *bounds* that optionally has lines that
+        connect the rectangle to an inset axes
+        (`.Axes.inset_axes_from_bounds`).
+
+        Warnings
+        --------
+
+        This method is experimental as of 3.0, and the API may change.
+
+
+        Parameters
+        ----------
+
+        bounds : [x0, y0, width, height]
+            Lower-left corner of rectangle to be marked, and its width
+            and height.
+
+        inset_ax : `.Axes`
+            An optional inset axes to draw connecting lines to.  Two lines are
+            drawn connecting the indicator box to the inset axes on corners
+            chosen so as to not overlap with the indicator box.
+
+        transform : `.Transform`
+            Transform for the rectangle co-ordinates. Defaults to
+            `ax.transAxes`, i.e. the units of *rect* are in axes-relative
+            coordinates.
+
+        facecolor : Matplotlib color
+            Facecolor of the rectangle (default 'none').
+
+        edgecolor : Matplotlib color
+            Color of the rectangle and color of the connecting lines.  Default
+            is '0.5'.
+
+        alpha : number
+            Transparency of the rectangle and connector lines.  Default is 0.5.
+
+        zorder : number
+            Drawing order of the rectangle and connector lines. Default is 4.99
+            (just below the default level of inset axes).
+
+        **kwargs
+            Other *kwargs* are passed on to the rectangle patch.
+
+        Returns
+        -------
+
+        rectangle_patch: `.Patches.Rectangle`
+             Rectangle artist.
+
+        connector_lines: 4-tuple of `.Patches.ConnectionPatch`
+            One for each of four connector lines.  Two are set with visibility
+            to *False*,  but the user can set the visibility to True if the
+            automatic choice is not deemed correct.
+
+        """
+
+        # to make the axes connectors work, we need to apply the aspect to
+        # the parent axes.
+        self.apply_aspect()
+
+        if transform is None:
+            transform = self.transData
+        label = kwargs.pop('label', 'indicate_inset_bounds')
+
+        xy = (bounds[0], bounds[1])
+        rectpatch = mpatches.Rectangle(xy, bounds[2], bounds[3],
+                facecolor=facecolor, edgecolor=edgecolor, alpha=alpha,
+                zorder=zorder,  label=label, transform=transform, **kwargs)
+        self.add_patch(rectpatch)
+
+        if inset_ax is not None:
+            # want to connect the indicator to the rect....
+
+            pos = inset_ax.get_position()  # this is in fig-fraction.
+            coordsA = 'axes fraction'
+            connects = []
+            xr = [bounds[0], bounds[0]+bounds[2]]
+            yr = [bounds[1], bounds[1]+bounds[3]]
+            for xc in range(2):
+                for yc in range(2):
+                    xyA = (xc, yc)
+                    xyB = (xr[xc], yr[yc])
+                    connects += [mpatches.ConnectionPatch(xyA, xyB,
+                            'axes fraction', 'data',
+                            axesA=inset_ax, axesB=self, arrowstyle="-",
+                            zorder=zorder, edgecolor=edgecolor, alpha=alpha)]
+                    self.add_patch(connects[-1])
+            # decide which two of the lines to keep visible....
+            pos = inset_ax.get_position()
+            bboxins = pos.transformed(self.figure.transFigure)
+            rectbbox = mtransforms.Bbox.from_bounds(
+                        *bounds).transformed(transform)
+            if rectbbox.x0 < bboxins.x0:
+                sig = 1
+            else:
+                sig = -1
+            if sig*rectbbox.y0 < sig*bboxins.y0:
+                connects[0].set_visible(False)
+                connects[3].set_visible(False)
+            else:
+                connects[1].set_visible(False)
+                connects[2].set_visible(False)
+
+        return rectpatch, connects
+
+    def indicate_inset_zoom(self, inset_ax, **kwargs):
+        """
+        Add an inset indicator rectangle to the axes based on the axis
+        limits for an *inset_ax* and draw connectors between *inset_ax*
+        and the rectangle.
+
+        Warnings
+        --------
+
+        This method is experimental as of 3.0, and the API may change.
+
+        Parameters
+        ----------
+
+        inset_ax : `.Axes`
+            Inset axes to draw connecting lines to.  Two lines are
+            drawn connecting the indicator box to the inset axes on corners
+            chosen so as to not overlap with the indicator box.
+
+        **kwargs
+            Other *kwargs* are passed on to `.Axes.inset_rectangle`
+
+        Returns
+        -------
+
+        rectangle_patch: `.Patches.Rectangle`
+             Rectangle artist.
+
+        connector_lines: 4-tuple of `.Patches.ConnectionPatch`
+            One for each of four connector lines.  Two are set with visibility
+            to *False*,  but the user can set the visibility to True if the
+            automatic choice is not deemed correct.
+
+        """
+
+        xlim = inset_ax.get_xlim()
+        ylim = inset_ax.get_ylim()
+        rect = [xlim[0], ylim[0], xlim[1] - xlim[0], ylim[1] - ylim[0]]
+        rectpatch, connects = self.indicate_inset_bounds(
+                rect, inset_ax, **kwargs)
+
+        return rectpatch, connects
+
     def text(self, x, y, s, fontdict=None, withdash=False, **kwargs):
         """
         Add text to the axes.
diff --git a/lib/matplotlib/axes/_base.py b/lib/matplotlib/axes/_base.py
@@ -1031,6 +1031,7 @@ def cla(self):
         self.artists = []
         self.images = []
         self._mouseover_set = _OrderedSet()
+        self.child_axes = []
         self._current_image = None  # strictly for pyplot via _sci, _gci
         self.legend_ = None
         self.collections = []  # collection.Collection instances
@@ -1807,6 +1808,27 @@ def add_artist(self, a):
         self.stale = True
         return a
 
+    def add_child_axes(self, ax):
+        """
+        Add a :class:`~matplotlib.axes.Axesbase` instance
+        as a child to the axes.
+
+        Returns the added axes.
+
+        This is the lowlevel version.  See `.axes.Axes.inset_axes`
+        """
+
+        # normally axes have themselves as the axes, but these need to have
+        # their parent...
+        # Need to bypass the getter...
+        ax._axes = self
+        ax.stale_callback = martist._stale_axes_callback
+
+        self.child_axes.append(ax)
+        ax._remove_method = self.child_axes.remove
+        self.stale = True
+        return ax
+
     def add_collection(self, collection, autolim=True):
         """
         Add a :class:`~matplotlib.collections.Collection` instance
@@ -4073,9 +4095,12 @@ def get_children(self):
         children.append(self._right_title)
         children.extend(self.tables)
         children.extend(self.images)
+        children.extend(self.child_axes)
+
         if self.legend_ is not None:
             children.append(self.legend_)
         children.append(self.patch)
+
         return children
 
     def contains(self, mouseevent):
diff --git a/lib/matplotlib/tests/test_axes.py b/lib/matplotlib/tests/test_axes.py
