egm_policy_iter.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "<a id='egm-policy-iter'></a>\n",
    "<div id=\"qe-notebook-header\" align=\"right\" style=\"text-align:right;\">\n",
    "        <a href=\"https://quantecon.org/\" title=\"quantecon.org\">\n",
    "                <img style=\"width:250px;display:inline;\" width=\"250px\" src=\"https://assets.quantecon.org/img/qe-menubar-logo.svg\" alt=\"QuantEcon\">\n",
    "        </a>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Optimal Growth III: The Endogenous Grid Method"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Contents\n",
    "\n",
    "- [Optimal Growth III: The Endogenous Grid Method](#Optimal-Growth-III:-The-Endogenous-Grid-Method)  \n",
    "  - [Overview](#Overview)  \n",
    "  - [Key Idea](#Key-Idea)  \n",
    "  - [Implementation](#Implementation)  \n",
    "  - [Speed](#Speed)  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In addition to what&rsquo;s in Anaconda, this lecture will need the following libraries:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": true
   },
   "outputs": [],
   "source": [
    "!pip install --upgrade quantecon\n",
    "!pip install interpolation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Overview\n",
    "\n",
    "We solved the stochastic optimal growth model using\n",
    "\n",
    "1. [value function iteration](https://lectures.quantecon.org/optgrowth.html)  \n",
    "1. [Euler equation based time iteration](https://lectures.quantecon.org/coleman_policy_iter.html)  \n",
    "\n",
    "\n",
    "We found time iteration to be significantly more accurate at each step.\n",
    "\n",
    "In this lecture, we&rsquo;ll look at an ingenious twist on the time iteration technique called the **endogenous grid method** (EGM).\n",
    "\n",
    "EGM is a numerical method for implementing policy iteration invented by [Chris Carroll](http://www.econ2.jhu.edu/people/ccarroll/).\n",
    "\n",
    "It is a good example of how a clever algorithm can save a massive amount of computer time.\n",
    "\n",
    "(Massive when we multiply saved CPU cycles on each implementation times the number of implementations worldwide)\n",
    "\n",
    "The original reference is [[Car06]](https://lectures.quantecon.org/zreferences.html#carroll2006).\n",
    "\n",
    "Let&rsquo;s start with some standard imports:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import quantecon as qe\n",
    "from interpolation import interp\n",
    "from numba import njit, prange\n",
    "from quantecon.optimize import brentq\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Key Idea\n",
    "\n",
    "Let&rsquo;s start by reminding ourselves of the theory and then see how the numerics fit in."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Theory\n",
    "\n",
    "Take the model set out in [the time iteration lecture](https://lectures.quantecon.org/coleman_policy_iter.html), following the same terminology and notation.\n",
    "\n",
    "The Euler equation is\n",
    "\n",
    "\n",
    "<a id='equation-egm-euler'></a>\n",
    "$$\n",
    "(u'\\circ \\sigma^*)(y)\n",
    "= \\beta \\int (u'\\circ \\sigma^*)(f(y - \\sigma^*(y)) z) f'(y - \\sigma^*(y)) z \\phi(dz) \\tag{1}\n",
    "$$\n",
    "\n",
    "As we saw, the Coleman-Reffett operator is a nonlinear operator $ K $ engineered so that $ \\sigma^* $ is a fixed point of $ K $.\n",
    "\n",
    "It takes as its argument a continuous strictly increasing consumption policy $ \\sigma \\in \\Sigma $.\n",
    "\n",
    "It returns a new function $ K \\sigma $,  where $ (K \\sigma)(y) $ is the $ c \\in (0, \\infty) $ that solves\n",
    "\n",
    "\n",
    "<a id='equation-egm-coledef'></a>\n",
    "$$\n",
    "u'(c)\n",
    "= \\beta \\int (u' \\circ \\sigma) (f(y - c) z ) f'(y - c) z \\phi(dz) \\tag{2}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exogenous Grid\n",
    "\n",
    "As discussed in [the lecture on time iteration](https://lectures.quantecon.org/coleman_policy_iter.html), to implement the method on a computer we need a numerical approximation.\n",
    "\n",
    "In particular, we represent a policy function by a set of values on a finite grid.\n",
    "\n",
    "The function itself is reconstructed from this representation when necessary, using interpolation or some other method.\n",
    "\n",
    "[Previously](https://lectures.quantecon.org/coleman_policy_iter.html), to obtain a finite representation of an updated consumption policy we\n",
    "\n",
    "- fixed a grid of income points $ \\{y_i\\} $  \n",
    "- calculated the consumption value $ c_i $ corresponding to each\n",
    "  $ y_i $ using [(2)](#equation-egm-coledef) and a root-finding routine  \n",
    "\n",
    "\n",
    "Each $ c_i $ is then interpreted as the value of the function $ K \\sigma $ at $ y_i $.\n",
    "\n",
    "Thus, with the points $ \\{y_i, c_i\\} $ in hand, we can reconstruct $ K \\sigma $ via approximation.\n",
    "\n",
    "Iteration then continues&hellip;"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Endogenous Grid\n",
    "\n",
    "The method discussed above requires a root-finding routine to find the\n",
    "$ c_i $ corresponding to a given income value $ y_i $.\n",
    "\n",
    "Root-finding is costly because it typically involves a significant number of\n",
    "function evaluations.\n",
    "\n",
    "As pointed out by Carroll [[Car06]](https://lectures.quantecon.org/zreferences.html#carroll2006), we can avoid this if\n",
    "$ y_i $ is chosen endogenously.\n",
    "\n",
    "The only assumption required is that $ u' $ is invertible on $ (0, \\infty) $.\n",
    "\n",
    "The idea is this:\n",
    "\n",
    "First, we fix an *exogenous* grid $ \\{k_i\\} $ for capital ($ k = y - c $).\n",
    "\n",
    "Then we obtain  $ c_i $ via\n",
    "\n",
    "\n",
    "<a id='equation-egm-getc'></a>\n",
    "$$\n",
    "c_i =\n",
    "(u')^{-1}\n",
    "\\left\\{\n",
    "    \\beta \\int (u' \\circ \\sigma) (f(k_i) z ) \\, f'(k_i) \\, z \\, \\phi(dz)\n",
    "\\right\\} \\tag{3}\n",
    "$$\n",
    "\n",
    "where $ (u')^{-1} $ is the inverse function of $ u' $.\n",
    "\n",
    "Finally, for each $ c_i $ we set $ y_i = c_i + k_i $.\n",
    "\n",
    "It is clear that each $ (y_i, c_i) $ pair constructed in this manner satisfies [(2)](#equation-egm-coledef).\n",
    "\n",
    "With the points $ \\{y_i, c_i\\} $ in hand, we can reconstruct $ K \\sigma $ via approximation as before.\n",
    "\n",
    "The name EGM comes from the fact that the grid $ \\{y_i\\} $ is  determined **endogenously**."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Implementation\n",
    "\n",
    "Let&rsquo;s implement this version of the Coleman-Reffett operator and see how it performs.\n",
    "\n",
    "First, we will construct a class `OptimalGrowthModel` to hold the parameters of the\n",
    "model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "class OptimalGrowthModel:\n",
    "\n",
    "    \"\"\"\n",
    "\n",
    "    The class holds parameters and true value and policy functions.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self,\n",
    "                 f,                # Production function\n",
    "                 f_prime,          # f'(k)\n",
    "                 u,                # Utility function\n",
    "                 u_prime,          # Marginal utility\n",
    "                 u_prime_inv,      # Inverse marginal utility\n",
    "                 &beta;=0.96,           # Discount factor\n",
    "                 &mu;=0,\n",
    "                 s=0.1,\n",
    "                 grid_max=4,\n",
    "                 grid_size=200,\n",
    "                 shock_size=250):\n",
    "\n",
    "        self.&beta;, self.&mu;, self.s = &beta;, &mu;, s\n",
    "        self.f, self.u = f, u\n",
    "        self.f_prime, self.u_prime, self.u_prime_inv = f_prime, u_prime, \\\n",
    "            u_prime_inv\n",
    "\n",
    "        # Set up grid\n",
    "        self.grid = np.linspace(1e-5, grid_max, grid_size)\n",
    "        # Store shocks\n",
    "        self.shocks = np.exp(&mu; + s * np.random.randn(shock_size))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### The Operator\n",
    "\n",
    "Here&rsquo;s an implementation of $ K $ using EGM as described above.\n",
    "\n",
    "Unlike the [previous lecture](https://lectures.quantecon.org/coleman_policy_iter.html), we do not just-in-time\n",
    "compile the operator because we want to return the policy function.\n",
    "\n",
    "Despite this, the EGM method is still faster than the standard Coleman-Reffett operator,\n",
    "as we will see later on."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def egm_operator_factory(og):\n",
    "    \"\"\"\n",
    "    A function factory for building the Coleman-Reffett operator\n",
    "\n",
    "    Here og is an instance of OptimalGrowthModel.\n",
    "    \"\"\"\n",
    "\n",
    "    f, u, &beta; = og.f, og.u, og.&beta;\n",
    "    f_prime, u_prime, u_prime_inv = og.f_prime, og.u_prime, og.u_prime_inv\n",
    "    grid, shocks = og.grid, og.shocks\n",
    "\n",
    "    def K(&sigma;):\n",
    "        \"\"\"\n",
    "        The Bellman operator\n",
    "\n",
    "        * &sigma; is a function\n",
    "        \"\"\"\n",
    "        # Allocate memory for value of consumption on endogenous grid points\n",
    "        c = np.empty_like(grid)\n",
    "\n",
    "        # Solve for updated consumption value\n",
    "        for i, k in enumerate(grid):\n",
    "            vals = u_prime(&sigma;(f(k) * shocks)) * f_prime(k) * shocks\n",
    "            c[i] = u_prime_inv(&beta; * np.mean(vals))\n",
    "\n",
    "        # Determine endogenous grid\n",
    "        y = grid + c  # y_i = k_i + c_i\n",
    "\n",
    "        # Update policy function and return\n",
    "        &sigma;_new = lambda x: interp(y, c, x)\n",
    "\n",
    "        return &sigma;_new\n",
    "\n",
    "    return K"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note the lack of any root-finding algorithm.\n",
    "\n",
    "We&rsquo;ll also run our original implementation, which uses an exogenous grid and requires root-finding, so we can perform some comparisons."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from interpolation import interp\n",
    "from numba import njit, prange\n",
    "from quantecon.optimize import brentq\n",
    "\n",
    "\n",
    "def time_operator_factory(og, parallel_flag=True):\n",
    "    \"\"\"\n",
    "    A function factory for building the Coleman-Reffett operator.\n",
    "     Here og is an instance of OptimalGrowthModel.\n",
    "    \"\"\"\n",
    "    &beta; = og.&beta;\n",
    "    f, u = og.f, og.u\n",
    "    f_prime, u_prime = og.f_prime, og.u_prime\n",
    "    grid, shocks = og.grid, og.shocks\n",
    "\n",
    "    @njit\n",
    "    def objective(c, &sigma;, y):\n",
    "        \"\"\"\n",
    "        The right hand side of the operator\n",
    "        \"\"\"\n",
    "        # First turn w into a function via interpolation\n",
    "        &sigma;_func = lambda x: interp(grid, &sigma;, x)\n",
    "        vals = u_prime(&sigma;_func(f(y - c) * shocks)) * f_prime(y - c) * shocks\n",
    "        return u_prime(c) - &beta; * np.mean(vals)\n",
    "\n",
    "    @njit(parallel=parallel_flag)\n",
    "    def K(&sigma;):\n",
    "        \"\"\"\n",
    "        The Coleman-Reffett operator\n",
    "        \"\"\"\n",
    "        &sigma;_new = np.empty_like(&sigma;)\n",
    "        for i in prange(len(grid)):\n",
    "            y = grid[i]\n",
    "            # Solve for optimal c at y\n",
    "            c_star = brentq(objective, 1e-10, y-1e-10, args=(&sigma;, y))[0]\n",
    "            &sigma;_new[i] = c_star\n",
    "\n",
    "        return &sigma;_new\n",
    "\n",
    "    return K"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let&rsquo;s test out the code above on some example parameterizations."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Testing on the Log / Cobb&ndash;Douglas Case\n",
    "\n",
    "As we [did for value function iteration](https://lectures.quantecon.org/optgrowth.html) and [time iteration](https://lectures.quantecon.org/coleman_policy_iter.html),\n",
    "let&rsquo;s start by testing our method with the log-linear benchmark.\n",
    "\n",
    "First, we generate an instance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "&alpha; = 0.4  # Production function parameter\n",
    "\n",
    "@njit\n",
    "def f(k):\n",
    "    \"\"\"\n",
    "    Cobb-Douglas production function\n",
    "    \"\"\"\n",
    "    return k**&alpha;\n",
    "\n",
    "@njit\n",
    "def f_prime(k):\n",
    "    \"\"\"\n",
    "    First derivative of the production function\n",
    "    \"\"\"\n",
    "    return &alpha; * k**(&alpha; - 1)\n",
    "\n",
    "@njit\n",
    "def u_prime(c):\n",
    "    return 1 / c\n",
    "\n",
    "og = OptimalGrowthModel(f=f,\n",
    "                        f_prime=f_prime,\n",
    "                        u=np.log,\n",
    "                        u_prime=u_prime,\n",
    "                        u_prime_inv=u_prime)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notice that we&rsquo;re passing `u_prime` twice.\n",
    "\n",
    "The reason is that, in the case of log utility, $ u'(c) = (u')^{-1}(c) = 1/c $.\n",
    "\n",
    "Hence `u_prime` and `u_prime_inv` are the same.\n",
    "\n",
    "As a preliminary test, let&rsquo;s see if $ K \\sigma^* = \\sigma^* $, as implied by the theory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "&beta;, grid = og.&beta;, og.grid\n",
    "\n",
    "def c_star(y):\n",
    "    \"True optimal policy\"\n",
    "    return (1 - &alpha; * &beta;) * y\n",
    "\n",
    "K = egm_operator_factory(og)  # Return the operator K with endogenous grid\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(9, 6))\n",
    "\n",
    "ax.plot(grid, c_star(grid), label=\"optimal policy $\\sigma^*$\")\n",
    "ax.plot(grid, K(c_star)(grid), label=\"$K\\sigma^*$\")\n",
    "\n",
    "ax.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can&rsquo;t really distinguish the two plots.\n",
    "\n",
    "In fact it&rsquo;s easy to see that the difference is essentially zero:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "max(abs(K(c_star)(grid) - c_star(grid)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, let&rsquo;s try iterating from an arbitrary initial condition and see if we\n",
    "converge towards $ \\sigma^* $.\n",
    "\n",
    "Let&rsquo;s start from the consumption policy that eats the whole pie: $ \\sigma(y) = y $"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "&sigma; = lambda x: x\n",
    "n = 15\n",
    "fig, ax = plt.subplots(figsize=(9, 6))\n",
    "\n",
    "ax.plot(grid, &sigma;(grid), color=plt.cm.jet(0),\n",
    "        alpha=0.6, label='initial condition $\\sigma(y) = y$')\n",
    "\n",
    "for i in range(n):\n",
    "    &sigma; = K(&sigma;)  # Update policy\n",
    "    ax.plot(grid, &sigma;(grid), color=plt.cm.jet(i / n),  alpha=0.6)\n",
    "\n",
    "ax.plot(grid, c_star(grid), 'k-',\n",
    "        alpha=0.8, label='true policy function $\\sigma^*$')\n",
    "\n",
    "ax.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We see that the policy has converged nicely, in only a few steps."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Speed\n",
    "\n",
    "Now let&rsquo;s compare the clock times per iteration for the standard Coleman-Reffett\n",
    "operator (with exogenous grid) and the EGM version.\n",
    "\n",
    "We&rsquo;ll do so using the CRRA model adopted in the exercises of the [Euler equation time iteration lecture](https://lectures.quantecon.org/coleman_policy_iter.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "&gamma; = 1.5   # Preference parameter\n",
    "\n",
    "@njit\n",
    "def u(c):\n",
    "    return (c**(1 - &gamma;) - 1) / (1 - &gamma;)\n",
    "\n",
    "@njit\n",
    "def u_prime(c):\n",
    "    return c**(-&gamma;)\n",
    "\n",
    "@njit\n",
    "def u_prime_inv(c):\n",
    "    return c**(-1 / &gamma;)\n",
    "\n",
    "og = OptimalGrowthModel(f=f,\n",
    "                        f_prime=f_prime,\n",
    "                        u=u,\n",
    "                        u_prime=u_prime,\n",
    "                        u_prime_inv=u_prime_inv)\n",
    "\n",
    "# Standard Coleman-Reffett operator\n",
    "K_time = time_operator_factory(og)\n",
    "# Call once to compile jitted version\n",
    "K_time(grid)\n",
    "# Coleman-Reffett operator with endogenous grid\n",
    "K_egm = egm_operator_factory(og)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here&rsquo;s the result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "sim_length = 20\n",
    "\n",
    "print(\"Timing standard Coleman policy function iteration\")\n",
    "&sigma; = grid    # Initial policy\n",
    "qe.util.tic()\n",
    "for i in range(sim_length):\n",
    "    &sigma;_new = K_time(&sigma;)\n",
    "    &sigma; = &sigma;_new\n",
    "qe.util.toc()\n",
    "\n",
    "print(\"Timing policy function iteration with endogenous grid\")\n",
    "&sigma; = lambda x: x  # Initial policy\n",
    "qe.util.tic()\n",
    "for i in range(sim_length):\n",
    "    &sigma;_new = K_egm(&sigma;)\n",
    "    &sigma; = &sigma;_new\n",
    "qe.util.toc()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We see that the EGM version is significantly faster, even without jit compilation!\n",
    "\n",
    "The absence of numerical root-finding means that it is typically more accurate at each step as well."
   ]
  }
 ],
 "metadata": {
  "filename": "egm_policy_iter.rst",
  "kernelspec": {
   "display_name": "Python",
   "language": "python3",
   "name": "python3"
  },
  "title": "Optimal Growth III: The Endogenous Grid Method"
 },
 "nbformat": 4,
 "nbformat_minor": 2
}