Provide axis('equal') for Axes3D

patquem · scottshambaugh · commit e59721eb67cc · 2022-07-21T19:07:24.000-06:00
Whats new for 3D plot equal aspect ratio

Code review updates

set_aspect('equal') now adopts current box aspect

Update test image

Update whats new

Update whats new

test image
diff --git a/doc/users/next_whats_new/3d_plot_aspects.rst b/doc/users/next_whats_new/3d_plot_aspects.rst
@@ -0,0 +1,28 @@
+Set equal aspect ratio for 3D plots
+-----------------------------------
+
+Users can set the aspect ratio for the X, Y, Z axes of a 3D plot to be 'equal',
+'equalxy', 'equalxz', or 'equalyz' rather than the default of 'auto'.
+
+.. plot::
+    :include-source: true
+
+    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), color='C0')
+
+    # Set the aspect ratios
+    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]}')")
+
+    plt.show()
diff --git a/lib/mpl_toolkits/mplot3d/axes3d.py b/lib/mpl_toolkits/mplot3d/axes3d.py
@@ -272,22 +272,16 @@ def set_aspect(self, aspect, adjustable=None, anchor=None, share=False):
         """
         Set the aspect ratios.
 
-        Axes 3D does not current support any aspect but 'auto' which fills
-        the Axes with the data limits.
-
-        To simulate having equal aspect in data space, set the ratio
-        of your data limits to match the value of `.get_box_aspect`.
-        To control box aspect ratios use `~.Axes3D.set_box_aspect`.
-
         Parameters
         ----------
-        aspect : {'auto'}
+        aspect : {'auto', 'equal'}
             Possible values:
 
             =========   ==================================================
             value       description
             =========   ==================================================
             'auto'      automatic; fill the position rectangle with data.
+            'equal'     adapt the axes to have equal aspect ratios.
             =========   ==================================================
 
         adjustable : None
@@ -321,14 +315,26 @@ def set_aspect(self, aspect, adjustable=None, anchor=None, share=False):
         --------
         mpl_toolkits.mplot3d.axes3d.Axes3D.set_box_aspect
         """
-        if aspect != 'auto':
+        if aspect not in ('auto', 'equal'):
             raise NotImplementedError(
                 "Axes3D currently only supports the aspect argument "
-                f"'auto'. You passed in {aspect!r}."
+                f"'auto' or 'equal'. You passed in {aspect!r}."
             )
         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))
+            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.)
+
     def set_box_aspect(self, aspect, *, zoom=1):
         """
         Set the Axes box aspect.
diff --git a/lib/mpl_toolkits/tests/baseline_images/test_mplot3d/aspects.png b/lib/mpl_toolkits/tests/baseline_images/test_mplot3d/aspects.png
diff --git a/lib/mpl_toolkits/tests/test_mplot3d.py b/lib/mpl_toolkits/tests/test_mplot3d.py
@@ -27,11 +27,27 @@ def test_invisible_axes(fig_test, fig_ref):
     ax.set_visible(False)
 
 
-def test_aspect_equal_error():
+@mpl3d_image_comparison(['aspect_equal.png'], remove_text=False)
+def test_aspect_equal():
     fig = plt.figure()
     ax = fig.add_subplot(projection='3d')
-    with pytest.raises(NotImplementedError):
-        ax.set_aspect('equal')
+
+    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(['bar3d.png'])
