<jupyter_start><jupyter_text>Introduction to Optimal Transport with PythonThis example gives an introduction on how to use Optimal Transport in Python.<jupyter_code># Author: <NAME>, <NAME>, <NAME>
#
# License: MIT License
# sphinx_gallery_thumbnail_number = 1<jupyter_output><empty_output><jupyter_text>POT Python Optimal Transport Toolbox POT installation* Install with pip::    pip install pot* Install with conda::    conda install -c conda-forge pot Import the toolbox<jupyter_code>import numpy as np  # always need it
import pylab as pl  # do the plots

import ot  # ot

import time<jupyter_output><empty_output><jupyter_text>Getting helpOnline  documentation : ``_Or inline help:<jupyter_code>help(ot.dist)<jupyter_output><empty_output><jupyter_text>First OT ProblemWe will solve the Bakery/Cafés problem of transporting croissants from anumber of Bakeries to Cafés in a City (in this case Manhattan). We did aquick google map search in Manhattan for bakeries and Cafés:We extracted from this search their positions and generated fictionalproduction and sale number (that both sum to the same value).We have acess to the position of Bakeries ``bakery_pos`` and theirrespective production ``bakery_prod`` which describe the sourcedistribution. The Cafés where the croissants are sold are defined also bytheir position ``cafe_pos`` and ``cafe_prod``, and describe the targetdistribution. For fun we also provide amap ``Imap`` that will illustrate the position of these shops in the city.Now we load the data<jupyter_code>data = np.load('../data/manhattan.npz')

bakery_pos = data['bakery_pos']
bakery_prod = data['bakery_prod']
cafe_pos = data['cafe_pos']
cafe_prod = data['cafe_prod']
Imap = data['Imap']

print('Bakery production: {}'.format(bakery_prod))
print('Cafe sale: {}'.format(cafe_prod))
print('Total croissants : {}'.format(cafe_prod.sum()))<jupyter_output><empty_output><jupyter_text>Plotting bakeries in the cityNext we plot the position of the bakeries and cafés on the map. The size ofthe circle is proportional to their production.<jupyter_code>pl.figure(1, (7, 6))
pl.clf()
pl.imshow(Imap, interpolation='bilinear')  # plot the map
pl.scatter(bakery_pos[:, 0], bakery_pos[:, 1], s=bakery_prod, c='r', ec='k', label='Bakeries')
pl.scatter(cafe_pos[:, 0], cafe_pos[:, 1], s=cafe_prod, c='b', ec='k', label='Cafés')
pl.legend()
pl.title('Manhattan Bakeries and Cafés')<jupyter_output><empty_output><jupyter_text>Cost matrixWe can now compute the cost matrix between the bakeries and the cafés, whichwill be the transport cost matrix. This can be done using the`ot.dist `_ function thatdefaults to squared Euclidean distance but can return other things such ascityblock (or Manhattan distance).<jupyter_code>C = ot.dist(bakery_pos, cafe_pos)

labels = [str(i) for i in range(len(bakery_prod))]
f = pl.figure(2, (14, 7))
pl.clf()
pl.subplot(121)
pl.imshow(Imap, interpolation='bilinear')  # plot the map
for i in range(len(cafe_pos)):
    pl.text(cafe_pos[i, 0], cafe_pos[i, 1], labels[i], color='b',
            fontsize=14, fontweight='bold', ha='center', va='center')
for i in range(len(bakery_pos)):
    pl.text(bakery_pos[i, 0], bakery_pos[i, 1], labels[i], color='r',
            fontsize=14, fontweight='bold', ha='center', va='center')
pl.title('Manhattan Bakeries and Cafés')

ax = pl.subplot(122)
im = pl.imshow(C, cmap="coolwarm")
pl.title('Cost matrix')
cbar = pl.colorbar(im, ax=ax, shrink=0.5, use_gridspec=True)
cbar.ax.set_ylabel("cost", rotation=-90, va="bottom")

pl.xlabel('Cafés')
pl.ylabel('Bakeries')
pl.tight_layout()<jupyter_output><empty_output><jupyter_text>The red cells in the matrix image show the bakeries and cafés that arefurther away, and thus more costly to transport from one to the other, whilethe blue ones show those that are very close to each other, with respect tothe squared Euclidean distance.  Solving the OT problem with `ot.emd `_<jupyter_code>start = time.time()
ot_emd = ot.emd(bakery_prod, cafe_prod, C)
time_emd = time.time() - start<jupyter_output><empty_output><jupyter_text>The function returns the transport matrix, which we can then visualize (next section).  Transportation plan vizualizationA good vizualization of the OT matrix in the 2D plane is to denote thetransportation of mass between a Bakery and a Café by a line. This can easilybe done with a double ``for`` loop.In order to make it more interpretable one can also use the ``alpha``parameter of plot and set it to ``alpha=G[i,j]/G.max()``.<jupyter_code># Plot the matrix and the map
f = pl.figure(3, (14, 7))
pl.clf()
pl.subplot(121)
pl.imshow(Imap, interpolation='bilinear')  # plot the map
for i in range(len(bakery_pos)):
    for j in range(len(cafe_pos)):
        pl.plot([bakery_pos[i, 0], cafe_pos[j, 0]], [bakery_pos[i, 1], cafe_pos[j, 1]],
                '-k', lw=3. * ot_emd[i, j] / ot_emd.max())
for i in range(len(cafe_pos)):
    pl.text(cafe_pos[i, 0], cafe_pos[i, 1], labels[i], color='b', fontsize=14,
            fontweight='bold', ha='center', va='center')
for i in range(len(bakery_pos)):
    pl.text(bakery_pos[i, 0], bakery_pos[i, 1], labels[i], color='r', fontsize=14,
            fontweight='bold', ha='center', va='center')
pl.title('Manhattan Bakeries and Cafés')

ax = pl.subplot(122)
im = pl.imshow(ot_emd)
for i in range(len(bakery_prod)):
    for j in range(len(cafe_prod)):
        text = ax.text(j, i, '{0:g}'.format(ot_emd[i, j]),
                       ha="center", va="center", color="w")
pl.title('Transport matrix')

pl.xlabel('Cafés')
pl.ylabel('Bakeries')
pl.tight_layout()<jupyter_output><empty_output><jupyter_text>The transport matrix gives the number of croissants that can be transportedfrom each bakery to each café. We can see that the bakeries only need totransport croissants to one or two cafés, the transport matrix being verysparse.  OT loss and dual variablesThe resulting wasserstein loss loss is of the form:\begin{align}W=\sum_{i,j}\gamma_{i,j}C_{i,j}\end{align}where $\gamma$ is the optimal transport matrix.<jupyter_code>W = np.sum(ot_emd * C)
print('Wasserstein loss (EMD) = {0:.2f}'.format(W))<jupyter_output><empty_output><jupyter_text>Regularized OT with SinkhornThe Sinkhorn algorithm is very simple to code. You can implement it directlyusing the following pseudo-codeIn this algorithm, $\oslash$ corresponds to the element-wise division.An alternative is to use the POT toolbox with`ot.sinkhorn `_Be careful of numerical problems. A good pre-processing for Sinkhorn is todivide the cost matrix ``C`` by its maximum value.  Algorithm<jupyter_code># Compute Sinkhorn transport matrix from algorithm
reg = 0.1
K = np.exp(-C / C.max() / reg)
nit = 100
u = np.ones((len(bakery_prod), ))
for i in range(1, nit):
    v = cafe_prod / np.dot(K.T, u)
    u = bakery_prod / (np.dot(K, v))
ot_sink_algo = np.atleast_2d(u).T * (K * v.T)  # Equivalent to np.dot(np.diag(u), np.dot(K, np.diag(v)))

# Compute Sinkhorn transport matrix with POT
ot_sinkhorn = ot.sinkhorn(bakery_prod, cafe_prod, reg=reg, M=C / C.max())

# Difference between the 2
print('Difference between algo and ot.sinkhorn = {0:.2g}'.format(np.sum(np.power(ot_sink_algo - ot_sinkhorn, 2))))<jupyter_output><empty_output><jupyter_text>Plot the matrix and the map<jupyter_code>print('Min. of Sinkhorn\'s transport matrix = {0:.2g}'.format(np.min(ot_sinkhorn)))

f = pl.figure(4, (13, 6))
pl.clf()
pl.subplot(121)
pl.imshow(Imap, interpolation='bilinear')  # plot the map
for i in range(len(bakery_pos)):
    for j in range(len(cafe_pos)):
        pl.plot([bakery_pos[i, 0], cafe_pos[j, 0]],
                [bakery_pos[i, 1], cafe_pos[j, 1]],
                '-k', lw=3. * ot_sinkhorn[i, j] / ot_sinkhorn.max())
for i in range(len(cafe_pos)):
    pl.text(cafe_pos[i, 0], cafe_pos[i, 1], labels[i], color='b',
            fontsize=14, fontweight='bold', ha='center', va='center')
for i in range(len(bakery_pos)):
    pl.text(bakery_pos[i, 0], bakery_pos[i, 1], labels[i], color='r',
            fontsize=14, fontweight='bold', ha='center', va='center')
pl.title('Manhattan Bakeries and Cafés')

ax = pl.subplot(122)
im = pl.imshow(ot_sinkhorn)
for i in range(len(bakery_prod)):
    for j in range(len(cafe_prod)):
        text = ax.text(j, i, np.round(ot_sinkhorn[i, j], 1),
                       ha="center", va="center", color="w")
pl.title('Transport matrix')

pl.xlabel('Cafés')
pl.ylabel('Bakeries')
pl.tight_layout()<jupyter_output><empty_output><jupyter_text>We notice right away that the matrix is not sparse at all with Sinkhorn,each bakery delivering croissants to all 5 cafés with that solution. Also,this solution gives a transport with fractions, which does not make sensein the case of croissants. This was not the case with EMD.  Varying the regularization parameter in Sinkhorn<jupyter_code>reg_parameter = np.logspace(-3, 0, 20)
W_sinkhorn_reg = np.zeros((len(reg_parameter), ))
time_sinkhorn_reg = np.zeros((len(reg_parameter), ))

f = pl.figure(5, (14, 5))
pl.clf()
max_ot = 100  # plot matrices with the same colorbar
for k in range(len(reg_parameter)):
    start = time.time()
    ot_sinkhorn = ot.sinkhorn(bakery_prod, cafe_prod, reg=reg_parameter[k], M=C / C.max())
    time_sinkhorn_reg[k] = time.time() - start

    if k % 4 == 0 and k > 0:  # we only plot a few
        ax = pl.subplot(1, 5, k / 4)
        im = pl.imshow(ot_sinkhorn, vmin=0, vmax=max_ot)
        pl.title('reg={0:.2g}'.format(reg_parameter[k]))
        pl.xlabel('Cafés')
        pl.ylabel('Bakeries')

    # Compute the Wasserstein loss for Sinkhorn, and compare with EMD
    W_sinkhorn_reg[k] = np.sum(ot_sinkhorn * C)
pl.tight_layout()<jupyter_output><empty_output><jupyter_text>This series of graph shows that the solution of Sinkhorn starts with somethingvery similar to EMD (although not sparse) for very small values of theregularization parameter, and tends to a more uniform solution as theregularization parameter increases.  Wasserstein loss and computational time<jupyter_code># Plot the matrix and the map
f = pl.figure(6, (4, 4))
pl.clf()
pl.title("Comparison between Sinkhorn and EMD")

pl.plot(reg_parameter, W_sinkhorn_reg, 'o', label="Sinkhorn")
XLim = pl.xlim()
pl.plot(XLim, [W, W], '--k', label="EMD")
pl.legend()
pl.xlabel("reg")
pl.ylabel("Wasserstein loss")<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.13125304729400292, 'fraction_numerical': 0.011506582155046319, 'mean_word_length': 5.0329411764705885, 'pattern_counts': {'":': 0, '<': 56, '<?xml version=': 0, '>': 57, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '13', 'dir': 'jupyter-structured-clean-dedup', 'id': '316931', 'license': 'MIT', 'n_tokens_mistral': 3756, 'n_tokens_neox': 3583, 'n_words': 1064, 'path': '_downloads/f7942777fc8bc11618d8908da9b54edc/plot_Intro_OT.ipynb', 'repo_name': 'FlopsKa/pythonot.github.io', 'size': '10280'}