{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# Optimal Transport between empirical distributions\n\nIllustration of optimal transport between distributions in 2D that are weighted\nsum of Diracs. The OT matrix is plotted with the samples.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# Author: Remi Flamary <remi.flamary@unice.fr>\n#         Kilian Fatras <kilian.fatras@irisa.fr>\n#\n# License: MIT License\n\n# sphinx_gallery_thumbnail_number = 4\n\nimport numpy as np\nimport matplotlib.pylab as pl\nimport ot\nimport ot.plot"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Generate data\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "n = 50  # nb samples\n\nmu_s = np.array([0, 0])\ncov_s = np.array([[1, 0], [0, 1]])\n\nmu_t = np.array([4, 4])\ncov_t = np.array([[1, -0.8], [-0.8, 1]])\n\nxs = ot.datasets.make_2D_samples_gauss(n, mu_s, cov_s)\nxt = ot.datasets.make_2D_samples_gauss(n, mu_t, cov_t)\n\na, b = np.ones((n,)) / n, np.ones((n,)) / n  # uniform distribution on samples\n\n# loss matrix\nM = ot.dist(xs, xt)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Plot data\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "pl.figure(1)\npl.plot(xs[:, 0], xs[:, 1], \"+b\", label=\"Source samples\")\npl.plot(xt[:, 0], xt[:, 1], \"xr\", label=\"Target samples\")\npl.legend(loc=0)\npl.title(\"Source and target distributions\")\n\npl.figure(2)\npl.imshow(M, interpolation=\"nearest\", cmap=\"gray_r\")\npl.title(\"Cost matrix M\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Compute EMD\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "G0 = ot.solve(M, a, b).plan\n\npl.figure(3)\npl.imshow(G0, interpolation=\"nearest\", cmap=\"gray_r\")\npl.title(\"OT matrix G0\")\n\npl.figure(4)\not.plot.plot2D_samples_mat(xs, xt, G0, c=[0.5, 0.5, 1])\npl.plot(xs[:, 0], xs[:, 1], \"+b\", label=\"Source samples\")\npl.plot(xt[:, 0], xt[:, 1], \"xr\", label=\"Target samples\")\npl.legend(loc=0)\npl.title(\"OT matrix with samples\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Compute Sinkhorn\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# reg term\nlambd = 1e-1\n\nGs = ot.sinkhorn(a, b, M, lambd)\n\npl.figure(5)\npl.imshow(Gs, interpolation=\"nearest\", cmap=\"gray_r\")\npl.title(\"OT matrix sinkhorn\")\n\npl.figure(6)\not.plot.plot2D_samples_mat(xs, xt, Gs, color=[0.5, 0.5, 1])\npl.plot(xs[:, 0], xs[:, 1], \"+b\", label=\"Source samples\")\npl.plot(xt[:, 0], xt[:, 1], \"xr\", label=\"Target samples\")\npl.legend(loc=0)\npl.title(\"OT matrix Sinkhorn with samples\")\n\npl.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Empirical Sinkhorn\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# reg term\nlambd = 1e-1\n\nGes = ot.bregman.empirical_sinkhorn(xs, xt, lambd)\n\npl.figure(7)\npl.imshow(Ges, interpolation=\"nearest\", cmap=\"gray_r\")\npl.title(\"OT matrix empirical sinkhorn\")\n\npl.figure(8)\not.plot.plot2D_samples_mat(xs, xt, Ges, color=[0.5, 0.5, 1])\npl.plot(xs[:, 0], xs[:, 1], \"+b\", label=\"Source samples\")\npl.plot(xt[:, 0], xt[:, 1], \"xr\", label=\"Target samples\")\npl.legend(loc=0)\npl.title(\"OT matrix Sinkhorn from samples\")\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
}