Plotting single sensor 3D animated trajectory

plot_3d_trajectory.py

# Lo Kwan Wai, Clive
# October 2, 2019
# Plotting 3D animated trajectory with single motion sensor

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from matplotlib.animation import FuncAnimation

# Plot 3d trajectory graph
# Input motion sensor sampling rate (fps) for real speed 
# (speed of trajectory is limited by computational speed)
def plot_3d_trajectory(x, y, z, fps=30, speed=1, line_trail=60, animate=True):
	
    def group_data(x, y, z):
        for i in range(len(x)):
            yield np.array([x[i], y[i], z[i]])

    # Update function
    def update_line(num, data, line):
        # line trail
        line_start = max((num - line_trail, 0))
        dot_start = max((num - 1, 0))

        # Update dot and line
        scat._offsets3d = (data[0, dot_start:num], data[1, dot_start:num], data[2, dot_start:num])
        line.set_data(data[:2, line_start:num])
        line.set_3d_properties(data[2, line_start:num])

        # Title displays frame number
        plt.title("Frame: {}".format(num + 1))

    # Convert to NumPy array
    x = np.array(x)
    y = np.array(y)
    z = np.array(z)

    # Set up the frame of the plot
    fig = plt.figure()
    ax = fig.add_subplot(111, projection="3d")

    # (this line will not work above Matplotlib v3.0.3)
    ax.set_aspect('equal')

    # Label axes
    ax.set_xlabel('X')
    ax.set_ylabel('Y')
    ax.set_zlabel('Z')

    # Create 3D bounding box to simulate equal aspect ratio
    # Based on https://stackoverflow.com/a/13701747
    max_range = np.array([x.max() - x.min(), y.max() - y.min(), z.max() - z.min()]).max()
    Xb = 0.5 * max_range * np.mgrid[-1:2:2, -1:2:2, -1:2:2][0].flatten() + 0.5 * (x.max() + x.min())
    Yb = 0.5 * max_range * np.mgrid[-1:2:2, -1:2:2, -1:2:2][1].flatten() + 0.5 * (y.max() + y.min())
    Zb = 0.5 * max_range * np.mgrid[-1:2:2, -1:2:2, -1:2:2][2].flatten() + 0.5 * (z.max() + z.min())
    for xb, yb, zb in zip(Xb, Yb, Zb):
        ax.plot([xb], [yb], [zb], 'w')
    plt.grid()

    # Generate an animated trajectory or a still image
    if animate:
        # Reshape data
        data = np.array(list(group_data(x, y, z))).T

        # Initialize first frame dot and line position
        line, = ax.plot(data[0, 0:1], data[1, 0:1], data[2, 0:1])
        scat = ax.scatter(data[0, 0], data[1, 0], data[2, 0], c='#771F1F')

        # Draw animation
        ani = FuncAnimation(fig, update_line, len(x), fargs=(data, line), interval=1000 / fps / speed, blit=False)
    else:
        ax.plot(x, y, z)

    plt.show()