Provide adjustable='box' for set_aspect('equal') for 3D axes.

tfpf · tfpf · commit 75e605ffe02b · 2022-08-07T20:02:46.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
@@ -2,7 +2,9 @@ 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'.
+'equalxy', 'equalxz', or 'equalyz' rather than the default of 'auto'. Just like
+in case of 2D plots, either the data limits or the bounding box can be adjusted
+to achieve the desired aspect.
 
 .. plot::
     :include-source: true
@@ -26,7 +28,9 @@ Users can set the aspect ratio for the X, Y, Z axes of a 3D plot to be 'equal',
     # Set the aspect ratios
     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_aspect(aspects[i], adjustable='datalim')
+        # Alternatively: ax.set_aspect(aspects[i], adjustable='box')
+        # which will modify the box aspect ratio instead of the data limits.
+        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
@@ -286,9 +286,7 @@ def set_aspect(self, aspect, adjustable=None, anchor=None, share=False):
             'equalyz'   adapt the y and z axes to have equal aspect ratios.
             =========   ==================================================
 
-        adjustable : None
-            Currently ignored by Axes3D
-
+        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.
@@ -319,34 +317,65 @@ def set_aspect(self, aspect, adjustable=None, anchor=None, share=False):
         """
         _api.check_in_list(('auto', 'equal', 'equalxy', 'equalyz', 'equalxz'),
                            aspect=aspect)
+        if adjustable is None:
+            adjustable = self._adjustable
+        _api.check_in_list(('box', 'datalim'), adjustable=adjustable)
         super().set_aspect(
             aspect='auto', adjustable=adjustable, anchor=anchor, share=share)
         self._aspect = aspect
 
         if aspect in ('equal', 'equalxy', 'equalxz', 'equalyz'):
-            if aspect == 'equal':
-                ax_indices = [0, 1, 2]
-            elif aspect == 'equalxy':
-                ax_indices = [0, 1]
-            elif aspect == 'equalxz':
-                ax_indices = [0, 2]
-            elif aspect == 'equalyz':
-                ax_indices = [1, 2]
+            ax_idx = self._equal_aspect_axis_indices(aspect)
 
             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)
+                delta = max(ptp[ax_idx])
+                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_idx:
+                        set_lim(mean[i] - deltas[i]/2., mean[i] + deltas[i]/2.)
+            else:  # 'box'
+                # Change the box aspect such that the ratio of the length of
+                # the unmodified axis to the length of the diagonal
+                # perpendicular to it remains unchanged.
+                box_aspect = np.array(self._box_aspect)
+                box_aspect[ax_idx] = ptp[ax_idx]
+                remaining_ax_idx = {0, 1, 2}.difference(ax_idx)
+                if remaining_ax_idx:
+                    remaining = remaining_ax_idx.pop()
+                    old_diag = np.linalg.norm(self._box_aspect[ax_idx])
+                    new_diag = np.linalg.norm(box_aspect[ax_idx])
+                    box_aspect[remaining] *= new_diag / old_diag
+                self.set_box_aspect(box_aspect)
+
+    def _equal_aspect_axis_indices(self, aspect):
+        """
+        Get the indices for which of the x, y, z axes are constrained to have
+        equal aspect ratios.
+
+        Parameters
+        ----------
+        aspect : {'auto', 'equal', 'equalxy', 'equalxz', 'equalyz'}
+            See descriptions in docstring for `.set_aspect()`.
+        """
+        ax_indices = []  # aspect == 'auto'
+        if aspect == 'equal':
+            ax_indices = [0, 1, 2]
+        elif aspect == 'equalxy':
+            ax_indices = [0, 1]
+        elif aspect == 'equalxz':
+            ax_indices = [0, 2]
+        elif aspect == 'equalyz':
+            ax_indices = [1, 2]
+        return ax_indices
 
     def set_box_aspect(self, aspect, *, zoom=1):
         """
diff --git a/lib/mpl_toolkits/tests/baseline_images/test_mplot3d/aspects_adjust_box.png b/lib/mpl_toolkits/tests/baseline_images/test_mplot3d/aspects_adjust_box.png
diff --git a/lib/mpl_toolkits/tests/test_mplot3d.py b/lib/mpl_toolkits/tests/test_mplot3d.py
@@ -44,7 +44,25 @@ def test_aspects():
                 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])
+        ax.set_aspect(aspects[i], adjustable='datalim')
+
+
+@mpl3d_image_comparison(['aspects_adjust_box.png'], remove_text=False)
+def test_aspects_adjust_box():
+    aspects = ('auto', 'equal', 'equalxy', 'equalyz', 'equalxz')
+    fig, axs = plt.subplots(1, len(aspects), subplot_kw={'projection': '3d'},
+                            figsize=(11, 3))
+
+    # Draw rectangular cuboid with side lengths [4, 3, 5]
+    r = [0, 1]
+    scale = np.array([4, 3, 5])
+    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_aspect(aspects[i], adjustable='box')
 
 
 def test_axes3d_repr():
