In [1]:
%matplotlib inline

Gromov-Wasserstein example

This example is designed to show how to use the Gromov-Wassertsein distance computation in POT.


In [2]:
# Author: Erwan Vautier <erwan.vautier@gmail.com>
#         Nicolas Courty <ncourty@irisa.fr>
#
# License: MIT License

import scipy as sp
import numpy as np
import matplotlib.pylab as pl
from mpl_toolkits.mplot3d import Axes3D  # noqa
import ot

Sample two Gaussian distributions (2D and 3D)

The Gromov-Wasserstein distance allows to compute distances with samples that do not belong to the same metric space. For demonstration purpose, we sample two Gaussian distributions in 2- and 3-dimensional spaces.


In [3]:
n_samples = 30  # nb samples

mu_s = np.array([0, 0])
cov_s = np.array([[1, 0], [0, 1]])

mu_t = np.array([4, 4, 4])
cov_t = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]])


xs = ot.datasets.get_2D_samples_gauss(n_samples, mu_s, cov_s)
P = sp.linalg.sqrtm(cov_t)
xt = np.random.randn(n_samples, 3).dot(P) + mu_t

Plotting the distributions


In [4]:
fig = pl.figure()
ax1 = fig.add_subplot(121)
ax1.plot(xs[:, 0], xs[:, 1], '+b', label='Source samples')
ax2 = fig.add_subplot(122, projection='3d')
ax2.scatter(xt[:, 0], xt[:, 1], xt[:, 2], color='r')
pl.show()


Compute distance kernels, normalize them and then display


In [5]:
C1 = sp.spatial.distance.cdist(xs, xs)
C2 = sp.spatial.distance.cdist(xt, xt)

C1 /= C1.max()
C2 /= C2.max()

pl.figure()
pl.subplot(121)
pl.imshow(C1)
pl.subplot(122)
pl.imshow(C2)
pl.show()


Compute Gromov-Wasserstein plans and distance


In [6]:
p = ot.unif(n_samples)
q = ot.unif(n_samples)

gw0, log0 = ot.gromov.gromov_wasserstein(
    C1, C2, p, q, 'square_loss', verbose=True, log=True)

gw, log = ot.gromov.entropic_gromov_wasserstein(
    C1, C2, p, q, 'square_loss', epsilon=5e-4, log=True, verbose=True)


print('Gromov-Wasserstein distances: ' + str(log0['gw_dist']))
print('Entropic Gromov-Wasserstein distances: ' + str(log['gw_dist']))


pl.figure(1, (10, 5))

pl.subplot(1, 2, 1)
pl.imshow(gw0, cmap='jet')
pl.title('Gromov Wasserstein')

pl.subplot(1, 2, 2)
pl.imshow(gw, cmap='jet')
pl.title('Entropic Gromov Wasserstein')

pl.show()


It.  |Loss        |Delta loss
--------------------------------
    0|3.731009e-02|0.000000e+00
    1|1.846414e-02|-1.020678e+00
    2|1.752056e-02|-5.385587e-02
    3|1.470479e-02|-1.914863e-01
    4|1.371582e-02|-7.210441e-02
    5|1.263823e-02|-8.526431e-02
    6|1.190590e-02|-6.151022e-02
    7|1.014664e-02|-1.733837e-01
    8|1.011396e-02|-3.230516e-03
    9|1.011363e-02|-3.245532e-05
   10|1.011363e-02|-3.245682e-07
   11|1.011363e-02|-3.245683e-09
   12|1.011363e-02|-3.245598e-11
It.  |Err         
-------------------
    0|7.847531e-02|
   10|2.424218e-03|
   20|4.250670e-02|
   30|5.679949e-05|
   40|1.219762e-08|
   50|1.121975e-11|
Gromov-Wasserstein distances: 0.01011363084946804
Entropic Gromov-Wasserstein distances: 0.0063386519916289385