ENH: inset_axes anchored-artist API

jklymak · jklymak · commit 5207b60a61a7 · 2018-07-21T16:08:07.000-07:00
diff --git a/examples/subplots_axes_and_figures/zoom_inset_axes.py b/examples/subplots_axes_and_figures/zoom_inset_axes.py
@@ -16,17 +16,18 @@ def get_demo_image():
     import numpy as np
     f = get_sample_data("axes_grid/bivariate_normal.npy", asfileobj=False)
     z = np.load(f)
+    Z2 = np.zeros([150, 150], dtype="d")
+    ny, nx = z.shape
+    Z2[30:30 + ny, 30:30 + nx] = z
+
     # z is a numpy array of 15x15
-    return z, (-3, 4, -4, 3)
+    extent = (-3, 4, -4, 3)
+    return Z2, 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
-
+Z2, extent = get_demo_image()
 ax.imshow(Z2, extent=extent, interpolation="nearest",
           origin="lower")
 
@@ -38,13 +39,61 @@ def get_demo_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('')
+axins.set_xticklabels(''); axins.set_yticklabels('')
 
 ax.indicate_inset_zoom(axins)
+fig.canvas.draw()
+plt.show()
+
+#############################################################################
+# There is a second interface that closely parallels the interface for
+# `~.axes.legend` whereby we specify a location for the inset axes using
+# a string code.
+
+fig, ax = plt.subplots(figsize=[5, 4])
+
+ax.imshow(Z2, extent=extent, interpolation="nearest",
+          origin="lower")
 
+# inset axes....
+axins = ax.inset_axes('NE', width=0.5, height=0.5)
+
+axins.imshow(Z2, extent=extent, interpolation="nearest",
+          origin="lower")
+# sub region of the original image
+axins.set_xlim(x1, x2)
+axins.set_ylim(y1, y2)
+axins.set_xticklabels('') ;axins.set_yticklabels('')
+
+ax.indicate_inset_zoom(axins)
+fig.canvas.draw()
 plt.show()
 
+#############################################################################
+# Its possible to use either form with a transform in data space instead of
+# in the axes-relative co-ordinates:
+
+fig, ax = plt.subplots(figsize=[5, 4])
+
+ax.imshow(Z2, extent=extent, interpolation="nearest",
+          origin="lower")
+
+# inset axes....
+axins = ax.inset_axes([-2.5, 0, 1.6, 1.6], transform=ax.transData)
+
+axins.imshow(Z2, extent=extent, interpolation="nearest",
+          origin="lower")
+# sub region of the original image
+axins.set_xlim(x1, x2)
+axins.set_ylim(y1, y2)
+axins.set_xticklabels(''); axins.set_yticklabels('')
+
+ax.indicate_inset_zoom(axins)
+fig.canvas.draw()
+plt.show()
+
+
+
 #############################################################################
 #
 # ------------
diff --git a/lib/matplotlib/axes/_axes.py b/lib/matplotlib/axes/_axes.py
@@ -110,6 +110,77 @@ def inset_locator(ax, renderer):
 
     return inset_locator
 
+def _make_inset_locator_anchored(loc, borderpad, width, height,
+                                bbox_to_anchor, transform, parent):
+    """
+    Helper function to locate inset axes, used in
+    `.Axes.inset_axes`.
+
+    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*.
+    """
+    codes = {'upper right':  'NE',
+             'upper left':   'NW',
+             'lower left':   'SW',
+             'lower right':  'SE',
+             'right':        'E',
+             'left':         'W',
+             'bottom':       'S',
+             'top':          'N',
+             'center left':  'W',
+             'center right': 'E',
+             'lower center': 'S',
+             'upper center': 'N',
+             }
+
+    if loc in codes:
+        loc = codes[loc]
+    if loc not in ['N', 'S', 'E', 'W', 'NE', 'NW', 'SE', 'SW']:
+        warnings.warn(f'inset_axes location "{loc}" not recognized; '
+                       'Set to "NE".')
+        loc = 'NE'
+    _loc = loc
+    _parent = parent
+    _borderpadder = borderpad
+    if _borderpadder is None:
+        _borderpadder = _parent.xaxis.get_ticklabels()[0].get_size() * 0.5
+
+    _basebox = mtransforms.Bbox.from_bounds(0, 0, width, height)
+
+    _transform = transform
+    if _transform is None:
+        _transform = BboxTransformTo(self.parent.bbox)
+
+    _bbox_to_anchor = bbox_to_anchor
+    if not isinstance(_bbox_to_anchor, mtransforms.BboxBase):
+        try:
+            l = len(_bbox_to_anchor)
+        except TypeError:
+            raise ValueError("Invalid argument for bbox : %s" % str(bbox))
+        if l == 2:
+            _bbox_to_anchor = [_bbox_to_anchor[0], _bbox_to_anchor[1],
+                                width, height]
+        _bbox_to_anchor = mtransforms.Bbox.from_bounds(*_bbox_to_anchor)
+        _bbox_to_anchor = mtransforms.TransformedBbox(_bbox_to_anchor,
+                                           _transform)
+
+    def inset_locator(ax, renderer):
+        bbox = mtransforms.TransformedBbox(_basebox, _transform)
+        borderpad = renderer.points_to_pixels(_borderpadder)
+        anchored_box = bbox.anchored(loc,
+                container=_bbox_to_anchor.padded(-borderpad))
+        tr = _parent.figure.transFigure.inverted()
+        anchored_box = mtransforms.TransformedBbox(anchored_box, tr)
+
+        return anchored_box
+
+    return inset_locator
+
 
 # The axes module contains all the wrappers to plotting functions.
 # All the other methods should go in the _AxesBase class.
@@ -419,7 +490,8 @@ def _remove_legend(self, legend):
         self.legend_ = None
 
     def inset_axes(self, bounds, *, transform=None, zorder=5,
-            **kwargs):
+            borderaxespad=None, width=None, height=None,
+            bbox_to_anchor=None, **kwargs):
         """
         Add a child inset axes to this existing axes.
 
@@ -431,13 +503,28 @@ def inset_axes(self, bounds, *, transform=None, zorder=5,
         Parameters
         ----------
 
-        bounds : [x0, y0, width, height]
-            Lower-left corner of inset axes, and its width and height.
+        bounds : [x0, y0, width, height] or string.
+            If four-tupple: lower-left corner of inset axes, and its width and
+            height.
+
+            If a string, then locations such as "NE", "N", "NW", "W", etc,
+            of "upper right", "top", "upper left", etc (see `~.axes.legend`)
+            for codes.  (Note we do *not* support the numerical codes).
 
         transform : `.Transform`
             Defaults to `ax.transAxes`, i.e. the units of *rect* are in
             axes-relative coordinates.
 
+        width, height : number
+            width and height of the inset axes.  Only used if ``bounds`` is
+            a string. Units are set by ``transform``, and default to
+            axes-relative co-ordinates.
+
+        borderaxespad : number
+            If ``bounds`` is a string, this is the padding between the inset
+            axes and the parent axes in points.  Defaults to half the fontsize
+            of the tick labels.
+
         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
@@ -470,9 +557,22 @@ def inset_axes(self, bounds, *, transform=None, zorder=5,
             transform = self.transAxes
         label = kwargs.pop('label', 'inset_axes')
 
-        # This puts the rectangle into figure-relative coordinates.
-        inset_locator = _make_inset_locator(bounds, transform, self)
-        bb = inset_locator(None, None)
+        if isinstance(bounds, str):
+            # i.e. NE, S, etc
+            if width is None:
+                width = 0.25
+            if height is None:
+                height = 0.25
+            if bbox_to_anchor is None:
+                bbox_to_anchor = self.bbox
+            inset_locator = _make_inset_locator_anchored(bounds,
+                    borderaxespad, width, height, bbox_to_anchor,
+                    transform, self)
+        else:
+            # This puts the rectangle into figure-relative coordinates.
+            inset_locator = _make_inset_locator(bounds, transform, self)
+
+        bb = inset_locator(None, self.figure.canvas.get_renderer())
 
         inset_ax = Axes(self.figure, bb.bounds, zorder=zorder,
                 label=label, **kwargs)
@@ -549,6 +649,7 @@ def indicate_inset(self, bounds, inset_ax=None, *, transform=None,
 
         # to make the axes connectors work, we need to apply the aspect to
         # the parent axes.
+        
         self.apply_aspect()
 
         if transform is None:
@@ -563,7 +664,7 @@ def indicate_inset(self, bounds, inset_ax=None, *, transform=None,
 
         if inset_ax is not None:
             # want to connect the indicator to the rect....
-
+            inset_ax.apply_aspect()
             pos = inset_ax.get_position()  # this is in fig-fraction.
             coordsA = 'axes fraction'
             connects = []
diff --git a/lib/matplotlib/tests/test_axes.py b/lib/matplotlib/tests/test_axes.py
@@ -5782,3 +5782,46 @@ def test_zoom_inset():
                    [0.8425, 0.907692]])
     np.testing.assert_allclose(axin1.get_position().get_points(),
             xx, rtol=1e-4)
+
+
+def test_inset_codes():
+    "Test that the inset codes put the inset where we want..."
+    fig, ax = plt.subplots()
+    poss = [[[0.415625, 0.686111],
+             [0.609375, 0.886111]],
+            [[0.695833, 0.686111],
+             [0.889583, 0.886111]],
+            [[0.695833, 0.4     ],
+             [0.889583, 0.6     ]],
+            [[0.695833, 0.113889],
+             [0.889583, 0.313889]],
+            [[0.415625, 0.113889],
+             [0.609375, 0.313889]],
+            [[0.135417, 0.113889],
+             [0.329167, 0.313889]],
+            [[0.135417, 0.4     ],
+             [0.329167, 0.6     ]],
+            [[0.135417, 0.686111],
+             [0.329167, 0.886111]]]
+    codes = ['N', 'NE', 'E', 'SE', 'S', 'SW', 'W', 'NW']
+    for pos, code in zip(poss, codes):
+        axin1 = ax.inset_axes(code)
+        np.testing.assert_allclose(axin1.get_position().get_points(),
+                pos, rtol=1e-4)
+        del axin1
+
+    # test synonyms
+    syns   = ['top', 'upper right', 'center right', 'lower right', 'bottom',
+             'lower left', 'center left', 'upper left']
+    codes = ['N', 'NE', 'E', 'SE', 'S', 'SW', 'W', 'NW']
+    for syn, code in zip(syns, codes):
+        axin1 = ax.inset_axes(code)
+        axin2 = ax.inset_axes(syn)
+        np.testing.assert_allclose(axin1.get_position().get_points(),
+                axin2.get_position().get_points(), rtol=1e-4)
+
+    # test the borderaxespad
+    axin1 = ax.inset_axes('NE', borderaxespad=20)
+    pos = [[0.671528, 0.653704], [0.865278, 0.853704]]
+    np.testing.assert_allclose(axin1.get_position().get_points(),
+                pos, rtol=1e-4)
