Plot one curve in different colors using Python

In order to plot a curve in different colors and the colors vary with some parameters, it seems that the best way is to use the LineCollection module of python.

In my problem I have an array residuals_resume_utdates_geoindex: the first column is dates (x), the second are residuals (y), and the third is the parameter to condition the residuals. I want to plot x, y in different colors when the values in the third column vary. In fact the values are between 0 and 5 (integer) in the third column. If the value in this column is 0, then the corresponding residuals should be plotted in one color; if the value is 1, a different color should be used for the corresponding residuals.

Here is my code:


import numpy as np
import matplotlib.pyplot as plt
from matplotlib.collections import LineCollection
from matplotlib.colors import ListedColormap, BoundaryNorm


ydata = residuals_resume_utdates_geoindex[:,1]
xdata = residuals_resume_utdates_geoindex[:,0]
# put x and y together as an array of N*1*2, where N is the number of data in ydata
points = np.array([xdata, ydata]).T.reshape(-1, 1, 2)
# create a line segment array in order to color them individually
# the segment array has the shape (N-1)*2*2 for line collection
# [ [[x0,y0],[x1,y1]],[[x1,y1],[x2,y2]],[].....]
segments = np.concatenate([points[:-1], points[1:]], axis=1)


# plot y with different color for each 6 indices (5,4,3,2,1,0)
# if 0: black, 1: yello, 2 cyron, 3 megan, 4 green, 5 red, biger than 5: blue
cmap = ListedColormap(['k','y','c','m','g', 'r', 'b'])
norm = BoundaryNorm([-1, 0.5, 1.5, 2.5, 3.5, 4.5, 5.5, 6], cmap.N)
lc = LineCollection(segments, cmap=cmap, norm=norm)
# define the color for each segment which depends on the value of the third column
color_seg = residuals_resume_utdates_geoindex[:,-1]
# set the linecollection property: color, linestyle
lc.set_array(color_seg)
lc.set_linewidth(3)
lc.set_linestyle('dashdot')

#plot the figure with linecollection
fig =  plt.figure()
plt.gca().add_collection(lc)
plt.xlim(np.min(xdata), np.max(xdata))
plt.ylim(0, 1)