OptimalTransport.jl Documentation

emd(μ, ν, C, optimizer)

Compute the optimal transport plan γ for the Monge-Kantorovich problem with source histogram μ, target histogram ν, and cost matrix C of size (length(μ), length(ν)) which solves

\[\inf_{γ ∈ Π(μ, ν)} \langle γ, C \rangle.\]

The corresponding linear programming problem is solved with the user-provided optimizer. Possible choices are Tulip.Optimizer() and Clp.Optimizer() in the Tulip and Clp packages, respectively.

emd2(μ, ν, C, optimizer; plan=nothing)

Compute the optimal transport cost (a scalar) for the Monge-Kantorovich problem with source histogram μ, target histogram ν, and cost matrix C of size (length(μ), length(ν)) which is given by

\[\inf_{γ ∈ Π(μ, ν)} \langle γ, C \rangle.\]

The corresponding linear programming problem is solved with the user-provided optimizer. Possible choices are Tulip.Optimizer() and Clp.Optimizer() in the Tulip and Clp packages, respectively.

A pre-computed optimal transport plan may be provided.

ot_plan(c, μ, ν; kwargs...)

Compute the optimal transport plan for the Monge-Kantorovich problem with source and target marginals μ and ν and cost c.

The optimal transport plan solves

\[\inf_{\gamma \in \Pi(\mu, \nu)} \int c(x, y) \, \mathrm{d}\gamma(x, y)\]

where $\Pi(\mu, \nu)$ denotes the couplings of $\mu$ and $\nu$.

See also: ot_cost

ot_plan(c, μ::ContinuousUnivariateDistribution, ν::UnivariateDistribution)

Compute the optimal transport plan for the Monge-Kantorovich problem with univariate distributions μ and ν as source and target marginals and cost function c of the form $c(x, y) = h(|x - y|)$ where $h$ is a convex function.

In this setting, the optimal transport plan is the Monge map

\[T = F_\nu^{-1} \circ F_\mu\]

where $F_\mu$ is the cumulative distribution function of μ and $F_\nu^{-1}$ is the quantile function of ν.

See also: ot_cost, emd

ot_plan(c, μ::DiscreteNonParametric, ν::DiscreteNonParametric)

Compute the optimal transport cost for the Monge-Kantorovich problem with univariate discrete distributions μ and ν as source and target marginals and cost function c of the form $c(x, y) = h(|x - y|)$ where $h$ is a convex function.

In this setting, the optimal transport plan can be computed analytically. It is returned as a sparse matrix.

See also: ot_cost, emd

ot_plan(::SqEuclidean, μ::Normal, ν::Normal)

Compute the optimal transport plan for the Monge-Kantorovich problem with normal distributions μ and ν as source and target marginals and cost function $c(x, y) = \|x - y\|_2^2$.

See also: ot_cost, emd

ot_plan(::SqEuclidean, μ::MvNormal, ν::MvNormal)

Compute the optimal transport plan for the Monge-Kantorovich problem with multivariate normal distributions μ and ν as source and target marginals and cost function $c(x, y) = \|x - y\|_2^2$.

In this setting, for $\mu = \mathcal{N}(m_\mu, \Sigma_\mu)$ and $\nu = \mathcal{N}(m_\nu, \Sigma_\nu)$, the optimal transport plan is the Monge map

\[T \colon x \mapsto m_\nu + \Sigma_\mu^{-1/2} {\big(\Sigma_\mu^{1/2} \Sigma_\nu \Sigma_\mu^{1/2}\big)}^{1/2}\Sigma_\mu^{-1/2} (x - m_\mu).\]

See also: ot_cost, emd

ot_cost(c, μ, ν; kwargs...)

Compute the optimal transport cost for the Monge-Kantorovich problem with source and target marginals μ and ν and cost c.

The optimal transport cost is the scalar value

\[\inf_{\gamma \in \Pi(\mu, \nu)} \int c(x, y) \, \mathrm{d}\gamma(x, y)\]

where $\Pi(\mu, \nu)$ denotes the couplings of $\mu$ and $\nu$.

See also: ot_plan

ot_cost(
    c, μ::ContinuousUnivariateDistribution, ν::UnivariateDistribution; plan=nothing
)

Compute the optimal transport cost for the Monge-Kantorovich problem with univariate distributions μ and ν as source and target marginals and cost function c of the form $c(x, y) = h(|x - y|)$ where $h$ is a convex function.

In this setting, the optimal transport cost can be computed as

\[\int_0^1 c(F_\mu^{-1}(x), F_\nu^{-1}(x)) \mathrm{d}x\]

where $F_\mu^{-1}$ and $F_\nu^{-1}$ are the quantile functions of μ and ν, respectively.

A pre-computed optimal transport plan may be provided.

See also: ot_plan, emd2

ot_cost(
    c, μ::DiscreteNonParametric, ν::DiscreteNonParametric; plan=nothing
)

Compute the optimal transport cost for the Monge-Kantorovich problem with discrete univariate distributions μ and ν as source and target marginals and cost function c of the form $c(x, y) = h(|x - y|)$ where $h$ is a convex function.

In this setting, the optimal transport cost can be computed analytically.

A pre-computed optimal transport plan may be provided.

See also: ot_plan, emd2

ot_cost(::SqEuclidean, μ::Normal, ν::Normal)

Compute the squared 2-Wasserstein distance between univariate normal distributions μ and ν as source and target marginals.

See also: ot_plan, emd2

ot_cost(::SqEuclidean, μ::MvNormal, ν::MvNormal)

Compute the squared 2-Wasserstein distance between normal distributions μ and ν as source and target marginals.

In this setting, the optimal transport cost can be computed as

\[W_2^2(\mu, \nu) = \|m_\mu - m_\nu \|^2 + \mathcal{B}(\Sigma_\mu, \Sigma_\nu)^2,\]

where $\mu = \mathcal{N}(m_\mu, \Sigma_\mu)$, $\nu = \mathcal{N}(m_\nu, \Sigma_\nu)$, and $\mathcal{B}$ is the Bures metric.

See also: ot_plan, emd2

wasserstein(μ, ν; metric=Euclidean(), p=Val(1), kwargs...)

Compute the p-Wasserstein distance with respect to the metric between measures μ and ν.

Order p can be provided as a scalar of type Real or as a parameter of a value type Val(p). For certain combinations of metric and p, such as metric=Euclidean() and p=Val(2), the computations are more efficient if p is specified as a value type. The remaining keyword arguments are forwarded to ot_cost.

See also: squared2wasserstein, ot_cost

squared2wasserstein(μ, ν; metric=Euclidean(), kwargs...)

Compute the squared 2-Wasserstein distance with respect to the metric between measures μ and ν.

The remaining keyword arguments are forwarded to ot_cost.

See also: wasserstein, ot_cost

discretemeasure(
    support::AbstractVector,
    probs::AbstractVector{<:Real}=FillArrays.Fill(inv(length(support)), length(support)),
)

Construct a finite discrete probability measure with support and corresponding probabilities. If the probability vector argument is not passed, then equal probability is assigned to each entry in the support.

Examples

using KernelFunctions
# rows correspond to samples
μ = discretemeasure(RowVecs(rand(7,3)), normalize!(rand(10),1))

# columns correspond to samples, each with equal probability
ν = discretemeasure(ColVecs(rand(3,12)))

sinkhorn(
    μ, ν, C, ε, alg=SinkhornGibbs();
    atol=0, rtol=atol > 0 ? 0 : √eps, check_convergence=10, maxiter=1_000,
)

Compute the optimal transport plan for the entropically regularized optimal transport problem with source and target marginals μ and ν, cost matrix C of size (length(μ), length(ν)), and entropic regularization parameter ε.

The optimal transport plan γ is of the same size as C and solves

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

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

Every check_convergence steps it is assessed if the algorithm is converged by checking if the iterate of the transport plan G satisfies

isapprox(sum(G; dims=2), μ; atol=atol, rtol=rtol, norm=x -> norm(x, 1))

The default rtol depends on the types of μ, ν, and C. After maxiter iterations, the computation is stopped.

Batch computations for multiple histograms with a common cost matrix C can be performed by passing μ or ν as matrices whose columns correspond to histograms. It is required that the number of source and target marginals is equal or that a single source or single target marginal is provided (either as matrix or as vector). The optimal transport plans are returned as three-dimensional array where γ[:, :, i] is the optimal transport plan for the ith pair of source and target marginals.

See also: sinkhorn2

sinkhorn2(
    μ, ν, C, ε, alg=SinkhornGibbs(); regularization=false, plan=nothing, kwargs...
)

Solve the entropically regularized optimal transport problem with source and target marginals μ and ν, cost matrix C of size (length(μ), length(ν)), and entropic regularization parameter ε, and return the optimal cost.

A pre-computed optimal transport plan may be provided. The other keyword arguments supported here are the same as those in the sinkhorn function.

See also: sinkhorn

sinkhorn_divergence(
    μ::AbstractVecOrMat,
    ν::AbstractVecOrMat,
    C,
    ε,
    alg::SinkhornDivergence=SinkhornDivergence(
        SinkhornGibbs(), SymmetricSinkhornGibbs(), SymmetricSinkhornGibbs()
    );
    kwargs...,
)

Compute the Sinkhorn Divergence between finite discrete measures μ and ν with respect to a common cost matrix C, entropic regularization parameter ε and algorithm alg.

In the default case where regularization = false, the Sinkhorn Divergence is that of ^[GPC18] and is computed as

\[\operatorname{S}_{ε}(μ,ν) := \operatorname{W}_{ε}(μ,ν) - \frac{1}{2}(\operatorname{W}_{ε}(μ,μ) + \operatorname{W}_{ε}(ν,ν)),\]

and $\operatorname{W}_{ε}$ is defined as

\[\operatorname{W}_{ε}(μ, ν) = \langle C, γ^\star \rangle,\]

where $γ^\star$ is the entropy-regularised transport plan between μ and ν. For regularization = true, the Sinkhorn Divergence is that of ^[FeydyP19] and is computed as above where $\operatorname{W}_{ε}$ is replaced by $\operatorname{OT}_{ε}$, the entropy-regularised optimal transport cost with regulariser penalty.

The default algorithm for computing the term $\operatorname{W}_{ε}(μ, ν)$ is the SinkhornGibbs algorithm. For the terms $\operatorname{W}_{ε}(μ, μ)$ and $\operatorname{W}_{ε}(ν, ν)$, the symmetric fixed point iteration of ^[FeydyP19] is used. Alternatively, a pre-computed optimal transport plan between μ and ν may be provided.

See also: sinkhorn2

sinkhorn_divergence(
    μ,
    ν,
    Cμν,
    Cμ,
    Cν,
    ε,
    alg::SinkhornDivergence=SinkhornDivergence(
        SinkhornGibbs(), SymmetricSinkhornGibbs(), SymmetricSinkhornGibbs()
    );
    kwargs...,
)

Compute the Sinkhorn Divergence between finite discrete measures μ and ν with respect to the precomputed cost matrices Cμν, Cμμ and Cνν, entropic regularization parameter ε and algorithm alg.

A pre-computed optimal transport plan between μ and ν may be provided.

See also: sinkhorn2, sinkhorn_divergence

sinkhorn_barycenter(μ, C, ε, w, alg = SinkhornGibbs(); kwargs...)

Compute the Sinkhorn barycenter for a collection of N histograms contained in the columns of μ, for a cost matrix C of size (size(μ, 1), size(μ, 1)), relative weights w of size N, and entropic regularisation parameter ε. Returns the entropically regularised barycenter of the μ, i.e. the histogram ρ of length size(μ, 1) that solves

\[\min_{\rho \in \Sigma} \sum_{i = 1}^N w_i \operatorname{OT}_{\varepsilon}(\mu_i, \rho)\]

where $\operatorname{OT}_{ε}(\mu, \nu) = \inf_{\gamma \Pi(\mu, \nu)} \langle \gamma, C \rangle + \varepsilon \Omega(\gamma)$ is the entropic optimal transport loss with cost $C$ and regularisation $\epsilon$.

SinkhornGibbs()

Vanilla Sinkhorn algorithm.

SinkhornStabilized(; absorb_tol::Real=1_000)

Construct a log-domain stabilized Sinkhorn algorithm with absorption tolerance absorb_tol for solving an entropically regularized optimal transport problem.

References

Schmitzer, B. (2019). Stabilized Sparse Scaling Algorithms for Entropy Regularized Transport Problems. SIAM Journal on Scientific Computing, 41(3), A1443–A1481.

SinkhornEpsilonScaling(algorithm::Sinkhorn; factor=1//2, steps=5)

Construct an ε-scaling Sinkhorn algorithm for solving an entropically regularized optimal transport problem.

The function uses the specified Sinkhorn algorithm with steps ε-scaling steps with scaling factor factor. It sequentially solves the entropically regularized optimal transport with regularization parameters

\[\varepsilon_i := \varepsilon \lambda^{i-k}, \qquad (i = 1,\ldots,k),\]

where $\lambda$ is the scaling factor and $k$ the number of scaling steps.

sinkhorn_stabilized(μ, ν, C, ε; absorb_tol=1_000, kwargs...)

This method is deprecated, please use

sinkhorn(
    μ, ν, C, ε, SinkhornStabilized(; absorb_tol=absorb_tol); kwargs...
)

instead.

See also: sinkhorn, SinkhornStabilized

sinkhorn_stabilized_epsscaling(
    μ, ν, C, ε;
    scaling_factor=1//2, scaling_steps=5, absorb_tol=1_000, kwargs...
)

This method is deprecated, please use

sinkhorn(
    μ,
    ν,
    C,
    ε,
    SinkhornEpsilonScaling(
        SinkhornStabilized(; absorb_tol=absorb_tol);
        factor=scaling_factor,
        steps=scaling_steps,
    );
    kwargs...,
)

instead.

See also: sinkhorn, SinkhornEpsilonScaling

sinkhorn_unbalanced(μ, ν, C, λ1::Real, λ2::Real, ε; kwargs...)

Compute the optimal transport plan for the unbalanced entropically regularized optimal transport problem with source and target marginals μ and ν, cost matrix C of size (length(μ), length(ν)), entropic regularization parameter ε, and marginal relaxation terms λ1 and λ2.

The optimal transport plan γ is of the same size as C and solves

\[\inf_{\gamma} \langle \gamma, C \rangle + \varepsilon \Omega(\gamma) + \lambda_1 \operatorname{KL}(\gamma 1 | \mu) + \lambda_2 \operatorname{KL}(\gamma^{\mathsf{T}} 1 | \nu),\]

where $\Omega(\gamma) = \sum_{i,j} \gamma_{i,j} \log \gamma_{i,j}$ is the entropic regularization term and $\operatorname{KL}$ is the Kullback-Leibler divergence.

The keyword arguments supported here are the same as those in the sinkhorn_unbalanced for unbalanced optimal transport problems with general soft marginal constraints.

sinkhorn_unbalanced(
    μ, ν, C, proxdivF1!, proxdivF2!, ε;
    atol=0, rtol=atol > 0 ? 0 : √eps, check_convergence=10, maxiter=1_000,
)

Compute the optimal transport plan for the unbalanced entropically regularized optimal transport problem with source and target marginals μ and ν, cost matrix C of size (length(μ), length(ν)), entropic regularization parameter ε, and soft marginal constraints $F_1$ and $F_2$ with "proxdiv" functions proxdivF! and proxdivG!.

The optimal transport plan γ is of the same size as C and solves

\[\inf_{\gamma} \langle \gamma, C \rangle + \varepsilon \Omega(\gamma) + F_1(\gamma 1, \mu) + F_2(\gamma^{\mathsf{T}} 1, \nu),\]

where $\Omega(\gamma) = \sum_{i,j} \gamma_{i,j} \log \gamma_{i,j}$ is the entropic regularization term and $F_1(\cdot, \mu)$ and $F_2(\cdot, \nu)$ are soft marginal constraints for the source and target marginals.

The functions proxdivF1!(s, p, ε) and proxdivF2!(s, p, ε) evaluate the "proxdiv" functions of $F_1(\cdot, p)$ and $F_2(\cdot, p)$ at $s$ for the entropic regularization parameter $\varepsilon$. They have to be mutating and overwrite the first argument s with the result of their computations.

Mathematically, the "proxdiv" functions are defined as

\[\operatorname{proxdiv}_{F_i}(s, p, \varepsilon) = \operatorname{prox}^{\operatorname{KL}}_{F_i(\cdot, p)/\varepsilon}(s) \oslash s\]

where $\oslash$ denotes element-wise division and $\operatorname{prox}_{F_i(\cdot, p)/\varepsilon}^{\operatorname{KL}}$ is the proximal operator of $F_i(\cdot, p)/\varepsilon$ for the Kullback-Leibler ($\operatorname{KL}$) divergence. It is defined as

\[\operatorname{prox}_{F}^{\operatorname{KL}}(x) = \operatorname{argmin}_{y} F(y) + \operatorname{KL}(y|x)\]

and can be computed in closed-form for specific choices of $F$. For instance, if $F(\cdot, p) = \lambda \operatorname{KL}(\cdot | p)$ ($\lambda > 0$), then

\[\operatorname{prox}_{F(\cdot, p)/\varepsilon}^{\operatorname{KL}}(x) = x^{\frac{\varepsilon}{\varepsilon + \lambda}} p^{\frac{\lambda}{\varepsilon + \lambda}},\]

where all operators are acting pointwise.^[CPSV18]

Every check_convergence steps it is assessed if the algorithm is converged by checking if the iterates of the scaling factor in the current and previous iteration satisfy isapprox(vcat(a, b), vcat(aprev, bprev); atol=atol, rtol=rtol) where a and b are the current iterates and aprev and bprev the previous ones. The default rtol depends on the types of μ, ν, and C. After maxiter iterations, the computation is stopped.

See also: sinkhorn_unbalanced2

sinkhorn_unbalanced2(μ, ν, C, λ1, λ2, ε; plan=nothing, kwargs...)
sinkhorn_unbalanced2(μ, ν, C, proxdivF1!, proxdivF2!, ε; plan=nothing, kwargs...)

Compute the optimal transport plan for the unbalanced entropically regularized optimal transport problem with source and target marginals μ and ν, cost matrix C of size (length(μ), length(ν)), entropic regularization parameter ε, and marginal relaxation terms λ1 and λ2 or soft marginal constraints with "proxdiv" functions proxdivF1! and proxdivF2!.

A pre-computed optimal transport plan may be provided. The other keyword arguments supported here are the same as those in the sinkhorn_unbalanced for unbalanced optimal transport problems with general soft marginal constraints.

See also: sinkhorn_unbalanced

quadreg(μ, ν, C, ε, alg::QuadraticOT; kwargs...)

Computes the optimal transport plan of histograms μ and ν with cost matrix C and quadratic regularization parameter ε.

The optimal transport plan γ is of the same size as C and solves

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

where $\Omega(\gamma) = \frac{1}{2} \sum_{i,j} \gamma_{i,j}^2$ is the quadratic regularization term.

Every check_convergence steps it is assessed if the algorithm is converged by checking if the iterate of the transport plan γ satisfies

    norm_diff < max(atol, rtol * max(norm(μ, Inf), norm(ν, Inf)))

where

\[ \text{normdiff} = \max\{ \| \gamma \mathbf{1} - \mu \|_\infty , \| \gamma^\top \mathbf{1} - \nu \|_\infty \} . \]

After maxiter iterations, the computation is stopped.

Note that unlike in the case of Sinkhorn's algorithm for the entropic regularisation, batch computation of optimal transport is not supported for the quadratic regularisation.

See also: sinkhorn, QuadraticOTNewton

QuadraticOTNewton

Semi-smooth Newton method (Algorithm 2 of Lorenz et al. 2019 ^[LMM19]) for solving quadratically regularised optimal transport

See also: QuadraticOTNewton, quadreg

ot_entropic_semidual(μ, v, eps, K)

Computes the semidual (in the second argument) of the entropic optimal transport loss, with source marginal μ, regularization parameter ε, and Gibbs kernel K.

That is,

\[ \operatorname{OT}_{\varepsilon}(\mu, \nu) = \inf_{\gamma \in \Pi(\mu, \nu)} \langle \gamma, C \rangle + \varepsilon \Omega(\gamma)\]

with $\Omega(\gamma) = \sum_{i,j} \gamma_{ij} \log \gamma_{ij}$, then the semidual in the second argument ν is ^[Z21]

\[\begin{aligned} \operatorname{OT}_{\varepsilon}^*(\mu, v) &= \sup_{\nu} \langle v, \nu \rangle - \operatorname{OT}_{\varepsilon}(\mu, \nu) \ \ &= -\varepsilon \left\langle \mu, \log\left( \dfrac{\mu}{K e^{v/\varepsilon}} \right) - 1\right\rangle. \end{aligned}\]

Notably, the semidual is computationally advantageous for solving variational problems since it is a smooth and unconstrained function of v since it admits a closed form gradient. See ^[CP16] for a detailed discussion of dual methods for variational problems in optimal transport.

ot_entropic_semidual_grad(μ, v, eps, K)

Computes the gradient with respect to v of the semidual of the entropic optimal transport loss. That is,

\[\nabla_v \operatorname{OT}^*_{\varepsilon}(\mu, v) = K^\top \left( \dfrac{\mu}{K e^{v/\varepsilon}} \right) \odot e^{v/\varepsilon}.\]

See also: ot_entropic_semidual

getprimal_ot_entropic_semidual(μ, v, eps, K)

Computes the the primal variable ν corresponding to the dual variable v at optimality. That is,

\[\nu^\star = e^{v^\star/\varepsilon} \odot K^\top \dfrac{\mu}{K e^{v^\star/\varepsilon}}. \]

See also: ot_entropic_semidual

ot_entropic_dual(u, v, eps, K)

Computes the dual in both arguments of entropic optimal transport loss, where u and v are the dual variables associated with the source and target marginals respectively.

That is,

\[ \begin{aligned} \operatorname{OT}_{\varepsilon}^*(u, v) &= \sup_{\mu, \nu} \langle u, \mu \rangle + \langle v, \nu \rangle - \operatorname{OT}_\varepsilon(\mu, \nu) \ \ &= \varepsilon \log \langle e^{u/\varepsilon}, K e^{v/\varepsilon} \rangle. \end{aligned}\]

ot_entropic_dual_grad(u, v, eps, K)

Computes the gradient with respect to u and v of the dual of the entropic optimal transport loss. That is,

\[\begin{aligned} \nabla_u \operatorname{OT}^*_{\varepsilon}(u, v) &= \dfrac{e^{u/\varepsilon} \odot K e^{v/\varepsilon}}{\langle e^{u/\varepsilon}, K e^{v/\varepsilon} \rangle} \ \ \nabla_v \operatorname{OT}^*_{\varepsilon}(u, v) &= \dfrac{e^{v/\varepsilon} \odot K^\top e^{u/\varepsilon}}{\langle e^{v/\varepsilon}, K^\top e^{u/\varepsilon} \rangle}. \end{aligned}\]

See also: ot_entropic_dual

getprimal_ot_entropic_dual(u, v, eps, K)

Computes the the primal variable γ corresponding to the dual variable u, v at optimality. That is,

\[ \gamma = \operatorname{softmax}(\mathrm{diag}(e^{u/\varepsilon}) K \mathrm{diag}(e^{v/\varepsilon}))\]

See also: ot_entropic_dual