Gromov-Wasserstein distance

Author: Nicolas Courty

README.md

@@ -16,7 +16,7 @@ It provides the following solvers:
 * Conditional gradient [6] and Generalized conditional gradient for regularized OT [7].
 * Joint OT matrix and mapping estimation [8].
 * Wasserstein Discriminant Analysis [11] (requires autograd + pymanopt).
-
+* Gromov-Wasserstein distances [12]
 
 Some demonstrations (both in Python and Jupyter Notebook format) are available in the examples folder.
 
@@ -182,3 +182,5 @@ You can also post bug reports and feature requests in Github issues. Make sure t
 [10] Chizat, L., Peyré, G., Schmitzer, B., & Vialard, F. X. (2016). [Scaling algorithms for unbalanced transport problems](https://arxiv.org/pdf/1607.05816.pdf). arXiv preprint arXiv:1607.05816.
 
 [11] Flamary, R., Cuturi, M., Courty, N., & Rakotomamonjy, A. (2016). [Wasserstein Discriminant Analysis](https://arxiv.org/pdf/1608.08063.pdf). arXiv preprint arXiv:1608.08063.
+
+[12] Peyré, Gabriel, Marco Cuturi, and Justin Solomon, [Gromov-Wasserstein averaging of kernel and distance matrices](http://proceedings.mlr.press/v48/peyre16.html)  International Conference on Machine Learning (ICML). 2016.

examples/plot_gromov.py

@@ -0,0 +1,98 @@
+# -*- coding: utf-8 -*-
+"""
+====================
+Gromov-Wasserstein example
+====================
+
+This example is designed to show how to use the Gromov-Wassertsein distance 
+computation in POT. 
+
+
+"""
+
+# Author: Erwan Vautier <erwan.vautier@gmail.com>
+#         Nicolas Courty <ncourty@irisa.fr>
+#
+# License: MIT License
+
+import scipy as sp
+import numpy as np
+
+import ot
+import matplotlib.pylab as pl
+from mpl_toolkits.mplot3d import Axes3D
+
+
+
+"""
+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. 
+
+"""
+n=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,mu_s,cov_s)
+P=sp.linalg.sqrtm(cov_t)
+xt= np.random.randn(n,3).dot(P)+mu_t
+
+
+
+"""
+Plotting the distributions
+====================
+"""
+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
+====================
+"""
+
+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
+====================
+"""
+
+p=ot.unif(n)
+q=ot.unif(n)
+
+gw=ot.gromov_wasserstein(C1,C2,p,q,'square_loss',epsilon=5e-4)
+gw_dist=ot.gromov_wasserstein2(C1,C2,p,q,'square_loss',epsilon=5e-4)
+
+print('Gromov-Wasserstein distances between the distribution: '+str(gw_dist))
+
+pl.figure()
+pl.imshow(gw,cmap='jet')
+pl.colorbar()
+pl.show()
+

ot/__init__.py

@@ -5,6 +5,7 @@
 """
 
 # Author: Remi Flamary <remi.flamary@unice.fr>
+#         Nicolas Courty <ncourty@irisa.fr>
 #
 # License: MIT License
 
@@ -17,11 +18,13 @@
 from . import datasets
 from . import plot
 from . import da
+from . import gromov
 
 # OT functions
 from .lp import emd, emd2
 from .bregman import sinkhorn, sinkhorn2, barycenter
 from .da import sinkhorn_lpl1_mm
+from .gromov import gromov_wasserstein, gromov_wasserstein2
 
 # utils functions
 from .utils import dist, unif, tic, toc, toq
@@ -30,4 +33,5 @@
 
 __all__ = ["emd", "emd2", "sinkhorn", "sinkhorn2", "utils", 'datasets',
            'bregman', 'lp', 'plot', 'tic', 'toc', 'toq',
-           'dist', 'unif', 'barycenter', 'sinkhorn_lpl1_mm', 'da', 'optim']
+           'dist', 'unif', 'barycenter', 'sinkhorn_lpl1_mm', 'da', 'optim', 
+           'gromov_wasserstein','gromov_wasserstein2']