Semi-Discrete Optimal Transport¶
Transport a continuous distribution toward a finite set of weighted points. Study the dual variables and the regions they induce, and connect geometric partitions with iterative optimization of the transport objective.
Run this tour¶
Run the cells in order with a Python 3 kernel. The first cell locates the companion data and toolbox and installs missing dependencies when needed. All worked examples include their implementation directly in this notebook. Random seeds make comparisons reproducible; you can change them to explore other samples.
# Locate the companion toolbox locally, or fetch it for a standalone/Colab copy.
from pathlib import Path
import importlib.util
import os
import subprocess
import sys
working = Path.cwd()
candidates = [working, working / "python", working.parent / "python"]
python_dir = next((p for p in candidates if (p / "nt_toolbox").is_dir()), None)
if python_dir is None:
checkout = working / "numerical-tours-support"
if not checkout.exists():
subprocess.run(
[
"git",
"clone",
"--depth",
"1",
"--branch",
"master",
"https://github.com/gpeyre/numerical-tours.git",
str(checkout),
],
check=True,
)
python_dir = checkout / "python"
os.chdir(python_dir)
if str(python_dir) not in sys.path:
sys.path.insert(0, str(python_dir))
requirements = python_dir / "requirements.txt"
if any(
importlib.util.find_spec(name) is None
for name in [
"numpy",
"scipy",
"matplotlib",
"skimage",
"sklearn",
"pywt",
"ipywidgets",
"cvxpy",
"skfmm",
"autograd",
"progressbar",
"celer",
]
):
subprocess.run(
[sys.executable, "-m", "pip", "install", "-r", str(requirements)], check=True
)
import numpy as np
import matplotlib.pyplot as plt
np.random.seed(0)
plt.rcParams.update(
{
"figure.figsize": (8, 4),
"figure.dpi": 100,
"axes.spines.top": False,
"axes.spines.right": False,
"font.size": 11,
"image.cmap": "gray",
}
)
%matplotlib inline
$\newcommand{\dotp}[2]{\langle #1, #2 \rangle}$ $\newcommand{\enscond}[2]{\lbrace #1, #2 \rbrace}$ $\newcommand{\pd}[2]{ \frac{ \partial #1}{\partial #2} }$ $\newcommand{\umin}[1]{\underset{#1}{\min}\;}$ $\newcommand{\umax}[1]{\underset{#1}{\max}\;}$ $\newcommand{\uargmin}[1]{\underset{#1}{argmin}\;}$ $\newcommand{\norm}[1]{\|#1\|}$ $\newcommand{\abs}[1]{\left|#1\right|}$ $\newcommand{\choice}[1]{ \left\{ \begin{array}{l} #1 \end{array} \right. }$ $\newcommand{\pa}[1]{\left(#1\right)}$ $\newcommand{\diag}[1]{{diag}\left( #1 \right)}$ $\newcommand{\qandq}{\quad\text{and}\quad}$ $\newcommand{\qwhereq}{\quad\text{where}\quad}$ $\newcommand{\qifq}{ \quad \text{if} \quad }$ $\newcommand{\qarrq}{ \quad \Longrightarrow \quad }$ $\newcommand{\ZZ}{\mathbb{Z}}$ $\newcommand{\CC}{\mathbb{C}}$ $\newcommand{\RR}{\mathbb{R}}$ $\newcommand{\EE}{\mathbb{E}}$ $\newcommand{\Zz}{\mathcal{Z}}$ $\newcommand{\Ww}{\mathcal{W}}$ $\newcommand{\Vv}{\mathcal{V}}$ $\newcommand{\Nn}{\mathcal{N}}$ $\newcommand{\NN}{\mathcal{N}}$ $\newcommand{\Hh}{\mathcal{H}}$ $\newcommand{\Bb}{\mathcal{B}}$ $\newcommand{\Ee}{\mathcal{E}}$ $\newcommand{\Cc}{\mathcal{C}}$ $\newcommand{\Gg}{\mathcal{G}}$ $\newcommand{\Ss}{\mathcal{S}}$ $\newcommand{\Pp}{\mathcal{P}}$ $\newcommand{\Ff}{\mathcal{F}}$ $\newcommand{\Xx}{\mathcal{X}}$ $\newcommand{\Yy}{\mathcal{Y}}$ $\newcommand{\Mm}{\mathcal{M}}$ $\newcommand{\Ii}{\mathcal{I}}$ $\newcommand{\Dd}{\mathcal{D}}$ $\newcommand{\Ll}{\mathcal{L}}$ $\newcommand{\Tt}{\mathcal{T}}$ $\newcommand{\si}{\sigma}$ $\newcommand{\al}{\alpha}$ $\newcommand{\la}{\lambda}$ $\newcommand{\ga}{\gamma}$ $\newcommand{\Ga}{\Gamma}$ $\newcommand{\La}{\Lambda}$ $\newcommand{\Si}{\Sigma}$ $\newcommand{\be}{\beta}$ $\newcommand{\de}{\delta}$ $\newcommand{\De}{\Delta}$ $\newcommand{\phi}{\varphi}$ $\newcommand{\th}{\theta}$ $\newcommand{\om}{\omega}$ $\newcommand{\Om}{\Omega}$ $\newcommand{\eqdef}{\equiv}$
This numerical tour studies semi-discrete optimal transport, i.e. when one of the two measure is discrete.
The initial papers that proposed this approach are [Oliker89,Aurenhammer98]. We refer to [Mérigot11,Lévy15] for modern references and fast implementations.
This tour is not inteded to show efficient algorithm but only conveys the main underlying idea (c-transform, Laguerre cells, connexion to optimal quantization). In the Euclidean case, there exists efficient algorithm to compute Laguerre cells leveraging computational geometry algorithm for convex hulls [Aurenhammer87].
import numpy as np
import matplotlib.pyplot as plt
Dual OT and c-transforms¶
The primal Kantorovitch OT problem reads $$ W_c(\al,\be) = \umin{\pi} \enscond{\int_{\Xx \times \Yy} c(x,y) \text{d}\pi(x,y)}{ \pi_1=\al,\pi_2=\be }. $$ It dual is $$ W_c(\al,\be) = \umax{f,g} \enscond{ \int_\Xx f \text{d} \al + \int_\Yy g \text{d} \be }{ f(x)+g(y) \leq c(x,y) }. $$
We consider the case where $\al=\sum_i a_i \de_{x_i}$ is a discrete measure, so that the function $f(x)$ can be replaced by a vector $(f_i)_{i=1}^n \in \RR^n$. The optimal $g(y)$ function can the be replaced by the $c$-transform of $f$ $$ f^c(y) \eqdef \umin{i} c(x_i,y) - f_i. $$
The function to maximize is then $$ W_c(\al,\be) = \umax{f \in \RR^n} \Ee(f) \eqdef \sum_i f_i a_i + \int f^c(y) \text{d}\be(y). $$
Semi-discret via Gradient Ascent¶
We now implement a gradient ascent scheme for the maximization of $\Ee$. The evaluation of $\Ee$) can be computed via the introduction of the partition of the domain in Laguerre cells $$ \Yy = \bigcup_{i} L_i(f) \qwhereq L_i(f) \eqdef \enscond{y}{ \forall j, c(x_i,y) - f_i \leq c(x_j,y) - f_j }. $$ When $f=0$, this corrsponds to the partition in Voronoi cells.
One has that $\forall y \in L_i(f)$, $f^c(y) = c(x_i,y) - f_i$, i.e. $f^c$ is piecewise smooth according to this partition.
The grid for evaluation of the "continuous measure".
p = 300 # size of the image for sampling, m=p*p
t = np.linspace(0, 1, p)
[V, U] = np.meshgrid(t, t)
Y = np.concatenate((U.flatten()[None, :], V.flatten()[None, :]))
First measure, sums of Dirac masses $\al = \sum_{i=1}^n a_i \de_{x_i}$.
n = 30
X = (
0.5
+ 0.5j
+ np.exp(1j * np.pi / 4)
* 1
* (0.1 * (np.random.rand(1, n) - 0.5) + 1j * (np.random.rand(1, n) - 0.5))
)
X = np.concatenate((np.real(X), np.imag(X)))
a = np.ones(n) / n
Second measure $\be$, potentially a continuous one (i.e. with a density), mixture of Gaussians. Here we discretize $\beta = \sum_{j=1}^m b_j \de_{y_j}$ on a very fine grid.
def Gauss(mx, my, s):
return np.exp((-((U - mx) ** 2) - (V - my) ** 2) / (2 * s**2))
Mx = [0.6, 0.4] # means
My = [0.9, 0.1]
S = [0.07, 0.09] # variance
W = [0.5, 0.5] # weights
b = W[0] * Gauss(Mx[0], My[0], S[0]) + W[1] * Gauss(Mx[1], My[1], S[1])
b = b / np.sum(b.flatten())
Display the two measures.
Col = np.random.rand(n, 3)
plt.imshow(-b[::-1, :], extent=[0, 1, 0, 1], cmap="gray")
plt.scatter(X[1, :], X[0, :], s=30, c=0.8 * Col);
Initial potentials.
f = np.zeros(n)
compute Laguerre cells and c-transform
def distmat(x, y):
return (
np.sum(x**2, 0)[:, None] + np.sum(y**2, 0)[None, :] - 2 * x.transpose().dot(y)
)
D = distmat(Y, X) - f[:].transpose()
fC = np.min(D, axis=1)
I = np.reshape(np.argmin(D, axis=1), [p, p])
Dual value of the OT, $\dotp{f}{a}+\dotp{f^c}{\be}$.
OT = np.sum(f * a) + np.sum(fC * b.flatten())
print(OT)
Display the Laguerre call partition (here this is equal to the Vornoi diagram since $f=0$).
plt.imshow(I[::-1, :], extent=[0, 1, 0, 1])
plt.scatter(X[1, :], X[0, :], s=20, c="k")
plt.contour(t, t, I, np.linspace(-0.5, n - 0.5, n), colors="k")
plt.axis("off")
Where $\be$ has a density with respect to Lebesgue measure, then $\Ee$ is smooth, and its gradient reads $$ \nabla \Ee(f)_i = a_i - \int_{L_i(f)} \text{d}\be(x). $$
sum area captured
Worked example 1
Implement a gradient ascent $$ f \leftarrow f + \tau \nabla \Ee(f). $$ Experiment on the impact of $\tau$, display the evolution of the OT value $\Ee$ and of the Laguerre cells.
tau = 0.02 # step size
niter = 200 # iteration for the descent
q = 6 # number of displays
ndisp = np.unique(np.round(1 + (niter / 4 - 1) * np.linspace(0, 1, q) ** 2))
kdisp = 0
f = np.zeros(n)
E = np.zeros(niter)
for it in range(niter):
# compute Laguerre cells and c-transform
D = distmat(Y, X) - f[:].transpose()
fC = np.min(D, axis=1)
I = np.reshape(np.argmin(D, axis=1), [p, p])
E[it] = np.sum(f * a) + np.sum(fC * b.flatten())
# display
if (kdisp < len(ndisp)) and (ndisp[kdisp] == it):
plt.subplot(2, 3, kdisp + 1)
plt.imshow(I[::-1, :], extent=[0, 1, 0, 1])
plt.scatter(X[1, :], X[0, :], s=20, c="k")
plt.contour(t, t, I, np.linspace(-0.5, n - 0.5, n), colors="k")
plt.axis("off")
kdisp = kdisp + 1
# gradient
R = (I[:, :, None] == np.arange(0, n)[None, None, :]) * b[:, :, None]
nablaE = a - np.sum(R, axis=(0, 1)).flatten()
f = f + tau * nablaE
Display the evolution of the estimated OT distance.
plt.plot(E, "-")
Stochastic Optimization¶
The function $\Ee$ to minimize can be written as an expectation over a random variable $Y \sim \be$ $$ \Ee(f)=\EE(E(f,Y)) \qwhereq E(f,y) = \dotp{f}{a} + f^c(y). $$
As proposed in [Genevay16], one can thus use a stochastic gradient ascent scheme to minimize this function, at iteration $\ell$ $$ f \leftarrow f + \tau_\ell \nabla E(f,y_\ell) $$ where $y_\ell \sim Y$ is a sample drawn according to $\be$ and the step size $\tau_\ell \sim 1/\ell$ should decay at a carefully chosen rate.
The gradient of the integrated functional reads $$ \nabla E(f,y)_i = a - 1_{L_i(f)}(y), $$ where $1_A$ is the binary indicator function of a set $A$.
Initialize the algorithm.
f = np.zeros(n)
Draw the sample.
k = int(np.random.rand() < W[1]) # select one of the two Gaussian
y = np.array((S[k] * np.random.randn() + Mx[k], S[k] * np.random.randn() + My[k]))
Compute the randomized gradient: detect Laguerre cell where $y$ is.
R = np.sum(y**2) + np.sum(X**2, axis=0) - 2 * y.transpose().dot(X) - f[:]
i = np.argmin(R)
Randomized gradient.
a = np.ones(n) / n
nablaEy = a.copy()
nablaEy[i] = nablaEy[i] - 1
Worked example 2
Implement the stochastic gradient descent. Test various step size selection rule.
niter = 300
q = 6
ndisp = np.unique(np.round(1 + (niter / 2 - 1) * np.linspace(0, 1, q) ** 2))
kdisp = 0
E = np.zeros(niter)
for it in range(niter):
# sample
k = int(np.random.rand() < W[1]) # select one of the two Gaussian
y = np.array((S[k] * np.random.randn() + Mx[k], S[k] * np.random.randn() + My[k]))
# detect Laguerre cell where y is
R = np.sum(y**2) + np.sum(X**2, axis=0) - 2 * y.transpose().dot(X) - f[:]
i = np.argmin(R)
# gradient
nablaEy = a.copy()
nablaEy[i] = nablaEy[i] - 1
# gradient ascent
l0 = 10 # warmup phase.
tau = 0.1 / (1 + it / l0)
f = f + tau * nablaEy
# compute Laguerre cells and c-transform
D = distmat(Y, X) - f[:].transpose()
fC = np.min(D, axis=1)
I = np.reshape(np.argmin(D, axis=1), [p, p])
E[it] = np.sum(f * a) + np.sum(fC * b.flatten())
# display
if (kdisp < len(ndisp)) and (ndisp[kdisp] == it):
plt.subplot(2, 3, kdisp + 1)
plt.imshow(I[::-1, :], extent=[0, 1, 0, 1])
plt.scatter(X[1, :], X[0, :], s=20, c="k")
plt.contour(t, t, I, np.linspace(-0.5, n - 0.5, n), colors="k")
plt.axis("off")
kdisp = kdisp + 1
Display the evolution of the estimated OT distance (warning: recording this takes lot of time).
plt.plot(E)
Optimal Quantization and Lloyd Algorithm¶
We consider the following optimal quantization problem [Gruber02] $$ \umin{ (a_i)_i,(x_i)_i } W_c\pa{ \sum_i a_i \de_{x_i},\be }. $$ This minimization is convex in $a$, and writing down the optimality condition, one has that the associated dual potential should be $f=0$, which means that the associated optimal Laguerre cells should be Voronoi cells $L_i(0)=V_i(x)$ associated to the sampling locations $$ V_i(x) = \enscond{y}{ \forall j, c(x_i,y) \leq c(x_j,y) }. $$
This problem is tightly connected to semi-discrete OT, and this connexion and its implications are studied in [Canas12].
The minimization is non-convex with respect to the positions $x=(x_i)_i$ and one needs to solve $$ \umin{x} \Ff(x) \eqdef \sum_{i=1}^n \int_{V_i(x)} c(x_i,y) \text{d} \be(y). $$ For the sake of simplicity, we consider the case where $c(x,y)=\frac{1}{2}\norm{x-y}^2$.
The gradient reads $$ \nabla \Ff(x)_i = x_i \int_{V_i(x)} \text{d}\be - \int_{V_i(x)} y \text{d}\be(y). $$ The usual algorithm to compute stationary point of this energy is Lloyd's algorithm [Lloyd82], which iterate the fixed point $$ x_i \leftarrow \frac{ \int_{V_i(x)} y \text{d}\be(y) }{ \int_{V_i(x)} \text{d}\be }, $$ i.e. one replaces the centroids by the barycenter of the cells.
Intialize the centroids positions.
X1 = X.copy()
Compute the Voronoi cells $V_i(x)$.
D = D = distmat(Y, X1)
fC = np.min(D, axis=1)
I = np.reshape(np.argmin(D, axis=1), [p, p])
Update the centroids to the barycenters.
A = (I[:, :, None] == np.arange(0, n)[None, None, :]) * b[:, :, None]
B = (
(I[:, :, None] == np.arange(0, n)[None, None, :])
* b[:, :, None]
* (U[:, :, None] + 1j * V[:, :, None])
)
X1 = np.sum(B, axis=(0, 1)) / np.sum(A, axis=(0, 1))
X1 = np.concatenate((np.real(X1)[None, :], np.imag(X1)[None, :]))
Worked example 3
Implement Lloyd algortihm.
niter = 60
q = 6
ndisp = np.unique(np.round(1 + (niter / 4 - 1) * np.linspace(0, 1, q) ** 2))
kdisp = 0
E = np.zeros(niter)
X1 = X.copy()
for it in range(niter):
# compute Voronoi cells
D = D = distmat(Y, X1)
fC = np.min(D, axis=1)
I = np.reshape(np.argmin(D, axis=1), [p, p])
E[it] = np.sum(fC * b.flatten())
# display
if (kdisp < len(ndisp)) and (ndisp[kdisp] == it):
plt.subplot(2, 3, kdisp + 1)
plt.imshow(I[::-1, :], extent=[0, 1, 0, 1])
plt.scatter(X[1, :], X[0, :], s=20, c="k")
plt.contour(t, t, I, np.linspace(-0.5, n - 0.5, n), colors="k")
plt.axis("off")
kdisp = kdisp + 1
# update barycenter
A = (I[:, :, None] == np.arange(0, n)[None, None, :]) * b[:, :, None]
B = (
(I[:, :, None] == np.arange(0, n)[None, None, :])
* b[:, :, None]
* (U[:, :, None] + 1j * V[:, :, None])
)
X1 = np.sum(B, axis=(0, 1)) / np.sum(A, axis=(0, 1))
X1 = np.concatenate((np.real(X1)[None, :], np.imag(X1)[None, :]))
Display the evolution of the estimated OT distance.
plt.plot(E[1:-1])
References and further reading¶
[Oliker89] Vladimir Oliker and Laird D Prussner. On the numerical solution of the equation $\frac{\partial^2 z}{\partial x^2} \frac{\partial^2 z}{\partial y^2} - \pa{\frac{\partial^2 z}{\partial x\partial y}}^2 = f$ and its discretizations, I. Numerische Mathematik, 54(3):271-293, 1989.
[Aurenhammer98] Franz Aurenhammer, Friedrich Hoffmann and Boris Aronov. Minkowski-type theorems and least-squares clustering. Algorithmica, 20(1):61-76, 1998.
[Mérigot11] Quentin Mérigot. A multiscale approach to optimal transport. Comput. Graph. Forum, 30(5):1583-1592, 2011.
[Lévy15] Bruno Lévy. A numerical algorithm for l2 semi-discrete optimal transport in 3D. ESAIM: Mathematical Modelling and Numerical Analysis, 49(6):1693-1715, 2015.
[Aurenhammer87] Franz Aurenhammer. Power diagrams: properties, algorithms and applications. SIAM Journal on Computing, 16(1):78-96, 1987.
[Canas12] Guillermo Canas, Lorenzo Rosasco, Learning probability measures with respect to optimal transport metrics. In Advances in Neural Information Processing Systems, pp. 2492--2500, 2012.
[Gruber02] Peter M. Gruber. Optimum quantization and its applications. Adv. Math, 186:2004, 2002.
[Lloyd82] Stuart P. Lloyd, Least squares quantization in PCM, IEEE Transactions on Information Theory, 28 (2): 129-137, 1982.
Gabriel Peyré and Marco Cuturi. Computational Optimal Transport. 2019, Foundations and Trends in Machine Learning 11(5–6), 355–607. Transport plans, duality, entropic regularization, and applications.
Marco Cuturi. Sinkhorn Distances: Lightspeed Computation of Optimal Transport. 2013, NeurIPS. Entropic regularization and scalable matrix scaling.