{ "cells": [ { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "%matplotlib inline" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "\n# Partial Wasserstein and Gromov-Wasserstein example\n\n\nThis example is designed to show how to use the Partial (Gromov-)Wassertsein\ndistance computation in POT.\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "# Author: Laetitia Chapel \n# License: MIT License\n\n# necessary for 3d plot even if not used\nfrom mpl_toolkits.mplot3d import Axes3D # noqa\nimport scipy as sp\nimport numpy as np\nimport matplotlib.pylab as pl\nimport ot" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Sample two 2D Gaussian distributions and plot them\n--------------------------------------------------\n\nFor demonstration purpose, we sample two Gaussian distributions in 2-d\nspaces and add some random noise.\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "n_samples = 20 # nb samples (gaussian)\nn_noise = 20 # nb of samples (noise)\n\nmu = np.array([0, 0])\ncov = np.array([[1, 0], [0, 2]])\n\nxs = ot.datasets.make_2D_samples_gauss(n_samples, mu, cov)\nxs = np.append(xs, (np.random.rand(n_noise, 2) + 1) * 4).reshape((-1, 2))\nxt = ot.datasets.make_2D_samples_gauss(n_samples, mu, cov)\nxt = np.append(xt, (np.random.rand(n_noise, 2) + 1) * -3).reshape((-1, 2))\n\nM = sp.spatial.distance.cdist(xs, xt)\n\nfig = pl.figure()\nax1 = fig.add_subplot(131)\nax1.plot(xs[:, 0], xs[:, 1], '+b', label='Source samples')\nax2 = fig.add_subplot(132)\nax2.scatter(xt[:, 0], xt[:, 1], color='r')\nax3 = fig.add_subplot(133)\nax3.imshow(M)\npl.show()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Compute partial Wasserstein plans and distance,\nby transporting 50% of the mass\n----------------------------------------------\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "p = ot.unif(n_samples + n_noise)\nq = ot.unif(n_samples + n_noise)\n\nw0, log0 = ot.partial.partial_wasserstein(p, q, M, m=0.5, log=True)\nw, log = ot.partial.entropic_partial_wasserstein(p, q, M, reg=0.1, m=0.5,\n log=True)\n\nprint('Partial Wasserstein distance (m = 0.5): ' + str(log0['partial_w_dist']))\nprint('Entropic partial Wasserstein distance (m = 0.5): ' +\n str(log['partial_w_dist']))\n\npl.figure(1, (10, 5))\npl.subplot(1, 2, 1)\npl.imshow(w0, cmap='jet')\npl.title('Partial Wasserstein')\npl.subplot(1, 2, 2)\npl.imshow(w, cmap='jet')\npl.title('Entropic partial Wasserstein')\npl.show()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Sample one 2D and 3D Gaussian distributions and plot them\n---------------------------------------------------------\n\nThe Gromov-Wasserstein distance allows to compute distances with samples that\ndo not belong to the same metric space. For demonstration purpose, we sample\ntwo Gaussian distributions in 2- and 3-dimensional spaces.\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "n_samples = 20 # nb samples\nn_noise = 10 # nb of samples (noise)\n\np = ot.unif(n_samples + n_noise)\nq = ot.unif(n_samples + n_noise)\n\nmu_s = np.array([0, 0])\ncov_s = np.array([[1, 0], [0, 1]])\n\nmu_t = np.array([0, 0, 0])\ncov_t = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]])\n\n\nxs = ot.datasets.make_2D_samples_gauss(n_samples, mu_s, cov_s)\nxs = np.concatenate((xs, ((np.random.rand(n_noise, 2) + 1) * 4)), axis=0)\nP = sp.linalg.sqrtm(cov_t)\nxt = np.random.randn(n_samples, 3).dot(P) + mu_t\nxt = np.concatenate((xt, ((np.random.rand(n_noise, 3) + 1) * 10)), axis=0)\n\nfig = pl.figure()\nax1 = fig.add_subplot(121)\nax1.plot(xs[:, 0], xs[:, 1], '+b', label='Source samples')\nax2 = fig.add_subplot(122, projection='3d')\nax2.scatter(xt[:, 0], xt[:, 1], xt[:, 2], color='r')\npl.show()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Compute partial Gromov-Wasserstein plans and distance,\nby transporting 100% and 2/3 of the mass\n-----------------------------------------------------\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "C1 = sp.spatial.distance.cdist(xs, xs)\nC2 = sp.spatial.distance.cdist(xt, xt)\n\nprint('-----m = 1')\nm = 1\nres0, log0 = ot.partial.partial_gromov_wasserstein(C1, C2, p, q, m=m,\n log=True)\nres, log = ot.partial.entropic_partial_gromov_wasserstein(C1, C2, p, q, 10,\n m=m, log=True)\n\nprint('Partial Wasserstein distance (m = 1): ' + str(log0['partial_gw_dist']))\nprint('Entropic partial Wasserstein distance (m = 1): ' +\n str(log['partial_gw_dist']))\n\npl.figure(1, (10, 5))\npl.title(\"mass to be transported m = 1\")\npl.subplot(1, 2, 1)\npl.imshow(res0, cmap='jet')\npl.title('Partial Wasserstein')\npl.subplot(1, 2, 2)\npl.imshow(res, cmap='jet')\npl.title('Entropic partial Wasserstein')\npl.show()\n\nprint('-----m = 2/3')\nm = 2 / 3\nres0, log0 = ot.partial.partial_gromov_wasserstein(C1, C2, p, q, m=m, log=True)\nres, log = ot.partial.entropic_partial_gromov_wasserstein(C1, C2, p, q, 10,\n m=m, log=True)\n\nprint('Partial Wasserstein distance (m = 2/3): ' +\n str(log0['partial_gw_dist']))\nprint('Entropic partial Wasserstein distance (m = 2/3): ' +\n str(log['partial_gw_dist']))\n\npl.figure(1, (10, 5))\npl.title(\"mass to be transported m = 2/3\")\npl.subplot(1, 2, 1)\npl.imshow(res0, cmap='jet')\npl.title('Partial Wasserstein')\npl.subplot(1, 2, 2)\npl.imshow(res, cmap='jet')\npl.title('Entropic partial Wasserstein')\npl.show()" ] } ], "metadata": { "kernelspec": { "display_name": "Python 3", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.6.9" } }, "nbformat": 4, "nbformat_minor": 0 }