Optimal Transport Tools (OTT) documentation
Code hosted on Github. To install,
clone that repo or simply run
pip install ott-jax.
OTT is a JAX package that bundles a few utilities to compute and differentiate the solution to optimal transport problems. OTT can help you compute Wasserstein distances between weighted clouds of points (or histograms), using a cost (e.g. a distance) between individual points.
To that end OTT uses a sturdy and versatile implementation of the Sinkhorn algorithm 1 2. This implementation takes advantage of several JAX features, such as Just-in-time (JIT) compilation, auto-vectorization (VMAP), and both automatic and/or implicit differentiation. Basic tutorial snippets are provided below, along with a more advanced example to process single-cell genomics data 3.
There are currently three packages,
tools, playing the following roles:
geometrydefines classes that describe two point clouds paired with a cost function (simpler geometries are also implemented, such as that defined by points supported on a multi-dimensional grids with a separable cost 4). A geometry, along with weight vectors
b, describes an OT problem.
corealgorithms can solve such problems, using subroutines provided by geometries;
corecontains the Sinkhorn algorithm, the main workhorse to solve OT in this package, as well as variants that can be used to compute barycenters of several measures;
toolsprovides an interface to exploit OT solutions, as produced by
corefunctions. Such tasks include instantiating OT matrices, computing approximations to Wasserstein distances 5 6, or computing differentiable sort and quantile operations 7.
Indices and tables
Cuturi, Sinkhorn Distances: Lightspeed Computation of Optimal Transport, NIPS’13.
Peyré, M. Cuturi, Computational Optimal Transport, FNT in ML, 2019.
Schiebinger et al., Optimal-Transport Analysis of Single-Cell Gene Expression Identifies Developmental Trajectories in Reprogramming, Cell 176, 928–943.
Solomon et al, Convolutional Wasserstein distances: efficient optimal transportation on geometric domains, ACM ToG, SIGGRAPH’15.
Genevay et al., Learning Generative Models with Sinkhorn Divergences, AISTATS’18.
Séjourné et al., Sinkhorn Divergences for Unbalanced Optimal Transport, arXiv:1910.12958.
Cuturi et al. Differentiable Ranking and Sorting using Optimal Transport, NeurIPS’19.