{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# The Diamond OLG Model\n",
    "\n",
    "[![badge](https://img.shields.io/badge/Launch%20using%20-Econ--ARK-blue)](https://econ-ark.org/materials/diamondolg#launch)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Convergence of OLG Economy to Steady State"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Some initial setup\n",
    "\n",
    "from matplotlib import pyplot as plt\n",
    "import numpy as np\n",
    "plt.style.use('seaborn-darkgrid')\n",
    "palette = plt.get_cmap('Dark2')\n",
    "from copy  import deepcopy\n",
    "\n",
    "from ipywidgets import interact, interactive, fixed, interact_manual, Layout\n",
    "import ipywidgets as widgets\n",
    "\n",
    "from HARK.ConsumptionSaving.ConsIndShockModel import (\n",
    "    PerfForesightConsumerType, init_perfect_foresight)\n",
    "\n",
    "years_per_gen = 30"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Define a function that plots something given some inputs\n",
    "def plot1(Epsilon, DiscFac, PopGrowth, YearsPerGeneration, Initialk):\n",
    "    '''Inputs:\n",
    "        Epsilon: Elasticity of output with respect to capital/labour ratio\n",
    "        DiscFac: One period discount factor\n",
    "        YearPerGeneration: No. of years per generation\n",
    "        PopGrowth: Gross growth rate of population in one period'''\n",
    "    \n",
    "    #\n",
    "    kMax = 0.1\n",
    "    # Define some parameters\n",
    "    Beta = DiscFac**YearsPerGeneration\n",
    "    xi = PopGrowth**YearsPerGeneration\n",
    "    Q = (1-Epsilon)*(Beta/(1+Beta))/xi\n",
    "    kBar = Q**(1/(1-Epsilon))\n",
    "    \n",
    "    # Create an agent that will solve the consumption problem\n",
    "    PFagent = PerfForesightConsumerType(**init_perfect_foresight)\n",
    "    PFagent.cycles = 1 # let the agent live the cycle of periods just once\n",
    "    PFagent.T_cycle = 2 # Number of periods in the cycle\n",
    "    PFagent.assign_parameters(PermGroFac = [0.000001]) # Income only in the first period\n",
    "    PFagent.LivPrb = [1.]\n",
    "    \n",
    "    PFagent.DiscFac = Beta\n",
    "    \n",
    "    # Hark seems to have trouble with log-utility\n",
    "    # so pass rho = 1 + something small.\n",
    "    PFagent.CRRA = 1.001\n",
    "    \n",
    "    PFagent.solve()\n",
    "    \n",
    "    # Plot the OLG capital accumulation curve and 45 deg line\n",
    "    plt.figure(figsize=(9,6))\n",
    "    kt_range = np.linspace(0, kMax, 300)\n",
    "    \n",
    "    # Analitical solution plot\n",
    "    ktp1 = Q*kt_range**Epsilon\n",
    "    plt.plot(kt_range, ktp1, 'b-', label='Capital accumulation curve')\n",
    "    plt.plot(kt_range, kt_range, 'k-', label='45 Degree line')\n",
    "    \n",
    "    # Plot the path\n",
    "    kt_ar = Initialk\n",
    "    ktp1_ar = 0.\n",
    "    for i in range(3):\n",
    "        \n",
    "        # Compute Next Period's capital using HARK\n",
    "        wage = (1-Epsilon)*kt_ar**Epsilon\n",
    "        c = PFagent.solution[0].cFunc(wage)\n",
    "        a = wage - c\n",
    "        k1 = a/xi\n",
    "        \n",
    "        plt.arrow(kt_ar, ktp1_ar, 0., k1-ktp1_ar,\n",
    "                  length_includes_head=True,\n",
    "                  lw=0.01,\n",
    "                  width=0.0005,\n",
    "                  color='black',\n",
    "                  edgecolor=None)\n",
    "        plt.arrow(kt_ar, k1, k1-kt_ar , 0.,\n",
    "                  length_includes_head=True,\n",
    "                  lw=0.01,\n",
    "                  width=0.0005,\n",
    "                  color='black',\n",
    "                  edgecolor=None)\n",
    "        \n",
    "        # Update arrow\n",
    "        kt_ar = k1\n",
    "        ktp1_ar = kt_ar\n",
    "    \n",
    "    # Plot kbar and initial k\n",
    "    plt.plot(kBar, kBar, 'ro', label=r'$\\bar{k}$')\n",
    "    plt.plot(Initialk, 0.0005, 'co', label = '$k_0$')\n",
    "    \n",
    "    plt.legend()\n",
    "    plt.xlim(0 ,kMax)\n",
    "    plt.ylim(0, kMax)\n",
    "    plt.xlabel('$k_t$')\n",
    "    plt.ylabel('$k_{t+1}$')\n",
    "    plt.show()\n",
    "\n",
    "    return None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Define some widgets to control the plot\n",
    "\n",
    "# Define a slider for Epsilon\n",
    "Epsilon_widget1 = widgets.FloatSlider(\n",
    "    min=0.2,\n",
    "    max=0.4,\n",
    "    step=0.01,\n",
    "    value=0.33,\n",
    "    continuous_update=False,\n",
    "    readout_format='.3f',\n",
    "    description=r'Capital Share $\\epsilon$',\n",
    "    style = {'description_width': 'initial'})\n",
    "\n",
    "# Define a slider for the discount factor\n",
    "DiscFac_widget1 = widgets.FloatSlider(\n",
    "    min=.94,\n",
    "    max=0.99,\n",
    "    step=0.001,\n",
    "    value=0.96,\n",
    "    continuous_update=False,\n",
    "    readout_format='.3f',\n",
    "    description=r'Discount Factor $\\beta$',\n",
    "    style = {'description_width': 'initial'})\n",
    "\n",
    "# Define a slider for pop. growth\n",
    "PopGrowth_widget1 = widgets.FloatSlider(\n",
    "    min=0.98,\n",
    "    max=1.05,\n",
    "    step=0.001,\n",
    "    value=1.01,\n",
    "    continuous_update=False,\n",
    "    readout_format='.3f',\n",
    "    description=r'Pop. Growth $\\Xi$',\n",
    "    style = {'description_width': 'initial'})\n",
    "\n",
    "# Define a slider for initial k\n",
    "Initialk_widget1 = widgets.FloatSlider(\n",
    "    min=0.01,\n",
    "    max=0.1,\n",
    "    step=0.01,\n",
    "    value=.05,\n",
    "    continuous_update=True,\n",
    "    readout_format='.3f',\n",
    "    description='Init. capital ratio',\n",
    "    style = {'description_width': 'initial'})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "fb7511e3f4704d048307f6203cd4f519",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "interactive(children=(FloatSlider(value=0.33, continuous_update=False, description='Capital Share $\\\\epsilon$'…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Make the widget\n",
    "interact(plot1,\n",
    "         Epsilon = Epsilon_widget1,\n",
    "         DiscFac = DiscFac_widget1,\n",
    "         PopGrowth = PopGrowth_widget1,\n",
    "         YearsPerGeneration = fixed(years_per_gen),\n",
    "         Initialk = Initialk_widget1);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Gross and Net Per Capita Output as a Function of k"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define a function that plots something given some inputs\n",
    "def plot2(Epsilon, PopGrowth, YearsPerGeneration):\n",
    "    '''Inputs:\n",
    "        Epsilon: Elasticity of output with respect to capital/labour ratio\n",
    "        DiscFac: One period discount factor\n",
    "        YearPerGeneration: No. of years per generation\n",
    "        PopGrowth: Gross growth rate of population in one period'''\n",
    "    \n",
    "    kMax = 5.5\n",
    "    \n",
    "    # Define some parameters\n",
    "    xi = PopGrowth**YearsPerGeneration\n",
    "    Xi = xi - 1\n",
    "    kBarForcZero = Xi**(1/(Epsilon-1))\n",
    "    \n",
    "    # Plot the production function and depreciation/dilution curves\n",
    "    kt_range = np.linspace(0, kMax, 500)\n",
    "    plt.figure(figsize=(18, 6))\n",
    "    plt.subplot(1,2,1)\n",
    "    plt.plot(kt_range, kt_range**Epsilon, 'b-', label = '$f(k)$')\n",
    "    plt.plot(kt_range, Xi*kt_range, 'k-', label = '$Xi * k$')\n",
    "    plt.legend()\n",
    "    plt.xlim(0, kMax)\n",
    "    plt.ylim(0, kMax)\n",
    "    plt.xlabel('$k_t$')\n",
    "    \n",
    "    plt.subplot(1,2,2)\n",
    "    plt.plot(kt_range, kt_range**Epsilon - Xi*kt_range, 'k-', label ='$f(k) - Xi * k$')\n",
    "    plt.legend()\n",
    "    plt.xlim(0, kMax)\n",
    "    plt.ylim(0, kMax*0.2)\n",
    "    plt.xlabel('$k_t$')\n",
    "    \n",
    "    plt.show()\n",
    "\n",
    "    return None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define some widgets to control the plot\n",
    "\n",
    "# Define a slider for Epsilon\n",
    "Epsilon_widget2 = widgets.FloatSlider(\n",
    "    min=0.2,\n",
    "    max=0.4,\n",
    "    step=0.01,\n",
    "    value=0.33,\n",
    "    continuous_update=False,\n",
    "    readout_format='.3f',\n",
    "    description=r'Capital Share $\\epsilon$',\n",
    "    style = {'description_width': 'initial'})\n",
    "\n",
    "# Define a slider for pop. growth\n",
    "PopGrowth_widget2 = widgets.FloatSlider(\n",
    "    min=0.98,\n",
    "    max=1.05,\n",
    "    step=0.001,\n",
    "    value=1.01,\n",
    "    continuous_update=False,\n",
    "    readout_format='.3f',\n",
    "    description=r'Pop. Growth $\\Xi$',\n",
    "    style = {'description_width': 'initial'})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "27a6b210d37d400a9b76dfef3cb570b9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "interactive(children=(FloatSlider(value=0.33, continuous_update=False, description='Capital Share $\\\\epsilon$'…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Make the widget\n",
    "interact(plot2,\n",
    "         Epsilon = Epsilon_widget2,\n",
    "         PopGrowth = PopGrowth_widget2,\n",
    "         YearsPerGeneration = fixed(years_per_gen)\n",
    "        );"
   ]
  }
 ],
 "metadata": {
  "jupytext": {
   "formats": "ipynb,py:percent",
   "notebook_metadata_filter": "all"
  },
  "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.8.8"
  },
  "latex_envs": {
   "LaTeX_envs_menu_present": true,
   "autoclose": false,
   "autocomplete": true,
   "bibliofile": "biblio.bib",
   "cite_by": "apalike",
   "current_citInitial": 1,
   "eqLabelWithNumbers": true,
   "eqNumInitial": 1,
   "hotkeys": {
    "equation": "Ctrl-E",
    "itemize": "Ctrl-I"
   },
   "labels_anchors": false,
   "latex_user_defs": false,
   "report_style_numbering": false,
   "user_envs_cfg": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}