dynamic_programming/egm_policy_iter.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "<a id='egm-policy-iter'></a>\n",
    "<div id=\"qe-notebook-header\" style=\"text-align:right;\">\n",
    "        <a href=\"https://quantecon.org/\" title=\"quantecon.org\">\n",
    "                <img style=\"width:250px;display:inline;\" 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": [
    "## Overview\n",
    "\n",
    "We solved the stochastic optimal growth model using\n",
    "\n",
    "1. [value function iteration](optgrowth.html)  \n",
    "1. [Euler equation based time iteration](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’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]](zreferences.html#carroll2006)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Key Idea\n",
    "\n",
    "Let’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](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 c^*)(y)\n",
    "= \\beta \\int (u'\\circ c^*)(f(y - c^*(y)) z) f'(y - c^*(y)) z \\phi(dz) \\tag{1}\n",
    "$$\n",
    "\n",
    "As we saw, the Coleman operator is a nonlinear operator $ K $ engineered so that $ c^* $ is a fixed point of $ K $\n",
    "\n",
    "It takes as its argument a continuous strictly increasing consumption policy $ g \\in \\Sigma $\n",
    "\n",
    "It returns a new function $ Kg $,  where $ (Kg)(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 g) (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](coleman_policy_iter.html), to implement the method on a computer we need 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](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 g $ at $ y_i $\n",
    "\n",
    "Thus, with the points $ \\{y_i, c_i\\} $ in hand, we can reconstruct $ Kg $ via approximation\n",
    "\n",
    "Iteration then continues…"
   ]
  },
  {
   "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]](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 g) (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 $ Kg $ 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’s implement this version of the Coleman operator and see how it performs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### The Operator\n",
    "\n",
    "Here’s an implementation of $ K $ using EGM as described above"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": true
   },
   "outputs": [],
   "source": [
    "using InstantiateFromURL\n",
    "github_project(\"QuantEcon/quantecon-notebooks-julia\", version = \"0.1.0\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "using LinearAlgebra, Statistics, Compat\n",
    "using BenchmarkTools, Interpolations, Parameters, Plots, QuantEcon, Random, Roots\n",
    "gr(fmt = :png);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "function coleman_egm(g, k_grid, β, u′, u′_inv, f, f′, shocks)\n",
    "\n",
    "    # Allocate memory for value of consumption on endogenous grid points\n",
    "    c = similar(k_grid)\n",
    "\n",
    "    # Solve for updated consumption value\n",
    "    for (i, k) in enumerate(k_grid)\n",
    "        vals = u′.(g.(f(k) * shocks)) .* f′(k) .* shocks\n",
    "        c[i] = u′_inv(β * mean(vals))\n",
    "    end\n",
    "\n",
    "    # Determine endogenous grid\n",
    "    y = k_grid + c  # y_i = k_i + c_i\n",
    "\n",
    "    # Update policy function and return\n",
    "    Kg = LinearInterpolation(y,c, extrapolation_bc=Line())\n",
    "    Kg_f(x) = Kg(x)\n",
    "    return Kg_f\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note the lack of any root finding algorithm\n",
    "\n",
    "We’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": [
    "function K!(Kg, g, grid, β, u′, f, f′, shocks)\n",
    "\n",
    "    # This function requires the container of the output value as argument Kg\n",
    "\n",
    "    # Construct linear interpolation object #\n",
    "    g_func = LinearInterpolation(grid, g, extrapolation_bc = Line())\n",
    "\n",
    "    # solve for updated consumption value #\n",
    "    for (i, y) in enumerate(grid)\n",
    "        function h(c)\n",
    "            vals = u′.(g_func.(f(y - c) * shocks)) .* f′(y - c) .* shocks\n",
    "            return u′(c) - β * mean(vals)\n",
    "        end\n",
    "        Kg[i] = find_zero(h, (1e-10, y - 1e-10))\n",
    "    end\n",
    "    return Kg\n",
    "end\n",
    "\n",
    "# The following function does NOT require the container of the output value as argument\n",
    "K(g, grid, β, u′, f, f′, shocks) =\n",
    "    K!(similar(g), g, grid, β, u′, f, f′, shocks)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let’s test out the code above on some example parameterizations, after the following imports"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Testing on the Log / Cobb–Douglas case\n",
    "\n",
    "As we [did for value function iteration](optgrowth.html) and [time iteration](coleman_policy_iter.html), let’s start by testing our method with the log-linear benchmark\n",
    "\n",
    "The first step is to bring in the model that we used in the [Coleman policy function iteration](coleman_policy_iter.html)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "# model\n",
    "\n",
    "Model = @with_kw (α = 0.65, # productivity parameter\n",
    "                  β = 0.95, # discount factor\n",
    "                  γ = 1.0,  # risk aversion\n",
    "                  μ = 0.0,  # lognorm(μ, σ)\n",
    "                  s = 0.1,  # lognorm(μ, σ)\n",
    "                  grid_min = 1e-6, # smallest grid point\n",
    "                  grid_max = 4.0,  # largest grid point\n",
    "                  grid_size = 200, # grid size\n",
    "                  u = γ == 1 ? log : c->(c^(1-γ)-1)/(1-γ), # utility function\n",
    "                  u′ = c-> c^(-γ), # u'\n",
    "                  f = k-> k^α, # production function\n",
    "                  f′ = k -> α*k^(α-1), # f'\n",
    "                  grid = range(grid_min, grid_max, length = grid_size)) # grid"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next we generate an instance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "mlog = Model(); # Log Linear model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We also need some shock draws for Monte Carlo integration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "Random.seed!(42); # For reproducible behavior.\n",
    "\n",
    "shock_size = 250     # Number of shock draws in Monte Carlo integral\n",
    "shocks = exp.(mlog.μ .+ mlog.s * randn(shock_size));"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As a preliminary test, let’s see if $ K c^* = c^* $, as implied by the theory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "c_star(y) = (1 - mlog.α * mlog.β) * y\n",
    "\n",
    "# some useful constants\n",
    "ab = mlog.α * mlog.β\n",
    "c1 = log(1 - ab) / (1 - mlog.β)\n",
    "c2 = (mlog.μ + mlog.α * log(ab)) / (1 - mlog.α)\n",
    "c3 = 1 / (1 - mlog.β)\n",
    "c4 = 1 / (1 - ab)\n",
    "\n",
    "v_star(y) = c1 + c2 * (c3 - c4) + c4 * log(y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "function verify_true_policy(m, shocks, c_star)\n",
    "    k_grid = m.grid\n",
    "    c_star_new = coleman_egm(c_star, k_grid, m.β, m.u′, m.u′, m.f, m.f′, shocks)\n",
    "\n",
    "    plt = plot()\n",
    "    plot!(plt, k_grid, c_star.(k_grid), lw = 2, label = \"optimal policy c*\")\n",
    "    plot!(plt, k_grid, c_star_new.(k_grid), lw = 2, label = \"Kc*\")\n",
    "    plot!(plt, legend = :topleft)\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "verify_true_policy(mlog, shocks, c_star)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notice that we’re passing u′ to coleman_egm twice\n",
    "\n",
    "The reason is that, in the case of log utility, $ u'(c) = (u')^{-1}(c) = 1/c $\n",
    "\n",
    "Hence u′ and u′_inv are the same\n",
    "\n",
    "We can’t really distinguish the two plots\n",
    "\n",
    "In fact it’s easy to see that the difference is essentially zero:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "c_star_new = coleman_egm(c_star, mlog.grid, mlog.β, mlog.u′,\n",
    "                         mlog.u′, mlog.f, mlog.f′, shocks)\n",
    "maximum(abs(c_star_new(g) - c_star(g)) for g in mlog.grid)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next let’s try iterating from an arbitrary initial condition and see if we\n",
    "converge towards $ c^* $\n",
    "\n",
    "Let’s start from the consumption policy that eats the whole pie: $ c(y) = y $"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "n = 15\n",
    "function check_convergence(m, shocks, c_star, g_init, n_iter)\n",
    "    k_grid = m.grid\n",
    "    g = g_init\n",
    "    plt = plot()\n",
    "    plot!(plt, m.grid, g.(m.grid),\n",
    "          color = RGBA(0,0,0,1), lw = 2, alpha = 0.6, label = \"initial condition c(y) = y\")\n",
    "    for i in 1:n_iter\n",
    "        new_g = coleman_egm(g, k_grid, m.β, m.u′, m.u′, m.f, m.f′, shocks)\n",
    "        g = new_g\n",
    "        plot!(plt, k_grid, new_g.(k_grid), alpha = 0.6, color = RGBA(0,0,(i / n_iter), 1),\n",
    "              lw = 2, label = \"\")\n",
    "    end\n",
    "\n",
    "    plot!(plt, k_grid, c_star.(k_grid),\n",
    "          color = :black, lw = 2, alpha = 0.8, label = \"true policy function c*\")\n",
    "    plot!(plt, legend = :topleft)\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "check_convergence(mlog, shocks, c_star, identity, n)"
   ]
  },
  {
   "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’s compare the clock times per iteration for the standard Coleman\n",
    "operator (with exogenous grid) and the EGM version\n",
    "\n",
    "We’ll do so using the CRRA model adopted in the exercises of the [Euler equation time iteration lecture](coleman_policy_iter.html)\n",
    "\n",
    "Here’s the model and some convenient functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "mcrra = Model(α = 0.65, β = 0.95, γ = 1.5)\n",
    "u′_inv(c) = c^(-1 / mcrra.γ)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here’s the result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "crra_coleman(g, m, shocks) = K(g, m.grid, m.β, m.u′, m.f, m.f′, shocks)\n",
    "crra_coleman_egm(g, m, shocks) = coleman_egm(g, m.grid, m.β, m.u′,\n",
    "                                             u′_inv, m.f, m.f′, shocks)\n",
    "function coleman(m = m, shocks = shocks; sim_length = 20)\n",
    "    g = m.grid\n",
    "    for i in 1:sim_length\n",
    "        g = crra_coleman(g, m, shocks)\n",
    "    end\n",
    "    return g\n",
    "end\n",
    "function egm(m, g = identity, shocks = shocks; sim_length = 20)\n",
    "    for i in 1:sim_length\n",
    "        g = crra_coleman_egm(g, m, shocks)\n",
    "    end\n",
    "    return g.(m.grid)\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "@benchmark coleman($mcrra)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "@benchmark egm($mcrra)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We see that the EGM version is about 30 times faster\n",
    "\n",
    "At the same time, the absence of numerical root finding means that it is\n",
    "typically more accurate at each step as well"
   ]
  }
 ],
 "metadata": {
  "filename": "egm_policy_iter.rst",
  "kernelspec": {
   "display_name": "Julia 1.2",
   "language": "julia",
   "name": "julia-1.2"
  },
  "title": "Optimal Growth III: The Endogenous Grid Method"
 },
 "nbformat": 4,
 "nbformat_minor": 2
}