Visualizing K-means Clustering on an Iris Dataset with Matplotlib

Topic Overview

Welcome to this lesson on "Visualizing Clusters with Matplotlib using an Iris Dataset". Previously, we introduced unsupervised learning, with a focus on clustering and the K-means clustering algorithm. In this unit, we will visualize the clusters resulting from the K-means algorithm using Python's Matplotlib library. The goal of this lesson is to illustrate the implementation of the K-means clustering algorithm and to demonstrate how to visualize the results using Matplotlib, utilizing a simple Iris dataset as an example.

Loading the Iris Dataset

In this lesson, we will focus on the scatter() function, which allows us to create scatter plots, and the show() function, which will enable us to display the plot.

We begin by loading the Iris dataset. It is a classic and widely used dataset in pattern recognition, consisting of 150 samples from three species of Iris flowers (Iris setosa, Iris virginica, and Iris versicolor), with four features measured for each sample: the length and the width of the sepals and petals.

Python
# Required Libraries
from sklearn.datasets import load_iris

# Load the Iris dataset
iris = load_iris()
data = iris.data

Utilizing Sklearn for K-means Clustering and Visualization

Let's see how to apply KMeans from sklearn and visualize results using Matplotlib. Horizontal and vertical indicators alone may not provide all the clarity we need in a plot. That's where Python's Matplotlib shines — its assortment of plot customization capabilities. Let's delve into customizing the colormap and other parameters for your plots for improving visualization and readability.

Before jumping into plotting, let's cluster our data first:

Python
from sklearn.cluster import KMeans
import matplotlib.pyplot as plt

# Perform K-means clustering and take cluster centers
kmeans_model = KMeans(n_clusters=2, random_state=1, n_init=10).fit(data)
labels = kmeans_model.labels_
clusters_sklearn = kmeans_model.cluster_centers_

A colormap is like an artist's palette, essentially mapping values to colors on a plot. Matplotlib offers a variety of default colormaps. The scatter() function uses a colormap to define the colors of markers.

The colormap can be specified using the parameter cmap of the scatter() method.

The c parameter passed to the scatter() method is a list of the same length as data that specifies the color of each point.

Let's customize our plot:

Python
# Function to plot final clusters
def plot_clusters_sklearn(data, labels, clusters):
    # Plot data points, choosing first two features for visualization
    plt.scatter(data[:, 0], data[:, 1], c=labels, cmap='viridis', label='Data points')
    # Plot cluster centers, also focusing on the same two features
    plt.scatter(clusters[:, 0], clusters[:, 1], s=200, color='red', marker='X', label='Centers')
    plt.title('Visualizing Clusters with Matplotlib using Iris Dataset')
    plt.xlabel('Sepal Length (cm)')
    plt.ylabel('Sepal Width (cm)')
    plt.legend()
    plt.grid(True)
    plt.show()

# Visualize the clusters
plot_clusters_sklearn(data, labels, clusters_sklearn)

With Sklearn's implementation, we instantiate the KMeans model with the desired number of clusters, then fit it to our data. The labels of the clusters for each data point are obtained via the labels_ attribute of the model, while the cluster_centers_ attribute gives us the centroid of these clusters. These are then visualized through the scatter function in Matplotlib just as we did before. Check out the plot presented below:

image

Sign up

Join the 1M+ learners on CodeSignal

Be a part of our community of 1M+ users who develop and demonstrate their skills on CodeSignal