Merge pull request #14917 from ImportanceOfBeingErnest/box-aspect

jklymak · web-flow · commit 805cb72c0d1f · 2019-10-01T09:44:49.000-07:00
ENH: box aspect for axes
diff --git a/doc/api/axes_api.rst b/doc/api/axes_api.rst
@@ -366,6 +366,9 @@ Aspect ratio
    Axes.set_aspect
    Axes.get_aspect
 
+   Axes.set_box_aspect
+   Axes.get_box_aspect
+
    Axes.set_adjustable
    Axes.get_adjustable
 
diff --git a/doc/users/next_whats_new/2019-07-31_axes-box-aspect.rst b/doc/users/next_whats_new/2019-07-31_axes-box-aspect.rst
@@ -0,0 +1,14 @@
+:orphan:
+
+Setting axes box aspect
+-----------------------
+
+It is now possible to set the aspect of an axes box directly via
+`~.Axes.set_box_aspect`. The box aspect is the ratio between axes height
+and axes width in physical units, independent of the data limits.
+This is useful to e.g. produce a square plot, independent of the data it
+contains, or to have a usual plot with the same axes dimensions next to
+an image plot with fixed (data-)aspect.
+
+For use cases check out the :doc:`Axes box aspect
+</gallery/subplots_axes_and_figures/axes_box_aspect>` example.
diff --git a/examples/subplots_axes_and_figures/axes_box_aspect.py b/examples/subplots_axes_and_figures/axes_box_aspect.py
@@ -0,0 +1,157 @@
+"""
+===============
+Axes box aspect
+===============
+
+This demo shows how to set the aspect of an axes box directly via
+`~.Axes.set_box_aspect`. The box aspect is the ratio between axes height
+and axes width in physical units, independent of the data limits.
+This is useful to e.g. produce a square plot, independent of the data it
+contains, or to have a usual plot with the same axes dimensions next to
+an image plot with fixed (data-)aspect.
+
+The following lists a few use cases for `~.Axes.set_box_aspect`.
+"""
+
+############################################################################
+# A square axes, independent of data
+# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+#
+# Produce a square axes, no matter what the data limits are.
+
+import matplotlib
+import numpy as np
+import matplotlib.pyplot as plt
+
+fig1, ax = plt.subplots()
+
+ax.set_xlim(300, 400)
+ax.set_box_aspect(1)
+
+plt.show()
+
+############################################################################
+# Shared square axes
+# ~~~~~~~~~~~~~~~~~~
+#
+# Produce shared subplots that are squared in size.
+#
+fig2, (ax, ax2) = plt.subplots(ncols=2, sharey=True)
+
+ax.plot([1, 5], [0, 10])
+ax2.plot([100, 500], [10, 15])
+
+ax.set_box_aspect(1)
+ax2.set_box_aspect(1)
+
+plt.show()
+
+############################################################################
+# Square twin axes
+# ~~~~~~~~~~~~~~~~
+#
+# Produce a square axes, with a twin axes. The twinned axes takes over the
+# box aspect of the parent.
+#
+
+fig3, ax = plt.subplots()
+
+ax2 = ax.twinx()
+
+ax.plot([0, 10])
+ax2.plot([12, 10])
+
+ax.set_box_aspect(1)
+
+plt.show()
+
+
+############################################################################
+# Normal plot next to image
+# ~~~~~~~~~~~~~~~~~~~~~~~~~
+#
+# When creating an image plot with fixed data aspect and the default
+# ``adjustable="box"`` next to a normal plot, the axes would be unequal in
+# height. `~.Axes.set_box_aspect` provides an easy solution to that by allowing
+# to have the normal plot's axes use the images dimensions as box aspect.
+#
+# This example also shows that ``constrained_layout`` interplays nicely with
+# a fixed box aspect.
+
+fig4, (ax, ax2) = plt.subplots(ncols=2, constrained_layout=True)
+
+im = np.random.rand(16, 27)
+ax.imshow(im)
+
+ax2.plot([23, 45])
+ax2.set_box_aspect(im.shape[0]/im.shape[1])
+
+plt.show()
+
+############################################################################
+# Square joint/marginal plot
+# ~~~~~~~~~~~~~~~~~~~~~~~~~~
+#
+# It may be desireable to show marginal distributions next to a plot of joint
+# data. The following creates a square plot with the box aspect of the
+# marginal axes being equal to the width- and height-ratios of the gridspec.
+# This ensures that all axes align perfectly, independent on the size of the
+# figure.
+
+fig5, axs = plt.subplots(2, 2, sharex="col", sharey="row",
+                        gridspec_kw=dict(height_ratios=[1, 3],
+                                         width_ratios=[3, 1]))
+axs[0, 1].set_visible(False)
+axs[0, 0].set_box_aspect(1/3)
+axs[1, 0].set_box_aspect(1)
+axs[1, 1].set_box_aspect(3/1)
+
+x, y = np.random.randn(2, 400) * np.array([[.5], [180]])
+axs[1, 0].scatter(x, y)
+axs[0, 0].hist(x)
+axs[1, 1].hist(y, orientation="horizontal")
+
+plt.show()
+
+############################################################################
+# Square joint/marginal plot
+# ~~~~~~~~~~~~~~~~~~~~~~~~~~
+#
+# When setting the box aspect, one may still set the data aspect as well.
+# Here we create an axes with a box twice as long as tall and use an "equal"
+# data aspect for its contents, i.e. the circle actually stays circular.
+
+fig6, ax = plt.subplots()
+
+ax.add_patch(plt.Circle((5, 3), 1))
+ax.set_aspect("equal", adjustable="datalim")
+ax.set_box_aspect(0.5)
+ax.autoscale()
+
+plt.show()
+
+############################################################################
+# Box aspect for many subplots
+# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+#
+# It is possible to pass the box aspect to an axes at initialization. The
+# following creates a 2 by 3 subplot grid with all square axes.
+
+fig7, axs = plt.subplots(2, 3, subplot_kw=dict(box_aspect=1),
+                        sharex=True, sharey=True, constrained_layout=True)
+
+for i, ax in enumerate(axs.flat):
+    ax.scatter(i % 3, -((i // 3) - 0.5)*200, c=[plt.cm.hsv(i / 6)], s=300)
+plt.show()
+
+#############################################################################
+#
+# ------------
+#
+# References
+# """"""""""
+#
+# The use of the following functions, methods and classes is shown
+# in this example:
+
+matplotlib.axes.Axes.set_box_aspect
diff --git a/lib/matplotlib/axes/_base.py b/lib/matplotlib/axes/_base.py
@@ -384,6 +384,7 @@ def __init__(self, fig, rect,
                  label='',
                  xscale=None,
                  yscale=None,
+                 box_aspect=None,
                  **kwargs
                  ):
         """
@@ -405,6 +406,10 @@ def __init__(self, fig, rect,
         frameon : bool, optional
             True means that the axes frame is visible.
 
+        box_aspect : None, or a number, optional
+            Sets the aspect of the axes box. See `~.axes.Axes.set_box_aspect`
+            for details.
+
         **kwargs
             Other optional keyword arguments:
 
@@ -438,7 +443,7 @@ def __init__(self, fig, rect,
             self._shared_y_axes.join(self, sharey)
         self.set_label(label)
         self.set_figure(fig)
-
+        self.set_box_aspect(box_aspect)
         self.set_axes_locator(kwargs.get("axes_locator", None))
 
         self.spines = self._gen_axes_spines()
@@ -1281,6 +1286,18 @@ def set_aspect(self, aspect, adjustable=None, anchor=None, share=False):
         self.stale = True
 
     def get_adjustable(self):
+        """
+        Returns the adjustable parameter, *{'box', 'datalim'}* that defines
+        which parameter the Axes will change to achieve a given aspect.
+
+        See Also
+        --------
+        matplotlib.axes.Axes.set_adjustable
+            defining the parameter to adjust in order to meet the required
+            aspect.
+        matplotlib.axes.Axes.set_aspect
+            for a description of aspect handling.
+        """
         return self._adjustable
 
     def set_adjustable(self, adjustable, share=False):
@@ -1332,6 +1349,55 @@ def set_adjustable(self, adjustable, share=False):
             ax._adjustable = adjustable
         self.stale = True
 
+    def get_box_aspect(self):
+        """
+        Get the axes box aspect.
+        Will be ``None`` if not explicitely specified.
+
+        See Also
+        --------
+        matplotlib.axes.Axes.set_box_aspect
+            for a description of box aspect.
+        matplotlib.axes.Axes.set_aspect
+            for a description of aspect handling.
+        """
+        return self._box_aspect
+
+    def set_box_aspect(self, aspect=None):
+        """
+        Set the axes box aspect. The box aspect is the ratio of the
+        axes height to the axes width in physical units. This is not to be
+        confused with the data aspect, set via `~Axes.set_aspect`.
+
+        Parameters
+        ----------
+        aspect : None, or a number
+            Changes the physical dimensions of the Axes, such that the ratio
+            of the axes height to the axes width in physical units is equal to
+            *aspect*. If *None*, the axes geometry will not be adjusted.
+
+        Note that calling this function with a number changes the *adjustable*
+        to *datalim*.
+
+        See Also
+        --------
+        matplotlib.axes.Axes.set_aspect
+            for a description of aspect handling.
+        """
+        axs = {*self._twinned_axes.get_siblings(self),
+               *self._twinned_axes.get_siblings(self)}
+
+        if aspect is not None:
+            aspect = float(aspect)
+            # when box_aspect is set to other than ´None`,
+            # adjustable must be "datalim"
+            for ax in axs:
+                ax.set_adjustable("datalim")
+
+        for ax in axs:
+            ax._box_aspect = aspect
+            ax.stale = True
+
     def get_anchor(self):
         """
         Get the anchor location.
@@ -1461,7 +1527,7 @@ def apply_aspect(self, position=None):
 
         aspect = self.get_aspect()
 
-        if aspect == 'auto':
+        if aspect == 'auto' and self._box_aspect is None:
             self._set_position(position, which='active')
             return
 
@@ -1481,11 +1547,20 @@ def apply_aspect(self, position=None):
             self._set_position(pb1.anchored(self.get_anchor(), pb), 'active')
             return
 
-        # self._adjustable == 'datalim'
+        # The following is only seen if self._adjustable == 'datalim'
+        if self._box_aspect is not None:
+            pb = position.frozen()
+            pb1 = pb.shrunk_to_aspect(self._box_aspect, pb, fig_aspect)
+            self._set_position(pb1.anchored(self.get_anchor(), pb), 'active')
+            if aspect == "auto":
+                return
 
         # reset active to original in case it had been changed by prior use
         # of 'box'
-        self._set_position(position, which='active')
+        if self._box_aspect is None:
+            self._set_position(position, which='active')
+        else:
+            position = pb1.anchored(self.get_anchor(), pb)
 
         x_trf = self.xaxis.get_transform()
         y_trf = self.yaxis.get_transform()
diff --git a/lib/matplotlib/tests/test_axes.py b/lib/matplotlib/tests/test_axes.py
@@ -6594,5 +6594,70 @@ def test_aspect_nonlinear_adjustable_datalim():
            aspect=1, adjustable="datalim")
     ax.margins(0)
     ax.apply_aspect()
+
     assert ax.get_xlim() == pytest.approx([1*10**(1/2), 100/10**(1/2)])
     assert ax.get_ylim() == (1 / 101, 1 / 11)
+
+
+def test_box_aspect():
+    # Test if axes with box_aspect=1 has same dimensions
+    # as axes with aspect equal and adjustable="box"
+
+    fig1, ax1 = plt.subplots()
+    axtwin = ax1.twinx()
+    axtwin.plot([12, 344])
+
+    ax1.set_box_aspect(1)
+
+    fig2, ax2 = plt.subplots()
+    ax2.margins(0)
+    ax2.plot([0, 2], [6, 8])
+    ax2.set_aspect("equal", adjustable="box")
+
+    fig1.canvas.draw()
+    fig2.canvas.draw()
+
+    bb1 = ax1.get_position()
+    bbt = axtwin.get_position()
+    bb2 = ax2.get_position()
+
+    assert_array_equal(bb1.extents, bb2.extents)
+    assert_array_equal(bbt.extents, bb2.extents)
+
+
+def test_box_aspect_custom_position():
+    # Test if axes with custom position and box_aspect
+    # behaves the same independent of the order of setting those.
+
+    fig1, ax1 = plt.subplots()
+    ax1.set_position([0.1, 0.1, 0.9, 0.2])
+    fig1.canvas.draw()
+    ax1.set_box_aspect(1.)
+
+    fig2, ax2 = plt.subplots()
+    ax2.set_box_aspect(1.)
+    fig2.canvas.draw()
+    ax2.set_position([0.1, 0.1, 0.9, 0.2])
+
+    fig1.canvas.draw()
+    fig2.canvas.draw()
+
+    bb1 = ax1.get_position()
+    bb2 = ax2.get_position()
+
+    assert_array_equal(bb1.extents, bb2.extents)
+
+
+def test_bbox_aspect_axes_init():
+    # Test that box_aspect can be given to axes init and produces
+    # all equal square axes.
+    fig, axs = plt.subplots(2, 3, subplot_kw=dict(box_aspect=1),
+                        constrained_layout=True)
+    fig.canvas.draw()
+    renderer = fig.canvas.get_renderer()
+    sizes = []
+    for ax in axs.flat:
+        bb = ax.get_window_extent(renderer)
+        sizes.extend([bb.width, bb.height])
+
+    assert_allclose(sizes, sizes[0])
