B05_entropic-regularization.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# <center>Block 5: Optimal transport with entropic regularization</center>\n",
    "### <center>Alfred Galichon (NYU & Sciences Po)</center>\n",
    "## <center>'math+econ+code' masterclass on optimal transport and economic applications</center>\n",
    "<center>© 2018-2021 by Alfred Galichon. Past and present support from NSF grant DMS-1716489, ERC grant CoG-866274, and contributions by Jules Baudet, Pauline Corblet, Gregory Dannay, and James Nesbit are acknowledged.</center>\n",
    "\n",
    "#### <center>With python code examples</center>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Learning objectives\n",
    "\n",
    "* Entropic regularization\n",
    "\n",
    "* The log-sum-exp trick\n",
    "\n",
    "* Gradient descent, coordinate descent\n",
    "\n",
    "* The Iterated Proportional Fitting Procedure (IPFP)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## References\n",
    "\n",
    "* Galichon, *Optimal Transport Methods in Economics*,  Ch. 7.3\n",
    "\n",
    "* Peyré, Cuturi, *Computational Optimal Transport*, Ch. 4.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Entropic regularization of the optimal transport problem\n",
    "\n",
    "Consider the problem\n",
    "\n",
    "\\begin{align*}\n",
    "\\max_{\\pi\\in\\mathcal{M}\\left(  p,q\\right)  }\\sum_{ij}\\pi_{ij}\\Phi_{ij}-\\sigma\\sum_{ij}\\pi_{ij}\\ln\\pi_{ij}\n",
    "\\end{align*}\n",
    "\n",
    "where $\\sigma>0$. The problem coincides with the optimal assignment problem when $\\sigma=0$. When $\\sigma\\rightarrow+\\infty$, the solution to this problem approaches the independent coupling, $\\pi_{ij}=p_{i}q_{j}$.\n",
    "\n",
    "Later on, we will provide microfoundations for this problem, and connect it with a number of important methods in economics (BLP, gravity model, Choo-Siow...). For now, let's just view this as an extension of the optimal transport problem."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We shall compute this problem using Python libraries that we have already met with. Let us start loading them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import os\n",
    "import pandas as pd\n",
    "import time\n",
    "import scipy.sparse as spr\n",
    "#!python -m pip install -i https://pypi.gurobi.com gurobipy ## only if Gurobi not here\n",
    "import gurobipy as grb"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that in the above, Gurobi is for benchmark purposes with the case $\\sigma=0$, but is not suited to compute the nonlinear optimization problem above.\n",
    "\n",
    "---\n",
    "\n",
    "Now, let's load up the `affinitymatrix.csv`, `Xvals.csv` and `Yvals.csv` that you will recall from the previous block. We will work on a smaller population, with `nbX` types of men and `nbY` types of women."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "nbX = 5\n",
    "nbY = 3\n",
    "tol = 1e-9\n",
    "maxite = 1e+06\n",
    "\n",
    "#thepath = os.path.join(os.getcwd(),'data_mec_optim/marriage_personality-traits/')\n",
    "thepath = 'https://raw.githubusercontent.com/math-econ-code/mec_optim_2021-01/master/data_mec_optim/marriage_personality-traits/'\n",
    "\n",
    "data_X = pd.read_csv(thepath + \"Xvals.csv\")\n",
    "data_Y = pd.read_csv(thepath + \"Yvals.csv\")\n",
    "affdf = pd.read_csv(thepath + \"affinitymatrix.csv\")\n",
    "\n",
    "nbcar = 10\n",
    "affmat = affdf.iloc[0:nbcar,1:nbcar+1].to_numpy()\n",
    "sdX = data_X.std().to_numpy()\n",
    "sdY = data_Y.std().to_numpy()\n",
    "mX = data_X.mean().to_numpy()\n",
    "mY = data_Y.mean().to_numpy()\n",
    "\n",
    "Xvals = ((data_X-mX)/sdX).to_numpy()\n",
    "Yvals = ((data_Y-mY)/sdY).to_numpy()\n",
    "nobs = Xvals.shape[0]\n",
    "Phi = (Xvals @ affmat @ Yvals.T)[:nbX,:nbY]\n",
    "obj = Phi.flatten()\n",
    "    \n",
    "p = 1/nbX * np.ones(nbX)\n",
    "q = 1/nbY * np.ones(nbY)\n",
    "\n",
    "nrow = min(8, nbX) # number of rows to display\n",
    "ncol = min(8, nbY) # number of cols to displayc"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As a warm-up, let us compute as in the previous lecture the solution to the problem for $\\sigma=0$ that we can compute with Gurobi:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Time elapsed (Gurobi) = 0.009987354278564453 s.\n",
      "Value of the problem (Gurobi) = 0.41095324822187485\n",
      "[-0.63237262 -0.57124993  1.634949   -1.24417523  0.        ]\n",
      "[0.62598494 0.61304671 0.48153736]\n",
      "*************************\n"
     ]
    }
   ],
   "source": [
    "ptm = time.time()\n",
    "A1 = spr.kron(np.ones((1, nbY)), spr.identity(nbX))\n",
    "A2 = spr.kron(spr.identity(nbY), np.ones((1, nbX)))\n",
    "A = spr.vstack([A1, A2])\n",
    "obj = Phi.flatten(order = 'F') # flatten order is important\n",
    "rhs = np.hstack([p,q])\n",
    "    \n",
    "m = grb.Model('marriage')\n",
    "x = m.addMVar(len(obj), name='couple')\n",
    "m.setObjective(obj @ x, grb.GRB.MAXIMIZE)\n",
    "m.addConstr(A @ x == rhs, name=\"Constr\")\n",
    "m.setParam( 'OutputFlag', False ) #quiet output\n",
    "m.optimize()\n",
    "diff = time.time() - ptm\n",
    "print(f'Time elapsed (Gurobi) = {diff} s.')\n",
    "if m.status == grb.GRB.Status.OPTIMAL:\n",
    "    val_gurobi = m.objval\n",
    "    x = m.getAttr('x')\n",
    "    x = np.array(x).reshape([nbX, nbY])\n",
    "    pi = m.getAttr('pi')\n",
    "    u_gurobi = pi[:nbX]\n",
    "    v_gurobi = pi[nbX:nbX + nbY]\n",
    "    print(f\"Value of the problem (Gurobi) = {val_gurobi}\")\n",
    "    print(np.subtract(u_gurobi[:nrow], u_gurobi[nrow - 1]))\n",
    "    print(np.add(v_gurobi[:ncol], u_gurobi[nrow - 1]))\n",
    "    print('*************************')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Dual of the regularized problem\n",
    "\n",
    "Let's compute the dual by the minimax approach. We have\n",
    "\n",
    "\\begin{align*}\n",
    "\\max_{\\pi\\geq0}\\min_{u,v}\\sum_{ij}\\pi_{ij}\\left(  \\Phi_{ij}-u_{i}-v_{j}%\n",
    "-\\sigma\\ln\\pi_{ij}\\right)  +\\sum_{i}u_{i}p_{i}+\\sum_{j}v_{j}q_{j}%\n",
    "\\end{align*}\n",
    "\n",
    "thus\n",
    "\n",
    "\\begin{align*}\n",
    "\\min_{u,v}\\sum_{i}u_{i}p_{i}+\\sum_{j}v_{j}q_{j}+\\max_{\\pi\\geq0}\\sum_{ij}%\n",
    "\\pi_{ij}\\left(  \\Phi_{ij}-u_{i}-v_{j}-\\sigma\\ln\\pi_{ij}\\right)\n",
    "\\end{align*}\n",
    "\n",
    "By FOC in the inner problem, one has $\\Phi_{ij}-u_{i}-v_{j}-\\sigma\\ln \\pi_{ij}-\\sigma=0,$thus\n",
    "\n",
    "\\begin{align*}\n",
    "\\pi_{ij}=\\exp\\left(  \\frac{\\Phi_{ij}-u_{i}-v_{j}-\\sigma}{\\sigma}\\right)\n",
    "\\end{align*}\n",
    "\n",
    "and $\\pi_{ij}\\left(  \\Phi_{ij}-u_{i}-v_{j}-\\sigma\\ln\\pi_{ij}\\right) =\\sigma\\pi_{ij}$, thus the dual problem is\n",
    "\n",
    "\\begin{align*}\n",
    "\\min_{u,v}\\sum_{i}u_{i}p_{i}+\\sum_{j}v_{j}q_{j}+\\sigma\\sum_{ij}\\exp\\left(\n",
    "\\frac{\\Phi_{ij}-u_{i}-v_{j}-\\sigma}{\\sigma}\\right)  .\n",
    "\\end{align*}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After replacing $v_{j}$ by $v_{j}+\\sigma$, the dual is\n",
    "\n",
    "\\begin{align*}\n",
    "\\min_{u,v}\\sum_{i}u_{i}p_{i}+\\sum_{j}v_{j}q_{j}+\\sigma\\sum_{ij}\\exp\\left(\n",
    "\\frac{\\Phi_{ij}-u_{i}-v_{j}}{\\sigma}\\right)  -\\sigma. \\tag{V1}\n",
    "\\end{align*}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Another expression of the dual\n",
    "\n",
    "**Claim:** the problem is equivalent to\n",
    "\n",
    "<a name='V2'></a>\n",
    "\\begin{align*}\n",
    "\\min_{u,v}\\sum_{i}u_{i}p_{i}+\\sum_{j}v_{j}q_{j}+\\sigma\\log\\sum_{i,j}\n",
    "\\exp\\left(  \\frac{\\Phi_{ij}-u_{i}-v_{j}}{\\sigma}\\right)  \\tag{V2}\n",
    "\\end{align*}\n",
    "\n",
    "Indeed, let us go back to the minimax expression\n",
    "\n",
    "\\begin{align*}\n",
    "\\min_{u,v}\\sum_{i}u_{i}p_{i}+\\sum_{j}v_{j}q_{j}+\\max_{\\pi\\geq0}\\sum_{ij}\\pi_{ij}\\left(  \\Phi_{ij}-u_{i}-v_{j}-\\sigma\\ln\\pi_{ij}\\right)\n",
    "\\end{align*}\n",
    "\n",
    "we see that the solution $\\pi$ has automatically $\\sum_{ij}\\pi_{ij}=1$; thus we can incorporate the constraint into\n",
    "\n",
    "\\begin{align*}\n",
    "\\min_{u,v}\\sum_{i}u_{i}p_{i}+\\sum_{j}v_{j}q_{j}+\\max_{\\pi\\geq0:\\sum_{ij}\\pi_{ij}=1}\\sum_{ij}\\pi_{ij}\\left(  \\Phi_{ij}-u_{i}-v_{j}-\\sigma\\ln\\pi_{ij}\\right)\n",
    "\\end{align*}\n",
    "\n",
    "which yields the [our desired result](#V2)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[This expression](#V2) is interesting because, taking *any* $\\hat{\\pi}\\in\n",
    "M\\left(  p,q\\right)$, it reexpresses as\n",
    "\n",
    "\\begin{align*}\n",
    "\\max_{u,v}\\sum_{ij}\\hat{\\pi}_{ij}\\left(  \\frac{\\Phi_{ij}-u_{i}-v_{j}}{\\sigma}\\right)  -\\log\\sum_{ij}\\exp\\left(  \\frac{\\Phi_{ij}-u_{i}-v_{j}}{\\sigma}\\right)\n",
    "\\end{align*}\n",
    "\n",
    "therefore if the parameter is $\\theta=\\left(  u,v\\right)$, observations are\n",
    "$ij$ pairs, and the likelihood of $ij$ is\n",
    "\n",
    "\\begin{align*}\n",
    "\\pi_{ij}^{\\theta}=\\frac{\\exp\\left(  \\frac{\\Phi_{ij}-u_{i}-v_{j}}{\\sigma\n",
    "}\\right)  }{\\sum_{ij}\\exp\\left(  \\frac{\\Phi_{ij}-u_{i}-v_{j}}{\\sigma}\\right)\n",
    "}\n",
    "\\end{align*}\n",
    "\n",
    "Hence, [our expression](#problem) will coincide with the maximum likelihood in this model."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### A third expression of the dual problem\n",
    "\n",
    "Consider\n",
    "\n",
    "<a name='V2'></a>\n",
    "\\begin{align*}\n",
    "\\min_{u,v}  &  \\sum_{i}u_{i}p_{i}+\\sum_{j}v_{j}q_{j} \\\\\n",
    "s.t. \\quad &  \\sum_{i,j}\\exp\\left(  \\frac{\\Phi_{ij}-u_{i}-v_{j}}{\\sigma}\\right)\n",
    "=1\n",
    "\\end{align*}\n",
    "\n",
    "It is easy to see that the solutions of this problem coincide with [version 2](#V2). Indeed, the Lagrange multiplier is forced to be one. In other words,\n",
    "\n",
    "\\begin{align*}\n",
    "\\min_{u,v}  &  \\sum_{i}u_{i}p_{i}+\\sum_{j}v_{j}q_{j}\\\\\n",
    "s.t. \\quad &  \\sigma\\log\\sum_{i,j}\\exp\\left(  \\frac{\\Phi_{ij}-u_{i}-v_{j}}{\\sigma\n",
    "}\\right)  =0\n",
    "\\end{align*}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Small-temperature limit and the log-sum-exp trick\n",
    "\n",
    "Recall that when $\\sigma\\rightarrow0$, one has\n",
    "\n",
    "\\begin{align*}\n",
    "\\sigma\\log\\left(  e^{a/\\sigma}+e^{b/\\sigma}\\right)  \\rightarrow\\max\\left(\n",
    "a,b\\right)\n",
    "\\end{align*}\n",
    "\n",
    "Indeed, letting $m=\\max\\left(  a,b\\right)$,\n",
    "\n",
    "<a name='lse'></a>\n",
    "\\begin{align*}\n",
    "\\sigma\\log\\left(  e^{a/\\sigma}+e^{b/\\sigma}\\right)  =m+\\sigma\\log\\left(\\exp\\left(  \\frac{a-m}{\\sigma}\\right)  +\\exp\\left(  \\frac{b-m}{\\sigma}\\right)\\right),\n",
    "\\end{align*}\n",
    "and the argument of the logarithm lies between $1$ and $2$.\n",
    "\n",
    "This simple remark is actually a useful numerical recipe called the *log-sum-exp trick*: when $\\sigma$ is small, using [the formula above](#lse) to compute $\\sigma\\log\\left(  e^{a/\\sigma}+e^{b/\\sigma}\\right)$ ensures the exponentials won't blow up.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### The log-sum-exp trick for regularized OT\n",
    "\n",
    "Back to the third expression, with $\\sigma\\rightarrow0$, one has\n",
    "\n",
    "\\begin{align*}\n",
    "\\min_{u,v}  &  \\sum_{i}u_{i}p_{i}+\\sum_{j}v_{j}q_{j}\\tag{V3}\\\\\n",
    "s.t.  &  \\max_{ij}\\left(  \\Phi_{ij}-u_{i}-v_{j}\\right)  =0\\nonumber\n",
    "\\end{align*}\n",
    "\n",
    "This is exactly equivalent with the classical Monge-Kantorovich expression\n",
    "\n",
    "\\begin{align*}\n",
    "\\min_{u,v}  &  \\sum_{i}u_{i}p_{i}+\\sum_{j}v_{j}q_{j}\\tag{V3}\\\\\n",
    "s.t.  &  \\Phi_{ij}-u_{i}-v_{j}\\leq0\\nonumber\n",
    "\\end{align*}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Back to the third expression of the dual, with $\\sigma\\rightarrow0$, one has\n",
    "\n",
    "\\begin{align*}\n",
    "\\min_{u,v}  &  \\sum_{i}u_{i}p_{i}+\\sum_{j}v_{j}q_{j}\\tag{V3}\\\\\n",
    "s.t.  &  \\max_{ij}\\left(  \\Phi_{ij}-u_{i}-v_{j}\\right)  =0\\nonumber\n",
    "\\end{align*}\n",
    "\n",
    "This is exactly equivalent with the classical Monge-Kantorovich expression\n",
    "\n",
    "\\begin{align*}\n",
    "\\min_{u,v}  &  \\sum_{i}u_{i}p_{i}+\\sum_{j}v_{j}q_{j}\\tag{V3}\\\\\n",
    "s.t.  &  \\Phi_{ij}-u_{i}-v_{j}\\leq0\\nonumber\n",
    "\\end{align*}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Computation\n",
    "\n",
    "We can compute $\\min F\\left(  x\\right)$ by two methods:\n",
    "\n",
    "Either by gradient descent: $x\\left(  t+1\\right)  =x_{t}-\\epsilon _{t}\\nabla F\\left(  x_{t}\\right)  $. (Steepest descent has $\\epsilon _{t}=1/\\left\\vert \\nabla F\\left(  x_{t}\\right)  \\right\\vert $.)\n",
    "\n",
    "Or by coordinate descent: $x_{i}\\left(  t+1\\right)  =\\arg\\min_{x_{i}}F\\left(  x_{i},x_{-i}\\left(  t\\right)  \\right)$.\n",
    "\n",
    "Why do these methods converge? Let's provide some justification. We will decrease $x_{t}$ by $\\epsilon d_{t}$, were $d_{t}$ is normalized by $\\left\\vert d_{t}\\right\\vert _{p}:=\\left(  \\sum_{i=1}^{n}d_{t}^{i}\\right) ^{1/p}=1$. At first order, we have \n",
    "\n",
    "\\begin{align*}\n",
    "F\\left(  x_{t}-\\epsilon d_{t}\\right)  =F\\left(  x_{t}\\right)  -\\epsilon d_{t}^{\\intercal}\\nabla F\\left(  x_{t}\\right)  +O\\left(  \\epsilon^{1}\\right).\n",
    "\\end{align*}\n",
    "\n",
    "We need to maximize $d_{t}^{\\intercal}\\nabla F\\left(  x_{t}\\right)$ over $\\left\\vert d_{t}\\right\\vert _{p}=1$.\n",
    "\n",
    "* For $p=2$, we get $d_{t}=\\nabla F\\left(  x_{t}\\right)  /\\left\\vert \\nabla F\\left(  x_{t}\\right)  \\right\\vert $\n",
    "\n",
    "* For $p=1$, we get $d_{t}=sign\\left(  \\partial F\\left(  x_{t}\\right)/\\partial x^{i}\\right)  $ if $\\left\\vert \\partial F\\left(  x_{t}\\right) /\\partial x^{i}\\right\\vert =\\max_{j}\\left\\vert \\partial F\\left(  x_{t}\\right) /\\partial x^{j}\\right\\vert $, $0$ otherwise.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In our context, gradient descent is\n",
    "\n",
    "\\begin{align*}\n",
    "u_{i}\\left(  t+1\\right)    & =u_{i}\\left(  t\\right)  -\\epsilon\\frac{\\partial\n",
    "F}{\\partial u_{i}}\\left(  u\\left(  t\\right)  ,v\\left(  t\\right)  \\right)\n",
    ",\\text{ and }\\\\\n",
    "v_{j}\\left(  t+1\\right)    & =v_{j}\\left(  t\\right)  -\\epsilon\\frac{\\partial\n",
    "F}{\\partial v_{j}}\\left(  u\\left(  t\\right)  ,v\\left(  t\\right)  \\right)\n",
    "\\end{align*}\n",
    "\n",
    "while coordinate descent is\n",
    "\n",
    "\\begin{align*}\n",
    "\\frac{\\partial F}{\\partial u_{i}}\\left(  u_{i}\\left(  t+1\\right)\n",
    ",u_{-i}\\left(  t\\right)  ,v\\left(  t\\right)  \\right)  =0,\\text{ and }\n",
    "\\frac{\\partial F}{\\partial v_{j}}\\left(  u\\left(  t\\right)  ,v_{j}\\left(\n",
    "t+1\\right)  ,v_{-j}\\left(  t\\right)  \\right)  =0.\n",
    "\\end{align*}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Gradient descent\n",
    "\n",
    "Gradient of objective function in version 1 of our problem:\n",
    "\n",
    "\\begin{align*}\n",
    "\\left(  p_{i}-\\sum_{j}\\exp\\left(  \\frac{\\Phi_{ij}-u_{i}-v_{j}}{\\sigma}\\right)\n",
    ",q_{j}-\\sum_{i}\\exp\\left(  \\frac{\\Phi_{ij}-u_{i}-v_{j}}{\\sigma}\\right)\n",
    "\\right)\n",
    "\\end{align*}\n",
    "\n",
    "Gradient of objective function in version 2\n",
    "\n",
    "\\begin{align*}\n",
    "\\left(  p_{i}-\\frac{\\sum_{j}\\exp\\left(  \\frac{\\Phi_{ij}-u_{i}-v_{j}}{\\sigma\n",
    "}\\right)  }{\\sum_{ij}\\exp\\left(  \\frac{\\Phi_{ij}-u_{i}-v_{j}}{\\sigma}\\right)\n",
    "},q_{j}-\\frac{\\sum_{i}\\exp\\left(  \\frac{\\Phi_{ij}-u_{i}-v_{j}}{\\sigma}\\right)\n",
    "}{\\sum_{ij}\\exp\\left(  \\frac{\\Phi_{ij}-u_{i}-v_{j}}{\\sigma}\\right)  }\\right)\n",
    "\\end{align*}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Coordinate descent\n",
    "\n",
    "Coordinate descent on objective function in version 1:\n",
    "\n",
    "\\begin{align*}\n",
    "p_{i}  & =\\sum_{j}\\exp\\left(  \\frac{\\Phi_{ij}-u_{i}\\left(  t+1\\right)\n",
    "-v_{j}\\left(  t\\right)  }{\\sigma}\\right)  ,\\\\\n",
    "q_{j}  & =\\sum_{i}\\exp\\left(  \\frac{\\Phi_{ij}-u_{i}\\left(  t\\right)\n",
    "-v_{j}\\left(  t+1\\right)  }{\\sigma}\\right)\n",
    "\\end{align*}\n",
    "\n",
    "that is\n",
    "\n",
    "\\begin{align*}\n",
    "\\left\\{\n",
    "\\begin{array}\n",
    "[c]{c}\n",
    "u_{i}\\left(  t+1\\right)  =\\sigma\\log\\left(  \\frac{1}{p_{i}}\\sum_{j}\\exp\\left(\n",
    "\\frac{\\Phi_{ij}-v_{j}\\left(  t\\right)  }{\\sigma}\\right)  \\right)  \\\\\n",
    "v_{j}\\left(  t+1\\right)  =\\sigma\\log\\left(  \\frac{1}{q_{j}}\\sum_{i}\\exp\\left(\n",
    "\\frac{\\Phi_{ij}-u_{i}\\left(  t\\right)  }{\\sigma}\\right)  \\right)\n",
    "\\end{array}\n",
    "\\right.\n",
    "\\end{align*}\n",
    "\n",
    "this is called the Iterated Fitting Proportional Procedure (IPFP), or Sinkhorn's algorithm.\n",
    "\n",
    "Coordinate descent on objective function in version 2 does not yield a closed-form expression."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### IPFP, matrix version\n",
    "\n",
    "Letting $a_{i}=\\exp\\left(  -u_{i}/\\sigma\\right)  $ and $b_{j}=\\exp\\left(  -v_{j}/\\sigma\\right)  $ and $K_{ij}=\\exp\\left(  \\Phi_{ij}/\\sigma\\right)  $, one has $\\pi_{ij}=a_{i}b_{j}K_{ij}$, and the procedure reexpresses as\n",
    "\n",
    "\\begin{align*}\n",
    "\\left\\{\n",
    "\\begin{array}\n",
    "[c]{l}%\n",
    "a_{i}\\left(  t+1\\right)  =p_{i}/\\left(  Kb\\left(  t\\right)  \\right)\n",
    "_{i}\\text{ and }\\\\\n",
    "b_{j}\\left(  t+1\\right)  =q_{j}/\\left(  K^{\\intercal}a\\left(  t\\right)\n",
    "\\right)  _{j}.\n",
    "\\end{array}\n",
    "\\right.\n",
    "\\end{align*}\n",
    "\n",
    "Because this algorithm involves matrix operations only, and is naturally suited for parallel computation, GPUs are a tool of choice for addressing is. See chap. 4 of Peyré and Cuturi."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Implementation**. Let's implement this algorithm. \n",
    "\n",
    "A quick note beforehand: Numpy doesn't exactly do what one would expect when broadcasting 1d arrays to 2 arrays (for example when adding vectors and matrices). Be really careful here because things can get messy. We advise you to change all vectors (1d objects) into matrices with a single column to be sure that broadcasting is done right. \n",
    "\n",
    "Hence we define:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "def two_d(X):\n",
    "    return np.reshape(X,(X.size, 1))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Returning to the matrix-IPFP algorithm:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "IPFP1 converged in 89 steps and 0.004000425338745117 s.\n",
      "Value of the problem (IPFP1) = 0.6045556509904391\n",
      "Sum(pi*Phi) (IPFP1) = 0.4001284575694894\n",
      "*************************\n"
     ]
    }
   ],
   "source": [
    "## Matrix IPFP\n",
    "ptm = time.time()\n",
    "ite = 0\n",
    "sigma = 0.1\n",
    "\n",
    "K = np.exp(Phi/sigma)\n",
    "B = two_d(np.repeat(1, nbY))\n",
    "error = tol + 1\n",
    "    \n",
    "while error > tol and ite < maxite:\n",
    "    A = two_d(p/(K @ B).flatten(order='F'))\n",
    "    KA = (A.T @ K)\n",
    "    error = np.max(abs(np.multiply(KA,B.flatten()/q)-1))\n",
    "    B = (q / KA).T\n",
    "    ite = ite + 1\n",
    "        \n",
    "u = - sigma * np.log(A)\n",
    "v = - sigma * np.log(B)\n",
    "pi = (K * A) * np.repeat(B, 5, axis = 1).T\n",
    "val = np.sum(pi * Phi) - sigma * np.sum(pi * np.log(pi))\n",
    "end = time.time() - ptm\n",
    "if ite >= maxite:\n",
    "    print('Maximum number of iterations reached in IPFP1.')\n",
    "else:\n",
    "    print(f'IPFP1 converged in {ite} steps and {end} s.')\n",
    "    print(f'Value of the problem (IPFP1) = {val}')\n",
    "    print(f'Sum(pi*Phi) (IPFP1) = {np.sum(pi*Phi)}')\n",
    "    print('*************************')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To see the benefit of the matrix version, let us recode the same algorithm as above, but in the log-domain, namely iterate over "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "IPFP1_logs converged in 156 steps and 0.006832122802734375 s.\n",
      "Value of the problem (IPFP1_logs) = {val}\n",
      "Sum(pi*Phi) (IPFP1_logs) = 0.4109531251395014\n",
      "*************************\n"
     ]
    }
   ],
   "source": [
    "# log-domain IPFP\n",
    "sigma = 0.01\n",
    "ptm = time.time()\n",
    "ite = 0\n",
    "v = np.repeat(0, nbY)\n",
    "mu = - sigma * np.log(p)\n",
    "nu = - sigma * np.log(q)\n",
    "error = tol + 1\n",
    "while error > tol and ite < maxite:\n",
    "    u = mu + sigma * np.log(np.sum(np.exp((Phi - np.repeat(two_d(v), nbX, axis = 1).T)/sigma), axis=1))\n",
    "    KA = np.sum(np.exp((Phi - two_d(u)) / sigma), axis=0)\n",
    "    error = np.max(np.abs(KA * np.exp(-v / sigma) / q - 1))\n",
    "    v = nu + sigma * np.log(KA)\n",
    "    ite = ite + 1\n",
    "pi = np.exp((Phi - two_d(u) - np.repeat(two_d(v), nbX, axis = 1).T) / sigma)\n",
    "val = np.sum(pi * Phi) - sigma * np.sum(pi * np.log(pi))\n",
    "end = time.time() - ptm\n",
    "\n",
    "if ite >= maxite:\n",
    "    print('Maximum number of iteations reached in IPFP1.')\n",
    "else:\n",
    "    print(f'IPFP1_logs converged in {ite} steps and {end} s.')\n",
    "    print('Value of the problem (IPFP1_logs) = {val}')\n",
    "    print(f'Sum(pi*Phi) (IPFP1_logs) = {np.sum(pi * Phi)}')\n",
    "    print('*************************')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We see that the log-domain IPFP, while  mathematically equivalent to matrix IPFP, it is noticeably slower. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### IPFP with the log-sum-exp trick\n",
    "\n",
    "The matrix IPFPis very fast, partly due to the fact that it involves linear algebra operations. However, it breaks down when $\\sigma$ is small; this is best seen taking a log transform and returning to $u^{k}=-\\sigma\\log a^{k}$ and $v^{k}=-\\sigma\\log b^{k}$, that is\n",
    "\n",
    "\\begin{align*}\n",
    "\\left\\{\n",
    "\\begin{array}\n",
    "[c]{l}%\n",
    "u_{i}^{k}=\\mu_{i}+\\sigma\\log\\sum_{j}\\exp\\left(  \\frac{\\Phi_{ij}-v_{j}^{k-1}%\n",
    "}{\\sigma}\\right) \\\\\n",
    "v_{j}^{k}=\\zeta_{j}+\\sigma\\log\\sum_{i}\\exp\\left(  \\frac{\\Phi_{ij}-u_{i}^{k}%\n",
    "}{\\sigma}\\right)\n",
    "\\end{array}\n",
    "\\right.\n",
    "\\end{align*}\n",
    "\n",
    "where $\\mu_{i}=-\\sigma\\log p_{i}$ and $\\zeta_{j}=-\\sigma\\log q_{j}$.\n",
    "\n",
    "One sees what may go wrong: if $\\Phi_{ij}-v_{j}^{k-1}$ is positive in the exponential in the first sum, then the exponential blows up due to the small $\\sigma$ at the denominator. However, the log-sum-exp trick can be used in order to avoid this issue.\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Consider\n",
    "\n",
    "\\begin{align*}\n",
    "\\left\\{\n",
    "\\begin{array}\n",
    "[c]{l}%\n",
    "\\tilde{v}_{i}^{k}=\\max_{j}\\left\\{  \\Phi_{ij}-v_{j}^{k}\\right\\} \\\\\n",
    "\\tilde{u}_{j}^{k}=\\max_{i}\\left\\{  \\Phi_{ij}-u_{i}^{k}\\right\\}\n",
    "\\end{array}\n",
    "\\right.\n",
    "\\end{align*}\n",
    "\n",
    "(the indexing is not a typo: $\\tilde{v}$ is indexed by $i$ and $\\tilde{u}$ by $j$).\n",
    "\n",
    "One has\n",
    "\n",
    "\\begin{align*}\n",
    "\\left\\{\n",
    "\\begin{array}\n",
    "[c]{l}%\n",
    "u_{i}^{k}=\\mu_{i}+\\tilde{v}_{i}^{k-1}+\\sigma\\log\\sum_{j}\\exp\\left(  \\frac\n",
    "{\\Phi_{ij}-v_{j}^{k-1}-\\tilde{v}_{i}^{k}}{\\sigma}\\right) \\\\\n",
    "v_{j}^{k}=\\zeta_{j}+\\tilde{u}_{j}^{k}+\\sigma\\log\\sum_{i}\\exp\\left(  \\frac\n",
    "{\\Phi_{ij}-u_{i}^{k}-\\tilde{u}_{j}^{k}}{\\sigma}\\right)\n",
    "\\end{array}\n",
    "\\right.\n",
    "\\end{align*}\n",
    "\n",
    "and now the arguments of the exponentials are always nonpositive, ensuring the exponentials don't blow up."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Both the matrix version and the log-domain version of the IPFP  will break down when $\\sigma$ is small, e.g. $\\sigma=0.001$ (Try!). However if we modify the second procedure using the log-sum-exp trick, things work again:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "IPFP1_logs converged in 315 steps and 0.02502751350402832 s.\n",
      "Value of the problem (IPFP1_logs) = nan\n",
      "Sum(pi*Phi) (IPFP1_logs) = 0.41095352613797476\n",
      "*************************\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\antoi\\Anaconda3\\lib\\site-packages\\ipykernel_launcher.py:23: RuntimeWarning: divide by zero encountered in log\n",
      "C:\\Users\\antoi\\Anaconda3\\lib\\site-packages\\ipykernel_launcher.py:23: RuntimeWarning: invalid value encountered in multiply\n"
     ]
    }
   ],
   "source": [
    "# IPFP with log-sum-exp trick\n",
    "\n",
    "sigma = 0.001\n",
    "ptm = time.time()\n",
    "ite = 0\n",
    "v = np.repeat(0, nbY)\n",
    "mu = - sigma * np.log(p)\n",
    "nu = - sigma * np.log(q)\n",
    "error = tol + 1\n",
    "uprec = np.NINF\n",
    "while error > tol and ite < maxite:\n",
    "    vstar = np.max(Phi.T - two_d(v), axis = 0)\n",
    "    u = mu + vstar + sigma * np.log(np.sum(np.exp((Phi - np.repeat(two_d(v), nbX, axis = 1).T - \n",
    "                                                   two_d(vstar))/sigma), axis=1))\n",
    "    error = np.max(abs(u - uprec))\n",
    "    uprec = u\n",
    "    ustar = np.max(Phi - two_d(u), axis = 0)\n",
    "    KA = np.sum(np.exp((Phi - two_d(u) - np.repeat(two_d(ustar), nbX, axis = 1).T) / sigma), axis=0)\n",
    "\n",
    "    v = nu + ustar + sigma * np.log(KA)\n",
    "    ite = ite + 1\n",
    "pi = np.exp((Phi - two_d(u) - np.repeat(two_d(v), nbX, axis = 1).T) / sigma)\n",
    "val = np.sum(pi * Phi) - sigma * np.sum(pi * np.log(pi))\n",
    "end = time.time() - ptm\n",
    "\n",
    "if ite >= maxite:\n",
    "    print('Maximum number of iteations reached in IPFP1.')\n",
    "else:\n",
    "    print(f'IPFP1_logs converged in {ite} steps and {end} s.')\n",
    "    print(f'Value of the problem (IPFP1_logs) = {val}')\n",
    "    print(f'Sum(pi*Phi) (IPFP1_logs) = {np.sum(pi *Phi)}')\n",
    "    print('*************************')"
   ]
  }
 ],
 "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.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}