umap_digits_dimensionality_reduction_2d_visualization.py

1import umap
2from sklearn.datasets import load_digits
3import matplotlib.pyplot as plt
4import seaborn as sns
5
6# Set up plotting aesthetics
7sns.set(style='white', context='notebook', rc={'figure.figsize':(14,10)})
8
9# Load the digits dataset (1797 samples, 64 features each)
10digits = load_digits()
11
12# Initialize UMAP and fit_transform the data
13# This reduces the data from 64 dimensions to 2
14reducer = umap.UMAP()
15embedding = reducer.fit_transform(digits.data)
16
17# Verify the shape of the resulting embedding
18print(f"Original shape: {digits.data.shape}")
19print(f"Reduced shape: {embedding.shape}")
20
21# Visualize the results
22plt.scatter(
23    embedding[:, 0],
24    embedding[:, 1],
25    c=digits.target,
26    cmap='Spectral',
27    s=5
28)
29plt.gca().set_aspect('equal', 'datalim')
30plt.colorbar(boundaries=range(11))
31plt.title('UMAP projection of the Digits dataset', fontsize=24)
32plt.show()