3

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?

Community
  • 1
  • 1
AviB
  • 95
  • 3
  • 11
  • possible duplicate of [Matplotlib discrete colorbar](http://stackoverflow.com/questions/14777066/matplotlib-discrete-colorbar) – tmdavison May 28 '15 at 12:15
  • Yes I agree its a duplicate. I had gone through http://stackoverflow.com/questions/14777066/matplotlib-discrete-colorbar 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 May 29 '15 at 05:18

1 Answers1

2

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

Community
  • 1
  • 1
tmdavison
  • 64,360
  • 12
  • 187
  • 165
  • Great! Now instead of showing 0,1....4 on the legend, can I show class0, class1,... class4? I know which number corresponds to which class. However my label array will be [0,0,1,2,4] only. Can I map this to class0, class1 etc? – AviB May 29 '15 at 05:02