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

tfpf · tfpf · commit aa9592c21b42 · 2022-08-04T14:34:13.000+05:30
diff --git a/lib/mpl_toolkits/mplot3d/axes3d.py b/lib/mpl_toolkits/mplot3d/axes3d.py
@@ -267,7 +267,7 @@ def tunit_edges(self, vals=None, M=None):
                  (tc[7], tc[4])]
         return edges
 
-    def set_aspect(self, aspect, adjustable=None, anchor=None, share=False):
+    def set_aspect(self, aspect, adjustable='datalim', anchor=None, share=False):
         """
         Set the aspect ratios.
 
@@ -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):
         """
