WIP: Add support for Axes3D.set_aspect

eric-wieser · Apr 22, 2020 · 0803457 · 0803457
1 parent 5922edd
commit 0803457
Show file tree

Hide file tree

Showing 2 changed files with 88 additions and 29 deletions.
diff --git a/lib/matplotlib/axes/_base.py b/lib/matplotlib/axes/_base.py
@@ -1268,7 +1268,7 @@ def set_aspect(self, aspect, adjustable=None, anchor=None):
                 and aspect in ('equal', 'auto')):
             self._aspect = aspect
         else:
-            self._aspect = float(aspect)  # raise ValueError if necessary
+            self._aspect = np(aspect)  # raise ValueError if necessary
 
         if adjustable is not None:
             self.set_adjustable(adjustable)
@@ -1532,11 +1532,16 @@ def axis(self, *v, **kwargs):
             The axis limits
 
         """
+        _axes = kwargs.pop('_axes', 'xy')
+
+        def _get_lims():
+            lims = ()
+            for ax in _axes:
+                lims += getattr(self, 'get_{}lim'.format(ax))()
+            return lims
 
         if len(v) == 0 and len(kwargs) == 0:
-            xmin, xmax = self.get_xlim()
-            ymin, ymax = self.get_ylim()
-            return xmin, xmax, ymin, ymax
+            return _get_lims()
 
         emit = kwargs.get('emit', True)
 
@@ -1584,27 +1589,23 @@ def axis(self, *v, **kwargs):
         try:
             v[0]
         except IndexError:
-            xmin = kwargs.get('xmin', None)
-            xmax = kwargs.get('xmax', None)
-            auto = False  # turn off autoscaling, unless...
-            if xmin is None and xmax is None:
-                auto = None  # leave autoscaling state alone
-            xmin, xmax = self.set_xlim(xmin, xmax, emit=emit, auto=auto)
-
-            ymin = kwargs.get('ymin', None)
-            ymax = kwargs.get('ymax', None)
-            auto = False  # turn off autoscaling, unless...
-            if ymin is None and ymax is None:
-                auto = None  # leave autoscaling state alone
-            ymin, ymax = self.set_ylim(ymin, ymax, emit=emit, auto=auto)
-            return xmin, xmax, ymin, ymax
-
+            for ax in _axes:
+                ax_min = kwargs.get('{}min'.format(ax), None)
+                ax_max = kwargs.get('{}max'.format(ax), None)
+                auto = False  # turn off autoscaling, unless...
+                if ax_min is None and ax_max is None:
+                    auto = None  # leave autoscaling state alone
+                ax_min, ax_max = self.set_xlim(ax_min, ax_max,
+                                               emit=emit, auto=auto)
         v = v[0]
-        if len(v) != 4:
-            raise ValueError('v must contain [xmin xmax ymin ymax]')
-
-        self.set_xlim([v[0], v[1]], emit=emit, auto=False)
-        self.set_ylim([v[2], v[3]], emit=emit, auto=False)
+        if len(v) != 2*len(_axes):
+            raise ValueError('v must contain [{}]'.format(
+                ' '.join('{}min {}max'.format(ax) for ax in _axes)
+            ))
+
+        for i, ax in enumerate(_axes):
+            getattr(self, 'set_{}lim'.format(ax))(
+                [v[i*2], v[i*2+1]], emit=emit, auto=False)
 
         return v
 

diff --git a/lib/mpl_toolkits/mplot3d/axes3d.py b/lib/mpl_toolkits/mplot3d/axes3d.py
@@ -78,6 +78,11 @@ def __init__(self, fig, rect=None, *args, **kwargs):
 
         self.initial_azim = kwargs.pop('azim', -60)
         self.initial_elev = kwargs.pop('elev', 30)
+        if 'pb_aspect' in kwargs:
+            self.pb_aspect = np.asarray(kwargs['pb_aspect'])
+        else:
+            # chosen for similarity with the previous initial view
+            self.pb_aspect = np.array([4, 4, 3]) / 3.5
         zscale = kwargs.pop('zscale', None)
         sharez = kwargs.pop('sharez', None)
         self.set_proj_type(kwargs.pop('proj_type', 'persp'))
@@ -247,6 +252,21 @@ def tunit_edges(self, vals=None, M=None):
                  (tc[7], tc[4])]
         return edges
 
+    def set_pb_aspect(self, pb_aspect, zoom=1):
+        self.pb_aspect = pb_aspect * 1.8 * zoom / proj3d.mod(pb_aspect)
+
+    def set_aspect(self, aspect, adjustable=None, anchor=None):
+        if isinstance(aspect, six.string_types):
+            return super(Axes3D, self).set_aspect(
+                aspect=aspect, adjustable=adjustable, anchor=anchor)
+        else:
+            super(Axes3D, self).set_aspect(
+                aspect='auto', adjustable=adjustable, anchor=anchor)
+            aspect = np.asarray(aspect)
+            if aspect.shape != (3,):
+                raise ValueError("aspect must be a vector of shape (3,)")
+            self._aspect = aspect
+
     def apply_aspect(self, position=None):
         if position is None:
             position = self.get_position(original=True)
@@ -261,6 +281,47 @@ def apply_aspect(self, position=None):
         pb1 = pb.shrunk_to_aspect(box_aspect, pb, fig_aspect)
         self.set_position(pb1.anchored(self.get_anchor(), pb), 'active')
 
+        # adjust the plot ratio to match the figure size
+        aspect_2d = (figW * pb.width) / (figH * pb.height)
+        self.set_pb_aspect(np.array([aspect_2d, aspect_2d, 1]))
+
+        aspect = self.get_aspect()
+        if isinstance(aspect, np.ndarray):
+            pass
+        elif aspect == 'auto':
+            return
+        elif aspect == 'equal':
+            aspect = np.array([1, 1, 1])
+
+        get_bound = [getattr(self, 'get_{}bound'.format(ax)) for ax in 'xyz']
+        set_bound = [getattr(self, 'set_{}bound'.format(ax)) for ax in 'xyz']
+
+        # get bounds
+        upper = np.empty(3)
+        lower = np.empty(3)
+        for i in range(3):
+            lower[i], upper[i] = get_bound[i]()
+
+        if self._adjustable == 'datalim':
+
+            center = (upper + lower) / 2.0
+            span = (upper - lower) / 2.0
+
+            # compute scaling
+            rel_aspect = self.pb_aspect / aspect
+            scale = rel_aspect / span
+            scale /= min(scale)
+
+            # rescale around center point
+            lower = center - span*scale
+            upper = center + span*scale
+            for i in range(3):
+                set_bound[i](lower[i], upper[i])
+        elif self._adjustable in ['box', 'box-forced']:
+            # override the plot ratio chosen earlier
+            scale = (upper - lower) * aspect
+            self.set_pb_aspect(scale)
+
     def draw(self, renderer):
         # draw the background patch
         self.patch.draw(renderer)
@@ -999,9 +1060,6 @@ def get_proj(self):
         point.
 
         """
-        # chosen for similarity with the previous initial view
-        pb_aspect = np.array([4, 4, 3]) / 3.5
-
         relev, razim = np.pi * self.elev/180, np.pi * self.azim/180
 
         xmin, xmax = self.get_xlim3d()
@@ -1011,10 +1069,10 @@ def get_proj(self):
         # transform to uniform world coordinates 0-1.0,0-1.0,0-1.0
         worldM = proj3d.world_transformation(xmin, xmax,
                                              ymin, ymax,
-                                             zmin, zmax, pb_aspect=pb_aspect)
+                                             zmin, zmax, pb_aspect=self.pb_aspect)
 
         # look into the middle of the new coordinates
-        R = pb_aspect / 2
+        R = self.pb_aspect / 2
 
         xp = R[0] + np.cos(razim) * np.cos(relev) * self.dist
         yp = R[1] + np.sin(razim) * np.cos(relev) * self.dist