POT

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 [73] 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].
• Sliced Unbalanced OT and Unbalanced Sliced OT [82]
• 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]
• Fast and accurate transport bijections using BSP-OT [84]
• Sliced Transport Plans [85, 86, 87]

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:

…

In Bibtex format:

…

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 git+https://github.com/PythonOT/POT.git # 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

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

# With the unified API :
T = ot.solve(M, a, b).plan # exact linear program
T_reg = ot.solve(M, a, b, reg=reg).plan # entropic regularized OT

# With the old API :
T = ot.emd(a, b, M) # exact linear program
T_reg = ot.sinkhorn(a, b, M, reg) # entropic regularized OT

• Compute OT on empirical distributions

# X and Y are two 2D arrays of shape (n_samples, n_features)

# with squared euclidean metric
T = ot.solve_sample(X, Y).plan # exact linear program
T_reg = ot.solve_sample(X, Y, reg=reg).plan # entropic regularized OT

Wass_2 = ot.solve_sample(X, Y).value # Squared Wasserstein_2
Wass_1 = ot.