{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# Gromov-Wasserstein example\n\n<div class=\"alert alert-info\"><h4>Note</h4><p>Example added in release: 0.8.0.</p></div>\n\nThis example is designed to show how to use the Gromov-Wasserstein distance\ncomputation in POT.\nWe first compare 3 solvers to estimate the distance based on\nConditional Gradient [24] or Sinkhorn projections [12, 51].\nThen we compare 2 stochastic solvers to estimate the distance with a lower\nnumerical cost [33].\n\n[12] Gabriel Peyr\u00e9, Marco Cuturi, and Justin Solomon (2016),\n\"Gromov-Wasserstein averaging of kernel and distance matrices\".\nInternational Conference on Machine Learning (ICML).\n\n[24] Vayer Titouan, Chapel Laetitia, Flamary R\u00e9mi, Tavenard Romain\nand Courty Nicolas\n\"Optimal Transport for structured data with application on graphs\"\nInternational Conference on Machine Learning (ICML). 2019.\n\n[33] Kerdoncuff T., Emonet R., Marc S. \"Sampled Gromov Wasserstein\",\nMachine Learning Journal (MJL), 2021.\n\n[51] Xu, H., Luo, D., Zha, H., & Carin, L. (2019).\n\"Gromov-wasserstein learning for graph matching and node embedding\".\nIn International Conference on Machine Learning (ICML), 2019.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# Author: Erwan Vautier <erwan.vautier@gmail.com>\n#         Nicolas Courty <ncourty@irisa.fr>\n#         C\u00e9dric Vincent-Cuaz <cedvincentcuaz@gmail.com>\n#         Tanguy Kerdoncuff <tanguy.kerdoncuff@laposte.net>\n#\n# License: MIT License\n\n# sphinx_gallery_thumbnail_number = 1\n\nimport scipy as sp\nimport numpy as np\nimport matplotlib.pylab as pl\nfrom mpl_toolkits.mplot3d import Axes3D  # noqa\nimport ot"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Sample two Gaussian distributions (2D and 3D)\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 = 30  # nb samples\n\nmu_s = np.array([0, 0])\ncov_s = np.array([[1, 0], [0, 1]])\n\nmu_t = np.array([4, 4, 4])\ncov_t = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]])\n\nnp.random.seed(0)\nxs = ot.datasets.make_2D_samples_gauss(n_samples, mu_s, cov_s)\nP = sp.linalg.sqrtm(cov_t)\nxt = np.random.randn(n_samples, 3).dot(P) + mu_t"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Plotting the distributions\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "fig = pl.figure(1)\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 distance kernels, normalize them and then display\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\nC1 /= C1.max()\nC2 /= C2.max()\n\npl.figure(2)\npl.subplot(121)\npl.imshow(C1)\npl.title(\"C1\")\n\npl.subplot(122)\npl.imshow(C2)\npl.title(\"C2\")\n\npl.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Compute Gromov-Wasserstein plans and distance\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "p = ot.unif(n_samples)\nq = ot.unif(n_samples)\n\n# Conditional Gradient algorithm\ngw0, log0 = ot.gromov.gromov_wasserstein(\n    C1, C2, p, q, \"square_loss\", verbose=True, log=True\n)\n\n# Proximal Point algorithm with Kullback-Leibler as proximal operator\ngw, log = ot.gromov.entropic_gromov_wasserstein(\n    C1, C2, p, q, \"square_loss\", epsilon=5e-4, solver=\"PPA\", log=True, verbose=True\n)\n\n# Projected Gradient algorithm with entropic regularization\ngwe, loge = ot.gromov.entropic_gromov_wasserstein(\n    C1, C2, p, q, \"square_loss\", epsilon=5e-4, solver=\"PGD\", log=True, verbose=True\n)\n\nprint(\n    \"Gromov-Wasserstein distance estimated with Conditional Gradient solver: \"\n    + str(log0[\"gw_dist\"])\n)\nprint(\n    \"Gromov-Wasserstein distance estimated with Proximal Point solver: \"\n    + str(log[\"gw_dist\"])\n)\nprint(\n    \"Entropic Gromov-Wasserstein distance estimated with Projected Gradient solver: \"\n    + str(loge[\"gw_dist\"])\n)\n\n# compute OT sparsity level\ngw0_sparsity = 100 * (gw0 == 0.0).astype(np.float64).sum() / (n_samples**2)\ngw_sparsity = 100 * (gw == 0.0).astype(np.float64).sum() / (n_samples**2)\ngwe_sparsity = 100 * (gwe == 0.0).astype(np.float64).sum() / (n_samples**2)\n\n# Methods using Sinkhorn projections tend to produce feasibility errors on the\n# marginal constraints\n\nerr0 = np.linalg.norm(gw0.sum(1) - p) + np.linalg.norm(gw0.sum(0) - q)\nerr = np.linalg.norm(gw.sum(1) - p) + np.linalg.norm(gw.sum(0) - q)\nerre = np.linalg.norm(gwe.sum(1) - p) + np.linalg.norm(gwe.sum(0) - q)\n\npl.figure(3, (10, 6))\ncmap = \"Blues\"\nfontsize = 12\npl.subplot(131)\npl.imshow(gw0, cmap=cmap)\npl.title(\n    \"(CG algo) GW=%s \\n  \\n OT sparsity=%s \\n feasibility error=%s\"\n    % (\n        np.round(log0[\"gw_dist\"], 4),\n        str(np.round(gw0_sparsity, 2)) + \" %\",\n        np.round(np.round(err0, 4)),\n    ),\n    fontsize=fontsize,\n)\n\npl.subplot(132)\npl.imshow(gw, cmap=cmap)\npl.title(\n    \"(PP algo) GW=%s \\n  \\n OT sparsity=%s \\nfeasibility error=%s\"\n    % (\n        np.round(log[\"gw_dist\"], 4),\n        str(np.round(gw_sparsity, 2)) + \" %\",\n        np.round(err, 4),\n    ),\n    fontsize=fontsize,\n)\n\npl.subplot(133)\npl.imshow(gwe, cmap=cmap)\npl.title(\n    \"Entropic GW=%s \\n  \\n OT sparsity=%s \\nfeasibility error=%s\"\n    % (\n        np.round(loge[\"gw_dist\"], 4),\n        str(np.round(gwe_sparsity, 2)) + \" %\",\n        np.round(erre, 4),\n    ),\n    fontsize=fontsize,\n)\n\npl.tight_layout()\npl.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Compute GW with scalable stochastic methods with any loss function\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "def loss(x, y):\n    return np.abs(x - y)\n\n\npgw, plog = ot.gromov.pointwise_gromov_wasserstein(\n    C1, C2, p, q, loss, max_iter=100, log=True\n)\n\nsgw, slog = ot.gromov.sampled_gromov_wasserstein(\n    C1, C2, p, q, loss, epsilon=0.1, max_iter=100, log=True\n)\n\nprint(\n    \"Pointwise Gromov-Wasserstein distance estimated: \" + str(plog[\"gw_dist_estimated\"])\n)\nprint(\"Variance estimated: \" + str(plog[\"gw_dist_std\"]))\nprint(\"Sampled Gromov-Wasserstein distance: \" + str(slog[\"gw_dist_estimated\"]))\nprint(\"Variance estimated: \" + str(slog[\"gw_dist_std\"]))\n\n\npl.figure(4, (10, 5))\n\npl.subplot(121)\npl.imshow(pgw.toarray(), cmap=cmap)\npl.title(\"Pointwise Gromov Wasserstein\")\n\npl.subplot(122)\npl.imshow(sgw, cmap=cmap)\npl.title(\"Sampled Gromov Wasserstein\")\n\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
}