coleman_policy_iter.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<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: Time Iteration"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Contents\n",
    "\n",
    "- [Optimal Growth III: Time Iteration](#Optimal-Growth-III:-Time-Iteration)  \n",
    "  - [Overview](#Overview)  \n",
    "  - [The Euler Equation](#The-Euler-Equation)  \n",
    "  - [Implementation](#Implementation)  \n",
    "  - [Exercises](#Exercises)  \n",
    "  - [Solutions](#Solutions)  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In addition to what’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",
    "In this lecture, we’ll continue our [earlier](https://python-programming.quantecon.org/optgrowth.html) study of the stochastic optimal growth model.\n",
    "\n",
    "In that lecture, we solved the associated dynamic programming\n",
    "problem using value function iteration.\n",
    "\n",
    "The beauty of this technique is its broad applicability.\n",
    "\n",
    "With numerical problems, however, we can often attain higher efficiency in\n",
    "specific applications by deriving methods that are carefully tailored to the\n",
    "application at hand.\n",
    "\n",
    "The stochastic optimal growth model has plenty of structure to exploit for\n",
    "this purpose, especially when we adopt some concavity and smoothness\n",
    "assumptions over primitives.\n",
    "\n",
    "We’ll use this structure to obtain an Euler equation based method.\n",
    "\n",
    "This will be an extension of the time iteration method considered\n",
    "in our elementary lecture on [cake eating](https://python-programming.quantecon.org/cake_eating_numerical.html).\n",
    "\n",
    "In a [subsequent lecture](https://python-programming.quantecon.org/egm_policy_iter.html), we’ll see that time\n",
    "iteration can be further adjusted to obtain even more efficiency.\n",
    "\n",
    "Let’s start with some imports:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import quantecon as qe\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "\n",
    "from interpolation import interp\n",
    "from quantecon.optimize import brentq\n",
    "from numba import njit, jitclass, float64"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The Euler Equation\n",
    "\n",
    "Our first step is to derive the Euler equation, which is a generalization of\n",
    "the Euler equation we obtained in the [lecture on cake eating](https://python-programming.quantecon.org/cake_eating_problem.html).\n",
    "\n",
    "We take the model set out in [the stochastic growth model lecture](https://python-programming.quantecon.org/optgrowth.html) and add the following assumptions:\n",
    "\n",
    "1. $ u $ and $ f $ are continuously differentiable and strictly concave  \n",
    "1. $ f(0) = 0 $  \n",
    "1. $ \\lim_{c \\to 0} u'(c) = \\infty $ and $ \\lim_{c \\to \\infty} u'(c) = 0 $  \n",
    "1. $ \\lim_{k \\to 0} f'(k) = \\infty $ and $ \\lim_{k \\to \\infty} f'(k) = 0 $  \n",
    "\n",
    "\n",
    "The last two conditions are usually called **Inada conditions**.\n",
    "\n",
    "Recall the Bellman equation\n",
    "\n",
    "\n",
    "<a id='equation-cpi-fpb30'></a>\n",
    "$$\n",
    "v^*(y) = \\max_{0 \\leq c \\leq y}\n",
    "    \\left\\{\n",
    "        u(c) + \\beta \\int v^*(f(y - c) z) \\phi(dz)\n",
    "    \\right\\}\n",
    "\\quad \\text{for all} \\quad\n",
    "y \\in \\mathbb R_+ \\tag{1}\n",
    "$$\n",
    "\n",
    "Let the optimal consumption policy be denoted by $ \\sigma^* $.\n",
    "\n",
    "We know that $ \\sigma^* $ is a $ v^* $-greedy policy so that $ \\sigma^*(y) $ is the maximizer in [(1)](#equation-cpi-fpb30).\n",
    "\n",
    "The conditions above imply that\n",
    "\n",
    "- $ \\sigma^* $ is the unique optimal policy for the stochastic optimal growth model  \n",
    "- the optimal policy is continuous, strictly increasing and also **interior**, in the sense that $ 0 < \\sigma^*(y) < y $ for all strictly positive $ y $, and  \n",
    "- the value function is strictly concave and continuously differentiable, with  \n",
    "\n",
    "\n",
    "\n",
    "<a id='equation-cpi-env'></a>\n",
    "$$\n",
    "(v^*)'(y) = u' (\\sigma^*(y) ) := (u' \\circ \\sigma^*)(y) \\tag{2}\n",
    "$$\n",
    "\n",
    "The last result is called the **envelope condition** due to its relationship with the [envelope theorem](https://en.wikipedia.org/wiki/Envelope_theorem).\n",
    "\n",
    "To see why [(2)](#equation-cpi-env) holds, write the Bellman equation in the equivalent\n",
    "form\n",
    "\n",
    "$$\n",
    "v^*(y) = \\max_{0 \\leq k \\leq y}\n",
    "    \\left\\{\n",
    "        u(y-k) + \\beta \\int v^*(f(k) z) \\phi(dz)\n",
    "    \\right\\},\n",
    "$$\n",
    "\n",
    "Differentiating with respect to $ y $,  and then evaluating at the optimum yields [(2)](#equation-cpi-env).\n",
    "\n",
    "(Section 12.1 of [EDTC](http://johnstachurski.net/edtc.html) contains full proofs of these results, and closely related discussions can be found in many other texts.)\n",
    "\n",
    "Differentiability of the value function and interiority of the optimal policy\n",
    "imply that optimal consumption satisfies the first order condition associated\n",
    "with [(1)](#equation-cpi-fpb30), which is\n",
    "\n",
    "\n",
    "<a id='equation-cpi-foc'></a>\n",
    "$$\n",
    "u'(\\sigma^*(y)) = \\beta \\int (v^*)'(f(y - \\sigma^*(y)) z) f'(y - \\sigma^*(y)) z \\phi(dz) \\tag{3}\n",
    "$$\n",
    "\n",
    "Combining [(2)](#equation-cpi-env) and the first-order condition [(3)](#equation-cpi-foc) gives the **Euler equation**\n",
    "\n",
    "\n",
    "<a id='equation-cpi-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{4}\n",
    "$$\n",
    "\n",
    "We can think of the Euler equation as a functional equation\n",
    "\n",
    "\n",
    "<a id='equation-cpi-euler-func'></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{5}\n",
    "$$\n",
    "\n",
    "over interior consumption policies $ \\sigma $, one solution of which is the optimal policy $ \\sigma^* $.\n",
    "\n",
    "Our aim is to solve the functional equation [(5)](#equation-cpi-euler-func) and hence obtain $ \\sigma^* $."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### The Coleman-Reffett Operator\n",
    "\n",
    "Recall the Bellman operator\n",
    "\n",
    "\n",
    "<a id='equation-fcbell20-coleman'></a>\n",
    "$$\n",
    "Tv(y) := \\max_{0 \\leq c \\leq y}\n",
    "\\left\\{\n",
    "    u(c) + \\beta \\int v(f(y - c) z) \\phi(dz)\n",
    "\\right\\} \\tag{6}\n",
    "$$\n",
    "\n",
    "Just as we introduced the Bellman operator to solve the Bellman equation, we\n",
    "will now introduce an operator over policies to help us solve the Euler\n",
    "equation.\n",
    "\n",
    "This operator $ K $ will act on the set of all $ \\sigma \\in \\Sigma $\n",
    "that are continuous, strictly increasing and interior.\n",
    "\n",
    "Henceforth we denote this set of policies by $ \\mathscr P $\n",
    "\n",
    "1. The operator $ K $ takes as its argument a $ \\sigma \\in \\mathscr P $ and  \n",
    "1. returns a new function $ K\\sigma $,  where $ K\\sigma(y) $ is the $ c \\in (0, y) $ that solves.  \n",
    "\n",
    "\n",
    "\n",
    "<a id='equation-cpi-coledef'></a>\n",
    "$$\n",
    "u'(c)\n",
    "= \\beta \\int (u' \\circ \\sigma) (f(y - c) z ) f'(y - c) z \\phi(dz) \\tag{7}\n",
    "$$\n",
    "\n",
    "We call this operator the **Coleman-Reffett operator** to acknowledge the work of\n",
    "[[Col90]](https://python-programming.quantecon.org/zreferences.html#coleman1990) and [[Ref96]](https://python-programming.quantecon.org/zreferences.html#reffett1996).\n",
    "\n",
    "In essence, $ K\\sigma $ is the consumption policy that the Euler equation tells\n",
    "you to choose today when your future consumption policy is $ \\sigma $.\n",
    "\n",
    "The  important thing to note about $ K $ is that, by\n",
    "construction, its fixed points coincide with solutions to the functional\n",
    "equation [(5)](#equation-cpi-euler-func).\n",
    "\n",
    "In particular, the optimal policy $ \\sigma^* $ is a fixed point.\n",
    "\n",
    "Indeed, for fixed $ y $, the value $ K\\sigma^*(y) $ is the $ c $ that\n",
    "solves\n",
    "\n",
    "$$\n",
    "u'(c)\n",
    "= \\beta \\int (u' \\circ \\sigma^*) (f(y - c) z ) f'(y - c) z \\phi(dz)\n",
    "$$\n",
    "\n",
    "In view of the Euler equation, this is exactly $ \\sigma^*(y) $."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Is the Coleman-Reffett Operator Well Defined?\n",
    "\n",
    "In particular, is there always a unique $ c \\in (0, y) $ that solves\n",
    "[(7)](#equation-cpi-coledef)?\n",
    "\n",
    "The answer is yes, under our assumptions.\n",
    "\n",
    "For any $ \\sigma \\in \\mathscr P $, the right side of [(7)](#equation-cpi-coledef)\n",
    "\n",
    "- is continuous and strictly increasing in $ c $ on $ (0, y) $  \n",
    "- diverges to $ +\\infty $ as $ c \\uparrow y $  \n",
    "\n",
    "\n",
    "The left side of [(7)](#equation-cpi-coledef)\n",
    "\n",
    "- is continuous and strictly decreasing in $ c $ on $ (0, y) $  \n",
    "- diverges to $ +\\infty $ as $ c \\downarrow 0 $  \n",
    "\n",
    "\n",
    "Sketching these curves and using the information above will convince you that they cross exactly once as $ c $ ranges over $ (0, y) $.\n",
    "\n",
    "With a bit more analysis, one can show in addition that $ K \\sigma \\in \\mathscr P $\n",
    "whenever $ \\sigma \\in \\mathscr P $."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Comparison with VFI (Theory)\n",
    "\n",
    "It is possible to prove that there is a tight relationship between iterates of\n",
    "$ K $ and iterates of the Bellman operator.\n",
    "\n",
    "Mathematically, the two operators are *topologically conjugate*.\n",
    "\n",
    "Loosely speaking, this means that if iterates of one operator converge then\n",
    "so do iterates of the other, and vice versa.\n",
    "\n",
    "Moreover, there is a sense in which they converge at the same rate, at least\n",
    "in theory.\n",
    "\n",
    "However, it turns out that the operator $ K $ is more stable numerically\n",
    "and hence more efficient in the applications we consider.\n",
    "\n",
    "Examples are given below."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Implementation\n",
    "\n",
    "As in our [previous study](https://python-programming.quantecon.org/optgrowth_fast.html), we continue to assume that\n",
    "\n",
    "- $ u(c) = \\ln c $  \n",
    "- $ f(k) = k^{\\alpha} $  \n",
    "- $ \\phi $ is the distribution of $ \\xi := \\exp(\\mu + s \\zeta) $ when $ \\zeta $ is standard normal  \n",
    "\n",
    "\n",
    "This will allow us to compare our results to the analytical solutions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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",
   "metadata": {},
   "source": [
    "As discussed above, our plan is to solve the model using time iteration, which\n",
    "means iterating with the operator $ K $.\n",
    "\n",
    "For this we need access to the functions $ u' $ and $ f, f' $.\n",
    "\n",
    "These are available in a class called `OptimalGrowthModel` that we\n",
    "constructed in an [earlier lecture](https://python-programming.quantecon.org/optgrowth_fast.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "\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",
   "metadata": {},
   "source": [
    "Now we implement a method called `euler_diff`, which returns\n",
    "\n",
    "\n",
    "<a id='equation-euler-diff'></a>\n",
    "$$\n",
    "u'(c) - \\beta \\int (u' \\circ \\sigma) (f(y - c) z ) f'(y - c) z \\phi(dz) \\tag{8}\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "@njit\n",
    "def euler_diff(c, σ, y, og):\n",
    "    \"\"\"\n",
    "    Set up a function such that the root with respect to c,\n",
    "    given y and σ, is equal to Kσ(y).\n",
    "\n",
    "    \"\"\"\n",
    "\n",
    "    β, shocks, grid = og.β, og.shocks, og.grid\n",
    "    f, f_prime, u_prime = og.f, og.f_prime, og.u_prime\n",
    "\n",
    "    # First turn σ into a function via interpolation\n",
    "    σ_func = lambda x: interp(grid, σ, x)\n",
    "\n",
    "    # Now set up the function we need to find the root of.\n",
    "    vals = u_prime(σ_func(f(y - c) * shocks)) * f_prime(y - c) * shocks\n",
    "    return u_prime(c) - β * np.mean(vals)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The function `euler_diff` evaluates integrals by Monte Carlo and\n",
    "approximates functions using linear interpolation.\n",
    "\n",
    "We will use a root-finding algorithm to solve [(8)](#equation-euler-diff) for $ c $ given\n",
    "state $ y $ and $ σ $, the current guess of the policy.\n",
    "\n",
    "Here’s the operator $ K $, that implements the root-finding step."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "@njit\n",
    "def K(σ, og):\n",
    "    \"\"\"\n",
    "    The Coleman-Reffett operator\n",
    "\n",
    "     Here og is an instance of OptimalGrowthModel.\n",
    "    \"\"\"\n",
    "\n",
    "    β = og.β\n",
    "    f, f_prime, u_prime = og.f, og.f_prime, og.u_prime\n",
    "    grid, shocks = og.grid, og.shocks\n",
    "\n",
    "    σ_new = np.empty_like(σ)\n",
    "    for i, y in enumerate(grid):\n",
    "        # Solve for optimal c at y\n",
    "        c_star = brentq(euler_diff, 1e-10, y-1e-10, args=(σ, y, og))[0]\n",
    "        σ_new[i] = c_star\n",
    "\n",
    "    return σ_new"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Testing\n",
    "\n",
    "Let’s generate an instance and plot some iterates of $ K $, starting from $ σ(y) = y $."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "og = OptimalGrowthModel()\n",
    "grid = og.grid\n",
    "\n",
    "n = 15\n",
    "σ = grid.copy()  # Set initial condition\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "lb = 'initial condition $\\sigma(y) = y$'\n",
    "ax.plot(grid, σ, color=plt.cm.jet(0), alpha=0.6, label=lb)\n",
    "\n",
    "for i in range(n):\n",
    "    σ = K(σ, og)\n",
    "    ax.plot(grid, σ, color=plt.cm.jet(i / n), alpha=0.6)\n",
    "\n",
    "# Update one more time and plot the last iterate in black\n",
    "σ = K(σ, og)\n",
    "ax.plot(grid, σ, color='k', alpha=0.8, label='last iterate')\n",
    "\n",
    "ax.legend()\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We see that the iteration process converges quickly to a limit\n",
    "that resembles the solution we obtained in [the previous lecture](https://python-programming.quantecon.org/optgrowth_fast.html).\n",
    "\n",
    "Here is a function called `solve_model_time_iter` that takes an instance of\n",
    "`OptimalGrowthModel` and returns an approximation to the optimal policy,\n",
    "using time iteration."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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 i == max_iter:\n",
    "        print(\"Failed to converge!\")\n",
    "\n",
    "    if verbose and i < max_iter:\n",
    "        print(f\"\\nConverged in {i} iterations.\")\n",
    "\n",
    "    return σ_new"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let’s call it:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "σ_init = np.copy(og.grid)\n",
    "σ = solve_model_time_iter(og, σ_init)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here is a plot of the resulting policy, compared with the true policy:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "\n",
    "ax.plot(og.grid, σ, lw=2,\n",
    "        alpha=0.8, label='approximate policy function')\n",
    "\n",
    "ax.plot(og.grid, σ_star(og.grid, og.α, og.β), 'k--',\n",
    "        lw=2, alpha=0.8, label='true policy function')\n",
    "\n",
    "ax.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Again, the fit is excellent.\n",
    "\n",
    "The maximal absolute deviation between the two policies is"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "np.max(np.abs(σ - σ_star(og.grid, og.α, og.β)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "How long does it take to converge?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "%%timeit -n 3 -r 1\n",
    "σ = solve_model_time_iter(og, σ_init, verbose=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Convergence is very fast, even compared to our [JIT-compiled value function iteration](https://python-programming.quantecon.org/optgrowth_fast.html).\n",
    "\n",
    "Overall, we find that time iteration provides a very high degree of efficiency\n",
    "and accuracy, at least for this model."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercises"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 1\n",
    "\n",
    "Solve the model with CRRA utility\n",
    "\n",
    "$$\n",
    "u(c) = \\frac{c^{1 - \\gamma}} {1 - \\gamma}\n",
    "$$\n",
    "\n",
    "Set `γ = 1.5`.\n",
    "\n",
    "Compute and plot the optimal policy."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Solutions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 1\n",
    "\n",
    "We use the class `OptimalGrowthModel_CRRA` from our [VFI lecture](https://python-programming.quantecon.org/optgrowth_fast.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "opt_growth_data = [\n",
    "    ('α', float64),          # Production parameter\n",
    "    ('β', float64),          # Discount factor\n",
    "    ('μ', float64),          # Shock location parameter\n",
    "    ('γ', float64),          # Preference 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_CRRA:\n",
    "\n",
    "    def __init__(self,\n",
    "                α=0.4, \n",
    "                β=0.96, \n",
    "                μ=0,\n",
    "                s=0.1,\n",
    "                γ=1.5, \n",
    "                grid_max=4,\n",
    "                grid_size=120,\n",
    "                shock_size=250,\n",
    "                seed=1234):\n",
    "\n",
    "        self.α, 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",
    "    def u(self, c):\n",
    "        \"The utility function.\"\n",
    "        return c**(1 - self.γ) / (1 - self.γ)\n",
    "\n",
    "    def f_prime(self, k):\n",
    "        \"Derivative of f.\"\n",
    "        return self.α * (k**(self.α - 1))\n",
    "\n",
    "    def u_prime(self, c):\n",
    "        \"Derivative of u.\"\n",
    "        return c**(-self.γ)\n",
    "\n",
    "    def u_prime_inv(c):\n",
    "        return c**(-1 / self.γ)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let’s create an instance:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "og_crra = OptimalGrowthModel_CRRA()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we solve and plot the policy:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "σ = solve_model_time_iter(og_crra, σ_init)\n",
    "\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "\n",
    "ax.plot(og.grid, σ, lw=2,\n",
    "        alpha=0.8, label='approximate policy function')\n",
    "\n",
    "ax.legend()\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "date": 1591697893.0653965,
  "filename": "coleman_policy_iter.rst",
  "kernelspec": {
   "display_name": "Python",
   "language": "python3",
   "name": "python3"
  },
  "title": "Optimal Growth III: Time Iteration"
 },
 "nbformat": 4,
 "nbformat_minor": 2
}