Plotting colored lines

Suppose we have a set of lines that we want to plot in matplotlib according to some color scheme. For example, each line might correspond to a contour and has some corresponding value which we would like to show on a color map. How can we go about doing this? This took me a while to figure out so this will be useful for my own future reference.

theta = np.linspace(0, 2*math.pi, 100)
r = np.arange(1, 11, 1)[:, np.newaxis]
x = r * np.cos(theta)
y = r * np.sin(theta)

fig, ax = plt.subplots()
ax.plot(x.T, y.T)
ax.set(xlabel="x", ylabel="y", aspect="equal")

# We want to color each circle according to the corresponding value of z
z = r[:, 0]

Here, x and y are 10x100 arrays representing 10 different circles as shown above. We want to color each line according to the value in z which is a length 10 array.

One way would just be to scatter plot things and let the scattered dots be colored according to z.

fig, ax = plt.subplots()
paths = ax.scatter(x.T, y.T, cmap="viridis", c=np.tile(z, (100, 1)))
ax.set(xlabel="x", ylabel="y", aspect="equal")
axcb = fig.colorbar(paths, ax=ax)
axcb.set_label("Radius")

But how can we connect these points together in a line? It turns out that we can't do the same thing with ax.plot as explained here. Instead we need to use LineCollections as explained in the answer. The code below is based on the solution provided in the answer linked.

from matplotlib.collections import LineCollection

fig, ax = plt.subplots()
lines = [np.column_stack([xi, yi]) for xi, yi in zip(x, y)]
lc = LineCollection(lines, cmap="viridis", lw=4)
# Set the values used to determine the color
lc.set_array(z)
ax.add_collection(lc)
ax.autoscale()
ax.set(xlabel="x", ylabel="y", aspect="equal")
axcb = fig.colorbar(lc, ax=ax)
axcb.set_label("Radius")

Animating colored lines

Now lets say we want to get fancy and we want to trace out the circles as a function of time. How can we make such an animation with our LineCollection? All we have to do is modify the above code and use FuncAnimation from matplotlib to update each frame.

import matplotlib.animation as animation


fig, ax = plt.subplots()

# First we initialize our line collection and add it to our Axes
lines = []
lc = LineCollection(lines, cmap="viridis", lw=4)
ax.add_collection(lc)


def update(num):
    new_x = x[:, :num]
    new_y = y[:, :num]

    lines = [np.column_stack([xi, yi]) for xi, yi in zip(new_x, new_y)]
    # Now we update the lines in our LineCollection
    lc.set_segments(lines)
    lc.set_array(z)
    return lc,
ani = animation.FuncAnimation(fig, update, x.shape[1], interval=100, blit=True)

# For some reason autoscaling doesn't work here so we need to set the limits manually
ax.set(xlabel="x", ylabel="y", aspect="equal", xlim=(-12, 12), ylim=(-12, 12))
# Put this after so that the colors are scaled properly
axcb = fig.colorbar(lc)
axcb.set_label("Radius")

I couldn't get rid of matplotlib showing this empty figure unfortunately... But here's the final animation!

from IPython.display import HTML
HTML(ani.to_html5_video())