Optimal Transport Tools (OTT) documentation
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 various implementation of the Sinkhorn algorithm [Cuturi, 2013, Peyré and Cuturi, 2019, Scetbon et al., 2021]. These implementation take advantage of several JAX features, such as Just-in-time (JIT) compilation, auto-vectorization (VMAP) and both automatic and/or implicit differentiation. A few tutorials are provided below, along with different use-cases, notably for single-cell genomics data [Schiebinger et al., 2019].
There are currently three packages,
tools, playing the following roles:
ott.geometry package defines 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 [Solomon et al., 2015]). The design choice in OTT is to state that cost functions and algorithms should operate independently: if a particular cost function allows for faster computations (e.g. squared-Euclidean distance when comparing grids), this should not be taken advantage of at the level of optimizers, but at the level of the problems description. Geometry objects are therefore only considered as arguments to describe OT problem handled in
core, using subroutines provided by geometries;
ott.core package help define first an OT problem (linear, quadratic, barycenters). These problems are then solved using Sinkhorn algorithm and its variants, the main workhorse to solve OT in this package, as well as variants that can comppute Gromov-Wasserstein distances or barycenters of several measures;
ott.tools package provides an interface to exploit OT solutions, as produced by
corefunctions. Such tasks include instantiating OT matrices, computing approximations to Wasserstein distances [Genevay et al., 2018, Séjourné et al., 2019], or computing differentiable sort and quantile operations [Cuturi et al., 2019].
- Sinkhorn Barycenters
- Gromov Wasserstein
- Low-Rank Gromov-Wasserstein
- Sinkhorn Divergence Hessians
- Soft Sorting
- Single-cell genomics
- Gromov-Wasserstein for multi-omics
- Fairness regularizers
- ICNN Dual Solver
- Initialization Schemes for Input Convex Neural Network
- Wasserstein Barycenters of Gaussian Mixture Models