Python Optimal Transport Library (POT)

Warnings

• breaking The Gromov-Wasserstein (GW) solvers underwent a major refactor in version 0.9.0, leading to significant performance gains and the ability to handle non-symmetric cost matrices. While the API generally remained consistent, users relying on specific internal behaviors or numerical properties of older GW implementations might observe changes in results or performance characteristics.
  Fix: Review the official documentation and examples for GW solvers if migrating from versions older than 0.9.0. Verify results, especially when dealing with non-symmetric matrices.
• gotcha POT supports multiple array backends (NumPy, PyTorch, JAX, CuPy) for computation. By default, it uses NumPy. To leverage GPU acceleration (e.g., with CuPy or PyTorch on CUDA), users must explicitly configure `ot.backend` or ensure their input arrays are of the desired backend type (e.g., CuPy arrays for `ot.gpu` functions). Mixing backends or incorrect setup can lead to errors or unexpected CPU-only computation.
  Fix: For GPU or specific backend usage, import and configure `ot.backend` (e.g., `import ot.backend as ob; ob.set_backend('torch', 'cuda')`) or ensure all input tensors are compatible with the desired backend.
• gotcha Input array dimensions are critical and frequently a source of errors. For example, marginal distributions `a` and `b` are typically 1D arrays, while coordinates `X` and `Y` are 2D arrays (n_samples, n_features), and cost matrices `M` are 2D (n_samples_source, n_samples_target). Mismatched dimensions (e.g., `(n,)` instead of `(n,1)` for single-feature coordinates or transposed cost matrices) will lead to runtime errors.
  Fix: Always consult the specific function's documentation for expected input shapes. Use `.reshape()` or `.T` carefully to ensure arrays conform to the required dimensions.
• gotcha For many optimal transport problems, particularly those with a probabilistic interpretation, the marginal distributions `a` and `b` are expected to sum to 1. While some solvers might handle unnormalized inputs, it's best practice to normalize them (e.g., `a = a / np.sum(a)`) to ensure correct interpretation and avoid potential numerical instabilities in certain algorithms.
  Fix: Normalize marginal distributions `a` and `b` such that `np.sum(a) == 1` and `np.sum(b) == 1` before passing them to POT functions, unless the specific function documentation explicitly states otherwise for unbalanced OT.
• gotcha Optimal transport problems, especially exact EMD, can be computationally very expensive for large numbers of samples. While POT provides efficient C/Cython implementations, exact solvers scale poorly (e.g., cubic complexity for EMD). For large-scale applications, consider using entropic regularized solvers (Sinkhorn) or specialized approximate methods which trade off accuracy for speed.
  Fix: For N > 1000 samples, favor `ot.sinkhorn` or other regularized/approximate solvers over `ot.emd`. Explore techniques like sub-sampling, multi-scale, or barycentric mapping for further scalability.

Install

• pip install pot Install stable release

Imports

• ot

  import ot

  The entire library is typically imported under the alias 'ot'.

Quickstart

This example demonstrates how to compute the Earth Mover's Distance (EMD) between two 1D samples using POT's core `ot.emd` function. It covers generating samples, defining uniform marginal distributions, computing a normalized cost matrix, and finally, calculating the optimal transport plan and its total cost.

import numpy as np
import ot

# Generate two 1D samples
n = 100
np.random.seed(0)
xs = np.random.normal(0, 1, n)
xt = np.random.normal(5, 1, n)

# Histogram counts (uniform distribution)
a = np.ones(n) / n
b = np.ones(n) / n

# Cost matrix: squared Euclidean distance
M = ot.dist(xs.reshape((n, 1)), xt.reshape((n, 1)))
M /= M.max() # Normalize cost matrix for stability

# Compute Earth Mover's Distance (EMD) / Wasserstein-1 distance
G = ot.emd(a, b, M)
print(f"Optimal Transport plan (first 5x5):
{G[:5,:5]}")
print(f"EMD cost: {np.sum(G * M)}")