D1.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# <center>Day 1: competitive equilibrium with gross substitutes</center>\n",
    "### <center>Alfred Galichon (NYU+ScPo)</center>\n",
    "## <center>'math+econ+code' masterclass on equilibrium transport and matching models in economics</center>\n",
    "<center>© 2020-2022 by Alfred Galichon. Support from  NSF DMS-1716489 and ERC CoG-866274 EQUIPRICE grants is acknowledged.</center>\n",
    "\n",
    "#### <center>with Python code</center>\n",
    "\n",
    "**If you reuse material from this masterclass, please cite as:**<br>\n",
    "Alfred Galichon, 'math+econ+code' masterclass on equilibrium transport and matching models in economics, June 2022. https://github.com/math-econ-code/mec_equil"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "# Reference for day 1\n",
    "\n",
    "## Textbooks\n",
    "\n",
    "* Dimitri Bertsekas and John Tsitsiklis (1989). *Parallel and Distributed Computation: Numerical Methods*. Prentice-Hall.\n",
    "* James Ortega and Werner Rheinboldt (1970). *Iterative Solution of Nonlinear Equations in Several Variables*. SIAM.\n",
    "\n",
    "\n",
    "## Papers\n",
    "\n",
    "* Steve Berry, Amit Gandhi and Philip Haile (2013). \"Connected Substitutes and Invertibility of Demand.\" *Econometrica* 81 no. 5, pp. 2087-2111.\n",
    "* Arnaud Dupuy, Alfred Galichon and Marc Henry (2014). \"Entropy methods for identifying hedonic models.\" *Mathematics and Financial Economics* no. 8, pp. 405–416.\n",
    "\n",
    "# Getting started\n",
    "\n",
    "See slides `D1a_getting-started.pdf`.\n",
    "\n",
    "\n",
    "# Generating demand and supply data\n",
    "\n",
    "## The pickup spots\n",
    "\n",
    "For each $z \\in \\{0,...,n-1\\}$, we define $h_z$ and $v_z$ the coordinates (horizontal and vertical) of $z$. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "def vec(X):\n",
    "    return np.reshape(X,(-1, 1))\n",
    "\n",
    "np.random.seed(777)\n",
    "nbz = 3\n",
    "z_df=pd.DataFrame ({'h':np.random.rand(nbz),'v':np.random.rand(nbz)}, columns = ['h','v'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generating supply\n",
    "For each driver $i$, we shall generate the position (horizontal and vertical coordinates), the concavity of preferences $\\tau_i$, and $\\lambda_i$ the value of $i$'s time. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "nbi = 200\n",
    "np.random.seed(778)\n",
    "i_df=pd.DataFrame ({'h':np.random.rand(nbi),\n",
    "                     'v':np.random.rand(nbi),\n",
    "                     'tau':np.random.rand(nbi),\n",
    "                     'lambda':10 * np.random.rand(nbi) }, \n",
    "                    columns = ['h','v','tau','lambda'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generating demand\n",
    "\n",
    "For each passenger $j$, we shall generate the position (horizontal and vertical coordinates), elasticity of price-time substitution $\\sigma_j$, and the valuation of $j$'s time $\\epsilon_j$.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "nbj = 500\n",
    "np.random.seed(779)\n",
    "j_df=pd.DataFrame ({'h':np.random.rand(nbj),\n",
    "                     'v':np.random.rand(nbj),\n",
    "                     'sigma':1/np.random.rand(nbj),\n",
    "                     'epsilon':20*np.random.rand(nbj),\n",
    "                     'eta':np.random.rand(nbj) },\n",
    "                    columns = ['h','v','sigma','epsilon','eta'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "The average speed walking is $4$ km/h and driving is $25$ km/h. Let's compute the time $T_{iz}$ that it takes driver each driver $i$ to drive to pickup at each $z$. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "avge_speed_drive = 25\n",
    "avge_speed_walk = 4\n",
    "\n",
    "from scipy.spatial.distance import cdist\n",
    "\n",
    "T_iz=cdist(i_df[['h','v']],z_df[['h','v']] ) / avge_speed_drive\n",
    "T_jz=cdist(j_df[['h','v']],z_df[['h','v']] ) / avge_speed_walk\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's compute the supply. The utility that each driver $i$ have for each option $z$ is $$u_{iz}\\left( p_z,\\varepsilon \\right)\n",
    "=p_{z}^{1-\\tau _{i}}-\\lambda _{i}T_{iz},$$ and the option that driver $i$ has for the exit option is $0$. We have $$s_z(p_z) = \\sum_i 1{\\{ z = argmax_{z^\\prime} (u_{iz^\\prime}(p_z))  \\}}.$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def s_z(p_z):\n",
    "    u_iz = (p_z[:,np.newaxis]**(1- i_df[['tau']].values.reshape(-1)) -  T_iz.transpose()*i_df[['lambda']].values.reshape(-1)).transpose()\n",
    "    u_i0 = np.zeros((nbi,1))\n",
    "    return np.asarray([np.count_nonzero(np.argmax(np.concatenate((u_iz,u_i0),axis=1), axis=1) == z) for z in range(nbz)] )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We now compute the demand. The cost that each passenger $j$ has for each option $z$, which is given by $$c_{jz}\\left( p_{z}\\right) =\\left( p_{z}^{1-\\frac{1}{%\n",
    "\\sigma _{j}}}+\\left( \\epsilon _{j}T_{jz}\\right) ^{1-\\frac{1}{\\sigma _{j}}%\n",
    "}\\right) ^{\\frac{\\sigma _{j}}{\\sigma _{j}-1}},$$\n",
    "and the cost associated with the exit option is $\\eta_j$. Similar to above, we have $$d_z(p_z) = \\sum_j 1{\\{ z = argmin_{z^\\prime} (v_{jz^\\prime}(p_z))  \\}}.$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def d_z(p_z):\n",
    "    expos_j = 1-1/j_df[['sigma']].values.reshape(-1)\n",
    "    c_jz = ((p_z[:,np.newaxis]**expos_j + ( T_jz.transpose()*j_df[['epsilon']].values.reshape(-1) )**expos_j )**(1/expos_j)).transpose()\n",
    "    c_j0 = j_df[['eta']].values.reshape((-1,1))\n",
    "    return np.asarray([np.count_nonzero(np.argmin(np.concatenate((c_jz,c_j0),axis=1), axis=1) == z) for z in range(nbz)])\n",
    "\n",
    "def Q_z(p_z):\n",
    "    return s_z(p_z)-d_z(p_z)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([199, -31,  -7])"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "p_z=np.random.rand(nbz)\n",
    "Q_z(p_z)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Smoothed supply and demand\n",
    "\n",
    "Actually, we would like to have a smooth approximation of supply and demand. Thus set $$s^{smooth}_z(p_z,T) = \\sum_i \\frac {e^{\\frac {u_{iz}} {T} }} {\\sum_{z^\\prime} e^{\\frac {c_{iz^\\prime}} {T} }},$$ and $$d^{smooth}_z(p_z,T) = \\sum_j \\frac {e^{\\frac {-c_{jz}} {T} }} {\\sum_{z^\\prime} e^{\\frac {-c_{jz^\\prime}} {T} }}.$$ \n",
    "Note (math exercise!) that as $T \\to 0$, we tend to the previous functions $s$ and $d$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def ssmooth_z(p_z,T):\n",
    "    u_iz = (p_z[:,np.newaxis]**(1- i_df[['tau']].values.reshape(-1)) -  T_iz.transpose()*i_df[['lambda']].values.reshape(-1)).transpose()\n",
    "    u_i0 = np.zeros((nbi,1))\n",
    "    max_i = np.max(np.concatenate((u_iz,u_i0),axis=1),axis=1)\n",
    "    utilde_i = max_i - T*np.log ((np.exp( (u_i0.reshape(-1) - max_i)/T) ).reshape(-1) + np.sum(np.exp( ( ( u_iz.transpose() - max_i).transpose())/T) ,axis=1))\n",
    "    return (np.sum( np.exp( ( (  u_iz.transpose() - utilde_i).transpose() )/T) ,axis = 0) )\n",
    "\n",
    "def dsmooth_z(p_z,T):\n",
    "    expos_j = 1-1/j_df[['sigma']].values.reshape(-1)\n",
    "    c_jz = ((p_z[:,np.newaxis]**expos_j + ( T_jz.transpose()*j_df[['epsilon']].values.reshape(-1) )**expos_j )**(1/expos_j)).transpose()\n",
    "    c_j0 = j_df[['eta']].values.reshape((-1,1))\n",
    "    min_j = np.min(np.concatenate((c_jz,c_j0),axis=1),axis=1)\n",
    "    ctilde_j = min_j-T*np.log ((np.exp( (- c_j0.reshape(-1) + min_j)/T) ).reshape(-1) + np.sum(np.exp( ( (- c_jz.transpose() + min_j).transpose())/T) ,axis=1))\n",
    "    return (np.sum( np.exp( ( ( - c_jz.transpose() + ctilde_j).transpose() )/T) ,axis = 0) )\n",
    "\n",
    "def Qsmooth_z(p_z,T):\n",
    "    return (ssmooth_z(p_z,T)-dsmooth_z(p_z,T))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setting up the `EquilibriumProblem` class\n",
    "\n",
    "We are now in a good place to create a Python class, with all the data that will be useful. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "import scipy.optimize as opt\n",
    "\n",
    "class EquilibriumProblem:\n",
    "    \n",
    "    def __init__(self,nbz=0,q_z=np.array([]),p0_z=np.array([]),pmin=-1e10,pmax = 1e10):\n",
    "        if q_z.size == 0:\n",
    "            if nbz != 0:\n",
    "                q_z = np.zeros(nbz)\n",
    "            else:\n",
    "                raise Exception(\"Either nbz or q_z needs to be provided.\")        \n",
    "        elif nbz == 0:\n",
    "            nb_z = q_z.size\n",
    "        if p0_z.size == 0:\n",
    "            p0_z = np.zeros(nbz)\n",
    "        self.nbz = nbz\n",
    "        self.q_z = q_z\n",
    "        self.p0_z = p0_z\n",
    "        self.Q_z = Q_z\n",
    "        self.pmin = pmin\n",
    "        self.pmax = pmax            \n",
    "        self.eq_nbsteps = []\n",
    "        self.eq_p_z = []\n",
    "        self.eq_diffp_z = []\n",
    "        self.eq_delta_z = []\n",
    "        self.eq_code = []\n",
    "    \n",
    "    def Q_z(self,p_z):\n",
    "        return Q_z(p_z)\n",
    "    \n",
    "\n",
    "            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([-83.12771217, -74.08853699, -15.33273283])"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Q_z = lambda p_z: Qsmooth_z(p_z,0.001)\n",
    "\n",
    "EqPb = EquilibriumProblem(nbz,pmin=0)\n",
    "p_z = np.zeros(nbz)\n",
    "EqPb.Q_z(p_z)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Equilibrium computation\n",
    "\n",
    "We now would like to construct a toolbox to solve the equilibrium problem $$Q(p)=q.$$\n",
    "We will see two related basic methods, the Gauss-Seidel algorithm and the Jacobi algorithm. Both of them are based ont the notion of coordinate updating.\n",
    "\n",
    "## Building coordinate update functions\n",
    "\n",
    "To start with, we need to introduce a *coordinate update function,* denoted $cu^z(p),$ such that \n",
    "$$Q_z(cu^z_z(p),p_{-z})=q_z\\\\\n",
    "cu^z_{-z}(p)=p_{-z}\n",
    "$$\n",
    "so that $cu^z_z$ finds the equilibrium price in the market for $z$, provided all the other are markets are at equilibrium. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "def replace_entry(ind,vect,value):\n",
    "    np.put(vect,ind,value)\n",
    "    return vect\n",
    "\n",
    "\n",
    "def cu(self,z,p_z):\n",
    "    oldpz = p_z[z]\n",
    "    newpz = opt.brentq(lambda thep : self.Q_z(replace_entry(z,p_z,thep))[z]-self.q_z[z],self.pmin,self.pmax)\n",
    "    p_z[z] = oldpz\n",
    "    return newpz, (newpz-oldpz)    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Above, `self` is intended to be an object of `EquilibriumProblem` class. We can just add the method `cu` to that class."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0.00016693 0.         0.        ]\n",
      "[0.00016693 0.00027716 0.        ]\n",
      "[0.00016693 0.00027716 0.00028675]\n"
     ]
    }
   ],
   "source": [
    "EquilibriumProblem.cu = cu\n",
    "\n",
    "# Try it:\n",
    "p_z = np.zeros(nbz)\n",
    "p_z[0],_=EqPb.cu(0,p_z)\n",
    "print(p_z)\n",
    "p_z[1],_=EqPb.cu(1,p_z)\n",
    "print(p_z)\n",
    "p_z[2],_=EqPb.cu(2,p_z)\n",
    "print(p_z)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0.00016693 0.         0.        ]\n",
      "[0.00000000e+00 1.73776256e-05 0.00000000e+00]\n",
      "[0.000000e+00 0.000000e+00 1.892368e-12]\n"
     ]
    }
   ],
   "source": [
    "p_z = np.zeros(nbz)\n",
    "pp_z = np.zeros(nbz)\n",
    "pp_z[0],_=EqPb.cu(0,p_z)\n",
    "print(pp_z)\n",
    "pp_z = np.zeros(nbz)\n",
    "pp_z[1],_=EqPb.cu(1,p_z)\n",
    "print(pp_z)\n",
    "pp_z = np.zeros(nbz)\n",
    "pp_z[2],_=EqPb.cu(2,p_z)\n",
    "print(pp_z)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Gauss-Seidel vs Jacobi algorithms\n",
    "\n",
    "The Gauss-Seidel algorithm consists of setting $p^t$ such that \n",
    "\n",
    "$\n",
    "p^1: Q_1(p^1_1,p^0_{-1})=q_1 , p^1_{-1} = p^0_{-1} \\\\\n",
    "p^2: Q_2(p^2_2,p^1_{-2})=q_2, p^2_{-2} = p^1_{-2} \\\\\n",
    "... \\\\\n",
    "p^Z: Q_z (p^Z_Z,p^{Z-1}_{-Z})=q_Z, p^Z_{-Z} = p^{Z-1}_{-Z} \\\\ \n",
    "p^{Z+1}: Q_{1} (p^{Z+1}_1,p^{Z}_{-1}) = q_1, p^{Z+1}_{-1} = p^{Z+1}_{-1} \\\\\n",
    "...\n",
    "$\n",
    "\n",
    "In other words, one full iteration of the Gauss-Seidel algorithm amounts to iterating the map\n",
    "\n",
    "$f^{GS}(p)=(cu^1(p),cu^2\\circ cu^1(p),...,cu^Z\\circ cu^{Z-1}\\circ ... \\circ cu^1(p)),$\n",
    "\n",
    "which we now implement."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fGS_z(self,p_z):\n",
    "    pp_z = np.copy(p_z)\n",
    "    diffp_z = np.zeros(self.nbz)\n",
    "    for z in range(nbz):\n",
    "        pp_z[z],diffp_z[z] = self.cu(z,pp_z)\n",
    "    return(pp_z,diffp_z)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The Jacobi algorithm consists of setting $p^t$ such that \n",
    "\n",
    "$\n",
    "Q_1(p^1_1,p^0_{-1}) = q_1  \\\\\n",
    "Q_2(p^1_2,p^0_{-2}) = q_2 \\\\\n",
    "... \\\\\n",
    "Q_z (p^1_Z,p^0_{-Z}) = q_Z \\\\ \n",
    "Q_{1} (p^2_1,p^1_{-1}) = q_1\\\\\n",
    "...\n",
    "$\n",
    "\n",
    "In other words, one full iteration of the Gauss-Seidel algorithm amounts to \n",
    "\n",
    "$f^{J}(p)=(cu^1(p),cu^2(p),...,cu^Z(p)).$\n",
    "\n",
    "Let's implement $f^J$:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fJ_z(self,p_z):\n",
    "    pp_z = np.zeros(self.nbz)\n",
    "    diffp_z = np.zeros(self.nbz)\n",
    "    for z in range(nbz):\n",
    "        pp_z[z],diffp_z[z]=self.cu(z,p_z)\n",
    "    return(pp_z,diffp_z)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The first five steps of the Gauss-Seidel algorithm will therefore be:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0.00016693 0.00027716 0.00028675]\n",
      "[0.0005505  0.00088081 0.00090476]\n",
      "[0.00122043 0.00164713 0.00162068]\n",
      "[0.00204997 0.00205102 0.00202434]\n",
      "[0.00239362 0.00240875 0.00237726]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "array([-6.97605609e+00, -1.45690873e+00,  4.84945417e-13])"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "p_z=np.zeros(nbz)\n",
    "for i in range(5):\n",
    "    p_z,_ = fGS_z(EqPb,p_z)\n",
    "    print(p_z)\n",
    "\n",
    "EqPb.Q_z(p_z) - EqPb.q_z"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The first five steps of the Jacobi algorithm will therefore be:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[1.66926158e-04 1.73776256e-05 1.89236800e-12]\n",
      "[0.00024788 0.00027722 0.00014836]\n",
      "[0.00050723 0.00047984 0.00028993]\n",
      "[0.00074001 0.00082783 0.00051344]\n",
      "[0.00113698 0.00115414 0.00085994]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "array([-14.92463311,  -7.67069782,  -9.16727684])"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "p_z=np.zeros(nbz)\n",
    "for i in range(5):\n",
    "    p_z,_ = fJ_z(EqPb,p_z)\n",
    "    print(p_z)\n",
    "    \n",
    "EqPb.Q_z(p_z) - EqPb.q_z"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Compared with Gauss-Seidel above, Jacobi does not make use of all the relevant information (i.e. latest price update) at a given point in time. But its structure makes it more naturally suited for parallelization -- we will return to that.\n",
    "\n",
    "Let's provide a `solve` method to `EquilibriumProblem` which computes both Gauss-Seidel and Jacobi."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "from time import time\n",
    "\n",
    "def solve(self,maxit = 1000, method = 'GS', valtol=1e-5, steptol=1e-9,output=0):\n",
    "    start_time = time()\n",
    "    code = 0\n",
    "    delta_z = np.zeros(self.nbz)\n",
    "    p_z = np.copy(self.p0_z)\n",
    "    if method=='GS':\n",
    "        method_name = 'Gauss-Seidel'\n",
    "    elif method=='J':\n",
    "        method_name = 'Jacobi'\n",
    "    else:\n",
    "        raise Exception('Method code '+method+' is not available.')\n",
    "    for i in range(maxit):\n",
    "        if method=='GS':\n",
    "            p_z,diffp_z = self.fGS_z(p_z)\n",
    "        elif method=='J':\n",
    "            p_z,diffp_z = self.fJ_z(p_z)    \n",
    "        else:\n",
    "            raise Exception('Method code '+method+' is not available.')\n",
    "        delta_z = self.Q_z(p_z) - self.q_z\n",
    "        if output > 1 :\n",
    "            print(\"p=\",p_z)\n",
    "        if np.max(np.abs(delta_z) ) < valtol :\n",
    "            code = 0\n",
    "            break\n",
    "        if np.max(np.abs(diffp_z) ) < steptol :\n",
    "            code = 1\n",
    "            break\n",
    "        code = 2    \n",
    "    comp_time = time() - start_time\n",
    "    if output > 0 :\n",
    "        print(method_name, 'method converged in', i, 'iterations.')\n",
    "        print('Value of p =',p_z)\n",
    "        print('Value of p\\'_z - p_z =',diffp_z)            \n",
    "        print('Value of e(p)-q =',delta_z)\n",
    "        print('Time elapsed =',comp_time,' s.')\n",
    "        print('Code =',code)\n",
    "    \n",
    "    self.eq_p_z = p_z\n",
    "    self.eq_diffp_z = diffp_z\n",
    "    self.eq_delta_z = delta_z\n",
    "    self.comp_code = code\n",
    "    self.comp_nbsteps = i\n",
    "    self.comp_time = comp_time\n",
    "    return code\n",
    "\n",
    "\n",
    "EquilibriumProblem.fGS_z = fGS_z\n",
    "EquilibriumProblem.fJ_z = fJ_z\n",
    "EquilibriumProblem.solve = solve\n",
    "        \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Gauss-Seidel method converged in 26 iterations.\n",
      "Value of p = [0.00404359 0.00398981 0.00398592]\n",
      "Value of p'_z - p_z = [-4.25764874e-10 -3.68276937e-11 -6.54447174e-10]\n",
      "Value of e(p)-q = [-3.54566518e-06  9.11272851e-06  1.52858348e-09]\n",
      "Time elapsed = 6.346354246139526  s.\n",
      "Code = 0\n"
     ]
    }
   ],
   "source": [
    "code = EqPb.solve(valtol=1e-5,steptol = 1e-12,method='GS',output=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Jacobi method converged in 64 iterations.\n",
      "Value of p = [0.00404359 0.00398981 0.00398593]\n",
      "Value of p'_z - p_z = [-4.68941770e-10  7.31033839e-10  1.58366912e-09]\n",
      "Value of e(p)-q = [3.16996350e-06 3.27073918e-06 9.37327900e-06]\n",
      "Time elapsed = 14.994502067565918  s.\n",
      "Code = 0\n"
     ]
    }
   ],
   "source": [
    "code = EqPb.solve(valtol=1e-5,steptol = 1e-12,method='J',output=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Convergence of the Gauss-Seidel and Jacobi algorithms\n",
    "\n",
    "See the basics in slides `D1.1_mathematical-results.pdf`.\n",
    "\n",
    "More on convergence will require Perron-Froebenius theory, which will be discussed in Advanced Lecture AL1.\n",
    "\n",
    "## Parallel implementation\n",
    "\n",
    "Parallel implementation on the Cloud will be addressed in Advanced Lecture AL5.  "
   ]
  }
 ],
 "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.7.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}