egm_policy_iter.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "11d4be5b",
   "metadata": {},
   "source": [
    "# Optimal Growth IV: The Endogenous Grid Method"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8ed231ad",
   "metadata": {},
   "source": [
    "## Contents\n",
    "\n",
    "- [Optimal Growth IV: The Endogenous Grid Method](#Optimal-Growth-IV:-The-Endogenous-Grid-Method)  \n",
    "  - [Overview](#Overview)  \n",
    "  - [Key Idea](#Key-Idea)  \n",
    "  - [Implementation](#Implementation)  "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7435b143",
   "metadata": {},
   "source": [
    "## Overview\n",
    "\n",
    "Previously, we solved the stochastic optimal growth model using\n",
    "\n",
    "1. [value function iteration](https://python.quantecon.org/optgrowth_fast.html)  \n",
    "1. [Euler equation based time iteration](https://python.quantecon.org/coleman_policy_iter.html)  \n",
    "\n",
    "\n",
    "We found time iteration to be significantly more accurate and efficient.\n",
    "\n",
    "In this lecture, we’ll look at a clever twist on time iteration 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",
    "The original reference is [[Carroll, 2006](https://python.quantecon.org/zreferences.html#id153)].\n",
    "\n",
    "Let’s start with some standard imports:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b077d200",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from numba import njit"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a6e00af2",
   "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",
   "id": "ba0e4cf7",
   "metadata": {},
   "source": [
    "### Theory\n",
    "\n",
    "Take the model set out in [the time iteration lecture](https://python.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{41.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{41.2}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "08b55436",
   "metadata": {},
   "source": [
    "### Exogenous Grid\n",
    "\n",
    "As discussed in [the lecture on time iteration](https://python.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://python.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 [(41.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…"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e7be62e1",
   "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 [[Carroll, 2006](https://python.quantecon.org/zreferences.html#id153)], 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",
    "Let $ (u')^{-1} $ be the inverse function of $ u' $.\n",
    "\n",
    "The idea is this:\n",
    "\n",
    "- First, we fix an *exogenous* grid $ \\{k_i\\} $ for capital ($ k = y - c $).  \n",
    "- Then we obtain  $ c_i $ via  \n",
    "\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{41.3}\n",
    "$$\n",
    "\n",
    "- Finally, for each $ c_i $ we set $ y_i = c_i + k_i $.  \n",
    "\n",
    "\n",
    "It is clear that each $ (y_i, c_i) $ pair constructed in this manner satisfies [(41.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",
   "id": "cfd68032",
   "metadata": {},
   "source": [
    "## Implementation\n",
    "\n",
    "As [before](https://python.quantecon.org/coleman_policy_iter.html), we will start with a simple setting\n",
    "where\n",
    "\n",
    "- $ u(c) = \\ln c $,  \n",
    "- production is Cobb-Douglas, and  \n",
    "- the shocks are lognormal.  \n",
    "\n",
    "\n",
    "This will allow us to make comparisons with the analytical solutions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a629798d",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "\n",
    "def v_star(y, α, β, μ):\n",
    "    \"\"\"\n",
    "    True value function\n",
    "    \"\"\"\n",
    "    c1 = np.log(1 - α * β) / (1 - β)\n",
    "    c2 = (μ + α * np.log(α * β)) / (1 - α)\n",
    "    c3 = 1 / (1 - β)\n",
    "    c4 = 1 / (1 - α * β)\n",
    "    return c1 + c2 * (c3 - c4) + c4 * np.log(y)\n",
    "\n",
    "def σ_star(y, α, β):\n",
    "    \"\"\"\n",
    "    True optimal policy\n",
    "    \"\"\"\n",
    "    return (1 - α * β) * y"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9be7fc5e",
   "metadata": {},
   "source": [
    "We reuse the `OptimalGrowthModel` class"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "99df1052",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "from numba import float64\n",
    "from numba.experimental import jitclass\n",
    "\n",
    "opt_growth_data = [\n",
    "    ('α', float64),          # Production parameter\n",
    "    ('β', float64),          # Discount factor\n",
    "    ('μ', float64),          # Shock location parameter\n",
    "    ('s', float64),          # Shock scale parameter\n",
    "    ('grid', float64[:]),    # Grid (array)\n",
    "    ('shocks', float64[:])   # Shock draws (array)\n",
    "]\n",
    "\n",
    "@jitclass(opt_growth_data)\n",
    "class OptimalGrowthModel:\n",
    "\n",
    "    def __init__(self,\n",
    "                α=0.4,\n",
    "                β=0.96,\n",
    "                μ=0,\n",
    "                s=0.1,\n",
    "                grid_max=4,\n",
    "                grid_size=120,\n",
    "                shock_size=250,\n",
    "                seed=1234):\n",
    "\n",
    "        self.α, self.β, self.μ, self.s = α, β, μ, s\n",
    "\n",
    "        # Set up grid\n",
    "        self.grid = np.linspace(1e-5, grid_max, grid_size)\n",
    "\n",
    "        # Store shocks (with a seed, so results are reproducible)\n",
    "        np.random.seed(seed)\n",
    "        self.shocks = np.exp(μ + s * np.random.randn(shock_size))\n",
    "\n",
    "\n",
    "    def f(self, k):\n",
    "        \"The production function\"\n",
    "        return k**self.α\n",
    "\n",
    "\n",
    "    def u(self, c):\n",
    "        \"The utility function\"\n",
    "        return np.log(c)\n",
    "\n",
    "    def f_prime(self, k):\n",
    "        \"Derivative of f\"\n",
    "        return self.α * (k**(self.α - 1))\n",
    "\n",
    "\n",
    "    def u_prime(self, c):\n",
    "        \"Derivative of u\"\n",
    "        return 1/c\n",
    "\n",
    "    def u_prime_inv(self, c):\n",
    "        \"Inverse of u'\"\n",
    "        return 1/c"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3b65e82c",
   "metadata": {},
   "source": [
    "### The Operator\n",
    "\n",
    "Here’s an implementation of $ K $ using EGM as described above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8e76d68e",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "@njit\n",
    "def K(σ_array, og):\n",
    "    \"\"\"\n",
    "    The Coleman-Reffett operator using EGM\n",
    "\n",
    "    \"\"\"\n",
    "\n",
    "    # Simplify names\n",
    "    f, β = og.f, og.β\n",
    "    f_prime, u_prime = og.f_prime, og.u_prime\n",
    "    u_prime_inv = og.u_prime_inv\n",
    "    grid, shocks = og.grid, og.shocks\n",
    "\n",
    "    # Determine endogenous grid\n",
    "    y = grid + σ_array  # y_i = k_i + c_i\n",
    "\n",
    "    # Linear interpolation of policy using endogenous grid\n",
    "    σ = lambda x: np.interp(x, y, σ_array)\n",
    "\n",
    "    # Allocate memory for new consumption array\n",
    "    c = np.empty_like(grid)\n",
    "\n",
    "    # Solve for updated consumption value\n",
    "    for i, k in enumerate(grid):\n",
    "        vals = u_prime(σ(f(k) * shocks)) * f_prime(k) * shocks\n",
    "        c[i] = u_prime_inv(β * np.mean(vals))\n",
    "\n",
    "    return c"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d20d9909",
   "metadata": {},
   "source": [
    "Note the lack of any root-finding algorithm."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3b1c6f34",
   "metadata": {},
   "source": [
    "### Testing\n",
    "\n",
    "First we create an instance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d0126cb3",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "og = OptimalGrowthModel()\n",
    "grid = og.grid"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d370ee6b",
   "metadata": {},
   "source": [
    "Here’s our solver routine:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3f937311",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def solve_model_time_iter(model,    # Class with model information\n",
    "                          σ,        # Initial condition\n",
    "                          tol=1e-4,\n",
    "                          max_iter=1000,\n",
    "                          verbose=True,\n",
    "                          print_skip=25):\n",
    "\n",
    "    # Set up loop\n",
    "    i = 0\n",
    "    error = tol + 1\n",
    "\n",
    "    while i < max_iter and error > tol:\n",
    "        σ_new = K(σ, model)\n",
    "        error = np.max(np.abs(σ - σ_new))\n",
    "        i += 1\n",
    "        if verbose and i % print_skip == 0:\n",
    "            print(f\"Error at iteration {i} is {error}.\")\n",
    "        σ = σ_new\n",
    "\n",
    "    if error > tol:\n",
    "        print(\"Failed to converge!\")\n",
    "    elif verbose:\n",
    "        print(f\"\\nConverged in {i} iterations.\")\n",
    "\n",
    "    return σ_new"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "306e65bb",
   "metadata": {},
   "source": [
    "Let’s call it:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f6c6ecf8",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "σ_init = np.copy(grid)\n",
    "σ = solve_model_time_iter(og, σ_init)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8d4cbc42",
   "metadata": {},
   "source": [
    "Here is a plot of the resulting policy, compared with the true policy:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5149a2af",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "y = grid + σ  # y_i = k_i + c_i\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "\n",
    "ax.plot(y, σ, lw=2,\n",
    "        alpha=0.8, label='approximate policy function')\n",
    "\n",
    "ax.plot(y, σ_star(y, og.α, og.β), 'k--',\n",
    "        lw=2, alpha=0.8, label='true policy function')\n",
    "\n",
    "ax.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "974a0088",
   "metadata": {},
   "source": [
    "The maximal absolute deviation between the two policies is"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "04ecfa35",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "np.max(np.abs(σ - σ_star(y, og.α, og.β)))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8e0429a4",
   "metadata": {},
   "source": [
    "How long does it take to converge?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3b9a9ee2",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "%%timeit -n 3 -r 1\n",
    "σ = solve_model_time_iter(og, σ_init, verbose=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a439aade",
   "metadata": {},
   "source": [
    "Relative to time iteration, which as already found to be highly efficient, EGM\n",
    "has managed to shave off still more run time without compromising accuracy.\n",
    "\n",
    "This is due to the lack of a numerical root-finding step.\n",
    "\n",
    "We can now solve the optimal growth model at given parameters extremely fast."
   ]
  }
 ],
 "metadata": {
  "date": 1714442504.7872095,
  "filename": "egm_policy_iter.md",
  "kernelspec": {
   "display_name": "Python",
   "language": "python3",
   "name": "python3"
  },
  "title": "Optimal Growth IV: The Endogenous Grid Method"
 },
 "nbformat": 4,
 "nbformat_minor": 5
}