show matplotlib colorbar instead of legend for multiple plots with gradually changing colors

Question

I am trying to make a simple plot that shows a lot of curves that belong to a large dataset with a certain order, let's say plot 1..n. The shape of the curves changes gradually with increasing n. It is not important that readers can see exactly which plot belongs to which value of n, but they should be able to guess in which order of magnitude n is.

I therefore do something like this:

nValues = range(0,30)
xValues = np.linspace(0,10)
dataset = [(xValues-5-0.5*n)**2 for n in nValues]
colors = {n: colorsys.hsv_to_rgb(hue,0.9,0.7) for n,hue in zip(nValues,np.linspace(0,0.7,len(nValues)))}
for n in nValues:
    plt.plot(dataset[n],color=colors[n])

(Just to be clear, this is just for the example, my data is actually stored in a nice pandas dataframe.)

Instead of a legend, I would like to add a colorbar next to the plot with perhaps a couple of tickmarks and labels (at least the minimum and the maximum), to indicate which color belongs to which value of n, but I don't know how to accomplish this. I figured things might be easier if I actually get my plot colors from a ColorMap, but I also don't know how to do that and I also wouldn't know how to proceed from there.

Any pointers are welcome!

Answer 1

Both @tom and @Joe Kington are right: this has been asked before. However, I tried to make an example with slighty less efforts as the linked answers. To use a colormap (which always maps values from [0,1] to color), you first need to normalize your data. For that you can use the Normalize class. The ScalarMappable instance is only required if you intend to create a colorbar.

import matplotlib.pyplot as plt
import matplotlib.colors as mcolors
import matplotlib.cm as cm
import numpy as np

# your dataset
nValues = np.arange(0,30)
xValues = np.linspace(0,10)
dataset = [(xValues-5-0.5*n)**2 for n in nValues]

# setup the normalization and the colormap
normalize = mcolors.Normalize(vmin=nValues.min(), vmax=nValues.max())
colormap = cm.jet

# plot
for n in nValues:
    plt.plot(dataset[n], color=colormap(normalize(n)))

# setup the colorbar
scalarmappaple = cm.ScalarMappable(norm=normalize, cmap=colormap)
scalarmappaple.set_array(nValues)
plt.colorbar(scalarmappaple)

# show the figure
plt.show()

Result:

Answer 2

This is almost a duplicate of a few other questions. The key is that matplotlib needs a ScalarMappable instance (usually an image, scatter plot, etc) to make a colormap from. It's straightforward to fake one if you're not using a plotting method that creates one. You'll need a Normalize instance to define the min/max/scaling/etc of the colormap and a Colormap instance to define the colors.

However, you have an additional wrinkle. You're not using a colormap, so you'll need to build one.

Here's an example with a discrete colormap:

import numpy as np
import matplotlib.colors as mcolors
import matplotlib.cm
import matplotlib.pyplot as plt

# Your example...
nvalues = range(0,30)
xvalues = np.linspace(0,10)

hsv2rgb = lambda hue: mcolors.hsv_to_rgb([hue,0.9,0.7])
hues = np.linspace(0, 0.7, len(nvalues))
colors = [hsv2rgb(hue) for hue in hues]

dataset = [(xvalues-5-0.5*n)**2 for n in nvalues]

fig, ax = plt.subplots()
for n in nvalues:
    ax.plot(dataset[n], color=colors[n])

# Fake a ScalarMappable so you can display a colormap
cmap, norm = mcolors.from_levels_and_colors(range(len(nvalues) + 1), colors)
sm = matplotlib.cm.ScalarMappable(cmap=cmap, norm=norm)
sm.set_array([])

fig.colorbar(sm)

plt.show()

And if you'd prefer a continuous colormap:

import numpy as np
import matplotlib.colors as mcolors
import matplotlib.cm
import matplotlib.pyplot as plt

# Your example...
nvalues = range(0,30)
xvalues = np.linspace(0,10)

hsv2rgb = lambda hue: mcolors.hsv_to_rgb([hue,0.9,0.7])
hues = np.linspace(0, 0.7, len(nvalues))
colors = [hsv2rgb(hue) for hue in hues]

dataset = [(xvalues-5-0.5*n)**2 for n in nvalues]

fig, ax = plt.subplots()
for n in nvalues:
    ax.plot(dataset[n], color=colors[n])

# Fake a ScalarMappable so you can display a colormap
cmap = mcolors.LinearSegmentedColormap.from_list('my_cmap', colors)
norm = mcolors.Normalize(min(nvalues), max(nvalues))
sm = matplotlib.cm.ScalarMappable(cmap=cmap, norm=norm)
sm.set_array([])

fig.colorbar(sm)

plt.show()