matplotlib function to set axes equal for 3D plots

cython/gtsam/utils/plot.py

@@ -3,6 +3,30 @@
 import numpy as np
 import matplotlib.pyplot as plt
 from matplotlib import patches
+from mpl_toolkits.mplot3d import Axes3D
+
+
+def set_axes_equal(ax):
+    """
+    Make axes of 3D plot have equal scale so that spheres appear as spheres,
+    cubes as cubes, etc..  This is one possible solution to Matplotlib's
+    ax.set_aspect('equal') and ax.axis('equal') not working for 3D.
+    Input
+      ax: a matplotlib axis, e.g., as output from plt.gca().
+    """
+
+    limits = np.array([
+        ax.get_xlim3d(),
+        ax.get_ylim3d(),
+        ax.get_zlim3d(),
+    ])
+
+    origin = np.mean(limits, axis=1)
+    radius = 0.5 * np.max(np.abs(limits[:, 1] - limits[:, 0]))
+
+    ax.set_xlim3d([origin[0] - radius, origin[0] + radius])
+    ax.set_ylim3d([origin[1] - radius, origin[1] + radius])
+    ax.set_zlim3d([origin[2] - radius, origin[2] + radius])
 
 
 def plot_pose2_on_axes(axes, pose, axis_length=0.1, covariance=None):
@@ -35,6 +59,7 @@ def plot_pose2_on_axes(axes, pose, axis_length=0.1, covariance=None):
                              np.rad2deg(angle), fill=False)
         axes.add_patch(e1)
 
+
 def plot_pose2(fignum, pose, axis_length=0.1, covariance=None):
     """Plot a 2D pose on given figure with given 'axis_length'."""
     # get figure object