Example plot_trajectory.py

#!/usr/bin/env python3

__copyright__ = """

  Slamcore Confidential
  ---------------------

  Slamcore Limited
  All Rights Reserved.
  (C) Copyright 2021

  NOTICE:

  All information contained herein is, and remains the property of Slamcore
  Limited and its suppliers, if any. The intellectual and technical concepts
  contained herein are proprietary to Slamcore Limited and its suppliers and
  may be covered by patents in process, and are protected by trade secret or
  copyright law. Dissemination of this information or reproduction of this
  material is strictly forbidden unless prior written permission is obtained
  from Slamcore Limited.
"""

__license__ = "Slamcore Confidential"

__doc__ = "Plot Slamcore trajectories in 3D space"


from mpl_toolkits.mplot3d import Axes3D
from pathlib import Path

import argparse
import matplotlib.pyplot as plt
import numpy as np

import sys


def valid_path(s: str) -> Path:
    p = Path(s)
    if not p.is_file():
        raise RuntimeError(f"Not a valid file: {p}")

    return p


def main():
    # Parse command line arguments. We expect a list of trajectory files
    # generated by any of the Slamcore tools.
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "trajectory_file",
        type=valid_path,
        nargs="+",
        help="Slamcore trajectory CSV files to plot",
    )
    args = parser.parse_args()

    # Create a generator to read trajectories into pairs like this:
    # (
    #   (filename_0, [[t_x_0, t_y_0, t_z_0], [t_x_1, t_y_1, t_z_1], ... ])
    #   (filename_1, [[t_x_0, t_y_0, t_z_0], [t_x_1, t_y_1, t_z_1], ... ])
    #   ...
    # )
    trajectories = (
        (
            p.stem,
            np.loadtxt(open(p, "rb"), delimiter=", ", skiprows=1, usecols=(3, 4, 5)),
        )
        for p in args.trajectory_file
    )

    # Plot trajectories.
    figure = plt.figure("Slamcore Plot Trajectory Example")
    ax = plt.axes(projection="3d")
    ax.set_xlabel("X Displacement [m]")
    ax.set_ylabel("Y Displacement [m]")
    ax.set_zlabel("Z Displacement [m]")

    max_dim = sys.float_info.min
    min_dim = sys.float_info.max

    for file_name, trajectory in trajectories:
        max_dim = max(max_dim, np.amax(trajectory))
        min_dim = min(min_dim, np.amin(trajectory))
        ax.plot(
            trajectory[:, 0],
            trajectory[:, 1],
            trajectory[:, 2],
            linewidth=0.5,
            label=file_name,
        )

    ax.set_xlim3d(min_dim, max_dim)
    ax.set_ylim3d(min_dim, max_dim)
    ax.set_zlim3d(min_dim, max_dim)

    ax.legend()
    plt.show()


if __name__ == "__main__":
    main()