{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# Optimal transport with factored couplings\n\n<div class=\"alert alert-info\"><h4>Note</h4><p>Example added in release: 0.8.2.</p></div>\n\nIllustration of the factored coupling OT between 2D empirical distributions\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# Author: Remi Flamary <remi.flamary@polytechnique.edu>\n#\n# License: MIT License\n\n# sphinx_gallery_thumbnail_number = 2\n\nimport numpy as np\nimport matplotlib.pylab as pl\nimport ot\nimport ot.plot"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Generate data an plot it\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# parameters and data generation\n\nnp.random.seed(42)\n\nn = 100  # nb samples\n\nxs = np.random.rand(n, 2) - 0.5\n\nxs = xs + np.sign(xs)\n\nxt = np.random.rand(n, 2) - 0.5\n\na, b = ot.unif(n), ot.unif(n)  # uniform distribution on samples"
      ]
    },
    {
      "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\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Compute Factored OT and exact OT solutions\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "M = ot.dist(xs, xt)\nG0 = ot.emd(a, b, M)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "Ga, Gb, xb = ot.factored_optimal_transport(xs, xt, a, b, r=4)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Plot factored OT and exact OT solutions\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "pl.figure(2, (14, 4))\n\npl.subplot(1, 3, 1)\not.plot.plot2D_samples_mat(xs, xt, G0, c=[0.2, 0.2, 0.2], alpha=0.1)\npl.plot(xs[:, 0], xs[:, 1], \"+b\", label=\"Source samples\")\npl.plot(xt[:, 0], xt[:, 1], \"xr\", label=\"Target samples\")\npl.title(\"Exact OT with samples\")\n\npl.subplot(1, 3, 2)\not.plot.plot2D_samples_mat(xs, xb, Ga, c=[0.6, 0.6, 0.9], alpha=0.5)\not.plot.plot2D_samples_mat(xb, xt, Gb, c=[0.9, 0.6, 0.6], alpha=0.5)\npl.plot(xs[:, 0], xs[:, 1], \"+b\", label=\"Source samples\")\npl.plot(xt[:, 0], xt[:, 1], \"xr\", label=\"Target samples\")\npl.plot(xb[:, 0], xb[:, 1], \"og\", label=\"Template samples\")\npl.title(\"Factored OT with template samples\")\n\npl.subplot(1, 3, 3)\not.plot.plot2D_samples_mat(xs, xt, Ga.dot(Gb), c=[0.2, 0.2, 0.2], alpha=0.1)\npl.plot(xs[:, 0], xs[:, 1], \"+b\", label=\"Source samples\")\npl.plot(xt[:, 0], xt[:, 1], \"xr\", label=\"Target samples\")\npl.title(\"Factored OT low rank OT plan\")"
      ]
    }
  ],
  "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
}