Create legend for scatter plot using the label of the samples in matplotlib

Question

I am using scatter plot in matplotlib to plot some points. I have two 1D arrays each storing the x and y coordinate of the samples. Also there is another 1D array that stores the label(to decide in which colour the point should be plotted). I programmed thus far:

import matplotlib.pyplot as plt
X = [1,2,3,4,5,6,7]
Y = [1,2,3,4,5,6,7]
label = [0,1,4,2,3,1,1]
plt.scatter(X, Y, c= label, s=50)
plt.show()

Now I want to be able to see which color corresponds to which label? I looked up the implementation of legends in matplotlib like the one here: how to add legend for scatter()? However they are suggesting to create a plot for each label of sample. However all my labels are in the same 1D array(label). How can I achieve this?

Yes I agree its a duplicate. I had gone through stackoverflow.com/questions/14777066/… but could not relate it to my requirement. I think my question is quite suggestive and would help narrow down the search for answer to this question in future. — AviB

tmdavison tmdavison · Accepted Answer · 2015-05-28T12:30:28

You could do it with a colormap. Some examples of how to do it are here.

import matplotlib.pyplot as plt
import numpy as np
import matplotlib.colors as colors
X = [1,2,3,4,5,6,7]
Y = [1,2,3,4,5,6,7]
label = [0,1,4,2,3,1,1]

# Define a colormap with the right number of colors
cmap = plt.cm.get_cmap('jet',max(label)-min(label)+1)

bounds = range(min(label),max(label)+2)
norm = colors.BoundaryNorm(bounds, cmap.N)

plt.scatter(X, Y, c= label, s=50, cmap=cmap, norm=norm)

# Add a colorbar. Move the ticks up by 0.5, so they are centred on the colour.
cb=plt.colorbar(ticks=np.array(label)+0.5)
cb.set_ticklabels(label)

plt.show()

You might need to play around to get the tick labels centred on their colours, but you get the idea.

enter image description here

Create legend for scatter plot using the label of the samples in matplotlib

1 Answers