Class_13_Inequality.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "https://lectures.quantecon.org/py/aiyagari.html"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Household's problem\n",
    "\n",
    "\\begin{align}\n",
    "            \\max  \\quad & U = E_{0} \\sum_{t=0}^{\\infty} \\beta^{t} u \\left( c_{t} \\right) \\\\\n",
    "\\text{subject to} \\quad & c_{t} + a_{t+1} = w \\cdot z_{t} + \\left( 1+r \\right) a_{t} \\\\\n",
    "                        & c_{t} \\geq 0 \\\\\n",
    "                        & a_{t+1} \\geq -b\n",
    "\\end{align}\n",
    "\n",
    "Exogenous process for $z_{t}$ follows a finite state Markov chain with given stochastic matrix $P$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# from time import time\n",
    "from bisect import bisect\n",
    "\n",
    "from scipy import sparse\n",
    "from scipy.sparse.linalg import spsolve, eigs\n",
    "\n",
    "from scipy import optimize\n",
    "from scipy.interpolate import InterpolatedUnivariateSpline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Household labor efficiency process\n",
    "z_grid = np.array([0.1, 1.0])\n",
    "z_len = 2\n",
    "\n",
    "P = np.array([[0.9, 0.1], \n",
    "              [0.1, 0.9]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Stationary distribution of states\n",
    "\n",
    "D, V = np.linalg.eig(P.T)\n",
    "\n",
    "col = np.where(np.abs(D-1) < 1e-3)[0][0]\n",
    "\n",
    "v = V[:, col]\n",
    "\n",
    "P_bar = v/sum(v)\n",
    "print('P_bar =', P_bar)\n",
    "\n",
    "N = P_bar @ z_grid.T\n",
    "print('N =', N)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Parameters\n",
    "α = 0.33\n",
    "β = 0.96\n",
    "δ = 0.05\n",
    "σ = 1\n",
    "A = 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def Capital_Demand(r):\n",
    "    return (α/(r+δ))**(1/(1-α)) * N\n",
    "\n",
    "def Wage(r):\n",
    "    K = Capital_Demand(r)\n",
    "    return (1-α) * (K/N)**α"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rr = np.linspace(0, 0.1, 100) #1/β-1\n",
    "\n",
    "plt.plot(Capital_Demand(rr), rr, label='Demand for capital')\n",
    "plt.xlabel('Aggregate capital stock $K$')\n",
    "plt.ylabel('Real interest rate $r$')\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "a_min = 0\n",
    "a_max = 20\n",
    "a_len = 1+20\n",
    "\n",
    "a_grid = np.linspace(a_min, a_max, a_len)\n",
    "# a_grid = np.geomspace(a_min+1, a_max+1, a_len)-1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initial guess for r\n",
    "r = 0.02\n",
    "R = 1+r\n",
    "w = Wage(r)\n",
    "w"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Endogenous Gridpoint Method\n",
    "\n",
    "# t0 = time()\n",
    "\n",
    "c_opt = np.zeros((z_len, a_len))\n",
    "c_new = np.zeros((z_len, a_len))\n",
    "\n",
    "a_new = np.zeros((z_len, a_len))\n",
    "a_opt = np.zeros((z_len, a_len))\n",
    "\n",
    "# Initial guess for c_opt\n",
    "for j, z in enumerate(z_grid):\n",
    "    c_opt[j, :] = r*a_grid[:] + z*w\n",
    "\n",
    "norm_pol = 1\n",
    "it = 0\n",
    "maxit = 5000\n",
    "\n",
    "while norm_pol > 1e-6*(1-β) and it < maxit:\n",
    "    ERHS = P @ (β*R * c_opt**(-σ))\n",
    "    IERHS = ERHS**(-1/σ)\n",
    "    for j, z in enumerate(z_grid):\n",
    "        c_end = IERHS[j, :]\n",
    "        a_end = (a_grid[:] + c_end - z*w) / R\n",
    "    \n",
    "        c_new[j, :] = InterpolatedUnivariateSpline(a_end, c_end, k=1)(a_grid)\n",
    "        \n",
    "        # Binding constraint\n",
    "        bind = np.where(a_grid < a_end[0])\n",
    "        c_new[j, bind] = R*a_grid[bind] + z*w - a_min\n",
    "        \n",
    "        a_new[j, :] = R*a_grid + z*w - c_new[j, :]\n",
    "    \n",
    "    norm_pol = np.max(np.abs(c_opt - c_new))\n",
    "    c_opt[:] = c_new[:]\n",
    "    a_opt[:] = a_new[:]\n",
    "    \n",
    "    it = it + 1\n",
    "    \n",
    "    if it % 100 == 0:\n",
    "        print('Iteration', it, '\\t polfun norm', norm_pol)\n",
    "\n",
    "# t1 = time()\n",
    "\n",
    "print('')\n",
    "# print('Elapsed time is', t1-t0, 'seconds.')\n",
    "print('Converged after', it, 'iterations.')\n",
    "print('Policy function L_inf norm is', norm_pol)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "a_EGM = np.zeros((z_len, a_len))\n",
    "c_EGM = np.zeros((z_len, a_len))\n",
    "\n",
    "for j, z in enumerate(z_grid):\n",
    "    c_EGM[j, :] = c_opt[j, :]\n",
    "    a_EGM[j, :] = R*a_grid + z*w - c_EGM[j, :]\n",
    "\n",
    "# plt.subplots(figsize=(4,3))\n",
    "plt.plot(a_grid, a_EGM[1, :], lw=2, label='employed')\n",
    "plt.plot(a_grid, a_EGM[0, :], lw=2, label='unemployed')\n",
    "plt.plot(a_grid, a_grid, 'k-', lw=1)\n",
    "plt.title('Policy function: assets')\n",
    "plt.xlabel('Current period assets $a$')\n",
    "plt.legend(frameon=False)\n",
    "plt.tight_layout()\n",
    "# plt.savefig('PE_polfun_A.pdf')\n",
    "plt.show()\n",
    "\n",
    "# plt.subplots(figsize=(4,3))\n",
    "plt.plot(a_grid, c_EGM[1, :], lw=2, label='employed')\n",
    "plt.plot(a_grid, c_EGM[0, :], lw=2, label='unemployed')\n",
    "plt.title('Policy function: consumption')\n",
    "plt.xlabel('Current period assets $a$')\n",
    "plt.legend(frameon=False)\n",
    "plt.tight_layout()\n",
    "# plt.savefig('PE_polfun_C.pdf')\n",
    "plt.show()\n",
    "\n",
    "# plt.subplots(figsize=(4,3))\n",
    "plt.plot(a_grid, a_EGM[1, :]-a_grid, lw=2)\n",
    "plt.plot(a_grid, np.zeros(len(a_grid)), 'k-', lw=1)\n",
    "plt.title('Employed: $a_{t+1} - a_{t}$')\n",
    "plt.xlabel('Current period assets $a$')\n",
    "plt.tight_layout()\n",
    "# plt.savefig('PE_polfun_AA.pdf')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dens_mult = 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "a_grid_dens = np.linspace(a_min, a_max, dens_mult*a_len)\n",
    "a_len_dens = len(a_grid_dens)\n",
    "index_dens = range(a_len_dens)\n",
    "\n",
    "pos_l = np.zeros((z_len, a_len_dens), dtype=int)\n",
    "pos_h = np.zeros((z_len, a_len_dens), dtype=int)\n",
    "\n",
    "val_l = np.zeros((z_len, a_len_dens))\n",
    "val_h = np.zeros((z_len, a_len_dens))\n",
    "\n",
    "a_EGM = np.zeros((z_len, a_len_dens))\n",
    "\n",
    "for j, z in enumerate(z_grid):\n",
    "    a_EGM[j, :] = InterpolatedUnivariateSpline(a_grid, a_opt[j, :], k=1)(a_grid_dens)\n",
    "\n",
    "for j, z in enumerate(z_grid):\n",
    "    for i, a in enumerate(a_grid_dens):\n",
    "        pos_h[j, i] = min(bisect(a_grid_dens, a_EGM[j, i]), a_len_dens-1)\n",
    "        pos_l[j, i] = max(pos_h[j, i]-1, 0)\n",
    "\n",
    "        val_h[j, i] = min(1, ((a_EGM[j, i] - a_grid_dens[pos_l[j, i]]) \n",
    "                              / (a_grid_dens[pos_h[j, i]] - a_grid_dens[pos_l[j, i]])) )\n",
    "        val_l[j, i] = 1 - val_h[j, i]\n",
    "\n",
    "G_0 = sparse.coo_matrix( ( np.concatenate((val_l[0, :], val_h[0, :])), \n",
    "                          ( np.concatenate((index_dens, index_dens)), \n",
    "                            np.concatenate((pos_l[0, :], pos_h[0, :])) ) ), \n",
    "                        shape=(a_len_dens, a_len_dens) ).tocsr()\n",
    "\n",
    "G_1 = sparse.coo_matrix( ( np.concatenate((val_l[1, :], val_h[1, :])), \n",
    "                          ( np.concatenate((index_dens, index_dens)), \n",
    "                            np.concatenate((pos_l[1, :], pos_h[1, :])) ) ), \n",
    "                        shape=(a_len_dens, a_len_dens) ).tocsr()\n",
    "\n",
    "G = sparse.vstack( (sparse.kron(P[0, :], G_0), \n",
    "                    sparse.kron(P[1, :], G_1)) )\n",
    "\n",
    "val, vec = eigs(G.transpose(), k=1, which='LR', return_eigenvectors=True)\n",
    "\n",
    "stat_EGM = np.real(vec.flat)\n",
    "stat_EGM = stat_EGM / np.sum(stat_EGM)\n",
    "\n",
    "K_sup = a_grid_dens @ stat_EGM[a_len_dens:] + a_grid_dens @ stat_EGM[:a_len_dens]\n",
    "\n",
    "print('Capital supplied =', K_sup)\n",
    "print('')\n",
    "\n",
    "# plt.subplots(figsize=(4,3))\n",
    "plt.plot(a_grid_dens, stat_EGM[a_len_dens:], lw=2, label='employed')\n",
    "plt.plot(a_grid_dens, stat_EGM[:a_len_dens], lw=2, label='unemployed')\n",
    "plt.title('Wealth distribution')\n",
    "plt.xlabel('Current period assets $a$')\n",
    "# plt.xlim(-5, 105)\n",
    "plt.legend(frameon=False)\n",
    "plt.tight_layout()\n",
    "# plt.savefig('PE_dist.pdf')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def K_sup(r, dens_mult=1):\n",
    "    K_dem = (α*A/(r+δ))**(1/(1-α)) * N\n",
    "    w = (1-α) * (K_dem/N)**α\n",
    "    R = 1+r\n",
    "    \n",
    "#     print('Interest rate pa =', r) #100*((1+r)**8-1)\n",
    "#     print('Capital demanded =', K_dem)\n",
    "    \n",
    "    # Endogenous Gridpoint Method\n",
    "\n",
    "#     t0 = time()\n",
    "\n",
    "    c_new = np.zeros((z_len, a_len))\n",
    "    c_opt = np.zeros((z_len, a_len))\n",
    "\n",
    "    a_new = np.zeros((z_len, a_len))\n",
    "    a_opt = np.zeros((z_len, a_len))\n",
    "    \n",
    "    index = range(a_len)\n",
    "\n",
    "    norm_pol = 1\n",
    "    it = 0\n",
    "    maxit = 5000\n",
    "\n",
    "    # Initial guess for c_opt\n",
    "    for j, z in enumerate(z_grid):\n",
    "        c_opt[j, :] = z*w + r*a_grid[:]\n",
    "\n",
    "    while norm_pol > 1e-6*(1-β) and it < maxit:\n",
    "        \n",
    "        ERHS = P @ (β*R * c_opt**(-σ))\n",
    "        IERHS = ERHS**(-1/σ)\n",
    "        \n",
    "        for j, z in enumerate(z_grid):\n",
    "            c_end = IERHS[j, :]\n",
    "            a_end = (a_grid[:] + c_end - z*w) / R\n",
    "\n",
    "            c_new[j, :] = InterpolatedUnivariateSpline(a_end, c_end, k=1)(a_grid)\n",
    "\n",
    "            # Binding constraint\n",
    "            bind = np.where(a_grid < a_end[0])\n",
    "            c_new[j, bind] = R*a_grid[bind] + z*w - a_min\n",
    "\n",
    "            a_new[j, :] = R*a_grid + z*w - c_new[j, :]\n",
    "\n",
    "        norm_pol = np.max(np.abs(a_opt - a_new))\n",
    "        c_opt[:, :] = c_new[:, :]\n",
    "        a_opt[:, :] = a_new[:, :]\n",
    "\n",
    "        it = it + 1\n",
    "\n",
    "#     t1 = time()\n",
    "\n",
    "#     print('Elapsed time is', t1-t0, 'seconds.')\n",
    "#     print('Converged after', it, 'iterations.')\n",
    "#     print('Policy function L_inf norm is', norm_pol)\n",
    "\n",
    "#     t0 = time()\n",
    "    \n",
    "#     c_EGM = c_opt[:, :]\n",
    "    \n",
    "    a_grid_dens = np.linspace(a_min, a_max, dens_mult*a_len)\n",
    "    a_len_dens = len(a_grid_dens)\n",
    "    index_dens = range(a_len_dens)\n",
    "    \n",
    "    pos_l = np.zeros((z_len, a_len_dens), dtype=int)\n",
    "    pos_h = np.zeros((z_len, a_len_dens), dtype=int)\n",
    "\n",
    "    val_l = np.zeros((z_len, a_len_dens))\n",
    "    val_h = np.zeros((z_len, a_len_dens))\n",
    "    \n",
    "    a_EGM = np.zeros((z_len, a_len_dens))\n",
    "\n",
    "    for j, z in enumerate(z_grid):\n",
    "        a_EGM[j, :] = InterpolatedUnivariateSpline(a_grid, a_opt[j, :], k=1)(a_grid_dens)\n",
    "\n",
    "    for j, z in enumerate(z_grid):\n",
    "        for i, a in enumerate(a_grid_dens):\n",
    "            pos_h[j, i] = min(bisect(a_grid_dens, a_EGM[j, i]), a_len_dens-1)\n",
    "            pos_l[j, i] = max(pos_h[j, i]-1, 0)\n",
    "\n",
    "            val_h[j, i] = min(1, ((a_EGM[j, i] - a_grid_dens[pos_l[j, i]]) \n",
    "                                  / (a_grid_dens[pos_h[j, i]] - a_grid_dens[pos_l[j, i]])) )\n",
    "            val_l[j, i] = 1 - val_h[j, i]\n",
    "\n",
    "    G_0 = sparse.coo_matrix( ( np.concatenate((val_l[0, :], val_h[0, :])), \n",
    "                              ( np.concatenate((index_dens, index_dens)), \n",
    "                                np.concatenate((pos_l[0, :], pos_h[0, :])) ) ), \n",
    "                            shape=(a_len_dens, a_len_dens) ).tocsr()\n",
    "\n",
    "    G_1 = sparse.coo_matrix( ( np.concatenate((val_l[1, :], val_h[1, :])), \n",
    "                              ( np.concatenate((index_dens, index_dens)), \n",
    "                                np.concatenate((pos_l[1, :], pos_h[1, :])) ) ), \n",
    "                            shape=(a_len_dens, a_len_dens) ).tocsr()\n",
    "\n",
    "    G = sparse.vstack( (sparse.kron(P[0, :], G_0), \n",
    "                        sparse.kron(P[1, :], G_1)) )\n",
    "\n",
    "    val, vec = eigs(G.transpose(), k=1, which='LR', return_eigenvectors=True)\n",
    "\n",
    "#     print(val)\n",
    "\n",
    "    stat_EGM = np.real(vec.flat)\n",
    "    stat_EGM = stat_EGM / np.sum(stat_EGM)\n",
    "    \n",
    "    K_sup = a_grid_dens @ stat_EGM[a_len_dens:] + a_grid_dens @ stat_EGM[:a_len_dens]\n",
    "    return K_sup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rr = np.linspace(0, 0.04, 100) #1/β-1\n",
    "\n",
    "KK = []\n",
    "\n",
    "for i, r in enumerate(rr):\n",
    "    KK.append(K_sup(r))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# rr = np.linspace(0, 0.1, 100) #1/β-1\n",
    "\n",
    "plt.plot(Capital_Demand(rr), rr*100, lw=2, label='Demand for capital')\n",
    "plt.plot(KK, rr*100, lw=2, label='Supply of capital')\n",
    "plt.xlabel('Aggregate capital stock $K$')\n",
    "plt.ylabel('Real interest rate $r$ (%)')\n",
    "plt.legend(frameon=False, loc='upper center')\n",
    "plt.xlim(None, 10.5)\n",
    "plt.tight_layout()\n",
    "# plt.savefig('Capital_eq.pdf')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def Capital_Market(r, dens_mult=1, distribution=False, xlim=None):\n",
    "    K_dem = (α*A/(r+δ))**(1/(1-α)) * N\n",
    "    w = (1-α) * (K_dem/N)**α\n",
    "    R = 1+r\n",
    "    \n",
    "    print('Interest rate pa =', r) #100*((1+r)**8-1)\n",
    "    print('Capital demanded =', K_dem)\n",
    "    \n",
    "    # Endogenous Gridpoint Method\n",
    "\n",
    "#     t0 = time()\n",
    "\n",
    "    c_new = np.zeros((z_len, a_len))\n",
    "    c_opt = np.zeros((z_len, a_len))\n",
    "\n",
    "    a_new = np.zeros((z_len, a_len))\n",
    "    a_opt = np.zeros((z_len, a_len))\n",
    "    \n",
    "    index = range(a_len)\n",
    "\n",
    "    norm_pol = 1\n",
    "    it = 0\n",
    "    maxit = 5000\n",
    "\n",
    "    # Initial guess for c_opt\n",
    "    for j, z in enumerate(z_grid):\n",
    "        c_opt[j, :] = z*w + r*a_grid[:]\n",
    "\n",
    "    while norm_pol > 1e-6*(1-β) and it < maxit:\n",
    "        \n",
    "        ERHS = P @ (β*R * c_opt**(-σ))\n",
    "        IERHS = ERHS**(-1/σ)\n",
    "        \n",
    "        for j, z in enumerate(z_grid):\n",
    "            c_end = IERHS[j, :]\n",
    "            a_end = (a_grid[:] + c_end - z*w) / R\n",
    "\n",
    "            c_new[j, :] = InterpolatedUnivariateSpline(a_end, c_end, k=1)(a_grid)\n",
    "\n",
    "            # Binding constraint\n",
    "            bind = np.where(a_grid < a_end[0])\n",
    "            c_new[j, bind] = R*a_grid[bind] + z*w - a_min\n",
    "\n",
    "            a_new[j, :] = R*a_grid + z*w - c_new[j, :]\n",
    "\n",
    "        norm_pol = np.max(np.abs(a_opt - a_new))\n",
    "        c_opt[:, :] = c_new[:, :]\n",
    "        a_opt[:, :] = a_new[:, :]\n",
    "\n",
    "        it = it + 1\n",
    "\n",
    "#     t1 = time()\n",
    "\n",
    "#     print('Elapsed time is', t1-t0, 'seconds.')\n",
    "#     print('Converged after', it, 'iterations.')\n",
    "#     print('Policy function L_inf norm is', norm_pol)\n",
    "\n",
    "#     t0 = time()\n",
    "    \n",
    "#     c_EGM = c_opt[:, :]\n",
    "    \n",
    "    a_grid_dens = np.linspace(a_min, a_max, dens_mult*a_len)\n",
    "    a_len_dens = len(a_grid_dens)\n",
    "    index_dens = range(a_len_dens)\n",
    "    \n",
    "    pos_l = np.zeros((z_len, a_len_dens), dtype=int)\n",
    "    pos_h = np.zeros((z_len, a_len_dens), dtype=int)\n",
    "\n",
    "    val_l = np.zeros((z_len, a_len_dens))\n",
    "    val_h = np.zeros((z_len, a_len_dens))\n",
    "    \n",
    "    a_EGM = np.zeros((z_len, a_len_dens))\n",
    "\n",
    "    for j, z in enumerate(z_grid):\n",
    "        a_EGM[j, :] = InterpolatedUnivariateSpline(a_grid, a_opt[j, :], k=1)(a_grid_dens)\n",
    "\n",
    "    for j, z in enumerate(z_grid):\n",
    "        for i, a in enumerate(a_grid_dens):\n",
    "            pos_h[j, i] = min(bisect(a_grid_dens, a_EGM[j, i]), a_len_dens-1)\n",
    "            pos_l[j, i] = max(pos_h[j, i]-1, 0)\n",
    "\n",
    "            val_h[j, i] = min(1, ((a_EGM[j, i] - a_grid_dens[pos_l[j, i]]) \n",
    "                                  / (a_grid_dens[pos_h[j, i]] - a_grid_dens[pos_l[j, i]])) )\n",
    "            val_l[j, i] = 1 - val_h[j, i]\n",
    "\n",
    "    G_0 = sparse.coo_matrix( ( np.concatenate((val_l[0, :], val_h[0, :])), \n",
    "                              ( np.concatenate((index_dens, index_dens)), \n",
    "                                np.concatenate((pos_l[0, :], pos_h[0, :])) ) ), \n",
    "                            shape=(a_len_dens, a_len_dens) ).tocsr()\n",
    "\n",
    "    G_1 = sparse.coo_matrix( ( np.concatenate((val_l[1, :], val_h[1, :])), \n",
    "                              ( np.concatenate((index_dens, index_dens)), \n",
    "                                np.concatenate((pos_l[1, :], pos_h[1, :])) ) ), \n",
    "                            shape=(a_len_dens, a_len_dens) ).tocsr()\n",
    "\n",
    "    G = sparse.vstack( (sparse.kron(P[0, :], G_0), \n",
    "                        sparse.kron(P[1, :], G_1)) )\n",
    "\n",
    "    val, vec = eigs(G.transpose(), k=1, which='LR', return_eigenvectors=True)\n",
    "\n",
    "#     print(val)\n",
    "\n",
    "    stat_EGM = np.real(vec.flat)\n",
    "    stat_EGM = stat_EGM / np.sum(stat_EGM)\n",
    "    \n",
    "    K_sup = a_grid_dens @ stat_EGM[a_len_dens:] + a_grid_dens @ stat_EGM[:a_len_dens]\n",
    "\n",
    "#     t1 = time()\n",
    "\n",
    "    print('Capital supplied =', K_sup)\n",
    "    print('')\n",
    "#     print('Elapsed time is', t1-t0, 'seconds.')\n",
    "\n",
    "    if distribution == True:\n",
    "        plt.plot(a_grid_dens, stat_EGM[a_len_dens:], lw=2, label='employed')\n",
    "        plt.plot(a_grid_dens, stat_EGM[:a_len_dens], lw=2, label='unemployed')\n",
    "        plt.title('Wealth distribution')\n",
    "        plt.xlabel('Current period assets $a$')\n",
    "        plt.xlim(xlim)\n",
    "        plt.legend(frameon=False)\n",
    "        plt.tight_layout()\n",
    "#         plt.savefig('GE_dist.pdf')\n",
    "        plt.show()\n",
    "    \n",
    "    return (K_sup - K_dem)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Capital_Market(0.02, True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# t0 = time()\n",
    "result = optimize.brentq(Capital_Market, 0, 1/β-1, full_output=True)\n",
    "# t1 = time()\n",
    "# print('Elapsed time is', t1-t0, 'seconds.')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "r_star = result[0]\n",
    "\n",
    "Capital_Market(result[0], 1, True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "stat_EGM_both = stat_EGM[:a_len_dens] + stat_EGM[a_len_dens:]\n",
    "\n",
    "stat_EGM_both[np.abs(stat_EGM_both) < 1e-15] = 0\n",
    "stat_EGM_both = stat_EGM_both/sum(stat_EGM_both)\n",
    "\n",
    "plt.plot(a_grid_dens, stat_EGM_both)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "a_dist_both = a_grid_dens * stat_EGM_both\n",
    "a_dist_both = a_dist_both / sum(a_dist_both)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# fig, ax = plt.subplots(figsize=(6, 6))\n",
    "\n",
    "plt.plot(np.linspace(0,1), np.linspace(0,1), lw=2, label='Perfect wealth equality')\n",
    "plt.plot(np.insert(np.cumsum(stat_EGM_both), 0, 0), \n",
    "         np.insert(np.cumsum(a_dist_both), 0, 0), \n",
    "         lw=2, label='Lorenz curve')\n",
    "\n",
    "plt.xlabel('Cumulative population')\n",
    "plt.ylabel('Cumulative wealth')\n",
    "plt.legend()\n",
    "\n",
    "# plt.plot(P_bar[0], P_bar[0]*z_grid[0]/t_inc, 'o')\n",
    "# plt.tight_layout()\n",
    "# plt.savefig('GE_Lorenz.pdf')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}