API and modules

Main ot functions

ot.barycenter(A, M, reg, weights=None, method='sinkhorn', numItermax=10000, stopThr=0.0001, verbose=False, log=False, warn=True, **kwargs)[source]

Compute the entropic regularized wasserstein barycenter of distributions \(\mathbf{A}\)

The function solves the following optimization problem:

\[\mathbf{a} = \mathop{\arg \min}_\mathbf{a} \quad \sum_i W_{reg}(\mathbf{a},\mathbf{a}_i)\]

where :

• \(W_{reg}(\cdot,\cdot)\) is the entropic regularized Wasserstein distance (see ot.bregman.sinkhorn()) if method is sinkhorn or sinkhorn_stabilized or sinkhorn_log.

• \(\mathbf{a}_i\) are training distributions in the columns of matrix \(\mathbf{A}\)

• reg and \(\mathbf{M}\) are respectively the regularization term and the cost matrix for OT

The algorithm used for solving the problem is the Sinkhorn-Knopp matrix scaling algorithm as proposed in [3]

Parameters:

• A (array-like, shape (dim, n_hists)) – n_hists training distributions \(\mathbf{a}_i\) of size dim

• M (array-like, shape (dim, dim)) – loss matrix for OT

• reg (float) – Regularization term > 0

• method (str (optional)) – method used for the solver either ‘sinkhorn’ or ‘sinkhorn_stabilized’ or ‘sinkhorn_log’

• weights (array-like, shape (n_hists,)) – Weights of each histogram \(\mathbf{a}_i\) on the simplex (barycentric coordinates)

• numItermax (int, optional) – Max number of iterations

• stopThr (float, optional) – Stop threshold on error (>0)

• verbose (bool, optional) – Print information along iterations

• log (bool, optional) – record log if True

• warn (bool, optional) – if True, raises a warning if the algorithm doesn’t convergence.

Returns:

• a ((dim,) array-like) – Wasserstein barycenter

• log (dict) – log dictionary return only if log==True in parameters

References

ot.barycenter_unbalanced(A, M, reg, reg_m, method='sinkhorn', weights=None, numItermax=1000, stopThr=1e-06, verbose=False, log=False, **kwargs)[source]

Compute the entropic unbalanced wasserstein barycenter of \(\mathbf{A}\).

The function solves the following optimization problem with \(\mathbf{a}\)

\[\mathbf{a} = \mathop{\arg \min}_\mathbf{a} \quad \sum_i W_{u_{reg}}(\mathbf{a},\mathbf{a}_i)\]

where :

• \(W_{u_{reg}}(\cdot,\cdot)\) is the unbalanced entropic regularized Wasserstein distance (see ot.unbalanced.sinkhorn_unbalanced())

• \(\mathbf{a}_i\) are training distributions in the columns of matrix \(\mathbf{A}\)

• reg and \(\mathbf{M}\) are respectively the regularization term and the cost matrix for OT

• reg_mis the marginal relaxation hyperparameter

The algorithm used for solving the problem is the generalized Sinkhorn-Knopp matrix scaling algorithm as proposed in [10]

Parameters:

• A (array-like, shape (dim, n_hists)) – n_hists training distributions \(\mathbf{a}_i\) of dimension dim

• M (array-like, shape (dim, dim)) – ground metric matrix for OT.

• reg (float) – Entropy regularization term > 0

• reg_m (float) – Marginal relaxation term > 0

• weights (array-like, shape (n_hists,) optional) – Weight of each distribution (barycentric coordinates) If None, uniform weights are used.

• numItermax (int, optional) – Max number of iterations

• stopThr (float, optional) – Stop threshold on error (> 0)

• verbose (bool, optional) – Print information along iterations

• log (bool, optional) – record log if True

Returns:

• a (array-like, shape (dim,)) – Unbalanced Wasserstein barycenter

• log (dict) – log dictionary return only if log==True in parameters

References

ot.binary_search_circle(u_values, v_values, u_weights=None, v_weights=None, p=1, Lm=10, Lp=10, tm=-1, tp=1, eps=1e-06, require_sort=True, log=False)[source]

Computes the Wasserstein distance on the circle using the Binary search algorithm proposed in [44]. Samples need to be in \(S^1\cong [0,1[\). If they are on \(\mathbb{R}\), takes the value modulo 1. If the values are on \(S^1\subset\mathbb{R}^2\), it is required to first find the coordinates using e.g. the atan2 function.

\[W_p^p(u,v) = \inf_{\theta\in\mathbb{R}}\int_0^1 |F_u^{-1}(q) - (F_v-\theta)^{-1}(q)|^p\ \mathrm{d}q\]

where:

• \(F_u\) and \(F_v\) are respectively the cdfs of \(u\) and \(v\)

For values \(x=(x_1,x_2)\in S^1\), it is required to first get their coordinates with

\[u = \frac{\pi + \mathrm{atan2}(-x_2,-x_1)}{2\pi}\]

using e.g. ot.utils.get_coordinate_circle(x)

The function runs on backend but tensorflow and jax are not supported.

Parameters:

• u_values (ndarray, shape (n, ...)) – samples in the source domain (coordinates on [0,1[)

• v_values (ndarray, shape (n, ...)) – samples in the target domain (coordinates on [0,1[)

• u_weights (ndarray, shape (n, ...), optional) – samples weights in the source domain

• v_weights (ndarray, shape (n, ...), optional) – samples weights in the target domain

• p (float, optional (default=1)) – Power p used for computing the Wasserstein distance

• Lm (int, optional) – Lower bound dC

• Lp (int, optional) – Upper bound dC

• tm (float, optional) – Lower bound theta

• tp (float, optional) – Upper bound theta

• eps (float, optional) – Stopping condition

• require_sort (bool, optional) – If True, sort the values.

• log (bool, optional) – If True, returns also the optimal theta

Returns:

• loss (float/array-like, shape (…)) – Batched cost associated to the optimal transportation

• log (dict, optional) – log dictionary returned only if log==True in parameters

Examples

>>> u = np.array([[0.2,0.5,0.8]])%1
>>> v = np.array([[0.4,0.5,0.7]])%1
>>> binary_search_circle(u.T, v.T, p=1)
array([0.1])

References

ot.compute_bspot_bijection(X, Y, n_plans=64, p=2, initial_perm=None, gaussian_slicing='auto', seed=0)[source]

This solver provides a good and fast approximation of the combinatorial problem of finding a bijection between two point clouds that minimizes the transport cost:

\[\min_{\sigma \in S_n} \sum_{i=1}^n \|X_i - Y_{\sigma(i)}\|_p^p\]

To do so, it generates \(n_{plans}\) random bijective BSP matchings, merges them together to obtain a bijection of low transport cost. Log-linear complexity in the number of points. Algorithm 2 & 3 from [84].

Note

There is no guarantee on the quality of the returned bijection, but the method is highly scalable on the CPU. Worst cases are obtained between point clouds that are very similar (e.g. two samples from the same distribution), where the solver can get stuck in local minima, but works well when the point clouds are very different. The method also works best for the standard squared euclidean cost (p=2), as this cost enables efficient BSP construction heuristic (with a cubic dependence on the dimension, this feature is disabled for dimensions larger than 64).

Parameters:

• X (array-like, shape (n_samples, dimension))

• Y (array-like, shape (n_samples, dimension))

• n_plans (int) – The number of BSP Matchings used to compute the final bijection.

• p (int, optional) – The power of the ground metric (default 2 for squared euclidean, -1 for infinity norm).

• initial_perm (array-like, shape (n_samples,), optional) – Bijection to use for initializing merging (optional).

• gaussian_slicing ("auto" or bool, optional) – If true then uses the Gaussian slicing heuristic to improve matching quality. Comes with a cubic complexity with dimension. If ‘auto’, then the heuristic is used for dimensions smaller than 64, and disabled for larger dimensions.

• seed (int, optional) – Random seed for reproducibility (default 0).

Returns:

• cost (float) – The transport cost of the final bijection.

• perm (array-like, shape (n_samples,)) – The final bijection, stored as a permutation (e.g. a list of numbers) such that X[i] is assigned to Y[perm[i]].

• perms (array-like, shape (n_plans,n_samples)) – The intermediary bijections used to compute the final one.

ot.dist(x1, x2=None, metric='sqeuclidean', p=2, w=None, backend='auto', nx=None, use_tensor=False)[source]

Compute distance between samples in \(\mathbf{x_1}\) and \(\mathbf{x_2}\)

Note

This function is backend-compatible and will work on arrays from all compatible backends for the following metrics: ‘sqeuclidean’, ‘euclidean’, ‘cityblock’, ‘minkowski’, ‘cosine’, ‘correlation’.

Parameters:

• x1 (array-like, shape (n1,d)) – matrix with n1 samples of size d

• x2 (array-like, shape (n2,d), optional) – matrix with n2 samples of size d (if None then \(\mathbf{x_2} = \mathbf{x_1}\))

• metric (str | callable, optional) – ‘sqeuclidean’ or ‘euclidean’ on all backends. On numpy the function also accepts from the scipy.spatial.distance.cdist function : ‘braycurtis’, ‘canberra’, ‘chebyshev’, ‘cityblock’, ‘correlation’, ‘cosine’, ‘dice’, ‘euclidean’, ‘hamming’, ‘jaccard’, ‘kulczynski1’, ‘mahalanobis’, ‘matching’, ‘minkowski’, ‘rogerstanimoto’, ‘russellrao’, ‘seuclidean’, ‘sokalmichener’, ‘sokalsneath’, ‘sqeuclidean’, ‘wminkowski’, ‘yule’.

• p (float, optional) – p-norm for the Minkowski and the Weighted Minkowski metrics. Default value is 2.

• w (array-like, rank 1) – Weights for the weighted metrics.

• backend (str, optional) – Backend to use for the computation. If ‘auto’, the backend is automatically selected based on the input data. if ‘scipy’, the scipy.spatial.distance.cdist function is used (and gradients are detached).

• use_tensor (bool, optional) – If true use tensorized computation for the distance matrix which can cause memory issues for large datasets. Default is False and the parameter is used only for the ‘cityblock’ and ‘minkowski’ metrics.

• nx (Backend, optional) – Backend to perform computations on. If omitted, the backend defaults to that of x1.

Returns:

M – distance matrix computed with given metric

Return type:

array-like, shape (n1, n2)

ot.dist_batch(X1, X2=None, metric='sqeuclidean', p=2, nx=None)[source]

Batched version of ot.dist, use it to compute many distance matrices in parallel.

Parameters:

• X1 (array-like, shape (b,n1,d)) – b matrices with n1 samples of size d

• X2 (array-like, shape (b,n2,d), optional) – b matrices with n2 samples of size d (if None then \(\mathbf{X_2} = \mathbf{X_1}\))

• metric (str, optional) – ‘sqeuclidean’, ‘euclidean’, ‘minkowski’ or ‘kl’

• p (float, optional) – p-norm for the Minkowski metrics. Default value is 2.

• nx (Backend, optional) – Backend to perform computations on. If omitted, the backend defaults to that of x1.

Returns:

M – distance matrix computed with given metric

Return type:

array-like, shape (b, n1, n2)

Examples

>>> import numpy as np
>>> from ot.batch import dist_batch
>>> X1 = np.random.randn(5, 10, 3)
>>> X2 = np.random.randn(5, 15, 3)
>>> M = dist_batch(X1, X2, metric="euclidean")
>>> M.shape
(5, 10, 15)

See also

ot.dist

equivalent non-batched function.

ot.emd(a, b, M, numItermax=100000, log=False, center_dual=True, numThreads=1, check_marginals=True, potentials_init=None)[source]

Solves the Earth Movers distance problem and returns the OT matrix

\[ \begin{align}\begin{aligned}\gamma = \mathop{\arg \min}_\gamma \quad \langle \gamma, \mathbf{M} \rangle_F\\s.t. \ \gamma \mathbf{1} = \mathbf{a}\\ \gamma^T \mathbf{1} = \mathbf{b}\\ \gamma \geq 0\end{aligned}\end{align} \]

where :

• \(\mathbf{M}\) is the metric cost matrix

• \(\mathbf{a}\) and \(\mathbf{b}\) are the sample weights

Warning

Note that the \(\mathbf{M}\) matrix in numpy needs to be a C-order numpy.array in float64 format. It will be converted if not in this format

Note

This function is backend-compatible and will work on arrays from all compatible backends. But the algorithm uses the C++ CPU backend which can lead to copy overhead on GPU arrays.

Note

This function will cast the computed transport plan to the data type of the provided input with the following priority: \(\mathbf{a}\), then \(\mathbf{b}\), then \(\mathbf{M}\) if marginals are not provided. Casting to an integer tensor might result in a loss of precision. If this behaviour is unwanted, please make sure to provide a floating point input.

Note

An error will be raised if the vectors \(\mathbf{a}\) and \(\mathbf{b}\) do not sum to the same value.

Uses the algorithm proposed in [1].

Parameters:

• a ((ns,) array-like, float) – Source histogram (uniform weight if empty list)

• b ((nt,) array-like, float) – Target histogram (uniform weight if empty list)

• M ((ns,nt) array-like or sparse matrix, float) –

  Loss matrix. Can be:

  • Dense array (c-order array in numpy with type float64)

  • Sparse matrix in backend’s format (scipy.sparse.coo_matrix for NumPy backend, torch.sparse_coo_tensor for PyTorch backend, etc.)

• numItermax (int, optional (default=100000)) – The maximum number of iterations before stopping the optimization algorithm if it has not converged.

• log (bool, optional (default=False)) – If True, returns a dictionary containing the cost and dual variables. Otherwise returns only the optimal transportation matrix.

• center_dual (boolean, optional (default=True)) – If True, centers the dual potential using function ot.lp.center_ot_dual().

• numThreads (int or "max", optional (default=1)) – Deprecated compatibility parameter. The network simplex solver no longer uses OpenMP, so this parameter is ignored.

• check_marginals (bool, optional (default=True)) – If True, checks that the marginals mass are equal. If False, skips the check.

• potentials_init (tuple of two arrays (alpha, beta), optional (default=None)) – Warmstart dual potentials to accelerate convergence. Should be a tuple (alpha, beta) where alpha is shape (ns,) and beta is shape (nt,). These potentials are used to guide initial pivots in the network simplex. Typically obtained from a previous EMD solve or Sinkhorn approximation.

• note: (..) –

  The solver automatically detects sparse format using the backend’s: issparse() method. For sparse inputs:

  • Uses a memory-efficient sparse EMD algorithm

  • Returns the transport plan as a sparse matrix in the backend’s format

  • Supports scipy.sparse (NumPy), torch.sparse (PyTorch), etc.

  • JAX and TensorFlow backends don’t support sparse matrices

Returns:

• gamma (array-like or sparse matrix, shape (ns, nt)) – Optimal transportation matrix for the given parameters.

  • For dense inputs: returns a dense array

  • For sparse inputs: returns a sparse matrix in the backend’s format (e.g., scipy.sparse.coo_matrix for NumPy, torch.sparse_coo_tensor for PyTorch)

• log (dict, optional) – If input log is True, a dictionary containing the cost, dual variables, and exit status.

Examples

Simple example with obvious solution. The function emd accepts lists and perform automatic conversion to numpy arrays

>>> import ot
>>> a=[.5,.5]
>>> b=[.5,.5]
>>> M=[[0.,1.],[1.,0.]]
>>> ot.emd(a, b, M)
array([[0.5, 0. ],
       [0. , 0.5]])

References

See also

ot.bregman.sinkhorn

Entropic regularized OT

ot.optim.cg

General regularized OT

ot.emd2(a, b, M, processes=1, numItermax=100000, log=False, return_matrix=False, center_dual=True, numThreads=1, check_marginals=True, potentials_init=None)[source]

Solves the Earth Movers distance problem and returns the loss

\[ \begin{align}\begin{aligned}\min_\gamma \quad \langle \gamma, \mathbf{M} \rangle_F\\s.t. \ \gamma \mathbf{1} = \mathbf{a}\\ \gamma^T \mathbf{1} = \mathbf{b}\\ \gamma \geq 0\end{aligned}\end{align} \]

where :

• \(\mathbf{M}\) is the metric cost matrix

• \(\mathbf{a}\) and \(\mathbf{b}\) are the sample weights

Note

This function is backend-compatible and will work on arrays from all compatible backends. But the algorithm uses the C++ CPU backend which can lead to copy overhead on GPU arrays.

Note

This function will cast the computed transport plan and transportation loss to the data type of the provided input with the following priority : \(\mathbf{a}\), then \(\mathbf{b}\), then \(\mathbf{M}\) if marginals are not provided. Casting to an integer tensor might result in a loss of precision. If this behaviour is unwanted, please make sure to provide a floating point input.

Note

An error will be raised if the vectors \(\mathbf{a}\) and \(\mathbf{b}\) do not sum to the same value.

Uses the algorithm proposed in [1].

Parameters:

• a ((ns,) array-like, float64) – Source histogram (uniform weight if empty list)

• b ((nt,) array-like, float64) – Target histogram (uniform weight if empty list)

• M ((ns,nt) array-like or sparse matrix, float64) –

  Loss matrix. Can be:

  • Dense array (c-order array in numpy with type float64)

  • Sparse matrix in backend’s format (scipy.sparse.coo_matrix for NumPy backend, torch.sparse_coo_tensor for PyTorch backend, etc.)

• processes (int, optional (default=1)) – Nb of processes used for multiple emd computation (deprecated)

• numItermax (int, optional (default=100000)) – The maximum number of iterations before stopping the optimization algorithm if it has not converged.

• log (boolean, optional (default=False)) – If True, returns a dictionary containing the cost and dual variables. Otherwise returns only the optimal transportation cost.

• return_matrix (boolean, optional (default=False)) – If True, returns the optimal transportation matrix in the log.

• center_dual (boolean, optional (default=True)) – If True, centers the dual potential using function ot.lp.center_ot_dual().

• numThreads (int or "max", optional (default=1)) – Deprecated compatibility parameter. The network simplex solver no longer uses OpenMP, so this parameter is ignored.

• check_marginals (bool, optional (default=True)) – If True, checks that the marginals mass are equal. If False, skips the check.

• potentials_init (tuple of two arrays (alpha, beta), optional (default=None)) – Warmstart dual potentials to accelerate convergence. Should be a tuple (alpha, beta) where alpha is shape (ns,) and beta is shape (nt,). These potentials are used to guide initial pivots in the network simplex. Typically obtained from a previous EMD solve or Sinkhorn approximation.

• note: (..) –

  The solver automatically detects sparse format using the backend’s: issparse() method. For sparse inputs:

  • Uses a memory-efficient sparse EMD algorithm

  • Edges not included are treated as having infinite cost (forbidden)

  • Supports scipy.sparse (NumPy), torch.sparse (PyTorch), etc.

  • JAX and TensorFlow backends don’t support sparse matrices

Returns:

• W (float, array-like) – Optimal transportation loss for the given parameters

• log (dict) – If input log is true, a dictionary containing the cost, dual variables (u, v), exit status, and optionally the optimal transportation matrix (G) if return_matrix is True

Examples

Simple example with obvious solution. The function emd accepts lists and perform automatic conversion to numpy arrays

>>> import ot
>>> a=[.5,.5]
>>> b=[.5,.5]
>>> M=[[0.,1.],[1.,0.]]
>>> ot.emd2(a,b,M)
0.0

References

See also

ot.bregman.sinkhorn

Entropic regularized OT

ot.optim.cg

General regularized OT

ot.emd2_1d(x_a, x_b, a=None, b=None, metric='sqeuclidean', p=1.0, dense=True, log=False)[source]

Solves the Earth Movers distance problem between 1d measures and returns the loss

\[ \begin{align}\begin{aligned}\gamma = arg\min_\gamma \sum_i \sum_j \gamma_{ij} d(x_a[i], x_b[j])\\s.t. \gamma 1 = a, \gamma^T 1= b, \gamma\geq 0\end{aligned}\end{align} \]

where :

• d is the metric

• \(x_a\) and \(x_b\) are the samples

• a and b are the sample weights

This implementation only supports metrics of the form \(d(x, y) = |x - y|^p\).

Uses the algorithm detailed in [1]_

Parameters:

• x_a (ndarray of float64, shape (ns,) or (ns, 1)) – Source dirac locations (on the real line)

• x_b (ndarray of float64, shape (nt,) or (ns, 1)) – Target dirac locations (on the real line)

• a (ndarray of float64, shape (ns,), optional) – Source histogram (default is uniform weight)

• b (ndarray of float64, shape (nt,), optional) – Target histogram (default is uniform weight)

• metric (str, optional (default='sqeuclidean')) – Metric to be used. Only works with either of the strings ‘sqeuclidean’, ‘minkowski’, ‘cityblock’, or ‘euclidean’.

• p (float, optional (default=1.0)) – The p-norm to apply for if metric=’minkowski’

• dense (boolean, optional (default=True)) – If True, returns \(\gamma\) as a dense ndarray of shape (ns, nt). Otherwise returns a sparse representation using scipy’s coo_matrix format. Only used if log is set to True. Due to implementation details, this function runs faster when dense is set to False.

• log (boolean, optional (default=False)) – If True, returns a dictionary containing the transportation matrix. Otherwise returns only the loss.

Returns:

• loss (float) – Cost associated to the optimal transportation

• log (dict) – If input log is True, a dictionary containing the Optimal transportation matrix for the given parameters

Examples

Simple example with obvious solution. The function emd2_1d accepts lists and performs automatic conversion to numpy arrays

>>> import ot
>>> a=[.5, .5]
>>> b=[.5, .5]
>>> x_a = [2., 0.]
>>> x_b = [0., 3.]
>>> ot.emd2_1d(x_a, x_b, a, b)
0.5
>>> ot.emd2_1d(x_a, x_b)
0.5

References

See also

ot.lp.emd2

EMD for multidimensional distributions

ot.lp.emd_1d

EMD for 1d distributions (returns the transportation matrix instead of the cost)

ot.emd2_lazy(X_a, X_b, a=None, b=None, metric='sqeuclidean', numItermax=100000, log=False, return_matrix=True, center_dual=True, check_marginals=True, potentials_init=None)[source]

Solves the Earth Movers distance problem with lazy cost computation and returns the loss

\[ \begin{align}\begin{aligned}\min_\gamma \quad \langle \gamma, \mathbf{M}(\mathbf{X}_a, \mathbf{X}_b) \rangle_F\\s.t. \ \gamma \mathbf{1} = \mathbf{a}\\ \gamma^T \mathbf{1} = \mathbf{b}\\ \gamma \geq 0\end{aligned}\end{align} \]

where :

• \(\mathbf{M}(\mathbf{X}_a, \mathbf{X}_b)\) is computed on-the-fly from coordinates

• \(\mathbf{a}\) and \(\mathbf{b}\) are the sample weights

Note

This function computes distances on-the-fly during the network simplex algorithm, avoiding the O(ns*nt) memory cost of pre-computing the full cost matrix. Memory usage is O(ns+nt) instead.

Note

This function is backend-compatible and will work on arrays from all compatible backends. But the algorithm uses the C++ CPU backend which can lead to copy overhead on GPU arrays.

Parameters:

• X_a ((ns, dim) array-like, float64) – Source sample coordinates

• X_b ((nt, dim) array-like, float64) – Target sample coordinates

• a ((ns,) array-like, float64, optional) – Source histogram (uniform weight if None)

• b ((nt,) array-like, float64, optional) – Target histogram (uniform weight if None)

• metric (str, optional (default='sqeuclidean')) –

  Distance metric for cost computation. Options:

  • ’sqeuclidean’: Squared Euclidean distance

  • ’euclidean’: Euclidean distance

  • ’cityblock’: Manhattan/L1 distance

• numItermax (int, optional (default=100000)) – Maximum number of iterations before stopping if not converged

• log (boolean, optional (default=False)) – If True, returns a dictionary containing the cost, dual variables, and optionally the transport plan (sparse format)

• return_matrix (boolean, optional (default=False)) – If True, returns the optimal transportation matrix in the log (sparse format)

• center_dual (boolean, optional (default=True)) – If True, centers the dual potential using ot.lp.center_ot_dual()

• check_marginals (bool, optional (default=True)) – If True, checks that the marginals mass are equal

• potentials_init (tuple of (ns,) and (nt,) arrays, optional) – Initial dual potentials (u, v) to warmstart the solver. If provided, the solver starts from these potentials instead of a cold start.

Returns:

• W (float) – Optimal transportation loss

• log (dict) – If input log is True, a dictionary containing:

  • cost: the optimal transportation cost

  • u, v: dual variables

  • warning: solver status message

  • result_code: solver return code

  • G: (optional) sparse transport plan if return_matrix=True

See also

ot.emd2

EMD with pre-computed cost matrix

ot.lp.emd_c_lazy

Low-level C++ lazy solver

ot.emd_1d(x_a, x_b, a=None, b=None, metric='sqeuclidean', p=1.0, dense=True, log=False, check_marginals=True)[source]

Solves the Earth Movers distance problem between 1d measures and returns the OT matrix

\[ \begin{align}\begin{aligned}\gamma = arg\min_\gamma \sum_i \sum_j \gamma_{ij} d(x_a[i], x_b[j])\\s.t. \gamma 1 = a, \gamma^T 1= b, \gamma\geq 0\end{aligned}\end{align} \]

where :

• d is the metric

• \(x_a\) and \(x_b\) are the samples

• a and b are the sample weights

This implementation only supports metrics of the form \(d(x, y) = |x - y|^p\).

Uses the algorithm detailed in [1]_

Parameters:

• x_a (ndarray of float64, shape (ns,) or (ns, 1)) – Source dirac locations (on the real line)

• x_b (ndarray of float64, shape (nt,) or (ns, 1)) – Target dirac locations (on the real line)

• a (ndarray of float64, shape (ns,), optional) – Source histogram (default is uniform weight)

• b (ndarray of float64, shape (nt,), optional) – Target histogram (default is uniform weight)

• metric (str, optional (default='sqeuclidean')) – Metric to be used. Only works with either of the strings ‘sqeuclidean’, ‘minkowski’, ‘cityblock’, or ‘euclidean’.

• p (float, optional (default=1.0)) – The p-norm to apply for if metric=’minkowski’

• dense (boolean, optional (default=True)) – If True, returns \(\gamma\) as a dense ndarray of shape (ns, nt). Otherwise returns a sparse representation using scipy’s coo_matrix format. Due to implementation details, this function runs faster when ‘sqeuclidean’, ‘minkowski’, ‘cityblock’, or ‘euclidean’ metrics are used.

• log (boolean, optional (default=False)) – If True, returns a dictionary containing the cost. Otherwise returns only the optimal transportation matrix.

• check_marginals (bool, optional (default=True)) – If True, checks that the marginals mass are equal. If False, skips the check.

Returns:

• gamma (ndarray, shape (ns, nt)) – Optimal transportation matrix for the given parameters

• log (dict) – If input log is True, a dictionary containing the cost and the indices of the non-zero elements of the transportation matrix

Examples

Simple example with obvious solution. The function emd_1d accepts lists and performs automatic conversion to numpy arrays

>>> import ot
>>> a=[.5, .5]
>>> b=[.5, .5]
>>> x_a = [2., 0.]
>>> x_b = [0., 3.]
>>> ot.emd_1d(x_a, x_b, a, b)
array([[0. , 0.5],
       [0.5, 0. ]])
>>> ot.emd_1d(x_a, x_b)
array([[0. , 0.5],
       [0.5, 0. ]])

References

See also

ot.lp.emd

EMD for multidimensional distributions

ot.lp.emd2_1d

EMD for 1d distributions (returns cost instead of the transportation matrix)

ot.emd_1d_dual_backprop(u_values, v_values, u_weights=None, v_weights=None, p=1, require_sort=True)[source]

Computes the 1 dimensional OT loss between two (batched) empirical distributions

\[OT_{loss} = \int_0^1 |cdf_u^{-1}(q) - cdf_v^{-1}(q)|^p dq\]

and returns the dual potentials and the loss, i.e. such that

\[OT_{loss}(u,v) = \int f(x)\mathrm{d}u(x) + \int g(y)\mathrm{d}v(y).\]

Warning

This function only works in pytorch or jax as it backpropagates through the wasserstein_1d function.

Parameters:

• u_values (array-like, shape (n, ...)) – locations of the first empirical distribution

• v_values (array-like, shape (m, ...)) – locations of the second empirical distribution

• u_weights (array-like, shape (n, ...), optional) – weights of the first empirical distribution, if None then uniform weights are used

• v_weights (array-like, shape (m, ...), optional) – weights of the second empirical distribution, if None then uniform weights are used

• p (int, optional) – order of the ground metric used, should be at least 1, default is 1

• require_sort (bool, optional) – sort the distributions atoms locations, if False we will consider they have been sorted prior to being passed to the function, default is True

Returns:

• f (array-like shape (n, …)) – First dual potential

• g (array-like shape (m, …)) – Second dual potential

• loss (float/array-like, shape (…)) – the batched EMD

ot.expected_sliced_plan(X_s, X_t, a=None, b=None, projections=None, metric='sqeuclidean', p=2, n_projections=None, beta=0.0, seed=None, dense=True, batch_size=None, log=False)[source]

Computes the Expected Sliced cost and plan between two datasets X_s and X_t of shapes (ns, d) and (nt, d). Given a set of n_projections projection directions, the expected sliced plan is obtained by averaging the n_projections 1d optimal transport plans between the projections of X_s and X_t on each direction. Expected Sliced was introduced in [87] and further studied in [86].

Note

The computation ignores potential ambiguities in the projections: if two points from a same measure have the same projection on a direction, then multiple sorting permutations are possible. To avoid combinatorial explosion, only one permutation is retained: this strays from theory in pathological cases.

Warning

Tensorflow and jax only returns dense plans, as they do not support well sparse matrices.

Parameters:

• X_s (array-like, shape (ns, d)) – The first set of vectors.

• X_t (array-like, shape (nt, d)) – The second set of vectors.

• a (ndarray of float64, shape (ns,), optional) – Source histogram (default is uniform weight)

• b (ndarray of float64, shape (nt,), optional) – Target histogram (default is uniform weight)

• projections (shape (dim, n_projections), optional) – Projection matrix (n_projections and seed are not used in this case). Default is None

• metric (str, optional (default='sqeuclidean')) – Metric to be used. Only works with either of the strings ‘sqeuclidean’, ‘minkowski’, ‘cityblock’, or ‘euclidean’.

• p (float, optional (default=2)) – The p-norm to apply for if metric=’minkowski’

• n_projections (int, optional) – The number of projection directions. Required if projections is None.

• seed (int, optional) – The seed for the random number generator for sampling projections, in case projections is None. Default is None.

• beta (float, optional) – Inverse-temperature parameter which weights each projection’s contribution to the expected plan. Default is 0 (uniform weighting).

• dense (boolean, optional (default=True)) – If True, returns \(\gamma\) as a dense ndarray of shape (ns, nt). Otherwise returns a sparse representation using scipy’s coo_matrix format.

• batch_size (int, optional) – If specified, compute the distance in batches of size batch_size to avoid memory issues for large datasets. Default is None (no batching).

• log (bool, optional) – If True, returns additional logging information. Default is False.

Returns:

• plan (ndarray, shape (ns, nt) or coo_matrix if dense is False) – Optimal transportation matrix for the given parameters.

• cost (float) – The cost associated to the optimal permutation.

• log_dict (dict, optional) – A dictionary containing intermediate computations for logging purposes. Returned only if log is True.

References

Examples

>>> import ot
>>> import numpy as np
>>> x=np.array([[3.,3.], [1.,1.]])
>>> y=np.array([[2.,2.5], [3.,2.]])
>>> projections=np.array([[1, 0], [0, 1]])
>>> plan, cost = ot.expected_sliced_plan(x, y, projections=projections)
>>> plan
array([[0.25, 0.25],
       [0.25, 0.25]])
>>> cost
2.625

ot.factored_optimal_transport(Xa, Xb, a=None, b=None, reg=0.0, r=100, X0=None, stopThr=1e-07, numItermax=100, verbose=False, log=False, **kwargs)[source]

Solves factored OT problem and return OT plans and intermediate distribution

This function solve the following OT problem [40]_

\[\mathop{\arg \min}_\mu \quad W_2^2(\mu_a,\mu)+ W_2^2(\mu,\mu_b)\]

where :

• \(\mu_a\) and \(\mu_b\) are empirical distributions.

• \(\mu\) is an empirical distribution with r samples

And returns the two OT plans between

Note

This function is backend-compatible and will work on arrays from all compatible backends. But the algorithm uses the C++ CPU backend which can lead to copy overhead on GPU arrays.

Uses the conditional gradient algorithm to solve the problem proposed in [39].

Parameters:

• Xa ((ns,d) array-like, float) – Source samples

• Xb ((nt,d) array-like, float) – Target samples

• a ((ns,) array-like, float) – Source histogram (uniform weight if empty list)

• b ((nt,) array-like, float) – Target histogram (uniform weight if empty list))

• numItermax (int, optional) – Max number of iterations

• stopThr (float, optional) – Stop threshold on the relative variation (>0)

• verbose (bool, optional) – Print information along iterations

• log (bool, optional) – record log if True

Returns:

• Ga (array-like, shape (ns, r)) – Optimal transportation matrix between source and the intermediate distribution

• Gb (array-like, shape (r, nt)) – Optimal transportation matrix between the intermediate and target distribution

• X (array-like, shape (r, d)) – Support of the intermediate distribution

• log (dict, optional) – If input log is true, a dictionary containing the cost and dual variables and exit status

References

See also

ot.bregman.sinkhorn

Entropic regularized OT

ot.optim.cg

General regularized OT

ot.fused_gromov_wasserstein(M, C1, C2, p=None, q=None, loss_fun='square_loss', symmetric=None, alpha=0.5, armijo=False, G0=None, log=False, max_iter=10000.0, tol_rel=1e-09, tol_abs=1e-09, **kwargs)[source]

Returns the Fused Gromov-Wasserstein transport between \((\mathbf{C_1}, \mathbf{Y_1}, \mathbf{p})\) and \((\mathbf{C_2}, \mathbf{Y_2}, \mathbf{q})\) with pairwise distance matrix \(\mathbf{M}\) between node feature matrices \(\mathbf{Y_1}\) and \(\mathbf{Y_2}\) (see [24]).

The function solves the following optimization problem using Conditional Gradient:

\[ \begin{align}\begin{aligned}\mathbf{T}^* \in\mathop{\arg\min}_\mathbf{T} \quad (1 - \alpha) \langle \mathbf{T}, \mathbf{M} \rangle_F + \alpha \sum_{i,j,k,l} L(\mathbf{C_1}_{i,k}, \mathbf{C_2}_{j,l}) \mathbf{T}_{i,j} \mathbf{T}_{k,l}\\s.t. \ \mathbf{T} \mathbf{1} &= \mathbf{p}\\ \mathbf{T}^T \mathbf{1} &= \mathbf{q}\\ \mathbf{T} &\geq 0\end{aligned}\end{align} \]

Where :

• \(\mathbf{M}\): metric cost matrix between features across domains

• \(\mathbf{C_1}\): Metric cost matrix in the source space

• \(\mathbf{C_2}\): Metric cost matrix in the target space

• \(\mathbf{p}\): distribution in the source space

• \(\mathbf{q}\): distribution in the target space

• L: loss function to account for the misfit between the similarity and feature matrices

• \(\alpha\): trade-off parameter

Note

This function is backend-compatible and will work on arrays from all compatible backends. But the algorithm uses the C++ CPU backend which can lead to copy overhead on GPU arrays.

Note

All computations in the conditional gradient solver are done with numpy to limit memory overhead.

Note

This function will cast the computed transport plan to the data type of the provided input \(\mathbf{M}\). Casting to an integer tensor might result in a loss of precision. If this behaviour is unwanted, please make sure to provide a floating point input.

Parameters:

• M (array-like, shape (ns, nt)) – Metric cost matrix between features across domains

• C1 (array-like, shape (ns, ns)) – Metric cost matrix representative of the structure in the source space

• C2 (array-like, shape (nt, nt)) – Metric cost matrix representative of the structure in the target space

• p (array-like, shape (ns,), optional) – Distribution in the source space. If let to its default value None, uniform distribution is taken.

• q (array-like, shape (nt,), optional) – Distribution in the target space. If let to its default value None, uniform distribution is taken.

• loss_fun (str, optional) – Loss function used for the solver

• symmetric (bool, optional) – Either C1 and C2 are to be assumed symmetric or not. If let to its default None value, a symmetry test will be conducted. Else if set to True (resp. False), C1 and C2 will be assumed symmetric (resp. asymmetric).

• alpha (float, optional) – Trade-off parameter (0 < alpha < 1)

• armijo (bool, optional) – If True the step of the line-search is found via an armijo research. Else closed form is used. If there are convergence issues use False.

• G0 (array-like, shape (ns,nt), optional) – If None the initial transport plan of the solver is pq^T. Otherwise G0 must satisfy marginal constraints and will be used as initial transport of the solver.

• log (bool, optional) – record log if True

• max_iter (int, optional) – Max number of iterations

• tol_rel (float, optional) – Stop threshold on relative error (>0)

• tol_abs (float, optional) – Stop threshold on absolute error (>0)

• **kwargs (dict) – parameters can be directly passed to the ot.optim.cg solver

Returns:

• T (array-like, shape (ns, nt)) – Optimal transportation matrix for the given parameters.

• log (dict) – Log dictionary return only if log==True in parameters.

References

ot.fused_gromov_wasserstein2(M, C1, C2, p=None, q=None, loss_fun='square_loss', symmetric=None, alpha=0.5, armijo=False, G0=None, log=False, max_iter=10000.0, tol_rel=1e-09, tol_abs=1e-09, **kwargs)[source]

Returns the Fused Gromov-Wasserstein distance between \((\mathbf{C_1}, \mathbf{Y_1}, \mathbf{p})\) and \((\mathbf{C_2}, \mathbf{Y_2}, \mathbf{q})\) with pairwise distance matrix \(\mathbf{M}\) between node feature matrices \(\mathbf{Y_1}\) and \(\mathbf{Y_2}\) (see [24]).

The function solves the following optimization problem using Conditional Gradient:

\[ \begin{align}\begin{aligned}\mathbf{FGW} = \mathop{\min}_\mathbf{T} \quad (1 - \alpha) \langle \mathbf{T}, \mathbf{M} \rangle_F + \alpha \sum_{i,j,k,l} L(\mathbf{C_1}_{i,k}, \mathbf{C_2}_{j,l}) \mathbf{T}_{i,j} \mathbf{T}_{k,l}\\s.t. \ \mathbf{T} \mathbf{1} &= \mathbf{p}\\ \mathbf{T}^T \mathbf{1} &= \mathbf{q}\\ \mathbf{T} &\geq 0\end{aligned}\end{align} \]

Where :

• \(\mathbf{M}\): metric cost matrix between features across domains

• \(\mathbf{C_1}\): Metric cost matrix in the source space

• \(\mathbf{C_2}\): Metric cost matrix in the target space

• \(\mathbf{p}\): distribution in the source space

• \(\mathbf{q}\): distribution in the target space

• L: loss function to account for the misfit between the similarity and feature matrices

• \(\alpha\): trade-off parameter

Note that when using backends, this loss function is differentiable wrt the matrices (C1, C2, M) and weights (p, q) for quadratic loss using the gradients from [38]_.

Note

This function is backend-compatible and will work on arrays from all compatible backends. But the algorithm uses the C++ CPU backend which can lead to copy overhead on GPU arrays.

Note

All computations in the conditional gradient solver are done with numpy to limit memory overhead.

Note

This function will cast the computed transport plan to the data type of the provided input \(\mathbf{M}\). Casting to an integer tensor might result in a loss of precision. If this behaviour is unwanted, please make sure to provide a floating point input.

Parameters:

• M (array-like, shape (ns, nt)) – Metric cost matrix between features across domains

• C1 (array-like, shape (ns, ns)) – Metric cost matrix representative of the structure in the source space.

• C2 (array-like, shape (nt, nt)) – Metric cost matrix representative of the structure in the target space.

• p (array-like, shape (ns,), optional) – Distribution in the source space. If let to its default value None, uniform distribution is taken.

• q (array-like, shape (nt,), optional) – Distribution in the target space. If let to its default value None, uniform distribution is taken.

• loss_fun (str, optional) – Loss function used for the solver.

• symmetric (bool, optional) – Either C1 and C2 are to be assumed symmetric or not. If let to its default None value, a symmetry test will be conducted. Else if set to True (resp. False), C1 and C2 will be assumed symmetric (resp. asymmetric).

• alpha (float, optional) – Trade-off parameter (0 < alpha < 1)

• armijo (bool, optional) – If True the step of the line-search is found via an armijo research. Else closed form is used. If there are convergence issues use False.

• G0 (array-like, shape (ns,nt), optional) – If None the initial transport plan of the solver is pq^T. Otherwise G0 must satisfy marginal constraints and will be used as initial transport of the solver.

• log (bool, optional) – Record log if True.

• max_iter (int, optional) – Max number of iterations

• tol_rel (float, optional) – Stop threshold on relative error (>0)

• tol_abs (float, optional) – Stop threshold on absolute error (>0)

• **kwargs (dict) – Parameters can be directly passed to the ot.optim.cg solver.

Returns:

• fgw-distance (float) – Fused Gromov-Wasserstein distance for the given parameters.

• log (dict) – Log dictionary return only if log==True in parameters.

References

ot.gromov_barycenters(N, Cs, ps=None, p=None, lambdas=None, loss_fun='square_loss', symmetric=True, armijo=False, max_iter=1000, tol=1e-09, stop_criterion='barycenter', warmstartT=False, verbose=False, log=False, init_C=None, random_state=None, **kwargs)[source]

Returns the Gromov-Wasserstein barycenters of S measured similarity matrices \((\mathbf{C}_s)_{1 \leq s \leq S}\)

The function solves the following optimization problem with block coordinate descent:

\[\mathbf{C}^* = \mathop{\arg \min}_{\mathbf{C}\in \mathbb{R}^{N \times N}} \quad \sum_s \lambda_s \mathrm{GW}(\mathbf{C}, \mathbf{C}_s, \mathbf{p}, \mathbf{p}_s)\]

Where :

• \(\mathbf{C}_s\): metric cost matrix

• \(\mathbf{p}_s\): distribution

Parameters:

• N (int) – Size of the targeted barycenter

• Cs (list of S array-like of shape (ns, ns)) – Metric cost matrices

• ps (list of S array-like of shape (ns,), optional) – Sample weights in the S spaces. If let to its default value None, uniform distributions are taken.

• p (array-like, shape (N,), optional) – Weights in the targeted barycenter. If let to its default value None, uniform distribution is taken.

• lambdas (list of float, optional) – List of the S spaces’ weights. If let to its default value None, uniform weights are taken.

• loss_fun (callable, optional) – tensor-matrix multiplication function based on specific loss function

• symmetric (bool, optional.) – Either structures are to be assumed symmetric or not. Default value is True. Else if set to True (resp. False), C1 and C2 will be assumed symmetric (resp. asymmetric).

• armijo (bool, optional) – If True the step of the line-search is found via an armijo research. Else closed form is used. If there are convergence issues use False.

• max_iter (int, optional) – Max number of iterations

• tol (float, optional) – Stop threshold on relative error (>0)

• stop_criterion (str, optional. Default is 'barycenter'.) – Stop criterion taking values in [‘barycenter’, ‘loss’]. If set to ‘barycenter’ uses absolute norm variations of estimated barycenters. Else if set to ‘loss’ uses the relative variations of the loss.

• warmstartT (bool, optional) – Either to perform warmstart of transport plans in the successive fused gromov-wasserstein transport problems.s

• verbose (bool, optional) – Print information along iterations.

• log (bool, optional) – Record log if True.

• init_C (bool | array-like, shape(N,N)) – Random initial value for the \(\mathbf{C}\) matrix provided by user.

• random_state (int or RandomState instance, optional) – Fix the seed for reproducibility

Returns:

• C (array-like, shape (N, N)) – Similarity matrix in the barycenter space (permutated arbitrarily)

• log (dict) – Only returned when log=True. It contains the keys:

  • \(\mathbf{T}\): list of (N, ns) transport matrices

  • \(\mathbf{p}\): (N,) barycenter weights

  • values used in convergence evaluation.

References

ot.gromov_wasserstein(C1, C2, p=None, q=None, loss_fun='square_loss', symmetric=None, log=False, armijo=False, G0=None, max_iter=10000.0, tol_rel=1e-09, tol_abs=1e-09, **kwargs)[source]

Returns the Gromov-Wasserstein transport between \((\mathbf{C_1}, \mathbf{p})\) and \((\mathbf{C_2}, \mathbf{q})\).

The function solves the following optimization problem using Conditional Gradient:

\[ \begin{align}\begin{aligned}\mathbf{T}^* \in \mathop{\arg \min}_\mathbf{T} \quad \sum_{i,j,k,l} L(\mathbf{C_1}_{i,k}, \mathbf{C_2}_{j,l}) \mathbf{T}_{i,j} \mathbf{T}_{k,l}\\s.t. \ \mathbf{T} \mathbf{1} &= \mathbf{p}\\ \mathbf{T}^T \mathbf{1} &= \mathbf{q}\\ \mathbf{T} &\geq 0\end{aligned}\end{align} \]

Where :

• \(\mathbf{C_1}\): Metric cost matrix in the source space.

• \(\mathbf{C_2}\): Metric cost matrix in the target space.

• \(\mathbf{p}\): Distribution in the source space.

• \(\mathbf{q}\): Distribution in the target space.

• L: Loss function to account for the misfit between the similarity matrices.

Note

This function is backend-compatible and will work on arrays from all compatible backends. But the algorithm uses the C++ CPU backend which can lead to copy overhead on GPU arrays.

Note

All computations in the conditional gradient solver are done with numpy to limit memory overhead.

Note

This function will cast the computed transport plan to the data type of the provided input \(\mathbf{C}_1\). Casting to an integer tensor might result in a loss of precision. If this behaviour is unwanted, please make sure to provide a floating point input.

Parameters:

• C1 (array-like, shape (ns, ns)) – Metric cost matrix in the source space.

• C2 (array-like, shape (nt, nt)) – Metric cost matrix in the target space.

• p (array-like, shape (ns,), optional) – Distribution in the source space. If let to its default value None, uniform distribution is taken.

• q (array-like, shape (nt,), optional) – Distribution in the target space. If let to its default value None, uniform distribution is taken.

• loss_fun (str, optional) – Loss function used for the solver either ‘square_loss’ or ‘kl_loss’.

• symmetric (bool, optional) – Either C1 and C2 are to be assumed symmetric or not. If let to its default None value, a symmetry test will be conducted. Else if set to True (resp. False), C1 and C2 will be assumed symmetric (resp. asymmetric).

• verbose (bool, optional) – Print information along iterations.

• log (bool, optional) – Record log if True.

• armijo (bool, optional) – If True, the step of the line-search is found via an armijo search. Else closed form is used. If there are convergence issues, use False.

• G0 (array-like, shape (ns,nt), optional) – If None, the initial transport plan of the solver is pq^T. Otherwise G0 must satisfy marginal constraints and will be used as initial transport of the solver.

• max_iter (int, optional) – Max number of iterations.

• tol_rel (float, optional) – Stop threshold on relative error (>0).

• tol_abs (float, optional) – Stop threshold on absolute error (>0).

• **kwargs (dict) – Parameters can be directly passed to the ot.optim.cg solver.

Returns:

• T (array-like, shape (ns, nt)) –

  Coupling between the two spaces that minimizes:

  \(\sum_{i,j,k,l} L(\mathbf{C_1}_{i,k}, \mathbf{C_2}_{j,l}) \mathbf{T}_{i,j} \mathbf{T}_{k,l}\)

• log (dict) – Convergence information and loss.

References

ot.gromov_wasserstein2(C1, C2, p=None, q=None, loss_fun='square_loss', symmetric=None, log=False, armijo=False, G0=None, max_iter=10000.0, tol_rel=1e-09, tol_abs=1e-09, **kwargs)[source]

Returns the Gromov-Wasserstein loss \(\mathbf{GW}\) between \((\mathbf{C_1}, \mathbf{p})\) and \((\mathbf{C_2}, \mathbf{q})\). To recover the Gromov-Wasserstein distance as defined in [13] compute \(d_{GW} = \frac{1}{2} \sqrt{\mathbf{GW}}\).

The function solves the following optimization problem using Conditional Gradient:

\[ \begin{align}\begin{aligned}\mathbf{GW} = \min_\mathbf{T} \quad \sum_{i,j,k,l} L(\mathbf{C_1}_{i,k}, \mathbf{C_2}_{j,l}) \mathbf{T}_{i,j} \mathbf{T}_{k,l}\\s.t. \ \mathbf{T} \mathbf{1} &= \mathbf{p}\\ \mathbf{T}^T \mathbf{1} &= \mathbf{q}\\ \mathbf{T} &\geq 0\end{aligned}\end{align} \]

Where :

• \(\mathbf{C_1}\): Metric cost matrix in the source space

• \(\mathbf{C_2}\): Metric cost matrix in the target space

• \(\mathbf{p}\): distribution in the source space

• \(\mathbf{q}\): distribution in the target space

• L: loss function to account for the misfit between the similarity matrices

Note that when using backends, this loss function is differentiable wrt the matrices (C1, C2) and weights (p, q) for quadratic loss using the gradients from [38]_.

Note

This function is backend-compatible and will work on arrays from all compatible backends. But the algorithm uses the C++ CPU backend which can lead to copy overhead on GPU arrays.

Note

All computations in the conditional gradient solver are done with numpy to limit memory overhead.

Note

This function will cast the computed transport plan to the data type of the provided input \(\mathbf{C}_1\). Casting to an integer tensor might result in a loss of precision. If this behaviour is unwanted, please make sure to provide a floating point input.

Parameters:

• C1 (array-like, shape (ns, ns)) – Metric cost matrix in the source space

• C2 (array-like, shape (nt, nt)) – Metric cost matrix in the target space

• p (array-like, shape (ns,), optional) – Distribution in the source space. If let to its default value None, uniform distribution is taken.

• q (array-like, shape (nt,), optional) – Distribution in the target space. If let to its default value None, uniform distribution is taken.

• loss_fun (str) – loss function used for the solver either ‘square_loss’ or ‘kl_loss’

• symmetric (bool, optional) – Either C1 and C2 are to be assumed symmetric or not. If let to its default None value, a symmetry test will be conducted. Else if set to True (resp. False), C1 and C2 will be assumed symmetric (resp. asymmetric).

• verbose (bool, optional) – Print information along iterations

• log (bool, optional) – record log if True

• armijo (bool, optional) – If True the step of the line-search is found via an armijo research. Else closed form is used. If there are convergence issues use False.

• G0 (array-like, shape (ns,nt), optional) – If None the initial transport plan of the solver is pq^T. Otherwise G0 must satisfy marginal constraints and will be used as initial transport of the solver.

• max_iter (int, optional) – Max number of iterations

• tol_rel (float, optional) – Stop threshold on relative error (>0)

• tol_abs (float, optional) – Stop threshold on absolute error (>0)

• **kwargs (dict) – parameters can be directly passed to the ot.optim.cg solver

Returns:

• gw_dist (float) – Gromov-Wasserstein distance

• log (dict) – convergence information and Coupling matrix

References

ot.linear_circular_ot(u_values, v_values=None, u_weights=None, v_weights=None)[source]

Computes the Linear Circular Optimal Transport distance from [78] using \(\eta=\mathrm{Unif}(S^1)\) as reference measure. Samples need to be in \(S^1\cong [0,1[\). If they are on \(\mathbb{R}\), takes the value modulo 1. If the values are on \(S^1\subset\mathbb{R}^2\), it is required to first find the coordinates using e.g. the atan2 function.

General loss returned:

\[\mathrm{LCOT}_2^2(\mu, \nu) = \int_0^1 d_{S^1}\big(\hat{\mu}(t), \hat{\nu}(t)\big)^2\ \mathrm{d}t\]

where \(\hat{\mu}(x)=F_{\mu}^{-1}(x-\int z\mathrm{d}\mu(z)+\frac12) - x\) for all \(x\in [0,1[\), and \(d_{S^1}(x,y)=\min(|x-y|, 1-|x-y|)\) for \(x,y\in [0,1[\).

Parameters:

• u_values (ndarray, shape (n, ...)) – samples in the source domain (coordinates on [0,1[)

• v_values (ndarray, shape (n, ...), optional) – samples in the target domain (coordinates on [0,1[), if None, compute distance against uniform distribution

• u_weights (ndarray, shape (n, ...), optional) – samples weights in the source domain

• v_weights (ndarray, shape (n, ...), optional) – samples weights in the target domain

Returns:

loss – Batched cost associated to the linear optimal transportation

Return type:

float/array-like, shape (…)

Examples

>>> u = np.array([[0.2,0.5,0.8]])%1
>>> v = np.array([[0.4,0.5,0.7]])%1
>>> linear_circular_ot(u.T, v.T)
array([0.0127])

References

ot.linear_sliced_wasserstein_sphere(X_s, X_t=None, a=None, b=None, n_projections=50, projections=None, seed=None, log=False)[source]

Computes the linear spherical sliced wasserstein distance from [79].

General loss returned:

\[\mathrm{LSSOT}_2(\mu, \nu) = \left(\int_{\mathbb{V}_{d,2}} \mathrm{LCOT}_2^2(P^U_\#\mu, P^U_\#\nu)\ \mathrm{d}\sigma(U)\right)^{\frac12},\]

where \(\mu,\nu\in\mathcal{P}(S^{d-1})\) are two probability measures on the sphere, \(\mathrm{LCOT}_2\) is the linear circular optimal transport distance, and \(P^U_\# \mu\) stands for the pushforwards of the projection \(\forall x\in S^{d-1},\ P^U(x) = \frac{U^Tx}{\|U^Tx\|_2}\).

Parameters:

• X_s (ndarray, shape (n_samples_a, dim)) – Samples in the source domain

• X_t (ndarray, shape (n_samples_b, dim), optional) – Samples in the target domain. If None, computes the distance against the uniform distribution on the sphere.

• a (ndarray, shape (n_samples_a,), optional) – samples weights in the source domain

• b (ndarray, shape (n_samples_b,), optional) – samples weights in the target domain

• n_projections (int, optional) – Number of projections used for the Monte-Carlo approximation

• projections (shape (n_projections, dim, 2), optional) – Projection matrix (n_projections and seed are not used in this case)

• seed (int or RandomState or None, optional) – Seed used for random number generator

• log (bool, optional) – if True, linear_sliced_wasserstein_sphere returns the projections used and their associated LCOT.

Returns:

• cost (float) – Linear Spherical Sliced Wasserstein Cost

• log (dict, optional) – log dictionary return only if log==True in parameters

Examples

>>> import ot
>>> import numpy as np
>>> n_samples_a = 20
>>> X = np.random.normal(0., 1., (n_samples_a, 5))
>>> X = X / np.sqrt(np.sum(X**2, -1, keepdims=True))
>>> ot.linear_sliced_wasserstein_sphere(X, X, seed=0)
0.0

References

ot.lowrank_gromov_wasserstein_samples(X_s, X_t, a=None, b=None, reg=0, rank=None, alpha=1e-10, gamma_init='rescale', rescale_cost=True, cost_factorized_Xs=None, cost_factorized_Xt=None, stopThr=0.0001, numItermax=1000, stopThr_dykstra=0.001, numItermax_dykstra=10000, seed_init=49, warn=True, warn_dykstra=False, log=False)[source]

Solve the entropic regularization Gromov-Wasserstein transport problem under low-nonnegative rank constraints on the couplings and cost matrices.

Squared euclidean distance matrices are considered for the target and source distributions.

The function solves the following optimization problem:

\[\mathop{\min_{(Q,R,g) \in \mathcal{C(a,b,r)}}} \mathcal{Q}_{A,B}(Q\mathrm{diag}(1/g)R^T) - \epsilon \cdot H((Q,R,g))\]

where :

• \(A\) is the (dim_a, dim_a) square pairwise cost matrix of the source domain.

• \(B\) is the (dim_b, dim_b) square pairwise cost matrix of the target domain.

• \(\mathcal{Q}_{A,B}\) is quadratic objective function of the Gromov Wasserstein plan.

• \(Q\) and R are the low-rank matrix decomposition of the Gromov-Wasserstein plan.

• \(g\) is the weight vector for the low-rank decomposition of the Gromov-Wasserstein plan.

• \(\mathbf{a}\) and \(\mathbf{b}\) are source and target weights (histograms, both sum to 1).

• \(r\) is the rank of the Gromov-Wasserstein plan.

• \(\mathcal{C(a,b,r)}\) are the low-rank couplings of the OT problem.

• \(H((Q,R,g))\) is the values of the three respective entropies evaluated for each term.

Parameters:

• X_s (array-like, shape (n_samples_a, dim_Xs)) – Samples in the source domain

• X_t (array-like, shape (n_samples_b, dim_Xt)) – Samples in the target domain

• a (array-like, shape (n_samples_a,), optional) – Samples weights in the source domain If let to its default value None, uniform distribution is taken.

• b (array-like, shape (n_samples_b,), optional) – Samples weights in the target domain If let to its default value None, uniform distribution is taken.

• reg (float, optional) – Regularization term >=0

• rank (int, optional. Default is None. (>0)) – Nonnegative rank of the OT plan. If None, min(ns, nt) is considered.

• alpha (float, optional. Default is 1e-10. (>0 and <1/r)) – Lower bound for the weight vector g.

• rescale_cost (bool, optional. Default is True) – Rescale the low rank factorization of the sqeuclidean cost matrix

• seed_init (int, optional. Default is 49. (>0)) – Random state for the ‘random’ initialization of low rank couplings

• gamma_init (str, optional. Default is "rescale".) – Initialization strategy for gamma. ‘rescale’, or ‘theory’ Gamma is a constant that scales the convergence criterion of the Mirror Descent optimization scheme used to compute the low-rank couplings (Q, R and g)

• numItermax (int, optional. Default is 1000.) – Max number of iterations for Low Rank GW

• stopThr (float, optional. Default is 1e-4.) – Stop threshold on error (>0) for Low Rank GW The error is the sum of Kullback Divergences computed for each low rank coupling (Q, R and g) and scaled using gamma.

• numItermax_dykstra (int, optional. Default is 10000.) – Max number of iterations for the Dykstra algorithm

• stopThr_dykstra (float, optional. Default is 1e-3.) – Stop threshold on error (>0) in Dykstra

• cost_factorized_Xs (tuple, optional. Default is None) – Tuple with two pre-computed low rank decompositions (A1, A2) of the source cost matrix. Both matrices should have a shape of (n_samples_a, dim_Xs + 2). If None, the low rank cost matrices will be computed as sqeuclidean cost matrices.

• cost_factorized_Xt (tuple, optional. Default is None) – Tuple with two pre-computed low rank decompositions (B1, B2) of the target cost matrix. Both matrices should have a shape of (n_samples_b, dim_Xt + 2). If None, the low rank cost matrices will be computed as sqeuclidean cost matrices.

• warn (bool, optional) – if True, raises a warning if the low rank GW algorithm doesn’t convergence.

• warn_dykstra (bool, optional) – if True, raises a warning if the Dykstra algorithm doesn’t convergence.

• log (bool, optional) – record log if True

Returns:

• Q (array-like, shape (n_samples_a, r)) – First low-rank matrix decomposition of the OT plan

• R (array-like, shape (n_samples_b, r)) – Second low-rank matrix decomposition of the OT plan

• g (array-like, shape (r, )) – Weight vector for the low-rank decomposition of the OT

• log (dict (lazy_plan, value and value_linear)) – log dictionary return only if log==True in parameters

References

ot.lowrank_sinkhorn(X_s, X_t, a=None, b=None, reg=0, rank=None, alpha=1e-10, rescale_cost=True, init='random', reg_init=0.1, seed_init=49, gamma_init='rescale', numItermax=2000, stopThr=1e-07, warn=True, log=False)[source]

Solve the entropic regularization optimal transport problem under low-nonnegative rank constraints on the couplings.

The function solves the following optimization problem:

\[\mathop{\inf_{(\mathbf{Q},\mathbf{R},\mathbf{g}) \in \mathcal{C}(\mathbf{a},\mathbf{b},r)}} \langle \mathbf{C}, \mathbf{Q}\mathrm{diag}(1/\mathbf{g})\mathbf{R}^\top \rangle - \mathrm{reg} \cdot H((\mathbf{Q}, \mathbf{R}, \mathbf{g}))\]

where :

• \(\mathbf{C}\) is the (dim_a, dim_b) metric cost matrix

• \(H((\mathbf{Q}, \mathbf{R}, \mathbf{g}))\) is the values of the three respective entropies evaluated for each term.

• \(\mathbf{Q}\) and \(\mathbf{R}\) are the low-rank matrix decomposition of the OT plan

• \(\mathbf{g}\) is the weight vector for the low-rank decomposition of the OT plan

• \(\mathbf{a}\) and \(\mathbf{b}\) are source and target weights (histograms, both sum to 1)

• \(r\) is the rank of the OT plan

• \(\mathcal{C}(\mathbf{a}, \mathbf{b}, r)\) are the low-rank couplings of the OT problem

Parameters:

• X_s (array-like, shape (n_samples_a, dim)) – samples in the source domain

• X_t (array-like, shape (n_samples_b, dim)) – samples in the target domain

• a (array-like, shape (n_samples_a,)) – samples weights in the source domain

• b (array-like, shape (n_samples_b,)) – samples weights in the target domain

• reg (float, optional) – Regularization term >0

• rank (int, optional. Default is None. (>0)) – Nonnegative rank of the OT plan. If None, min(ns, nt) is considered.

• alpha (int, optional. Default is 1e-10. (>0 and <1/r)) – Lower bound for the weight vector g.

• rescale_cost (bool, optional. Default is True) – Rescale the low rank factorization of the sqeuclidean cost matrix

• init (str, optional. Default is 'random'.) – Initialization strategy for the low rank couplings. ‘random’, ‘deterministic’ or ‘kmeans’

• reg_init (float, optional. Default is 1e-1. (>0)) – Regularization term for a ‘kmeans’ init. If None, 1 is considered.

• seed_init (int, optional. Default is 49. (>0)) – Random state for a ‘random’ or ‘kmeans’ init strategy.

• gamma_init (str, optional. Default is "rescale".) – Initialization strategy for gamma. ‘rescale’, or ‘theory’ Gamma is a constant that scales the convergence criterion of the Mirror Descent optimization scheme used to compute the low-rank couplings (Q, R and g)

• numItermax (int, optional. Default is 2000.) – Max number of iterations for the Dykstra algorithm

• stopThr (float, optional. Default is 1e-7.) – Stop threshold on error (>0) in Dykstra

• warn (bool, optional) – if True, raises a warning if the algorithm doesn’t convergence.

• log (bool, optional) – record log if True

Returns:

• Q (array-like, shape (n_samples_a, r)) – First low-rank matrix decomposition of the OT plan

• R (array-like, shape (n_samples_b, r)) – Second low-rank matrix decomposition of the OT plan

• g (array-like, shape (r, )) – Weight vector for the low-rank decomposition of the OT

• log (dict (lazy_plan, value and value_linear)) – log dictionary return only if log==True in parameters

References

ot.max_sliced_wasserstein_distance(X_s, X_t, a=None, b=None, n_projections=50, p=2, projections=None, seed=None, log=False, scaler=None)[source]

Computes a Monte-Carlo approximation of the max p-Sliced Wasserstein distance

\[\mathcal{Max-SWD}_p(\mu, \nu) = \underset{\theta \in \mathcal{U}(\mathbb{S}^{d-1})}{\max} [\mathcal{W}_p^p(\theta_\# \mu, \theta_\# \nu)]^{\frac{1}{p}}\]

where :

• \(\theta_\# \mu\) stands for the pushforwards of the projection \(\mathbb{R}^d \ni X \mapsto \langle \theta, X \rangle\)

Parameters:

• X_s (ndarray, shape (n_samples_a, dim)) – samples in the source domain

• X_t (ndarray, shape (n_samples_b, dim)) – samples in the target domain

• a (ndarray, shape (n_samples_a,), optional) – samples weights in the source domain

• b (ndarray, shape (n_samples_b,), optional) – samples weights in the target domain

• n_projections (int, optional) – Number of projections used for the Monte-Carlo approximation

• p (float, optional =) – Power p used for computing the sliced Wasserstein

• projections (shape (dim, n_projections), optional) – Projection matrix (n_projections and seed are not used in this case)

• seed (int or RandomState or None, optional) – Seed used for random number generator

• log (bool, optional) – if True, sliced_wasserstein_distance returns the projections used and their associated EMD.

• scaler (None, object with .transform(), or callable, optional) –

  Preprocessing applied to X_s and X_t before computing the distance. Useful for normalizing inputs when features have very different scales.

  • None : no preprocessing (default)

  • Object with .transform() method : e.g. an ot.utils.DataScaler fitted on a representative sample. This is the recommended way to get stable, consistent normalization across multiple calls (e.g. when using SWD as a loss in mini-batch training).

  • Callable : any function, lambda, or PyTorch transform applied directly as scaler(X_s) and scaler(X_t).

  See ot.utils.DataScaler for a backend-aware scaler that supports joint fitting on multiple distributions.

Returns:

• cost (float) – Sliced Wasserstein Cost

• log (dict, optional) – log dictionary return only if log==True in parameters

Examples

>>> import numpy as np
>>> n_samples_a = 20
>>> X = np.random.normal(0., 1., (n_samples_a, 5))
>>> sliced_wasserstein_distance(X, X, seed=0)
0.0

References

ot.merge_bijections(X, Y, perms, p=2)[source]

Merge several bijections between two point clouds to obtain a new one with low transport cost. The new bijection is guaranteed to have a transport cost no greater than the cost of any of the input bijections. Based on simple local/global swapping strategy, with a linear complexity in the number of points. Algorithm 3 from [84].

Parameters:

• X (array-like, shape (n_samples, dimension))

• Y (array-like, shape (n_samples, dimension))

• perms (array-like, shape (n_plans,n_samples)) – The bijections to merge, stored as permutations (e.g. a list of numbers)

• p (int, optional) – The power of the ground metric (default 2 for squared euclidean). If p is -1, the infinity norm is used.

Returns:

• cost (float) – The transport cost of the merged bijection.

• perm (array-like, shape (n_samples,)) – The merged bijection, stored as a permutation (e.g. a list of numbers) such that X[i] is assigned to Y[perm[i]].

ot.min_sliced_transport_plan(X_s, X_t, a=None, b=None, projections=None, metric='sqeuclidean', p=2, n_projections=None, seed=None, batch_size=None, dense=True, log=False)[source]

Computes the cost and permutation associated to the min-Pivot Sliced Discrepancy (introduced as min-SWGG in [85] and studied further in [86]). Given the supports X_s and X_t of two discrete uniform measures with ns and nt atoms in dimension d, the min-Pivot Sliced Discrepancy goes through n_projections different projections of the measures on random directions, and retains the couplings that yields the lowest cost between X_s and X_t (compared in \(\mathbb{R}^d\)). When ns=nt, it gives

\[\mathrm{min\text{-}PS}_p^p(X_s, X_t) \approx \min_{k \in [1, n_{\mathrm{proj}}]} \left( \frac{1}{n_s} \sum_{i=1}^{n_s} \|X_{s,i} - X_{t,\sigma_k(i)}\|_2^p \right),\]

where \(\sigma_k\) is a permutation such that ordering the projections on the axis projections[k, :] matches \(X_s[i, :]\) to \(X_t[\sigma_k(i), :]\).

Note

The computation ignores potential ambiguities in the projections: if two points from a same measure have the same projection on a direction, then multiple sorting permutations are possible. To avoid combinatorial explosion, only one permutation is retained: this strays from theory in pathological cases.

Warning

Tensorflow and jax only returns dense plans, as they do not support well sparse matrices.

Parameters:

• X_s (array-like, shape (ns, d)) – The first set of vectors.

• X_t (array-like, shape (nt, d)) – The second set of vectors.

• a (ndarray of float64, shape (ns,), optional) – Source histogram (default is uniform weight)

• b (ndarray of float64, shape (nt,), optional) – Target histogram (default is uniform weight)

• projections (shape (dim, n_projections), optional) – Projection matrix (n_projections and seed are not used in this case). Default is None

• metric (str, optional (default='sqeuclidean')) – Metric to be used. Only works with either of the strings ‘sqeuclidean’, ‘minkowski’, ‘cityblock’, or ‘euclidean’.

• p (float, optional (default=1.0)) – The p-norm to apply for if metric=’minkowski’

• n_projections (int, optional) – The number of projection directions. Required if projections is None.

• seed (int, optional) – The seed for the random number generator for sampling projections, in case projections is None. Default is None.

• batch_size (int, optional) – If specified, compute the distance in batches of size batch_size to avoid memory issues for large datasets. Default is None (no batching).

• dense (boolean, optional (default=True)) – If True, returns \(\gamma\) as a dense ndarray of shape (ns, nt). Otherwise returns a sparse representation using scipy’s coo_matrix format.

• log (bool, optional) – If True, returns additional logging information. Default is False.

Returns:

• plan (ndarray, shape (ns, nt) or coo_matrix if dense is False) – Optimal transportation matrix for the given parameters.

• cost (float) – The cost associated to the optimal permutation.

• log_dict (dict, optional) – A dictionary containing intermediate computations for logging purposes. Returned only if log is True.

References

Examples

>>> import ot
>>> import numpy as np
>>> x=np.array([[3.,3.], [1.,1.]])
>>> y=np.array([[2.,2.5], [3.,2.]])
>>> projections=np.array([[1, 0], [0, 1]])
>>> plan, cost = ot.min_sliced_transport_plan(x, y, projections=projections)
>>> plan
array([[0. , 0.5],
       [0.5, 0. ]])
>>> cost
2.125

ot.semidiscrete_wasserstein2_unif_circle(u_values, u_weights=None)[source]

Computes the closed-form for the 2-Wasserstein distance between samples and a uniform distribution on \(S^1\) Samples need to be in \(S^1\cong [0,1[\). If they are on \(\mathbb{R}\), takes the value modulo 1. If the values are on \(S^1\subset\mathbb{R}^2\), it is required to first find the coordinates using e.g. the atan2 function.

\[W_2^2(\mu_n, \nu) = \sum_{i=1}^n \alpha_i x_i^2 - \left(\sum_{i=1}^n \alpha_i x_i\right)^2 + \sum_{i=1}^n \alpha_i x_i \left(1-\alpha_i-2\sum_{k=1}^{i-1}\alpha_k\right) + \frac{1}{12}\]

where:

• \(\nu=\mathrm{Unif}(S^1)\) and \(\mu_n = \sum_{i=1}^n \alpha_i \delta_{x_i}\)

For values \(x=(x_1,x_2)\in S^1\), it is required to first get their coordinates with

\[u = \frac{\pi + \mathrm{atan2}(-x_2,-x_1)}{2\pi},\]

using e.g. ot.utils.get_coordinate_circle(x).

Parameters:

• u_values (ndarray, shape (n, ...)) – Samples

• u_weights (ndarray, shape (n, ...), optional) – samples weights in the source domain

Returns:

loss – Batched cost associated to the optimal transportation

Return type:

float/array-like, shape (…)

Examples

>>> x0 = np.array([[0], [0.2], [0.4]])
>>> semidiscrete_wasserstein2_unif_circle(x0)
array([0.02111111])

References

ot.sinkhorn(a, b, M, reg, method='sinkhorn', numItermax=1000, stopThr=1e-09, verbose=False, log=False, warn=True, warmstart=None, **kwargs)[source]

Solve the entropic regularization optimal transport problem and return the OT matrix

The function solves the following optimization problem:

\[ \begin{align}\begin{aligned}\gamma = \mathop{\arg \min}_\gamma \quad \langle \gamma, \mathbf{M} \rangle_F + \mathrm{reg}\cdot\Omega(\gamma)\\s.t. \ \gamma \mathbf{1} &= \mathbf{a}\\ \gamma^T \mathbf{1} &= \mathbf{b}\\ \gamma &\geq 0\end{aligned}\end{align} \]

where :

• \(\mathbf{M}\) is the (dim_a, dim_b) metric cost matrix

• \(\Omega\) is the entropic regularization term \(\Omega(\gamma)=\sum_{i,j} \gamma_{i,j}\log(\gamma_{i,j})\)

• \(\mathbf{a}\) and \(\mathbf{b}\) are source and target weights (histograms, both sum to 1)

Note

This function is backend-compatible and will work on arrays from all compatible backends.

The algorithm used for solving the problem is the Sinkhorn-Knopp matrix scaling algorithm as proposed in [2]

Choosing a Sinkhorn solver

By default and when using a regularization parameter that is not too small the default sinkhorn solver should be enough. If you need to use a small regularization to get sharper OT matrices, you should use the ot.bregman.sinkhorn_stabilized() solver that will avoid numerical errors. This last solver can be very slow in practice and might not even converge to a reasonable OT matrix in a finite time. This is why ot.bregman.sinkhorn_epsilon_scaling() that relies on iterating the value of the regularization (and using warm start) sometimes leads to better solutions. Note that the greedy version of the sinkhorn ot.bregman.greenkhorn() can also lead to a speedup and the screening version of the sinkhorn ot.bregman.screenkhorn() aim at providing a fast approximation of the Sinkhorn problem. For use of GPU and gradient computation with small number of iterations we strongly recommend the ot.bregman.sinkhorn_log() solver that will no need to check for numerical problems.

Parameters:

• a (array-like, shape (dim_a,)) – samples weights in the source domain

• b (array-like, shape (dim_b,) or ndarray, shape (dim_b, n_hists)) – samples in the target domain, compute sinkhorn with multiple targets and fixed \(\mathbf{M}\) if \(\mathbf{b}\) is a matrix (return OT loss + dual variables in log)

• M (array-like, shape (dim_a, dim_b)) – loss matrix

• reg (float) – Regularization term >0

• method (str) – method used for the solver either ‘sinkhorn’,’sinkhorn_log’, ‘greenkhorn’, ‘sinkhorn_stabilized’ or ‘sinkhorn_epsilon_scaling’, see those function for specific parameters

• numItermax (int, optional) – Max number of iterations

• stopThr (float, optional) – Stop threshold on error (>0)

• verbose (bool, optional) – Print information along iterations

• log (bool, optional) – record log if True

• warn (bool, optional) – if True, raises a warning if the algorithm doesn’t convergence.

• warmstart (tuple of arrays, shape (dim_a, dim_b), optional) – Initialization of dual potentials. If provided, the dual potentials should be given (that is the logarithm of the u,v sinkhorn scaling vectors)

Returns:

• gamma (array-like, shape (dim_a, dim_b)) – Optimal transportation matrix for the given parameters

• log (dict) – log dictionary return only if log==True in parameters

Examples

>>> import ot
>>> a=[.5, .5]
>>> b=[.5, .5]
>>> M=[[0., 1.], [1., 0.]]
>>> ot.sinkhorn(a, b, M, 1)
array([[0.36552929, 0.13447071],
       [0.13447071, 0.36552929]])

References

See also

ot.lp.emd

Unregularized OT

ot.optim.cg

General regularized OT

ot.bregman.sinkhorn_knopp

Classic Sinkhorn [2]

ot.bregman.sinkhorn_stabilized

Stabilized sinkhorn [9] [10]

ot.bregman.sinkhorn_epsilon_scaling

Sinkhorn with epsilon scaling [9] [10]

ot.sinkhorn2(a, b, M, reg, method='sinkhorn', numItermax=1000, stopThr=1e-09, verbose=False, log=False, warn=False, warmstart=None, **kwargs)[source]

Solve the entropic regularization optimal transport problem and return the loss

The function solves the following optimization problem:

\[ \begin{align}\begin{aligned}W = \min_\gamma \quad \langle \gamma, \mathbf{M} \rangle_F + \mathrm{reg}\cdot\Omega(\gamma)\\s.t. \ \gamma \mathbf{1} &= \mathbf{a}\\ \gamma^T \mathbf{1} &= \mathbf{b}\\ \gamma &\geq 0\end{aligned}\end{align} \]

where :

• \(\mathbf{M}\) is the (dim_a, dim_b) metric cost matrix

• \(\Omega\) is the entropic regularization term \(\Omega(\gamma)=\sum_{i,j} \gamma_{i,j}\log(\gamma_{i,j})\)

• \(\mathbf{a}\) and \(\mathbf{b}\) are source and target weights (histograms, both sum to 1)

and returns \(\langle \gamma^*, \mathbf{M} \rangle_F\) (without the entropic contribution).

Note

This function is backend-compatible and will work on arrays from all compatible backends.

The algorithm used for solving the problem is the Sinkhorn-Knopp matrix scaling algorithm as proposed in [2]

Choosing a Sinkhorn solver

By default and when using a regularization parameter that is not too small the default sinkhorn solver should be enough. If you need to use a small regularization to get sharper OT matrices, you should use the ot.bregman.sinkhorn_log() solver that will avoid numerical errors. This last solver can be very slow in practice and might not even converge to a reasonable OT matrix in a finite time. This is why ot.bregman.sinkhorn_epsilon_scaling() that relies on iterating the value of the regularization (and using warm start) sometimes leads to better solutions. Note that the greedy version of the sinkhorn ot.bregman.greenkhorn() can also lead to a speedup and the screening version of the sinkhorn ot.bregman.screenkhorn() aim a providing a fast approximation of the Sinkhorn problem. For use of GPU and gradient computation with small number of iterations we strongly recommend the ot.bregman.sinkhorn_log() solver that will no need to check for numerical problems.

Parameters:

• a (array-like, shape (dim_a,)) – samples weights in the source domain

• b (array-like, shape (dim_b,) or ndarray, shape (dim_b, n_hists)) – samples in the target domain, compute sinkhorn with multiple targets and fixed \(\mathbf{M}\) if \(\mathbf{b}\) is a matrix (return OT loss + dual variables in log)

• M (array-like, shape (dim_a, dim_b)) – loss matrix

• reg (float) – Regularization term >0

• method (str) – method used for the solver either ‘sinkhorn’,’sinkhorn_log’, ‘sinkhorn_stabilized’, see those function for specific parameters

• numItermax (int, optional) – Max number of iterations

• stopThr (float, optional) – Stop threshold on error (>0)

• verbose (bool, optional) – Print information along iterations

• log (bool, optional) – record log if True

• warn (bool, optional) – if True, raises a warning if the algorithm doesn’t convergence.

• warmstart (tuple of arrays, shape (dim_a, dim_b), optional) – Initialization of dual potentials. If provided, the dual potentials should be given (that is the logarithm of the u,v sinkhorn scaling vectors)

Returns:

• W ((n_hists) float/array-like) – Optimal transportation loss for the given parameters

• log (dict) – log dictionary return only if log==True in parameters

Examples

>>> import ot
>>> a=[.5, .5]
>>> b=[.5, .5]
>>> M=[[0., 1.], [1., 0.]]
>>> ot.sinkhorn2(a, b, M, 1)
0.26894142136999516

References

See also

ot.lp.emd

Unregularized OT

ot.optim.cg

General regularized OT

ot.bregman.sinkhorn_knopp

Classic Sinkhorn [2]

ot.bregman.greenkhorn

Greenkhorn [21]

ot.bregman.sinkhorn_stabilized

Stabilized sinkhorn [9] [10]

ot.sinkhorn_lpl1_mm(a, labels_a, b, M, reg, eta=0.1, numItermax=10, numInnerItermax=200, stopInnerThr=1e-09, verbose=False, log=False)[source]

Solve the entropic regularization optimal transport problem with non-convex group lasso regularization

The function solves the following optimization problem:

\[ \begin{align}\begin{aligned}\gamma = \mathop{\arg \min}_\gamma \quad \langle \gamma, \mathbf{M} \rangle_F + \mathrm{reg} \cdot \Omega_e(\gamma) + \eta \ \Omega_g(\gamma)\\s.t. \ \gamma \mathbf{1} = \mathbf{a}\\ \gamma^T \mathbf{1} = \mathbf{b}\\ \gamma \geq 0\end{aligned}\end{align} \]

where :

• \(\mathbf{M}\) is the (ns, nt) metric cost matrix

• \(\Omega_e\) is the entropic regularization term \(\Omega_e (\gamma)=\sum_{i,j} \gamma_{i,j}\log(\gamma_{i,j})\)

• \(\Omega_g\) is the group lasso regularization term \(\Omega_g(\gamma)=\sum_{i,c} \|\gamma_{i,\mathcal{I}_c}\|^{1/2}_1\) where \(\mathcal{I}_c\) are the index of samples from class c in the source domain.

• \(\mathbf{a}\) and \(\mathbf{b}\) are source and target weights (sum to 1)

The algorithm used for solving the problem is the generalized conditional gradient as proposed in [5, 7].

Parameters:

• a (array-like (ns,)) – samples weights in the source domain

• labels_a (array-like (ns,)) – labels of samples in the source domain

• b (array-like (nt,)) – samples weights in the target domain

• M (array-like (ns,nt)) – loss matrix

• reg (float) – Regularization term for entropic regularization >0

• eta (float, optional) – Regularization term for group lasso regularization >0

• numItermax (int, optional) – Max number of iterations

• numInnerItermax (int, optional) – Max number of iterations (inner sinkhorn solver)

• stopInnerThr (float, optional) – Stop threshold on error (inner sinkhorn solver) (>0)

• verbose (bool, optional) – Print information along iterations

• log (bool, optional) – record log if True

Returns:

• gamma ((ns, nt) array-like) – Optimal transportation matrix for the given parameters

• log (dict) – log dictionary return only if log==True in parameters

References

See also

ot.lp.emd

Unregularized OT

ot.bregman.sinkhorn

Entropic regularized OT

ot.optim.cg

General regularized OT