Computational Optimal Transport with OptimalTransport.jl
Stephen Zhang
2 November 2020
Introduction
This file contains a collection of example codes for various functions offered by the OptimalTransport.jl package, and can be treated as an informal mini-tutorial for using the package.
Installation
Install the package as follows,
# using Pkg; # Pkg.add("OptimalTransport")
and load into our environment:
using OptimalTransport using Distances using LinearAlgebra using Plots, LaTeXStrings using StatsPlots using Random, Distributions
ERROR: ArgumentError: Package StatsPlots not found in current path:
- Run `import Pkg; Pkg.add("StatsPlots")` to install the StatsPlots package.
Problem setup
First, let us initialise two random probability measures $\mu$ (source measure) and $\nu$ (target measure) in 1D:
N = 200; M = 200 μ_spt = rand(N) ν_spt = rand(M) μ = fill(1/N, N) ν = fill(1/M, M);
Now we compute the quadratic cost matrix $C_{ij} = \| x_i - x_j \|^2$
C = pairwise(SqEuclidean(), μ_spt', ν_spt');
Earth mover's distance (EMD)
The earth mover's distance is defined as the optimal value of the Monge-Kantorovich problem:
\[ \min_{\gamma \in \Pi(\mu, \nu)} \langle \gamma, C \rangle = \min_{\gamma \in \Pi(\mu, \nu)} \sum_{i, j} \gamma_{ij} C_{ij} \]
where $\Pi(\mu, \nu)$ denotes the set of joint distributions whose marginals agree with $\mu$ and $\nu$. In the case where $C$ is the quadratic cost, this corresponds to what is known as the 2-Wasserstein distance.
N.B. At the moment, this functionality is available through PyCall and the Python OT library.
Using emd() returns the optimal transport plan $\gamma$:
γ = OptimalTransport.emd(μ, ν, C);
ERROR: MethodError: no method matching emd(::Array{Float64,1}, ::Array{Float64,1}, ::Array{Float64,2})
Closest candidates are:
emd(::Any, ::Any, ::Any, !Matched::Any) at /zfs/users/syz/syz/.julia/packages/OptimalTransport/45wZO/src/OptimalTransport.jl:47
whilst using emd2() returns the optimal transport cost at the minimum:
OptimalTransport.emd2(μ, ν, C)
ERROR: MethodError: no method matching emd2(::Array{Float64,1}, ::Array{Float64,1}, ::Array{Float64,2})
Closest candidates are:
emd2(::Any, ::Any, ::Any, !Matched::Any) at /zfs/users/syz/syz/.julia/packages/OptimalTransport/45wZO/src/OptimalTransport.jl:74
Entropically regularised transport using Sinkhorn scaling algorithm
We may add an entropy term to the Monge-Kantorovich problem to obtain the entropically regularised optimal transport problem:
\[ \min_{\gamma \in \Pi(\mu, \nu)} \langle \gamma, C \rangle - \epsilon H(\gamma) \]
where $H(\gamma) = -\sum_{i, j} \gamma_{ij} \log(\gamma_{ij})$ is the entropy of the coupling $\gamma$ and $\epsilon$ controls the strength of regularisation.
This problem is strictly convex and admits a very efficient iterative scaling scheme for its solution known as the Sinkhorn algorithm [Peyre 2019].
Compute the transport plan using native Julia vs. POT
ϵ = 0.01 γ_ = OptimalTransport.pot_sinkhorn(μ, ν, C, ϵ); γ = OptimalTransport.sinkhorn(μ, ν, C, ϵ)
200×200 Array{Float64,2}:
1.00006e-24 6.56364e-6 9.51197e-17 0.000131254 0.00013603 1.69017e-
23 2.38572e-10 3.77072e-21 … 0.000109709 1.41985e-5 1.14677e-12 6.3
2456e-7 7.43872e-8 2.85166e-7 1.08118e-6
0.000193881 2.04971e-14 6.25759e-5 3.32854e-27 4.83878e-23 0.0001935
46 1.57911e-7 0.000171382 3.19241e-21 1.47463e-36 3.4518e-6 9.5
228e-12 4.25691e-10 2.75338e-45 1.66218e-42
1.84201e-5 7.73023e-11 0.000133929 2.23481e-21 7.06935e-18 2.88922e-
5 9.24514e-6 6.18718e-5 2.30812e-16 2.16438e-29 5.26422e-5 9.2
682e-9 1.60426e-7 5.36534e-37 1.44115e-34
1.30544e-6 6.00187e-9 7.43768e-5 5.33135e-18 6.15443e-15 2.76028e-
6 4.57418e-5 1.05967e-5 1.26223e-13 3.96649e-25 0.00010705 2.9
3942e-7 2.71763e-6 5.43242e-32 8.54362e-30
6.81945e-9 8.38366e-7 7.79228e-6 1.0927e-13 2.9055e-11 2.22772e-
8 0.000115866 2.00123e-7 3.02739e-10 1.58397e-19 7.42707e-5 1.1
1452e-5 4.13372e-5 2.61525e-25 1.88618e-23
9.03194e-5 9.89365e-13 0.000114152 1.5567e-24 1.15988e-20 0.0001099
09 1.22822e-6 0.000143317 … 5.6222e-19 2.66531e-33 1.48893e-5 2.5
3874e-10 7.48799e-9 1.5457e-41 6.54338e-39
3.7098e-33 5.233e-9 1.30639e-23 4.25771e-5 7.52844e-6 1.05875e-
31 1.5194e-15 6.57655e-29 2.68585e-6 0.000164685 1.53508e-18 1.0
4841e-10 4.10418e-12 6.63338e-5 9.83409e-5
1.88774e-26 1.89361e-6 4.00428e-18 0.000143809 0.000100634 3.58409e-
25 2.35465e-11 1.00378e-22 6.77151e-5 3.44265e-5 8.00462e-14 1.2
8733e-7 1.18591e-8 1.34573e-6 4.14179e-6
5.75199e-28 5.7234e-7 2.39616e-19 0.000133597 6.71796e-5 1.20442e-
26 2.88605e-12 4.08455e-24 3.88131e-5 6.24011e-5 7.33038e-15 2.9
0123e-8 2.17599e-9 4.27182e-6 1.10315e-5
6.02812e-8 1.5749e-7 2.21872e-5 3.1174e-15 1.44349e-12 1.67078e-
7 9.90213e-5 1.08858e-6 1.94257e-11 1.4716e-21 0.000103532 3.4
2663e-6 1.79468e-5 9.48587e-28 9.18471e-26
9.8553e-25 6.53528e-6 9.40224e-17 0.000131348 0.000135925 1.66634e-
23 2.3658e-10 3.72078e-21 … 0.00010955 1.42513e-5 1.1357e-12 6.2
8893e-7 7.39e-8 2.86948e-7 1.08708e-6
1.84455e-38 2.12598e-11 4.66925e-28 4.60426e-6 3.09916e-7 7.00819e-
37 4.41455e-19 7.61538e-34 7.08195e-8 0.000125607 1.90273e-22 1.8
0636e-13 3.87755e-15 0.000260204 0.000230988
1.4819e-12 3.70797e-5 6.16956e-8 1.91464e-9 8.75432e-8 8.15567e-
12 4.18224e-5 2.03064e-10 4.04955e-7 9.76778e-14 5.67383e-6 0.0
00103207 0.000128029 3.19161e-18 9.03835e-17
3.93975e-10 4.46781e-6 1.70359e-6 5.33022e-12 7.38689e-10 1.56109e-
9 0.000104153 2.04522e-8 5.69867e-9 2.88669e-17 3.75761e-5 3.3
2962e-5 8.23356e-5 1.43891e-22 7.34228e-21
2.54812e-26 2.0896e-6 5.09405e-18 0.000143792 0.000103586 4.79648e-
25 2.81269e-11 1.32094e-22 7.06402e-5 3.24602e-5 9.8096e-14 1.4
5766e-7 1.36728e-8 1.20794e-6 3.77546e-6
9.23926e-17 0.00013132 9.62852e-11 1.43857e-6 1.35928e-5 8.11254e-
16 1.99702e-6 5.03319e-14 … 3.03844e-5 1.78091e-9 6.74321e-8 9.0
1008e-5 4.19114e-5 8.43231e-13 1.03377e-11
0.000144048 1.26291e-13 8.66196e-5 5.77459e-26 6.1899e-22 0.0001573
84 4.23354e-7 0.00016625 3.54843e-20 4.73783e-35 7.07336e-6 4.4
7635e-11 1.65553e-9 1.48294e-43 7.61509e-41
1.05526e-8 6.18987e-7 9.69943e-6 5.61673e-14 1.66146e-11 3.34008e-
8 0.00011445 2.82093e-7 1.81838e-10 6.56318e-20 8.05942e-5 9.0
4574e-6 3.58502e-5 9.04479e-26 6.90313e-24
7.93428e-6 3.91486e-10 0.000119101 3.78044e-20 8.38547e-17 1.38251e-
5 1.77587e-5 3.63616e-5 2.32437e-15 7.50632e-28 7.39372e-5 3.4
2457e-8 4.7532e-7 3.39685e-35 7.55673e-33
0.000193815 2.05473e-14 6.26061e-5 3.34127e-27 4.85533e-23 0.0001935
04 1.58125e-7 0.000171384 3.20274e-21 1.48148e-36 3.45525e-6 9.5
4276e-12 4.26476e-10 2.76804e-45 1.67067e-42
1.36118e-5 1.4326e-10 0.000129703 6.49511e-21 1.79975e-17 2.22045e-
5 1.19336e-5 5.134e-5 … 5.52798e-16 8.22248e-29 6.04594e-5 1.5
2704e-8 2.43417e-7 2.55144e-36 6.38791e-34
1.19091e-6 6.75988e-9 7.21907e-5 6.65707e-18 7.45504e-15 2.54085e-
6 4.74198e-5 9.92723e-6 1.50772e-13 5.26643e-25 0.000108045 3.2
2259e-7 2.92369e-6 7.59376e-32 1.17518e-29
2.72856e-26 2.13686e-6 5.38138e-18 0.000143767 0.000104257 5.12606e-
25 2.92888e-11 1.4063e-22 7.13157e-5 3.20221e-5 1.02748e-13 1.4
9943e-7 1.41228e-8 1.17831e-6 3.69583e-6
1.41412e-21 4.21564e-5 2.75182e-14 6.01677e-5 0.000135613 1.89855e-
20 1.27908e-8 2.70106e-18 0.000156511 1.35215e-6 1.22003e-10 8.0
99e-6 1.54458e-6 7.27359e-9 4.16596e-8
4.0133e-5 1.23956e-11 0.000135449 9.99531e-23 4.60401e-19 5.63166e-
5 4.13741e-6 9.70167e-5 1.78769e-17 4.52684e-31 3.2817e-5 2.0
7502e-9 4.53765e-8 5.93377e-39 1.94582e-36
⋮ ⋮
⋱ ⋮
1.18903e-38 1.72764e-11 3.2221e-28 4.19063e-6 2.7282e-7 4.56249e-
37 3.275e-19 5.05448e-34 6.13906e-8 0.000122303 1.37064e-22 1.4
2505e-13 2.9962e-15 0.000268103 0.000233822
2.32368e-29 1.76564e-7 1.76024e-20 0.000111456 4.182e-5 5.30653e-
28 4.00175e-13 2.13211e-25 2.11093e-5 9.41226e-5 7.85078e-16 6.9
0066e-9 4.31379e-10 1.05859e-5 2.34008e-5
0.000184419 2.90868e-14 6.69549e-5 5.74485e-27 7.88383e-23 0.0001872
71 1.9146e-7 0.000171451 5.065e-21 2.85969e-36 3.97776e-6 1.2
8394e-11 5.53743e-10 5.88746e-45 3.4471e-42
2.63602e-8 3.11943e-7 1.51303e-5 1.29321e-14 4.81729e-12 7.79257e-
8 0.000108263 5.75888e-7 5.86384e-11 9.47946e-21 9.3429e-5 5.5
9458e-6 2.5592e-5 8.8367e-27 7.62263e-25
4.4922e-6 1.01429e-9 0.000105475 2.06163e-19 3.6734e-16 8.35225e-
6 2.52952e-5 2.49381e-5 … 9.2038e-15 6.37536e-27 8.68126e-5 7.3
0399e-8 8.84626e-7 4.1827e-34 8.28326e-32
2.09916e-7 4.89487e-8 3.79573e-5 2.96995e-16 1.94764e-13 5.24813e-
7 7.96175e-5 2.79531e-6 3.0774e-12 6.93507e-23 0.000113152 1.4
5075e-6 9.43487e-6 2.4778e-29 2.88607e-27
9.95715e-23 2.29927e-5 3.53825e-15 8.93671e-5 0.000149991 1.45884e-
21 3.1181e-9 2.46189e-19 0.000151093 3.64618e-6 2.29306e-11 3.3
996e-6 5.3969e-7 3.23361e-8 1.58365e-7
0.000158426 7.49365e-14 7.94519e-5 2.53334e-26 2.96794e-22 0.0001685
95 3.20219e-7 0.000169158 1.7726e-20 1.73653e-35 5.78648e-6 2.8
7427e-11 1.12345e-9 4.67495e-44 2.51666e-41
4.1703e-21 5.23537e-5 6.29245e-14 4.89375e-5 0.000124933 5.39602e-
20 2.23219e-8 7.14269e-18 0.000152712 8.54848e-7 2.37638e-10 1.1
2347e-5 2.31523e-6 3.72295e-9 2.27812e-8
4.85646e-17 0.000131913 6.12319e-11 1.98396e-6 1.70765e-5 4.38371e-
16 1.55487e-6 2.87069e-14 … 3.65641e-5 2.96606e-9 4.83555e-8 8.3
3111e-5 3.65683e-5 1.64503e-12 1.91929e-11
0.000159057 7.32432e-14 7.91412e-5 2.4438e-26 2.87412e-22 0.0001690
72 3.16304e-7 0.000169258 1.71963e-20 1.66211e-35 5.73525e-6 2.8
1898e-11 1.10449e-9 4.44538e-44 2.39799e-41
4.82403e-17 0.000131913 6.0943e-11 1.99047e-6 1.71161e-5 4.35568e-
16 1.55077e-6 2.85393e-14 3.66325e-5 2.98162e-9 4.81872e-8 8.3
2398e-5 3.6515e-5 1.65637e-12 1.93153e-11
4.66307e-8 1.96097e-7 1.97424e-5 4.89979e-15 2.11848e-12 1.31896e-
7 0.000102241 8.94157e-7 2.76219e-11 2.65703e-21 0.000100626 4.0
1458e-6 2.01483e-5 1.92385e-27 1.79617e-25
6.745e-8 1.42778e-7 2.33356e-5 2.54963e-15 1.21692e-12 1.85276e-
7 9.75186e-5 1.18615e-6 1.66071e-11 1.13201e-21 0.000104723 3.1
9129e-6 1.70324e-5 6.93136e-28 6.82018e-26
0.000253436 1.88108e-15 3.74461e-5 8.32577e-29 1.77442e-24 0.0002258
87 4.12027e-8 0.000160273 … 1.39662e-22 1.70135e-38 1.26215e-6 1.2
2758e-12 6.96217e-11 1.66052e-47 1.22823e-44
2.63125e-8 3.12389e-7 1.51173e-5 1.29711e-14 4.82956e-12 7.77955e-
8 0.00010828 5.75081e-7 5.87751e-11 9.51701e-21 9.34049e-5 5.6
0025e-6 2.56105e-5 8.87872e-27 7.65698e-25
6.09726e-6 6.15834e-10 0.000112995 8.44388e-20 1.68934e-16 1.0954e-5
2.10753e-5 3.05839e-5 4.46505e-15 2.06567e-27 8.01127e-5 4.9
1535e-8 6.39821e-7 1.1135e-34 2.34506e-32
9.76366e-7 8.69555e-9 6.75706e-5 1.06751e-17 1.12038e-14 2.12354e-
6 5.10935e-5 8.6144e-6 2.19908e-13 9.62928e-25 0.000109941 3.9
1326e-7 3.40987e-6 1.54992e-31 2.31738e-29
0.000316266 6.39022e-17 1.61893e-5 4.84878e-31 1.73652e-26 0.0002417
08 5.77679e-9 0.000126979 1.73648e-24 3.46799e-41 2.79701e-7 6.6
1247e-14 5.17951e-12 1.40383e-50 1.36787e-47
1.50328e-32 9.44692e-9 4.17322e-23 5.17454e-5 1.02796e-5 4.14476e-
31 3.77003e-15 2.40668e-28 … 3.86973e-6 0.000158146 4.22096e-18 2.0
9889e-10 8.8338e-12 5.22855e-5 8.24585e-5
3.13823e-15 0.000109084 1.09788e-9 1.94394e-7 3.13456e-6 2.35198e-
14 7.1412e-6 1.07081e-12 8.96564e-6 8.16417e-11 3.86359e-7 0.0
00120314 7.80406e-5 1.56183e-14 2.5431e-13
4.85779e-17 0.000131913 6.12438e-11 1.98369e-6 1.70749e-5 4.38486e-
16 1.55503e-6 2.87138e-14 3.65613e-5 2.96542e-9 4.83624e-8 8.3
314e-5 3.65704e-5 1.64457e-12 1.91879e-11
8.50901e-25 6.25662e-6 8.36831e-17 0.000132268 0.000134855 1.44507e-
23 2.17473e-10 3.25461e-21 0.000107945 1.47896e-5 1.03037e-12 5.9
4174e-7 6.91769e-8 3.05396e-7 1.14786e-6
7.30777e-9 7.99786e-7 8.0695e-6 9.84759e-14 2.66271e-11 2.37542e-
8 0.000115707 2.11333e-7 2.79592e-10 1.37996e-19 7.52724e-5 1.0
7916e-5 4.0445e-5 2.21463e-25 1.61151e-23
Now we can check that both implementations arrive at the same result:
norm(γ - γ_, Inf)
4.5747369867982224e-11
Again, we can compute the optimal value (a scalar) of the entropic OT problem using sinkhorn2():
OptimalTransport.sinkhorn2(μ, ν, C, ϵ)
0.009683588757086085
Quadratically regularised transport.
Try computing the transport plan for the same problem, this time using a quadratic regularisation [Lorenz 2019] instead of the more common entropic regulariser term. We solve the problem
\[ \min_{\gamma \in \Pi(\mu, \nu)} \langle \gamma, C \rangle + \epsilon \frac{\| \gamma \|_F^2}{2} \]
One property of quadratically regularised OT is that the resulting transport plan $\gamma$ is sparse. We take advantage of this and represent it as a sparse matrix.
γ = OptimalTransport.quadreg(μ, ν, C, ϵ); γ
200×200 SparseArrays.SparseMatrixCSC{Float64,Int64} with 688 stored entries
:
[26 , 1] = 0.00235892
[61 , 1] = 0.000785196
[111, 1] = 6.00416e-5
[191, 1] = 0.00179472
[16 , 2] = 0.00330926
[50 , 2] = 0.00404066
[57 , 2] = 0.00279057
[113, 2] = 0.00194348
[149, 2] = 0.00390947
[3 , 3] = 0.00239125
[29 , 3] = 0.00238983
[114, 3] = 0.000217794
[8 , 4] = 0.00226896
[15 , 4] = 0.000140892
[175, 4] = 0.00259186
[60 , 5] = 0.00145916
[77 , 5] = 0.00161994
[97 , 5] = 0.00182218
[183, 5] = 0.00145827
[2 , 6] = 0.00157347
[20 , 6] = 0.00156927
[116, 6] = 0.00151072
[179, 6] = 0.000345422
[5 , 7] = 0.00141462
⋮
[136, 193] = 0.00101561
[161, 193] = 0.00398601
[24 , 194] = 0.000323591
[110, 194] = 0.00121398
[185, 194] = 0.00640302
[27 , 195] = 0.00181813
[55 , 195] = 0.00141802
[65 , 195] = 0.00176559
[41 , 196] = 0.000450519
[49 , 196] = 0.00323492
[69 , 196] = 0.0011165
[92 , 196] = 0.000196945
[157, 197] = 0.00320402
[172, 197] = 0.00562157
[44 , 198] = 0.00147756
[71 , 198] = 0.00282114
[120, 198] = 0.00550232
[163, 198] = 0.000784002
[62 , 199] = 0.00141127
[82 , 199] = 0.00108189
[107, 199] = 0.000859831
[118, 199] = 0.00164875
[78 , 200] = 0.00257453
[90 , 200] = 0.00202903
[124, 200] = 0.000398191
Stabilized Sinkhorn
This is a log-stabilised version of the Sinkhorn algorithm which is useful when $\epsilon$ is very small [Schmitzer 2019]
ϵ = 0.005 γ = sinkhorn_stabilized(μ, ν, C, ϵ, max_iter = 5000); γ_ = pot_sinkhorn(μ, ν, C, ϵ, method = "sinkhorn_stabilized", max_iter = 5000); norm(γ - γ_, Inf) # Check that we get the same result as POT
3.2784934475885673e-10
Stabilized Sinkhorn eps-scaling
In addition to log-stabilisation, we can additionally use $\epsilon$-scaling [Schmitzer 2019]
γ = sinkhorn_stabilized_epsscaling(μ, ν, C, ϵ, max_iter = 5000) γ_ = pot_sinkhorn(μ, ν, C, ϵ, method = "sinkhorn_epsilon_scaling", max_iter = 5000) norm(γ - γ_, Inf) # Check that we get the same result as POT
3.5134034066406847e-10
Unbalanced transport
Unbalanced transport was introduced by [Chizat 2019] to interpolate between general positive measures which do not have the same total mass. That is, for $\mu, \nu \in \mathcal{M}_+$ and a cost function (e.g. quadratic) $C$, we solve the following problem:
\[ \min_{\gamma \ge 0} \epsilon \mathrm{KL}(\gamma | \exp(-C/\epsilon)) + \lambda_1 \mathrm{KL}(\gamma_1 | \mu) + \lambda_2 \mathrm{KL}(\gamma_2 | \nu) \]
where $\epsilon$ controls the strength of entropic regularisation, and $\lambda_1, \lambda_2$ control how strongly we enforce the marginal constraints.
We construct two random measures, now with different total masses:
N = 100; M = 200 μ_spt = rand(N) ν_spt = rand(M) μ = fill(1/N, N) ν = fill(1/N, M);
Set up quadratic cost matrix:
C = pairwise(SqEuclidean(), μ_spt', ν_spt')
100×200 Array{Float64,2}:
8.16407e-5 0.180375 0.00362636 0.14397 0.0225504 0.014613
0.512335 0.0259424 … 0.0417258 0.115198 0.554281 0.006
17441 0.0101448 0.398316 0.191951
0.00101127 0.147356 0.0102122 0.114649 0.0119535 0.00640771
0.455543 0.0144553 0.0267103 0.0891457 0.495143 0.014
2596 0.00358628 0.348438 0.157836
0.0167408 0.297086 0.00361578 0.249785 0.0731803 0.0581942
0.699107 0.0791955 0.105378 0.211379 0.747967 0.001
745 0.0488729 0.564712 0.311892
0.102078 0.00924928 0.151128 0.00259099 0.0318139 0.0431178
0.149957 0.0280449 0.0154414 0.000118293 0.173029 0.165
738 0.0518978 0.0915609 0.0120099
0.00309923 0.222162 0.000184527 0.181535 0.0387315 0.0280627
0.58127 0.04314 0.0629529 0.14903 0.625895 0.001
02031 0.021714 0.459356 0.23499
0.240117 0.00552754 0.312786 0.0143087 0.121721 0.143012
0.0469686 0.114235 … 0.0868976 0.0254863 0.0602443 0.333
657 0.158668 0.0174424 0.00371259
0.612718 0.134757 0.725934 0.170044 0.41169 0.450126
0.00577951 0.397823 0.345191 0.204657 0.00223713 0.757
554 0.477587 0.0258167 0.125087
0.00103464 0.200557 0.0013756 0.162058 0.0300323 0.0207402
0.545982 0.0339284 0.0517104 0.131434 0.589257 0.003
07439 0.0153393 0.428047 0.212754
0.612959 0.13487 0.726197 0.170171 0.411888 0.450333
0.00580297 0.398017 0.345372 0.204796 0.00225173 0.757
823 0.4778 0.0258663 0.125196
0.204947 0.001372 0.272448 0.00677527 0.097081 0.116187
0.0645309 0.0904083 0.0662945 0.0149665 0.0799498 0.291
95 0.130339 0.0286884 0.000558101
0.0270285 0.336485 0.00905325 0.286013 0.0933521 0.0763152
0.758891 0.100131 … 0.129339 0.244804 0.809763 0.005
89676 0.0655818 0.618567 0.35223
0.00041104 0.156338 0.00801543 0.122587 0.0146067 0.0083859
0.471235 0.0173598 0.0306108 0.0961614 0.511498 0.011
6396 0.00509967 0.362179 0.167128
0.0331285 0.357224 0.0127143 0.305158 0.104423 0.0863544
0.789881 0.111585 0.142315 0.262539 0.841764 0.008
91127 0.0749109 0.646576 0.373442
0.0176227 0.0800436 0.0408061 0.0564767 7.02576e-5 0.00043688
3 0.329464 0.000371736 0.00390413 0.0390548 0.363264 0.048
56 0.00168628 0.23945 0.0878151
0.204581 0.00134224 0.272026 0.00670893 0.0968295 0.115912
0.0647362 0.0901656 0.0660867 0.0148678 0.0801784 0.291
513 0.130047 0.0288254 0.000539178
0.0823265 0.0165751 0.126865 0.0069677 0.0212557 0.030652
0.176244 0.0181967 … 0.0084075 0.00188767 0.201187 0.140
279 0.0381186 0.112333 0.0202096
0.428934 0.0572457 0.524445 0.0809586 0.263989 0.294938
0.00268418 0.252908 0.211322 0.105338 0.0064857 0.551
371 0.317245 0.00107871 0.0510057
0.474951 0.0748004 0.575203 0.101613 0.300342 0.333296
0.000308817 0.288515 0.24397 0.128733 0.0021435 0.603
387 0.356983 0.00449966 0.0676417
0.635258 0.145436 0.75045 0.182015 0.430202 0.469474
0.00815238 0.416024 0.36216 0.217769 0.00379034 0.782
594 0.497511 0.0306051 0.135383
0.276419 0.0121188 0.354038 0.0241358 0.147935 0.171319
0.0327553 0.13967 0.109245 0.0381743 0.0439778 0.376
221 0.188416 0.00927975 0.00934494
0.0341381 0.0533172 0.0645261 0.0344599 0.00190381 0.00531679
0.272458 0.00107152 … 0.000109589 0.0212017 0.30327 0.074
1898 0.00866369 0.191251 0.059693
0.436604 0.060069 0.532922 0.0843097 0.270013 0.301304
0.00211416 0.258805 0.216715 0.109155 0.0055808 0.560
062 0.323845 0.00149558 0.0536726
0.0776455 0.0187747 0.121037 0.00841793 0.0189109 0.0278225
0.183261 0.0160323 0.00695822 0.00267535 0.20868 0.134
148 0.0349553 0.11795 0.0226313
0.538832 0.101367 0.645301 0.132243 0.351553 0.387136
0.000745926 0.338748 0.290325 0.162956 1.99692e-6 0.675
133 0.412633 0.012536 0.0930037
0.228 0.00382215 0.298934 0.0114693 0.113139 0.133696
0.0525539 0.105926 0.0796708 0.0216444 0.0665491 0.319
345 0.148847 0.0209073 0.00234325
⋮ ⋮
⋱ ⋮
0.153645 0.000561429 0.212733 0.000465564 0.0629223 0.0784711
0.0990766 0.0575735 0.0387075 0.00379486 0.117985 0.230
005 0.0901738 0.0529513 0.00137722
0.217623 0.00258377 0.287034 0.00923565 0.105865 0.125778
0.0577148 0.0988914 0.073586 0.0185308 0.0723415 0.307
042 0.140486 0.0242071 0.00139984
0.488473 0.0802238 0.590074 0.107919 0.311114 0.344639
6.13361e-5 0.299075 0.253688 0.135818 0.00133638 0.618
616 0.368719 0.00590146 0.0728037
0.295803 0.0164373 0.375932 0.0300951 0.162204 0.186649
0.0265239 0.153544 0.121553 0.0455844 0.0367054 0.398
781 0.204478 0.00611665 0.0131771
0.481145 0.0772705 0.582017 0.104489 0.305271 0.338488
0.000171455 0.293347 … 0.248415 0.131967 0.00174882 0.610
365 0.362356 0.00512064 0.0699916
0.0857158 0.0151038 0.131064 0.0060258 0.0229948 0.0327335
0.171369 0.0198082 0.00951388 0.0014138 0.195976 0.144
693 0.0404358 0.108448 0.0185815
0.157012 0.000377261 0.216692 0.000668131 0.0650834 0.0808824
0.0964059 0.0596415 0.0404064 0.00433935 0.115069 0.234
12 0.0927573 0.0510038 0.00107844
0.321674 0.0229501 0.405028 0.0387163 0.181502 0.207311
0.0194817 0.172335 0.138332 0.0560695 0.0283254 0.428
732 0.226078 0.00301665 0.0190652
0.461678 0.0695899 0.560587 0.0955245 0.289806 0.322193
0.000743701 0.278191 0.234484 0.121868 0.00313551 0.588
415 0.345489 0.00329267 0.0626914
0.462342 0.0698476 0.561319 0.0958265 0.290332 0.322747
0.000717317 0.278706 … 0.234957 0.122209 0.00308109 0.589
164 0.346063 0.00334893 0.0629361
0.0468246 0.0397126 0.081593 0.0237187 0.00566371 0.0109289
0.240443 0.0041421 0.000447598 0.0129922 0.269439 0.092
4178 0.0155511 0.16459 0.0452398
1.83316e-5 0.176359 0.00422152 0.140385 0.0211452 0.0134862
0.505552 0.0244336 0.0398062 0.111994 0.547224 0.006
94413 0.00920978 0.392338 0.187808
0.266459 0.0101057 0.342754 0.0212572 0.140674 0.163498
0.0363066 0.132617 0.103018 0.0345305 0.0480783 0.364
587 0.18021 0.0112127 0.00758826
0.325059 0.0238608 0.408825 0.0398964 0.184047 0.210031
0.0186597 0.174815 0.140555 0.0574879 0.0273324 0.432
638 0.228918 0.00269856 0.019896
0.0247684 0.066714 0.0513632 0.0453771 0.000263973 0.00207307
0.301797 2.86093e-5 … 0.00143291 0.0299268 0.334182 0.060
0214 0.00431565 0.215953 0.0738246
0.0644047 0.0262082 0.104352 0.0135997 0.0126896 0.0201447
0.205172 0.010353 0.00342778 0.00586641 0.232019 0.116
55 0.0262747 0.13565 0.0307321
0.0125659 0.278539 0.00183551 0.232803 0.0641254 0.050152
0.670495 0.0697639 0.0944524 0.19578 0.718362 0.000
599499 0.0415278 0.539027 0.29288
0.0136859 0.0892118 0.0346859 0.0642176 0.000583016 2.6399e-5
0.347809 0.0012281 0.00612256 0.0455338 0.382515 0.041
861 0.00064012 0.255127 0.0974063
0.00286625 0.13114 0.0150779 0.100401 0.0076729 0.00340021
0.426674 0.00970098 0.0200778 0.0766381 0.465025 0.019
9234 0.00145533 0.323249 0.141037
0.0169288 0.297877 0.00370345 0.25051 0.0735729 0.0585444
0.700319 0.0796039 … 0.105849 0.212045 0.749221 0.001
80606 0.0491938 0.565802 0.312702
0.138205 0.00192812 0.194494 1.85265e-6 0.053189 0.0675538
0.112212 0.0482808 0.0311615 0.00171285 0.132281 0.211
023 0.0784413 0.0626638 0.00328635
9.61334e-5 0.181029 0.00353431 0.144555 0.022782 0.0147996
0.513436 0.0261908 0.0420407 0.115721 0.555427 0.006
05412 0.0103004 0.399287 0.192626
0.0283124 0.340978 0.0098026 0.290156 0.0957255 0.0784625
0.76563 0.102588 0.13213 0.248638 0.816724 0.006
50441 0.0675735 0.624653 0.356826
0.602596 0.130033 0.714913 0.164732 0.4034 0.441456
0.00483448 0.389675 0.337604 0.198824 0.0016651 0.746
294 0.468656 0.0237724 0.120537
Now we solve the corresponding unbalanced, entropy-regularised transport problem.
ϵ = 0.01 λ = 1.0 γ_ = pot_sinkhorn_unbalanced(μ, ν, C, ϵ, λ) γ = sinkhorn_unbalanced(μ, ν, C, λ, λ, ϵ)
100×200 Array{Float64,2}:
0.00040934 1.51504e-11 0.000268736 5.19769e-10 5.44264e-5 0.0001228
92 2.80471e-25 3.82536e-5 … 7.60805e-6 7.72267e-9 4.98563e-27 0.0
00211731 0.000190054 1.41528e-20 4.83778e-12
0.000349993 3.86162e-10 0.000130514 9.15335e-9 0.00014736 0.0002619
53 7.70285e-23 0.000113213 3.20433e-5 9.80772e-8 1.73128e-24 8.8
5117e-5 0.000343605 1.94691e-18 1.37588e-10
8.79255e-5 1.46973e-16 0.000305712 1.49924e-14 3.913e-7 1.78789e-
6 2.46604e-33 2.11566e-7 1.4874e-8 5.83762e-13 2.19561e-35 0.0
00374697 4.49226e-6 9.54713e-28 3.39791e-17
8.51487e-9 0.000229122 5.90429e-11 0.000401411 1.20574e-5 3.97492e-
6 8.58005e-10 1.73427e-5 5.89582e-5 0.000429913 1.00695e-10 1.3
9238e-11 1.63426e-6 1.66264e-7 0.00017666
0.000322724 2.47414e-13 0.000404206 1.29477e-11 1.15047e-5 3.41356e-
5 3.0329e-28 7.30453e-6 9.70941e-7 2.79434e-10 4.12415e-30 0.0
00377942 6.37137e-5 3.37066e-23 6.9711e-14
5.51912e-15 0.000212982 3.6067e-18 7.96784e-5 9.62254e-10 1.16852e-
10 1.63256e-5 2.00711e-9 … 2.97772e-8 2.17924e-5 5.10316e-6 4.5
4773e-19 2.41544e-11 0.000176346 0.000259496
7.66209e-32 1.09722e-10 8.68177e-37 2.89834e-12 5.18106e-23 1.1328e-2
4 0.000211842 2.04561e-22 3.80777e-20 7.60963e-14 0.000355943 3.7
3671e-38 7.19161e-26 1.61063e-5 2.93242e-10
0.000386083 2.08887e-12 0.000349189 8.83593e-11 2.67216e-5 6.90891e-
5 1.006e-26 1.78579e-5 2.90826e-6 1.57988e-9 1.56574e-28 0.0
00299504 0.000117292 7.50976e-22 6.26874e-13
7.48733e-32 1.08602e-10 8.46566e-37 2.86475e-12 5.08496e-23 1.11078e-
24 0.000211569 2.00834e-22 3.74338e-20 7.51219e-14 0.000355801 3.6
4163e-38 7.04748e-26 1.60437e-5 2.90369e-10
2.06012e-13 0.000357615 2.25709e-16 0.000187548 1.25309e-8 1.89331e-
9 3.12428e-6 2.40962e-8 2.5898e-7 6.91476e-5 7.88205e-7 3.2
6378e-17 4.54901e-10 6.34685e-5 0.00039421
3.56314e-5 3.24095e-18 0.000201221 4.53942e-16 5.90155e-8 3.31029e-
7 7.08153e-36 2.95633e-8 … 1.53579e-9 2.3394e-14 5.15603e-38 0.0
00280465 9.57887e-7 4.95995e-30 6.82085e-19
0.000377542 1.59779e-10 0.000165158 4.20392e-9 0.000114812 0.0002183
48 1.62926e-23 8.60192e-5 2.20385e-5 4.93986e-8 3.42728e-25 0.0
00116848 0.000300035 5.00538e-19 5.51947e-11
2.16803e-5 4.56176e-19 0.000156254 7.49384e-17 2.18441e-8 1.35839e-
7 3.57597e-37 1.05306e-8 4.69842e-10 4.4466e-15 2.35314e-39 0.0
00232335 4.21999e-7 3.37459e-31 9.15754e-20
6.26556e-5 3.05092e-7 5.77169e-6 2.89924e-6 0.00045581 0.0004485
98 2.17035e-17 0.000436381 0.000295476 1.38456e-5 8.71197e-19 2.7
0194e-6 0.000391652 9.93012e-14 1.42522e-7
2.13845e-13 0.000358955 2.35606e-16 0.00018894 1.28599e-8 1.94762e-
9 3.0631e-6 2.47071e-8 2.64621e-7 6.98864e-5 7.70978e-7 3.4
1195e-17 4.68713e-10 6.26526e-5 0.000395258
6.99767e-8 0.000125567 7.61852e-10 0.000295445 3.95143e-5 1.57643e-
5 7.06057e-11 5.29399e-5 … 0.000135828 0.000410679 6.87219e-12 2.0
2488e-10 7.3911e-6 2.37485e-8 8.87142e-5
8.70475e-24 3.02252e-7 5.79399e-28 2.54032e-8 1.59509e-16 7.37335e-
18 0.000342164 4.76626e-16 2.9399e-14 1.85576e-9 0.000275843 3.9
8746e-29 7.83812e-19 0.000226542 5.73237e-7
6.97945e-26 4.17357e-8 2.89135e-30 2.57279e-9 3.3615e-18 1.27155e-
19 0.000346677 1.08227e-17 8.9735e-16 1.42894e-10 0.00034023 1.7
548e-31 1.1774e-20 0.000128561 8.67701e-8
9.00584e-33 4.22272e-11 8.37454e-38 9.80268e-13 9.10951e-24 1.83205e-
25 0.000187077 3.71026e-23 7.81254e-21 2.29593e-14 0.000341183 3.4
2051e-39 1.09806e-26 1.11712e-5 1.17261e-10
1.20706e-16 9.08905e-5 4.80827e-20 2.46028e-5 5.77098e-11 5.68452e-
12 5.57928e-5 1.30125e-10 2.62895e-9 5.05479e-6 2.14148e-5 5.3
1724e-21 1.01734e-12 0.000329077 0.000121885
1.13383e-5 4.16857e-6 5.08161e-7 2.47347e-5 0.000358094 0.0002598
78 6.12517e-15 0.000383988 … 0.000407531 7.78937e-5 3.31499e-16 1.9
6532e-7 0.000183958 1.16168e-11 2.23899e-6
3.86439e-24 2.17846e-7 2.37267e-28 1.73679e-8 8.34765e-17 3.72923e-
18 0.000346246 2.52625e-16 1.63872e-14 1.21092e-9 0.00028864 1.5
9831e-29 3.8722e-19 0.000207702 4.1967e-7
1.15327e-7 0.000104 1.4081e-9 0.000263741 5.15551e-5 2.15896e-
5 3.61202e-11 6.7837e-5 0.000162037 0.000391722 3.3524e-12 3.8
5797e-10 1.04658e-5 1.39766e-8 7.18629e-5
1.0484e-28 2.61677e-9 2.33242e-33 1.0744e-10 1.79252e-20 5.21312e-
22 0.000296456 6.36439e-20 7.77704e-18 4.16609e-12 0.000376526 1.2
005e-34 4.02792e-23 5.14178e-5 6.13667e-9
1.93609e-14 0.000263772 1.50488e-17 0.000110528 2.37039e-9 3.09758e-
10 9.75271e-6 4.81113e-9 6.40565e-8 3.34178e-5 2.83692e-6 1.9
8681e-18 6.73466e-11 0.000130229 0.000310757
⋮ ⋮
⋱ ⋮
3.87428e-11 0.000431417 9.84448e-14 0.000392118 4.2438e-7 9.15116e-
8 1.09832e-7 7.1487e-7 4.54605e-6 0.000235089 1.95471e-8 1.7
7915e-14 2.8089e-8 6.23894e-6 0.000404047
5.63089e-14 0.000307605 5.09658e-17 0.000142385 5.05495e-9 7.04497e-
10 5.99754e-6 1.0017e-8 1.21284e-7 4.70094e-5 1.63783e-6 7.0
0587e-18 1.60112e-10 9.6468e-5 0.000351864
1.72818e-26 2.32267e-8 6.25562e-31 1.31093e-9 1.09577e-18 3.91508e-
20 0.000340163 3.60365e-18 3.2503e-16 6.73474e-11 0.000353053 3.6
6323e-32 3.48554e-21 0.000106966 4.9568e-8
1.51292e-17 5.13936e-5 4.68881e-21 1.18063e-5 1.2064e-11 1.06866e-
12 9.06025e-5 2.82983e-11 6.6863e-10 2.09811e-6 3.85915e-5 4.8
5134e-22 1.77777e-13 0.000393195 7.23538e-5
3.67764e-26 3.19133e-8 1.43188e-30 1.88911e-9 2.00997e-18 7.40622e-
20 0.000344057 6.53479e-18 … 5.63204e-16 1.01226e-10 0.000346459 8.5
485e-32 6.73507e-21 0.000118271 6.71512e-8
4.87284e-8 0.000142167 4.89248e-10 0.000317254 3.2453e-5 1.25114e-
5 1.12351e-10 4.40371e-5 0.00011884 0.000420831 1.13089e-11 1.2
7272e-10 5.72928e-6 3.42283e-8 0.00010203
2.74147e-11 0.000435415 6.56586e-14 0.000380739 3.38772e-7 7.12466e-
8 1.42141e-7 5.76001e-7 3.80066e-6 0.000220593 2.59258e-8 1.1
6815e-14 2.14954e-8 7.511e-6 0.000412492
9.22577e-19 2.17167e-5 2.07101e-22 4.04051e-6 1.41945e-12 1.09709e-
13 0.000148496 3.50293e-12 1.01213e-10 5.95934e-7 7.23041e-5 1.9
6722e-23 1.66149e-14 0.000434483 3.25447e-5
2.77495e-25 7.40986e-8 1.31483e-29 4.98701e-9 1.01645e-17 4.06962e-
19 0.000349985 3.20413e-17 2.44312e-15 2.9934e-10 0.000324863 8.2
6915e-31 3.91857e-20 0.000152946 1.50097e-7
2.58932e-25 7.20052e-8 1.21861e-29 4.82475e-9 9.61615e-18 3.83907e-
19 0.0003499 3.03449e-17 … 2.32355e-15 2.8847e-10 0.000325696 7.6
5027e-31 3.68931e-20 0.000151651 1.46047e-7
3.01228e-6 1.53513e-5 8.71191e-8 6.8408e-5 0.000232286 0.0001400
73 1.42173e-13 0.000266859 0.000372219 0.000167246 9.22755e-15 2.9
9997e-8 8.72811e-5 1.57854e-10 8.97559e-6
0.000408618 2.24542e-11 0.000251167 7.3789e-10 6.21328e-5 0.0001364
4 5.48219e-25 4.4125e-5 9.14375e-6 1.05541e-8 1.00149e-26 0.0
00194465 0.000207 2.55243e-20 7.26208e-12
3.47741e-16 0.000118286 1.5813e-19 3.4913e-5 1.26941e-10 1.32239e-
11 4.1623e-5 2.80327e-10 5.21456e-9 7.7435e-6 1.51224e-5 1.8
1115e-20 2.45963e-12 0.000288626 0.000154607
6.38962e-19 1.92631e-5 1.37642e-22 3.48872e-6 1.06925e-12 8.12142e-
14 0.000156638 2.65588e-12 7.87373e-11 5.02437e-7 7.7584e-5 1.2
9323e-23 1.21527e-14 0.000435783 2.90993e-5
3.00002e-5 1.13195e-6 1.96478e-6 8.60658e-6 0.00043739 0.0003726
46 3.37754e-16 0.00044184 … 0.000370121 3.37469e-5 1.56184e-17 8.4
0257e-7 0.000294582 1.01843e-12 5.64908e-7
4.71889e-7 5.38354e-5 8.13089e-9 0.000171002 0.000104551 5.0647e-5
4.39597e-12 0.000130311 0.000251077 0.000309927 3.53682e-13 2.4
4072e-9 2.71416e-5 2.59157e-9 3.47984e-5
0.000129984 9.14554e-16 0.000355707 7.97738e-14 9.42365e-7 3.89114e-
6 4.19815e-32 5.29079e-7 4.31894e-8 2.7047e-12 4.12809e-34 0.0
00409158 9.11847e-6 1.21286e-26 2.21478e-16
9.37159e-5 1.23067e-7 1.07394e-5 1.3489e-6 0.000436914 0.0004715
91 3.49703e-18 0.000404162 0.000238813 7.30822e-6 1.2822e-19 5.3
2709e-6 0.000438748 2.08945e-14 5.51088e-8
0.000284372 1.9116e-9 7.84744e-5 3.72164e-8 0.000221141 0.0003461
19 1.35154e-21 0.000178142 6.08376e-5 3.35083e-7 3.44176e-23 4.9
1373e-5 0.000415904 2.36427e-17 7.22002e-10
8.64185e-5 1.36009e-16 0.000303504 1.39653e-14 3.76807e-7 1.729e-6
2.18781e-33 2.03408e-7 … 1.42112e-8 5.46933e-13 1.9398e-35 0.0
00372981 4.35698e-6 8.57453e-28 3.13832e-17
1.90911e-10 0.00039596 6.41858e-13 0.00043218 1.18189e-6 2.86892e-
7 3.10732e-8 1.90508e-6 1.01732e-5 0.000304622 4.92381e-9 1.2
4938e-13 9.55395e-8 2.48549e-6 0.00035126
0.000409281 1.42098e-11 0.000271575 4.90908e-10 5.32497e-5 0.0001207
78 2.51539e-25 3.73639e-5 7.38187e-6 7.3389e-9 4.45161e-27 0.0
00214574 0.000187365 1.28593e-20 4.52808e-12
3.20213e-5 2.11315e-18 0.000190764 3.06503e-16 4.75619e-8 2.72882e-
7 3.68815e-36 2.36261e-8 1.1871e-9 1.62914e-14 2.62644e-38 0.0
00269684 8.02017e-7 2.75763e-30 4.40148e-19
2.02271e-31 1.68837e-10 2.50756e-36 4.7299e-12 1.13872e-22 2.5862e-2
4 0.000223377 4.43274e-22 7.80113e-20 1.3081e-13 0.000361583 1.1
0531e-37 1.68542e-25 1.89567e-5 4.43439e-10
Check agreement with POT:
norm(γ - γ_, Inf) # Check that we get the same result as POT
1.8257663176451944e-12
Example plots
Entropically regularised transport
Let's construct source and target measures again
μ_spt = ν_spt = LinRange(-2, 2, 100) C = pairwise(Euclidean(), μ_spt', ν_spt').^2 μ = exp.((-(μ_spt).^2)/0.5^2) μ /= sum(μ) ν = ν_spt.^2 .*exp.((-(ν_spt).^2)/0.5^2) ν /= sum(ν) plot(μ_spt, μ, label = L"\mu") plot!(ν_spt, ν, label = L"\nu") current()
Now compute the entropic transport plan:
γ = OptimalTransport.sinkhorn(μ, ν, C, 0.01) heatmap(γ, title = "Entropically regularised OT", size = (500, 500)) current()
Quadratically regularised transport
Using a quadratic regularisation, notice how the 'edges' of the transport plan here are sharper compared to the entropic OT transport plan.
γ_quad = Matrix(OptimalTransport.quadreg(μ, ν, C, 5, maxiter = 500)) heatmap(γ_quad, title = "Quadratically regularised OT", size = (500, 500)) current()
Sinkhorn barycenters
For a collection of probability measures $\{ \mu_i \}_{i = 1}^N$ and corresponding cost matrices $C_i$, we define the barycenter to be the measure $\mu$ that solves the following:
\[ \min_{\mu \text{ a distribution}} \sum_{i = 1}^N \lambda_i \mathrm{entropicOT}^{\epsilon}_{C_i}(\mu, \mu_i) \]
where $\mathrm{entropicOT}^\epsilon_{C_i}(\mu, \mu_i)$ denotes the entropic optimal transport cost between measures $(\mu, \mu_i)$ with cost $C_i$ and entropic regularisation strength $\epsilon$.
For instance, below we set up two measures $\mu_0$ and $\mu_1$ and compute the weighted barycenters. The weights are $(w, 1-w)$ where $w = 0.25, 0.5, 0.75$.
spt = LinRange(-1, 1, 250) f(x, σ) = exp.(-x.^2/σ^2) normalize(x) = x./sum(x) mu_all = hcat([normalize(f(spt .- z, 0.1)) for z in [-0.5, 0.5]]...)' C_all = [pairwise(SqEuclidean(), spt', spt') for i = 1:size(mu_all, 1)] plot(size = (400, 200)); plot!(spt, mu_all[1, :], label = L"\mu_0") plot!(spt, mu_all[2, :], label = L"\mu_1") for s = [0.25, 0.5, 0.75] a = sinkhorn_barycenter(mu_all, C_all, 0.01, [s, 1-s]; max_iter = 1000); plot!(spt, a, label= latexstring("\\mu_{$s}")) end current()
