change qplot to xplot

petercorke · petercorke · commit 65559a914fed · 2021-09-21T09:05:02.000+10:00
put it in its own module
diff --git a/roboticstoolbox/tools/__init__.py b/roboticstoolbox/tools/__init__.py
@@ -10,7 +10,6 @@
     ctraj,
     lspb,
     lspb_func,
-    qplot,
     mstraj,
 )
 from roboticstoolbox.tools.numerical import jacobian_numerical, hessian_numerical
@@ -21,7 +20,7 @@
     rtb_load_jsonfile,
     rtb_path_to_datafile,
 )
-
+from roboticstoolbox.tools.plot import xplot
 from roboticstoolbox.tools.params import rtb_set_param, rtb_get_param
 
 __all__ = [
@@ -34,7 +33,7 @@
     "ctraj",
     "lspb",
     "lspb_func",
-    "qplot",
+    "xplot",
     "mtraj",
     "mstraj",
     "jsingu",
diff --git a/roboticstoolbox/tools/plot.py b/roboticstoolbox/tools/plot.py
@@ -0,0 +1,104 @@
+import numpy as np
+import matplotlib.pyplot as plt
+
+def xplot(
+    x,
+    y=None,
+    wrist=False,
+    unwrap=False,
+    block=False,
+    labels=None,
+    loc=None,
+    grid=True,
+    stack=False,
+    **kwargs,
+):
+    """
+    Plot trajectory data
+
+    :param x: trajectory, one row per timestep
+    :type x: ndarray(m,n)
+    :param t: time vector, optional
+    :type t: numpy ndarray, shape=(M,)
+    :param wrist: distinguish arm and wrist joints with line styles
+    :type wrist: bool
+    :param unwrap: unwrap joint angles so that they smoothly increase or
+        decrease when they pass through :math:`\pm \pi`
+    :type unwrap: bool
+    :param block: block until the plot is closed
+    :type block: bool
+    :param labels: legend labels
+    :type labels: list of str, or single string with space separated labels
+    :param kwargs: options passed to pyplot.plot
+    :param loc: legend location as per pyplot.legend
+    :type loc: str
+
+    This is a convenience function to plot trajectories, where each row represents one time step.
+
+    - ``xplot(q)`` plots the joint angles versus row number.  If N==6 a
+      conventional 6-axis robot is assumed, and the first three joints are
+      shown as solid lines, the last three joints (wrist) are shown as dashed
+      lines. A legend is also displayed.
+
+    - ``xplot(t, q)`` as above but displays the joint angle trajectory versus
+      time given the time vector T (Mx1).
+
+    Example::
+
+        >>> qplot(q, x, labels='x y z')
+
+    :seealso: :func:`jtraj`, :func:`numpy.unwrap`
+    """
+    if y is None:
+        q = x
+        t = np.arange(0, q.shape[0])
+    else:
+        t = x
+        q = y
+
+    if t.ndim != 1 or q.shape[0] != t.shape[0]:
+        raise ValueError("dimensions of arguments are not consistent")
+
+    if unwrap:
+        q = np.unwrap(q, axis=0)
+
+    n = q.shape[1]
+
+    if labels is None:
+        labels = [f"q{i}" for i in range(n)]
+    elif isinstance(labels, str):
+        labels = labels.split(" ")
+    elif not isinstance(labels, (tuple, list)):
+        raise TypeError("wrong type for labels")
+
+    fig, ax = plt.subplots()
+
+    if stack:
+        for i in range(n):
+            ax = plt.subplot(n, 1, i + 1)
+
+            plt.plot(t, q[:, i], **kwargs)
+
+            plt.grid(grid)
+            ax.set_ylabel(labels[i])
+            ax.set_xlim(t[0], t[-1])
+
+        ax.set_xlabel("Time (s)")
+
+    else:
+        if n == 6 and wrist:
+            plt.plot(t, q[:, 0:3], **kwargs)
+            plt.plot(t, q[:, 3:6], "--", **kwargs)
+        else:
+            plt.plot(t, q, **kwargs)
+
+        ax.legend(labels, loc=loc)
+
+        plt.grid(grid)
+        ax.set_xlabel("Time (s)")
+        ax.set_ylabel("Joint coordinates (rad,m)")
+        ax.set_xlim(t[0], t[-1])
+
+    plt.show(block=block)
+
+    return fig.get_axes()
