Colormap a 3D curve in matplotlib

Question

I have 4 arrays x, y, z and T of length n and I want to plot a 3D curve using matplotlib. The (x, y, z) are the points positions and T is the value of each point (which is plotted as color), like the temperature of each point. How can I do it?

Example code:

import numpy as np
from matplotlib import pyplot as plt

fig = plt.figure()
ax = fig.gca(projection='3d')
n = 100
cmap = plt.get_cmap("bwr")
theta = np.linspace(-4 * np.pi, 4 * np.pi, n)
z = np.linspace(-2, 2, n)
r = z**2 + 1
x = r * np.sin(theta)
y = r * np.cos(theta)
T = (2*np.random.rand(n) - 1)  # All the values are in [-1, 1]

What I found over the internet:

It's possible to use cmap with scatter like shown in the docs and in this stackoverflow question

ax = plt.gca()
ax.scatter(x, y, z, cmap=cmap, c=T)

The problem is that scatter is a set of points, not a curve.

In this stackoverflow question the solution was divide in n-1 intervals and each interval we use a different color like

t = (T - np.min(T))/(np.max(T)-np.min(T))  # Normalize
for i in range(n-1):
    plt.plot(x[i:i+2], y[i:i+2], z[i:i+2], c=cmap(t[i])

The problem is that each segment has only one color, but it should be an gradient. The last value is not even used.

Useful links:

What's the problem with the first solution with scatter and cmap? Isn't that exactly what you want ? — Rabinzel, Jul 15 '22 at 16:38

score 0 · Answer 1 · answered Jul 16 '22 at 07:55

This is a case where you probably need to use Line3DCollection. This is the recipe:

create segments from your array of coordinates.
create a Line3DCollection object.
add that collection to the axis.
set the axis limits.

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d.art3d import Line3DCollection
from matplotlib.cm import ScalarMappable
from matplotlib.colors import Normalize

def get_segments(x, y, z):
    """Convert lists of coordinates to a list of segments to be used
    with Matplotlib's Line3DCollection.
    """
    points = np.ma.array((x, y, z)).T.reshape(-1, 1, 3)
    return np.ma.concatenate([points[:-1], points[1:]], axis=1)

fig = plt.figure()
ax = fig.add_subplot(projection='3d')
n = 100
cmap = plt.get_cmap("bwr")
theta = np.linspace(-4 * np.pi, 4 * np.pi, n)
z = np.linspace(-2, 2, n)
r = z**2 + 1
x = r * np.sin(theta)
y = r * np.cos(theta)
T = np.cos(theta)

segments = get_segments(x, y, z)
c = Line3DCollection(segments, cmap=cmap, array=T)
ax.add_collection(c)
fig.colorbar(c)

ax.set_xlim(x.min(), x.max())
ax.set_ylim(y.min(), y.max())
ax.set_zlim(z.min(), z.max())
plt.show()

I tested and this solution is the same as the second one. The difference is: this one is faster cause we don't call ```plt.plot``` many times. With a small value of ```n``` or when ```T``` is very random, we still see the color segments. — Carlos Adir, Jul 16 '22 at 17:23

Colormap a 3D curve in matplotlib

1 Answers1