Style fixes. Addressed review comments.

tfpf · tfpf · commit dd8e12e07d45 · 2022-08-04T17:24:14.000+05:30
diff --git a/lib/mpl_toolkits/mplot3d/axes3d.py b/lib/mpl_toolkits/mplot3d/axes3d.py
@@ -267,8 +267,8 @@ 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='datalim', anchor=None,
+                   share=False):
         """
         Set the aspect ratios.
 
@@ -350,28 +350,18 @@ def set_aspect(
                     if i in ax_indices:
                         set_lim(mean[i] - deltas[i]/2., mean[i] + deltas[i]/2.)
             elif adjustable == '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.
                 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] *= np.sqrt(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] *= np.sqrt(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] *= np.sqrt(new_diag / old_diag)
-                    box_aspect[1] = deltas[1]
-                    box_aspect[2] = deltas[2]
+                box_aspect = np.array(self.get_box_aspect())
+                box_aspect[ax_indices] = deltas[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(deltas[ax_indices])
+                    box_aspect[remaining_ax_index] *= new_diag / old_diag
                 self.set_box_aspect(box_aspect)
 
     def set_box_aspect(self, aspect, *, zoom=1):
