ot06_gravity_equation.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# <center>The gravity equation</center>\n",
    "### <center>Alfred Galichon (NYU & Sciences Po)</center>\n",
    "## <center>'math+econ+code' masterclass series</center>\n",
    "#### <center>With python code examples</center>\n",
    "© 2018-2023 by Alfred Galichon. Past and present support from NSF grant DMS-1716489, ERC grant CoG-866274 are acknowledged, as well as inputs from contributors listed [here](http://www.math-econ-code.org/theteam), in particular Giovanni Montanari.\n",
    "\n",
    "**If you reuse material from this masterclass, please cite as:**<br>\n",
    "Alfred Galichon, 'math+econ+code' masterclass series. https://github.com/math-econ-code/mec_optim"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Learning objectives\n",
    "\n",
    "* Regularized optimal transport\n",
    "\n",
    "* The gravity equation\n",
    "\n",
    "* Generalized linear models\n",
    "\n",
    "* Pseudo-Poisson maximum likelihood estimation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## References\n",
    "\n",
    "* Anderson and van Wincoop (2003). \"Gravity with Gravitas: A Solution to the Border Puzzle\". *American Economic Review*.\n",
    "\n",
    "* Head and Mayer (2014). \"Gravity Equations: Workhorse, Toolkit and Cookbook\". *Handbook of International Economics*.\n",
    "\n",
    "* Choo and Siow (2005). \"Who marries whom and why\". *Journal of Political Economy*.\n",
    "\n",
    "* Gourieroux, Trognon, Monfort (1984). \"Pseudo Maximum Likelihood Methods: Theory\". *Econometrica*.\n",
    "\n",
    "* McCullagh and Nelder (1989). *Generalized Linear Models*. Chapman and Hall/CRC.\n",
    "\n",
    "* Santos Silva and Tenreyro (2006). \"The Log of Gravity\". *Review of Economics and Statistics*.\n",
    "\n",
    "* Yotov et al. (2011). *An advanced guide to trade policy analysis*. WTO.\n",
    "\n",
    "* Guimares and Portugal (2012). \"Real Wages and the Business Cycle: Accounting for Worker, Firm, and Job Title Heterogeneity\". *AEJ: Macro*.\n",
    "\n",
    "* Dupuy and G (2014), \"Personality traits and the marriage market\". *Journal of Political Economy*.\n",
    "\n",
    "* Dupuy, G and Sun (2019), \"Estimating matching affinity matrix under low-rank constraints\". *Information and Inference*.\n",
    "\n",
    "* Carlier, Dupuy, Galichon and Sun \"SISTA: learning optimal transport costs under sparsity constraints.\" *Communications on Pure and Applied Mathematics* (forthcoming)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Motivation\n",
    "\n",
    "The gravity equation is a very useful tool for explaining trade flows by various measures of proximity between countries.\n",
    "\n",
    "A number of regressors have been proposed. They include: geographic distance, common official languague, common colonial past, share of common religions, etc.\n",
    "\n",
    "The dependent variable is the volume of exports from country $x$ to country $y$, for each pair of country $\\left(  x, y\\right)$.\n",
    "\n",
    "Today, we shall see a close connection between gravity models of international trade and separable matching models."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# The gravity equation\n",
    "\n",
    "\"Structural gravity equation\" (Anderson and van Wincoop, 2003) as exposited in Head and Mayer (2014)\n",
    "handbook chapter:\n",
    "\n",
    "\\begin{align*}\n",
    "\\mu_{xy}=\\frac{n_x}{\\Psi_{x}} \\frac{m_y}{\\Omega_{y}}  \\Phi_{xy}%\n",
    "\\end{align*}\n",
    "\n",
    "where $x$=exporter,  $y$=importer, $\\mu_{xy}$=trade flow from $x$ to $y$, $n_x=\\sum_{y}\\mu_{xy}$ is value of production,  $m_y=\\sum_{x}\\mu_{xy}$ is importers' expenditures, and $\\phi_{xy}$=bilateral accessibility of $x$ to $y$.\n",
    "\n",
    "$\\Omega_{y}$ and $\\Psi_{x}$ are *multilateral resistances*, satisfying the set of implicit equations\n",
    "\n",
    "\\begin{align*}\n",
    "\\Psi_{x}=\\sum_{y}\\frac{\\Phi_{xy}m_y}{\\Omega_{y}}\\text{ and }\\Omega_{y}%\n",
    "=\\sum_{x}\\frac{\\Phi_{xy}n_x}{\\Psi_{x}}%\n",
    "\\end{align*}\n",
    "\n",
    "We will see that these are exactly the same equations as those of the regularized OT."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Explaining trade\n",
    "\n",
    "Parameterize $\\Phi_{xy}=\\exp\\left(  \\sum_{k=1}^{K}\\beta_{k}D_{xy}^{k}\\right)  $, where the $D_{xy}^{k}$ are $K$ pairwise measures of distance between $x$ and $y$. We have\n",
    "\n",
    "\\begin{align*}\n",
    "\\mu_{xy}=\\exp\\left(  \\sum_{k=1}^{K}\\beta_{k}D_{xy}^{k}-a_{x}-b_{y}\\right)\n",
    "\\end{align*}\n",
    "\n",
    "where fixed effects $b_{y}=-\\ln \\frac{m_y}{\\Omega_{y}}$ and $a_{x}=-\\ln \\frac{n_x}{\\Psi_{x}}$ are adjusted by\n",
    "\n",
    "\\begin{align*}\n",
    "\\sum_{x}\\mu_{xy}=n_x\\text{ and }\\sum_{y}\\mu_{xy}=m_y.\n",
    "\\end{align*}\n",
    "\n",
    "Standard choices of $D_{xy}^{k}$'s:\n",
    "\n",
    "* Logarithm of bilateral distance between $x$ and $y$\n",
    "\n",
    "* Indicator of contiguous borders; of common official language; of\n",
    "colonial ties\n",
    "\n",
    "* Trade policy variables: presence of a regional trade agreement; tariffs\n",
    "\n",
    "* Could include many other measures of proximity, e.g. measure of genetic/cultural distance, intensity of communications, etc."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Regularized optimal transport\n",
    "\n",
    "Consider the optimal transport duality\n",
    "\n",
    "\\begin{align*}\n",
    "\\max_{\\mu\\in\\mathcal{M}\\left(  P,Q\\right)  }\\sum_{xy}\\mu_{xy}\\Phi_{xy}=\\min_{u_{x}+v_{y}\\geq\\Phi_{xy}}\\sum_{x\\in\\mathcal{X}}n_xu_{x}+\\sum_{y\\in\\mathcal{Y}}m_yv_{y}\n",
    "\\end{align*}\n",
    "\n",
    "Now let's assume that we are adding an entropy to the primal objective function. For any $\\sigma>0$, we get\n",
    "\n",
    "\\begin{align*}\n",
    "&  \\max_{\\mu\\in\\mathcal{M}\\left(  P,Q\\right)  }\\sum_{xy}\\mu_{xy}\\Phi_{xy}-\\sigma\\sum_{xy}\\mu_{xy}\\ln\\mu_{xy}\\\\\n",
    "&  =\\min_{u,v}\\sum_{x\\in\\mathcal{X}}n_xu_{x}+\\sum_{y\\in\\mathcal{Y}}m_y v_{y}+\\sigma\\sum_{xy}\\exp\\left(  \\frac{\\Phi_{xy}-u_{x}-v_{y}-\\sigma}{\\sigma}\\right)\n",
    "\\end{align*}\n",
    "\n",
    "The latter problem is an unconstrained convex optimization problem. But the most efficient numerical computation technique is often coordinate descent, i.e. alternate between minimization in $u$ and minimization in $v$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Iterated fitting\n",
    "\n",
    "Maximize wrt to $u$ yields\n",
    "\n",
    "\\begin{align*}\n",
    "e^{-u_{x}/\\sigma}=\\frac{n_x}{\\sum_{y}\\exp\\left(  \\frac{\\Phi_{xy}-v_{y}-\\sigma}{\\sigma}\\right)  }\n",
    "\\end{align*}\n",
    "\n",
    "and wrt $v$ yields\n",
    "\n",
    "\\begin{align*}\n",
    "e^{-v_{y}/\\sigma}=\\frac{m_y}{\\sum_{x}\\exp\\left(  \\frac{\\Phi_{xy}-v_{y}-\\sigma}{\\sigma}\\right)  }\n",
    "\\end{align*}\n",
    "\n",
    "It is called the \"iterated projection fitting procedure\" (ipfp), aka \"matrix scaling\", \"RAS algorithm\", \"Sinkhorn-Knopp algorithm\", \"Kruithof's method\", \"Furness procedure\", \"biproportional fitting procedure\", \"Bregman's procedure\". See survey in Idel (2016).\n",
    "\n",
    "Maybe the most often reinvented algorithm in applied mathematics. Recently rediscovered in a machine learning context."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Econometrics of matching\n",
    "\n",
    "The goal is to estimate the matching surplus $\\Phi_{xy}$. For this, take a linear parameterization\n",
    "\n",
    "\\begin{align*}\n",
    "\\Phi_{xy}^{\\beta}=\\sum_{k=1}^{K}\\beta_{k}\\phi_{xy}^{k}.\n",
    "\\end{align*}\n",
    "\n",
    "Following Choo and Siow (2006), Galichon and Salanie (2011) introduce logit heterogeneity in individual preferences and show that the equilibrium now maximizes the *regularized Monge-Kantorovich problem*\n",
    "\n",
    "\\begin{align*}\n",
    "W\\left(  \\beta\\right)  =\\max_{\\mu\\in\\mathcal{M}\\left(  P,Q\\right)  }\\sum_{xy}\\mu_{xy}\\Phi_{xy}^{\\beta}-\\sigma\\sum_{xy}\\mu_{xy}\\ln\\mu_{xy}\n",
    "\\end{align*}\n",
    "\n",
    "By duality, $W\\left(  \\beta\\right)  $ can be expressed\n",
    "\n",
    "\\begin{align*}\n",
    "W\\left(  \\beta\\right)  =\\min_{u,v}\\sum_{x}n_xu_{x}+\\sum_{y}m_yv_{y}+\\sigma\\sum_{xy}\\exp\\left(  \\frac{\\Phi_{xy}^{\\beta}-u_{x}-v_{y}-\\sigma}{\\sigma}\\right)\n",
    "\\end{align*}\n",
    "\n",
    "and w.l.o.g. can set $\\sigma=1$ and drop the additive constant $-\\sigma$ in the $\\exp$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Estimation\n",
    "\n",
    "We observe the actual matching $\\hat{\\mu}_{xy}$. Note that $\\partial W/ \\partial\\beta^{k}=\\sum_{xy}\\mu_{xy}\\phi_{xy}^{k},$ hence $\\beta$ is estimated by running\n",
    "\n",
    "<a name='objFun'></a>\n",
    "\\begin{align*}\n",
    "\\min_{u,v,\\beta}\\sum_{x}n_xu_{x}+\\sum_{y}m_yv_{y}+\\sum_{xy}\\exp\\left(\\Phi_{xy}^{\\beta}-u_{x}-v_{y}\\right)  -\\sum_{xy,k}\\hat{\\mu}_{xy}\\beta_{k}\\phi_{xy}^{k}\n",
    "\\end{align*}\n",
    "\n",
    "which is still a convex optimization problem.\n",
    "\n",
    "As we will show later, this is actually the objective function of the log-likelihood in a Poisson regression with $x$ and $y$ fixed effects, where we assume\n",
    "\n",
    "\\begin{align*}\n",
    "\\mu_{xy}|xy\\sim Poisson\\left(  \\exp\\left(  \\sum_{k=1}^{K}\\beta_{k}\\phi\n",
    "_{xy}^{k}-u_{x}-v_{y}\\right)  \\right)  .\n",
    "\\end{align*}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "To start working with our application, let's load some of the libraries we shall need."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np  # used to work with arrays\n",
    "import pandas as pd # used to load and work with dataframes\n",
    "#import math  # used to work with logs and infinities\n",
    "import time  # used to time the execution of the code\n",
    "\n",
    "import scipy.sparse as spr  # used to work with sparse matrices (when working with the Poisson matrix representation)\n",
    "from sklearn import linear_model  # used to implement the Poisson regression"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And let's load our data, which comes from the book *An Advanced Guide to Trade Policy Analysis: The Structural Gravity Mode*, by Yotov et al. We will estimate the gravity model using optimal transport as well as using Poisson regression.\n",
    "\n",
    "While the table of trade data includes several possible regressors, we focus on four types of regressors: the logarithm of the distance between countries and dummy variables for whether any two countries are contiguous, share a common official language, or share colonial ties. These are regressors known in the literature to have explanatory power, and they are the same as the ones used in Yotov et al.\n",
    "\n",
    "Our data look as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>exporter</th>\n",
       "      <th>importer</th>\n",
       "      <th>year</th>\n",
       "      <th>trade</th>\n",
       "      <th>DIST</th>\n",
       "      <th>ln_DIST</th>\n",
       "      <th>CNTG</th>\n",
       "      <th>LANG</th>\n",
       "      <th>CLNY</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>ARG</td>\n",
       "      <td>ARG</td>\n",
       "      <td>1986</td>\n",
       "      <td>61288.590263</td>\n",
       "      <td>533.908240</td>\n",
       "      <td>6.280224</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>ARG</td>\n",
       "      <td>AUS</td>\n",
       "      <td>1986</td>\n",
       "      <td>27.764874</td>\n",
       "      <td>12044.574134</td>\n",
       "      <td>9.396370</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>ARG</td>\n",
       "      <td>AUT</td>\n",
       "      <td>1986</td>\n",
       "      <td>3.559843</td>\n",
       "      <td>11751.146521</td>\n",
       "      <td>9.371706</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>ARG</td>\n",
       "      <td>BEL</td>\n",
       "      <td>1986</td>\n",
       "      <td>96.102567</td>\n",
       "      <td>11305.285764</td>\n",
       "      <td>9.333026</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>ARG</td>\n",
       "      <td>BGR</td>\n",
       "      <td>1986</td>\n",
       "      <td>3.129231</td>\n",
       "      <td>12115.572046</td>\n",
       "      <td>9.402246</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "  exporter importer  year         trade          DIST   ln_DIST  CNTG  LANG  \\\n",
       "0      ARG      ARG  1986  61288.590263    533.908240  6.280224     0     0   \n",
       "1      ARG      AUS  1986     27.764874  12044.574134  9.396370     0     0   \n",
       "2      ARG      AUT  1986      3.559843  11751.146521  9.371706     0     0   \n",
       "3      ARG      BEL  1986     96.102567  11305.285764  9.333026     0     0   \n",
       "4      ARG      BGR  1986      3.129231  12115.572046  9.402246     0     0   \n",
       "\n",
       "   CLNY  \n",
       "0     0  \n",
       "1     0  \n",
       "2     0  \n",
       "3     0  \n",
       "4     0  "
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "thepath = 'https://raw.githubusercontent.com/math-econ-code/mec_optim_2021-01/master/data_mec_optim/gravity_wtodata/'\n",
    "\n",
    "tradedata = pd.read_csv(thepath + '1_TraditionalGravity_from_WTO_book.csv') # load full table\n",
    "tradedata = tradedata[['exporter', 'importer','year', 'trade', 'DIST','ln_DIST', 'CNTG', 'LANG', 'CLNY']]   # focus on a subset of regressors\n",
    "\n",
    "tradedata.sort_values(['year','exporter','importer'], inplace = True)\n",
    "tradedata.reset_index(inplace = True, drop = True)\n",
    "\n",
    "nbt = len(tradedata['year'].unique())  # number of periods\n",
    "nbi = len(tradedata['importer'].unique())  # number of countries (we have the same number of importers and exporters)\n",
    "nbk = 4  # number of regressors we are interested in using \n",
    "\n",
    "tradedata.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's extract the data from the table and format it using multidimensional tensors. Note that we are storing both absolute and normalized trade flows, as we need the former (properly cleaned of the flows within a country) to normalize the latter. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "years = tradedata['year'].unique()  # array of calendar years reported in the data\n",
    "distances = np.array(['ln_DIST', 'CNTG', 'LANG', 'CLNY'])  # array of trade \"distances\" we are interested in using as regressors\n",
    "\n",
    "D_x_y_t_k = np.zeros((nbi,nbi,nbt,nbk)) # Initialize an empty tensor of dimensions nbi x nbi x nbt x nbk to store distances\n",
    "tradevol_x_y_t = np.zeros((nbi,nbi,nbt)) # Initialize empty tensor nbi x nbi x nbt to store trade volume\n",
    "muhat_x_y_t = np.zeros((nbi,nbi,nbt)) # Initialize empty tensor nbi x nbi x nbt to store normalized trade flows\n",
    "\n",
    "# fill tensors with distance, contiguity, language, and colony variables, as well as the trade flow\n",
    "for t, year in enumerate(years):\n",
    "    tradevol_x_y_t[:, :, t] = np.array(tradedata.loc[tradedata['year'] == year, 'trade']).reshape((nbi, nbi))  # store trade flows\n",
    "    np.fill_diagonal(tradevol_x_y_t[:, :, t], 0)  # set to zero the trade within a country; we will repeat this operation within the estimation functions\n",
    "    for k, distance in enumerate(distances):\n",
    "        D_x_y_t_k[:, :, t, k] = np.array(tradedata.loc[tradedata['year'] == year, distance]).reshape((nbi, nbi))  # store distances\n",
    "\n",
    "# normalize and store trade flows\n",
    "muhat_x_y_t = tradevol_x_y_t / (tradevol_x_y_t.sum() / len(years))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We define a class `GravityModel` to store the data relevant for estimation. Later on we will populate it with our estimation methods."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "class GravityModel():\n",
    "    def __init__(self, muhat_x_y_t, D_x_y_t_k):\n",
    "        self.nbi, _, self.nbt, self.nbk = D_x_y_t_k.shape  # number of countries, periods, and regressors\n",
    "        self.muhat_x_y_t = muhat_x_y_t    # tensor of trade flows over time\n",
    "        self.D_x_y_t_k = D_x_y_t_k  # tensor of bilateral resistances in each time period "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will solve this model by fixing a $\\beta$ and solving the matching problem using IPFP. Then in an outer loop we will solve for the $\\beta$ which minimizes the distance between model and empirical moments, where the optimization is based on gradient descent:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fit_ipfp(self, sigma = 1, maxiterIpfp = 1000, maxiter = 500, tolIpfp = 1e-12, tolDescent = 1e-6, t_s = 0.03):\n",
    "\n",
    "    iterCount = 0\n",
    "    contIter = True\n",
    "    meanD_k =self.D_x_y_t_k.mean(axis=(0,1,2))\n",
    "    sdD_k = self.D_x_y_t_k.std(axis=(0, 1), ddof = 1).mean(axis = 0)\n",
    "    D_x_y_t_k = (self.D_x_y_t_k-meanD_k[None,None,None,:]) / sdD_k[None,None,None,:]\n",
    "    n_x_t = self.muhat_x_y_t.sum(axis=1)\n",
    "    m_y_t = self.muhat_x_y_t.sum(axis=0)\n",
    "\n",
    "    beta_k = np.zeros(self.nbk)\n",
    "\n",
    "    ptm = time.time()\n",
    "    while(contIter):\n",
    "        iterCount += 1\n",
    "        thegrad = np.zeros(nbk)\n",
    "        for t in range(self.nbt):\n",
    "            v_y = np.zeros(self.nbi)\n",
    "            D_xy_k = D_x_y_t_k[:,:,t,:].reshape((-1,self.nbk))\n",
    "            K_x_y = np.exp(D_xy_k @ beta_k / sigma).reshape((self.nbi,self.nbi))\n",
    "            np.fill_diagonal(K_x_y, 0)  # no self-flow: having already exponentiated Phi, we set the diagonal to zero\n",
    "            contIpfp = True\n",
    "            iterIpfp = 0\n",
    "            \n",
    "            while(contIpfp):\n",
    "                iterIpfp += 1\n",
    "                u_x = sigma * np.log(  ( K_x_y @  np.exp(-v_y / sigma)  ) / n_x_t[:,t] ).flatten() \n",
    "                v_y = sigma * np.log(  ( np.exp(-u_x / sigma) @ K_x_y   ) / m_y_t[:,t] ).flatten()\n",
    "                mu_x_y = (K_x_y *  np.exp(-(u_x[:,None] +v_y[None,:] ) / sigma))\n",
    "                if (np.max(np.abs(   mu_x_y.sum(axis=1) /  n_x_t[:,t] - 1)) < tolIpfp or iterIpfp >= maxiterIpfp):\n",
    "                    contIpfp = False\n",
    "                    #print(iterIpfp)\n",
    "            thegrad = thegrad + ((mu_x_y - self.muhat_x_y_t[:, :, t]).flatten().dot(D_xy_k)).flatten()\n",
    "        beta_k = beta_k - t_s * thegrad\n",
    "        #print(beta_k)\n",
    "        \n",
    "        if (iterCount > maxiter or np.sum(np.abs(thegrad)) < tolDescent):  # measure distance against value of the problem\n",
    "            contIter = False\n",
    "\n",
    "    diff = time.time() - ptm\n",
    "    print('Time elapsed = ', diff, 's.')\n",
    "\n",
    "    return np.asarray(beta_k / sdD_k).round(3)\n",
    "\n",
    "GravityModel.fit_ipfp = fit_ipfp"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's test this solution method by initializing an instance of the `GravityModel` class with the data from Yotov et al. :"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "trade_yotov = GravityModel(muhat_x_y_t, D_x_y_t_k)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's run our estimation method:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Time elapsed =  2.7788071632385254 s.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "array([-0.841,  0.437,  0.247, -0.222])"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trade_yotov.fit_ipfp()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We recover the PPML estimates on Table 1 p. 42 of [Yotov et al.'s book](https://www.wto.org/english/res_e/booksp_e/advancedwtounctad2016_e.pdf).\n",
    "\n",
    "We now proceed to show how this problem can be recast as an instance of Poisson regression with fixed effects.\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Poisson regression with fixed effects\n",
    "\n",
    "Let $\\theta=\\left(  \\beta,u,v\\right)  $ and $Z=\\left(  \\phi,D^{x},D^{y}\\right)  $ where $D_{x^{\\prime}y^{\\prime}}^{x}=1\\left\\{  x=x^{\\prime}\\right\\}  $ and $D_{x^{\\prime}y^{\\prime}}^{y}=1\\left\\{  y=y^{\\prime}\\right\\}$ are $x$-and $y$-dummies. Let $\\lambda_{xy}\\left(  Z;\\theta\\right)  =\\exp\\left(\\theta^{\\intercal}Z_{xy}\\right)  $ be the parameter of the Poisson distribution.\n",
    "\n",
    "The conditional likelihood of $\\hat{\\mu}_{xy}$ given $Z_{xy}$ is\n",
    "\n",
    "\\begin{align*}\n",
    "l_{xy}\\left(  \\hat{\\mu}_{xy};\\theta\\right)   &  =\\hat{\\mu}_{xy}\\log \\lambda_{xy}\\left(  Z;\\theta\\right)  -\\lambda_{xy}\\left(  Z;\\theta\\right) \\\\\n",
    "&  =\\hat{\\mu}_{xy}\\left(  \\theta^{\\intercal}Z_{xy}\\right)  -\\exp\\left(\\theta^{\\intercal}Z_{xy}\\right) \\\\\n",
    "&  =\\hat{\\mu}_{xy}\\left(  \\sum_{k=1}^{K}\\beta_{k}\\phi_{xy}^{k}-u_{x}-v_{y}\\right)  -\\exp\\left(  \\sum_{k=1}^{K}\\beta_{k}\\phi_{xy}^{k}-u_{x}-v_{y}\\right)\n",
    "\\end{align*}\n",
    "\n",
    "Summing over $x$ and $y$, the sample log-likelihood is\n",
    "\n",
    "\\begin{align*}\n",
    "\\sum_{xy}\\hat{\\mu}_{xy}\\sum_{k=1}^{K}\\beta_{k}\\phi_{xy}^{k}-\\sum_{x}n_xu_{x}-\\sum_{y}m_yv_{y}-\\sum_{xy}\\exp\\left(  \\sum_{k=1}^{K}\\beta_{k}\\phi_{xy}^{k}-u_{x}-v_{y}\\right)\n",
    "\\end{align*}\n",
    "\n",
    "hence we recover the [objective function](#objFun)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### From Poisson to pseudo-Poisson\n",
    "\n",
    "If $\\mu_{xy}|xy$ is Poisson, then $\\mathbb{E}\\left[\\mu_{xy}\\right]=\\lambda_{xy}\\left(  Z_{xy};\\theta\\right)  =\\mathbb{V}ar\\left(  \\mu_{xy}\\right)  $. While it makes sense to assume the former equality, the latter is a rather strong assumption.\n",
    "\n",
    "For estimation purposes, $\\hat{\\theta}$ is obtained by\n",
    "\n",
    "\\begin{align*}\n",
    "\\max_{\\theta}\\sum_{xy}l\\left(  \\hat{\\mu}_{xy};\\theta\\right)  =\\sum_{xy}\\left(\\hat{\\mu}_{xy}\\left(  \\theta^{\\intercal}Z_{xy}\\right)  -\\exp\\left(\\theta^{\\intercal}Z_{xy}\\right)  \\right)\n",
    "\\end{align*}\n",
    "\n",
    "however, for inference purposes, one shall not assume the Poisson distribution. Instead\n",
    "\n",
    "\\begin{align*}\n",
    "\\sqrt{N}\\left(  \\hat{\\theta}-\\theta\\right)  \\Longrightarrow\\left(A_{0}\\right)  ^{-1}B_{0}\\left(  A_{0}\\right)  ^{-1}\n",
    "\\end{align*}\n",
    "\n",
    "where $N=\\left\\vert \\mathcal{X}\\right\\vert \\times\\left\\vert \\mathcal{Y}\\right\\vert $ and $A_{0}$ and $B_{0}$ are estimated by\n",
    "\n",
    "\\begin{align*}\n",
    "\\hat{A}_{0}  &  =N^{-1}\\sum_{xy}D_{\\theta\\theta}^{2}l\\left(  \\hat{\\mu}_{xy};\\hat{\\theta}\\right)  =N^{-1}\\sum_{xy}\\exp\\left(  \\hat{\\theta}^{\\intercal}Z_{xy}\\right)  Z_{xy}Z_{xy}^{\\intercal}\\\\\n",
    "\\hat{B}_{0}  &  =N^{-1}\\sum_{xy}\\left(  \\hat{\\mu}_{xy}-\\exp\\left(  \\hat{\\theta}^{\\intercal}Z_{xy}\\right)  \\right)  ^{2}Z_{xy}Z_{xy}^{\\intercal}.\n",
    "\\end{align*}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Implementation: Poisson regression with fixed effects"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We now introduce an additional method to recover the coefficients of interest through Poisson regression. Notice that this approach naturally recovers the fixed effects too, although in this application we are not directly interested in them. \n",
    "\n",
    "Consistent with common practice and as previously described, we do not want to consider the trade flow between a country and itself. To incorporate the restriction in the Poisson regression, we use weights and assign zero weight to the trade flow between a country and itself. We can then recover the Poisson objective function expressed in matrix formulation by appropriately stacking the matrix of bilateral resistancies and the matrices of fixed effects."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fit_glm(self, verbosity=0, max_iter = 8000, tol=1e-12, pretest = False):\n",
    "    \"\"\"\n",
    "    fit_glm(args) estimates the gravity equation via weighted Poisson regression.\n",
    "    \"\"\"\n",
    "\n",
    "    kr = spr.kron   # shorthand to implement the Kronecker product\n",
    "\n",
    "    M1 = kr(spr.identity(self.nbi), kr(np.ones((self.nbi, 1)), spr.identity(self.nbt)))\n",
    "    M2 = kr(np.ones((self.nbi, 1)), kr(spr.identity(self.nbi), spr.identity(self.nbt)))\n",
    "    C_a_k = spr.hstack([self.D_x_y_t_k.reshape((-1, self.nbk)), -M1, -M2])\n",
    "    muhat_a = self.muhat_x_y_t.flatten()\n",
    "\n",
    "    weighting_matrix_xyt = kr(np.eye(self.nbi**2), np.ones((self.nbt, 1)))@(np.ones((self.nbi, self.nbi)) - np.eye(self.nbi)).flatten()\n",
    "\n",
    "    clf = linear_model.PoissonRegressor(fit_intercept=False, tol=tol , max_iter=max_iter, verbose=verbosity, alpha=0)\n",
    "    clf.fit(C_a_k, muhat_a, sample_weight=weighting_matrix_xyt)\n",
    "\n",
    "    return clf.coef_[:self.nbk].round(3)\n",
    "\n",
    "GravityModel.fit_glm = fit_glm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Solving the model using the newly added method:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([-0.841,  0.438,  0.248, -0.223])"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "estimates_glm = trade_yotov.fit_glm()\n",
    "estimates_glm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Again, we recover the same estimates as in the book by Yotov et al. "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.12"
  },
  "vscode": {
   "interpreter": {
    "hash": "beda123ca6d46414026d3c59f732de1f5fb19d6ba2f32753cc4223591eed0a9d"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}