CODEWITHSUNDEEP
pythonmatplotlibseaborn
akshay_rao
akshay_rao

Reputation: 75

ScatterPlot colouring and labelling with Clustering in Python

I am trying to cluster my results. I get into 3 clusters along with label names using matplotlib:

Y_sklearn - 2 dimensional array contains X and Y coordinates

ent_names - contains label names

I have a logic to display scatter plot as follows:

from sklearn.cluster import KMeans
model = KMeans(n_clusters = 3)
model.fit(Y_sklearn)
plt.scatter(Y_sklearn[:,0],Y_sklearn[:,1], c=model.labels_); 
plt.show()

Now the above code does display the scatter plot as follows: enter image description here

However, along with this plot I want to display the label names as well. I tried something like this but it displays only one color for all:

with plt.style.context('seaborn-whitegrid'):
   plt.figure(figsize=(8, 6))
   for lab, col in zip(ent_names,
                   model.labels_):
       plt.scatter(Y_sklearn[y==lab, 0],
               Y_sklearn[y==lab, 1],
               label=lab,
               c=model.labels_)
   plt.xlabel('Principal Component 1')
   plt.ylabel('Principal Component 2')
   plt.legend(loc='lower center')
   plt.tight_layout()
   plt.show()

Upvotes: 3

Views: 6692

Answers (1)

Serenity
Serenity

Reputation: 36635

You have to plot in the same axes ax to put scatter plots together like in this example:

import matplotlib.pyplot as plt
import numpy as np

XY = np.random.rand(10,2,3)

labels = ['test1', 'test2', 'test3']
colors = ['r','b','g']

with plt.style.context('seaborn-whitegrid'):
    plt.figure(figsize=(8, 6))
    ax = plt.gca()

    i=0
    for lab,col in zip(labels, colors):
        ax.scatter(XY[:,0,i],XY[:,1,i],label=lab, c=col)
        i+=1

    plt.xlabel('Principal Component 1')
    plt.ylabel('Principal Component 2')
    plt.legend(loc='lower center')
    plt.tight_layout()
    plt.show()

enter image description here

Upvotes: 1

Related Questions