UMAP-Learn
Overview
UMAP (Uniform Manifold Approximation and Projection) is a dimensionality reduction algorithm for visualization and general non-linear dimensionality reduction. It is faster than t-SNE, scales to larger datasets, preserves both local and global structure, and supports supervised learning and embedding of new data points.
When to Use
- Reducing high-dimensional data to 2D/3D for visualization
- Preprocessing for density-based clustering (HDBSCAN, DBSCAN)
- Feature engineering in ML pipelines (transform new data into learned embedding)
- Supervised/semi-supervised embedding with partial labels
- Tracking embeddings across time points or batches (AlignedUMAP)
- Density-preserving embeddings (DensMAP)
- Neural network-based embedding with custom architectures (Parametric UMAP)
- For linear dimensionality reduction use PCA (scikit-learn)
- For neighborhood-graph construction without embedding use scikit-learn NearestNeighbors
Prerequisites
pip install umap-learn
# For Parametric UMAP (neural network variant)
pip install umap-learn[parametric_umap] # requires TensorFlow 2.x
Critical: Always standardize features before applying UMAP to ensure equal weighting across dimensions.
Quick Start
import umap
import numpy as np
from sklearn.preprocessing import StandardScaler
from sklearn.datasets import load_digits
# Load and scale data
X, y = load_digits(return_X_y=True)
X_scaled = StandardScaler().fit_transform(X)
# Fit and transform
embedding = umap.UMAP(random_state=42).fit_transform(X_scaled)
print(f"Input: {X_scaled.shape}, Output: {embedding.shape}")
# Input: (1797, 64), Output: (1797, 2)
Core API
1. Standard UMAP
Basic dimensionality reduction following scikit-learn conventions.
import umap
from sklearn.preprocessing import StandardScaler
X_scaled = StandardScaler().fit_transform(data)
# Method 1: fit_transform (single step)
embedding = umap.UMAP(
n_neighbors=15, # local neighborhood size (2-200)
min_dist=0.1, # min distance between embedded points (0.0-0.99)
n_components=2, # output dimensions
metric='euclidean', # distance metric
random_state=42, # reproducibility
).fit_transform(X_scaled)
print(f"Embedding shape: {embedding.shape}")
# Method 2: fit + access (for reuse)
reducer = umap.UMAP(random_state=42)
reducer.fit(X_scaled)
embedding = reducer.embedding_ # trained embedding
graph = reducer.graph_ # fuzzy simplicial set (sparse matrix)
# Visualization
import matplotlib.pyplot as plt
plt.figure(figsize=(8, 6))
plt.scatter(embedding[:, 0], embedding[:, 1], c=labels, cmap='Spectral', s=5)
plt.colorbar()
plt.title('UMAP Embedding')
plt.tight_layout()
plt.savefig('umap_embedding.png', dpi=150)
2. Supervised & Semi-Supervised UMAP
Incorporate label information to guide embedding via the y parameter.
import umap
# Supervised — all labels known
embedding = umap.UMAP(random_state=42).fit_transform(X_scaled, y=labels)
# Semi-supervised — partial labels (mark unlabeled as -1)
semi_labels = labels.copy()
semi_labels[unlabeled_indices] = -1
embedding = umap.UMAP(random_state=42).fit_transform(X_scaled, y=semi_labels)
# Control label influence with target_weight (0.0=unsupervised, 1.0=fully supervised)
reducer = umap.UMAP(
target_weight=0.7, # emphasize labels
target_metric='categorical', # for classification; use distance metric for regression
random_state=42
)
embedding = reducer.fit_transform(X_scaled, y=labels)
print(f"Supervised embedding: {embedding.shape}")
3. Transform New Data
Project unseen data into the trained embedding space.
import umap
from sklearn.preprocessing import StandardScaler
scaler = StandardScaler()
X_train_scaled = scaler.fit_transform(X_train)
X_test_scaled = scaler.transform(X_test)
# Fit on training data
reducer = umap.UMAP(n_components=10, random_state=42)
X_train_emb = reducer.fit_transform(X_train_scaled)
# Transform test data
X_test_emb = reducer.transform(X_test_scaled)
print(f"Train: {X_train_emb.shape}, Test: {X_test_emb.shape}")
# Works in sklearn Pipelines
from sklearn.pipeline import Pipeline
from sklearn.svm import SVC
pipeline = Pipeline([
('scaler', StandardScaler()),
('umap', umap.UMAP(n_components=10, random_state=42)),
('classifier', SVC())
])
pipeline.fit(X_train, y_train)
accuracy = pipeline.score(X_test, y_test)
print(f"Pipeline accuracy: {accuracy:.3f}")
4. Parametric UMAP
Neural network-based embedding via TensorFlow/Keras. Enables efficient transform, reconstruction, and custom architectures.
from umap.parametric_umap import ParametricUMAP
# Default architecture (3-layer, 100-neuron FC network)
embedder = ParametricUMAP(n_components=2, random_state=42)
embedding = embedder.fit_transform(X_scaled)
new_emb = embedder.transform(new_data) # fast neural network inference
print(f"Parametric embedding: {embedding.shape}")
import tensorflow as tf
from umap.parametric_umap import ParametricUMAP
# Custom encoder/decoder for autoencoder mode
input_dim = X_scaled.shape[1]
encoder = tf.keras.Sequential([
tf.keras.layers.InputLayer(input_shape=(input_dim,)),
tf.keras.layers.Dense(128, activation='relu'),
tf.keras.layers.Dense(64, activation='relu'),
tf.keras.layers.Dense(2),
])
decoder = tf.keras.Sequential([
tf.keras.layers.InputLayer(input_shape=(2,)),
tf.keras.layers.Dense(64, activation='relu'),
tf.keras.layers.Dense(128, activation='relu'),
tf.keras.layers.Dense(input_dim),
])
embedder = ParametricUMAP(
encoder=encoder, decoder=decoder, dims=(input_dim,),
parametric_reconstruction=True, autoencoder_loss=True,
n_training_epochs=10, batch_size=128,
n_neighbors=15, min_dist=0.1, random_state=42
)
embedding = embedder.fit_transform(X_scaled)
reconstructed = embedder.inverse_transform(embedding)
print(f"Reconstruction error: {np.mean((X_scaled - reconstructed)**2):.4f}")
5. DensMAP
Variant preserving local density information in the embedding.
import umap
reducer = umap.UMAP(
densmap=True, # enable DensMAP
dens_lambda=2.0, # density preservation weight
dens_frac=0.3, # fraction for density estimation
output_dens=True, # output density estimates
n_neighbors=15,
min_dist=0.1,
random_state=42
)
embedding = reducer.fit_transform(X_scaled)
# Access density estimates
original_density = reducer.rad_orig_ # density in original space
embedded_density = reducer.rad_emb_ # density in embedded space
print(f"DensMAP embedding: {embedding.shape}")
print(f"Density correlation: {np.corrcoef(original_density, embedded_density)[0,1]:.3f}")
6. AlignedUMAP
Align embeddings across multiple related datasets (time points, batches).
from umap import AlignedUMAP
# Multiple related datasets
datasets = [day1_data, day2_data, day3_data]
mapper = AlignedUMAP(
n_neighbors=15,
alignment_regularisation=1e-2, # alignment strength
alignment_window_size=2, # align with N adjacent datasets
n_components=2,
random_state=42
)
mapper.fit(datasets)
aligned_embeddings = mapper.embeddings_ # list of aligned embedding arrays
print(f"Aligned {len(aligned_embeddings)} datasets")
for i, emb in enumerate(aligned_embeddings):
print(f" Dataset {i}: {emb.shape}")
Key Concepts
Parameter Tuning Guide
| Parameter | Low | Medium (default) | High | Effect | |-----------|-----|-------------------|------|--------| | n_neighbors | 2-5 | 15 | 50-200 | Local detail vs global structure | | min_dist | 0.0 | 0.1 | 0.5-0.99 | Tight clusters vs spread out | | n_components | 2 | 2 | 5-50 | Visualization vs ML/clustering | | spread | 0.5 | 1.0 | 2.0 | Embedding scale (with min_dist) |
Configuration by Use-Case
| Use-Case | n_neighbors | min_dist | n_components | metric | |----------|