YES! This is actually a perfect use case for Wasserstein, and it solves a major problem in single-cell analysis!

The problem you're highlighting:

UMAP and t-SNE are non-linear projections that distort distances:

• t-SNE preserves local structure but destroys global distances
• UMAP is better but still compresses/expands regions non-uniformly
• Same cells can look very different in different embeddings!
• Comparing cell types directly in these spaces is unreliable

The Wasserstein solution:

Don't compute Wasserstein in the visualization space! Instead:

Option 1: Compute in the original high-dimensional space

# Cell type A and B as distributions in gene expression space
# Each cell is a point in ~20,000 dimensional gene space

cell_type_A = adata[adata.obs['cell_type'] == 'A'].X  # n_cells × n_genes
cell_type_B = adata[adata.obs['cell_type'] == 'B'].X

# Wasserstein distance in original space
W_dist = ot.wasserstein_distance(cell_type_A, cell_type_B)

Pros: True biological distance, invariant to visualization Cons: Curse of dimensionality, noisy genes matter

Option 2: Use PCA space (better!)

# PCA preserves global structure linearly
pca_embeddings = adata.obsm['X_pca']  # e.g., 50 dimensions

A_pca = pca_embeddings[adata.obs['cell_type'] == 'A']
B_pca = pca_embeddings[adata.obs['cell_type'] == 'B']

W_dist = ot.wasserstein_distance(A_pca, B_pca)

Pros: Denoised, consistent across visualizations, computationally tractable Cons: Linear assumption

Option 3: Graph-based geodesic distances (most robust!)

This is really clever:

# Build a k-NN graph in gene expression space
# Compute geodesic (shortest path) distances on the graph
# Use these as the ground metric for Wasserstein

import scanpy as sc

# Build neighbor graph
sc.pp.neighbors(adata, use_rep='X_pca')

# For Wasserstein, use graph distances as the metric
# This respects the manifold structure!

Why this is brilliant: