How to add colorbar in matplotlib

Question

I don't really understand the documentation on colorbar I wanted explanation on a basic example. So below I am plotting exp(-x*a) for a={1, 2, 3, 4}. How do I add color bar with the values of a.

import numpy as np
import matplotlib.pyplot as plt

def func(x,a):
    return np.exp(-x*a)
x = np.linspace(0, 5)

for a in range(1,5):
    plt.plot(func(x,a))
plt.show()

I'm confused with imshow, subplot, ax.

Answer 1

Colour schemes are defined in the range 0 to 1, therefore you first need to normalise your values (0 to 5) to the range 0 to 1. Then you can pull the colour from the colormap. At the end you have to plot a color bar using the colour map and norm that you chose on the axis you used for plotting.

import numpy as np
import matplotlib.pyplot as plt
from matplotlib import cm, colors


def func(x, a):
    return np.exp(-x * a)


x = np.linspace(0, 5)

fig, ax = plt.subplots()

# define color map
cmap = cm.get_cmap("Spectral")

# need to normalize because color maps are defined in [0, 1]
norm = colors.Normalize(0, 5)

for a in range(1, 5):
    ax.plot(x, func(x, a), 
        color=cmap(norm(a)))  # get color from color map 

# plot colorbar
fig.colorbar(cm.ScalarMappable(norm=norm, cmap=cmap), ax=ax)

plt.show()

The plot looks like this: