Merge pull request #341 from borglab/fix/python-plotting

Plotting updates

cython/gtsam/utils/plot.py

@@ -13,8 +13,9 @@ def set_axes_equal(fignum):
     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().
+
+    Args:
+      fignum (int): An integer representing the figure number for Matplotlib.
     """
     fig = plt.figure(fignum)
     ax = fig.gca(projection='3d')
@@ -34,7 +35,21 @@ def set_axes_equal(fignum):
 
 
 def ellipsoid(xc, yc, zc, rx, ry, rz, n):
-    """Numpy equivalent of Matlab's ellipsoid function"""
+    """
+    Numpy equivalent of Matlab's ellipsoid function.
+
+    Args:
+        xc (double): Center of ellipsoid in X-axis.
+        yc (double): Center of ellipsoid in Y-axis.
+        zc (double): Center of ellipsoid in Z-axis.
+        rx (double): Radius of ellipsoid in X-axis.
+        ry (double): Radius of ellipsoid in Y-axis.
+        rz (double): Radius of ellipsoid in Z-axis.
+        n (int): The granularity of the ellipsoid plotted. 
+
+    Returns:
+        tuple[numpy.ndarray]: The points in the x, y and z axes to use for the surface plot.
+    """
     u = np.linspace(0, 2*np.pi, n+1)
     v = np.linspace(0, np.pi, n+1)
     x = -rx * np.outer(np.cos(u), np.sin(v)).T
@@ -51,6 +66,14 @@ def plot_covariance_ellipse_3d(axes, origin, P, scale=1, n=8, alpha=0.5):
     Based on Maybeck Vol 1, page 366
     k=2.296 corresponds to 1 std, 68.26% of all probability
     k=11.82 corresponds to 3 std, 99.74% of all probability
+
+    Args:
+        axes (matplotlib.axes.Axes): Matplotlib axes.
+        origin (gtsam.Point3): The origin in the world frame.
+        P (numpy.ndarray): The marginal covariance matrix of the 3D point which will be represented as an ellipse.
+        scale (float): Scaling factor of the radii of the covariance ellipse.
+        n (int): Defines the granularity of the ellipse. Higher values indicate finer ellipses.
+        alpha (float): Transparency value for the plotted surface in the range [0, 1].
     """
     k = 11.82
     U, S, _ = np.linalg.svd(P)
@@ -74,7 +97,15 @@ def plot_covariance_ellipse_3d(axes, origin, P, scale=1, n=8, alpha=0.5):
 
 
 def plot_pose2_on_axes(axes, pose, axis_length=0.1, covariance=None):
-    """Plot a 2D pose on given axis 'axes' with given 'axis_length'."""
+    """
+    Plot a 2D pose on given axis `axes` with given `axis_length`.
+
+    Args:
+        axes (matplotlib.axes.Axes): Matplotlib axes.
+        pose (gtsam.Pose2): The pose to be plotted.
+        axis_length (float): The length of the camera axes.
+        covariance (numpy.ndarray): Marginal covariance matrix to plot the uncertainty of the estimation.
+    """
     # get rotation and translation (center)
     gRp = pose.rotation().matrix()  # rotation from pose to global
     t = pose.translation()
@@ -105,22 +136,47 @@ def plot_pose2_on_axes(axes, pose, axis_length=0.1, covariance=None):
 
 
 def plot_pose2(fignum, pose, axis_length=0.1, covariance=None):
-    """Plot a 2D pose on given figure with given 'axis_length'."""
+    """
+    Plot a 2D pose on given figure with given `axis_length`.
+
+    Args:
+        fignum (int): Integer representing the figure number to use for plotting.
+        pose (gtsam.Pose2): The pose to be plotted.
+        axis_length (float): The length of the camera axes.
+        covariance (numpy.ndarray): Marginal covariance matrix to plot the uncertainty of the estimation.
+    """
     # get figure object
     fig = plt.figure(fignum)
     axes = fig.gca()
-    plot_pose2_on_axes(axes, pose, axis_length, covariance)
+    plot_pose2_on_axes(axes, pose, axis_length=axis_length,
+                       covariance=covariance)
 
 
 def plot_point3_on_axes(axes, point, linespec, P=None):
-    """Plot a 3D point on given axis 'axes' with given 'linespec'."""
+    """
+    Plot a 3D point on given axis `axes` with given `linespec`.
+
+    Args:
+        axes (matplotlib.axes.Axes): Matplotlib axes.
+        point (gtsam.Point3): The point to be plotted.
+        linespec (string): String representing formatting options for Matplotlib.
+        P (numpy.ndarray): Marginal covariance matrix to plot the uncertainty of the estimation.
+    """
     axes.plot([point.x()], [point.y()], [point.z()], linespec)
     if P is not None:
         plot_covariance_ellipse_3d(axes, point.vector(), P)
 
 
 def plot_point3(fignum, point, linespec, P=None):
-    """Plot a 3D point on given figure with given 'linespec'."""
+    """
+    Plot a 3D point on given figure with given `linespec`.
+
+    Args:
+        fignum (int): Integer representing the figure number to use for plotting.
+        point (gtsam.Point3): The point to be plotted.
+        linespec (string): String representing formatting options for Matplotlib.
+        P (numpy.ndarray): Marginal covariance matrix to plot the uncertainty of the estimation.
+    """
     fig = plt.figure(fignum)
     axes = fig.gca(projection='3d')
     plot_point3_on_axes(axes, point, linespec, P)
@@ -128,10 +184,16 @@ def plot_point3(fignum, point, linespec, P=None):
 
 def plot_3d_points(fignum, values, linespec="g*", marginals=None):
     """
-    Plots the Point3s in 'values', with optional covariances.
+    Plots the Point3s in `values`, with optional covariances.
     Finds all the Point3 objects in the given Values object and plots them.
     If a Marginals object is given, this function will also plot marginal
     covariance ellipses for each point.
+
+    Args:
+        fignum (int): Integer representing the figure number to use for plotting.
+        values (gtsam.Values): Values dictionary consisting of points to be plotted.
+        linespec (string): String representing formatting options for Matplotlib.
+        covariance (numpy.ndarray): Marginal covariance matrix to plot the uncertainty of the estimation.
     """
 
     keys = values.keys()
@@ -142,11 +204,11 @@ def plot_3d_points(fignum, values, linespec="g*", marginals=None):
             key = keys.at(i)
             point = values.atPoint3(key)
             if marginals is not None:
-                P = marginals.marginalCovariance(key);
+                covariance = marginals.marginalCovariance(key)
             else:
-                P = None
+                covariance = None
 
-            plot_point3(fignum, point, linespec, P)
+            plot_point3(fignum, point, linespec, covariance)
 
         except RuntimeError:
             continue
@@ -154,7 +216,15 @@ def plot_3d_points(fignum, values, linespec="g*", marginals=None):
 
 
 def plot_pose3_on_axes(axes, pose, axis_length=0.1, P=None, scale=1):
-    """Plot a 3D pose on given axis 'axes' with given 'axis_length'."""
+    """
+    Plot a 3D pose on given axis `axes` with given `axis_length`.
+
+    Args:
+        axes (matplotlib.axes.Axes): Matplotlib axes.
+        point (gtsam.Point3): The point to be plotted.
+        linespec (string): String representing formatting options for Matplotlib.
+        P (numpy.ndarray): Marginal covariance matrix to plot the uncertainty of the estimation.
+    """
     # get rotation and translation (center)
     gRp = pose.rotation().matrix()  # rotation from pose to global
     origin = pose.translation().vector()
@@ -182,13 +252,33 @@ def plot_pose3_on_axes(axes, pose, axis_length=0.1, P=None, scale=1):
 
 
 def plot_pose3(fignum, pose, axis_length=0.1, P=None):
-    """Plot a 3D pose on given figure with given 'axis_length'."""
+    """
+    Plot a 3D pose on given figure with given `axis_length`.
+
+    Args:
+        fignum (int): Integer representing the figure number to use for plotting.
+        pose (gtsam.Pose3): 3D pose to be plotted.
+        linespec (string): String representing formatting options for Matplotlib.
+        P (numpy.ndarray): Marginal covariance matrix to plot the uncertainty of the estimation.
+    """
     # get figure object
     fig = plt.figure(fignum)
     axes = fig.gca(projection='3d')
-    plot_pose3_on_axes(axes, pose, P=P, axis_length=axis_length)
+    plot_pose3_on_axes(axes, pose, P=P,
+                       axis_length=axis_length)
+
 
 def plot_trajectory(fignum, values, scale=1, marginals=None):
+    """
+    Plot a complete 3D trajectory using poses in `values`.
+
+    Args:
+        fignum (int): Integer representing the figure number to use for plotting.
+        values (gtsam.Values): Values dict containing the poses.
+        scale (float): Value to scale the poses by.
+        marginals (gtsam.Marginals): Marginalized probability values of the estimation.
+            Used to plot uncertainty bounds.
+    """
     pose3Values = gtsam.allPose3s(values)
     keys = gtsam.KeyVector(pose3Values.keys())
     lastIndex = None
@@ -209,11 +299,12 @@ def plot_trajectory(fignum, values, scale=1, marginals=None):
                 pass
 
             if marginals:
-                P = marginals.marginalCovariance(lastKey)
+                covariance = marginals.marginalCovariance(lastKey)
             else:
-                P = None
+                covariance = None
 
-            plot_pose3(fignum, lastPose, P, scale)
+            plot_pose3(fignum, lastPose,  P=covariance,
+                       axis_length=scale)
 
         lastIndex = i
 
@@ -223,11 +314,12 @@ def plot_trajectory(fignum, values, scale=1, marginals=None):
         try:
             lastPose = pose3Values.atPose3(lastKey)
             if marginals:
-                P = marginals.marginalCovariance(lastKey)
+                covariance = marginals.marginalCovariance(lastKey)
             else:
-                P = None
+                covariance = None
 
-            plot_pose3(fignum, lastPose, P, scale)
+            plot_pose3(fignum, lastPose, P=covariance,
+                       axis_length=scale)
 
         except:
             pass