UMAP (Uniform Manifold Approximation and Projection) › Implementation in Scikit-learn
Machine Learning from ScratchAll chapters UMAP (Uniform Manifold Approximation and Projection) › Implementation in Scikit-learnOne section · Sections in this chapter5 / 8 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
UMAP is implemented in the umap-learn library, which provides a scikit-learn compatible interface. Let's demonstrate how to use it effectively with a real-world example.
Loading and Preparing Data
We'll use the digits dataset, which contains 8×8 pixel images of handwritten digits (0-9). This is a good example because it has high-dimensional data (64 features) that benefits from dimensionality reduction.
from sklearn.datasets import load_digits
# Load the digits dataset
digits = load_digits()
X, y = digits.data, digits.targetDigits dataset shape: (1797, 64) Number of classes: 10
The dataset contains 1,797 samples with 64 features each (representing the 8×8 pixel values). We have 10 classes corresponding to digits 0 through 9.
Preprocessing
UMAP is sensitive to feature scales, so we standardize the features before applying dimensionality reduction.
# Preprocess the data
scaler = StandardScaler()
X_scaled = scaler.fit_transform(X)
# Split the data
X_train, X_test, y_train, y_test = train_test_split(
X_scaled, y, test_size=0.2, random_state=42
)Training set shape: (1437, 64) Test set shape: (360, 64)
We've split the data into 80% training (1,437 samples) and 20% test (360 samples) sets. Standardization ensures all features have zero mean and unit variance.
Applying UMAP
Now we apply UMAP to reduce the 64-dimensional data to 2 dimensions for visualization. We set n_neighbors=15 to balance local and global structure, and min_dist=0.1 to allow moderate clustering.
# Apply UMAP for dimensionality reduction
reducer = UMAP(
n_components=2, # Reduce to 2D for visualization
n_neighbors=15, # Number of neighbors to consider
min_dist=0.1, # Minimum distance between points
metric="euclidean", # Distance metric
random_state=42,
)
# Fit UMAP on training data
X_train_umap = reducer.fit_transform(X_train)
X_test_umap = reducer.transform(X_test)UMAP training embedding shape: (1437, 2) UMAP test embedding shape: (360, 2)
The embedding successfully reduces the dimensionality from 64 to 2 for both training and test sets. Note that we use fit_transform on training data to learn the embedding, then transform on test data to project it into the same embedding space.
Visualizing the Embedding
Let's visualize the 2D UMAP embedding to see how well it separates the different digit classes.
The visualization shows that UMAP successfully groups similar digits together in the 2D space. We can see distinct clusters for each digit class, with some expected overlaps (e.g., between digits that look similar like 3 and 8, or 4 and 9). The test data embedding matches the training data structure, indicating that the learned manifold generalizes well.
Using UMAP for Classification
To demonstrate UMAP's value as a preprocessing step, we'll train a classifier on the reduced 2D representation and compare it to the original 64-dimensional data.
# Train a classifier on the UMAP embedding
rf_umap = RandomForestClassifier(n_estimators=100, random_state=42)
rf_umap.fit(X_train_umap, y_train)
# Make predictions
y_pred_umap = rf_umap.predict(X_test_umap)
accuracy_umap = accuracy_score(y_test, y_pred_umap)
# For comparison, train on original data
rf_original = RandomForestClassifier(n_estimators=100, random_state=42)
rf_original.fit(X_train, y_train)
y_pred_original = rf_original.predict(X_test)
accuracy_original = accuracy_score(y_test, y_pred_original)Classification Performance: Accuracy on UMAP embedding (2D): 0.947 Accuracy on original data (64D): 0.972 Dimensionality reduction: 64 → 2 features (96.9% reduction)
The classifier achieves competitive accuracy even with 97% fewer features. While the original 64-dimensional data provides slightly better accuracy, the 2D UMAP embedding retains most of the discriminative information while being much faster to train and easier to visualize.
Detailed Classification Report (UMAP embedding):
precision recall f1-score support
0 0.97 1.00 0.99 33
1 0.90 1.00 0.95 28
2 1.00 0.91 0.95 33
3 0.92 1.00 0.96 34
4 1.00 1.00 1.00 46
5 0.95 0.89 0.92 47
6 0.97 0.97 0.97 35
7 1.00 0.94 0.97 34
8 0.82 0.90 0.86 30
9 0.92 0.88 0.90 40
accuracy 0.95 360
macro avg 0.95 0.95 0.95 360
weighted avg 0.95 0.95 0.95 360The classification report shows per-class performance. Most digits achieve high precision and recall, with the model performing well across all classes. Any confusion between classes reflects the inherent similarity between certain digit shapes in the 2D projection.
Alternative Implementation
We can also implement a simplified version of UMAP from scratch to better understand the algorithm:
def simple_umap(
X, n_components=2, n_neighbors=15, min_dist=0.1, random_state=None
):
"""
Simplified UMAP implementation for educational purposes.
This is a basic version that demonstrates the core concepts.
"""
if random_state is not None:
np.random.seed(random_state)
n_samples, n_features = X.shape
# Step 1: Find nearest neighbors
from sklearn.neighbors import NearestNeighbors
nbrs = NearestNeighbors(n_neighbors=n_neighbors).fit(X)
distances, indices = nbrs.kneighbors(X)
# Step 2: Compute local distances and sigmas
rhos = distances[:, 1] # Distance to nearest neighbor (excluding self)
sigmas = np.zeros(n_samples)
for i in range(n_samples):
# Binary search for sigma
low, high = 0, 10
for _ in range(20): # Max iterations
mid = (low + high) / 2
membership_strengths = np.exp(-(distances[i, 1:] - rhos[i]) / mid)
if np.sum(membership_strengths) < np.log2(n_neighbors):
low = mid
else:
high = mid
sigmas[i] = (low + high) / 2
# Step 3: Compute fuzzy simplicial set
fuzzy_simplicial_set = np.zeros((n_samples, n_samples))
for i in range(n_samples):
for j in range(1, n_neighbors): # Skip self
neighbor_idx = indices[i, j]
if distances[i, j] - rhos[i] <= 0:
fuzzy_simplicial_set[i, neighbor_idx] = 1.0
else:
fuzzy_simplicial_set[i, neighbor_idx] = np.exp(
-(distances[i, j] - rhos[i]) / sigmas[i]
)
# Step 4: Symmetrize the fuzzy simplicial set
fuzzy_simplicial_set = np.maximum(
fuzzy_simplicial_set, fuzzy_simplicial_set.T
)
# Step 5: Initialize low-dimensional embedding
Y = np.random.normal(0, 0.1, (n_samples, n_components))
# Step 6: Optimize embedding (simplified gradient descent)
learning_rate = 0.01
n_epochs = 200
for epoch in range(n_epochs):
# Compute pairwise distances in low-dimensional space
Y_distances = np.sqrt(
np.sum((Y[:, np.newaxis, :] - Y[np.newaxis, :, :]) ** 2, axis=2)
)
# Compute membership strengths in low-dimensional space
a, b = 1.0, 1.0 # Simplified parameters
low_dim_membership = 1 / (1 + a * Y_distances ** (2 * b))
# Compute gradient (simplified)
gradient = np.zeros_like(Y)
for i in range(n_samples):
for j in range(n_samples):
if i != j:
diff = Y[i] - Y[j]
dist = Y_distances[i, j]
if dist > 0:
weight = (
fuzzy_simplicial_set[i, j]
- low_dim_membership[i, j]
)
gradient[i] += weight * diff / dist
# Update embedding
Y -= learning_rate * gradient
return Y
# Test the simplified implementation
X_simple = np.random.randn(100, 10)
Y_simple = simple_umap(
X_simple, n_components=2, n_neighbors=15, random_state=42
)
print(f"Simplified UMAP embedding shape: {Y_simple.shape}")
Key Parameters
Below are the main parameters that control how UMAP works and performs. Understanding these parameters helps you tune UMAP for your specific use case.
n_components: Number of dimensions in the low-dimensional embedding (default: 2). Common values are 2 for visualization or 10-50 for preprocessing before other algorithms. Higher values preserve more information but lose visualization benefits.
n_neighbors: Number of nearest neighbors to consider when building the high-dimensional graph (default: 15). This parameter controls the balance between local and global structure:
- Smaller values (5-10): Focus on local structure, revealing fine details and small clusters
- Medium values (15-30): Balance local and global structure, suitable for most applications
- Larger values (50-100): Emphasize global structure, showing broad patterns
min_dist: Minimum distance between points in the low-dimensional embedding (default: 0.1). Controls how tightly points pack together:
- Small values (0.0-0.1): Allow tight clustering, good for finding detailed structure
- Medium values (0.1-0.5): Balance clustering and separation, suitable for most visualizations
- Large values (0.5-0.99): Spread points out more evenly, favoring broad organization
metric: Distance metric for computing distances in the high-dimensional space (default: 'euclidean'). Options include 'euclidean', 'manhattan', 'cosine', 'correlation', and many others. Choose based on your data type:
- Euclidean: General-purpose, works well for continuous features
- Cosine: Good for text data or when magnitude doesn't matter
- Manhattan: Less sensitive to outliers in high dimensions
random_state: Seed for reproducibility (default: None). Set to an integer to ensure consistent results across runs.
n_epochs: Number of optimization iterations (default: auto-determined based on dataset size). More epochs improve quality but increase computation time. The default usually works well.
learning_rate: Step size for gradient descent optimization (default: 1.0). Rarely needs adjustment, but lower values (0.1-0.5) can help if the embedding quality is poor.
low_memory: Whether to use a lower memory, but longer-running implementation (default: False). Set to True for very large datasets that don't fit in memory.
Key Methods
The following are the most commonly used methods for interacting with UMAP.
fit(X, y=None): Learns the UMAP embedding from training data X. The optional y parameter can provide supervised information to guide the embedding.
fit_transform(X, y=None): Fits the model and returns the transformed data in one step. Use this for the initial embedding of your training data.
transform(X): Projects new data into the existing embedding space. Use this for test data or new observations after fitting on training data. Note that transform may be less accurate than the original embedding.
inverse_transform(X_embedded): Attempts to reconstruct the original high-dimensional data from the low-dimensional embedding. This is approximate and works better for higher-dimensional embeddings.