Add equalxy, equalyz, equalxz aspect ratios

scottshambaugh · scottshambaugh · commit cc9e7ba80963 · 2022-07-14T18:24:00.000-05:00
Update docstrings
diff --git a/doc/users/next_whats_new/3d_plot_aspects.rst b/doc/users/next_whats_new/3d_plot_aspects.rst
@@ -7,13 +7,17 @@ Users can set the aspect ratio for the X, Y, Z axes of a 3D plot to be 'equal',
 .. plot::
     :include-source: true
 
+    import matplotlib.pyplot as plt
+    import numpy as np
+    from itertools import combinations, product
+
     aspects = ('auto', 'equal', 'equalxy', 'equalyz', 'equalxz')
     fig, axs = plt.subplots(1, len(aspects), subplot_kw={'projection': '3d'})
 
     # Draw rectangular cuboid with side lengths [1, 1, 2]
     r = [0, 1]
     scale = np.array([1, 1, 2])
-    pts = itertools.combinations(np.array(list(itertools.product(r, r, r))), 2)
+    pts = combinations(np.array(list(product(r, r, r))), 2)
     for start, end in pts:
         if np.sum(np.abs(start - end)) == r[1] - r[0]:
             for ax in axs:
@@ -23,6 +27,6 @@ Users can set the aspect ratio for the X, Y, Z axes of a 3D plot to be 'equal',
     for i, ax in enumerate(axs):
         ax.set_box_aspect((3, 4, 5))
         ax.set_aspect(aspects[i])
-        ax.title(f"set_aspect('{aspects[i]}')")
+        ax.set_title("set_aspect('{aspects[i]}')")
 
     plt.show()
diff --git a/lib/mpl_toolkits/mplot3d/axes3d.py b/lib/mpl_toolkits/mplot3d/axes3d.py
@@ -274,14 +274,17 @@ def set_aspect(self, aspect, adjustable=None, anchor=None, share=False):
 
         Parameters
         ----------
-        aspect : {'auto', 'equal'}
+        aspect : {'auto', 'equal', 'equalxy', 'equalxz', 'equalyz'}
             Possible values:
 
             =========   ==================================================
             value       description
             =========   ==================================================
             'auto'      automatic; fill the position rectangle with data.
-            'equal'     adapt the axes to have equal aspect ratios.
+            'equal'     adapt all the axes to have equal aspect ratios.
+            'equalxy'   adapt the x and y axes to have equal aspect ratios.
+            'equalxz'   adapt the x and z axes to have equal aspect ratios.
+            'equalyz'   adapt the y and z axes to have equal aspect ratios.
             =========   ==================================================
 
         adjustable : None
@@ -315,25 +318,33 @@ def set_aspect(self, aspect, adjustable=None, anchor=None, share=False):
         --------
         mpl_toolkits.mplot3d.axes3d.Axes3D.set_box_aspect
         """
-        if aspect not in ('auto', 'equal'):
-            raise NotImplementedError(
-                "Axes3D currently only supports the aspect argument "
-                f"'auto' or 'equal'. You passed in {aspect!r}."
-            )
+        _api.check_in_list(('auto', 'equal', 'equalxy', 'equalyz', 'equalxz'),
+                           aspect=aspect)
         super().set_aspect(
-            aspect, adjustable=adjustable, anchor=anchor, share=share)
-
-        if aspect == 'equal':
-            v_intervals = np.vstack((self.xaxis.get_view_interval(),
-                                     self.yaxis.get_view_interval(),
-                                     self.zaxis.get_view_interval()))
-            mean = np.mean(v_intervals, axis=1)
-            delta = np.max(np.ptp(v_intervals, axis=1))
+            aspect='auto', adjustable=adjustable, anchor=anchor, share=share)
+
+        if aspect in ('equal', 'equalxy', 'equalxz', 'equalyz'):
+            if aspect == 'equal':
+                axis_indices = [0, 1, 2]
+            elif aspect == 'equalxy':
+                axis_indices = [0, 1]
+            elif aspect == 'equalxz':
+                axis_indices = [0, 2]
+            elif aspect == 'equalyz':
+                axis_indices = [1, 2]
+
+            view_intervals = np.array([self.xaxis.get_view_interval(),
+                                       self.yaxis.get_view_interval(),
+                                       self.zaxis.get_view_interval()])
+            mean = np.mean(view_intervals, axis=1)
+            delta = np.max(np.ptp(view_intervals, axis=1))
             deltas = delta * self._box_aspect / min(self._box_aspect)
 
-            self.set_xlim3d(mean[0] - deltas[0] / 2., mean[0] + deltas[0] / 2.)
-            self.set_ylim3d(mean[1] - deltas[1] / 2., mean[1] + deltas[1] / 2.)
-            self.set_zlim3d(mean[2] - deltas[2] / 2., mean[2] + deltas[2] / 2.)
+            for i, set_lim in enumerate((self.set_xlim3d,
+                                         self.set_ylim3d,
+                                         self.set_zlim3d)):
+                if i in axis_indices:
+                    set_lim(mean[i] - deltas[i]/2., mean[i] + deltas[i]/2.)
 
     def set_box_aspect(self, aspect, *, zoom=1):
         """
diff --git a/lib/mpl_toolkits/tests/test_mplot3d.py b/lib/mpl_toolkits/tests/test_mplot3d.py
@@ -27,27 +27,22 @@ def test_invisible_axes(fig_test, fig_ref):
     ax.set_visible(False)
 
 
-@mpl3d_image_comparison(['aspect_equal.png'], remove_text=False)
-def test_aspect_equal():
-    fig = plt.figure()
-    ax = fig.add_subplot(projection='3d')
-
-    points = np.array([[0, 0, 0], [1, 0, 0], [0, 1, 0], [1, 1, 0],
-                       [0, 0, 3], [1, 0, 3], [0, 1, 3], [1, 1, 3]])
-
-    x = np.asarray([coord[0] for coord in points])
-    y = np.asarray([coord[1] for coord in points])
-    z = np.asarray([coord[2] for coord in points])
-
-    def plot_edge(i, j):
-        ax.plot([x[i], x[j]], [y[i], y[j]], [z[i], z[j]], c='r')
-
-    # hexaedron creation
-    plot_edge(0, 1), plot_edge(1, 3), plot_edge(3, 2), plot_edge(2, 0)
-    plot_edge(4, 5), plot_edge(5, 7), plot_edge(7, 6), plot_edge(6, 4)
-    plot_edge(0, 4), plot_edge(1, 5), plot_edge(3, 7), plot_edge(2, 6)
-
-    ax.set_aspect('equal')
+@mpl3d_image_comparison(['aspects.png'], remove_text=False)
+def test_aspects():
+    aspects = ('auto', 'equal', 'equalxy', 'equalyz', 'equalxz')
+    fig, axs = plt.subplots(1, len(aspects), subplot_kw={'projection': '3d'})
+
+    # Draw rectangular cuboid with side lengths [1, 1, 2]
+    r = [0, 1]
+    scale = np.array([1, 1, 2])
+    pts = itertools.combinations(np.array(list(itertools.product(r, r, r))), 2)
+    for start, end in pts:
+        if np.sum(np.abs(start - end)) == r[1] - r[0]:
+            for ax in axs:
+                ax.plot3D(*zip(start*scale, end*scale))
+    for i, ax in enumerate(axs):
+        ax.set_box_aspect((3, 4, 5))
+        ax.set_aspect(aspects[i])
 
 
 @mpl3d_image_comparison(['bar3d.png'])
