Plot "time-ordered point clouds" with colors on the rainbow to indicate the time

makeColorScatterplots.py

"""
Programmer: Chris Tralie
Purpose: To show how to plot "time-ordered point clouds"
with colors on the rainbow to indicate the time
(red beginning, magenta end)
"""
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

"""
STEP 1: Initialize a 3D Figure 8
"""
N = 100
t = np.linspace(0, 2*np.pi, N+1)[0:N]
X = np.zeros((N, 3))
X[:, 0] = np.cos(t)
X[:, 1] = np.sin(2*t)
X[:, 2] = (t/(2*np.pi)-0.5)**2



"""
STEP 2: Do the "easy way," choosing a colormap that
automatically colors things according to the time parameter
Show in 2D and 3D
"""
plt.figure(figsize=(12, 5))
plt.subplot(121)
plt.scatter(X[:, 0], X[:, 1], c=t, cmap='Spectral')
plt.axis('equal')
ax = plt.gcf().add_subplot(122, projection='3d')
ax.scatter(X[:, 0], X[:, 1], X[:, 2], c=t, cmap='Spectral')
ax.view_init(azim=-140, elev=9)
plt.show()


"""
STEP 3: Do the "still easy but slightly more flexible way," 
setting up an array that's parallel to the points to be plotted
where the colors are declared explicitly.  This makes it easier
to plot subsets of the points with consistent colors, such
as an evolving video of the point cloud in time, which is what
I use as an example below
Show in 2D and 3D, and make a video with points gradually added
"""
plt.figure(figsize=(12, 5))
# Precompute an Nx4 array of RGBA colors the same as the
# automatically generated colors when c=t
c = plt.get_cmap('Spectral')
C = c(np.interp(t, (t.min(), t.max()), (0, 1)))
for i in range(X.shape[0]):
    plt.clf()
    plt.subplot(121)
    # Plot transparent full point cloud to get axes range right
    plt.scatter(X[:, 0], X[:, 1], alpha=0) 
    plt.scatter(X[0:i+1, 0], X[0:i+1, 1], c=C[0:i+1, :])
    plt.axis('equal')
    ax = plt.gcf().add_subplot(122, projection='3d')
    ax.scatter(X[:, 0], X[:, 1], X[:, 2], alpha=0)
    ax.scatter(X[0:i+1, 0], X[0:i+1, 1], X[0:i+1, 2], c=C[0:i+1, :])
    ax.view_init(azim=-140, elev=9)
    plt.savefig("%i.png"%i, bbox_inches='tight')