Fixed calculation errors. Updated tests and image.

tfpf · tfpf · commit aa5228b80147 · 2022-08-05T13:15:19.000+05:30
diff --git a/lib/mpl_toolkits/mplot3d/axes3d.py b/lib/mpl_toolkits/mplot3d/axes3d.py
@@ -267,8 +267,7 @@ def tunit_edges(self, vals=None, M=None):
                  (tc[7], tc[4])]
         return edges
 
-    def set_aspect(self, aspect, adjustable='datalim', anchor=None,
-                   share=False):
+    def set_aspect(self, aspect, adjustable=None, anchor=None, share=False):
         """
         Set the aspect ratios.
 
@@ -287,9 +286,10 @@ def set_aspect(self, aspect, adjustable='datalim', anchor=None,
             'equalyz'   adapt the y and z axes to have equal aspect ratios.
             =========   ==================================================
 
-        adjustable : {'datalim', 'box'}, default: 'datalim'
-            This defines which parameter will be adjusted to meet the required
-            aspect. See `.set_adjustable` for further details.
+        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.
 
         anchor : None or str or 2-tuple of float, optional
             If not *None*, this defines where the Axes will be drawn if there
@@ -317,50 +317,65 @@ def set_aspect(self, aspect, adjustable='datalim', anchor=None,
         """
         _api.check_in_list(('auto', 'equal', 'equalxy', 'equalyz', 'equalxz'),
                            aspect=aspect)
-        _api.check_in_list(('datalim', 'box'), adjustable=adjustable)
+        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()])
             ptp = np.ptp(view_intervals, axis=1)
             if adjustable == 'datalim':
                 mean = np.mean(view_intervals, axis=1)
-                delta = max(ptp[ax_indices])
+                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_indices:
+                    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_indices] = ptp[ax_indices]
-                remaining_ax_indices = {0, 1, 2}.difference(ax_indices)
-                if remaining_ax_indices:
-                    remaining_ax_index = remaining_ax_indices.pop()
-                    old_diag = np.linalg.norm(box_aspect[ax_indices])
-                    new_diag = np.linalg.norm(ptp[ax_indices])
-                    box_aspect[remaining_ax_index] *= new_diag / old_diag
+                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):
         """
         Set the Axes box aspect.
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,24 @@ 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'})
+
+    # Draw rectangular cuboid with side lengths [1, 1, 5]
+    r = [0, 1]
+    scale = np.array([1, 1, 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():
