|Downloads| |Tests| |Docs| |Coverage|
Optimal Transport Tools (OTT) documentation
===========================================
`Code `_ on GitHub.
To install, simply run ``pip install ott-jax``.
Intro
-----
`OTT` is a `JAX `_ package that bundles a few utilities to compute, and
differentiate as needed, the solution to optimal transport (OT) problems, taken in a fairly wide sense.
For instance, `OTT` can of course compute Wasserstein (or Gromov-Wasserstein) distances between
weighted clouds of points (or histograms) in a wide variety of scenarios,
but also estimate Monge maps, Wasserstein barycenters, and help with simpler tasks
such as differentiable approximations to ranking or even clustering.
To achieve this, `OTT` rests on two families of tools:
The first family consists in *discrete* solvers computing transport between point clouds,
using the Sinkhorn :cite:`cuturi:13` and low-rank Sinkhorn :cite:`scetbon:21` algorithms,
and moving up towards Gromov-Wasserstein :cite:`memoli:11,peyre:16`;
the second family consists in *continuous* solvers, using suitable neural architectures :cite:`amos:17` coupled
with SGD type estimators :cite:`makkuva:20,korotin:21`.
Design Choices
--------------
`OTT` is designed with the following choices:
- Take advantage whenever possible of JAX features, such as `Just-in-time (JIT) compilation`_,
`auto-vectorization (VMAP)`_ and both `automatic`_ but most importantly `implicit`_ differentiation.
- Split geometry from OT solvers in the discrete case: We argue that there
should be one, and one implementation only, of every major OT algorithm
(Sinkhorn, Gromov-Wasserstein, barycenters, etc...), regardless of the
geometric setup that is considered. To give a concrete example, any
speedups one may benefit from by using a specific cost
(e.g. Sinkhorn being faster when run on a separable cost on histograms supported
on a separable grid :cite:`solomon:15`) should not require a separate
reimplementation of a Sinkhorn routine.
- As a consequence, and to minimize code copy/pasting, use as often as possible
object hierarchies, and interleave outer solvers (such as quadratic,
aka Gromov-Wasserstein solvers) with inner solvers (e.g. Low-Rank Sinkhorn).
This choice ensures that speedups achieved at lower computation levels
(e.g. low-rank factorization of squared Euclidean distances) propagate seamlessly and
automatically in higher level calls (e.g. updates in Gromov-Wasserstein),
without requiring any attention from the user.
.. TODO(marcocuturi): add missing package descriptions below
Packages
--------
- :ref:`geometry` contains classes to instantiate objects that describe
*two point clouds* paired with a *cost* function. Geometry objects are used to
describe OT problems, handled by solvers in the :ref:`solvers`.
- :ref:`problems`
- :ref:`solvers`
- :ref:`initializers`
- :ref:`tools` provides an interface to exploit OT solutions, as produced by
solvers in the :ref:`solvers`. Such tasks include computing approximations
to Wasserstein distances :cite:`genevay:18,sejourne:19`, approximating OT
between GMMs, or computing differentiable sort and quantile operations
:cite:`cuturi:19`.
- :ref:`math`
.. toctree::
:maxdepth: 1
:caption: Tutorials:
notebooks/point_clouds.ipynb
notebooks/introduction_grid.ipynb
.. toctree::
:maxdepth: 1
:caption: Benchmarks:
notebooks/OTT_&_POT.ipynb
notebooks/One_Sinkhorn.ipynb
notebooks/LRSinkhorn.ipynb
.. toctree::
:maxdepth: 1
:caption: Advanced Applications:
notebooks/Sinkhorn_Barycenters.ipynb
notebooks/gromov_wasserstein.ipynb
notebooks/GWLRSinkhorn.ipynb
notebooks/Hessians.ipynb
notebooks/soft_sort.ipynb
notebooks/application_biology.ipynb
notebooks/gromov_wasserstein_multiomics.ipynb
notebooks/fairness.ipynb
notebooks/neural_dual.ipynb
notebooks/icnn_inits.ipynb
notebooks/wasserstein_barycenters_gmms.ipynb
notebooks/gmm_pair_demo.ipynb
notebooks/MetaOT.ipynb
.. toctree::
:maxdepth: 1
:caption: Public API: ott packages
geometry
problems/index
solvers/index
initializers/index
tools
math
.. toctree::
:maxdepth: 1
:caption: References:
references
.. |Downloads| image:: https://pepy.tech/badge/ott-jax
:target: https://pypi.org/project/ott-jax/
:alt: Documentation
.. |Tests| image:: https://img.shields.io/github/workflow/status/ott-jax/ott/tests/main
:target: https://github.com/ott-jax/ott/actions/workflows/tests.yml
:alt: Documentation
.. |Docs| image:: https://img.shields.io/readthedocs/ott-jax/latest
:target: https://ott-jax.readthedocs.io/en/latest/
:alt: Documentation
.. |Coverage| image:: https://img.shields.io/codecov/c/github/ott-jax/ott/main
:target: https://app.codecov.io/gh/ott-jax/ott
:alt: Coverage
.. _Just-in-time (JIT) compilation: https://jax.readthedocs.io/en/latest/jax.html#just-in-time-compilation-jit
.. _auto-vectorization (VMAP): https://jax.readthedocs.io/en/latest/jax.html#vectorization-vmap
.. _automatic: https://jax.readthedocs.io/en/latest/jax.html#automatic-differentiation
.. _implicit: https://jax.readthedocs.io/en/latest/_autosummary/jax.custom_jvp.html#jax.custom_jvp