0

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()
JohanC
  • 71,591
  • 8
  • 33
  • 66
  • 3
    Apart from implementation details - how do you imagine the mapping here? You have `n-1` lines but `n` color values. The linked question is about a scatter plot - `n` markers, `n` color values. – Mr. T Feb 22 '21 at 12:10
  • @Mr.T Oh, I see what you mean. Then my other question is, is it possible to implement [this](https://matplotlib.org/3.3.2/gallery/lines_bars_and_markers/multicolored_line.html) in this 3d case? – Amirhossein Rezaei Feb 22 '21 at 12:35

1 Answers1

2

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()

example plot

JohanC
  • 71,591
  • 8
  • 33
  • 66