{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# Partial Wasserstein and Gromov-Wasserstein example\n\n<div class=\"alert alert-info\"><h4>Note</h4><p>Example added in release: 0.7.0.</p></div>\n\nThis example is designed to show how to use the Partial (Gromov-)Wasserstein\ndistance computation in POT [29].\n\n[29] Chapel, L., Alaya, M., Gasso, G. (2020). \"Partial Optimal\nTransport with Applications on Positive-Unlabeled Learning\". NeurIPS.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# Author: Laetitia Chapel <laetitia.chapel@irisa.fr>\n# License: MIT License\n\n# sphinx_gallery_thumbnail_number = 2\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\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\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, log=True)\n\nprint(\"Partial Wasserstein distance (m = 0.5): \" + str(log0[\"partial_w_dist\"]))\nprint(\"Entropic partial Wasserstein distance (m = 0.5): \" + str(log[\"partial_w_dist\"]))\n\npl.figure(1, (10, 5))\npl.subplot(1, 2, 1)\npl.imshow(w0, cmap=\"gray_r\")\npl.title(\"Partial Wasserstein\")\npl.subplot(1, 2, 2)\npl.imshow(w, cmap=\"gray_r\")\npl.title(\"Entropic partial Wasserstein\")\npl.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Sample one 2D and 3D Gaussian distributions and plot them\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\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\n# transport 100% of the mass\nprint(\"------m = 1\")\nm = 1\nres0, log0 = ot.gromov.partial_gromov_wasserstein(C1, C2, p, q, m=m, log=True)\nres, log = ot.gromov.entropic_partial_gromov_wasserstein(\n    C1, C2, p, q, 10, m=m, log=True, verbose=True\n)\n\nprint(\"Wasserstein distance (m = 1): \" + str(log0[\"partial_gw_dist\"]))\nprint(\"Entropic Wasserstein distance (m = 1): \" + str(log[\"partial_gw_dist\"]))\n\npl.figure(1, (10, 5))\npl.title(\"mass to be transported m = 1\")\npl.axis(\"off\")\npl.subplot(1, 2, 1)\npl.imshow(res0, cmap=\"gray_r\")\npl.title(\"Gromov-Wasserstein\")\npl.subplot(1, 2, 2)\npl.imshow(res, cmap=\"gray_r\")\npl.title(\"Entropic Gromov-Wasserstein\")\npl.show()\n\n# transport 2/3 of the mass\nprint(\"------m = 2/3\")\nm = 2 / 3\nres0, log0 = ot.gromov.partial_gromov_wasserstein(\n    C1, C2, p, q, m=m, log=True, verbose=True\n)\nres, log = ot.gromov.entropic_partial_gromov_wasserstein(\n    C1, C2, p, q, 10, m=m, log=True, verbose=True\n)\n\nprint(\"Partial Wasserstein distance (m = 2/3): \" + str(log0[\"partial_gw_dist\"]))\nprint(\"Entropic partial Wasserstein distance (m = 2/3): \" + str(log[\"partial_gw_dist\"]))\n\npl.figure(2, (10, 5))\npl.title(\"mass to be transported m = 2/3\")\npl.axis(\"off\")\npl.subplot(1, 2, 1)\npl.imshow(res0, cmap=\"gray_r\")\npl.title(\"Partial Gromov-Wasserstein\")\npl.subplot(1, 2, 2)\npl.imshow(res, cmap=\"gray_r\")\npl.title(\"Entropic partial Gromov-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.10.18"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}