Provide adjustable='box' for aspect('equal') in 3D.

tfpf · tfpf · commit 30724fb6f66b · 2022-08-04T14:19:21.000+05:30
diff --git a/doc/users/next_whats_new/3d_plot_aspects.rst b/doc/users/next_whats_new/3d_plot_aspects.rst
@@ -27,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.set_title("set_aspect('{aspects[i]}')")
+        ax.set_title(f"set_aspect('{aspects[i]}')")
 
     plt.show()
diff --git a/lib/mpl_toolkits/mplot3d/axes3d.py b/lib/mpl_toolkits/mplot3d/axes3d.py
@@ -336,17 +336,42 @@ def set_aspect(self, aspect, adjustable=None, anchor=None, share=False):
             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)
-            ptp = np.ptp(view_intervals, axis=1)
-            delta = max(ptp[ax_indices])
-            scale = self._box_aspect[ptp == delta][0]
-            deltas = delta * self._box_aspect / scale
-
-            for i, set_lim in enumerate((self.set_xlim3d,
-                                         self.set_ylim3d,
-                                         self.set_zlim3d)):
-                if i in ax_indices:
-                    set_lim(mean[i] - deltas[i]/2., mean[i] + deltas[i]/2.)
+            if adjustable == 'datalim':
+                mean = np.mean(view_intervals, axis=1)
+                ptp = np.ptp(view_intervals, axis=1)
+                delta = max(ptp[ax_indices])
+                scale = self._box_aspect[ptp == delta][0]
+                deltas = delta * self._box_aspect / scale
+
+                for i, set_lim in enumerate((self.set_xlim3d,
+                                             self.set_ylim3d,
+                                             self.set_zlim3d)):
+                    if i in ax_indices:
+                        set_lim(mean[i] - deltas[i]/2., mean[i] + deltas[i]/2.)
+            elif adjustable == 'box':
+                deltas = np.ptp(view_intervals, axis=1)
+                box_aspect = self.get_box_aspect()
+                if aspect == 'equal':
+                    box_aspect = np.ptp(view_intervals, axis=1)
+                elif aspect == 'equalxy':
+                    old_diag = box_aspect[0] ** 2 + box_aspect[1] ** 2
+                    new_diag = deltas[0] ** 2 + deltas[1] ** 2
+                    box_aspect[0] = deltas[0]
+                    box_aspect[1] = deltas[1]
+                    box_aspect[2] *= new_diag / old_diag
+                elif aspect == 'equalxz':
+                    old_diag = box_aspect[0] ** 2 + box_aspect[2] ** 2
+                    new_diag = deltas[0] ** 2 + deltas[2] ** 2
+                    box_aspect[0] = deltas[0]
+                    box_aspect[1] *= new_diag / old_diag
+                    box_aspect[2] = deltas[2]
+                elif aspect == 'equalyz':
+                    old_diag = box_aspect[1] ** 2 + box_aspect[2] ** 2
+                    new_diag = deltas[1] ** 2 + deltas[2] ** 2
+                    box_aspect[0] *= new_diag / old_diag
+                    box_aspect[1] = deltas[1]
+                    box_aspect[2] = deltas[2]
+                self.set_box_aspect(box_aspect)
 
     def set_box_aspect(self, aspect, *, zoom=1):
         """
