L01_optimal-assignment.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# <center>Block 1b: discrete optimal transport</center>\n",
    "### <center>Alfred Galichon (NYU & Sciences Po)</center>\n",
    "## <center>'math+econ+code' masterclass on optimal transport and economic applications</center>\n",
    "#### <center>With python code examples</center>\n",
    "© 2018-2022 by Alfred Galichon. Past and present support from NSF grant DMS-1716489, ERC grant CoG-866274 are acknowledged, as well as inputs from contributors listed [here](http://www.math-econ-code.org/theteam).\n",
    "\n",
    "**If you reuse material from this masterclass, please cite as:**<br>\n",
    "Alfred Galichon, 'math+econ+code' masterclass on optimal transport and economic applications, January 2022. https://github.com/math-econ-code/mec_optim"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Learning Objectives\n",
    "\n",
    "* Optimal assignment problem\n",
    "\n",
    "* Pairwise stability, Walrasian equilibrium\n",
    "\n",
    "* Computation\n",
    "\n",
    "## References\n",
    "\n",
    "* Galichon, *Optimal Transport Methods in Economics*, Ch. 3\n",
    "\n",
    "* Roth, Sotomayor(1990). *Two-Sided Matching*. Cambridge.\n",
    "\n",
    "* Koopmans and Beckmann (1957). \"Assignment problems and the location of economic activities.\" *Econometrica*.\n",
    "\n",
    "* Shapley and Shubik (1972). The assignment game I: The core.\" *IJGT*.\n",
    "\n",
    "* Becker (1993). *A Treatise of the Family*. Harvard.\n",
    "\n",
    "* Gretsky, Ostroy, and Zame (1992). \"The nonatomic assignment model.\" *Economic Theory*.\n",
    "\n",
    "* Burkard, Dell'Amico, and Martello (2012). *Assignment Problems*. SIAM.\n",
    "\n",
    "* Dupuy and Galichon (2014). \"Personality traits and the marriage market.\" *JPE*."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Motivation\n",
    "\n",
    "## Optimal Transport\n",
    "\n",
    "Consider the problem of assigning a possibly infinite number of workers\n",
    "and firms.\n",
    "\n",
    "* Each worker should work for one firm, and each firm should hire one worker.\n",
    "\n",
    "* Workers and firms have heterogenous characteristics; let $x\\in \\mathcal{X}$ and $y\\in\\mathcal{Y}$ be the characteristics of workers and firms respectively.\n",
    "\n",
    "* Workers and firms are in equal mass, which is normalized to one. Both worker's types and firm's types have a distribution on $\\mathcal{X}$ and $\\mathcal{Y}$ which is assumed known.\n",
    "\n",
    "It is assumed that if a worker $x$ matches with a firm $y$, the total output generated is $\\Phi_{xy}$. The questions are then:\n",
    "\n",
    "* **Optimality:** what is the optimal assignment in the sense that it maximizes the overal output generated?\n",
    "\n",
    "* **Equilibrium:** what are the equilibrium assignment and the equilibrium wages?\n",
    "\n",
    "* **Efficiency:** do these two notions coincide?\n",
    "\n",
    "The same tools have been used by Gary Becker to study the heterosexual marriage market, where $x$ is the man's characteristics, and $y$ is the woman's characteristics, and \"wages\" are replaced by \"transfers\".\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data\n",
    "\n",
    "We shall take a first look at marriage data. Dupuy and Galichon (JPE, 2014) study a marriage dataset where, in addition to usual socio-demographic features (such as education and age), other features that measure personality traits are reported.\n",
    "\n",
    "* The literature on quantitative psychology argues that one can capture relatively well an individual's personality along five dimensions, the \"big 5\" - consciousness, extraversion, agreableness, emotional stability, autonomy - assessed though a standardized questionaire.\n",
    "\n",
    "* In addition to this, we observed a (self-assessed) measure of health, risk-aversion, education, height and body mass index = weight in kg/ (height in m)$^2$. In total, the available characteristics $x_{i}$ of man $i$ and $y_{j}$ of woman $j$ are both $10$-dimensional vectors.\n",
    "\n",
    "* It is assumed that the surplus of interaction is given by $\\Phi\\left(x,y\\right)  =x ^{\\intercal}Ay$, where $A$ is a *given* $10 \\times 10$ matrix. (later in this course, we'll see how to estimate $A$ based on matched marital data).\n",
    "\n",
    "Today, we solve a central planner's problem (a stylized version of the problem OKCupids would solve): given a population of men and a population of women, how do we mutually assign these in order to:\n",
    "\n",
    "1. maximize matching surplus \n",
    "\n",
    "2. attain a (hopefully) stable assignment."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# !python -m pip install -i https://pypi.gurobi.com gurobipy ## only if Gurobi not here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "import scipy.sparse as spr\n",
    "import gurobipy as grb\n",
    "from tabulate import tabulate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "thepath = 'https://raw.githubusercontent.com/math-econ-code/mec_optim_2021-01/master/data_mec_optim/marriage_personality-traits/'\n",
    "data_X = pd.read_csv(thepath + \"Xvals.csv\")\n",
    "data_Y = pd.read_csv(thepath + \"Yvals.csv\")\n",
    "data_X.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We standardize the data using:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sdX,sdY = data_X.std().values, data_Y.std().values\n",
    "mX,mY = data_X.mean().values, data_Y.mean().values\n",
    "\n",
    "feats_x_k, feats_y_l = ((data_X-mX)/sdX).values, ((data_Y-mY)/sdY).values\n",
    "nbx,nbk = feats_x_k.shape\n",
    "nby,nbl = feats_y_l.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We postulate that the form of the surplus function is\n",
    "\\begin{align*}\n",
    "\\Phi_{xy}=x^{\\intercal} Ay\n",
    "\\end{align*}\n",
    "where $x$ and $y$ are the 10-dimensional features of a man and a woman, and the form of $A$, a 10x10 matrix, is given and stored in the file `affinitymatrix.csv`. Again, we'll see later how to solve the econometrics problem of estimating $A$.\n",
    "\n",
    "We load $A$:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "aff_data = pd.read_csv(thepath + \"affinitymatrix.csv\")\n",
    "A_k_l = aff_data.iloc[0:nbk,1:nbl+1].values\n",
    "print(tabulate(A_k_l))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we build the surplus matrix $\\Phi_{xy}$ using:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Φ_x_y = feats_x_k @ A_k_l @ feats_y_l.T"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we have one individuals of each type, so:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We are ready to build a class `OTProblem` which summarize the data of the problem."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class OTProblem():\n",
    "    def __init__(self,Φ_x_y,n_x=None,m_y=None):\n",
    "        self.Φ_a = Φ_x_y.flatten()\n",
    "        self.nbx , self.nby = Φ_x_y.shape # number of types on each parts\n",
    "        self.nbz = self.nbx + self.nby # total number of types \n",
    "        self.nba = self.nbx * self.nby # number of arcs\n",
    "        if n_x is None: \n",
    "            self.n_x=np.ones(self.nbx)\n",
    "        else:\n",
    "            self.n_x = n_x\n",
    "        if m_y is None: \n",
    "            self.m_y=np.ones(self.nby)\n",
    "        else:\n",
    "            self.m_y = m_y    \n",
    "        return\n",
    "    \n",
    "marriageEx = OTProblem(Φ_x_y)\n",
    "twoBytwoEx = OTProblem(np.eye(2))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The Discrete Monge-Kantorovich Theorem\n",
    "\n",
    "Assume that the type spaces $\\mathcal{X}$ and $\\mathcal{Y}$ are finite. \n",
    "\n",
    "The mass of workers of type $x$ is $n_{x}$; the mass of jobs of type $y$ is $m_{y}$, with $\\sum_{x}n_{x}=\\sum_{y}m_{y}$.\n",
    "\n",
    "Let $\\mu_{xy}$ be the mass of workers of type $x$ assigned to jobs of type $y$. Every worker is busy and every job is filled, thus\n",
    "\n",
    "<a name=\"dicsr-constraints\"></a>\n",
    "\\begin{align*}\n",
    "\\mu\\in\\mathcal{M}\\left(n,m\\right) :=\\left\\{ \\mu_{xy}\\geq 0, \\sum_{y\\in\\mathcal{Y}}\\mu_{xy}=n_{x}\\text{ and }\\sum_{x\\in\\mathcal{X}}\\mu\n",
    "_{xy}=m_{y} \\right\\}. \n",
    "\\end{align*}\n",
    "\n",
    "(Note that this formulation allows for mixing, i.e. it allows for $\\mu_{xy}>0$ and $\\mu_{xy^{\\prime}}>0$ to hold simultaneously with $y\\neq y^{\\prime}$.)\n",
    "\n",
    "\n",
    "Assume the economic output created when assigning worker $x$ to job $y$ is $\\Phi_{xy}$. Hence, the optimal assignment is\n",
    "\n",
    "<a name=\"OAP\"></a>\n",
    "\\begin{align*}\n",
    "\\max_{\\mu\\geq0}  &  \\sum_{xy}\\mu_{xy}\\Phi_{xy} \\\\\n",
    "s.t.~  &  \\sum_{y\\in\\mathcal{Y}}\\mu_{xy}=n_{x} \\left[u_{x}\\right] \\\\\n",
    "&  \\sum_{x\\in\\mathcal{X}}\\mu_{xy}=m_{y}~\\left[v_{y}\\right]\n",
    "\\end{align*}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Duality\n",
    "\n",
    "---\n",
    "**Theorem**\n",
    "1. The value of the [primal problem](#OAP) *coincides with the value of the dual problem*\n",
    "\n",
    "    <a name=\"dual-discr\"></a>\n",
    "    \\begin{align*}\n",
    "    \\min_{u,v}  &  \\sum_{x\\in\\mathcal{X}}n_{x}u_{x}+\\sum_{y\\in\\mathcal{Y}}\n",
    "    m_{y}v_{y}\\\\\n",
    "    s.t.  &  u_{x}+v_{y}\\geq\\Phi_{xy}~\\left[\\mu_{xy}\\geq0\\right]\n",
    "    \\end{align*}\n",
    "\n",
    "2. Both the primal and the dual problems have optimal solutions. If $\\mu$ *is a solution to the primal problem* and $\\left(u,v\\right)$ is *a solution to the dual problem, then by complementary slackness*,\n",
    "\n",
    "    <a name=\"complSlack\"> </a>\n",
    "    \\begin{align*}\n",
    "    \\mu_{xy}>0\\text{ implies }u_{x}+v_{y}=\\Phi_{xy}.\n",
    "    \\end{align*}\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Proof**. The proof follows from the min-cost flow duality result, but let us rewrite it anyway. \n",
    "\n",
    "1. The value of the [primal problem](#OAP) can be written as $\\max_{\\mu\\geq0}\\min_{u,v}S\\left(  \\mu,u,v\\right)$, where\n",
    "\n",
    "    \\begin{align*}\n",
    "    S\\left(  \\mu,u,v\\right)  :=\\sum_{xy}\\mu_{xy}\\Phi_{xy}+\\sum_{x\\in\\mathcal{X}}u_{x}(n_{x}-\\sum_{y\\in\\mathcal{Y}}\\mu_{xy})+\\sum_{y\\in\\mathcal{Y}}v_{y}(m_{y}-\\sum_{x\\in\\mathcal{X}}\\mu_{xy})\n",
    "    \\end{align*}\n",
    "\n",
    "    but by the minmax theorem, this value is equal to $\\min_{u,v}\\max_{\\mu\\geq 0}S\\left(  \\mu,u,v\\right)$, which is the value of the [dual problem](#ual-discr}).\n",
    "\n",
    "2. Follows by noting that, for a primal solution $\\mu$ and a dual solution $\\left(  u,v\\right)  $, then $S\\left( \\mu,u,v\\right)  =\\sum_{xy}\\mu_{xy} \\Phi_{xy}$. $\\blacksquare$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The following statements are equivalent:\n",
    "\n",
    "* $\\mu$ is an optimal solution to the primal problem, and $\\left(\n",
    "u,v\\right)  $ is an optimal solution to the dual problem, and\n",
    "\n",
    "* (i) $\\mu\\in \\mathcal{M}\\left(  n,m\\right)  $\n",
    "\n",
    "* (ii) $u_{x}+v_{y}\\geq\\Phi_{xy}$\n",
    "\n",
    "* (iii) $\\mu_{xy}>0$ implies $u_{x}+v_{y}\\leq\\Phi_{xy}$.\n",
    "\n",
    "We saw the direct implication. But the converse is easy: take $\\mu$ and $\\left(  u,v\\right)  $ satisfying (i)--(iii), Then one has\n",
    "\n",
    "\\begin{align*}\n",
    "dual\\leq\\sum_{x}n_{x}u_{x}+\\sum_{y}m_{y}v_{y}=\\sum_{xy}\\mu_{xy}\\left(\n",
    "u_{x}+v_{y}\\right)  \\leq\\sum_{xy}\\mu_{xy}\\Phi_{xy}\\leq primal\n",
    "\\end{align*}\n",
    "\n",
    "but by the MK duality theorem, both ends coincide. Thus $\\mu$ is optimal for the primal and $\\left(  u,v\\right)  $ for the dual."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Pairwise stability\n",
    "\n",
    "In a marriage context, an important concept is stability:\n",
    "\n",
    "An outcome is a vector $\\left(  \\mu,u,v\\right)  $, where $u_{x}$ and $v_{y}$ are $x$'s and $y$'s payoffs, and $\\mu$ is a matching that is\n",
    "\n",
    "<a name=\"primFeas\"></a>\n",
    "\\begin{align*}\n",
    "\\mu\\in\\mathcal{M}\\left(  n,m\\right).\n",
    "\\end{align*}\n",
    "\n",
    "A pair $xy$ is blocking if $x$ and $y$ can find a way of sharing their joint surplus $\\Phi_{xy}$ in such a way that $x$ gets more than $u_{x}$ and $y$ gets more than $v_{y}$. Hence there is no blocking pair if and only if for every $x$ and $y$, one has\n",
    "\n",
    "<a name=\"noBlocking\"></a>\n",
    "\\begin{align*}\n",
    "u_{x}+v_{y}\\geq\\Phi_{xy}.\n",
    "\\end{align*}\n",
    "\n",
    "If $x$ and $y$ are actually matched, their utilities $u_{x}$ and $v_{y}$ need to be feasible, i.e. the above inequality should be saturated. Hence\n",
    "\n",
    "<a name=\"cplSlck\"></a>\n",
    "\\begin{align*}\n",
    "\\mu_{xy}>0\\text{ implies }u_{x}+v_{y}=\\Phi_{xy}.\n",
    "\\end{align*}\n",
    "\n",
    "---\n",
    "**Definition:**\n",
    "\n",
    "A matching that satisfies [primal feasbilitity](#primFeas), [no blocking](#noBlocking), and [complementary slackness](#cplSlck) is called a stable matching.\n",
    "\n",
    "As it turns out, these conditions are precisely the conditions that express complementarity slackness in the Monge-Kantorovich problem. Therefore, outcome$\\left(  \\mu,u,v\\right)  $ is stable if and only if $\\mu$ is a solution to the primal problem, and $\\left(  u,v\\right)  $ is a solution to the dual problem."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Wage market interpretation 1\n",
    "\n",
    "Back to the workers / firms interpretation and assume for now that workers are indifferent between any two firms that offer the same salary. We argue that $u_x  $ can be interpreted as the equilibrium wage of worker $x$, while $v_y $ can be interpreted as the equilibrium profit of firm $y$. Indeed:\n",
    "\n",
    "---\n",
    "**Proposition**\n",
    "If $\\left(u,v\\right)$ is a solution to the dual of the Kantorovich problem, then\n",
    "\n",
    "\\begin{align*}\n",
    "u_{x}  &  =\\max_{y\\in\\mathcal{Y}}\\left(  \\Phi_{xy}-v_{y}\\right)\n",
    "\\label{conjug1}\\\\\n",
    "v_{y}  &  =\\max_{x\\in\\mathcal{X}}\\left(  \\Phi_{xy}-u_{x}\\right)  .\n",
    "\\label{conjug2}\n",
    "\\end{align*}\n",
    "\n",
    "---\n",
    "\n",
    "Therefore, $u_{x}$ can be interpreted as equilibrium wage of worker $x$, and $v_{y}$ as equilibrium profit of firm $y$. In this interpretation, all workers get the same wage at equilibrium.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Wage market interpretation 2\n",
    "\n",
    "Assume now that if a worker of type $x$ works for a firm of type $y$ for wage $w_{xy}$, then gets $\\alpha_{xy}+w_{xy}$, where $\\alpha_{xy}$ is the nonmonetary payoff associated with working with a firm of type $y$. The firm's profit is $\\gamma_{xy}-w_{xy}$, where $\\gamma_{xy}$ is the economic output.\n",
    "\n",
    "If an employee of type $x$ matches with a firm of type $y$, they generate joint surplus $\\Phi_{xy}$, given by%\n",
    "\n",
    "\\begin{align*}\n",
    "\\Phi_{xy}=\\underset{\\text{employee's payoff}}{\\underbrace{\\alpha_{xy}+w_{xy}}}+\\underset{\\text{firm's payoff}}{\\underbrace{\\gamma_{xy}-w_{xy}}}=\\alpha_{xy}+\\gamma_{xy}\n",
    "\\end{align*}\n",
    "\n",
    "which is independent from $w$.\n",
    "\n",
    "Workers choose firms which maximize their utility, i.e. solve \n",
    "\n",
    "<a name=\"equivWalrasStable1\"></a>\n",
    "\\begin{align*}\n",
    "u_{x}=\\max_{y}\\left\\{  \\alpha_{xy}+w_{xy}\\right\\}\n",
    "\\end{align*}\n",
    "\n",
    "and $u_{x}=\\alpha_{xy}+w_{xy}$ if $x$ and $y$ are matched. Similarly, the indirect payoff vector of firms is\n",
    "\n",
    "<a name=\"equivWalrasStable2\"></a>\n",
    "\\begin{align*}\n",
    "v_{y}=\\max_{x}\\left\\{  \\gamma_{xy}-w_{xy}\\right\\}\n",
    "\\end{align*}\n",
    "\n",
    "and, again, $v_{y}=\\gamma_{xy}-w_{xy}$ if $x$ and $y$ are matched."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As a result,\n",
    "\n",
    "\\begin{align*}\n",
    "u_{x}+v_{y}\\geq\\alpha_{xy}+\\gamma_{xy}=\\Phi_{xy}\n",
    "\\end{align*}\n",
    "\n",
    "and equality holds if $x$ and $y$ are matched. Thus, if $w_{xy}$ is an equilibrium wage, then the triple $\\left(  \\mu,u,v\\right)$ where $\\mu$ is the corresponding matching, and $u_{x}$ and $v_{y}$ are defined by [this](#equivWalrasStable1) and [this](#equivWalrasStable2) defines a stable outcome.\n",
    "\n",
    "Conversely, let $\\left(\\mu,u,v\\right)$ be a stable outcome. Then let $\\bar{w}_{xx}$ and \\b{w}$_{xy}$ be defined by\n",
    "\n",
    "\\begin{align*}\n",
    "\\bar{w}_{xy}=u_{x}-\\alpha_{xy}\\text{ and }w^l_{xy}=\\gamma_{xy}-v_{y}.\n",
    "\\end{align*}\n",
    "\n",
    "One has $\\bar{w}_{xy}\\geq b^l_{xy}$. Any $w_{xy}$ such that $\\bar{w}_{xy}\\geq w_{xy}\\geq b^l_{xy}$ is an equilibrium wage. Indeed, $\\mu_{xy}>0$ implies $\\bar{w}_{xy} \\geq b^l_{xy}$, thus [this](#equivWalrasStable1) and [this](#equivWalrasStable2) hold. Given $u$ and $v$, $w_{xy}$ is uniquely defined on the equilibrium path (ie. when $x$ and $y$ are such that $\\mu_{xy}>0$), but there are multiple choices of $w$ outside the equilibrium path.\n",
    "\n",
    "Note that all workers of the same type get the same indirect utility, but not necessarly the same wage."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Computation\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This problem of computation of the Optimal Assignment Problem, more specifically of $\\left(\\mu,u,v\\right)$, is arguably the most studied problem in Computer Science, and dozens, if not hundreds of algorithms exist, whose running time is polynomial in $\\max\\left(n,m\\right)$, typically a power less than three of the latter.\n",
    "\n",
    "Famous algorithms include: the Hungarian algorithm (Kuhn-Munkres); Bertsekas' auction algorithm; Goldberg and Kennedy's pseudoflow algorithm. For more on these, see the book by Burkard, Dell'Amico, and Martello, and a\n",
    "introductory presentation in http://www.assignmentproblems.com/doc/LSAPIntroduction.pdf.\n",
    "\n",
    "Here, we will show how to solve the problem with the help of a Linear Programming solver used as a black box; our challenge here will be to carefully set up the constraint matrix as a sparse matrix in order to let a large scale Linear Programming solvers such as Gurobi recognize and exploit the sparsity of the problem."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let $\\mu$ and $\\Phi$ be the matrices with typical elements $\\left(\\mu_{xy}\\right)  $ and $\\left(  \\Phi_{xy}\\right)  $. We let $n$, $m$, $u$,$v$, and $1$ the column vectors with entries $\\left(  n_{x}\\right)$, $\\left(  m_{y}\\right)  $, $\\left(  u_{x}\\right)  $, $\\left(  v_{y}\\right)$, and $1$, respectively. The optimal assignment problem\n",
    "\n",
    "\\begin{align*}\n",
    "\\max_{\\mu\\geq0}  &  \\sum_{xy}\\mu_{xy}\\Phi_{xy} \\\\\n",
    "s.t.~  &  \\sum_{y\\in\\mathcal{Y}}\\mu_{xy}=n_{x}~\\left[  u_{x}\\right]\\\\\n",
    "&  \\sum_{x\\in\\mathcal{X}}\\mu_{xy}=m_{y}~\\left[  v_{y}\\right] \n",
    "\\end{align*}\n",
    "\n",
    "Can be rewritten writes using matrix algebra as\n",
    "\n",
    "\\begin{align*}\n",
    "&  \\max_{\\mu\\geq0}Tr\\left(  \\Pi^{\\top}\\mu\\right) \\\\\n",
    "&  \\mu 1_{\\mathcal{Y}}=n\\\\\n",
    "&  1_{\\mathcal{X}}^{\\top}\\mu=m^{\\top}.\n",
    "\\end{align*}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Recall that if $A$ is a $M\\times P$ matrix and $B$ a $N\\times Q$ matrix, and if $vec(A)$ vectorizes in the row-major order (i.e. concatenates the rows of $A$) then the Kronecker product $A\\otimes B$ of $A$ and $B$ is a $MN\\times PQ$ matrix such that\n",
    "\n",
    "\\begin{align*}\n",
    "vec\\left(  AXB^{\\top}\\right)  =\\left(  A\\otimes B\\right)  vec\\left(\n",
    "X\\right). \\label{VecAndKronecker}\n",
    "\\end{align*}\n",
    "\n",
    "In python, $A\\otimes B$ is implemented by `sparse.kron(A,B)` of the library `sparse` of scipy."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Vectorizing the constraints"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The first constraint $I_{\\mathcal{X}}\\mu 1_{\\mathcal{Y}}=n$,  therefore vectorizes under the row-major order as\n",
    "\n",
    "\\begin{align*}\n",
    "\\left( I_{\\mathcal{X}} \\otimes1_{\\mathcal{Y}}^{\\top}\\right)  vec\\left(  \\mu\\right)  =vec\\left(\n",
    "n\\right),\n",
    "\\end{align*}\n",
    "\n",
    "and similarly, the second constraint $1_{\\mathcal{X}}^{\\top}\\mu I_{\\mathcal{Y}}=m^{\\top}$, vectorizes as\n",
    "\n",
    "\\begin{align*}\n",
    "\\left( 1_{\\mathcal{X}}^{\\top} \\otimes I_{\\mathcal{Y}}\\right)  vec\\left(  \\mu\\right)  =vec\\left(\n",
    "m\\right)  .\n",
    "\\end{align*}\n",
    "\n",
    "Note that the matrix $I_{\\mathcal{X}} \\otimes1_{\\mathcal{Y}}^{\\top}$ is of size $| \\mathcal{X}| \\times |\\mathcal{X}| |\\mathcal{Y}|$, and the matrix $ 1_{\\mathcal{X}}^{\\top} \\otimes I_{\\mathcal{Y}}$ is of size $ |\\mathcal{Y}|\\times |\\mathcal{X}| |\\mathcal{Y}|$; hence the full matrix involved in the left-hand side of the constraints is of size $ (|\\mathcal{X}|+|\\mathcal{Y}|)\\times |\\mathcal{X}| |\\mathcal{Y}|$. In spite of its large size, this matrix is *sparse*. In python, the identity matrix $I_{N}$ is coded as `scipy.sparse.identity(N)`.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Introducing the MOM (Margining-Out Matrix)\n",
    "\n",
    "We therefore introduce the *margining-out matrix* as:\n",
    "\n",
    "$M=\\binom{M_\\mathcal{X}}{M_\\mathcal{Y}}$\n",
    "\n",
    "where<br>\n",
    "$\n",
    "\\begin{eqnarray*}\n",
    "M_\\mathcal{X} := I_{\\mathcal{X}}\\otimes 1_{\\mathcal{Y}}^{\\top }\\\\\n",
    "M_\\mathcal{Y}=1_{\\mathcal{X}}^{\\top }\\otimes I_{\\mathcal{Y}}\n",
    "\\end{eqnarray*}\n",
    "$<br>\n",
    "have the effect of taking the $x$ and $y$ margins, respectively.\n",
    "\n",
    "We implement using:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def M_z_a(self): \n",
    "    return spr.vstack([spr.kron(spr.identity(self.nbx), np.ones((1,self.nby))),spr.kron(np.ones((1,self.nbx)),spr.identity(self.nby))])\n",
    "\n",
    "OTProblem.M_z_a = M_z_a"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For instance, let us see what the margining matrix is in the case of a 2x2 example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "twoBytwoEx.M_z_a().todense()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We set $q=(n^\\top,m^\\top)^\\top$, and implement using:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def q_z(self):\n",
    "    return(np.concatenate([self.n_x,self.m_y]))\n",
    "\n",
    "OTProblem.q_z = q_z"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The linear programming problem then becomes<br>\n",
    "$\\begin{eqnarray}\n",
    "&  \\max_{z\\geq0}vec\\left(  \\Phi\\right)  ^{\\top} vec(\\mu) \\label{LPvectorized}\\\\\n",
    "s.t.~  &  M vec(\\mu) = q \n",
    "\\end{eqnarray}\n",
    "$\n",
    "\n",
    "which is ready to be passed on to a linear programming solver such as Gurobi."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def solve_full_lp(self,OutputFlag=True):\n",
    "    m=grb.Model()\n",
    "    μ_a = m.addMVar(self.nba)\n",
    "    m.setObjective(self.Φ_a @ μ_a, grb.GRB.MAXIMIZE)\n",
    "    m.addConstr(self.M_z_a() @ μ_a == self.q_z())\n",
    "    m.setParam( 'OutputFlag', OutputFlag )\n",
    "    m.optimize()\n",
    "    if m.status == grb.GRB.Status.OPTIMAL:\n",
    "        μopt_a = np.array(m.getAttr('x')).reshape(self.nbx,self.nby)\n",
    "        popt_z = m.getAttr('pi')\n",
    "        uopt_x, vopt_y = popt_z[:self.nbx],popt_z[self.nbx:]\n",
    "    return (μopt_a.reshape((self.nbx,-1)),uopt_x,vopt_y)\n",
    "\n",
    "OTProblem.solve_full_lp = solve_full_lp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "(μ_x_y,u_x,v_y) = marriageEx.solve_full_lp()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Who does man $0$ match with?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Man 0 matches with woman '+str(np.argwhere(μ_x_y[0,:] != 0)[0][0])+'.'  )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Allowing singles: matching with partial assignment\n",
    "\n",
    "A important variant of the problem exists with $\\sum_{x\\in\\mathcal{X}}n_{x}\\neq\\sum_{y\\in\\mathcal{Y}}m_{y}$ and the primal constraints become inequality constraints. The duality then becomes\n",
    "\n",
    "\\begin{align*}\n",
    "\\begin{array}\n",
    "[c]{rrr}\n",
    "\\max_{\\mu\\geq0}\\sum\\mu_{xy}\\Phi_{xy} & = & \\min_{u,v}\\sum_{x\\in\\mathcal{X}\n",
    "}n_{x}u_{x}+\\sum_{y\\in\\mathcal{Y}}m_{y}v_{y}\\\\\n",
    "s.t.~\\sum_{y\\in\\mathcal{Y}}\\mu_{xy}\\leq n_{x} &  & u\\geq0,~v\\geq0 \\\\\n",
    "\\sum_{x\\in\\mathcal{X}}\\mu_{xy}\\leq m_{y} &  & u_{x}+v_{y}\\geq\\Phi_{xy}\n",
    "\\end{array}\n",
    "\\end{align*}\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def solve_partial_lp(self, OutputFlag = True):\n",
    "    m=grb.Model()\n",
    "    μ_a = m.addMVar(self.nba)\n",
    "    m.setObjective(self.Φ_a @ μ_a, grb.GRB.MAXIMIZE)\n",
    "    m.addConstr(self.M_z_a() @ μ_a <= self.q_z())\n",
    "    m.setParam( 'OutputFlag', OutputFlag )\n",
    "    m.optimize()\n",
    "    if m.status == grb.GRB.Status.OPTIMAL:\n",
    "        μopt_a = np.array(m.getAttr('x')).reshape(self.nbx,self.nby)\n",
    "        popt_z = np.array(m.getAttr('pi'))\n",
    "        uopt_x, vopt_y = popt_z[:self.nbx],popt_z[self.nbx:]\n",
    "    return (μopt_a.reshape((self.nbx,-1)),uopt_x,vopt_y)\n",
    "\n",
    "OTProblem.solve_partial_lp = solve_partial_lp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "(μ_x_y,u_x,v_y) = marriageEx.solve_partial_lp()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "μ_x_y.sum()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We saw that in the present example, even though everyone had the possibility of being unmatched, no individual actually took that option."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also study the payoff vectors."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "u_x.min(),v_y.min()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We see that at least one woman gets a payoff of zero, which makes her indifferent between being matched and being unmatched. It looks like this solution is not the preferred one for women. Let's see if there is another one which is better for women."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Multiple objectives\n",
    "\n",
    "Let's see if, among the solutions to the problem above, there is one which maximizes the welfare of women. Using https://www.gurobi.com/documentation/9.5/refman/py_model_setobjectiven.html\n",
    "we write:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def solve_dual_partial_multiobj_lp(self, favor_Y = True):\n",
    "    if favor_Y:\n",
    "        welfare_weights_z = np.concatenate([np.ones(self.nbx),-np.ones(self.nby) ]) # welfare weights for each individuals\n",
    "    else:\n",
    "        welfare_weights_z = np.concatenate([-np.ones(self.nbx),np.ones(self.nby) ]) # welfare weights for each individuals\n",
    "    m=grb.Model()\n",
    "    p_z = m.addMVar(self.nbz)\n",
    "    m.ModelSense = grb.GRB.MINIMIZE\n",
    "    m.setParam('OutputFlag', 0) \n",
    "    m.setObjectiveN(self.q_z() @ p_z, 0, priority = 2,weight = 1)\n",
    "    m.setObjectiveN(welfare_weights_z @ p_z,1,priority = 1, weight = 1)\n",
    "    m.addConstr(self.M_z_a().T @ p_z >= self.Φ_a)\n",
    "    m.optimize()  \n",
    "    if m.status == grb.GRB.Status.OPTIMAL:\n",
    "        #μopt_a = np.array(m.getAttr('Pi')).reshape(self.nbx,self.nby)\n",
    "        popt_z = np.array(m.getAttr('x'))\n",
    "        uopt_x, vopt_y = popt_z[:self.nbx],popt_z[self.nbx:]\n",
    "    return (uopt_x,vopt_y) # μopt_a.reshape((self.nbx,-1))\n",
    "\n",
    "OTProblem.solve_dual_partial_multiobj_lp = solve_dual_partial_multiobj_lp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "u_x,v_y = marriageEx.solve_dual_partial_multiobj_lp(favor_Y=True)\n",
    "u_x.min(),v_y.min()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "u_x,v_y = marriageEx.solve_dual_partial_multiobj_lp(favor_Y=False)\n",
    "u_x.min(),v_y.min()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Exercises"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Exercise.** Integral primal solutions. ([OTME], ex. 3.3) A *doubly stochastic matrix* is a $n \\times n$ matrix $\\mu$ which is a solution to<br>\n",
    "$M \\mu = 1_{2n}$\n",
    "where $M$ is the margining matrix seen above.<br>\n",
    "A *permutation matrix* is a doubly stochastic matrix whose entries are all integer.<br>\n",
    "The solution to $\\max_{\\mu \\geq 0} \\sum_{xy} \\mu_{xy} \\Phi_{xy} : M \\mu = 1_{2n}$ is therefore a doubly stochastic matrix.<br>\n",
    "Write a program that takes as an input a doubly stochastic matrix solution to the optimal assignment problem and returns a permutation matrix solution to the same problem."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# answer:\n",
    "\n",
    "import networkx as nx\n",
    "\n",
    "def improve(μ_x_y):\n",
    "    nbx,nby = μ_x_y.shape\n",
    "    matched_pairs = [(x,nbx+y) for x in range(nbx) for y in range(nby) if ((μ_x_y[x,y] > 1e-9 ) and (μ_x_y[x,y] < 1-1e-9 ) ) ]\n",
    "    if len(matched_pairs)==0:\n",
    "        return False\n",
    "    g = nx.Graph()\n",
    "    g.add_nodes_from(list(range(nbx))+list(range(nbx,nby)))\n",
    "    g.add_edges_from(matched_pairs)\n",
    "    cycle = nx.find_cycle(g)\n",
    "    Δμ_x_y = np.zeros((nbx,nby ))\n",
    "    m = np.inf\n",
    "    for i,j in cycle:\n",
    "        if i < nbx:\n",
    "            Δμ_x_y[i,j-nbx] = 1\n",
    "            m = min(m,μ_x_y[i,j-nbx])\n",
    "        else:\n",
    "            Δμ_x_y[j,i-nbx] = -1\n",
    "            m = min(m,μ_x_y[j,i-nbx])\n",
    "    μ_x_y -= m* Δμ_x_y\n",
    "    return True\n",
    "\n",
    "def purify(theμ_x_y):\n",
    "    μ_x_y = theμ_x_y.copy()\n",
    "    cont = True\n",
    "    while cont:\n",
    "        cont = improve(μ_x_y)\n",
    "    return(μ_x_y.round())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.random.seed(107)\n",
    "d = 10\n",
    "theμ_x_y = 0.3*np.eye(d)[:,np.random.permutation(d)]+0.2*np.eye(d)[:,np.random.permutation(d)]+0.1*np.eye(d)[:,np.random.permutation(d)]+0.25*np.eye(d)[:,np.random.permutation(d)]+0.15*np.eye(d)[:,np.random.permutation(d)]\n",
    "purify(theμ_x_y)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Exercise.** Integral dual solutions. ([OTME], ex. 3.2) Assume $c_\\{xy\\}∈{0,1}$ and consider the problem<br>\n",
    "$\\begin{eqnarray}\n",
    "V &=&\\max_{u,v}\\sum_{x}u_{x}-\\sum_{y}v_{y}.  \\\\\n",
    "&&s.t.~u_{x}-v_{y}\\leq c_{xy}  \\notag\n",
    "\\end{eqnarray}\n",
    "$<br>\n",
    "\n",
    "(i) Show that this is a variant of the optimal transport problem seen in class.<br>\n",
    "(ii) Show that the value of problem is unchanged if one restricts the entries of u and v to be between 0 and 1.<br>\n",
    "(iii) Show that if $(u_{x},v_{y})\\in [0,1]^2$, then<br>\n",
    "$\t(u_{x},v_{y})=\\int_0^1(1\\{t\\leq u_x\\},1\\{t\\leq v_y\\})dt.$<br>\n",
    "(iv) Assume $(u,v)$ is feasible for problem above (i.e. it satisfies the constraints). Show that for each $t\\in[0,1]$, the vectors $(u^{t},v^{t})$ defined by $u_{x}^{t}=1\\{t≤u_{x}\\}$ and $v_{y}^{t}=1\\{t≤v_{y}\\}$ are also feasible.<br>\n",
    "(v) Write a program that in the above setting, takes a solution  $(u_{x},v_{y})\\in [0,1]^2$ and returns a solution  $(u_{x},v_{y})\\in \\{0,1\\}^2$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# answer:\n",
    "\n",
    "# t=0.5\n",
    "\n",
    "# ut_x = np.int64(t <= u_x)\n",
    "# vt_y = np.int64(t <= v_y)"
   ]
  }
 ],
 "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.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}