ENH/FIX: make "auto" aspect with Axes3D match Axes better

tacaswell · tacaswell · commit 5451b25977d4 · 2020-05-26T23:34:16.000-04:00
To make bbox_inches='tight' work correctly, we cache the positions of all of the axes, change their aspect to 'auto', adjust the transforms on the figure, render the output, and then restore the previous setting of aspect to each of the figures. This prevent the Axes from trying to adjust their size on draw. This matters because for the render that is emitted the figure transforms are out of sync with the figure size. As part of cleaning fixing the rendering of Axes3D in #8896 / #16472 we started to use apply_aspect to re-size the area the Axes3D takes up to maintain a fixed ratio between the axes when the aspect is "auto". However this conflicts with the expectation in tight_bbox.adjust_bbox as it assumes setting the aspect to "auto" will prevent any axes resizing. This commit addresses this by: - exiting the Axes3D.apply_aspect early if aspect is auto - adding a new aspect mode 'auto_pb' which is the default for Axes3D which maintains the current master branch behavior of maintaining a fixed ratio between the sizes of the x, y z axis independent of the data limits. - re implement several functions on Axes3D to make sure they handle sharez correctly. closes #16463.
diff --git a/lib/matplotlib/axes/_base.py b/lib/matplotlib/axes/_base.py
@@ -1339,10 +1339,6 @@ def set_aspect(self, aspect, adjustable=None, anchor=None, share=False):
         if cbook._str_equal(aspect, 'equal'):
             aspect = 1
         if not cbook._str_equal(aspect, 'auto'):
-            if self.name == '3d':
-                raise NotImplementedError(
-                    'It is not currently possible to manually set the aspect '
-                    'on 3D axes')
             aspect = float(aspect)  # raise ValueError if necessary
 
         if share:
diff --git a/lib/mpl_toolkits/mplot3d/axes3d.py b/lib/mpl_toolkits/mplot3d/axes3d.py
@@ -24,6 +24,7 @@
 import matplotlib.colors as mcolors
 import matplotlib.docstring as docstring
 import matplotlib.scale as mscale
+import matplotlib.transforms as mtransforms
 from matplotlib.axes import Axes, rcParams
 from matplotlib.axes._base import _axis_method_wrapper
 from matplotlib.transforms import Bbox
@@ -99,6 +100,7 @@ def __init__(
             self._shared_z_axes.join(self, sharez)
             self._adjustable = 'datalim'
 
+        kwargs.setdefault('aspect', 'auto_pb')
         super().__init__(fig, rect, frameon=True, *args, **kwargs)
         # Disable drawing of axes by base class
         super().set_axis_off()
@@ -261,10 +263,99 @@ def tunit_edges(self, vals=None, M=None):
                  (tc[7], tc[4])]
         return edges
 
+    def set_aspect(self, aspect, adjustable=None, anchor=None, share=False):
+        """
+        Set the aspect of the axis scaling.
+
+        Parameters
+        ----------
+        aspect : {'auto_pb', 'auto'}
+            Possible values:
+
+            ========   =================================================
+            value      description
+            ========   =================================================
+            'auto_pb'  This will make the size of the axes have a fixed
+                       ratio
+            'auto'     automatic; fill the position rectangle with data.
+                       This will let the ratio between the size of the
+                       (pseudo-) bounding box be free and will not adjust
+                       anything at draw time.
+            ========   =================================================
+
+        adjustable : None or {'box', 'datalim'}, optional
+            If not ``None``, this defines which parameter will be adjusted to
+            meet the required aspect. See `.set_adjustable` for further
+            details.
+
+            Currently ignored by Axes3D
+
+        anchor : None or str or 2-tuple of float, optional
+            If not ``None``, this defines where the Axes will be drawn if there
+            is extra space due to aspect constraints. The most common way to
+            to specify the anchor are abbreviations of cardinal directions:
+
+            =====   =====================
+            value   description
+            =====   =====================
+            'C'     centered
+            'SW'    lower left corner
+            'S'     middle of bottom edge
+            'SE'    lower right corner
+            etc.
+            =====   =====================
+
+            See `.set_anchor` for further details.
+
+        share : bool, default: False
+            If ``True``, apply the settings to all shared Axes.
+
+        """
+        if aspect not in {'auto', 'auto_pb'}:
+            raise NotImplementedError(
+                "Axes3D currently only support the aspect arguments "
+                "{'auto', 'auto_pb'}. " + f"You passed in {aspect!r}."
+            )
+
+        if share:
+            axes = {*self._shared_x_axes.get_siblings(self),
+                    *self._shared_y_axes.get_siblings(self),
+                    *self._shared_z_axes.get_siblings(self),
+                    }
+        else:
+            axes = {self}
+
+        for ax in axes:
+            ax._aspect = aspect
+            ax.stale = True
+
+        if anchor is not None:
+            self.set_anchor(anchor, share=share)
+
+    def set_anchor(self, anchor, share=False):
+        # docstring inherited
+        if not (anchor in mtransforms.Bbox.coefs or len(anchor) == 2):
+            raise ValueError('argument must be among %s' %
+                             ', '.join(mtransforms.Bbox.coefs))
+        if share:
+            axes = {*self._shared_x_axes.get_siblings(self),
+                    *self._shared_y_axes.get_siblings(self),
+                    *self._shared_z_axes.get_siblings(self),
+                    }
+        else:
+            axes = {self}
+        for ax in axes:
+            ax._anchor = anchor
+            ax.stale = True
+
     def apply_aspect(self, position=None):
         if position is None:
             position = self.get_position(original=True)
 
+        aspect = self.get_aspect()
+        if aspect == 'auto':
+            return
+
         # in the superclass, we would go through and actually deal with axis
         # scales and box/datalim. Those are all irrelevant - all we need to do
         # is make sure our coordinate system is square.
