coleman_policy_iter.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "<a id='coleman-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: 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",
    "  - [Comparison with Value Function Iteration](#Comparison-with-Value-Function-Iteration)  \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 study](https://python.quantecon.org/optgrowth.html) of the stochastic optimal growth model.\n",
    "\n",
    "In that lecture, we solved the associated discounted dynamic programming 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 specific\n",
    "applications by deriving methods that are carefully tailored to the application at hand.\n",
    "\n",
    "The stochastic optimal growth model has plenty of structure to exploit for this purpose,\n",
    "especially when we adopt some concavity and smoothness assumptions over primitives.\n",
    "\n",
    "We’ll use this structure to obtain an **Euler equation**  based method that’s more efficient\n",
    "than value function iteration for this and some other closely related applications.\n",
    "\n",
    "In a [subsequent lecture](https://python.quantecon.org/egm_policy_iter.html), we’ll see that the numerical implementation\n",
    "part of the Euler equation method 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",
    "from quantecon.optimize import brentq"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The Euler Equation\n",
    "\n",
    "Let’s take the model set out in [the stochastic growth model lecture](https://python.quantecon.org/optgrowth.html) and add the assumptions that\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) might be valid, 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",
    "differentiate naively with respect to $ y $,  and then  evaluate at the optimum.\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 famous **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",
    "Tw(y) := \\max_{0 \\leq c \\leq y}\n",
    "\\left\\{\n",
    "    u(c) + \\beta \\int w(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 (i.e., $ 0 < \\sigma(y) < y $ for all strictly positive $ y $).\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.quantecon.org/zreferences.html#coleman1990) and [[Ref96]](https://python.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 Value Function Iteration\n",
    "\n",
    "How does Euler equation time iteration compare with value function iteration?\n",
    "\n",
    "Both can be used to compute the optimal policy, but is one faster or more\n",
    "accurate?\n",
    "\n",
    "There are two parts to this story.\n",
    "\n",
    "First, on a theoretical level, the two methods are essentially isomorphic.\n",
    "\n",
    "In particular, they converge at  the same rate.\n",
    "\n",
    "We’ll prove this in just a moment.\n",
    "\n",
    "The other side of the story is the accuracy of the numerical implementation.\n",
    "\n",
    "It turns out that, once we actually implement these two routines, time iteration is more accurate than value function iteration.\n",
    "\n",
    "More on this below."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Equivalent Dynamics\n",
    "\n",
    "Let’s talk about the theory first.\n",
    "\n",
    "To explain the connection between the two algorithms, it helps to understand\n",
    "the notion of equivalent dynamics.\n",
    "\n",
    "(This concept is very helpful in many other contexts as well)\n",
    "\n",
    "Suppose that we have a function $ g \\colon X \\to X $ where $ X $ is a given set.\n",
    "\n",
    "The pair $ (X, g) $ is sometimes called a **dynamical system** and we\n",
    "associate it with trajectories of the form\n",
    "\n",
    "$$\n",
    "x_{t+1} = g(x_t), \\qquad x_0 \\text{ given}\n",
    "$$\n",
    "\n",
    "Equivalently, $ x_t = g^t(x_0) $, where $ g $ is the $ t $-th\n",
    "composition of $ g $ with itself.\n",
    "\n",
    "Here’s the picture\n",
    "\n",
    "<img src=\"https://s3-ap-southeast-2.amazonaws.com/python.quantecon.org/_static/lecture_specific/coleman_policy_iter/col_pol_composition.png\" style=\"\">\n",
    "\n",
    "  \n",
    "Now let another function $ h \\colon Y \\to Y $ where $ Y $ is another set.\n",
    "\n",
    "Suppose further that\n",
    "\n",
    "- there exists a bijection $ \\tau $ from $ X $ to $ Y $  \n",
    "- the two functions **commute** under $ \\tau $, which is to say that\n",
    "  $ \\tau(g(x)) = h (\\tau(x)) $ for all $ x \\in X $  \n",
    "\n",
    "\n",
    "The last statement can be written more simply as\n",
    "\n",
    "$$\n",
    "\\tau \\circ g = h \\circ \\tau\n",
    "$$\n",
    "\n",
    "or, by applying $ \\tau^{-1} $ to both sides\n",
    "\n",
    "\n",
    "<a id='equation-cpi-ghcom'></a>\n",
    "$$\n",
    "g = \\tau^{-1} \\circ h \\circ \\tau \\tag{8}\n",
    "$$\n",
    "\n",
    "Here’s a commutative diagram that illustrates\n",
    "\n",
    "<img src=\"https://s3-ap-southeast-2.amazonaws.com/python.quantecon.org/_static/lecture_specific/coleman_policy_iter/col_pol_bij1.png\" style=\"\">\n",
    "\n",
    "  \n",
    "Here’s a similar figure that traces out the action of the maps on a point\n",
    "$ x \\in X $\n",
    "\n",
    "<img src=\"https://s3-ap-southeast-2.amazonaws.com/python.quantecon.org/_static/lecture_specific/coleman_policy_iter/col_pol_bij2.png\" style=\"\">\n",
    "\n",
    "  \n",
    "Now, it’s easy to check from [(8)](#equation-cpi-ghcom) that $ g^2 = \\tau^{-1} \\circ h^2 \\circ \\tau $ holds.\n",
    "\n",
    "In fact, if you like proofs by induction, you won’t have trouble showing that\n",
    "\n",
    "$$\n",
    "g^n = \\tau^{-1} \\circ h^n \\circ \\tau\n",
    "$$\n",
    "\n",
    "is valid for all $ n $.\n",
    "\n",
    "What does this tell us?\n",
    "\n",
    "It tells us that the following are equivalent\n",
    "\n",
    "- iterate $ n $ times with $ g $, starting at $ x $  \n",
    "- shift $ x $ to $ Y $ using $ \\tau $,  iterate $ n $ times with $ h $ starting at $ \\tau(x) $ and shift the result $ h^n(\\tau(x)) $ back to $ X $ using $ \\tau^{-1} $  \n",
    "\n",
    "\n",
    "We end up with exactly the same object."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Back to Economics\n",
    "\n",
    "Have you guessed where this is leading?\n",
    "\n",
    "What we’re going to show now is that the operators $ T $ and $ K $\n",
    "commute under a certain bijection.\n",
    "\n",
    "The implication is that they have exactly the same rate of convergence.\n",
    "\n",
    "To make life a little easier, we’ll assume in the following analysis (although not\n",
    "always in our applications) that $ u(0) = 0 $."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### A Bijection\n",
    "\n",
    "Let $ \\mathscr V $ be all strictly concave, continuously differentiable functions $ v $ mapping $ \\mathbb R_+ $ to itself and satisfying $ v(0) = 0 $ and $ v'(y) > u'(y) $ for all positive $ y $.\n",
    "\n",
    "For $ v \\in \\mathscr V $ let\n",
    "\n",
    "$$\n",
    "M v := h \\circ v' \\qquad \\text{where } h := (u')^{-1}\n",
    "$$\n",
    "\n",
    "Although we omit details, $ \\sigma := M v $ is actually the unique\n",
    "$ v $-greedy policy.\n",
    "\n",
    "- See proposition 12.1.18 of [EDTC](http://johnstachurski.net/edtc.html).  \n",
    "\n",
    "\n",
    "It turns out that $ M $ is a bijection from $ \\mathscr V $ to $ \\mathscr P $.\n",
    "\n",
    "A (solved) exercise below asks you to confirm this."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Commutative Operators\n",
    "\n",
    "It is an additional solved exercise (see below) to show that $ T $ and $ K $ commute under $ M $, in the sense that\n",
    "\n",
    "\n",
    "<a id='equation-cpi-ed-tk'></a>\n",
    "$$\n",
    "M \\circ T = K \\circ M \\tag{9}\n",
    "$$\n",
    "\n",
    "In view of the preceding discussion, this implies that\n",
    "\n",
    "$$\n",
    "T^n = M^{-1} \\circ K^n \\circ M\n",
    "$$\n",
    "\n",
    "Hence, $ T $ and $ K $ converge at exactly the same rate!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Implementation\n",
    "\n",
    "We’ve just shown that the operators $ T $ and $ K $ have the same rate of convergence.\n",
    "\n",
    "However, it turns out that, once numerical approximation is taken into account, significant differences arise.\n",
    "\n",
    "In particular, the image of policy functions under $ K $ can be calculated faster and with greater accuracy than the image of value functions under $ T $.\n",
    "\n",
    "Our intuition for this result is that\n",
    "\n",
    "- the Coleman-Reffett operator exploits more information because it uses first order and envelope conditions  \n",
    "- policy functions generally have less curvature than value functions, and hence admit more accurate approximations based on grid point information  \n",
    "\n",
    "\n",
    "First, we’ll store the parameters of the model in a class `OptimalGrowthModel`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "class OptimalGrowthModel:\n",
    "\n",
    "    def __init__(self,\n",
    "                 f,\n",
    "                 f_prime,\n",
    "                 u,\n",
    "                 u_prime,\n",
    "                 β=0.96,\n",
    "                 μ=0,\n",
    "                 s=0.1,\n",
    "                 grid_max=4,\n",
    "                 grid_size=200,\n",
    "                 shock_size=250):\n",
    "\n",
    "        self.β, self.μ, self.s = β, μ, s\n",
    "        self.f, self.u = f, u\n",
    "        self.f_prime, self.u_prime = f_prime, u_prime\n",
    "\n",
    "        self.grid = np.linspace(1e-5, grid_max, grid_size)  # Set up grid\n",
    "        # Store shocks\n",
    "        self.shocks = np.exp(μ + s * np.random.randn(shock_size))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here’s some code that returns the Coleman-Reffett operator, $ K $."
   ]
  },
  {
   "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",
    "    β = og.β\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, σ, y):\n",
    "        \"\"\"\n",
    "        The right hand side of the operator\n",
    "        \"\"\"\n",
    "        # First turn w into a function via interpolation\n",
    "        σ_func = lambda x: interp(grid, σ, x)\n",
    "        vals = u_prime(σ_func(f(y - c) * shocks)) * f_prime(y - c) * shocks\n",
    "        return u_prime(c) - β * np.mean(vals)\n",
    "\n",
    "    @njit(parallel=parallel_flag)\n",
    "    def K(σ):\n",
    "        \"\"\"\n",
    "        The Coleman-Reffett operator\n",
    "        \"\"\"\n",
    "        σ_new = np.empty_like(σ)\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=(σ, y))[0]\n",
    "            σ_new[i] = c_star\n",
    "\n",
    "        return σ_new\n",
    "\n",
    "    return K"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It has some similarities to the code for the Bellman operator in our [optimal growth lecture](https://python.quantecon.org/optgrowth.html).\n",
    "\n",
    "For example, it evaluates integrals by Monte Carlo and approximates functions using linear interpolation.\n",
    "\n",
    "Here’s that Bellman operator code again, which needs to be executed because we’ll use it in some tests below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false,
    "html-class": "collapse"
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from interpolation import interp\n",
    "from numba import njit, prange\n",
    "from quantecon.optimize.scalar_maximization import brent_max\n",
    "\n",
    "\n",
    "def operator_factory(og, parallel_flag=True):\n",
    "    \"\"\"\n",
    "    A function factory for building the Bellman operator, as well as\n",
    "    a function that computes greedy policies.\n",
    "\n",
    "    Here og is an instance of OptimalGrowthModel.\n",
    "    \"\"\"\n",
    "\n",
    "    f, u, β = og.f, og.u, og.β\n",
    "    grid, shocks = og.grid, og.shocks\n",
    "\n",
    "    @njit\n",
    "    def objective(c, v, y):\n",
    "        \"\"\"\n",
    "        The right-hand side of the Bellman equation\n",
    "        \"\"\"\n",
    "        # First turn v into a function via interpolation\n",
    "        v_func = lambda x: interp(grid, v, x)\n",
    "        return u(c) + β * np.mean(v_func(f(y - c) * shocks))\n",
    "\n",
    "    @njit(parallel=parallel_flag)\n",
    "    def T(v):\n",
    "        \"\"\"\n",
    "        The Bellman operator\n",
    "        \"\"\"\n",
    "        v_new = np.empty_like(v)\n",
    "        for i in prange(len(grid)):\n",
    "            y = grid[i]\n",
    "            # Solve for optimal v at y\n",
    "            v_max = brent_max(objective, 1e-10, y, args=(v, y))[1]\n",
    "            v_new[i] = v_max\n",
    "        return v_new\n",
    "\n",
    "    @njit\n",
    "    def get_greedy(v):\n",
    "        \"\"\"\n",
    "        Computes the v-greedy policy of a given function v\n",
    "        \"\"\"\n",
    "        σ = np.empty_like(v)\n",
    "        for i in range(len(grid)):\n",
    "            y = grid[i]\n",
    "            # Solve for optimal c at y\n",
    "            c_max = brent_max(objective, 1e-10, y, args=(v, y))[0]\n",
    "            σ[i] = c_max\n",
    "        return σ\n",
    "\n",
    "    return T, get_greedy"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Testing on the Log / Cobb–Douglas Case\n",
    "\n",
    "As we [did for value function iteration](https://python.quantecon.org/optgrowth.html), let’s start by\n",
    "testing our method in the presence of a model that does have an analytical\n",
    "solution.\n",
    "\n",
    "First, we generate an instance of `OptimalGrowthModel` and return the corresponding\n",
    "Coleman-Reffett operator."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "α = 0.3\n",
    "\n",
    "@njit\n",
    "def f(k):\n",
    "    \"Deterministic part of production function\"\n",
    "    return k**α\n",
    "\n",
    "@njit\n",
    "def f_prime(k):\n",
    "    return α * k**(α - 1)\n",
    "\n",
    "og = OptimalGrowthModel(f=f, f_prime=f_prime,\n",
    "                        u=np.log, u_prime=njit(lambda x: 1/x))\n",
    "\n",
    "K = time_operator_factory(og)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As a preliminary test, let’s see if $ K \\sigma^* = \\sigma^* $, as implied by the\n",
    "theory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "@njit\n",
    "def σ_star(y, α, β):\n",
    "    \"True optimal policy\"\n",
    "    return (1 - α * β) * y\n",
    "\n",
    "grid, β = og.grid, og.β\n",
    "σ_star_new = K(σ_star(grid, α, β))\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "ax.plot(grid, σ_star(grid, α, β), label=\"optimal policy $\\sigma^*$\")\n",
    "ax.plot(grid, σ_star_new, label=\"$K\\sigma^*$\")\n",
    "\n",
    "ax.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can’t really distinguish the two plots, so we are looking good, at least\n",
    "for this test.\n",
    "\n",
    "Next, let’s try iterating from an arbitrary initial condition and see if we\n",
    "converge towards $ \\sigma^* $.\n",
    "\n",
    "The initial condition we’ll use is the one that eats the whole pie: $ \\sigma(y) = y $."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "n = 15\n",
    "σ = grid.copy()  # Set initial condition\n",
    "fig, ax = plt.subplots(figsize=(9, 6))\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(σ)\n",
    "    ax.plot(grid, σ, color=plt.cm.jet(i / n), alpha=0.6)\n",
    "\n",
    "lb = 'true policy function $\\sigma^*$'\n",
    "ax.plot(grid, σ_star(grid, α, β), 'k-', alpha=0.8, label=lb)\n",
    "ax.legend()\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We see that the policy has converged nicely, in only a few steps.\n",
    "\n",
    "Now let’s compare the accuracy of iteration between the operators.\n",
    "\n",
    "We’ll generate\n",
    "\n",
    "1. $ K^n \\sigma $ where $ \\sigma(y) = y $  \n",
    "1. $ (M \\circ T^n \\circ M^{-1}) \\sigma $ where $ \\sigma(y) = y $  \n",
    "\n",
    "\n",
    "In each case, we’ll compare the resulting policy to $ \\sigma^* $.\n",
    "\n",
    "The theory on equivalent dynamics says we will get the same policy function\n",
    "and hence the same errors.\n",
    "\n",
    "But in fact we expect the first method to be more accurate for reasons\n",
    "discussed above"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "T, get_greedy = operator_factory(og)  # Return the Bellman operator\n",
    "\n",
    "σ = grid          # Set initial condition for σ\n",
    "v = og.u(grid)    # Set initial condition for v\n",
    "sim_length = 20\n",
    "\n",
    "for i in range(sim_length):\n",
    "    σ = K(σ)  # Time iteration\n",
    "    v = T(v)  # Value function iteration\n",
    "\n",
    "# Calculate difference with actual solution\n",
    "σ_error = σ_star(grid, α, β) - σ\n",
    "v_error = σ_star(grid, α, β) - get_greedy(v)\n",
    "\n",
    "plt.plot(grid, σ_error, alpha=0.6, label=\"policy iteration error\")\n",
    "plt.plot(grid, v_error, alpha=0.6, label=\"value iteration error\")\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As you can see, time iteration is much more accurate for a given\n",
    "number of iterations."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercises"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 1\n",
    "\n",
    "Show that [(9)](#equation-cpi-ed-tk) is valid.  In particular,\n",
    "\n",
    "- Let $ v $ be strictly concave and continuously differentiable on $ (0, \\infty) $.  \n",
    "- Fix $ y \\in (0, \\infty) $ and show that $ MTv(y) = KMv(y) $.  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 2\n",
    "\n",
    "Show that $ M $ is a bijection from $ \\mathscr V $ to $ \\mathscr P $."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 3\n",
    "\n",
    "Consider the same model as above but with the CRRA utility function\n",
    "\n",
    "$$\n",
    "u(c) = \\frac{c^{1 - \\gamma} - 1}{1 - \\gamma}\n",
    "$$\n",
    "\n",
    "Iterate 20 times with Bellman iteration and Euler equation time iteration\n",
    "\n",
    "- start time iteration from $ \\sigma(y) = y $  \n",
    "- start value function iteration from $ v(y) = u(y) $  \n",
    "- set $ \\gamma = 1.5 $  \n",
    "\n",
    "\n",
    "Compare the resulting policies and check that they are close."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 4\n",
    "\n",
    "Solve the above model as we did in [the previous lecture](https://python.quantecon.org/optgrowth.html) using\n",
    "the operators $ T $ and $ K $, and check the solutions are similiar by plotting."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Solutions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 1\n",
    "\n",
    "Let $ T, K, M, v $ and $ y $ be as stated in the exercise.\n",
    "\n",
    "Using the envelope theorem, one can show that $ (Tv)'(y) = u'(\\sigma(y)) $\n",
    "where $ \\sigma(y) $ solves\n",
    "\n",
    "\n",
    "<a id='equation-cpi-foo'></a>\n",
    "$$\n",
    "u'(\\sigma(y))\n",
    "= \\beta \\int v' (f(y - \\sigma(y)) z ) f'(y - \\sigma(y)) z \\phi(dz) \\tag{10}\n",
    "$$\n",
    "\n",
    "Hence $ MTv(y) = (u')^{-1} (u'(\\sigma(y))) = \\sigma(y) $.\n",
    "\n",
    "On the other hand, $ KMv(y) $ is the $ \\sigma(y) $ that solves\n",
    "\n",
    "$$\n",
    "\\begin{aligned}\n",
    "    u'(\\sigma(y))\n",
    "    & = \\beta \\int (u' \\circ (Mv)) (f(y - \\sigma(y)) z ) f'(y - \\sigma(y)) z \\phi(dz)\n",
    "    \\\\\n",
    "    & = \\beta \\int (u' \\circ ((u')^{-1} \\circ v'))\n",
    "        (f(y - \\sigma(y)) z ) f'(y - \\sigma(y)) z \\phi(dz)\n",
    "    \\\\\n",
    "    & = \\beta \\int v'(f(y - \\sigma(y)) z ) f'(y - \\sigma(y)) z \\phi(dz)\n",
    "\\end{aligned}\n",
    "$$\n",
    "\n",
    "We see that $ \\sigma(y) $ is the same in each case."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 2\n",
    "\n",
    "We need to show that $ M $ is a bijection from $ \\mathscr V $ to $ \\mathscr P $.\n",
    "\n",
    "To see this, first observe that, in view of our assumptions above, $ u' $ is a strictly decreasing continuous bijection from $ (0,\\infty) $ to itself.\n",
    "\n",
    "It [follows](https://math.stackexchange.com/questions/672174/continuity-of-an-inverse-function) that $ h $ has the same properties.\n",
    "\n",
    "Moreover, for fixed $ v \\in \\mathscr V $, the derivative $ v' $ is\n",
    "a continuous, strictly decreasing function.\n",
    "\n",
    "Hence, for fixed $ v \\in \\mathscr V $, the map $ M v = h \\circ v' $\n",
    "is strictly increasing and continuous, taking values in $ (0, \\infty) $.\n",
    "\n",
    "Moreover, interiority holds because $ v' $ strictly dominates $ u' $, implying that\n",
    "\n",
    "$$\n",
    "(M v)(y) = h(v'(y)) < h(u'(y)) = y\n",
    "$$\n",
    "\n",
    "In particular, $ \\sigma(y) := (Mv)(y) $ is an element of $ \\mathscr\n",
    "P $.\n",
    "\n",
    "To see that each $ \\sigma \\in \\mathscr P $ has a preimage $ v \\in \\mathscr V $ with $ Mv = \\sigma $, fix any $ \\sigma \\in \\mathscr P $.\n",
    "\n",
    "Let $ v(y) := \\int_0^y u'(\\sigma(x)) dx $ with $ v(0) = 0 $.\n",
    "\n",
    "With a small amount of effort, you will be able to show that $ v \\in \\mathscr V $ and $ Mv = \\sigma $.\n",
    "\n",
    "It’s also true that $ M $ is one-to-one on $ \\mathscr V $.\n",
    "\n",
    "To see this, suppose that  $ v $ and $ w $ are elements of $ \\mathscr V $\n",
    "satisfying $ Mv = Mw $.\n",
    "\n",
    "Then $ v(0) = w(0) = 0 $ and $ v' = w' $ on $ (0, \\infty) $.\n",
    "\n",
    "The fundamental theorem of calculus then implies that $ v = w $ on $ \\mathbb R_+ $."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 3\n",
    "\n",
    "Here’s the code, which will execute if you’ve run all the code above"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "γ = 1.5   # Preference parameter\n",
    "\n",
    "@njit\n",
    "def u(c):\n",
    "    return (c**(1 - γ) - 1) / (1 - γ)\n",
    "\n",
    "@njit\n",
    "def u_prime(c):\n",
    "    return c**(-γ)\n",
    "\n",
    "og = OptimalGrowthModel(f=f, f_prime=f_prime, u=u, u_prime=u_prime)\n",
    "\n",
    "T, get_greedy = operator_factory(og)\n",
    "K = time_operator_factory(og)\n",
    "\n",
    "σ = grid        # Initial condition for σ\n",
    "v = u(grid)     # Initial condition for v\n",
    "sim_length = 20\n",
    "\n",
    "for i in range(sim_length):\n",
    "    σ = K(σ)  # Time iteration\n",
    "    v = T(v)  # Value function iteration\n",
    "\n",
    "\n",
    "plt.plot(grid, σ, alpha=0.6, label=\"policy iteration\")\n",
    "plt.plot(grid, get_greedy(v), alpha=0.6, label=\"value iteration\")\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The policies are indeed close."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 4\n",
    "\n",
    "Here’s is the function we need to solve the model using value function iteration,\n",
    "copied from the previous lecture"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "\n",
    "def solve_model(og,\n",
    "                use_parallel=True,\n",
    "                tol=1e-4,\n",
    "                max_iter=1000,\n",
    "                verbose=True,\n",
    "                print_skip=25):\n",
    "\n",
    "    T, _ = operator_factory(og, parallel_flag=use_parallel)\n",
    "\n",
    "    # Set up loop\n",
    "    v = np.log(og.grid)  # Initial condition\n",
    "    i = 0\n",
    "    error = tol + 1\n",
    "\n",
    "    while i < max_iter and error > tol:\n",
    "        v_new = T(v)\n",
    "        error = np.max(np.abs(v - v_new))\n",
    "        i += 1\n",
    "        if verbose and i % print_skip == 0:\n",
    "            print(f\"Error at iteration {i} is {error}.\")\n",
    "        v = v_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 v_new"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Similarly, we can write a function that uses `K` to solve the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def solve_model_time(og,\n",
    "                     use_parallel=True,\n",
    "                     tol=1e-4,\n",
    "                     max_iter=1000,\n",
    "                     verbose=True,\n",
    "                     print_skip=25):\n",
    "\n",
    "    K = time_operator_factory(og, parallel_flag=use_parallel)\n",
    "\n",
    "    # Set up loop\n",
    "    σ = og.grid  # Initial condition\n",
    "    i = 0\n",
    "    error = tol + 1\n",
    "\n",
    "    while i < max_iter and error > tol:\n",
    "        σ_new = K(σ)\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": [
    "Solving both models and plotting"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "v_star = solve_model(og)\n",
    "σ_star = solve_model_time(og)\n",
    "\n",
    "plt.plot(grid, get_greedy(v_star), alpha=0.6, label='Bellman operator')\n",
    "plt.plot(grid, σ_star, alpha=0.6, label='Coleman-Reffett operator')\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Time iteration is numerically far more accurate for a given number of iterations."
   ]
  }
 ],
 "metadata": {
  "date": 1584334735.9181814,
  "filename": "coleman_policy_iter.rst",
  "kernelspec": {
   "display_name": "Python",
   "language": "python3",
   "name": "python3"
  },
  "title": "Optimal Growth III: Time Iteration"
 },
 "nbformat": 4,
 "nbformat_minor": 2
}