∫ Optimal Transport & Wasserstein Distance in Single-Cell Biology

🎯 Core Concepts

What is Optimal Transport?

Optimal transport theory provides a mathematical framework for comparing probability distributions by computing the minimal "cost" of transforming one distribution into another. Originally formulated by Gaspard Monge in 1781 and generalized by Leonid Kantorovich in 1942, the theory asks: What is the most efficient way to move a pile of dirt from one location to match a desired shape elsewhere?

The Wasserstein Distance

The Wasserstein distance (also called Earth Mover's Distance) quantifies the dissimilarity between two probability distributions as the minimal cost to transport mass from one to the other. For single-cell biology:

• Cell-Cell Similarity: Treating gene expression profiles as probability distributions, Wasserstein distance captures how much "work" is needed to transform one cell's expression into another's
• Population Alignment: Compares entire cell populations by finding optimal couplings that preserve local neighborhood structure
• Trajectory Inference: Models cellular differentiation as probability distributions flowing through gene expression space over time

Why Optimal Transport for Biology?

• Natural Framework: Biological processes (cell differentiation, population dynamics) naturally involve mass transport with growth and death
• Handles Unpaired Data: Unlike supervised methods, OT aligns datasets without requiring matched cells or shared features
• Preserves Structure: Gromov-Wasserstein formulation compares metric spaces directly, preserving manifold geometry across modalities
• Probabilistic Interpretation: Coupling matrices provide interpretable cell-cell correspondence probabilities for downstream analysis

Key Challenges Addressed

• Computational Tractability: Classical OT is computationally expensive (O(n³)); entropic regularization (Sinkhorn algorithm) reduces complexity to O(n² log n)
• Unbalanced Mass: Biological systems have growth/death - unbalanced OT with KL penalties accommodates varying cell numbers
• Cross-Modality Integration: Different omics have distinct feature spaces - Gromov-Wasserstein compares via pairwise distances instead of direct feature matching
• Scalability to Atlases: Linear memory implementations (moscot) enable analysis of 1.7M+ cell datasets

📈 Major Application Domains

1. Trajectory Inference & Temporal Dynamics

2. Multi-Omics Integration & Cross-Modality Alignment

3. Cell Similarity & Clustering

4. Neural Optimal Transport & Deep Learning

📚 Method Comparison

💡 Practical Implementation Guide

Common Pitfalls & Best Practices

• Data Preprocessing: Normalize to probability distributions (sum to 1) for OT; log-transform for scRNA-seq, appropriate transform for other modalities
• Growth Rate Estimation: Use proliferation/apoptosis signatures or cell cycle scoring; iterative refinement improves accuracy
• Hyperparameter Tuning: Entropy regularization ε (0.001-0.1) trades off between exact OT and regularized smoothness; use grid search with interpolation validation
• Ground Cost Selection: Pearson correlation for RNA-seq, Cosine for ATAC-seq/methylation; geodesic distances on manifolds when available
• Computational Efficiency: Use local PCA (30-50D) instead of full gene space; GPU acceleration for Sinkhorn iterations; hierarchical coarse-graining for very large datasets
• Validation: Geodesic interpolation for trajectory methods; Label Transfer Accuracy and Alignment Score for integration; compare to simple baselines (mean, linear models)

Software Ecosystem

• POT (Python Optimal Transport): Core library with Sinkhorn, GW, unbalanced variants
• Waddington-OT: Trajectory inference from time-series
• moscot: Unified framework for temporal/spatial/spatiotemporal OT at scale
• SCOT: Multi-omics alignment with GW
• Pamona: Partial alignment for heterogeneous datasets
• CellOT: Neural OT for perturbation prediction

📖 Learning Resources

Theoretical Foundations

• Villani (2008) "Optimal Transport: Old and New" - Comprehensive mathematical treatment (Chapters 1-2, 6-7)
• Peyré & Cuturi (2019) "Computational Optimal Transport" - Modern algorithmic perspective with code
• Mémoli (2011) "Gromov-Wasserstein Distances" - Foundation for comparing metric spaces
• Chapel et al. (2020) "Partial Optimal Transport" - Theory of unbalanced mass with virtual points

Tutorials & Code Examples

• POT Library Tutorials: Sinkhorn, GW, unbalanced OT with executable notebooks
• Waddington-OT Documentation: Step-by-step trajectory inference pipeline
• moscot Tutorials: Temporal, spatial, and spatiotemporal mapping workflows
• Scanpy Tutorials: Single-cell preprocessing for OT methods (normalization, PCA, k-NN graphs)

Benchmark Datasets

• iPSC Reprogramming (GSE122662): ~65K cells, 10 timepoints over 16 days - Waddington-OT landmark dataset

🛠️ Hands-On Practice

The walkthrough below runs a complete optimal transport analysis on a single-cell time course — from a preprocessed AnnData object with a day annotation, through a PCA-space cost matrix, an entropic-regularized Sinkhorn solve, and finally to ancestor/descendant probabilities for a chosen cell population. It uses POT directly for the low-level mechanics (so the coupling is visible as a plain matrix) and moscot for the scalable, biology-aware version of the same computation. Everything is retrospective and purely computational: the inputs are existing count matrices, the outputs are transport plans and probabilities.

Environment & packages

POT provides the core solvers (Sinkhorn, unbalanced, Gromov-Wasserstein); moscot wraps them in an OTT-JAX backend with single-cell problem classes (TemporalProblem, SpatialAlignmentProblem, TranslationProblem). Scanpy handles the preprocessing that produces the latent space the cost is computed in. Install into a clean environment — moscot pins JAX/OTT versions and will fight an existing deep-learning stack.

# conda / mamba recommended
conda create -n otsc python=3.10 -y
conda activate otsc

pip install scanpy anndata
pip install pot                 # POT: ot.sinkhorn, ot.gromov, ot.unbalanced
pip install moscot              # pulls in ott-jax + jax (CPU build by default)

# optional GPU backend for large problems (CUDA 12)
# pip install --upgrade "jax[cuda12]"

# optional: the original Waddington-OT implementation
# pip install wot

Hardware. The dense coupling is n × m in float64, so two timepoints of 5,000 cells each is ~200 MB and fits comfortably on a laptop; 30,000 × 30,000 is ~7 GB and needs a compute node. moscot's low-rank solvers (rank=) and the OTT-JAX GPU backend are what make 10⁵-cell problems tractable — a GPU is optional for the toy case below but effectively required for a full atlas-scale time course.

Data structures & formats

• AnnData — the input: adata.X (normalized, log1p expression), adata.obs["day"] (numeric timepoint — must be numeric, not categorical, for TemporalProblem), and adata.obsm["X_pca"] (the space the cost is computed in).
• Marginals a and b — probability vectors over source and target cells, each summing to 1. Uniform by default; weighted by an estimated proliferation/apoptosis rate when modelling growth.
• Cost matrix M — an n × m array of pairwise distances (usually squared Euclidean in PCA or latent space). Normalize it (e.g. divide by M.max()) so that reg has a comparable meaning across datasets.
• Transport plan / coupling π — the n × m output. Row i sums to a_i, column j to b_j. Row-normalizing gives descendant probabilities; column-normalizing gives ancestor probabilities.
• moscot solution object — tp.solutions[(t0, t1)] holds the coupling for each consecutive pair; tp.pull() / tp.push() propagate a cell mask backward/forward through the chain of couplings and write the result into adata.obs.
• Gromov-Wasserstein inputs — for unpaired modalities you supply two intra-domain cost matrices (C₁: n × n, C₂: m × m) instead of a single cross-domain M, because no shared feature space exists.

Minimal code walkthrough

Load the time-course object, reduce to a shared PCA space, build and normalize the cost matrix between two adjacent days, solve the entropic problem with Sinkhorn, read ancestor/descendant probabilities off the coupling, then repeat the whole thing at scale with moscot's TemporalProblem including growth-aware unbalanced marginals.

import numpy as np
import scanpy as sc
import ot                      # POT

# 1. Load a preprocessed time course. adata.obs["day"] must be NUMERIC:
#    OT needs an ordering, and moscot uses the gap between timepoints to
#    scale the growth term.
adata = sc.read_h5ad("timecourse.h5ad")
adata.obs["day"] = adata.obs["day"].astype(float)

# 2. Build ONE shared latent space across all timepoints. Never compute a
#    cost on raw counts: Euclidean distance on counts is dominated by library
#    size and by a handful of high-expression genes.
sc.pp.normalize_total(adata, target_sum=1e4)
sc.pp.log1p(adata)
sc.pp.highly_variable_genes(adata, n_top_genes=2000)
sc.pp.pca(adata, n_comps=30, use_highly_variable=True)

# 3. Take two adjacent timepoints as source (t0) and target (t1)
t0, t1 = 2.0, 4.0
src = adata[adata.obs["day"] == t0]
tgt = adata[adata.obs["day"] == t1]
X0 = src.obsm["X_pca"]
X1 = tgt.obsm["X_pca"]
n, m = X0.shape[0], X1.shape[0]
print(f"day {t0}: {n} cells  ->  day {t1}: {m} cells")

# 4. Cost matrix: squared Euclidean in PCA space, then rescale to [0, 1].
#    Rescaling matters because the meaning of `reg` below is relative to the
#    magnitude of M -- an unnormalised M makes epsilon dataset-dependent.
M = ot.dist(X0, X1, metric="sqeuclidean")
M /= M.max()

# 5. Marginals. Uniform is the assumption "no cell is more likely than any
#    other to have a descendant" -- revisited in step 8.
a = np.ones(n) / n
b = np.ones(m) / m

# 6. Entropic-regularised OT (Sinkhorn). reg is epsilon: small = sharp,
#    near-deterministic plan but slow / numerically unstable; large = blurred
#    plan where every cell maps a little onto every other cell.
#    Use the log-domain solver -- it is the one that survives small reg.
pi = ot.sinkhorn(a, b, M, reg=0.05, method="sinkhorn_log", numItermax=2000)
print("coupling shape:", pi.shape, " mass:", pi.sum())   # ~1.0

# 7. Read biology off the coupling.
#    Row-normalise -> descendant distribution of each day-2 cell.
#    Column-normalise -> ancestor distribution of each day-4 cell.
descendants = pi / pi.sum(axis=1, keepdims=True)
ancestors   = pi / pi.sum(axis=0, keepdims=True)

# Ancestor profile of one terminal cluster at t1: how much of its mass
# came from each source cell type?
mask_t1 = (tgt.obs["cell_type"] == "Neuron").to_numpy()
anc_mass = ancestors[:, mask_t1].sum(axis=1)
anc_mass /= anc_mass.sum()
src_ct = src.obs["cell_type"].to_numpy()
for ct in np.unique(src_ct):
    print(f"  {ct:<20s} {anc_mass[src_ct == ct].sum():.3f}")

# 8. Growth and death break the "equal mass in, equal mass out" assumption:
#    proliferating cells contribute MORE than 1/n of the day-4 mass, apoptotic
#    cells less. Unbalanced OT relaxes the marginal constraints; reg_m controls
#    how strongly (small reg_m = marginals free to change, large = back to balanced).
g = src.obs["proliferation_score"].to_numpy()   # e.g. sc.tl.score_genes on a cycle signature
a_growth = g / g.sum()
pi_unb = ot.unbalanced.sinkhorn_unbalanced(
    a_growth, b, M, reg=0.05, reg_m=1.0, numItermax=2000
)
print("unbalanced transported mass:", pi_unb.sum())     # != 1.0 by design

# ---------------------------------------------------------------------------
# 9. The same analysis at scale, with moscot: solves ALL consecutive day pairs,
#    estimates growth rates from proliferation/apoptosis signatures, and keeps
#    couplings on GPU via the OTT-JAX backend.
# ---------------------------------------------------------------------------
from moscot.problems.time import TemporalProblem

tp = TemporalProblem(adata)
tp = tp.score_genes_for_marginals(
    gene_set_proliferation="human", gene_set_apoptosis="human"
)
tp = tp.prepare(time_key="day", joint_attr="X_pca")
tp = tp.solve(epsilon=1e-2, tau_a=0.95, tau_b=0.999)   # tau < 1 => unbalanced

# Pull a terminal population backwards through the whole chain of couplings
tp.pull(source=t0, target=t1, data="cell_type", subset="Neuron")
sc.pl.embedding(adata, basis="umap", color="pull")

# 10. Wasserstein distance as a summary statistic: how far apart are two
#     conditions in expression space? Use the un-regularised (exact) cost of
#     the Sinkhorn plan, or ot.emd2 for small problems.
w2 = np.sum(pi * M)
print(f"Entropic W2 (squared, normalised cost) between {t0} and {t1}: {w2:.4f}")

Common pitfalls & tips

• Epsilon controls sharpness, and there is no default that works everywhere. Entropic regularization blurs the plan: as reg grows, π converges toward the independent coupling ab^T, where every cell is a partial ancestor of every other cell and the ancestor probabilities become uninformative. As reg shrinks, the scaling vectors underflow and the solver returns NaNs. Normalize M to [0, 1], start around 0.01–0.1, and sanity-check by plotting the row entropy of π — if the effective number of descendants per cell is on the order of the whole target population, epsilon is too large.
• Always compute the cost in a denoised latent space, never on raw counts. Euclidean distance on raw or even log-normalized full-gene expression is dominated by library-size differences and by sparsity; 30–50 PCs (or an scVI latent space) gives distances that track biological similarity. Critically, the PCA must be fit on all timepoints jointly — a per-timepoint PCA produces two incomparable coordinate systems and the cost matrix becomes meaningless.
• Balanced OT silently assumes no proliferation and no death. If a progenitor population expands tenfold between timepoints, balanced OT must still move exactly 1/n of mass out of each source cell, so it fabricates ancestors for the expanded population from unrelated cells. Use unbalanced OT (ot.unbalanced.sinkhorn_unbalanced, or moscot's tau_a/tau_b < 1) with growth rates estimated from proliferation and apoptosis gene signatures — this is exactly what Waddington-OT's growth-rate iteration does.
• Memory is O(n·m) and it bites earlier than you expect. The dense coupling for 50,000 × 50,000 cells is ~20 GB in float64 before the solver's intermediate buffers. Subsample per timepoint, use moscot/OTT low-rank solvers (rank=200), or work on metacells. Do not build the cost matrix with ot.dist on the full atlas and hope.
• Timepoints must be close enough that the transport assumption holds. OT assumes cells move a short distance in expression space between snapshots; if the interval is long enough that populations fully reorganize, the minimal-cost matching is not the biologically correct one. Couple adjacent timepoints and compose the couplings, rather than solving day 0 → day 16 directly.
• Use Gromov-Wasserstein only when the two spaces are genuinely incomparable. For unpaired modalities (scRNA vs. scATAC with no shared feature space), GW matches intra-domain distance structure instead of cross-domain distances — this is the SCOT/Pamona setting. But GW is non-convex, solved by alternating minimization, and sensitive to initialization; run several seeds and check for consistency. If a shared feature space does exist (e.g. gene activity scores), fused GW or plain Wasserstein is more stable and far cheaper.
• The Sinkhorn cost is not the Wasserstein distance. The value returned by ot.sinkhorn2 includes the entropic penalty and is biased — it does not vanish when the two distributions are identical. For a distance you can compare across pairs, use the Sinkhorn divergence (ot.bregman.empirical_sinkhorn_divergence), which debiases by subtracting the self-transport terms, or ot.emd2 for exact W₂ on small problems.