[python] Matplotlib color according to class labels

I have two vectors, one with values and one with class labels like 1,2,3 etc.

I would like to plot all the points that belong to class 1 in red, to class 2 in blue, to class 3 in green etc. How can I do that?

Assuming that you have your data in a 2d array, this should work:

import numpy
import pylab
xy = numpy.zeros((2, 1000))
xy[0] = range(1000)
xy[1] = range(1000)
colors = [int(i % 23) for i in xy[0]]
pylab.scatter(xy[0], xy[1], c=colors)
pylab.show()

You can also set a cmap attribute to control which colors will appear through use of a colormap; i.e. replace the pylab.scatter line with:

pylab.scatter(xy[0], xy[1], c=colors, cmap=pylab.cm.cool)

A list of color maps can be found here

A simple solution is to assign color for each class. This way, we can control how each color is for each class. For example:

arr1 = [1, 2, 3, 4, 5]
arr2 = [2, 3, 3, 4, 4]
labl = [0, 1, 1, 0, 0]
color= ['red' if l == 0 else 'green' for l in labl]
plt.scatter(arr1, arr2, color=color)