{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# Regularized OT with generic solver\n\nIllustrates the use of the generic solver for regularized OT with\nuser-designed regularization term. It uses Conditional gradient as in [6] and\ngeneralized Conditional Gradient as proposed in [5,7].\n\n\n[5] N. Courty; R. Flamary; D. Tuia; A. Rakotomamonjy, Optimal Transport for\nDomain Adaptation, in IEEE Transactions on Pattern Analysis and Machine\nIntelligence , vol.PP, no.99, pp.1-1.\n\n[6] Ferradans, S., Papadakis, N., Peyr\u00e9, G., & Aujol, J. F. (2014).\nRegularized discrete optimal transport. SIAM Journal on Imaging\nSciences, 7(3), 1853-1882.\n\n[7] Rakotomamonjy, A., Flamary, R., & Courty, N. (2015). Generalized\nconditional gradient: analysis of convergence and applications.\narXiv preprint arXiv:1510.06567.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# sphinx_gallery_thumbnail_number = 5\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 = 100  # nb bins\n\n# bin positions\nx = np.arange(n, dtype=np.float64)\n\n# Gaussian distributions\na = ot.datasets.make_1D_gauss(n, m=20, s=5)  # m= mean, s= std\nb = ot.datasets.make_1D_gauss(n, m=60, s=10)\n\n# loss matrix\nM = ot.dist(x.reshape((n, 1)), x.reshape((n, 1)))\nM /= M.max()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Solve EMD\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "G0 = ot.emd(a, b, M)\n\npl.figure(1, figsize=(5, 5))\not.plot.plot1D_mat(a, b, G0, \"OT matrix G0\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Solve EMD with Frobenius norm regularization\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "def f(G):\n    return 0.5 * np.sum(G**2)\n\n\ndef df(G):\n    return G\n\n\nreg = 1e-1\n\nGl2 = ot.optim.cg(a, b, M, reg, f, df, verbose=True)\n\npl.figure(2)\not.plot.plot1D_mat(a, b, Gl2, \"OT matrix Frob. reg\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Solve EMD with entropic regularization\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "def f(G):\n    return np.sum(G * np.log(G))\n\n\ndef df(G):\n    return np.log(G) + 1.0\n\n\nreg = 1e-3\n\nGe = ot.optim.cg(a, b, M, reg, f, df, verbose=True)\n\npl.figure(3, figsize=(5, 5))\not.plot.plot1D_mat(a, b, Ge, \"OT matrix Entrop. reg\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Solve EMD with Frobenius norm + entropic regularization\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "def f(G):\n    return 0.5 * np.sum(G**2)\n\n\ndef df(G):\n    return G\n\n\nreg1 = 1e-3\nreg2 = 1e-1\n\nGel2 = ot.optim.gcg(a, b, M, reg1, reg2, f, df, verbose=True)\n\npl.figure(4, figsize=(5, 5))\not.plot.plot1D_mat(a, b, Gel2, \"OT entropic + matrix Frob. reg\")\npl.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Comparison of the OT matrices\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "nvisu = 40\n\npl.figure(5, figsize=(10, 4))\n\npl.subplot(2, 2, 1)\npl.imshow(G0[:nvisu, :], cmap=\"gray_r\")\npl.axis(\"off\")\npl.title(\"Exact OT\")\n\npl.subplot(2, 2, 2)\npl.imshow(Gl2[:nvisu, :], cmap=\"gray_r\")\npl.axis(\"off\")\npl.title(\"Frobenius reg.\")\n\npl.subplot(2, 2, 3)\npl.imshow(Ge[:nvisu, :], cmap=\"gray_r\")\npl.axis(\"off\")\npl.title(\"Entropic reg.\")\n\npl.subplot(2, 2, 4)\npl.imshow(Gel2[:nvisu, :], cmap=\"gray_r\")\npl.axis(\"off\")\npl.title(\"Entropic + Frobenius reg.\")"
      ]
    }
  ],
  "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
}