python - gnuplot linecolor variable in matplotlib?

Question

I have an array of y-values that form a line. Additionally, I have an array with the same number of elements as the y-array of values ranging from 0 to 1. We'll call this array 'z'. I want to plot the array of y-values so that the color of each point corresponds with the z-value.

In gnuplot, you can do this using the 'lc variable':

plot ’data’ using 1:2:3 with points lc variable  

Using the advice from here: Matplotlib scatterplot; colour as a function of a third variable , I was able to use a scatter plot, which did work:

import matplotlib as mpl  
import matplotlib.pyplot as plt  

plt.scatter(x, y, c=z, s=1, edgecolors='none', cmap=mpl.cm.jet)  
plt.colorbar()  
plt.show()  

Is there a way to do this with the plot method in matplotlib, similar to this?

plt.plot(x, y, c=z)

When I tried the above code, all of the lines just appeared black.

Answer 1

I had the same problem: wanted to plot line(s) with non-uniform color, which I wanted to be dependent on a third variable (z).

But I definitelly wanted to use a line, not markers (as in @joaquin's answer). I found a solution in a matplotlib gallery example, using the class matplotlib.collections.LineCollection (link here).

Here is my example, which plots trajectories in a Basemap, coloring them according to its height:

import matplotlib.pyplot as plt
from mpl_toolkits.basemap import Basemap
from matplotlib.collections import LineCollection
import numpy as np

m = Basemap(llcrnrlon=-42,llcrnrlat=0,urcrnrlon=5,urcrnrlat=50, resolution='h')
fig = plt.figure()
m.drawcoastlines()
m.drawcountries()

for i in trajectorias:
    # for each i, the x (longitude), y (latitude) and z (height)
    # are read from a file and stored as numpy arrays

    points = np.array([x, y]).T.reshape(-1, 1, 2)
    segments = np.concatenate([points[:-1], points[1:]], axis=1)

    lc = LineCollection(segments, cmap=plt.get_cmap('Spectral'),
                        norm=plt.Normalize(250, 1500))
    lc.set_array(z)
    lc.set_linewidth(2)

    plt.gca().add_collection(lc)

axcb = fig.colorbar(lc)
axcb.set_label('cota (m)')

plt.show()