Reproduces the spectral clustering example of Figure 14.29 in Hastie, Tibshirani and Friedman (2nd ed).

spectral_clustering.py

# Reproduce the spectral clustering example of Figure 14.29.

import numpy as np
import matplotlib.pyplot as plt
from sklearn.cluster import SpectralClustering


def sample(n):
    assert n % 3 == 0
    
    angles = np.random.uniform(0, 2*np.pi, size=(n,1))
    radii = np.vstack((np.ones((n/3,1)), 2.8*np.ones((n/3,1)), 5*np.ones((n/3,1))))
    x = radii*np.cos(angles) + np.random.normal(scale=0.25, size=(n,1))
    y = radii*np.sin(angles) + np.random.normal(scale=0.25, size=(n,1))
    
    return np.hstack((x, y))


if __name__ == "__main__":
    data = sample(450)
    clusters = 3
    sc = SpectralClustering(n_clusters=clusters, n_neighbors=10,
                            affinity="nearest_neighbors")
    labels = sc.fit_predict(data)
    
    pts = []
    for c in xrange(clusters):
        selector = (labels == c)
        pts.append(data[selector, :])
    
    for c in xrange(clusters):
        plt.plot(pts[c][:,0], pts[c][:,1], 'o')
    
    plt.axes().set_aspect("equal")
    plt.tight_layout()
    plt.show()