[FIXED] How are the values from c variable in scatter converted?

Issue

I am creating a PCA plot using matplotlib Python library, I choose the color of each point according to a class value (which is 0, 1 or 2). To do so I am using the parameter called c:

plt.scatter(pca_data[:, 0], pca_data[:, 1], c=[0,1,1,0,2])

What I would like to do is add a legend linking the colour from the plot with a label, however, I cannot find how the values I give (0, 1 or 2) are converted to actual colors.

I thought about giving colors directly but I would like the process to be automated so it would work no matter the actual number of classes.

I tried using to_rgb (from matplotlib.colors) but since the values are not between 0 and 1 it does not work and if I scale them I end up with an odd color vector (black, white, black).

Solution

Here is an example where there are 3 labels (but the method is scalable to any number):

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

pca_data = np.random.randn(100,2) #some fake data
labels = np.random.randint(low=0,high=3,size=100) #some fake class labels

cNorm = colors.Normalize(vmin=0,vmax=2) #normalise the colormap
scalarMap = cm.ScalarMappable(norm=cNorm,cmap='hot') #map numbers to colors

fig,ax = plt.subplots()
for i in np.unique(labels):
    ax.scatter(pca_data[labels==i,0],pca_data[labels==i,1],\
               c=scalarMap.to_rgba(i),label="class {}".format(i),s=100)

ax.legend(scatterpoints=1)

Answered By - Angus Williams