High-dimensional data appears across many real-world problems, from clinical datasets to economic indicators. While these rich feature spaces capture important detail, they also make it difficult to understand relationships between observations. Uniform Manifold Approximation and Projection, or UMAP, provides a practical way to uncover the big-picture structure by mapping complex data into a lower-dimensional space that can be plotted and explored directly.
UMAP in practice
Exploratory data analysis often begins with a simple question. How are observations related across many variables? In high-dimensional settings, this question becomes difficult to answer. Distance metrics lose interpretability, and visualization is no longer feasible beyond a few dimensions.
UMAP addresses this challenge by producing embeddings where nearby points in the original space remain close in the reduced space. At the same time, meaningful groupings can emerge, though the method optimizes only for those local neighborhoods and leaves the geometry between groups to fall where it will. That focus makes it especially useful for cluster discovery and for comparing subpopulations within a dataset.
In applied settings such as country-level socioeconomic analysis, UMAP can reveal how nations with similar profiles across indicators like income, health, and demographics naturally group together, even when those relationships are not obvious in the raw data.
Figure 1 traces this journey one stage at a time, from a table with many features per observation to a two-dimensional layout where related observations gather together.
Figure 1 · Interactive
How UMAP transforms high-dimensional data
Follow the method from many features per observation through to a two-dimensional layout. Move through the five stages with the numbered steps or the Previous and Next buttons.
Key idea: UMAP connects each observation to its nearest neighbors, weights those links, merges them into one weighted graph, then places the points in two dimensions so that strongly linked observations sit close together. The interactive version with animated stages loads when JavaScript is available.
The stages preview the rest of this section. A nearest-neighbor search and a weighting step build a graph of local relationships, and an optimization step arranges the points in two dimensions. A two-dimensional embedding is a compact summary of those relationships, so it supports exploration and is best read alongside quantitative checks of the original distances.
The UMAP method
UMAP is based on manifold learning. The central assumption is that high-dimensional data lies on a lower-dimensional manifold embedded within the larger space. The goal is to learn that structure.
The method proceeds in two main steps:
• Constructing a neighborhood graph
Each data point is connected to its nearest neighbors, forming a weighted graph that represents local relationships in the original space.
• Optimizing a low-dimensional embedding
UMAP then finds a configuration in lower dimensions that preserves this graph structure as closely as possible. Points that are similar are pulled together, while dissimilar points are pushed apart.
This combination of attraction and repulsion is what settles the layout. What it optimizes for is the local neighborhoods. The geometry between groups is whatever the optimization happened to produce. Unlike linear methods such as principal component analysis (PCA), UMAP can represent non-linear relationships, which are common in real-world data.
The weighted graph rewards a closer look, because the weights are where UMAP adapts to the shape of the data. Figure 2 shows how the distance to each neighbor becomes a connection strength, and how two one-way strengths combine into a single shared weight.
Figure 2 · Interactive
From local distances to a fuzzy neighborhood graph
See how the distance to each neighbor becomes a connection strength, how the local scale adapts to dense and sparse regions, and how two one-way strengths merge into a single shared weight. Adjust the neighbor count, switch regions, and select a neighbor to inspect its numbers.
Key idea: point i connects to its nearest neighbor with full strength at the radius rho, and the strength fades with distance at a rate set by a local scale sigma. The scale is chosen so the neighbor strengths sum to log2 of the neighbor count, which keeps dense and sparse regions comparable. Two one-way strengths p(j|i) and p(i|j) are combined into one shared weight with the fuzzy union p(i,j) = p(j|i) + p(i|j) minus p(j|i) times p(i|j). The interactive controls load when JavaScript is available.
Two ideas carry the most weight here. A local scale adjusts each neighborhood so dense and sparse regions contribute comparable strengths, and the two directed strengths on a link merge into one symmetric weight. The result is a single weighted graph that records which observations sit close together in the original space.
The role of preprocessing
The quality of a UMAP embedding depends heavily on preprocessing. Since the algorithm relies on nearest neighbor relationships, the input feature space must be carefully prepared.
A typical workflow includes:
• Imputation of missing values
Instead of simple strategies such as median imputation, a more expressive approach is to use the IterativeImputer from scikit-learn. This method models each feature with missing values as a function of other features and iteratively refines those estimates. It helps most in structured datasets where variables are correlated, because it preserves multivariate relationships rather than collapsing missing values to a single summary statistic.
• Feature scaling
Standardizing features to zero mean and unit variance stops a variable with a large range from dominating the distance calculations simply because of its units. It levels the spreads without deciding how much any one feature ends up mattering. Without scaling, variables with larger numeric ranges can distort the neighborhood graph. Poor preprocessing can lead to misleading embeddings that show scale artifacts dressed up as structure.
Key parameters and their effects
Two parameters play a central role in shaping the UMAP embedding:
• n_neighbors
This controls the size of the local neighborhood used to build the graph.
Smaller values emphasize fine-grained local structure and can separate tight clusters. Larger values incorporate more global context and produce smoother layouts.
• min_dist
This determines how tightly points can cluster in the embedding.
Lower values create compact clusters, which can make dense groupings easier to see. Higher values spread points out. The setting controls how tightly the layout may pack points. It does not improve how faithfully the map preserves the original structure.
A commonly effective setting for min_dist is around 0.1, which balances separation and interpretability. Another important detail is the use of a fixed random_state. UMAP includes stochastic elements, and setting a seed ensures reproducibility, which is important for both research and production workflows.
Figure 3 makes these parameters tangible. It runs a small attraction and repulsion layout, in which n_neighbors and min_dist can be varied and the result replayed to watch the arrangement respond.
Figure 3 · Interactive
Optimization and the effects of UMAP parameters
Strong links pull related points together while a repulsion term keeps unrelated points apart, and the layout settles into groups. Adjust n_neighbors and min_dist, replay the layout, and try a new seed to see how the picture responds.
Key idea: attraction along strong links and repulsion between unrelated points settle the layout into groups. Smaller n_neighbors emphasizes local groups, larger n_neighbors emphasizes broad structure, smaller min_dist packs points tightly, and larger min_dist spreads them out. These parameters shape the view. Questions about cluster structure, global distances, and causation are settled on the data itself. The interactive controls and animation load when JavaScript is available.
The controls shape the view without fixing any ground truth. Group sizes, the gaps between groups, and the axis orientation all follow from the optimization, so an embedding is best treated as an exploratory map that is confirmed with quantitative analysis. The worked example below compares several views and reports how often near neighbors share a continent, a figure measured on the embedding rather than on the original indicators.
Example: UMAP dimensionality reduction on country data
UMAPEmbeddings
How UMAP turns nine socioeconomic and health indicators for 167 countries into a three-dimensional map that keeps similar nations close while broader structure emerges.
UMAP Embeddings. How UMAP turns nine socioeconomic and health indicators for 167 countries into a three-dimensional map that keeps similar nations close while broader structure emerges. Key topics covered: Country indicator data, Optimize the layout, 2D embedding views, Impute and scale, Neighbor graph, Key parameters, Interactive 3D view, Reading the embedding.
Data
The dataset used in this notebook contains country-level development indicators, with each row representing one country and each feature capturing a different aspect of national conditions such as mortality, trade, health expenditure, income, inflation, life expectancy, fertility, and gross domestic product (GDP) per capita.
This notebook treats the file as an analysis-ready tabular dataset for dimensionality reduction and visualization. Further information about the dataset can be found by accessing the link: https://www.kaggle.com/datasets/vipulgohel/clustering-pca-assignment/data
Notebook overview
This notebook follows a simple analytical sequence:
- Import the required packages
- Define reusable constants and configuration values
- Load and validate the country dataset
- Build a clean feature matrix for modeling
- Compute a three-dimensional UMAP embedding
- Visualize the embedding from multiple views
The emphasis is on a transparent workflow where each computational step is documented close to the code that performs it.
Import dependencies
This section collects all required libraries in one place so the notebook can be reviewed and rerun more easily.
The imported packages serve distinct purposes:
pandasfor tabular data handlingIterativeImputerandStandardScalerfor preprocessingumap-learnfor nonlinear dimensionality reductionplotlyfor interactive visualization
import pandas as pd
from sklearn.experimental import enable_iterative_imputer # noqa: F401
from sklearn.impute import IterativeImputer
from sklearn.preprocessing import StandardScaler
import umap
import plotly.express as px
import plotly.io as pioGrouping imports at the top is especially helpful in notebooks intended for publication or reuse because readers can immediately see the required software stack and package choices without scanning the entire file.
Define constants
Rather than scattering hard-coded values throughout the notebook, the main analysis settings are defined once here.
This improves readability and makes the notebook easier to modify. The number of neighbors, the random seed, the plotted dimensions, and the feature list can be changed from a single location without searching through later cells.
# Set Plotly theme and figure dimensions
pio.templates.default = "plotly_white"
FIG_WIDTH = 800
FIG_HEIGHT = 700
# Define feature columns and other constants
FEATURE_COLS = [
"child_mort", "exports", "health", "imports", "income", "inflation", "life_expec", "total_fer", "gdpp",
]
REQUIRED_COLS = {"country", "continent", *FEATURE_COLS}
IMPUTER_MAX_ITER = 25
DEFAULT_N_NEIGHBORS = 20
# For reproducibility
RANDOM_SEED = 42The feature list is defined explicitly instead of selecting columns by position. That reduces the risk of accidentally including identifiers or metadata in the embedding if the source CSV changes shape in the future.
Load and validate data
Before any modeling begins, the notebook checks that the input file has the expected structure.
These checks matter because downstream dimensionality-reduction results are only as trustworthy as the data entering the pipeline. In practical terms, this section verifies that:
- all required columns are present
- country identifiers are populated
- each country appears only once
- modeled features can be interpreted numerically
Catching structural issues early prevents silent failures later in the analysis.
# Path to data file
csv_path = "Country-data.csv"
# Load CSV into DataFrame
df = pd.read_csv(csv_path)
# Check for required columns
missing = [c for c in REQUIRED_COLS if c not in df.columns]
if missing:
raise ValueError(
f"CSV is missing required columns: {missing}\n"
f"Found columns: {list(df.columns)}"
)
# Work with a copy to avoid side effects
df = df.copy()
# Ensure categorical identifiers are present and unique
if df[["country", "continent"]].isna().any().any():
raise ValueError("Columns 'country' and 'continent' must not contain missing values.")
duplicate_countries = df["country"].duplicated(keep=False)
if duplicate_countries.any():
duplicate_labels = sorted(df.loc[duplicate_countries, "country"].unique())
raise ValueError(f"Duplicate country entries detected: {duplicate_labels}")
# Coerce feature columns to numeric so any parsing issues are surfaced as missing values
df[FEATURE_COLS] = df[FEATURE_COLS].apply(pd.to_numeric, errors="coerce")
missing_by_feature = df[FEATURE_COLS].isna().sum().sort_values(ascending=False)
missing_total = int(missing_by_feature.sum())
# Print dataset diagnostics
print(f"Rows: {len(df):,}")
print(f"Countries: {df['country'].nunique():,}")
if missing_total:
print("Missing values by feature:")
print(missing_by_feature[missing_by_feature > 0].to_string())
else:
print("Missing values by feature: none")Rows: 167
Countries: 167
Missing values by feature: noneResult note. The validation output indicates a clean one-row-per-country dataset with 167 observations and no missing values in the modeled features. In this file the later imputation logic is effectively dormant, which is a useful sanity check before fitting the embedding.
Prepare feature matrix and preprocess
This section creates the numeric matrix that will be passed into UMAP.
Two preprocessing steps are kept explicit:
- missing-value handling
- feature scaling
Even when the current file is complete, it is useful to keep the preprocessing pipeline in place so the notebook remains robust if the dataset is updated later with partially observed values.
# Select features for UMAP
X_raw = df[FEATURE_COLS].copy()
feature_bounds = pd.DataFrame({"min": X_raw.min(), "max": X_raw.max()})
# Iterative, model-based imputation preserves multivariate structure better than a univariate fill
imputer = IterativeImputer(
random_state=RANDOM_SEED,
initial_strategy="median",
max_iter=IMPUTER_MAX_ITER,
skip_complete=True,
min_value=feature_bounds["min"].to_numpy(),
max_value=feature_bounds["max"].to_numpy(),
)
X_imputed = pd.DataFrame(
imputer.fit_transform(X_raw),
columns=FEATURE_COLS,
index=df.index,
)
# Standardize features before distance-based manifold learning
scaler = StandardScaler()
X = scaler.fit_transform(X_imputed)
if X_raw.isna().any().any():
print(f"Iterative imputer converged in {imputer.n_iter_} round(s).")
else:
print("No missing values detected; the iterative imputer remains in place for reproducible future runs.")No missing values detected; the iterative imputer remains in place for reproducible future runs.The imputer and scaler are applied in separate, visible steps rather than being hidden inside a pipeline. For an educational notebook, that makes it easier to inspect intermediate objects and explain what happened to the data before the embedding was learned. Because the current dataset has no missing feature values, the iterative imputer leaves the numeric matrix unchanged and preprocessing reduces to scaling. That means the notebook is ready for incomplete data without altering the present analysis unnecessarily.
Compute UMAP embedding and prepare plot dataframe
Once the features are prepared, the notebook fits a three-dimensional UMAP model and stores the resulting coordinates in a plotting dataframe.
A few interpretation points are worth keeping in mind:
- the UMAP axes are learned coordinates with no scientific meaning of their own
- which observations share a neighborhood is the informative part, while the distances themselves, the gaps between groups, the group sizes and the axis orientation all come out of the optimization
- the embedding is built for exploration and will not support formal inference
Attaching the coordinates back to country and continent labels makes the visualization stage much easier to interpret.
# Configure UMAP
n_neighbors = min(DEFAULT_N_NEIGHBORS, len(df) - 1)
if n_neighbors < 2:
raise ValueError("UMAP requires at least two observations.")
umap_3d = umap.UMAP(
n_neighbors=n_neighbors,
min_dist=0.1,
n_components=3,
metric="euclidean",
random_state=RANDOM_SEED,
transform_seed=RANDOM_SEED,
)
# Fit and transform features
embedding_3d = umap_3d.fit_transform(X)
# Prepare DataFrame for plotting
plot_df = df[["country", "continent"]].copy()
# Add UMAP1
plot_df["UMAP1"] = embedding_3d[:, 0]
# Add UMAP2
plot_df["UMAP2"] = embedding_3d[:, 1]
# Add UMAP3
plot_df["UMAP3"] = embedding_3d[:, 2]
Choosing three components provides some flexibility during interpretation. It allows the notebook to show pairwise projections as well as a fully interactive 3D view without refitting a second model. The embedding has the expected shape of 167 countries by 3 coordinates, so every observation is carried through into the visualization.
Plot UMAP interactive results
The block generates three two-dimensional views plus a three-dimensional interactive plot.
Using multiple views is helpful because any single projection can hide separation that becomes visible from another angle. The saved HTML outputs also make it easy to share the figures outside the notebook while preserving hover labels and rotation controls.
# Tooltip for UMAP1 vs UMAP2
hovertemplate_12 = "<b>%{customdata[2]}</b><br>UMAP1: %{x:.3f}<br>UMAP2: %{y:.3f}<extra></extra>"
# Tooltip for UMAP2 vs UMAP3
hovertemplate_23 = "<b>%{customdata[2]}</b><br>UMAP2: %{x:.3f}<br>UMAP3: %{y:.3f}<extra></extra>"
# Tooltip for UMAP1 vs UMAP3
hovertemplate_13 = "<b>%{customdata[2]}</b><br>UMAP1: %{x:.3f}<br>UMAP3: %{y:.3f}<extra></extra>"
# Tooltip for 3D plot
hovertemplate_3d = "<b>%{customdata[3]}</b><br>UMAP1: %{x:.3f}<br>UMAP2: %{y:.3f}<br>UMAP3: %{z:.3f}<extra></extra>"
# Plot UMAP 1 vs UMAP 2
fig_12 = px.scatter(
plot_df, x="UMAP1", y="UMAP2", color="continent", custom_data=["UMAP1", "UMAP2", "country"], title="UMAP 1 vs UMAP 2"
)
fig_12.update_traces(marker=dict(size=9, line=dict(width=0.5, color="rgba(0,0,0,0.4)")), hovertemplate=hovertemplate_12)
fig_12.update_layout(legend_title_text="Continent", xaxis_title="UMAP 1", yaxis_title="UMAP 2", dragmode="pan", width=FIG_WIDTH, height=FIG_HEIGHT)
fig_12.update_yaxes(scaleanchor="x", scaleratio=1)
fig_12.show()
# Save plot as HTML
fig_12.write_html("umap_1_vs_2_nt.html")
# Plot UMAP 2 vs UMAP 3
fig_23 = px.scatter(
plot_df, x="UMAP2", y="UMAP3", color="continent", custom_data=["UMAP2", "UMAP3", "country"], title="UMAP 2 vs UMAP 3"
)
fig_23.update_traces(marker=dict(size=9, line=dict(width=0.5, color="rgba(0,0,0,0.4)")), hovertemplate=hovertemplate_23)
fig_23.update_layout(legend_title_text="Continent", xaxis_title="UMAP 2", yaxis_title="UMAP 3", dragmode="pan", width=FIG_WIDTH, height=FIG_HEIGHT)
fig_23.update_yaxes(scaleanchor="x", scaleratio=1)
fig_23.show()
# Save plot as HTML
fig_23.write_html("umap_2_vs_3_nt.html")
# Plot UMAP 1 vs UMAP 3
fig_13 = px.scatter(
plot_df, x="UMAP1", y="UMAP3", color="continent", custom_data=["UMAP1", "UMAP3", "country"], title="UMAP 1 vs UMAP 3"
)
fig_13.update_traces(marker=dict(size=9, line=dict(width=0.5, color="rgba(0,0,0,0.4)")), hovertemplate=hovertemplate_13)
fig_13.update_layout(legend_title_text="Continent", xaxis_title="UMAP 1", yaxis_title="UMAP 3", dragmode="pan", width=FIG_WIDTH, height=FIG_HEIGHT)
fig_13.update_yaxes(scaleanchor="x", scaleratio=1)
fig_13.show()
# Save plot as HTML
fig_13.write_html("umap_1_vs_3_nt.html")
# Plot UMAP 3D embedding
fig3d = px.scatter_3d(
plot_df, x="UMAP1", y="UMAP2", z="UMAP3", color="continent", custom_data=["UMAP1", "UMAP2", "UMAP3", "country"], title="UMAP of Countries (3D)"
)
fig3d.update_traces(marker=dict(size=4), hovertemplate=hovertemplate_3d)
fig3d.update_layout(legend_title_text="Continent", width=FIG_WIDTH, height=FIG_HEIGHT)
fig3d.show()
# Save plot as HTML
fig3d.write_html("umap_3d_nt.html")UMAP 1 against UMAP 2 by continent
UMAP 2 against UMAP 3 by continent
UMAP 1 against UMAP 3 by continent
Rotatable three-dimensional UMAP embedding
These figures are a map of multivariable similarity, and they are neither a map of geography nor the output of a supervised classifier.
A few practical interpretation habits help:
- look for broad neighborhoods before focusing on isolated points
- compare patterns across the different 2D views
- treat continent coloring as a label laid over the result, since the model never saw it
- use hover labels to inspect countries that appear unusually separated or unexpectedly close
The plots suggest continent explains part of the structure, but not all of it. Europe tends to occupy a relatively distinct higher UMAP2 and UMAP3 region, Africa is often lower on UMAP2, and Asia, North America, Oceania, and South America overlap more substantially. A simple diagnostic finds that about 54% of each country’s five nearest neighbors share its continent, well above the roughly 22% expected if neighbors were drawn at random from the same table. Because that count is measured on the embedding and not on the nine indicators, it restates the pattern the layout already shows without checking it against the original features. Read the map as a broad socioeconomic similarity structure. It captures continent only loosely.
Key Takeaways
The graph is built from each point’s nearest neighbors, and the layout pulls linked points together, so observations close in the original feature space are drawn close in the map. That is the relationship the method supports.
Group sizes, the gaps between groups, and axis direction all follow from the optimization, so a wide gap does not measure how different two groups are. Treat the embedding as an exploratory map and check any distance claim against the original features.
UMAP finds neighbors with a distance metric, Euclidean in this example, so an unscaled feature with a large range would dominate the graph. The notebook standardizes all nine indicators to zero mean and unit variance before fitting.
Smaller values emphasize fine local structure and can separate tight groups, while larger values fold in more context and smooth the layout. The example asks for 20 neighbors, falling back to one less than the row count on a smaller table.
min_dist sets how tightly points may sit in the embedding, and it is 0.1 here. Lower values make groups look compact and higher values spread them out, but neither setting changes which observations are actually similar.
UMAP has stochastic elements, so two runs on the same matrix can produce different pictures. The example sets random_state and transform_seed to 42. Only random_state fixes this layout, while transform_seed applies when transform is called on new data.
Data & License
Country-level socio-economic and health indicators for 167 countries (the “HELP International” data), published on Kaggle by Vipul Gohel. Source: kaggle.com/datasets/vipulgohel/clustering-pca-assignment.License: CC0 1.0 Universal (Public Domain Dedication) — no rights reserved; attribution is not legally required and is given here as a courtesy. CC0 1.0
© 2025 Philip Sarajlic. All rights reserved for the article’s original text and figures; the third-party data remains under the license shown above.
Code examples in this article are licensed under the Common Public Attribution License Version 1.0 (CPAL-1.0), an OSI-approved copyleft license based on the Mozilla Public License 1.1. Initial Developer: Philip Sarajlic.
Attribution required by CPAL Exhibit B: © 2025 Philip Sarajlic · “Based on code by Philip Sarajlic” · philipsarajlic.com · no graphic image. This attribution must be displayed in Larger Works.
Modifications must be released in Source Code form under CPAL-1.0. Making the code usable by anyone other than you over a network is External Deployment under the license and is treated as distribution, so the Source Code must be made available to those users.
Full text: opensource.org/license/cpal-1-0 (SPDX identifier CPAL-1.0)
