match_transport.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "173b88d3",
   "metadata": {},
   "source": [
    "# Composite Sorting"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7dc438f8",
   "metadata": {},
   "source": [
    "## Overview\n",
    "\n",
    "Optimal transport theory is studies how one (marginal) probabilty measure can be related to another (marginal) probability measure in an ideal way.\n",
    "\n",
    "The output of such a theory is a **coupling** of the two probability measures, i.e., a joint probabilty\n",
    "measure having those two  marginal probability measures.\n",
    "\n",
    "This lecture describes how Job Boerma, Aleh Tsyvinski, Ruodo Wang,\n",
    "and Zhenyuan Zhang  [[Boerma *et al.*, 2024](https://python-advanced.quantecon.org/zreferences.html#id3)] used optimal transport theory to formulate and solve an equilibrium of a model in which wages and allocations of workers across jobs  adjust to match measures of  different types with measures of different types of occupations.\n",
    "\n",
    "Production technologies allow firms to affect  shape costs of mismatch with the consequence\n",
    "that costs of mismatch can be concave.\n",
    "\n",
    "That means that it possible that equilibrium there is neither **positive assortive** nor **negative assorting**  matching, an outcome that   [[Boerma *et al.*, 2024](https://python-advanced.quantecon.org/zreferences.html#id3)] call **composite assortive** matching.\n",
    "\n",
    "For example, in  an equilibrium with composite matching,  identical **workers** can sort into different **occupations**, some positively and some negatively.\n",
    "\n",
    "[[Boerma *et al.*, 2024](https://python-advanced.quantecon.org/zreferences.html#id3)]\n",
    "show how this can generate distinct distributions  of labor earnings  within and across occupations.\n",
    "\n",
    "This lecture describes the [[Boerma *et al.*, 2024](https://python-advanced.quantecon.org/zreferences.html#id3)] model and  presents  Python code for computing equilibria.\n",
    "\n",
    "The lecture  applies the code to the [[Boerma *et al.*, 2024](https://python-advanced.quantecon.org/zreferences.html#id3)] model of labor markets.\n",
    "\n",
    "As with an earlier QuantEcon lecture on optimal transport ([https://python.quantecon.org/opt_transport.html](https://python.quantecon.org/opt_transport.html)), a key tool will be **linear programming**."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0a05795d",
   "metadata": {},
   "source": [
    "## Setup\n",
    "\n",
    "$ X $ and $ Y $ are finite sets that represent two distinct types of people to be matched.\n",
    "\n",
    "For each $ x \\in X, $ let a positive integer $ n_x $ be the number  of agents of type $ x $.\n",
    "\n",
    "Similarly, let a positive integer $ m_y $ be the agents of agents of type $ y \\in Y $.\n",
    "\n",
    "We refer to these two measures as *marginals*.\n",
    "\n",
    "We assume that\n",
    "\n",
    "$$\n",
    "\\sum_{x \\in X} n_x = \\sum_{y \\in Y} m_y  =: N\n",
    "$$\n",
    "\n",
    "so that  the matching  problem is *balanced*.\n",
    "\n",
    "Given a *cost function* $ c:X \\times Y \\rightarrow \\mathbb{R} $, the (discrete) *optimal transport problem* is\n",
    "\n",
    "$$\n",
    "\\begin{aligned}\n",
    "\\min_{\\mu \\geq 0}& \\sum_{(x,y) \\in X \\times Y} \\mu_{xy}c_{xy} \\\\\n",
    "\\text{s.t. }& \\sum_{x \\in X} \\mu_{xy} = n_x \\\\\n",
    "& \\sum_{y \\in Y} \\mu_{xy} = m_y \n",
    "\\end{aligned}\n",
    "$$\n",
    "\n",
    "Given our discreteness  assumptions about $ n $ and $ m $, the problem admits an integer solution $ \\mu \\in \\mathbb{Z}_+^{X \\times Y} $, i.e. $ \\mu_{xy} $ is a non-negative integer for each $ x\\in X, y\\in Y $.\n",
    "\n",
    "We will study integer solutions.\n",
    "\n",
    "Two points about restricting ourselves to integer solutions are worth mentioning:\n",
    "\n",
    "- it is without loss of generality for computational purposes, since every problem with float marginals can be transformed into an equivalent problem with integer marginals;  \n",
    "- although the mathematical structure that we present actually   wors for arbitrary real marginals, some of our Python  implementations would  fail to work with float arithmetic.  \n",
    "\n",
    "\n",
    "We focus on  a specific instance of an  optimal transport problem:\n",
    "\n",
    "We assume that $ X $ and $ Y $ are finite subsets of $ \\mathbb{R} $ and that the cost function satisfies $ c_{xy} = h(|x - y|) $ for all $ x,y \\in \\mathbb{R}, $ for an $ h: \\mathbb{R}_+ \\rightarrow \\mathbb{R}_+ $ that  is **strictly concave** and **strictly increasing** and **grounded** (i.e., $ h(0)=0 $).\n",
    "\n",
    "Such an  $ h $ satisfies the following\n",
    "\n",
    "**Lemma.** Let $ h: \\mathbb{R}_+ \\rightarrow \\mathbb{R}_+ $ be **strictly concave** and **grounded**. Then $ h $ is strictly subadditive, i.e. for all $ x,y\\in \\mathbb{R}_+, 0< x < y, $ we have\n",
    "\n",
    "$$\n",
    "h(x+y) < h(x) + h(y)\n",
    "$$\n",
    "\n",
    "*Proof.* For $ \\alpha \\in (0,1) $ and $ x >0 $ we have, by strict concavity and groundedness, $ h(\\alpha x) > \\alpha h(x) + (1-\\alpha) h(0)=\\alpha h(x). $\n",
    "\n",
    "Now fix $ x,y\\in \\mathbb{R}_+, 0< x < y, $ and let $ \\alpha = \\frac{x}{x+y}; $ the previous observation gives $ h(x) = h(\\alpha (x+y)) > \\alpha h(x+y) $ and $ h(y) = h((1-\\alpha) (x+y)) > (1-\\alpha) h(x+y) $;  summing  these inequality delivers the result. $ \\square $\n",
    "\n",
    "In the following implementation we assume that the cost function is $ c_{xy} = |x-y|^{1/\\zeta} $ for $ \\zeta>1, $ i.e. $ h(z) =  z^{1/\\zeta} $ for $ z \\in \\mathbb{R}_+. $\n",
    "\n",
    "Hence, our problem is\n",
    "\n",
    "$$\n",
    "\\begin{aligned}\n",
    "\\min_{\\mu \\in \\mathbb{Z}_+^{X \\times Y}}& \\sum_{(x,y) \\in X \\times Y} \\mu_{xy}|x-y|^{1/\\zeta} \\\\\n",
    "\\text{s.t. }& \\sum_{x \\in X} \\mu_{xy} = n_x \\\\\n",
    "& \\sum_{y \\in Y} \\mu_{xy} = m_y \n",
    "\\end{aligned}\n",
    "$$\n",
    "\n",
    "Let’s start setting up some Python code.\n",
    "\n",
    "We  use the following imports:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8af6b529",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from scipy.optimize import linprog\n",
    "from itertools import chain\n",
    "import pandas as pd\n",
    "from collections import namedtuple\n",
    "\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.patches as patches\n",
    "from matplotlib.ticker import MaxNLocator\n",
    "from matplotlib import cm\n",
    "from matplotlib.colors import Normalize"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "79467e74",
   "metadata": {},
   "source": [
    "The following Python class takes as inputs sets of types $ X,Y \\subset \\mathbb{R}, $ marginals $ n, m $ with positive integer entries such that $ \\sum_{x \\in X} n_x = \\sum_{y \\in Y} m_y $ and cost parameter $ \\zeta>1 $.\n",
    "\n",
    "The cost function is stored as an $ |X| \\times |Y| $ matrix with $ (x,y) $-entry equal to $ |x-y|^{1/\\zeta}, $ i.e., the cost of matching an agent of type $ x \\in X $ with an agent of type $ y \\in Y. $"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ec037bf1",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "class ConcaveCostOT():\n",
    "    def __init__(self, X_types=None, Y_types=None, n_x =None, m_y=None, ζ=2):\n",
    "        \n",
    "        # Sets of types \n",
    "        self.X_types, self.Y_types = X_types, Y_types\n",
    "        \n",
    "        # Marginals\n",
    "        if X_types is not None and Y_types is not None:\n",
    "            non_empty_types = True\n",
    "            self.n_x = np.ones(len(X_types), dtype=int) if n_x is None else n_x\n",
    "            self.m_y = np.ones(len(Y_types), dtype=int) if m_y is None else m_y\n",
    "        else:\n",
    "            non_empty_types = False\n",
    "            self.n_x, self.m_y = n_x, m_y\n",
    "\n",
    "        # Cost function: |X|x|Y| matrix\n",
    "        self.ζ = ζ\n",
    "        if non_empty_types:\n",
    "            self.cost_x_y = np.abs(X_types[:, None] - Y_types[None, :]) \\\n",
    "            ** (1 / ζ)\n",
    "        else:\n",
    "            self.cost_x_y = None "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "07fb1139",
   "metadata": {},
   "source": [
    "Let’s consider a random instance with given numbers of types $ |X| $ and $ |Y| $ and a given number of agents.\n",
    "\n",
    "First, we generate random types $ X $ and $ Y. $\n",
    "\n",
    "Then we generate random quantities for each type so that there are $ N $ agents for each side."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "622542f4",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "number_of_x_types = 20\n",
    "number_of_y_types = 20\n",
    "N_agents_per_side = 60\n",
    "\n",
    "np.random.seed(1)\n",
    "\n",
    "## Genetate random types\n",
    "# generate random support for distributions of types\n",
    "support_size = 50\n",
    "random_support = np.unique(np.random.uniform(0,200, size=support_size))\n",
    "\n",
    "# generate types\n",
    "X_types_example = np.random.choice(random_support, \n",
    "                    size=number_of_x_types, replace=False)\n",
    "Y_types_example = np.random.choice(random_support, \n",
    "                    size=number_of_y_types, replace=False)\n",
    "\n",
    "## Generate random integer types quantities summing to N_agents_per_side\n",
    "\n",
    "# generate integer vectors of lenght n_types summing to n_agents\n",
    "def random_marginal(n_types, n_agents):\n",
    "    cuts = np.sort(np.random.choice(np.arange(1,n_agents), \n",
    "                                    size= n_types-1, replace=False))\n",
    "    segments = np.diff(np.concatenate(([0], cuts, [n_agents])))\n",
    "    return segments\n",
    "\n",
    "# Create a method to assign random marginals to our class\n",
    "def assign_random_marginals(self,random_seed):\n",
    "    np.random.seed(random_seed)\n",
    "    self.n_x = random_marginal(len(self.X_types), N_agents_per_side)\n",
    "    self.m_y = random_marginal(len(self.Y_types), N_agents_per_side)\n",
    "\n",
    "ConcaveCostOT.assign_random_marginals = assign_random_marginals\n",
    "\n",
    "# Create an instance of our class and generate random marginals\n",
    "example_pb = ConcaveCostOT(X_types_example, Y_types_example, ζ=2)\n",
    "example_pb.assign_random_marginals(random_seed=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "421a16cf",
   "metadata": {},
   "source": [
    "We use  $ F $ (resp. $ G $) to denote the cumulative distribution function associated to the measure $ n $ (resp. $ m $)\n",
    "\n",
    "Thus, $ F(z) =\\sum_{x \\leq z: n_x > 0} n_x $ and $ G(z) =\\sum_{y \\leq z: m_y > 0} m_y $ for $ z\\in \\mathbb{R}. $\n",
    "\n",
    "Notice that we not normalizing the measures so $ F(\\infty) = G(\\infty) =N. $\n",
    "\n",
    "The following method plots the marginals on the real line\n",
    "\n",
    "- blue for $ X $ types,  \n",
    "- red for $ Y $ types.  \n",
    "\n",
    "\n",
    "Note that there are possible overlaps between $ X $ and $ Y. $"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c586c87f",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def plot_marginals(self, figsize=(15, 8), title='Distributions of types'):\n",
    "\n",
    "    plt.figure(figsize=figsize)\n",
    "    \n",
    "    # Scatter plot n_x\n",
    "    plt.scatter(self.X_types, self.n_x, color='blue', label='n_x')\n",
    "    plt.vlines(self.X_types, ymin=0, ymax= self.n_x, \n",
    "               color='blue', linestyles='dashed')\n",
    "    \n",
    "    # Scatter plot m_y\n",
    "    plt.scatter(self.Y_types, - self.m_y, color='red', label='m_y')\n",
    "    plt.vlines(self.Y_types, ymin=0, ymax=- self.m_y, \n",
    "               color='red', linestyles='dashed')\n",
    "\n",
    "    # Add grid and y=0 axis\n",
    "    plt.grid(True)\n",
    "    plt.axhline(0, color='black', linewidth=1)\n",
    "    plt.gca().spines['bottom'].set_position(('data', 0))\n",
    "\n",
    "    # Labeling the axes and the title\n",
    "    plt.ylabel('frequency')\n",
    "    plt.title(title)\n",
    "    plt.gca().yaxis.set_major_locator(MaxNLocator(integer=True))\n",
    "    plt.legend()\n",
    "    plt.show()\n",
    "    \n",
    "ConcaveCostOT.plot_marginals = plot_marginals"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1dd612c3",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "example_pb.plot_marginals()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6bedddac",
   "metadata": {},
   "source": [
    "## Characterization of primal solution"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9b70a70f",
   "metadata": {},
   "source": [
    "### Three properties of an optimal solution\n",
    "\n",
    "We now indicate important properties that  are satisfied by an optimal solution.\n",
    "\n",
    "1. Maximal number of perfect pairs  \n",
    "1. No intersecting pairs  \n",
    "1. Layering  \n",
    "\n",
    "\n",
    "**(Maximal number of perfect pairs)**\n",
    "\n",
    "If $ (z,z) \\in X \\times Y $ for some $ z \\in \\mathbb{R} $ then in each optimal solution there are $ \\min\\{n_z,m_z\\} $ matches between type $ z \\in X $ and $ z \\in Y $.\n",
    "\n",
    "Indeed, assume by contradiction that at an optimal solution we have $ (z,y) $ and $ (x,z) $ matched in positive amounts for $ y,x \\neq z $.\n",
    "\n",
    "We can verify that reassigning the minimum of such quantities to the pairs $ (z,z) $ and $ (x,y) $ improves upon the current matching since\n",
    "\n",
    "$$\n",
    "h(|x-y|) \\leq h(|x-z| +|z - y|) < h(|x-z|)+ h(|z - y|)\n",
    "$$\n",
    "\n",
    "where the first inequality follows from triangle inequality and the fact that $ h $ is increasing and the strict inequality from strict subadditivity.\n",
    "\n",
    "We can then repeat the operation for any other analogous pair of matches involving $ z, $ while improving  the value, until we have  mass $ \\min\\{n_z,m_z\\} $ on  match $ (z,z). $\n",
    "\n",
    "Viewing the matching $ \\mu $ as a measure on $ X \\times Y $ with marginals $ n $ and $ m $, this property says that in any optimal $ \\mu $ we have $ \\mu_{zz} = n_z \\wedge m_z $ for $ (z,z) $ in the diagonal $ \\{(x,y) \\in X \\times Y: x=y\\} $ of $ \\mathbb{R} \\times \\mathbb{R} $.\n",
    "\n",
    "The following method finds perfect pairs and returns the on-diagonal matchings as well as the residual off-diagonal marginals."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "737c5da4",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def match_perfect_pairs(self):\n",
    "\n",
    "    # Find pairs on diagonal and related mass\n",
    "    perfect_pairs_x, perfect_pairs_y = np.where(\n",
    "                        self.X_types[:,None] == self.Y_types[None,:])\n",
    "    Δ_q = np.minimum(self.n_x[perfect_pairs_x] ,self.m_y[perfect_pairs_y])\n",
    "\n",
    "    # Compute off-diagonal residual masses for each side\n",
    "    n_x_off_diag = self.n_x.copy()\n",
    "    n_x_off_diag[perfect_pairs_x]-= Δ_q\n",
    "\n",
    "    m_y_off_diag = self.m_y.copy()\n",
    "    m_y_off_diag[perfect_pairs_y] -= Δ_q\n",
    "\n",
    "    # Compute on-diagonal matching\n",
    "    matching_diag = np.zeros((len(self.X_types), len(self.Y_types)), dtype=int)\n",
    "    matching_diag[perfect_pairs_x, perfect_pairs_y] = Δ_q\n",
    "    \n",
    "    return n_x_off_diag, m_y_off_diag , matching_diag\n",
    "\n",
    "ConcaveCostOT.match_perfect_pairs = match_perfect_pairs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5c45f369",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "n_x_off_diag, m_y_off_diag , matching_diag = example_pb.match_perfect_pairs()\n",
    "print(f\"On-diagonal matches: {matching_diag.sum()}\")\n",
    "print(f\"Residual types in X: {len(n_x_off_diag[n_x_off_diag >0])}\")\n",
    "print(f\"Residual types in Y: {len(m_y_off_diag[m_y_off_diag >0])}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f1e7d8b7",
   "metadata": {},
   "source": [
    "We can therefore create a new instance with the residual marginals that will feature no perfect pairs.\n",
    "\n",
    "Later we shall   add the on-diagonal matching to the solution of this new instance.\n",
    "\n",
    "We refer to this instance as “off-diagonal” since the product measure of the residual marginals $ n \\otimes m $ feature zeros mass on the diagonal of $ \\mathbb{R} \\times \\mathbb{R}. $\n",
    "\n",
    "In the rest of this section, we will focus on  this instance.\n",
    "\n",
    "We create a subclass to study the residual off-diagonal problem.\n",
    "\n",
    "The subclass inherits the attributes and the modules from the original class.\n",
    "\n",
    "We let $ Z := X \\sqcup Y , $ where $ \\sqcup $ denotes the union of disjoint sets. We will\n",
    "\n",
    "- index types $ X $ as $ \\{0, \\dots,|X|-1\\} $ and types $ Y $ as $ \\{|X|, \\dots,|X| + |Y|-1\\}; $  \n",
    "- store the cost function as a $ |Z| \\times |Z| $ matrix with entry $ (z,z') $ equal to $ c_{xy} $ if $ z=x \\in X $ and $ z' =y\\in Y $ or $ z=y \\in Y $ and $ z' =x\\in X $ or equal to $ +\\infty $ if $ z $ and $ z' $ belong to the same side  \n",
    "  - (the latter is just customary, since these “infinitely penalized” entries are actually never accessed in the implementation);  \n",
    "- let $ q $ be a vector of size $ |Z| $ whose $ z $-th entry equals $ n_x $ if type $ x $ is the $ z $-th smallest type in $ Z $ and $ -m_y $ if type $ y $ is the $ z $-th smallest type in $ Z $; hence $ q $ encodes capacities of both sides on the (ascending) sorted set of types.  \n",
    "\n",
    "\n",
    "Finally, we add a method to flexibly add a pair $ (i,j) $ with $ i \\in \\{0, \\dots,|X|-1\\},j \\in \\{|X|, \\dots,|X| + |Y|-1\\} $ or $ j \\in \\{0, \\dots,|X|-1\\},i \\in \\{|X|, \\dots,|X| + |Y|-1\\} $ to a matching matrix of size $ |X| \\times |Y| $."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "93b922c6",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "class OffDiagonal(ConcaveCostOT):\n",
    "    def __init__(self, X_types, Y_types, n_x, m_y, ζ):\n",
    "        super().__init__(X_types, Y_types, n_x, m_y, ζ)\n",
    "\n",
    "        # Types (unsorted)\n",
    "        self.types_list = np.concatenate((X_types,Y_types))\n",
    "\n",
    "        # Cost function: |Z|x|Z| matrix \n",
    "        self.cost_z_z = np.ones((len(self.types_list),\n",
    "                                 len(self.types_list))) * np.inf\n",
    "\n",
    "        # upper-right block\n",
    "        self.cost_z_z[:len(self.X_types), len(self.X_types):] = self.cost_x_y \n",
    "        \n",
    "        # lower-left block\n",
    "        self.cost_z_z[len(self.X_types):, :len(self.X_types)] = self.cost_x_y.T\n",
    "\n",
    "        ## Distributions of types\n",
    "        # sorted types and index identifier for each z in support\n",
    "        self.type_z = np.argsort(self.types_list)\n",
    "        self.support_z = self.types_list[self.type_z]\n",
    "\n",
    "        # signed quantity for each type z\n",
    "        self.q_z = np.concatenate([n_x, - m_y])[self.type_z] \n",
    "\n",
    "    # Mathod that adds to matching matrix a pair (i,j)\n",
    "    def add_pair_to_matching(self, pair_ids, matching):\n",
    "        if pair_ids[0] < pair_ids[1]:\n",
    "            # the pair of indices correspond to a pair (x,y)\n",
    "            matching[pair_ids[0], pair_ids[1]-len(self.X_types)] = 1\n",
    "        else:\n",
    "            # the pair of indices correspond to a pair (y,x)\n",
    "            matching[pair_ids[1], pair_ids[0]-len(self.X_types)] = 1  "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "41e1f7f0",
   "metadata": {},
   "source": [
    "We add a function that returns an instance of the off-diagonal subclass as well as the on-diagonal matching and the indices of the residual off-diagonal types.\n",
    "\n",
    "These indices will come handy for adding the off-diagonal matching matrix to the diagonal matching matrix we just found, since the former will have a smaller size if there are perfect pairs in the original problem."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4f96765f",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def generate_offD_onD_matching(self):\n",
    "    # Match perfect pairs and compute on-diagonal matching\n",
    "    n_x_off_diag, m_y_off_diag , matching_diag = self.match_perfect_pairs()\n",
    "\n",
    "    # Find indices of residual non-zero quantities for each side\n",
    "    nonzero_id_x = np.flatnonzero(n_x_off_diag)\n",
    "    nonzero_id_y = np.flatnonzero(m_y_off_diag)\n",
    "\n",
    "    # Create new instance with off-diagonal types\n",
    "    off_diagonal = OffDiagonal(self.X_types[nonzero_id_x], \n",
    "                                    self.Y_types[nonzero_id_y], \n",
    "                                    n_x_off_diag[nonzero_id_x], \n",
    "                                    m_y_off_diag[nonzero_id_y], \n",
    "                                    self.ζ)\n",
    "        \n",
    "    return off_diagonal, (nonzero_id_x, nonzero_id_y, matching_diag)\n",
    "\n",
    "ConcaveCostOT.generate_offD_onD_matching = generate_offD_onD_matching"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e2095837",
   "metadata": {},
   "source": [
    "We apply it to our example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "806db2b0",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "example_off_diag, _ = example_pb.generate_offD_onD_matching()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9f17b51a",
   "metadata": {},
   "source": [
    "Let’s plot the residual marginals to verify visually that there are no overlappings between types from distinct sides in the off-diagonal instance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4c3754c3",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "example_off_diag.plot_marginals(title='Distributions of types: off-diagonal')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a3e07441",
   "metadata": {},
   "source": [
    "**(No intersecting pairs)** This  property summarizes  the following fact:\n",
    "\n",
    "- represent both types  on the real line and draw a semicirle joining $ (x,y) $ for all pairs $ (x,y) \\in X \\times Y $ that are matched in a solution  \n",
    "- these semicirles do not intersect (unless they share one of the endpoints).  \n",
    "\n",
    "\n",
    "A proof proceeds by contradiction.\n",
    "\n",
    "Let’s consider types $ x,x' \\in X $ and $ y,y' \\in Y. $\n",
    "\n",
    "Matched pairs cain “intersect” (or be  tangent).\n",
    "\n",
    "We will show that in both cases the partial matching among types $ x,x',y,y' $ can be improved by *uncrossing*, i.e. reassigning the quantities while improving on the solution and reducing the number of intersecting pairs.\n",
    "\n",
    "The first case of intersecting pairs is\n",
    "\n",
    "$$\n",
    "x < y < y' < x'\n",
    "$$\n",
    "\n",
    "with pairs $ (x,y') $ and $ (x',y) $ matched in positive quantities.\n",
    "\n",
    "Then it follows from strict monotonicity of $ h $ that $ h(|x-y|) < h(|x-y'|) $ and $ h(|x'-y'|) < h(|x'-y|), $ hence $ h(|x-y|)+ h(|x'-y'|) < h(|x-y'|) + h(|x'-y|). $\n",
    "\n",
    "Therefore, we can take the minimum of the masses of the matched pairs $ (x,y') $ and $ (x',y) $ and  reallocate it to the pairs $ (x,y) $ and $ (x',y') $,\n",
    "therby  strictly improving the cost among $ x,y,x',y'. $\n",
    "\n",
    "The second case of intersecting pairs is\n",
    "\n",
    "$$\n",
    "x < x' < y' < y\n",
    "$$\n",
    "\n",
    "with pairs $ (x,y') $ and $ (x',y) $ matched.\n",
    "\n",
    "In this case we have\n",
    "\n",
    "$$\n",
    "|x - y'| + |x' - y| = |x - y| + |x' - y'|\n",
    "$$\n",
    "\n",
    "Letting $ \\alpha := \\frac{|x - y|+|x' - y|}{|x - y'| - |x' - y|} \\in (0,1), $ we have $ |x - y| = \\alpha|x - y'| +(1-\\alpha) |x' - y| $ and $ |x' - y'| = (1-\\alpha)|x - y'| +\\alpha |x' - y|. $\n",
    "\n",
    "Hence, by strict concavity of $ h, $\n",
    "\n",
    "$$\n",
    "h(|x-y|)+ h(|x'-y'|) <\\alpha h(|x - y'|) +(1-\\alpha) h(|x' - y|) +  (1-\\alpha) h(|x - y'|) +\\alpha h(|x' - y|) = h(|x-y'|) + h(|x'-y|).\n",
    "$$\n",
    "\n",
    "Therefore, as in the first case, we can strictly improve the cost among $ x,y,x',y' $ by uncrossing the pairs.\n",
    "\n",
    "Finally, it remains to argue that in both cases  *uncrossing* operations do not increase the number of intersections with other matched pairs.\n",
    "\n",
    "It can indeed be shown on a case-by-case basis that, in both of the above cases, for any other matched pair $ (x'',y'') $ the number of intersections between pairs $ (x,y), (x',y') $ and the pair $ (x'',y'') $ (i.e., after uncrossing) is not larger than the number of intersections between pairs $ (x,y'), (x',y) $ and the pair $ (x'',y'') $ (i.e., before uncrossing), hence the uncrossing operations above reduce the number of intersections.\n",
    "\n",
    "We conclude that if a matching features intersecting pairs, it can be modified via a sequence of uncrossing operations into a matching without intersecting pairs while improving on the value.\n",
    "\n",
    "**(Layering)** Recall that there are $ 2N $ individual agents, each agent $ i $ having type $ z_i \\in X \\sqcup Y. $\n",
    "\n",
    "When we introduce the off diagonal matching, to stress that the types sets are disjoint now.\n",
    "\n",
    "To simplify our explanation  of this property, assume for now that each agent has its own distinct type (i.e., |X|=|Y| =N and $ n=m= \\mathbf{1}_N $), in which case the optimal transport problem is also referred to as *assignment problem*.\n",
    "\n",
    "Let’s index  agents according to their types:\n",
    "\n",
    "$$\n",
    "z_1 < z_2 \\dots<  z_{2N-1} < z_{2N}.\n",
    "$$\n",
    "\n",
    "Suppose that agents $ i $ of type $ z_i $ and $ j $ of type $ z_j $, with $ z_i < z_j, $ are matched in a particular optimal solution.\n",
    "\n",
    "Then there is an equal number of agents from each side in $ \\{i+1, \\dots, j-1\\}, $ if this set is not empty.\n",
    "\n",
    "Indeed, if this were not the case, then some agent $ k \\in \\{i+1,j-1\\} $ would be  matched with some agent $ \\ell $ with $ \\ell \\notin \\{i,\\dots, j\\}, $ i.e., there would be  types\n",
    "\n",
    "$$\n",
    "z_i < z_k < z_j < z_\\ell\n",
    "$$\n",
    "\n",
    "with matches $ (z_i,z_j) $ and $ (z_k, z_\\ell), $ violating the no intersecting pairs property.\n",
    "\n",
    "We conclude that we can define a binary relation on $ [N] $ such that $ i \\sim j $ if there is an equal number of agents of each side in $ \\{i,i+1,\\dots, j\\} $ (or if this set is empty).\n",
    "\n",
    "This is an equivalence relation, so we can find associated equivalence classes that we call *layers*.\n",
    "\n",
    "By the reasoning above, in an optimal solution all pairs $ i,j $ (of opposite sides) which are matched belong to the same layer, hence we can solve the assignment problem associated to each layer and then add up the solutions.\n",
    "\n",
    "In terms of distributions, $ i $ and $ j, $ of types $ x \\in X $ and $ y \\in Y $ respectively, belong to the same layer (i.e., $ x \\sim y $) if and only if $ F(y-) - F(x) = G(y-) - G(x). $\n",
    "\n",
    "If  $ F $ and $ G $ were continuous, then  $ F(y) - F(x) = G(y) - G(x) \\iff F(x) - G(x) = F(y) - G(y). $\n",
    "\n",
    "This suggests that the following quantity plays an important role:\n",
    "\n",
    "$$\n",
    "H(z) := F(z) - G(z), \\text{ for } z \\in \\mathbb{R}.\n",
    "$$\n",
    "\n",
    "Returning to our general (integer) discrete setting, let’s plot $ H $.\n",
    "\n",
    "Notice that $ H $ is right-continuous (being the difference of right-continuous functions) and that upward (resp. downward) jumps correspond to point masses of agents with types from $ X $ (resp. $ Y $)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7689df67",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def plot_H_z(self, figsize=(15, 8), range_x_axis=None, scatter=True):\n",
    "    # Determine H(z) = F(z) - G(z)\n",
    "    H_z = np.cumsum(self.q_z)\n",
    "    \n",
    "    # Plot H(z)\n",
    "    plt.figure(figsize=figsize)\n",
    "    plt.axhline(0, color='black', linewidth=1)\n",
    "    \n",
    "    # determine the step points for horizontal lines\n",
    "    step = np.concatenate(([self.support_z.min() - .05 * self.support_z.ptp()],\n",
    "                            self.support_z,\n",
    "                           [self.support_z.max() + .05 * self.support_z.ptp()]))\n",
    "    height = np.concatenate(([0], H_z, [0]))\n",
    "    \n",
    "    # plot the horizontal lines of the step function\n",
    "    for i in range(len(step) - 1):\n",
    "        plt.plot([step[i], step[i+1]], [height[i], height[i]], color='black')\n",
    "    \n",
    "    # draw dashed vertical lines for the step function\n",
    "    for i in range(1, len(step) - 1):\n",
    "        plt.plot([step[i], step[i]], [height[i-1], height[i]], \n",
    "                color='black', linestyle='--')\n",
    "    \n",
    "    # plot discontinuities points of H(z)\n",
    "    if scatter:\n",
    "        plt.scatter(np.sort(self.X_types), H_z[self.q_z > 0], color='blue')\n",
    "        plt.scatter(np.sort(self.Y_types), H_z[self.q_z < 0], color='red')\n",
    "    \n",
    "    if range_x_axis is not None:\n",
    "        plt.xlim(range_x_axis)\n",
    "\n",
    "    # Add labels and title\n",
    "    plt.title('Underqualification Measure (Off-Diagonal)')\n",
    "    plt.xlabel('$z$')\n",
    "    plt.ylabel('$H(z)$')\n",
    "    plt.grid(False)\n",
    "    plt.gca().yaxis.set_major_locator(MaxNLocator(integer=True))\n",
    "    plt.show()\n",
    "\n",
    "OffDiagonal.plot_H_z = plot_H_z"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2a2196ca",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "example_off_diag.plot_H_z()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6be932c0",
   "metadata": {},
   "source": [
    "The layering property extends to the general discrete setting.\n",
    "\n",
    "There are $ |H(\\mathbb{R})|-1 $ layers in total.\n",
    "\n",
    "Enumerating the range of $ H $ as $ H(\\mathbb{R}) = \\{h_1,h_2, \\dots, h_{|H(\\mathbb{R})|}\\} $ with $ h_1 < h_2 < \\dots < h_{|H(\\mathbb{R})|}, $ we can define layer $ L_\\ell, $ for $ \\ell \\in \\{ 1,\\dots,|H(\\mathbb{R})|-1\\} $ as the collection of types $ z \\in Z $ such that\n",
    "\n",
    "$$\n",
    "H(z-) \\leq h_{\\ell -1} < h_{\\ell } \\leq H(z),\n",
    "$$\n",
    "\n",
    "(which are types in $ X $), or\n",
    "\n",
    "$$\n",
    "H(z) \\leq h_{\\ell -1} < h_{\\ell } \\leq H(z-),\n",
    "$$\n",
    "\n",
    "which are types in $ Y $.\n",
    "\n",
    "The *mass* associated with  layer $ L_\\ell $ is $ M_\\ell = h_{\\ell+1}- h_{\\ell}. $\n",
    "\n",
    "Intuitively, a layer $ L_\\ell $ consists of some mass $ M_\\ell, $ of multiple types in $ Z, $ i.e. the problem within the layer is *unitary*.\n",
    "\n",
    "A unitary problem is essentially an assignment problem up to a constant: we can solve the problem with unit mass and then rescale a solution by $ M_\\ell. $\n",
    "\n",
    "Moreover, each layer $ L_\\ell $ contains an even number of types $ N_\\ell \\in 2\\mathbb{N}, $ which are alternating, i.e., ordering them as $ z_1 < z_2\\dots < z_{ N_\\ell-1} < z_{ N_\\ell} $ all odd (or even, respectively) indexed types belong to the same side.\n",
    "\n",
    "The following method finds the layers associated with distributions $ F $ and $ G $.\n",
    "\n",
    "Again, types in $ X $ are indexed with $ \\{0, \\dots,|X|-1\\} $ and types in $ Y $ with $ \\{|X|, \\dots,|X| + |Y|-1\\} $.\n",
    "\n",
    "Using these indices (instead of the types themselves) to represent the layers allows  keeping  track of  sides  types in each layer, without adding an additional bit of information that would  identify the side of the first type in the layer, which, because  a layer is alternating,  would then allow identifying  sides of all  types in the layer.\n",
    "\n",
    "In addition, using indices will let us  extract the cost function within a layer from the cost function $ c_{zz'} $ computed offline."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11463e5f",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def find_layers(self):\n",
    "    # Compute H(z) on the joint support\n",
    "    H_z = np.concatenate([[0], np.cumsum(self.q_z)])\n",
    "\n",
    "    # Compute the range of H, i.e. H(R), stored in ascending order\n",
    "    layers_height = np.unique(H_z)\n",
    "    \n",
    "    # Compute the mass of each layer\n",
    "    layers_mass = np.diff(layers_height)\n",
    "\n",
    "    # Compute layers\n",
    "    # the following |H(R)|x|Z| matrix has entry (z,l) equal to 1 iff type z belongs to layer l\n",
    "    layers_01 = ((H_z[None, :-1] <= layers_height[:-1, None]) \n",
    "                 * (layers_height[1:, None] <= H_z[None, 1:]) |\n",
    "                (H_z[None, 1:] <= layers_height[:-1, None]) \n",
    "                 * (layers_height[1:, None] <= H_z[None, :-1]))\n",
    "    \n",
    "    # each layer is reshaped as a list of indices correponding to types\n",
    "    layers = [self.type_z[layers_01[ell]] \n",
    "                for ell in range(len(layers_height)-1)]\n",
    "\n",
    "    return layers, layers_mass, layers_height, H_z\n",
    "\n",
    "OffDiagonal.find_layers = find_layers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a8de9a2e",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "layers_list_example, layers_mass_example, _, _ = example_off_diag.find_layers()\n",
    "print(layers_list_example)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "25e8a256",
   "metadata": {},
   "source": [
    "The following method gives a graphical representation of the layers.\n",
    "\n",
    "From the  picture it is easy to spot  two key features described  above:\n",
    "\n",
    "- types are alternating  \n",
    "- the layer problem is unitary  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "af486880",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def plot_layers(self, figsize=(15, 8)):\n",
    "    # Find layers\n",
    "    layers, layers_mass , layers_height, H_z = self.find_layers()\n",
    "    \n",
    "    plt.figure(figsize=figsize)\n",
    "\n",
    "    # Plot H(z)\n",
    "    step = np.concatenate(([self.support_z.min() - .05 * self.support_z.ptp()],\n",
    "                           self.support_z,\n",
    "                           [self.support_z.max() + .05 * self.support_z.ptp()]))\n",
    "    height = np.concatenate((H_z, [0]))\n",
    "    plt.step(step, height, where='post', color='black', label='CDF', zorder=1)\n",
    "    \n",
    "    # Plot layers\n",
    "    colors = cm.viridis(np.linspace(0, 1, len(layers))) \n",
    "    for  ell, layer  in enumerate(layers):\n",
    "        plt.vlines(self.types_list[layer], layers_height[ell] , \n",
    "                   layers_height[ell] + layers_mass[ell], \n",
    "                   color=colors[ell], linewidth=2)\n",
    "        plt.scatter(self.types_list[layer], \n",
    "                    np.ones(len(layer)) * layers_height[ell] \n",
    "                    +.5 * layers_mass[ell], \n",
    "                    color=colors[ell], s=50)\n",
    "        \n",
    "        plt.axhline(layers_height[ell], color=colors[ell], \n",
    "                    linestyle=':', linewidth=1.5, zorder=0)\n",
    "\n",
    "    # Add labels and title\n",
    "    plt.xlabel('$z$')\n",
    "    plt.title('Layers')\n",
    "    plt.gca().yaxis.set_major_locator(MaxNLocator(integer=True))\n",
    "    plt.show()\n",
    "\n",
    "OffDiagonal.plot_layers = plot_layers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "28c62963",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "example_off_diag.plot_layers()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "60fc02d3",
   "metadata": {},
   "source": [
    "### Solving a layer\n",
    "\n",
    "Recall that layer $ L_\\ell $ consists of a list of distinct types from $ Y \\sqcup X $\n",
    "\n",
    "$$\n",
    "z_1 < z_2\\dots < z_{N_\\ell-1} < z_{N_\\ell},\n",
    "$$\n",
    "\n",
    "which is alternating.\n",
    "\n",
    "The problem within a layer is unitary.\n",
    "\n",
    "Hence we can solve the problem with unit masses and later rescale the solution by the layer’s mass $ M_\\ell $.\n",
    "\n",
    "Let us select a layer from the example above (we pick the one with maximum number of types) and plot the types on the real line"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6bbf357f",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "# Pick layer with maximum number of types\n",
    "layer_id_example = max(enumerate(layers_list_example), \n",
    "                    key = lambda x: len(x[1]))[0]\n",
    "layer_example = layers_list_example[layer_id_example]\n",
    "\n",
    "\n",
    "# Plot layer types\n",
    "def plot_layer_types(self, layer, mass, figsize=(15, 3)):\n",
    "\n",
    "    plt.figure(figsize=figsize)\n",
    "    \n",
    "    # Scatter plot n_x\n",
    "    x_layer = layer[layer < len(self.X_types)]\n",
    "    y_layer = layer[layer >= len(self.X_types)] - len(self.X_types)\n",
    "    M_ell = np.ones(len(x_layer))* mass\n",
    "\n",
    "    plt.scatter(self.X_types[x_layer], M_ell, color='blue', label='X types')\n",
    "    plt.vlines(self.X_types[x_layer], ymin=0, ymax= M_ell, \n",
    "               color='blue', linestyles='dashed')\n",
    "    \n",
    "    # Scatter plot m_y\n",
    "    plt.scatter(self.Y_types[y_layer], - M_ell, color='red', label='Y types')\n",
    "    plt.vlines(self.Y_types[y_layer], ymin=0, ymax=- M_ell,\n",
    "               color='red', linestyles='dashed')\n",
    "\n",
    "    # Add grid and y=0 axis\n",
    "    # plt.grid(True)\n",
    "    plt.axhline(0, color='black', linewidth=1)\n",
    "    plt.gca().spines['bottom'].set_position(('data', 0))\n",
    "\n",
    "    # Labeling the axes and the title\n",
    "    plt.ylabel('mass')\n",
    "    plt.title('Distributions of types in the layer')\n",
    "    plt.gca().yaxis.set_major_locator(MaxNLocator(integer=True))\n",
    "    plt.gca().spines['top'].set_visible(False)\n",
    "    plt.gca().spines['right'].set_visible(False)\n",
    "    plt.legend()\n",
    "    plt.show()\n",
    "    \n",
    "ConcaveCostOT.plot_layer_types = plot_layer_types\n",
    "\n",
    "example_off_diag.plot_layer_types(layer_example, \n",
    "                                layers_mass_example[layer_id_example])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dc6a1bec",
   "metadata": {},
   "source": [
    "Given the structure of a layer and the *no intersecting pairs* property, the optimal matching and value of the layer can be found recursively.\n",
    "\n",
    "Indeed, if in certain optimal matching $ 1 $ and $ j \\in [N_\\ell], $ $ j-1 $ odd, are paired, then there is no matching between agents in $ [2,j-1] $ and those in $ [j+1,N_\\ell] $ (if both are non empty, i.e., $ j $ is not $ 2 $ or $ N_\\ell $).\n",
    "\n",
    "Hence  in such optimal solution agents in $ [2,j-1] $ are matched among themselves.\n",
    "\n",
    "Since $ [z_2,z_{j-1}] $ (as well as $ [z_{j+1},z_{N_\\ell}] $) is alternating, we can reason recursively.\n",
    "\n",
    "Let $ V_{ij} $ be the optimal value of matching agents in $ [i,j] $ with  $ i,j \\in [N_\\ell], $ $ j -i \\in \\{1,3,\\dots,N_\\ell-1\\} $.\n",
    "\n",
    "Suppose that we computed the value $ V_{ij} $ for all $ i,j \\in [N_\\ell] $ with $ i-j \\in \\{1,3,\\dots,t-2\\} $ for some odd natural number $ t $.\n",
    "\n",
    "Then, for $ i,j \\in [N_\\ell] $ with $ i-j= t $ we have\n",
    "\n",
    "$$\n",
    "V_{ij} = \\min_{k \\in \\{i+1,i+3,\\dots,j\\}} \\left\\{ c_{ik} + V_{i+1,k-1} + V_{k+1,j}\\right\\}\n",
    "$$\n",
    "\n",
    "with the RHS depending only on previously computed values.\n",
    "\n",
    "We set the boundary conditions at $ t=-1 $: $ V_{i+1,i} = 0 $ for each $ i \\in [N_\\ell], $ so that we can apply the same Bellman equation at $ t =1. $\n",
    "\n",
    "The following method takes as input the layer types indices and computes the value function as a matrix $ [V_{ij}]_{ i \\in [N_\\ell+1], j \\in [N_\\ell ]} $.\n",
    "\n",
    "In order to  distinguish  entries that  are relevant for our computations from those that  are never accessed, we initialize this matrix as full of  NaN values."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "48cfd92d",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def solve_bellman_eqs(self,layer):\n",
    "    # Recover cost function within the layer\n",
    "    cost_i_j = self.cost_z_z[layer[:,None],layer[None,:]]\n",
    "\n",
    "    # Initialize value function\n",
    "    V_i_j = np.full((len(layer)+1,len(layer)), np.nan)\n",
    "\n",
    "    # Add boundary conditions\n",
    "    i_bdry = np.arange(len(layer))\n",
    "    V_i_j[i_bdry+1, i_bdry] = 0\n",
    "\n",
    "    t = 1\n",
    "    while t < len(layer):\n",
    "        # Select agents i in [n_L-t] (with potential partners j's in [t,n_L])\n",
    "        i_t = np.arange(len(layer)-t)\n",
    "\n",
    "        # For each i, select each k with |k-i| <= t \n",
    "        # (potential partners of i within segment)\n",
    "        index_ik =  i_t[:,None] + np.arange(1, t+1, 2)[None,:]\n",
    "\n",
    "        # Compute optimal value for pairs with |i-j| = t \n",
    "        V_i_j[i_t, i_t + t] = (cost_i_j[i_t[:,None], index_ik] + \n",
    "                                V_i_j[i_t[:,None] + 1, index_ik - 1] + \n",
    "                                V_i_j[index_ik + 1, i_t[:,None] + t]).min(1)\n",
    "        # Go to next odd integer\n",
    "        t += 2\n",
    "\n",
    "    return V_i_j\n",
    "\n",
    "OffDiagonal.solve_bellman_eqs = solve_bellman_eqs"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14f6974a",
   "metadata": {},
   "source": [
    "Let’s compute values for the layer from our example.\n",
    "\n",
    "Only non-NaN entries are actually used in the computations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5f3138bc",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "# Compute layer value function\n",
    "V_i_j = example_off_diag.solve_bellman_eqs(layer_example)\n",
    "\n",
    "print(f\"Type indices in the layer: {layer_example}\")\n",
    "print('##########################')\n",
    "print(\"Section of the Value function of the layer:\")\n",
    "print(V_i_j.round(2)[:min(10, V_i_j.shape[0]), \n",
    "                         :min(10, V_i_j.shape[1])])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6e6d2657",
   "metadata": {},
   "source": [
    "Having computed the value function, we can proceed to compute the optimal matching as the *policy* that attains the value function that solves the  Bellman equation (*policy evaluation*).\n",
    "\n",
    "We start from agent $ 1 $ and match it with the $ k $ that achieves the minimum in the equation associated with $ V_{1,2N_\\ell}. $\n",
    "\n",
    "Then we store  segments $ [2,k-1] $ and $ [k+1,2N_\\ell] $ (if not empty).\n",
    "\n",
    "In general, given a segment $ [i,j], $ we match $ i $ with $ k $ that achieves the minimum in the equation associated with  $ V_{ij} $ and store the segments $ [i,k-1] $ and $ [k+1,j] $ (if not empty).\n",
    "\n",
    "The algorithm proceeds until there are no segments left."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "125ba739",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def find_layer_matching(self, V_i_j, layer):\n",
    "    # Initialize\n",
    "    segments_to_process  = [np.arange(len(layer))]\n",
    "    matching = np.zeros((len(self.X_types),len(self.Y_types)), bool)\n",
    "\n",
    "    while segments_to_process:\n",
    "        # Pick i, first agent of the segment \n",
    "        # and potential partners i+1,i+3,..., in the segment\n",
    "        segment = segments_to_process[0]\n",
    "        i_0 = segment[0]\n",
    "        potential_matches = np.arange(i_0, segment[-1], 2) + 1\n",
    "\n",
    "        # Compute optimal partner j_i \n",
    "        obj = (self.cost_z_z[layer[i_0],layer[potential_matches]] + \n",
    "                V_i_j[i_0 +1, potential_matches -1] +\n",
    "                V_i_j[potential_matches +1,segment[-1]])\n",
    "        \n",
    "        j_i_0 = np.argmin(obj)*2 + (i_0 + 1)\n",
    "     \n",
    "        # Add matched pair (i,j_i)\n",
    "        self.add_pair_to_matching(layer[[i_0,j_i_0]], matching)\n",
    "\n",
    "        # Update segments to process: \n",
    "        # remove current segment\n",
    "        segments_to_process = segments_to_process[1:]\n",
    "        \n",
    "        # add [i+1,j-1] and [j+1,last agent of the segment]\n",
    "        if j_i_0 > i_0 + 1:\n",
    "            segments_to_process.append(np.arange(i_0 + 1, j_i_0))\n",
    "        if j_i_0 < segment[-1]:\n",
    "            segments_to_process.append(np.arange(j_i_0 + 1, segment[-1] + 1))\n",
    "    \n",
    "    return matching \n",
    "\n",
    "OffDiagonal.find_layer_matching = find_layer_matching"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "750537c7",
   "metadata": {},
   "source": [
    "Lets apply this method our example to find the matching within the layer and then rescale it by $ M_\\ell $.\n",
    "\n",
    "Note that the unscaled value equals $ V_{1,N_\\ell}. $"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a90b9061",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "matching_layer = example_off_diag.find_layer_matching(V_i_j,layer_example)\n",
    "print(f\"Value of the layer (unscaled): {(matching_layer * example_off_diag.cost_x_y).sum()}\")\n",
    "print(f\"Value of the layer (scaled by the mass = {layers_mass_example[layer_id_example]}): \"\n",
    "     f\"{layers_mass_example[layer_id_example] * (matching_layer * example_off_diag.cost_x_y).sum()}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b854d7fa",
   "metadata": {},
   "source": [
    "The following method plots the matching within a layer.\n",
    "\n",
    "We apply it to the layer from our example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bbe24cc2",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def plot_layer_matching(self, layer, matching_layer):\n",
    "    # Create the figure and axis\n",
    "    fig, ax = plt.subplots(figsize=(15, 15))\n",
    "\n",
    "    # Plot the points on the x-axis\n",
    "    X_types_layer = self.X_types[layer[layer < len(self.X_types)]]\n",
    "    Y_types_layer = self.Y_types[layer[layer >= len(self.X_types)] \n",
    "                    - len(self.X_types)]\n",
    "    ax.scatter(X_types_layer, np.zeros_like(X_types_layer), color='blue', \n",
    "                s = 20, zorder=5)\n",
    "    ax.scatter(Y_types_layer, np.zeros_like(Y_types_layer), color='red', \n",
    "                s = 20, zorder=5)\n",
    "\n",
    "    # Draw semicircles for each row in matchings\n",
    "    matched_types = np.where(matching_layer >0)\n",
    "    matched_types_x = self.X_types[matched_types[0]]\n",
    "    matched_types_y = self.Y_types[matched_types[1]]\n",
    "    \n",
    "    for iter in range(len(matched_types_x)):\n",
    "        width = abs(matched_types_x[iter] - matched_types_y[iter])\n",
    "        center = (matched_types_x[iter] + matched_types_y[iter]) / 2        \n",
    "        height = width \n",
    "        semicircle = patches.Arc((center, 0), width, height, theta1=0, \n",
    "                    theta2=180, lw=3)\n",
    "        ax.add_patch(semicircle)\n",
    "\n",
    "    # Add title and layout settings\n",
    "    plt.title('Optimal Layer Matching' )\n",
    "    ax.set_aspect('equal')\n",
    "    plt.gca().spines['bottom'].set_position(('data', 0))\n",
    "    ax.spines['left'].set_color('none') \n",
    "    ax.spines['top'].set_color('none') \n",
    "    ax.spines['right'].set_color('none')  \n",
    "    ax.yaxis.set_ticks([])\n",
    "    ax.set_ylim(bottom= -self.support_z.ptp() / 100)\n",
    "\n",
    "    plt.show()\n",
    "\n",
    "ConcaveCostOT.plot_layer_matching = plot_layer_matching"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5d544d26",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "example_off_diag.plot_layer_matching(layer_example, matching_layer)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "18788372",
   "metadata": {},
   "source": [
    "#### Solving a layer in a smarter way\n",
    "\n",
    "We  now present two key results in the context of OT with concave type costs.\n",
    "\n",
    "We refer [[Boerma *et al.*, 2024](https://python-advanced.quantecon.org/zreferences.html#id3)] and [[Delon *et al.*, 2011](https://python-advanced.quantecon.org/zreferences.html#id4)] for proofs.\n",
    "\n",
    "Consider the problem faced within a layer, i.e., types from $ Y \\sqcup X $\n",
    "\n",
    "$$\n",
    "z_1 < z_2\\dots < z_{N_\\ell-1} < z_{N_\\ell}, \\quad N_\\ell \\in 2 \\mathbb{N}\n",
    "$$\n",
    "\n",
    "are alternating and the problem is unitary.\n",
    "\n",
    "Given a matching on $ [1,k], $ $ k \\in [N_\\ell], $ $ k $ even, we say that a matched pair $ (i,j) $ within this matching is *hidden* if there is a matched pair $ (i',j') $ with $ i' < i <j <j'. $\n",
    "\n",
    "Visually, the arc joining $ (i',j') $ surmounts the arc joining $ (i,j). $\n",
    "\n",
    "**Theorem (DSS)** Given an optimal matching on $ [1,k], $ if $ (i,j) $ is hidden in this matching, then the pair $ (i,j) $ belongs to every optimal matching on $ [1, 2 N_\\ell] $ and is hidden in this matching too.\n",
    "\n",
    "As a consequence, there exists a more efficient way to compute the value function within a layer.\n",
    "\n",
    "It can be shown that the solving the following second-order difference  equations delivers the same result as the Bellman equations above:\n",
    "\n",
    "$$\n",
    "V_{ij} = \\min \\{ c_{ij} + V_{i+1,j-1}, V_{i+2,j} + V_{i,j-2} - V_{i+2,j-2}\\}\n",
    "$$\n",
    "\n",
    "for $ i,j \\in [N_\\ell], $ $ j-i $ odd, with boundary conditions $ V_{i+1,i}= 0 $ for $ i \\in [0,N_\\ell ] $ and $ V_{i+2, i-1} = - c_{i,i+1} $ for $ i \\in [N_\\ell -1] $ .\n",
    "\n",
    "The following method uses these equations to compute the value function that  is stored as a matrix $ [V_{ij}]_{ i \\in [N_\\ell+1], j \\in [N_\\ell +1]}. $"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d99c373a",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def solve_bellman_eqs_DSS(self,layer):\n",
    "    # Recover cost function within the layer\n",
    "    cost_i_j = self.cost_z_z[layer[:,None],layer[None,:]]\n",
    "\n",
    "    # Initialize value function\n",
    "    V_i_j = np.full((len(layer)+1,len(layer)+1), np.nan)\n",
    "    \n",
    "    # Add boundary conditions\n",
    "    V_i_j[np.arange(len(layer)+1), np.arange(len(layer)+1)] = 0\n",
    "    i_bdry = np.arange(len(layer)-1)\n",
    "    V_i_j[i_bdry+2,i_bdry] = - cost_i_j[i_bdry, i_bdry+1]\n",
    "\n",
    "    t = 1\n",
    "    while t < len(layer):\n",
    "        # Select agents i in [n_l-t] and potential partner j=i+t for each i\n",
    "        i_t = np.arange(len(layer)-t)\n",
    "        j_t = i_t + t +1       \n",
    "\n",
    "        # Compute optimal values for ij with j-i = t \n",
    "        V_i_j[i_t, j_t] = np.minimum(cost_i_j[i_t, j_t-1] \n",
    "                        + V_i_j[i_t + 1, j_t - 1],\n",
    "                        V_i_j[i_t, j_t - 2] + V_i_j[i_t + 2, j_t] \n",
    "                        - V_i_j[i_t + 2, j_t - 2])\n",
    "        \n",
    "        ## Go to next odd integer\n",
    "        t += 2\n",
    "    \n",
    "    return V_i_j\n",
    "\n",
    "OffDiagonal.solve_bellman_eqs_DSS = solve_bellman_eqs_DSS"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7358859a",
   "metadata": {},
   "source": [
    "Let’s apply the algorithm to our example and compare outcomes with those attained  with the Bellman equations above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a9864017",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "V_i_j_DSS = example_off_diag.solve_bellman_eqs_DSS(layer_example)\n",
    "\n",
    "print(f\"Type indices of the layer: {layer_example}\")\n",
    "print('##########################')\n",
    "\n",
    "print(\"Section of Value function of the layer:\")\n",
    "print(V_i_j_DSS.round(2)[:min(10, V_i_j_DSS.shape[0]), \n",
    "                         :min(10, V_i_j_DSS.shape[1])])\n",
    "\n",
    "print('##########################')\n",
    "print(f\"Difference with previous Bellman equations: \\\n",
    "      {(V_i_j_DSS[:,1:] - V_i_j)[V_i_j >= 0].sum()}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a5a4b058",
   "metadata": {},
   "source": [
    "We can actually compute the optimal matching within the layer simultaneously with computing the value function, rather than sequentially.\n",
    "\n",
    "The key idea is that, if at some step of the computation of the values the left branch of the minimum above achieves the minimum, say $ V_{ij}= c_{ij} + V_{i+1,j-1}, $ then $ (i,j) $ are optimally matched on $ [i,j] $ and by the theorem above we get that a matching on $ [i+1,j-1] $ which achieves $ V_{i+1,j-1} $ belongs to an optimal matching on the whole layer (since it is covered by the arc $ (i,j) $ in $ [i,j] $).\n",
    "\n",
    "We can therefore proceed as follows\n",
    "\n",
    "We initialize an empty matching and a list with all the agents in the layer (representing the agents which are not matched yet).\n",
    "\n",
    "Then whenever the left branch of the minimum is achieved for some $ (i,j) $ in the computation of $ V, $ we take the collections of agents $ k_1,\\dots,k_M $ in $ [i+1,j-1] $ (in ascending order, i.e. with $ z_{k_{p}} < z_{k_{p+1}} $) that are not matched yet (if any) and add to the matching the pairs $ (k_1,k_2), (k_3,k_4),\\dots,(k_{M-1},k_M). $\n",
    "\n",
    "Thus, we match each unmatched agent $ k_p $ in $ [i+1,j-1] $ with the closest unmatched right neighbour $ k_{p+1} $ (starting from $ k_1 $).\n",
    "\n",
    "Intuitively, if $ k_p $ were optimally matched with some $ k_{q} $ in $ [i+1,j-1] $ and not with $ k_{p+1} $, then $ k_{p+1} $ would have already been hidden by the match $ (k_p,k_{q}) $ from some previous computation (because $ |k_p - k_q|< |i-j| $) and it would therefore be matched.\n",
    "\n",
    "Finally, if the process above leaves some umatched agents, we proceed by matching each of these agent with the closest unmatched right neighbour, starting again from the leftmost of these collection.\n",
    "\n",
    "To gain understanding, note that this situation happens  when the left branch is achieved only for pairs $ i,j $ with $ |i-j|=1, $ which leads to the optimal matching $ (1,2), (2,3), \\dots, (n_\\ell -1, n_\\ell). $"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d1a37a3d",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def find_layer_matching_DSS(self,layer):\n",
    "    # Recover cost function within the layer\n",
    "    cost_i_j = self.cost_z_z[layer[:,None],layer[None,:]]\n",
    "    \n",
    "    # Add boundary conditions\n",
    "    V_i_j = np.zeros((len(layer)+1,len(layer)+1))\n",
    "    i_bdry = np.arange(len(layer)-1)\n",
    "    V_i_j[i_bdry+2,i_bdry] = - cost_i_j[i_bdry, i_bdry+1]\n",
    "\n",
    "    # Initialize matching and list of to-match agents\n",
    "    unmatched = np.ones(len(layer), dtype = bool)\n",
    "    matching = np.zeros((len(self.X_types),len(self.Y_types)), bool)\n",
    "\n",
    "    t = 1\n",
    "    while t < len(layer):\n",
    "        # Compute optimal value for pairs with |i-j| = t \n",
    "        i_t = np.arange(len(layer)-t)\n",
    "        j_t = i_t + t + 1\n",
    "        \n",
    "        left_branch = cost_i_j[i_t, j_t-1] + V_i_j[i_t + 1, j_t - 1]\n",
    "        V_i_j[i_t, j_t] = np.minimum(left_branch, V_i_j[i_t, j_t - 2] \n",
    "                        + V_i_j[i_t + 2, j_t] - V_i_j[i_t + 2, j_t - 2])\n",
    "        \n",
    "        # Select each i for which left branch achieves minimum in the V_{i,i+t}\n",
    "        left_branch_achieved = i_t[left_branch == V_i_j[i_t, j_t]]\n",
    "        \n",
    "        # Update matching\n",
    "        for i in left_branch_achieved:\n",
    "            # for each agent k in [i+1,i+t-1]\n",
    "            for k in np.arange(i+1,i+t)[unmatched[range(i+1,i+t)]]:   \n",
    "                # if k is unmatched\n",
    "                if unmatched[k] == True:                               \n",
    "                    # find unmatched right neighbour  \n",
    "                    j_k = np.arange(k+1,len(layer))[unmatched[k+1:]][0]\n",
    "                    # add pair to matching\n",
    "                    self.add_pair_to_matching(layer[[k, j_k]], matching) \n",
    "                    # remove pair from unmatched agents list\n",
    "                    unmatched[[k, j_k]] = False                         \n",
    "                    \n",
    "        # go to next odd integer\n",
    "        t += 2\n",
    "    \n",
    "    # Each umatched agent is matched with next unmatched agent\n",
    "    for i in np.arange(len(layer))[unmatched]:                  \n",
    "        # if i is unmatched\n",
    "        if unmatched[i] == True:                                 \n",
    "            # find unmatched right neighbour\n",
    "            j_i = np.arange(i+1,len(layer))[unmatched[i+1:]][0]  \n",
    "            # add pair to matching\n",
    "            self.add_pair_to_matching(layer[[i, j_i]], matching) \n",
    "            # remove pair from unmatched agents list \n",
    "            unmatched[[i, j_i]] = False                          \n",
    "\n",
    "    return matching\n",
    "\n",
    "OffDiagonal.find_layer_matching_DSS = find_layer_matching_DSS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7f1965b2",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "matching_layer_DSS = example_off_diag.find_layer_matching_DSS(layer_example)\n",
    "print(f\" Value of layer with DSS recursive equations \\\n",
    "{(matching_layer_DSS * example_off_diag.cost_x_y).sum()}\")\n",
    "print(f\" Value of layer with Bellman equations \\\n",
    "{(matching_layer * example_off_diag.cost_x_y).sum()}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17c976d3",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "example_off_diag.plot_layer_matching(layer_example, matching_layer_DSS)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5a8a4776",
   "metadata": {},
   "source": [
    "## Solving primal problem\n",
    "\n",
    "The following method assembles  our components in order to solve the primal problem.\n",
    "\n",
    "First, if matches are perfect pairs, we store the on-diagonal matching and create an off-diagonal instance with the residual marginals.\n",
    "\n",
    "Then we compute the set of layers of the residual distributions.\n",
    "\n",
    "Finally, we solve each layer and put together  matchings within each layer with the on-diagonal matchings.\n",
    "\n",
    "We then return the full matching, the off-diagonal matching, and the off-diagonal instance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a7916cd1",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def solve_primal_pb(self):\n",
    "    # Compute on-diagonal matching, create new instance with resitual types\n",
    "    off_diagoff_diagonal, match_tuple = self.generate_offD_onD_matching()\n",
    "    nonzero_id_x, nonzero_id_y, matching_diag = match_tuple\n",
    "    \n",
    "    # Compute layers\n",
    "    layers_list, layers_mass, _, _ = off_diagoff_diagonal.find_layers()\n",
    "\n",
    "    # Solve layers to compute off-diagonal matching\n",
    "    matching_off_diag = np.zeros_like(off_diagoff_diagonal.cost_x_y, dtype=int)\n",
    "\n",
    "    for ell, layer  in enumerate(layers_list):\n",
    "        V_i_j = off_diagoff_diagonal.solve_bellman_eqs(layer)\n",
    "        matching_off_diag += layers_mass[ell] \\\n",
    "                    * off_diagoff_diagonal.find_layer_matching(V_i_j, layer)\n",
    "\n",
    "    # Add together on- and off-diagonal matchings\n",
    "    matching = matching_diag.copy()\n",
    "    matching[np.ix_(nonzero_id_x, nonzero_id_y)] += matching_off_diag\n",
    "\n",
    "    return matching, matching_off_diag, off_diagoff_diagonal\n",
    "\n",
    "ConcaveCostOT.solve_primal_pb = solve_primal_pb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fb6fbc15",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "matching, matching_off_diag, off_diagoff_diagonal = example_pb.solve_primal_pb()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3bcf1308",
   "metadata": {},
   "source": [
    "We implement a similar method that  adopts the DSS algorithm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a33c0bb5",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def solve_primal_DSS(self):\n",
    "    # Compute on-diagonal matching, create new instance with resitual types\n",
    "    off_diagoff_diagonal, match_tuple = self.generate_offD_onD_matching()\n",
    "    nonzero_id_x, nonzero_id_y, matching_diag = match_tuple\n",
    "\n",
    "    # Find layers\n",
    "    layers, layers_mass, _, _ = off_diagoff_diagonal.find_layers()\n",
    "\n",
    "    # Solve layers to compute off-diagonal matching\n",
    "    matching_off_diag = np.zeros_like(off_diagoff_diagonal.cost_x_y, dtype=int)\n",
    "\n",
    "    for ell, layer  in enumerate(layers):\n",
    "        matching_off_diag += layers_mass[ell] \\\n",
    "                          * off_diagoff_diagonal.find_layer_matching_DSS(layer)\n",
    "\n",
    "    # Add together on- and off-diagonal matchings\n",
    "    matching = matching_diag.copy()\n",
    "    matching[np.ix_(nonzero_id_x, nonzero_id_y)] += matching_off_diag\n",
    "\n",
    "    return matching, matching_off_diag, off_diagoff_diagonal\n",
    "\n",
    "ConcaveCostOT.solve_primal_DSS = solve_primal_DSS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b4242d73",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "DSS_tuple = example_pb.solve_primal_DSS()\n",
    "matching_DSS, matching_off_diag_DSS, off_diagoff_diagonal_DSS = DSS_tuple"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "49c33247",
   "metadata": {},
   "source": [
    "By drawing semicircles joining the matched agents (with distinct types), we can visualize the off-diagonal matching.\n",
    "\n",
    "In the following figure,  widths and colors of  semicirles indicate relative numbers of agents that  are “transported” along an  arc."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "827f6e6f",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def plot_matching(self, matching_off_diag, title, figsize=(15, 15), \n",
    "                    add_labels=False, plot_H_z=False, scatter=True):\n",
    "    \n",
    "    # Create the figure and axis\n",
    "    fig, ax = plt.subplots(figsize=figsize)\n",
    "\n",
    "    # Plot types on the real line\n",
    "    if scatter:\n",
    "        ax.scatter(self.X_types, np.zeros_like(self.X_types), color='blue', \n",
    "                    s=20, zorder=5)\n",
    "        ax.scatter(self.Y_types, np.zeros_like(self.Y_types), color='red', \n",
    "                    s=20, zorder=5)\n",
    "\n",
    "    # Add labels for X_types and Y_types if add_labels is True\n",
    "    if add_labels:\n",
    "        # Remove x-axis ticks\n",
    "        ax.set_xticks([])\n",
    "\n",
    "        # Add labels\n",
    "        for i, x in enumerate(self.X_types):\n",
    "            ax.annotate(f'$x_{{{i }}}$', (x, 0), textcoords=\"offset points\", \n",
    "                        xytext=(0, -15), ha='center', color='blue', fontsize=12)\n",
    "        for j, y in enumerate(self.Y_types):\n",
    "            ax.annotate(f'$y_{{{j }}}$', (y, 0), textcoords=\"offset points\", \n",
    "                        xytext=(0, -15), ha='center', color='red', fontsize=12)\n",
    "            \n",
    "    # Draw semicircles for each pair of matched types\n",
    "    matched_types = np.where(matching_off_diag > 0)\n",
    "    matched_types_x = self.X_types[matched_types[0]]\n",
    "    matched_types_y = self.Y_types[matched_types[1]]\n",
    "   \n",
    "    count = matching_off_diag[matched_types]\n",
    "    colors = plt.cm.Greys(np.linspace(0.5, 1.5, count.max() + 1))\n",
    "    max_height = 0\n",
    "    for iter in range(len(count)):\n",
    "        width = abs(matched_types_x[iter] - matched_types_y[iter])\n",
    "        center = (matched_types_x[iter] + matched_types_y[iter]) / 2        \n",
    "        height = width \n",
    "        max_height = max(max_height, height)\n",
    "        semicircle = patches.Arc((center, 0), width, height, \n",
    "                                 theta1=0, theta2=180, \n",
    "                                 color=colors[count[iter]], \n",
    "                                 lw=count[iter] * (2.2 / count.max()))\n",
    "        ax.add_patch(semicircle)\n",
    "\n",
    "    # Title and layout settings for the main plot\n",
    "    plt.title(title)\n",
    "    ax.set_aspect('equal')\n",
    "    plt.axhline(0, color='black', linewidth=1)\n",
    "    ax.spines['bottom'].set_position(('data', 0))\n",
    "    ax.spines['left'].set_color('none')\n",
    "    ax.spines['top'].set_color('none')\n",
    "    ax.spines['right'].set_color('none')\n",
    "    ax.yaxis.set_ticks([])\n",
    "    ax.set_ylim(- self.X_types.ptp() / 10, \n",
    "                (max_height / 2) + self.X_types.ptp()*.01)  \n",
    "\n",
    "    # Plot H_z on the main axis if enabled\n",
    "    if plot_H_z:\n",
    "        H_z = np.cumsum(self.q_z)  \n",
    "\n",
    "        step = np.concatenate(([self.support_z.min() \n",
    "                                - .02 * self.support_z.ptp()],\n",
    "                                self.support_z,\n",
    "                                [self.support_z.max() \n",
    "                                + .02 * self.support_z.ptp()]))\n",
    "\n",
    "        H_z = H_z/H_z.ptp() * self.support_z.ptp() /2\n",
    "        height = np.concatenate(([0], H_z, [0]))\n",
    "\n",
    "        # Plot the compressed H_z on the same main x-axis\n",
    "        ax.step(step, height, color='green', lw=2, \n",
    "                                label='$H_z$', where='post')\n",
    "        \n",
    "        # Set the y-limit to keep H_z and maximum circle size in the plot\n",
    "        ax.set_ylim(np.min(H_z) - H_z.ptp() *.01,\n",
    "                    np.maximum(np.max(H_z), max_height / 2) + H_z.ptp() *.01) \n",
    "\n",
    "        # Add label and legend for H_z\n",
    "        ax.legend(loc=\"upper right\")\n",
    "\n",
    "        plt.show()\n",
    "\n",
    "ConcaveCostOT.plot_matching = plot_matching"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f5b463d0",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "off_diagoff_diagonal.plot_matching(matching_off_diag, \n",
    "            title='Optimal Matching (off-diagonal)', plot_H_z=True)\n",
    "off_diagoff_diagonal_DSS.plot_matching(matching_off_diag_DSS, \n",
    "            title='Optimal Matching (off-diagonal) with DSS algorithm')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ae7ffc11",
   "metadata": {},
   "source": [
    "### Verify with linear programming\n",
    "\n",
    "Let’s verify some of the proceeding findings using linear programming."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "54dc4dbd",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def solve_1to1(c_i_j, n_x, m_y, return_dual=False):\n",
    "    n, m = np.shape(c_i_j)\n",
    "\n",
    "    # Constraint matrix \n",
    "    M_z_a = np.vstack([np.kron(np.eye(n), np.ones(m)),\n",
    "                       np.kron(np.ones(n), np.eye(m))])\n",
    "    # Constraint vector\n",
    "    q = np.concatenate((n_x, m_y))\n",
    "    \n",
    "    # Solve the linear programming problem using linprog from scipy\n",
    "    result = linprog(c_i_j.flatten(), A_eq=M_z_a, b_eq=q, \n",
    "                        bounds=(0, None), method='highs')\n",
    "\n",
    "    if return_dual:\n",
    "        return (np.round(result.x).astype(int).reshape([n, m]), \n",
    "                result.eqlin.marginals)\n",
    "    else:\n",
    "        return np.round(result.x).astype(int).reshape([n, m])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "96dff11c",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "mu_x_y_LP = solve_1to1(example_pb.cost_x_y,\n",
    "                example_pb.n_x,\n",
    "                example_pb.m_y)\n",
    "print(f\"Value of LP (scipy): {(mu_x_y_LP * example_pb.cost_x_y).sum()}\")\n",
    "print(f\"Value (plain Bellman equations): {(matching * example_pb.cost_x_y).sum()}\")\n",
    "print(f\"Value (DSS): {(matching_DSS * example_pb.cost_x_y).sum()}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "903b493b",
   "metadata": {},
   "source": [
    "## Examples"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3e75dfab",
   "metadata": {},
   "source": [
    "### Example 1\n",
    "\n",
    "We study optimal transport problems on the real line with cost $ c(x,y)= h(|x-y|) $ for a strictly concave and increasing function $ h: \\mathbb{R}_+ \\rightarrow \\mathbb{R}_+. $\n",
    "\n",
    "The outcome  is called *composite sorting*.\n",
    "\n",
    "Here, we  will focus on $ c(x,y)=|x-y|^{\\frac{1}{\\zeta}} $ for $ \\zeta>1 $\n",
    "\n",
    "To appreciate  differences with *positive assortative matching* (PAM) note that the latter is induced by a cost of the form $ h(x-y) $ for some strictly convex $ h: \\mathbb{R} \\rightarrow \\mathbb{R}_+. $\n",
    "\n",
    "See Santambrogio 2015, Ch. 2.2.\n",
    "\n",
    "For example, the cost function $ |x-y|^{p},p>1 $ induces PAM.\n",
    "\n",
    "On the other hand, *negative assortative matching* (NAM) arises if $ c(x,y)= h(x-y) $ with  $ h: \\mathbb{R} \\rightarrow \\mathbb{R}_+ $ strictly concave.\n",
    "\n",
    "For example,  the cost function $ -|x-y|^{p},p>1, $ induces NAM.\n",
    "\n",
    "Thus, NAM corresponds to a matching that *maximizes* a transport problem criterion with *gain* function $ g(x,y)=|x-y|^{p} $.\n",
    "\n",
    "Note how PAM and NAM differ from  **composite sorting**\n",
    "\n",
    "**Composite sorting**   is  induced by a cost that is the composition of a strictly concave increasing function $ h $ and a convex function $ |\\cdot| $ applied to  displacement $ x-y. $\n",
    "\n",
    "Different functions $ h $ potentially induce different matchings.\n",
    "\n",
    "The following example shows that composite matching can feature both positive and negative assortative patterns.\n",
    "\n",
    "Suppose that  there are two agents per side and types\n",
    "\n",
    "$$\n",
    "\\textcolor{blue}{x_0} <  \\textcolor{red}{y_0} <  \\textcolor{blue}{x_1} <  \\textcolor{red}{y_1}\n",
    "$$\n",
    "\n",
    "There are  two feasible matchings, one corresponding  to PAM, the other to  NAM.\n",
    "\n",
    "- The first features two displacements $ |\\textcolor{blue}{x_0} -  \\textcolor{red}{y_0}| ,|  \\textcolor{blue}{x_1} -  \\textcolor{red}{y_1}| $  \n",
    "- The second features a large displacement $ |\\textcolor{blue}{x_0} -  \\textcolor{red}{y_1}| $ and a small displacement $ |  \\textcolor{blue}{x_1} -  \\textcolor{red}{y_0}|. $  \n",
    "\n",
    "\n",
    "Evidently,\n",
    "\n",
    "- PAM corresponds to the matching with two medium side displacement because the correponding cost is strictly convex and increasing in the the displacement.  \n",
    "- NAM corresponds to the matching with a small displacement and a large displacement because the gain is strictly convex and increasing in the displacement.  \n",
    "\n",
    "\n",
    "In this example, composite sorting ends up coinciding  with NAM, but this is something of a coincidence\n",
    "\n",
    "- Thus,  note that  in composite matching the cost function is strictly concave and increasing in the displacement.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ef064de7",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "N = 2\n",
    "p = 2\n",
    "ζ = 2\n",
    "\n",
    "# Solve composite sorting problem\n",
    "example_1 = ConcaveCostOT(np.array([0,5]), \n",
    "                          np.array([4,10]),\n",
    "                          ζ=ζ)\n",
    "matching_CS, _ ,_ = example_1.solve_primal_DSS()\n",
    "\n",
    "# Solve PAM and NAM\n",
    "# I use the linear programs to compute PAM and NAM,\n",
    "# but of course they can be computed directly\n",
    "\n",
    "convex_cost = np.abs(example_1.X_types[:,None] - example_1.Y_types[None,:])**p\n",
    "\n",
    "#PAM: |x-y|^p , p>1\n",
    "matching_PAM = solve_1to1(convex_cost, example_1.n_x, example_1.m_y)\n",
    "\n",
    "#NAM: -|x-y|^p , p>1\n",
    "matching_NAM = solve_1to1(-convex_cost, example_1.n_x, example_1.m_y)\n",
    "\n",
    "# Plot the matchings\n",
    "example_1.plot_matching(matching_CS, \n",
    "                title=f'Composite Sorting: $|x-y|^{{1/{ζ}}}$', \n",
    "                figsize=(5,5), add_labels=True)\n",
    "example_1.plot_matching(matching_PAM, title='PAM', \n",
    "                figsize=(5,5), add_labels=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "076be401",
   "metadata": {},
   "source": [
    "To explore the coincidental resemblence to a NAM outcome,  let’s  shift left type $ \\textcolor{red}{y_0} $ while  keeping it in between $ \\textcolor{blue}{x_0} $ and $ \\textcolor{blue}{x_1} $.\n",
    "\n",
    "PAM and NAM are invariant to any such shift.\n",
    "\n",
    "However, for a large enough shift, composite sorting now coindices with PAM."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8a15c893",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "N = 2\n",
    "ζ = 2\n",
    "p = 2\n",
    "\n",
    "# Solve composite sorting problem\n",
    "example_1 = ConcaveCostOT(np.array([0,5]), \n",
    "                          np.array([1,10]) ,\n",
    "                           ζ = ζ)\n",
    "matching_CS, _ ,_ = example_1.solve_primal_DSS()\n",
    "\n",
    "# Solve PAM and NAM\n",
    "convex_cost = np.abs(example_1.X_types[:,None] - example_1.Y_types[None,:])**p\n",
    "\n",
    "matching_PAM = solve_1to1(convex_cost, example_1.n_x, example_1.m_y)\n",
    "matching_NAM = solve_1to1(-convex_cost, example_1.n_x, example_1.m_y)\n",
    "\n",
    "# Plot the matchings\n",
    "example_1.plot_matching(matching_CS, \n",
    "        title = f'Composite Sorting: $|x-y|^{{1/{ζ}}}$', \n",
    "        figsize = (5,5), add_labels = True)\n",
    "example_1.plot_matching(matching_PAM, title = 'PAM', \n",
    "        figsize = (5,5), add_labels = True)\n",
    "example_1.plot_matching(matching_NAM, title = 'NAM', \n",
    "        figsize = (5,5), add_labels = True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c96fb086",
   "metadata": {},
   "source": [
    "Finally, notice that the  **Monge problem**  cost function $ |x-y| $  equals the limit of the  composite sorting cost $ |x-y|^{1/\\zeta} $ as $ \\zeta \\downarrow 1 $ and also  the limit of $ |x-y|^p $ as $ p \\downarrow 1. $\n",
    "\n",
    "Evidently, the Monge problem is solved by both the PAM and the composite sorting assignment that arises for $ \\zeta \\downarrow 1. $\n",
    "\n",
    "In the following example, the Monge cost of the composite sorting assignment equals the Monge cost of PAM.\n",
    "\n",
    "Consequently, it is optimal for the Monge problem."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2103cff4",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "N = 10\n",
    "\n",
    "ζ = 1.01\n",
    "p = 2\n",
    "np.random.seed(1)\n",
    "X_types = np.random.uniform(0,10, size=N)\n",
    "Y_types = np.random.uniform(0,10, size=N)\n",
    "\n",
    "# Solve composite sorting problem\n",
    "example_1 = ConcaveCostOT(X_types, Y_types, ζ=ζ)\n",
    "\n",
    "matching_CS, _ ,_ = example_1.solve_primal_DSS()\n",
    "\n",
    "# Solve PAM and NAM\n",
    "convex_cost = np.abs(X_types[:,None] - Y_types[None,:])** p\n",
    "\n",
    "matching_PAM = solve_1to1(convex_cost, example_1.n_x, example_1.m_y)\n",
    "matching_NAM = solve_1to1(-convex_cost, example_1.n_x, example_1.m_y)\n",
    "\n",
    "\n",
    "example_1.plot_matching(matching_CS, \n",
    "            title=f'Composite Sorting: $|x-y|^{{1/{ζ}}}$', figsize=(5,5))\n",
    "example_1.plot_matching(matching_PAM, title = 'PAM', figsize=(5,5)) \n",
    "\n",
    "monge_cost_comp = (matching_CS * np.abs(X_types[:,None] - Y_types[None,:])).sum()\n",
    "monge_cost_PAM = (matching_PAM * np.abs(example_1.X_types[:,None] \n",
    "                               - example_1.Y_types[None,:])).sum()\n",
    "print(\"Monge cost of the composite matching assignment:\")\n",
    "print(monge_cost_comp)\n",
    "print(\"Monge cost of PAM:\")\n",
    "print(monge_cost_PAM)   "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eed9912b",
   "metadata": {},
   "source": [
    "### Example 2\n",
    "\n",
    "The following example has five agents per side.\n",
    "\n",
    "The composite sorting assignment differs from both PAM and NAM.\n",
    "\n",
    "Composite sorting features a hierarchical structure, with each hierarchy positively sorted.\n",
    "\n",
    "Indeed, consider the composite sorting assignment and note that\n",
    "\n",
    "- the only arcs *visible from above* are the ones corresponding to pairings $ (\\textcolor{blue}{x_0},\\textcolor{red}{y_3}) $ and $ (\\textcolor{red}{y_4},\\textcolor{blue}{x_4}); $  \n",
    "- after removing these agents, the only arcs visible from above correspond to $ (\\textcolor{blue}{x_1},\\textcolor{red}{y_1}) $ and $ (\\textcolor{blue}{x_3},\\textcolor{red}{y_2}) $ ;  \n",
    "- after removing these agents, the only arc/pairing left is $ (\\textcolor{blue}{x_2},\\textcolor{red}{y_0}). $  \n",
    "\n",
    "\n",
    "Note that, at each iteration, the partial assignment corresponding to the arcs visible from above features positive assortativeness.\n",
    "\n",
    "Another distinct feature of composite matching  stands out from the figures:\n",
    "\n",
    "- **arcs do not intersect**  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "53c3a7b0",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "N = 5\n",
    "ζ = 2\n",
    "p = 2\n",
    "\n",
    "X_types_example_2 = np.array([-2,0,2,9, 15])\n",
    "Y_types_example_2 = np.array([3,6,10,12, 14]) \n",
    "\n",
    "# Solve composite sorting problem\n",
    "example_2 = ConcaveCostOT(X_types_example_2, Y_types_example_2, ζ=ζ)\n",
    "\n",
    "matching_CS, _ ,_ = example_2.solve_primal_DSS()\n",
    "\n",
    "# Solve PAM and NAM\n",
    "convex_cost = np.abs(X_types_example_2[:,None] - Y_types_example_2[None,:])** p\n",
    "\n",
    "matching_PAM = solve_1to1(convex_cost, example_2.n_x, example_2.m_y)\n",
    "matching_NAM = solve_1to1(-convex_cost, example_2.n_x, example_2.m_y)\n",
    "\n",
    "\n",
    "example_2.plot_matching(matching_CS, title = 'Composite Sorting: $|x-y|^{1/2}$', \n",
    "                        figsize = (5,5), add_labels=True)\n",
    "example_2.plot_matching(matching_PAM, title = 'PAM', \n",
    "                        figsize = (5,5), add_labels=True) \n",
    "example_2.plot_matching(matching_NAM, title = 'NAM', \n",
    "                        figsize = (5,5), add_labels=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b5f34904",
   "metadata": {},
   "source": [
    "### Example 3\n",
    "\n",
    "[[Boerma *et al.*, 2024](https://python-advanced.quantecon.org/zreferences.html#id3)] provide the following  example.\n",
    "\n",
    "There are four agents per side and three types per side (so the problem is not unitary, as opposed to the examples above)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bad55a25",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "X_types_example_3 = np.array([0,5,9])\n",
    "Y_types_example_3 = np.array([1,6,10]) \n",
    "n_x_example_3 =  np.array([2,1,1], dtype= int)\n",
    "m_y_example_3 =  np.array([1,1,2], dtype= int)\n",
    "\n",
    "\n",
    "example_3 = ConcaveCostOT(X_types_example_3, Y_types_example_3, \n",
    "                          n_x_example_3, m_y_example_3, ζ = 2)\n",
    "example_3.plot_marginals(figsize = (5,5))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9995a337",
   "metadata": {},
   "source": [
    "In the case of positive assortative matching (PAM), the two agents with lowest value $ \\textcolor{blue}{x_0} $ are matched with the lowest valued agents on the other side $ \\textcolor{red}{y_0},\\textcolor{red}{y_1}. $\n",
    "\n",
    "Similarly, the agents with highest value $ \\textcolor{red}{y_2} $ are matched with the highest valued types on the other side, $ \\textcolor{blue}{x_1} $ and $ \\textcolor{blue}{x_2}. $\n",
    "\n",
    "Composite sorting features both negative and positive sorting patterns: agents of type $ \\textcolor{blue}{x_0} $ are matched with both the bottom $ \\textcolor{red}{y_0} $ and the top $ \\textcolor{red}{y_2} $ of the distribution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c49400b4",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "matching_CS, _ ,_ = example_3.solve_primal_DSS()\n",
    "\n",
    "convex_cost = np.abs(example_3.X_types[:,None] - example_3.Y_types[None,:])**2\n",
    "matching_PAM = solve_1to1(convex_cost, example_3.n_x, example_3.m_y)\n",
    "matching_NAM = solve_1to1(-convex_cost, example_3.n_x, example_3.m_y)\n",
    "\n",
    "example_3.plot_matching(matching_PAM, title = 'PAM', \n",
    "                        figsize = (5,5), add_labels= True)\n",
    "example_3.plot_matching(matching_CS, title = 'Composite Sorting', \n",
    "                        figsize = (5,5), add_labels= True)\n",
    "example_3.plot_matching(matching_NAM, title = 'NAM', \n",
    "                        figsize = (5,5), add_labels= True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8045aed9",
   "metadata": {},
   "source": [
    "## Dual Solution\n",
    "\n",
    "Let’s recall the formulation\n",
    "\n",
    "$$\n",
    "\\begin{aligned}\n",
    "V_P = \\min_{\\mu \\geq 0}& \\sum_{(x,y) \\in X \\times Y} \\mu_{xy}c_{xy} \\\\\n",
    "\\text{s.t. }& \\sum_{x \\in X} \\mu_{xy} = n_x \\\\\n",
    "& \\sum_{y \\in Y} \\mu_{xy} = m_y \n",
    "\\end{aligned}\n",
    "$$\n",
    "\n",
    "The *dual problem* is\n",
    "\n",
    "$$\n",
    "\\begin{aligned}\n",
    "V_D = \\max_{\\phi,\\psi}& \\sum_{x \\in X }n_x \\phi_x + \\sum_{y \\in Y} m_y \\psi_y \\\\\n",
    "\\text{s.t. }&  \\phi_x + \\psi_y \\leq c_{xy}\n",
    "\\end{aligned}\n",
    "$$\n",
    "\n",
    "where $ (\\phi , \\psi) $ are dual variables, which can be interpreted as shadow cost of agents in $ X $ and $ Y $, respectively.\n",
    "\n",
    "Since the dual is feasible and bounded,  $ V_P = V_D $ (*strong duality* prevails).\n",
    "\n",
    "Assume now that $ y_{xy} = \\alpha_x + \\gamma_y - c_{xy} $ is the output generated by matching $ x $ and $ y. $\n",
    "\n",
    "It includes the sum of $ x $ and $ y $ specific amenities/outputs minus the cost $ c_{xy}. $\n",
    "\n",
    "Then we  can formulate the following problem and its dual\n",
    "\n",
    "$$\n",
    "\\begin{aligned}\n",
    "W_P = \\max_{\\mu \\geq 0}& \\sum_{(x,y) \\in X \\times Y} \\mu_{xy}y_{xy} \\\\\n",
    "\\text{s.t. }& \\sum_{x \\in X} \\mu_{xy} = n_x \\\\\n",
    "& \\sum_{y \\in Y} \\mu_{xy} = m_y \n",
    "\\end{aligned}\n",
    "$$\n",
    "\n",
    "$$\n",
    "\\begin{aligned}\n",
    " W_D = \\min_{u,v}& \\sum_{x \\in X }n_x u_x + \\sum_{y \\in Y} m_y v_y \\\\\n",
    "\\text{s.t. }&  u_x + v_y \\geq y_{xy}\n",
    "\\end{aligned}\n",
    "$$\n",
    "\n",
    "Given the constraints, the primal problem $ W_P $ does not depend on $ \\alpha,\\gamma $ and it has the same solutions as the cost minimization problem $ V_P $.\n",
    "\n",
    "The values are related by $ W_P =  \\sum_{x \\in X}n_x \\alpha_x +  \\sum_{y \\in Y}m_y \\gamma_y - V_P. $\n",
    "\n",
    "The dual solutions of $ V_D $ and $ W_D $ are related by $ u_x = \\alpha_x - \\phi_x $ and $ v_y = \\gamma_y - \\psi_y. $\n",
    "\n",
    "The dual solution $ (u,v) $ of $ W_D $ can be interpreted as equilibrium utilities of the agents, which include the individual specific amenities and equilibrium shadow costs.\n",
    "\n",
    "[[Boerma *et al.*, 2024](https://python-advanced.quantecon.org/zreferences.html#id3)] propose an efficient method to compute the dual variables from the optimal matching (primal solution) in the case of composite sorting.\n",
    "\n",
    "Let’s generate an instance and compute the optimal matching."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c8c578ec",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "num_agents = 8\n",
    "\n",
    "np.random.seed(1)\n",
    "\n",
    "X_types_assignment_pb = np.random.uniform(0, 10, size=num_agents)\n",
    "Y_types_assignment_pb = np.random.uniform(0, 10, size=num_agents)\n",
    "\n",
    "\n",
    "# Create instance of the problem\n",
    "exam_assign = ConcaveCostOT(X_types_assignment_pb, Y_types_assignment_pb)\n",
    "\n",
    "# Solve primal problem\n",
    "assignment, assignment_OD, exam_assign_OD = exam_assign.solve_primal_DSS()\n",
    "\n",
    "# Plot matching\n",
    "add_labels = True if num_agents < 16 else False\n",
    "exam_assign_OD.plot_matching(assignment_OD, title = f'Composite Sorting', \n",
    "                            figsize=(10,10), add_labels=add_labels)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "beee650b",
   "metadata": {},
   "source": [
    "Having computed the optimal matching, we say that a pair $ (x_0,y_0) $ is a *subpair* of a matched pair $ (x,y) $ if $ x_0,y_0 $ are in the open interval between $ x $ and $ y $ and the pair $ (x_0,y_0) $ is not nested.\n",
    "\n",
    "The following method computes the subpairs of the optimal matching of the off-diagonal instance.\n",
    "\n",
    "The output of this method is a dictionary with keys corresponding to matched pairs and an “artificial pair” which collects all arcs which are visible from above.\n",
    "\n",
    "Values of each key $ (x_0,y_0) $ are the subpairs ordered so that the first subpair is the subpair with the $ x $ type closest to $ x_0 $ and the last subpair is the subpair with the $ y $ type closest to $ y_0. $"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3a4d93ef",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def sort_subpairs(self, subpairs, x_smaller_y=True ):\n",
    "\n",
    "    x_key = min if x_smaller_y else max\n",
    "    y_key = max if x_smaller_y else min\n",
    "\n",
    "    first_pair = x_key(subpairs, key=lambda pair: self.X_types[pair[0]])\n",
    "    last_pair = y_key(subpairs, key=lambda pair: self.Y_types[pair[1]])\n",
    "\n",
    "    intermediate_pairs = [pair for pair in subpairs \n",
    "                            if pair != first_pair and pair != last_pair]\n",
    "\n",
    "    return [first_pair] + intermediate_pairs + [last_pair]\n",
    "\n",
    "ConcaveCostOT.sort_subpairs = sort_subpairs\n",
    "\n",
    "def find_subpairs(self, matching, return_pairs_between = False):\n",
    "        \n",
    "    # Create set of matched pairs of types and add an artificial pair \n",
    "    matched_pairs = set( zip(* np.where(matching > 0))) \n",
    "\n",
    "    # Initialize dictionary to store subpairs\n",
    "    subpairs = {} \n",
    "    pairs_between = {}\n",
    "\n",
    "    # Find subpairs (both nested and non-nested) for each matched pair \n",
    "    for matched_pair in matched_pairs | {'artificial_pair'}:\n",
    "        # Determine the interval of the matched pair\n",
    "        if matched_pair != 'artificial_pair':\n",
    "            min_type, max_type = sorted([self.X_types[matched_pair[0]], \n",
    "                                         self.Y_types[matched_pair[1]]]) \n",
    "        else:\n",
    "            min_type, max_type = (-np.inf, np.inf)\n",
    "            \n",
    "        # Add all pairs in the interval to the list of nested_subpairs\n",
    "        pairs_between[matched_pair] = {\n",
    "                    pair for pair in matched_pairs if pair != matched_pair and\n",
    "                    min_type <= self.X_types[pair[0]] <= max_type and\n",
    "                    min_type <= self.Y_types[pair[1]] <= max_type}\n",
    "    \n",
    "    subpairs = {key: value.copy() for key, value in pairs_between.items()}\n",
    "\n",
    "    # Remove nested pairs\n",
    "    for matched_pair in matched_pairs | {'artificial_pair'}:\n",
    "        # Compute all nested subpairs \n",
    "        nested_subpairs = set(chain.from_iterable(subpairs[pair] \n",
    "                                for pair in subpairs[matched_pair])) \n",
    "        # Remove nested pairs from subpairs[matched_pair]\n",
    "        subpairs[matched_pair] -= nested_subpairs\n",
    "        # subpairs[matched_pair].discard(matched_pair)\n",
    "        subpairs[matched_pair] = list(subpairs[matched_pair])\n",
    "\n",
    "        # Order the subpairs: \n",
    "        # the first (last) pair should have x (y) close to pair_x (pair_y)\n",
    "        if matched_pair != 'artificial_pair' and len(subpairs[matched_pair]) > 1:\n",
    "            subpairs[matched_pair] = self.sort_subpairs(\n",
    "                subpairs[matched_pair], \n",
    "                x_smaller_y=self.X_types[matched_pair[0]] \n",
    "                            < self.Y_types[matched_pair[1]])\n",
    "\n",
    "    if return_pairs_between:\n",
    "        return subpairs, pairs_between\n",
    "    return subpairs\n",
    "\n",
    "OffDiagonal.find_subpairs = find_subpairs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5db1af85",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "subpairs, pairs_between = exam_assign_OD.find_subpairs(assignment, \n",
    "                                                return_pairs_between = True)\n",
    "subpairs"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1c0dfce3",
   "metadata": {},
   "source": [
    "The algorithm to compute the dual variables has a hierarchical structure: it starts from the matched pairs with no subpairs and then moves to those pairs whose subpairs have been already processed.\n",
    "\n",
    "We can visualize the hierarchical structure by computing the order in which he pairs will be processed and plotting the matching with color of the arcs corresponding the hierarchy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "41614a69",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "## Compute Hierarchies\n",
    "\n",
    "def find_hierarchies(subpairs):  \n",
    "    \n",
    "    # Initialize sets for faster membership checks\n",
    "    pairs_to_process = set(subpairs.keys())  # All pairs to process\n",
    "    processed_pairs = set()  # Pairs that have been processed\n",
    "\n",
    "    # Initialize ready_to_process with pairs that have no subpairs\n",
    "    ready_to_process = {pair for pair, sublist in subpairs.items() \n",
    "                        if len(sublist) == 0}\n",
    "\n",
    "    # Initialize hierarchies with the first level\n",
    "    hierarchies = [list(ready_to_process)]\n",
    "\n",
    "    # Continue processing while there are unprocessed pairs\n",
    "    while len(processed_pairs) < len(subpairs):\n",
    "        # Mark ready_to_process pairs as processed\n",
    "        processed_pairs.update(ready_to_process)\n",
    "\n",
    "        # Remove ready_to_process pairs from pairs_to_process\n",
    "        pairs_to_process -= ready_to_process\n",
    "\n",
    "        # Find new ready_to_process pairs that have all their subpairs processed\n",
    "        ready_to_process = {\n",
    "            pair for pair in pairs_to_process\n",
    "            if all(subpair in processed_pairs for subpair in subpairs[pair])}\n",
    "\n",
    "        # Append the new ready_to_process to hierarchies\n",
    "        hierarchies.append(list(ready_to_process))\n",
    "\n",
    "    return hierarchies\n",
    "\n",
    "## Plot Hierarchies\n",
    "\n",
    "def plot_hierarchies(self, subpairs, scatter=True, range_x_axis=None):\n",
    "    # Compute hierarchies\n",
    "    hierarchies = find_hierarchies(subpairs)\n",
    "\n",
    "    # Create the figure and axis\n",
    "    fig, ax = plt.subplots(figsize=(15, 15))\n",
    "\n",
    "    # Plot types on the real line (blue for X_types, red for Y_types)\n",
    "    size_marker = 20 if scatter else 0\n",
    "    ax.scatter(self.X_types, np.zeros_like(self.X_types), color='blue', \n",
    "                    s=size_marker, zorder=5, label='X_types')\n",
    "    ax.scatter(self.Y_types, np.zeros_like(self.Y_types), color='red', \n",
    "                    s=size_marker, zorder=5, label='Y_types')\n",
    "\n",
    "    # Plot arcs\n",
    "    # Create a colormap ('viridis' or 'coolwarm', 'plasma')\n",
    "    cmap = plt.colormaps['plasma'] \n",
    "    for level, hierarchy in enumerate(hierarchies):\n",
    "        color = (cmap(level / (len(hierarchies) - 1)) \n",
    "                        if len(hierarchies) > 1 else cmap(0))\n",
    "        for pair in hierarchy:\n",
    "            if pair == 'artificial_pair':\n",
    "                continue  \n",
    "\n",
    "            min_type, max_type = sorted([self.X_types[pair[0]], \n",
    "                                         self.Y_types[pair[1]]])\n",
    "            width = max_type - min_type\n",
    "            center = (max_type + min_type) / 2\n",
    "            # Semicircle height can be the same as the width for a perfect arc\n",
    "            height = width  \n",
    "            semicircle = patches.Arc((center, 0), width, height, \n",
    "                            theta1=0, theta2=180, \n",
    "                            color=color, lw = 3)\n",
    "            ax.add_patch(semicircle)\n",
    "\n",
    "    if range_x_axis is not None:\n",
    "        ax.set_xlim(range_x_axis)\n",
    "        ax.set_ylim(- self.X_types.ptp() / 10, \n",
    "                    (range_x_axis[1] - range_x_axis[0]) / 2 )\n",
    "\n",
    "    # Title and layout settings for the main plot\n",
    "    plt.title('Hierarchies of the optimal matching (off-diagonal)')\n",
    "    ax.set_aspect('equal')\n",
    "    plt.axhline(0, color='black', linewidth=1)\n",
    "    ax.spines['bottom'].set_position(('data', 0))\n",
    "    ax.spines['left'].set_color('none')\n",
    "    ax.spines['top'].set_color('none')\n",
    "    ax.spines['right'].set_color('none')\n",
    "    ax.yaxis.set_ticks([])  # Hide the y-axis ticks\n",
    "\n",
    "    # Add a colorbar to represent hierarchy levels\n",
    "    sm = cm.ScalarMappable(cmap=cmap, \n",
    "                    norm=Normalize(vmin=0, vmax= len(hierarchies) - 1))\n",
    "    sm.set_array([])\n",
    "    cbar = plt.colorbar(sm, ax=ax, orientation='vertical', pad=0.1, shrink=0.2)   \n",
    "    # Show only min and max levels \n",
    "    cbar.set_ticks([0, len(hierarchies) - 1])\n",
    "    # Label the ticks for clarity\n",
    "    cbar.set_ticklabels(['Lowest', 'Highest']) \n",
    "\n",
    "    plt.show()\n",
    "\n",
    "OffDiagonal.plot_hierarchies = plot_hierarchies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6c44f729",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "exam_assign_OD.plot_hierarchies(subpairs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "665ff690",
   "metadata": {},
   "source": [
    "We proceed to describe and implement the algorithm to compute the dual solution.\n",
    "\n",
    "As already mentioned, the algorithm starts from the matched pairs $ (x_0,y_0) $ with no subpairs and assigns the (temporary) values $ \\psi_{x_0} = c_{x_0 y_0} $ and $ \\psi_{y_0} = 0, $ i.e. the $ x $ type sustains the whole cost of matching.\n",
    "\n",
    "The algorithm then proceeds sequentially  by processing any matched pair whose subpairs have already been processed.\n",
    "\n",
    "After picking any such matched pair $ (x_0,y_0) $, the dual variables already computed for the processed subpairs need to be made “comparable”.\n",
    "\n",
    "Indeed, for any subpair $ (x_1,y_1) $ of $ (x_0,y_0) $, the dual variables of all the types between the $ x_1 $ and $ y_1 $ satisfy dual feasibility and complementary slackness *locally*, i.e. $ \\phi_x + \\psi_y \\leq c_{xy} $ with equality if $ (x,y) $ is a matched pair for all types $ x,y $ between  $ x_0 $ and $ y_0. $\n",
    "\n",
    "But dual feasibility is not satisfied globally in general, for instance it might not be satisfied for two subpairs $ (x_1,y_1) $ and  $ (x_2,y_2) $ of $ (x_0,y_0). $\n",
    "\n",
    "Therefore, letting  $ (x_1,y_1), \\dots,  (x_p,y_p) $ be the subpairs of $ (x_0,y_0), $ we compute the solution $ (\\beta_2, \\dots, \\beta_p) $ of the linear system\n",
    "\n",
    "$$\n",
    "\\max (c_{x_0 y_0} - c_{x_0 y_i} - c_{x_j y_0} , - c_{x_j y_i}) + c_{x_i y_i} \n",
    "\\leq \\sum_{k=i+1}^{j} \\beta_k \n",
    "\\leq \\min (c_{x_0 y_j} + c_{x_i y_0} - c_{x_0 y_0} , c_{x_i y_j}) -  c_{x_j y_j} , \\quad \\text{for all } 1 \\leq i < j \\leq p.\n",
    "$$\n",
    "\n",
    "Then for all $ i \\in [p] $ compute the adjustment $ \\Delta_i = \\sum_{k = i+1}^p \\beta_k + \\phi_{x_p} - \\phi_{x_1} $ and modify the dual variables\n",
    "\n",
    "$$\n",
    "\\begin{aligned}\n",
    "\\phi_{x} &\\leftarrow \\phi_{x} + \\Delta_i \\\\\n",
    "\\psi_{y} &\\leftarrow \\psi_{y} - \\Delta_i,\n",
    "\\end{aligned}\n",
    "$$\n",
    "\n",
    "for all matched pairs $ (x,y) $ between $ x_i $ and $ y_i. $\n",
    "\n",
    "After this step, the dual variables of the types between $ x_0 $ and $ y_0 $ satisfy dual feasibility and complementary slackness; we can then proceed to compute the dual variables for $ x_0 $ and $ y_0 $ by setting\n",
    "\n",
    "$$\n",
    "\\begin{aligned}\n",
    "&\\psi_{y_0} = \\min_{i \\in [p]} \\{ c_{x_i y_0} - \\phi_{x_i} \\} \\\\\n",
    "&\\phi_{x_0} = c_{x_0 y_0}  - \\psi_{y_0}.\n",
    "\\end{aligned}\n",
    "$$\n",
    "\n",
    "The pair $ (x_0,y_0) $ is now processed.\n",
    "\n",
    "The following method computes the solution $ \\beta $ of the linear system of inequalities above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b0331c2a",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def compute_betas(self, pair, subpairs):\n",
    "    types_subpairs = np.array(subpairs)\n",
    "\n",
    "    # Define the bounds of the linear inequality system\n",
    "    if pair == 'artificial_pair':\n",
    "        bounds = (- self.cost_x_y[types_subpairs[:,0][:,None], \n",
    "                types_subpairs[:,1][None,:]]\n",
    "                + self.cost_x_y[types_subpairs[:,0], \n",
    "                                types_subpairs[:,1]][None,:])\n",
    "    else:\n",
    "        bounds = (np.maximum(\n",
    "        self.cost_x_y[pair]  \n",
    "        - self.cost_x_y[pair[0], types_subpairs[:,1]][None,:] \n",
    "        - self.cost_x_y[types_subpairs[:,0],pair[1]][:,None],\n",
    "        - self.cost_x_y[types_subpairs[:,0][:,None], \n",
    "        types_subpairs[:,1][None,:]]\n",
    "        )\n",
    "        + self.cost_x_y[types_subpairs[:,0], types_subpairs[:,1]][None,:])\n",
    "    \n",
    "    # Define linear inequality system \n",
    "    num_subpairs = len(types_subpairs)\n",
    "    c_1 = (np.arange(num_subpairs)[:, None, None] \n",
    "            >= np.arange(num_subpairs)[None, None, :])\n",
    "    c_2 = (np.arange(num_subpairs)[None, None, :] \n",
    "            > np.arange(num_subpairs)[ None,:, None])\n",
    "    sum_tensor = (c_1 & c_2).astype(int)\n",
    "\n",
    "    sum_tensor -= sum_tensor.transpose(1, 0, 2) \n",
    "    \n",
    "    # Solve the system of linear inequalities\n",
    "    result = linprog(c = np.zeros(num_subpairs), \n",
    "                    A_ub= - sum_tensor.reshape(num_subpairs**2, num_subpairs), \n",
    "                    b_ub= - bounds.flatten(), \n",
    "                    bounds=(None,None), \n",
    "                    method='highs')\n",
    "    \n",
    "    beta = result.x\n",
    "    beta[0] = 0 \n",
    "\n",
    "    return beta\n",
    "    \n",
    "OffDiagonal.compute_betas = compute_betas"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7a2d5265",
   "metadata": {},
   "source": [
    "The following method iteratively processes the matched pairs of the off-diagonal matching as explained above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3ec7e6b5",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def compute_dual_off_diagonal(self, subpairs, pairs_between):\n",
    "\n",
    "    # Initialize dual variables\n",
    "    ϕ_x = np.zeros(len(self.X_types))\n",
    "    ψ_y = np.zeros(len(self.Y_types))\n",
    "\n",
    "    # Initialize sets for faster membership checks\n",
    "    pairs_to_process = set(subpairs.keys())  # All pairs to process\n",
    "    processed_pairs = set()  # Pairs that have been processed\n",
    "\n",
    "    # Initialize ready_to_process with pairs that have no subpairs\n",
    "    ready_to_process = {pair for pair, sublist in subpairs.items() \n",
    "                        if len(sublist) == 0}\n",
    "\n",
    "    while len(processed_pairs) < len(subpairs): \n",
    "        \n",
    "        # 1. Pick any subpair which is ready to process \n",
    "        for pair in ready_to_process:\n",
    "\n",
    "            # 2. If there are no subpairs, φ_x = c_{xy} and ψ_y = 0\n",
    "            if len(subpairs[pair]) == 0:\n",
    "                ϕ_x[pair[0]] = self.cost_x_y[pair]\n",
    "                ψ_y[pair[1]] = 0\n",
    "\n",
    "            # 3. If there are subpairs:\n",
    "            else:\n",
    "                # (a) compute betas\n",
    "                beta = self.compute_betas(pair, subpairs[pair])\n",
    "\n",
    "                # (b) adjust potentials of types between each subpair of the pair\n",
    "                for i, subpair in enumerate(subpairs[pair]):\n",
    "                    # update potentials of these types\n",
    "                    types_between_subpair = np.array(\n",
    "                                list(pairs_between[subpair]) + [subpair])\n",
    "\n",
    "                    Δ_subpair = (beta[np.arange(i+1,len(subpairs[pair]))].sum() \n",
    "                                 + ϕ_x[subpairs[pair][-1][0]] \n",
    "                                 - ϕ_x[subpair[0]])\n",
    "             \n",
    "                    ϕ_x[ types_between_subpair[:,0]] += Δ_subpair\n",
    "                    ψ_y[ types_between_subpair[:,1]] -= Δ_subpair\n",
    "      \n",
    "                # (c) compute potentials of the pair\n",
    "                subpairs_x = np.array(subpairs[pair])[:,0]\n",
    "                subpairs_y = np.array(subpairs[pair])[:,1]\n",
    "                \n",
    "                if pair != 'artificial_pair':\n",
    "                    if pair[0] == subpairs_x[0]:\n",
    "                        ψ_y[pair[1]] = np.min(self.cost_x_y[pair[0], subpairs_y] \n",
    "                                    - ψ_y[subpairs_y]) + self.cost_x_y[pair]\n",
    "                    else:\n",
    "                        ψ_y[pair[1]] = np.min(self.cost_x_y[subpairs_x, \n",
    "                                         pair[1]] - ϕ_x[subpairs_x] )\n",
    "\n",
    "\n",
    "                    ϕ_x[pair[0]] =  self.cost_x_y[pair] - ψ_y[pair[1]] \n",
    "\n",
    "            # Add pair to processed pairs\n",
    "            processed_pairs.add(pair)\n",
    "      \n",
    "        # Remove ready_to_process from pairs_to_process\n",
    "        pairs_to_process -= ready_to_process\n",
    "\n",
    "        # Add to ready_to_process pairs for which all subpairs are in processed_pairs\n",
    "        ready_to_process = {pair for pair in pairs_to_process \n",
    "            if all(subpair in processed_pairs for subpair in subpairs[pair])}\n",
    "\n",
    "    return ϕ_x, ψ_y\n",
    "\n",
    "OffDiagonal.compute_dual_off_diagonal = compute_dual_off_diagonal"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2ee1fb77",
   "metadata": {},
   "source": [
    "We apply the algorithm to our example and check that dual feasibility ($ \\phi_x + \\psi_y \\leq c_{xy} $ for all $ x \\in X $ and $ y \\in Y $) as well as strong duality ($ V_P = V_D $) are satisfied."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aa94490c",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "ϕ_x , ψ_y = exam_assign_OD.compute_dual_off_diagonal(subpairs, pairs_between)\n",
    "\n",
    "# Check dual feasibility\n",
    "dual_feasibility_i_j = ϕ_x[:,None] + ψ_y[None,:] - exam_assign_OD.cost_x_y \n",
    "print('Violations of dual feasibility:' , np.sum(dual_feasibility_i_j > 1e-10))\n",
    "\n",
    "dual_sol = (exam_assign_OD.n_x * ϕ_x).sum() + (exam_assign_OD.m_y* ψ_y).sum()\n",
    "primal_sol = (assignment_OD * exam_assign_OD.cost_x_y).sum()\n",
    "\n",
    "# Check strong duality\n",
    "print('Value of dual solution: ', dual_sol)\n",
    "print('Value of primal solution: ', primal_sol)\n",
    "\n",
    "# # Check the value of the primal problem\n",
    "if len(exam_assign_OD.n_x) * len(exam_assign_OD.m_y) < 1000:\n",
    "    mu_x_y , p_z= solve_1to1(exam_assign_OD.cost_x_y,\n",
    "                            exam_assign_OD.n_x,\n",
    "                            exam_assign_OD.m_y,\n",
    "                            return_dual = True)\n",
    "    print('Value of primal solution (scipy)', \n",
    "    (mu_x_y * exam_assign_OD.cost_x_y).sum())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aea8a062",
   "metadata": {},
   "source": [
    "Having computed the dual variables of the off-diagonal types, we compute the dual variables for perfecly matched pairs by setting\n",
    "\n",
    "$$\n",
    "\\begin{aligned}\n",
    "\\phi_{x} = \\min_{y \\in Y^{OD}} \\{ c_{xy} -\\psi_{y} \\} \\\\\n",
    "\\psi_{y} = \\min_{x \\in X^{OD}} \\{ c_{xy} -\\phi_{x} \\}\n",
    "\\end{aligned}\n",
    "$$\n",
    "\n",
    "where $ X^{OD} $ and $ Y^{OD} $ are the types of the off-diagonal instance, for which the dual variables have already been computed.\n",
    "\n",
    "The following method computes the full dual solution from the primal solution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5087f190",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def compute_dual_solution(self, matching_off_diag):\n",
    "\n",
    "    # Compute the dual solution for the off-diagonal types\n",
    "    off_diag, match_tuple = self.generate_offD_onD_matching()\n",
    "    nonzero_id_x, nonzero_id_y, matching_diag = match_tuple\n",
    "\n",
    "    subpairs, pairs_between = off_diag.find_subpairs(matching_off_diag, \n",
    "                                                return_pairs_between = True)\n",
    "    ϕ_x_off_diag, ψ_x_off_diag = off_diag.compute_dual_off_diagonal(\n",
    "                                        subpairs,pairs_between)\n",
    "    \n",
    "    # Compute the dual solution for the on-diagonal types\n",
    "    ϕ_x = np.ones(len(self.X_types)) * np.inf\n",
    "    ψ_y = np.ones(len(self.Y_types)) * np.inf\n",
    "\n",
    "    ϕ_x[nonzero_id_x] = ϕ_x_off_diag\n",
    "    ψ_y[nonzero_id_y] = ψ_x_off_diag\n",
    "    \n",
    "    ϕ_x = np.min( self.cost_x_y - ψ_y[None,:] , axis = 1)\n",
    "    ψ_y = np.min( self.cost_x_y - ϕ_x[:,None] , axis = 0)\n",
    "\n",
    "    return ϕ_x, ψ_y\n",
    "\n",
    "ConcaveCostOT.compute_dual_solution = compute_dual_solution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fc74b749",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "ϕ_x, ψ_y = exam_assign.compute_dual_solution(assignment_OD)\n",
    "\n",
    "dual_feasibility_i_j = ϕ_x[:,None] + ψ_y[None,:] - exam_assign.cost_x_y \n",
    "print('Violations of dual feasibility:' , np.sum(dual_feasibility_i_j > 1e-10))\n",
    "print('Value of dual solution: ', (exam_assign.n_x * ϕ_x).sum() \n",
    "                                + (exam_assign.m_y * ψ_y).sum())\n",
    "print('Value of primal solution: ', (assignment * exam_assign.cost_x_y).sum())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bc4695d6",
   "metadata": {},
   "source": [
    "## Application"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c65a8b0b",
   "metadata": {},
   "source": [
    "### Data\n",
    "\n",
    "We now replicate the empirical analysis carried out by [[Boerma *et al.*, 2024](https://python-advanced.quantecon.org/zreferences.html#id3)].\n",
    "\n",
    "The dataset is obtained from the American Community Survey and contains individual level data on income, age and occupation.\n",
    "\n",
    "The occupation of each individual consists of a Standard Occupational Classification (SOC) code.\n",
    "\n",
    "There are 497 codes in total.\n",
    "\n",
    "We consider only employed (civilian) individuals with ages between 25 and 60 from 2010 to 2017.\n",
    "\n",
    "To visualize log-wage dispersion, we group the individuals by occupation and compute the mean and standard deviation of the wages within each occupation.\n",
    "\n",
    "Then we sort  occupations by average log-earnings within each occupation.\n",
    "\n",
    "The resulting dataset is included in the dataset `acs_data_summary.csv`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d040435a",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "data_path = '_static/lecture_specific/match_transport/'\n",
    "occupation_df = pd.read_csv(data_path + 'acs_data_summary.csv')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "afc5cfe1",
   "metadata": {},
   "source": [
    "We plot the wage standard deviation for the sorted occupations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cb06563d",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "# Scatter plot wage dispersion for each occupation\n",
    "plt.figure(figsize=(10, 6))\n",
    "\n",
    "# Scatter plot with marker size proportional to count\n",
    "plt.scatter(\n",
    "    occupation_df.index,\n",
    "    occupation_df['std_Earnings'],\n",
    "    # marker_sizes\n",
    "    s = 1000 * (occupation_df['count'] / occupation_df['count'].max()), \n",
    "    # transparency\n",
    "    alpha = 0.5, \n",
    "    label = 'Occupations'\n",
    ")\n",
    "\n",
    "# Polynomial interpolation\n",
    "x = np.arange(len(occupation_df))\n",
    "y = occupation_df['std_Earnings']\n",
    "degree = 5\n",
    "p = np.poly1d(np.polyfit(x, y, degree)  )\n",
    "plt.plot(x, p(x), color='red')\n",
    "\n",
    "# Add labels and title\n",
    "plt.xlabel(\"Occupations\", fontsize=12)\n",
    "plt.ylabel(\"Wage Dispersion\", fontsize=12)\n",
    "plt.xticks([], fontsize=8)\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cf9d6d26",
   "metadata": {},
   "source": [
    "We also plot the average wages for each occupation (SOC code). Again, occupations are ordered by increasing average wage."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "33e216d3",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "# Scatter plot average wage for each occupation\n",
    "plt.figure(figsize=(10, 6))\n",
    "\n",
    "# Scatter plot with marker size proportional to count\n",
    "plt.scatter(\n",
    "    occupation_df.index,\n",
    "    occupation_df['mean_Earnings'],\n",
    "    alpha = 0.5, # transparency\n",
    "    label = 'Occupations'\n",
    ")\n",
    "\n",
    "# Polynomial interpolation\n",
    "x = np.arange(len(occupation_df))\n",
    "y = occupation_df['mean_Earnings']\n",
    "degree = 5\n",
    "p = np.poly1d(np.polyfit(x, y, degree)  )\n",
    "plt.plot(x, p(x), color='red')\n",
    "\n",
    "# Add labels and title\n",
    "plt.xlabel(\"Occupations\", fontsize=12)\n",
    "plt.ylabel(\"Average Wage\", fontsize=12)\n",
    "plt.xticks([], fontsize=8)\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a47e61cf",
   "metadata": {},
   "source": [
    "### Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8331ec8f",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "parameters_1980 = namedtuple('Params_Jobs', [\n",
    "    'mean_1', 'var_1', 'mean_2', 'var_2', 'mixing_weight', 'var_workers'\n",
    "])(\n",
    "    mean_1=0.38,\n",
    "    var_1=0.06,\n",
    "    mean_2=0.0,\n",
    "    var_2=0.75,\n",
    "    mixing_weight=0.36,\n",
    "    var_workers=0.2\n",
    ")\n",
    "\n",
    "num_agents=1500"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "69a58102",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def generate_types_application(self, num_agents, params, random_seed=1):\n",
    "\n",
    "    mean_1, var_1, mean_2, var_2, mixing_weight, var_workers = params\n",
    "\n",
    "    np.random.seed(random_seed)\n",
    "\n",
    "    # Job types\n",
    "    job_types = np.where(np.random.rand(num_agents) < mixing_weight, \n",
    "                     np.random.lognormal(mean_1, var_1, num_agents), \n",
    "                     np.random.lognormal(mean_2, var_2, num_agents))\n",
    "\n",
    "    # Worker types\n",
    "    mean_workers = - var_workers/ 2\n",
    "    worker_types = np.random.lognormal(mean_workers, var_workers, num_agents)\n",
    "\n",
    "    # Check that worker and job types have distinct values\n",
    "    assert len(np.unique(worker_types)) == num_agents\n",
    "    assert len(np.unique(job_types)) == num_agents\n",
    "\n",
    "    # Assign types to the instance\n",
    "    self.X_types = worker_types\n",
    "    self.Y_types = job_types\n",
    "\n",
    "    # Assign unitary marginals\n",
    "    self.n_x = np.ones(num_agents, dtype=int)\n",
    "    self.m_y = np.ones(num_agents, dtype=int)\n",
    "\n",
    "    # Assign cost matrix\n",
    "    self.cost_x_y = np.abs(worker_types[:, None] \\\n",
    "                    - job_types[None, :]) ** (1/self.ζ)\n",
    "\n",
    "ConcaveCostOT.generate_types_application = generate_types_application\n",
    "\n",
    "# Create an instance of ConcaveCostOT class and generate types\n",
    "model_1980 = ConcaveCostOT()\n",
    "model_1980.generate_types_application(num_agents, parameters_1980)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "55ca0149",
   "metadata": {},
   "source": [
    "Since we will consider examples with a large number of agents, it will be convenient to visualize the distributions as histograms approximating the pdfs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "523fa6d5",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def plot_marginals_pdf(self, bins, figsize=(15, 8), \n",
    "                        range_x_axis=None, title='Distributions of types'):\n",
    "\n",
    "    plt.figure(figsize=figsize)\n",
    "\n",
    "    # Plotting histogram for X_types (approximating PDF)\n",
    "    plt.hist(self.X_types, bins=bins, density=True, color='blue', alpha=0.7, \n",
    "             label='PDF of worker types', \n",
    "             edgecolor='blue', range = range_x_axis)\n",
    "\n",
    "    # Plotting histogram for Y_types (approximating PDF)\n",
    "    counts, edges = np.histogram(self.Y_types, bins=bins, \n",
    "                                density=True, range=range_x_axis)\n",
    "    plt.bar(edges[:-1], -counts, width=np.diff(edges), color='red', alpha=0.7, \n",
    "            label='PDF of job types ', align='edge', edgecolor='red')\n",
    "\n",
    "    # Add grid and y=0 axis\n",
    "    plt.grid(False)\n",
    "    plt.axhline(0, color='black', linewidth=1)\n",
    "    plt.gca().spines['bottom'].set_position(('data', 0))\n",
    "\n",
    "    # Set the x-axis limits based on the range argument\n",
    "    if range_x_axis is not None:\n",
    "        plt.xlim(range_x_axis)\n",
    "\n",
    "    # Labeling the axes and the title\n",
    "    plt.ylabel('Density')\n",
    "    plt.title(title)\n",
    "    plt.gca().yaxis.set_major_locator(MaxNLocator(integer=True))\n",
    "    plt.legend()\n",
    "    plt.yticks([])\n",
    "\n",
    "    plt.show()\n",
    "\n",
    "ConcaveCostOT.plot_marginals_pdf = plot_marginals_pdf"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e9e4ecec",
   "metadata": {},
   "source": [
    "We plot the hystograms and the measure of underqualification for the worker types and job types. We then compute the primal solution and plot the matching."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2cbe9585",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "# Plot pdf\n",
    "range_x_axis = (0, 4)\n",
    "model_1980.plot_marginals_pdf(figsize=(8, 5), \n",
    "                              bins=300, range_x_axis=range_x_axis)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "40d2a3cc",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "# Plot H_z\n",
    "model_OD_1980 , _ = model_1980.generate_offD_onD_matching()\n",
    "model_OD_1980.plot_H_z(figsize=(8, 5), range_x_axis=range_x_axis, scatter=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "44b5ebe3",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "# Compute optimal matching and plot off diagonal matching\n",
    "matching_1980, matching_OD_1980, model_OD_1980 = model_1980.solve_primal_DSS()\n",
    "model_OD_1980.plot_matching(matching_OD_1980, \n",
    "                            title = 'Optimal Matching (off-diagonal)', \n",
    "                            figsize=(10, 10), plot_H_z=True, scatter=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8ce3933f",
   "metadata": {},
   "source": [
    "From the optimal matching we compute and visualize the hierarchies.\n",
    "\n",
    "We then find the dual solution $ (\\phi,\\psi) $ and compute the wages as $ w_x = g(x) - \\phi_x, $ assuming that the type-specific productivity of type $ x $ is $ g(x) = x $."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "455a3d8e",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "# Find subpairs and plot hierarchies\n",
    "subpairs, pairs_between = model_OD_1980.find_subpairs(matching_OD_1980, \n",
    "                                        return_pairs_between=True)\n",
    "model_OD_1980.plot_hierarchies(subpairs, scatter=False, \n",
    "                                        range_x_axis=range_x_axis)\n",
    "\n",
    "# Compute dual solution: φ_x and ψ_y\n",
    "ϕ_worker_x_1980 , ψ_firm_y_1980 = model_OD_1980.compute_dual_off_diagonal(\n",
    "                                        subpairs, pairs_between)\n",
    "\n",
    "# Check dual feasibility\n",
    "dual_feasibility_i_j = ϕ_worker_x_1980[:,None] + ψ_firm_y_1980[None,:] \\\n",
    "                       - model_OD_1980.cost_x_y \n",
    "print('Dual feasibility violation:', dual_feasibility_i_j.max())\n",
    "\n",
    "# Check strong duality\n",
    "dual_sol = (model_OD_1980.n_x * ϕ_worker_x_1980).sum() \\\n",
    "            + (model_OD_1980.m_y * ψ_firm_y_1980).sum()\n",
    "primal_sol = (matching_OD_1980 * model_OD_1980.cost_x_y).sum()\n",
    "\n",
    "print('Value of dual solution: ', dual_sol)\n",
    "print('Value of primal solution: ', primal_sol)\n",
    "\n",
    "# Compute wages: wage_x = x - φ_x \n",
    "wage_worker_x_1980 = model_1980.X_types - ϕ_worker_x_1980"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a907935f",
   "metadata": {},
   "source": [
    "Let’s plot  average wages and wage dispersion generated by the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4ba44414",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def plot_wages_application(wages):\n",
    "\n",
    "    plt.figure(figsize=(10, 6))\n",
    "    plt.plot(np.sort(wages), label='Wages')\n",
    "    plt.xlabel(\"Occupations\", fontsize=12)\n",
    "    plt.ylabel(\"Wages\", fontsize=12)\n",
    "    plt.grid(True)\n",
    "    plt.show()\n",
    "\n",
    "\n",
    "def plot_wage_dispersion_model(wage_worker_x, bins=100, \n",
    "                                title='Wage Dispersion', figsize=(10, 6)):\n",
    "    # Compute the percentiles \n",
    "    percentiles = np.linspace(0, 100, bins + 1)\n",
    "    bin_edges = np.percentile(wage_worker_x, percentiles)\n",
    "    \n",
    "    # Compute the standard deviation within each percentile range\n",
    "    stds = []\n",
    "    for i in range(bins):\n",
    "        # Compute the standard deviation for the current bin \n",
    "        bin_data = wage_worker_x[\n",
    "        (wage_worker_x >= bin_edges[i]) & (wage_worker_x < bin_edges[i + 1])]\n",
    "        if len(bin_data) > 1:\n",
    "            stds.append(np.std(bin_data))\n",
    "        else:\n",
    "            stds.append(0)\n",
    "    \n",
    "    # Plot the standard deviations for each percentile as bars\n",
    "    plt.figure(figsize=figsize)\n",
    "    plt.bar(range(bins), stds, width=1.0, color='grey', \n",
    "                            alpha=0.7, edgecolor='white')\n",
    "    plt.xlabel('Percentile', fontsize=12)\n",
    "    plt.ylabel('Standard Deviation', fontsize=12)\n",
    "    plt.title(title, fontsize=14)\n",
    "    plt.grid(axis='y', linestyle='-', alpha=0.6)\n",
    "    \n",
    "\n",
    "plot_wages_application(wage_worker_x_1980)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bdad9718",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "plot_wage_dispersion_model(wage_worker_x_1980, bins=100)"
   ]
  }
 ],
 "metadata": {
  "date": 1735958032.2017026,
  "filename": "match_transport.md",
  "kernelspec": {
   "display_name": "Python",
   "language": "python3",
   "name": "python3"
  },
  "title": "Composite Sorting"
 },
 "nbformat": 4,
 "nbformat_minor": 5
}