From cd1ba206c80399459a8c1384d72d9452c5d04cbd Mon Sep 17 00:00:00 2001
From: Michael Grupp <grupp@magazino.eu>
Date: Tue, 7 May 2024 20:54:29 +0200
Subject: [PATCH] Add 'speeds' plot tab to evo_traj (#664)

---
 evo/main_traj.py  | 26 ++++++++++++++++++++++++--
 evo/tools/plot.py | 24 +++++++++++++++++++++++-
 2 files changed, 47 insertions(+), 3 deletions(-)

diff --git a/evo/main_traj.py b/evo/main_traj.py
index 8acefc25..ae71448e 100755
--- a/evo/main_traj.py
+++ b/evo/main_traj.py
@@ -312,6 +312,7 @@ def run(args):
         fig_rpy, axarr_rpy = plt.subplots(3, sharex="col",
                                           figsize=tuple(SETTINGS.plot_figsize))
         fig_traj = plt.figure(figsize=tuple(SETTINGS.plot_figsize))
+        fig_speed = plt.figure()
 
         plot_mode = plot.PlotMode[args.plot_mode]
         length_unit = Unit(SETTINGS.plot_trajectory_length_unit)
@@ -349,6 +350,16 @@ def run(args):
                           label=short_traj_name,
                           alpha=SETTINGS.plot_reference_alpha,
                           start_timestamp=start_time)
+            if isinstance(ref_traj, trajectory.PoseTrajectory3D):
+                try:
+                    plot.speeds(fig_speed.gca(), ref_traj,
+                                style=SETTINGS.plot_reference_linestyle,
+                                color=SETTINGS.plot_reference_color,
+                                alpha=SETTINGS.plot_reference_alpha,
+                                label=short_traj_name)
+                except trajectory.TrajectoryException as error:
+                    logger.error(
+                        f"Can't plot speeds of {short_traj_name}: {error}")
         elif args.plot_relative_time:
             # Use lower bound timestamp as the 0 time if there's no reference.
             if len(trajectories) > 1:
@@ -390,6 +401,16 @@ def run(args):
                           color, short_traj_name,
                           alpha=SETTINGS.plot_trajectory_alpha,
                           start_timestamp=start_time)
+            if isinstance(traj, trajectory.PoseTrajectory3D):
+                try:
+                    plot.speeds(fig_speed.gca(), traj,
+                                style=SETTINGS.plot_trajectory_linestyle,
+                                color=color,
+                                alpha=SETTINGS.plot_trajectory_alpha,
+                                label=short_traj_name)
+                except trajectory.TrajectoryException as error:
+                    logger.error(
+                        f"Can't plot speeds of {short_traj_name}: {error}")
             if not SETTINGS.plot_usetex:
                 fig_rpy.text(
                     0., 0.005, "euler_angle_sequence: {}".format(
@@ -399,8 +420,9 @@ def run(args):
             plot.ros_map(ax_traj, args.ros_map_yaml, plot_mode)
 
         plot_collection.add_figure("trajectories", fig_traj)
-        plot_collection.add_figure("xyz_view", fig_xyz)
-        plot_collection.add_figure("rpy_view", fig_rpy)
+        plot_collection.add_figure("xyz", fig_xyz)
+        plot_collection.add_figure("rpy", fig_rpy)
+        plot_collection.add_figure("speeds", fig_speed)
         if args.plot:
             plot_collection.show()
         if args.save_plot:
diff --git a/evo/tools/plot.py b/evo/tools/plot.py
index 8b98a336..d01a4f40 100644
--- a/evo/tools/plot.py
+++ b/evo/tools/plot.py
@@ -472,7 +472,7 @@ def traj_colormap(ax: Axes, traj: trajectory.PosePath3D, array: ListOrArray,
     ax.autoscale_view(True, True, True)
     if plot_mode == PlotMode.xyz and isinstance(ax, Axes3D):
         min_z = np.amin(traj.positions_xyz[:, 2])
-        max_z = np.amax(traj.positions_xyz[:, 2])                
+        max_z = np.amax(traj.positions_xyz[:, 2])
         # Only adjust limits if there are z values to suppress mpl warning.
         if min_z != max_z:
             ax.set_zlim(min_z, max_z)
@@ -650,6 +650,28 @@ def traj_rpy(axarr: np.ndarray, traj: trajectory.PosePath3D, style: str = '-',
         axarr[0].legend(frameon=True)
 
 
+def speeds(ax: Axes, traj: trajectory.PoseTrajectory3D, style: str = '-',
+           color="black", label: str = "", alpha: float = 1.):
+    """
+    Plots the speed between poses of a trajectory.
+    Note that a speed value is shown at the timestamp of the newer pose.
+    :param ax: matplotlib axis
+    :param traj: trajectory.PoseTrajectory3D object
+    :param style: matplotlib line style
+    :param color: matplotlib color
+    :param label: label (for legend)
+    :param alpha: alpha value for transparency
+    """
+    if not isinstance(traj, trajectory.PoseTrajectory3D):
+        raise PlotException("speeds can only be plotted with trajectories")
+    ax.plot(traj.timestamps[1:], traj.speeds, style, color=color, alpha=alpha,
+            label=label)
+    ax.set_xlabel("$t$ (s)")
+    ax.set_ylabel("$v$ (m/s)")
+    if label and SETTINGS.plot_show_legend:
+        ax.legend(frameon=True)
+
+
 def trajectories(fig: Figure, trajectories: typing.Union[
         trajectory.PosePath3D, typing.Sequence[trajectory.PosePath3D],
         typing.Dict[str, trajectory.PosePath3D]], plot_mode=PlotMode.xy,