POT: Python Optimal Transport

This open source Python library provides several solvers for optimization problems related to Optimal Transport for signal, image processing and machine learning.

Website and documentation: https://PythonOT.github.io/

Source Code (MIT): https://github.com/PythonOT/POT

POT has the following main features: * A large set of differentiable solvers for optimal transport problems, including: * Exact linear OT, entropic and quadratic regularized OT, * Gromov-Wasserstein (GW) distances, Fused GW distances and variants of quadratic OT, * Unbalanced and partial OT for different divergences, * OT barycenters (Wasserstein and GW) for fixed and free support, * Fast OT solvers in 1D, on the circle and between Gaussian Mixture Models (GMMs), * Many ML related solvers, such as domain adaptation, optimal transport mapping estimation, subspace learning, Graph Neural Networks (GNNs) layers. * Several backends for easy use with Pytorch, Jax, Tensorflow, Numpy and Cupy arrays.

Implemented Features

POT provides the following generic OT solvers:

• OT Network Simplex solver for the linear program/ Earth Movers Distance [1] .
• Conditional gradient [6] and Generalized conditional gradient for regularized OT [7].
• Entropic regularization OT solver with Sinkhorn Knopp Algorithm [2] , stabilized version [9] [10] [34], lazy CPU/GPU solver from geomloss [60] [61], greedy Sinkhorn [22] and Screening Sinkhorn [26].
• Bregman projections for Wasserstein barycenter [3], convolutional barycenter [21] and unmixing [4].
• Sinkhorn divergence [23] and entropic regularization OT from empirical data.
• Debiased Sinkhorn barycenters Sinkhorn divergence barycenter [37]
• Smooth optimal transport solvers (dual and semi-dual) for KL and squared L2 regularizations [17].
• Weak OT solver between empirical distributions [39]
• Non regularized Wasserstein barycenters [16] with LP solver (only small scale).
• Gromov-Wasserstein distances and GW barycenters (exact [13] and regularized [12,51]), differentiable using gradients from Graph Dictionary Learning [38]
• Fused-Gromov-Wasserstein distances solver and FGW barycenters (exact [24] and regularized [12,51]).
• Stochastic solver and differentiable losses for Large-scale Optimal Transport (semi-dual problem [18] and dual problem [19])
• Sampled solver of Gromov Wasserstein for large-scale problem with any loss functions [33]
• Non regularized free support Wasserstein barycenters [20].
• One dimensional Unbalanced OT with KL relaxation and barycenter [10, 25]. Also exact unbalanced OT with KL and quadratic regularization and the regularization path of UOT [41]
• Partial Wasserstein and Gromov-Wasserstein and Partial Fused Gromov-Wasserstein (exact [29] and entropic [3] formulations).
• Sliced Wasserstein [31, 32] and Max-sliced Wasserstein [35] that can be used for gradient flows [36].
• Wasserstein distance on the circle [44, 45] and Spherical Sliced Wasserstein [46]
• Graph Dictionary Learning solvers [38].
• Semi-relaxed (Fused) Gromov-Wasserstein divergences with corresponding barycenter solvers (exact and regularized [48]).
• Quantized (Fused) Gromov-Wasserstein distances [68].
• Efficient Discrete Multi Marginal Optimal Transport Regularization [50].
• Several backends for easy use of POT with Pytorch/jax/Numpy/Cupy/Tensorflow arrays.
• Smooth Strongly Convex Nearest Brenier Potentials [58], with an extension to bounding potentials using [59].
• Gaussian Mixture Model OT [69].
• Co-Optimal Transport [49] and unbalanced Co-Optimal Transport [71].
• Fused unbalanced Gromov-Wasserstein [70].
• Optimal Transport Barycenters for Generic Costs [77]
• Barycenters between Gaussian Mixture Models [69, 77]

POT provides the following Machine Learning related solvers:

• Optimal transport for domain adaptation with group lasso regularization, Laplacian regularization [5] [30] and semi supervised setting.
• Linear OT mapping [14] and Joint OT mapping estimation [8].
• Wasserstein Discriminant Analysis [11] (requires autograd + pymanopt).
• JCPOT algorithm for multi-source domain adaptation with target shift [27].
• Graph Neural Network OT layers TFGW [52] and TW (OT-GNN) [53]

Some other examples are available in the documentation.

Using and citing the toolbox

If you use this toolbox in your research and find it useful, please cite POT using the following references from the current version and from our JMLR paper:

Flamary R., Vincent-Cuaz C., Courty N., Gramfort A., Kachaiev O., Quang Tran H., David L., Bonet C., Cassereau N., Gnassounou T., Tanguy E., Delon J., Collas A., Mazelet S., Chapel L., Kerdoncuff T., Yu X., Feickert M., Krzakala P., Liu T., Fernandes Montesuma E. POT Python Optimal Transport (version 0.9.5). URL: https://github.com/PythonOT/POT

Rémi Flamary, Nicolas Courty, Alexandre Gramfort, Mokhtar Z. Alaya, Aurélie Boisbunon, Stanislas Chambon, Laetitia Chapel, Adrien Corenflos, Kilian Fatras, Nemo Fournier, Léo Gautheron, Nathalie T.H. Gayraud, Hicham Janati, Alain Rakotomamonjy, Ievgen Redko, Antoine Rolet, Antony Schutz, Vivien Seguy, Danica J. Sutherland, Romain Tavenard, Alexander Tong, Titouan Vayer, POT Python Optimal Transport library, Journal of Machine Learning Research, 22(78):1−8, 2021. URL: https://pythonot.github.io/

In Bibtex format:

@misc{flamary2024pot,
  author = {Flamary, R{\'e}mi and Vincent-Cuaz, C{\'e}dric and Courty, Nicolas and Gramfort, Alexandre and Kachaiev, Oleksii and Quang Tran, Huy and David, Laurène and Bonet, Cl{\'e}ment and Cassereau, Nathan and Gnassounou, Th{\'e}o and Tanguy, Eloi and Delon, Julie and Collas, Antoine and Mazelet, Sonia and Chapel, Laetitia and Kerdoncuff, Tanguy and Yu, Xizheng and Feickert, Matthew and Krzakala, Paul and Liu, Tianlin and Fernandes Montesuma, Eduardo},
  title = {POT Python Optimal Transport (version 0.9.5)},
  url = {https://github.com/PythonOT/POT},
  year = {2024}
}

@article{flamary2021pot,
  author  = {R{\'e}mi Flamary and Nicolas Courty and Alexandre Gramfort and Mokhtar Z. Alaya and Aur{\'e}lie Boisbunon and Stanislas Chambon and Laetitia Chapel and Adrien Corenflos and Kilian Fatras and Nemo Fournier and L{\'e}o Gautheron and Nathalie T.H. Gayraud and Hicham Janati and Alain Rakotomamonjy and Ievgen Redko and Antoine Rolet and Antony Schutz and Vivien Seguy and Danica J. Sutherland and Romain Tavenard and Alexander Tong and Titouan Vayer},
  title   = {POT: Python Optimal Transport},
  journal = {Journal of Machine Learning Research},
  year    = {2021},
  volume  = {22},
  number  = {78},
  pages   = {1-8},
  url     = {http://jmlr.org/papers/v22/20-451.html}
}

Installation

The library has been tested on Linux, MacOSX and Windows. It requires a C++ compiler for building/installing the EMD solver and relies on the following Python modules:

• Numpy (>=1.16)
• Scipy (>=1.0)
• Cython (>=0.23) (build only, not necessary when installing from pip or conda)

Pip installation

You can install the toolbox through PyPI with:

pip install POT

or get the very latest version by running:

pip install -U https://github.com/PythonOT/POT/archive/master.zip # with --user for user install (no root)

Optional dependencies may be installed with

pip install POT[all]

Note that this installs cvxopt, which is licensed under GPL 3.0. Alternatively, if you cannot use GPL-licensed software, the specific optional dependencies may be installed individually, or per-submodule. The available optional installations are backend-jax, backend-tf, backend-torch, cvxopt, dr, gnn, all.

Anaconda installation with conda-forge

If you use the Anaconda python distribution, POT is available in conda-forge. To install it and the required dependencies:

conda install -c conda-forge pot

Post installation check

After a correct installation, you should be able to import the module without errors:

import ot

Note that for easier access the module is named ot instead of pot.

Dependencies

Some sub-modules require additional dependencies which are discussed below

• ot.dr (Wasserstein dimensionality reduction) depends on autograd and pymanopt that can be installed with:

pip install pymanopt autograd

Examples

Short examples

• Import the toolbox

import ot

• Compute Wasserstein distances

# a,b are 1D histograms (sum to 1 and positive)
# M is the ground cost matrix

# With the unified  API :
Wd = ot.solve(M, a, b).value # exact linear program
Wd_reg = ot.solve(M, a, b, reg=reg).value # entropic regularized OT

# With the old API :
Wd = ot.emd2(a, b, M) # exact linear program
Wd_reg = ot.sinkhorn2(a, b, M, reg) # entropic regularized OT
# if b is a matrix compute all distances to a and return a vector

• Compute OT matrix

```python

a,b are 1D histograms (sum to 1 and positive)

M is the ground cost matrix