I’m trying to plot a 3d curve that has different colors depending on one of its parameters. I tried this method similar to this question, but it doesn’t work. Can anyone point me in the right direction?
import matplotlib.pyplot as plt
from matplotlib import cm
T=100
N=5*T
x=np.linspace(0,T,num=N)
y=np.cos(np.linspace(0,T,num=N))
z=np.sin(np.linspace(0,T,num=N))
fig = plt.figure()
ax = fig.gca(projection='3d')
ax.plot(x,y,z,cmap = cm.get_cmap("Spectral"),c=z)
plt.show()
Advertisement
Answer
To extend the approach in this tutorial to 3D, use x,y,z instead of x,y.
The desired shape for the segments is (number of segments, 2 points, 3 coordinates per point), so N-1,2,3. First the array of points is created with shape N, 3. Then start (xyz[:-1, :]) and end points (xyz[1:, :]) are stacked together.
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d.art3d import Line3DCollection
T = 100
N = 5 * T
x = np.linspace(0, T, num=N)
y = np.cos(np.linspace(0, T, num=N))
z = np.sin(np.linspace(0, T, num=N))
xyz = np.array([x, y, z]).T
segments = np.stack([xyz[:-1, :], xyz[1:, :]], axis=1) # shape is 499,2,3
cmap = plt.cm.get_cmap("Spectral")
norm = plt.Normalize(z.min(), z.max())
lc = Line3DCollection(segments, linewidths=2, colors=cmap(norm(z[:-1])))
fig = plt.figure()
ax = fig.gca(projection='3d')
ax.add_collection(lc)
ax.set_xlim(-10, 110)
ax.set_ylim(-1.1, 1.1)
ax.set_zlim(-1.1, 1.1)
plt.show()
