OptimalTransport.jl Documentation

Exact optimal transport (Kantorovich) problem

OptimalTransport.emdFunction
emd(mu, nu, C)

Compute transport map for Monge-Kantorovich problem with source and target marginals mu and nu and a cost matrix C of dimensions (length(mu), length(nu)).

Return optimal transport coupling γ of the same dimensions as C which solves

\[\inf_{\gamma \in \Pi(\mu, \nu)} \langle \gamma, C \rangle\]

This function is a wrapper of the function emd in the Python Optimal Transport package.

OptimalTransport.emd2Function
emd2(mu, nu, C)

Compute exact transport cost for Monge-Kantorovich problem with source and target marginals mu and nu and a cost matrix C of dimensions (length(mu), length(nu)).

Returns optimal transport cost (a scalar), i.e. the optimal value

\[\inf_{\gamma \in \Pi(\mu, \nu)} \langle \gamma, C \rangle\]

This function is a wrapper of the function emd2 in the Python Optimal Transport package.

Entropically regularised optimal transport

OptimalTransport.sinkhornFunction
sinkhorn(mu, nu, C, eps; tol=1e-9, check_marginal_step=10, maxiter=1000)

Compute entropically regularised transport map of histograms mu and nu with cost matrix C and entropic regularization parameter eps.

Return optimal transport coupling γ of the same dimensions as C which solves

\[\inf_{\gamma \in \Pi(\mu, \nu)} \langle \gamma, C \rangle - \epsilon H(\gamma)\]

where $H$ is the entropic regulariser, $H(\gamma) = -\sum_{i, j} \gamma_{ij} \log(\gamma_{ij})$.

OptimalTransport.sinkhorn2Function
sinkhorn2(mu, nu, C, eps; tol=1e-9, check_marginal_step=10, maxiter=1000)

Compute entropically regularised transport cost of histograms mu and nu with cost matrix C and entropic regularization parameter eps.

Return optimal value of

\[\inf_{\gamma \in \Pi(\mu, \nu)} \langle \gamma, C \rangle - \epsilon H(\gamma)\]

where $H$ is the entropic regulariser, $H(\gamma) = -\sum_{i, j} \gamma_{ij} \log(\gamma_{ij})$.

OptimalTransport.sinkhorn_stabilized_epsscalingFunction
sinkhorn_stabilized_epsscaling(mu, nu, C, eps; absorb_tol = 1e3, max_iter = 1000, tol = 1e-9, lambda = 0.5, k = 5, verbose = false)

Compute optimal transport map of histograms mu and nu with cost matrix C and entropic regularisation parameter eps. Uses stabilized Sinkhorn algorithm with epsilon-scaling (Schmitzer et al., 2019).

k epsilon-scaling steps are used with scaling factor lambda, i.e. sequentially solve Sinkhorn with regularisation parameters [lambda^(1-k), ..., lambda^(-1), 1]*eps.

OptimalTransport.sinkhorn_stabilizedFunction
sinkhorn_stabilized(mu, nu, C, eps; absorb_tol = 1e3, max_iter = 1000, tol = 1e-9, alpha = nothing, beta = nothing, return_duals = false, verbose = false)

Compute optimal transport map of histograms mu and nu with cost matrix C and entropic regularisation parameter eps. Uses stabilized Sinkhorn algorithm (Schmitzer et al., 2019).

OptimalTransport.sinkhorn_barycenterFunction
sinkhorn_barycenter(mu_all, C_all, eps, lambda_all; tol = 1e-9, check_marginal_step = 10, max_iter = 1000)

Compute the entropically regularised (i.e. Sinkhorn) barycenter for a collection of `N` histograms `mu_all` with respective cost matrices `C_all`, relative weights `lambda_all`, and entropic regularisation parameter `eps`.

mu_all is taken to contain N histograms mu_all[i, :] for i = 1, ..., N C_all is taken to be a list of N cost matrices corresponding to the mu_all[i, :] eps is the scalar regularisation parameter lambda_all are positive weights.

Returns the entropically regularised barycenter of the mu_all, i.e. the distribution that minimises

\[\min_{\gamma \in \Sigma} \sum_{i = 1}^N \lambda_i \mathrm{entropicOT}^{\epsilon}_{C_i}(\mu, \mu_i)\]

where $\mathrm{entropicOT}^{\epsilon}_{C}$ denotes the entropic optimal transport cost with cost $C$ and entropic regularisation level $\epsilon$.

Wrapper functions to POT library:

OptimalTransport.pot_sinkhornFunction
pot_sinkhorn(mu, nu, C, eps; tol=1e-9, max_iter = 1000, method = "sinkhorn", verbose = false)

Compute optimal transport map of histograms mu and nu with cost matrix C and entropic regularization parameter eps.

Method can be a choice of "sinkhorn", "greenkhorn", "sinkhorn_stabilized", or "sinkhorn_epsilon_scaling" (Flamary et al., 2017)

This function is a wrapper of the function sinkhorn in the Python Optimal Transport package.

OptimalTransport.pot_sinkhorn2Function
pot_sinkhorn2(mu, nu, C, eps; tol=1e-9, max_iter = 1000, method = "sinkhorn", verbose = false)

Compute optimal transport cost of histograms mu and nu with cost matrix C and entropic regularization parameter eps.

Method can be a choice of "sinkhorn", "greenkhorn", "sinkhorn_stabilized", or "sinkhorn_epsilon_scaling" (Flamary et al., 2017)

This function is a wrapper of the function sinkhorn2 in the Python Optimal Transport package.

Unbalanced optimal transport

OptimalTransport.sinkhorn_unbalancedFunction
sinkhorn_unbalanced(mu, nu, C, lambda1, lambda2, eps; tol = 1e-9, max_iter = 1000, verbose = false, proxdiv_F1 = nothing, proxdiv_F2 = nothing)

Computes the optimal transport map of histograms mu and nu with cost matrix C and entropic regularization parameter eps, using the unbalanced Sinkhorn algorithm [Chizat 2016] with KL-divergence terms for soft marginal constraints, with weights (lambda1, lambda2) for the marginals mu, nu respectively.

In general, the user can specify the soft marginal constraints $(F_1(\cdot | \mu), F_2(\cdot | \nu))$ to the problem

\[min_\gamma \epsilon \mathrm{KL}(\gamma | \exp(-C/\epsilon)) + F_1(\gamma_1 | \mu) + F_2(\gamma_2 | \nu)\]

via \mathrm{proxdiv}_{F_1}(s, p) and \mathrm{proxdiv}_{F_2}(s, p) (see Chizat et al., 2016 for details on this). If specified, the algorithm will use the user-specified F1, F2 rather than the default (a KL-divergence).

OptimalTransport.sinkhorn_unbalanced2Function
sinkhorn_unbalanced2(mu, nu, C, lambda1, lambda2, eps; tol = 1e-9, max_iter = 1000, verbose = false, proxdiv_F1 = nothing, proxdiv_F2 = nothing)

Computes the optimal transport cost of histograms mu and nu with cost matrix C and entropic regularization parameter eps, using the unbalanced Sinkhorn algorithm [Chizat 2016] with KL-divergence terms for soft marginal constraints, with weights (lambda1, lambda2) for the marginals mu, nu respectively.

See documentation for sinkhorn_unbalanced for additional details.

OptimalTransport.pot_sinkhorn_unbalancedFunction
pot_sinkhorn_unbalanced(mu, nu, C, eps, lambda; tol = 1e-9, max_iter = 1000, method = "sinkhorn", verbose = false)

Compute optimal transport map of histograms mu and nu with cost matrix C, using entropic regularisation parameter eps and marginal weighting functions lambda.

This function is a wrapper of the function sinkhorn_unbalanced in the Python Optimal Transport package.

OptimalTransport.pot_sinkhorn_unbalanced2Function
pot_sinkhorn_unbalanced2(mu, nu, C, eps, lambda; tol = 1e-9, max_iter = 1000, method = "sinkhorn", verbose = false)

Compute optimal transport cost of histograms mu and nu with cost matrix C, using entropic regularisation parameter eps and marginal weighting functions lambda.

This function is a wrapper of the function sinkhorn_unbalanced2 in the Python Optimal Transport package.

Quadratically regularised optimal transport

OptimalTransport.quadregFunction
quadreg(mu, nu, C, ϵ; θ = 0.1, tol = 1e-5,maxiter = 50,κ = 0.5,δ = 1e-5)

Computes the optimal transport map of histograms mu and nu with cost matrix C and quadratic regularization parameter ϵ, using the semismooth Newton algorithm [Lorenz 2016].

This implementation makes use of IterativeSolvers.jl and SparseArrays.jl.

Parameters:

θ: starting Armijo parameter.

tol: tolerance of marginal error.

maxiter: maximum interation number.

κ: control parameter of Armijo.

δ: small constant for the numerical stability of conjugate gradient iterative solver.

Tips: If the algorithm does not converge, try some different values of θ.

Reference: Lorenz, D.A., Manns, P. and Meyer, C., 2019. Quadratically regularized optimal transport. arXiv preprint arXiv:1903.01112v4.