python - Set line colors according to colormap

Question

I have a series of lines stored in a list like so:

line_list = [line_1, line_2, line_3, ..., line_M]

where each line_i is a sub-list composed of two sub-sub-lists, one for the x coordinates and the other for the y coordinates:

line_i = [[x_1i, x2_i, .., x_Ni], [y_1i, y_2i, .., y_Ni]]

I also have a list of the same length as line_list composed of floats,:

floats_list = [0.23, 4.5, 1.6, ..., float_M]

I want to plot each line giving it a color taken from a color map and related to the position of its index in the floats_list list. So line_j will have its color determined by the number floats_list[j]. I also need a colorbar shown to the side

The code would like something like this, except it should work :)

import matplotlib.pyplot as plt

line1 = [[0.5,3.6,4.5],[1.2,2.0,3.6]]
line2 = [[1.5,0.4,3.1,4.9],[5.1,0.2,7.4,0.3]]
line3 = [[1.5,3.6],[8.4,2.3]]

line_list = [line1,line2,line3]
floats_list = [1.2,0.3,5.6]

# Define colormap.
cm = plt.cm.get_cmap('RdYlBu')

# plot all lines.
for j,lin in enumerate(line_list): 
    plt.plot(lin[0], lin[1], c=floats_list[j])

# Show colorbar.
plt.colorbar()

plt.show()

Answer 1

It's easiest to use a LineCollection for this. In fact, it expects the lines to be in a similar format to what you already have. To color the lines by a third variable, just specify the array=floats_list. As an example:

import numpy
import matplotlib.pyplot as plt
from matplotlib.collections import LineCollection

# The line format you curently have:
lines = [[(0, 1, 2, 3, 4), (4, 5, 6, 7, 8)],
         [(0, 1, 2, 3, 4), (0, 1, 2, 3, 4)],
         [(0, 1, 2, 3, 4), (8, 7, 6, 5, 4)],
         [(4, 5, 6, 7, 8), (0, 1, 2, 3, 4)]]

# Reformat it to what `LineCollection` expects:
lines = [zip(x, y) for x, y in lines]

z = np.array([0.1, 9.4, 3.8, 2.0])

fig, ax = plt.subplots()
lines = LineCollection(lines, array=z, cmap=plt.cm.rainbow, linewidths=5)
ax.add_collection(lines)
fig.colorbar(lines)

# Manually adding artists doesn't rescale the plot, so we need to autoscale
ax.autoscale()

plt.show()

There are two main advantages of this over repeatedly calling plot.

1. Rendering speed. Collections render much faster than a large number of similar artists.
2. It's easier to color the data by another variable according to a colormap (and/or update the colormap later).