Package: ott-jax
Optimal Transport Tools (OTT)See thefull documentation.What is OTT-JAX?AJAXpowered library to compute optimal transport at scale and on accelerators,OTT-JAXincludes the fastest
implementation of the Sinkhorn algorithm you will find around. We have implemented all tweaks (scheduling, momentum, acceleration, initializations) and extensions (low-rank, entropic maps). They can be used directly between two datasets, or within more advanced problems
(Gromov-Wasserstein, barycenters). Some ofJAXfeatures, includingJIT,auto-vectorizationandimplicit differentiationwork towards the goal of having end-to-end differentiable outputs.OTT-JAXis led by a team of researchers at Apple, with contributions from Google and Meta researchers, as well as many academic partners, including TU München, Oxford, ENSAE/IP Paris, ENS Paris and the Hebrew University.InstallationInstallOTT-JAXfromPyPIas:pipinstallott-jaxor withcondaviaconda-forgeas:condainstall-cconda-forgeott-jaxWhat is optimal transport?Optimal transport can be loosely described as the branch of mathematics and optimization that studiesmatching problems: given two families of points, and a cost function on pairs of points, find a "good" (low cost) way
to associate bijectively to every point in the first family another in the second.Such problems appear in all areas of science, are easy to describe, yet hard to solve. Indeed, while matching optimally
two sets of $n$ points using a pairwise cost can be solved with theHungarian algorithm, solving it costs an order of $O(n^3)$
operations, and lacks flexibility, since one may want to couple families of different sizes.Optimal transport extends all of this, through faster algorithms (in $n^2$ or even linear in $n$) along with numerous
generalizations that can help it handle weighted sets of different size, partial matchings, and even more evolved
so-called quadratic matching problems.In the simple toy example below, we compute the optimal coupling matrix between two point clouds sampled randomly
(2D vectors, compared with the squared Euclidean distance):Exampleimportjaximportjax.numpyasjnpfromott.geometryimportpointcloudfromott.problems.linearimportlinear_problemfromott.solvers.linearimportsinkhorn# sample two point clouds and their weights.rngs=jax.random.split(jax.random.PRNGKey(0),4)n,m,d=12,14,2x=jax.random.normal(rngs[0],(n,d))+1y=jax.random.uniform(rngs[1],(m,d))a=jax.random.uniform(rngs[2],(n,))b=jax.random.uniform(rngs[3],(m,))a,b=a/jnp.sum(a),b/jnp.sum(b)# Computes the couplings using the Sinkhorn algorithm.geom=pointcloud.PointCloud(x,y)prob=linear_problem.LinearProblem(geom,a,b)solver=sinkhorn.Sinkhorn()out=solver(prob)The call tosolver(prob)above works out the optimal transport solution. Theoutobject contains a transport matrix
(here of size $12\times 14$) that quantifies the association strength between each point of the first point cloud, to one or
more points from the second, as illustrated in the plot below. We provide more flexibility to define custom cost
functions, objectives, and solvers, as detailed in thefull documentation.CitationIf you have found this work useful, please consider citing this reference:@article{cuturi2022optimal,
  title={Optimal Transport Tools (OTT): A JAX Toolbox for all things Wasserstein},
  author={Cuturi, Marco and Meng-Papaxanthos, Laetitia and Tian, Yingtao and Bunne, Charlotte and
          Davis, Geoff and Teboul, Olivier},
  journal={arXiv preprint arXiv:2201.12324},
  year={2022}
}See alsoThemoscotpackage for OT analysis of multi-omics data also uses OTT as a backbone.