ifp.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "<a id='ifp'></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": [
    "# The Income Fluctuation Problem"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Contents\n",
    "\n",
    "- [The Income Fluctuation Problem](#The-Income-Fluctuation-Problem)  \n",
    "  - [Overview](#Overview)  \n",
    "  - [The Optimal Savings Problem](#The-Optimal-Savings-Problem)  \n",
    "  - [Computation](#Computation)  \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",
    "Next, we study an optimal savings problem for an infinitely lived consumer—the “common ancestor” described in [[LS18]](https://lectures.quantecon.org/zreferences.html#ljungqvist2012), section 1.3.\n",
    "\n",
    "This is an essential sub-problem for many representative macroeconomic models\n",
    "\n",
    "- [[Aiy94]](https://lectures.quantecon.org/zreferences.html#aiyagari1994)  \n",
    "- [[Hug93]](https://lectures.quantecon.org/zreferences.html#huggett1993)  \n",
    "- etc.  \n",
    "\n",
    "\n",
    "It is related to the decision problem in the [stochastic optimal growth\n",
    "model](https://lectures.quantecon.org/optgrowth.html) and yet differs in important ways.\n",
    "\n",
    "For example, the choice problem for the agent includes an additive income term that leads to an occasionally binding constraint.\n",
    "\n",
    "Our presentation of the model will be relatively brief.\n",
    "\n",
    "- For further details on economic intuition, implication and models, see [[LS18]](https://lectures.quantecon.org/zreferences.html#ljungqvist2012).  \n",
    "- Proofs of all mathematical results stated below can be found in <a href=_static/lecture_specific/ifp/pi2.pdf download>this paper</a>.  \n",
    "\n",
    "\n",
    "To solve the model we will use Euler equation based time iteration, similar to [this lecture](https://lectures.quantecon.org/coleman_policy_iter.html).\n",
    "\n",
    "This method turns out to be globally convergent under mild assumptions, even when utility is unbounded (both above and below).\n",
    "\n",
    "We’ll need the following imports:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from quantecon.optimize import brent_max, brentq\n",
    "from interpolation import interp\n",
    "from numba import njit\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "from quantecon import MarkovChain"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### References\n",
    "\n",
    "Other useful references include [[Dea91]](https://lectures.quantecon.org/zreferences.html#deaton1991), [[DH10]](https://lectures.quantecon.org/zreferences.html#denhaan2010), [[Kuh13]](https://lectures.quantecon.org/zreferences.html#kuhn2013), [[Rab02]](https://lectures.quantecon.org/zreferences.html#rabault2002),  [[Rei09]](https://lectures.quantecon.org/zreferences.html#reiter2009)  and [[SE77]](https://lectures.quantecon.org/zreferences.html#schechtmanescudero1977)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The Optimal Savings Problem\n",
    "\n",
    "\n",
    "<a id='index-1'></a>\n",
    "Let’s write down the model and then discuss how to solve it."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Set-Up\n",
    "\n",
    "Consider a household that chooses a state-contingent consumption plan $ \\{c_t\\}_{t \\geq 0} $ to maximize\n",
    "\n",
    "$$\n",
    "\\mathbb{E} \\, \\sum_{t=0}^{\\infty} \\beta^t u(c_t)\n",
    "$$\n",
    "\n",
    "subject to\n",
    "\n",
    "\n",
    "<a id='equation-eqst'></a>\n",
    "$$\n",
    "c_t + a_{t+1} \\leq  Ra_t  + z_t,\n",
    "\\qquad c_t \\geq 0,\n",
    "\\qquad a_t \\geq -b\n",
    "\\qquad t = 0, 1, \\ldots \\tag{1}\n",
    "$$\n",
    "\n",
    "Here\n",
    "\n",
    "- $ \\beta \\in (0,1) $ is the discount factor  \n",
    "- $ a_t $ is asset holdings at time $ t $, with ad-hoc borrowing constraint $ a_t \\geq -b $  \n",
    "- $ c_t $ is consumption  \n",
    "- $ z_t $ is non-capital income (wages, unemployment compensation, etc.)  \n",
    "- $ R := 1 + r $, where $ r > 0 $ is the interest rate on savings  \n",
    "\n",
    "\n",
    "Non-capital income $ \\{z_t\\} $ is assumed to be a Markov process taking values in $ Z\\subset (0,\\infty) $ with stochastic kernel $ \\Pi $.\n",
    "\n",
    "This means that $ \\Pi(z, B) $ is the probability that $ z_{t+1} \\in\n",
    "B $ given $ z_t = z $.\n",
    "\n",
    "The expectation of $ f(z_{t+1}) $ given $ z_t = z $ is written as\n",
    "\n",
    "$$\n",
    "\\int f( \\acute z) \\, \\Pi(z, d \\acute z)\n",
    "$$\n",
    "\n",
    "We further assume that\n",
    "\n",
    "1. $ r > 0 $ and $ \\beta R < 1 $  \n",
    "1. $ u $ is smooth, strictly increasing and strictly concave with $ \\lim_{c \\to 0} u'(c) = \\infty $ and $ \\lim_{c \\to \\infty} u'(c) = 0 $  \n",
    "\n",
    "\n",
    "The asset space is $ [-b, \\infty) $ and the state is the pair $ (a,z) \\in S := [-b,\\infty) \\times Z $.\n",
    "\n",
    "A *feasible consumption path* from $ (a,z) \\in S $ is a consumption\n",
    "sequence $ \\{c_t\\} $ such that $ \\{c_t\\} $ and its induced asset path $ \\{a_t\\} $ satisfy\n",
    "\n",
    "1. $ (a_0, z_0) = (a, z) $  \n",
    "1. the feasibility constraints in [(1)](#equation-eqst), and  \n",
    "1. measurability of $ c_t $ w.r.t. the filtration generated by $ \\{z_1, \\ldots, z_t\\} $  \n",
    "\n",
    "\n",
    "The meaning of the third point is just that consumption at time $ t $ can only be\n",
    "a function of outcomes that have already been observed."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Value Function and Euler Equation\n",
    "\n",
    "The *value function* $ V \\colon S \\to \\mathbb{R} $ is defined by\n",
    "\n",
    "\n",
    "<a id='equation-eqvf'></a>\n",
    "$$\n",
    "V(a, z) := \\sup \\, \\mathbb{E}\n",
    "\\left\\{\n",
    "\\sum_{t=0}^{\\infty} \\beta^t u(c_t)\n",
    "\\right\\} \\tag{2}\n",
    "$$\n",
    "\n",
    "where the supremum is overall feasible consumption paths from $ (a,z) $.\n",
    "\n",
    "An *optimal consumption path* from $ (a,z) $ is a feasible consumption path from $ (a,z) $ that attains the supremum in [(2)](#equation-eqvf).\n",
    "\n",
    "To pin down such paths we can use a version of the Euler equation, which in the present setting is\n",
    "\n",
    "\n",
    "<a id='equation-ee00'></a>\n",
    "$$\n",
    "u' (c_t)\n",
    "\\geq \\beta R \\,  \\mathbb{E}_t [ u'(c_{t+1}) ] \\tag{3}\n",
    "$$\n",
    "\n",
    "and\n",
    "\n",
    "\n",
    "<a id='equation-ee01'></a>\n",
    "$$\n",
    "u' (c_t) = \\beta R \\,  \\mathbb{E}_t [ u'(c_{t+1}) ]\n",
    "\\quad \\text{whenever }\n",
    "c_t < Ra_t + z_t + b \\tag{4}\n",
    "$$\n",
    "\n",
    "In essence, this says that the natural “arbitrage” relation $ u' (c_t) = \\beta R \\,  \\mathbb{E}_t [ u'(c_{t+1}) ] $ holds when the choice of current consumption is interior.\n",
    "\n",
    "Interiority means that $ c_t $ is strictly less than its upper bound $ Ra_t + z_t + b $.\n",
    "\n",
    "(The lower boundary case $ c_t = 0 $ never arises at the optimum because\n",
    "$ u'(0) = \\infty $)\n",
    "\n",
    "When $ c_t $ does hit the upper bound $ Ra_t + z_t + b $, the\n",
    "strict inequality $ u' (c_t) > \\beta R \\,  \\mathbb{E}_t [ u'(c_{t+1}) ] $\n",
    "can occur because $ c_t $ cannot increase sufficiently to attain equality.\n",
    "\n",
    "With some thought and effort, one can show that [(3)](#equation-ee00) and [(4)](#equation-ee01) are\n",
    "equivalent to\n",
    "\n",
    "\n",
    "<a id='equation-eqeul0'></a>\n",
    "$$\n",
    "u' (c_t)\n",
    "= \\max \\left\\{\n",
    "\\beta R \\,  \\mathbb{E}_t [ u'(c_{t+1}) ] \\,,\\;  u'(Ra_t + z_t + b)\n",
    "\\right\\} \\tag{5}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Optimality Results\n",
    "\n",
    "Given our assumptions, it is <a href=_static/lecture_specific/ifp/pi2.pdf download>known</a> that\n",
    "\n",
    "1. For each $ (a,z) \\in S $, a unique optimal consumption path from $ (a,z) $ exists  \n",
    "1. This path is the unique feasible path from $ (a,z) $ satisfying the\n",
    "  Euler equality [(5)](#equation-eqeul0) and the transversality condition  \n",
    "\n",
    "\n",
    "\n",
    "<a id='equation-eqtv'></a>\n",
    "$$\n",
    "\\lim_{t \\to \\infty} \\beta^t \\, \\mathbb{E} \\, [ u'(c_t) a_{t+1} ] = 0 \\tag{6}\n",
    "$$\n",
    "\n",
    "Moreover, there exists an *optimal consumption function*\n",
    "$ \\sigma^* \\colon S \\to [0, \\infty) $ such that the path from $ (a,z) $ generated by\n",
    "\n",
    "$$\n",
    "(a_0, z_0) = (a, z),\n",
    "\\quad\n",
    "z_{t+1} \\sim \\Pi(z_t, dy),\n",
    "\\quad\n",
    "c_t = \\sigma^*(a_t, z_t)\n",
    "\\quad \\text{and} \\quad\n",
    "a_{t+1} = R a_t + z_t - c_t\n",
    "$$\n",
    "\n",
    "satisfies both [(5)](#equation-eqeul0) and [(6)](#equation-eqtv), and hence is the unique optimal\n",
    "path from $ (a,z) $.\n",
    "\n",
    "In summary, to solve the optimization problem, we need to compute $ \\sigma^* $.\n",
    "\n",
    "\n",
    "<a id='ifp-computation'></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Computation\n",
    "\n",
    "\n",
    "<a id='index-2'></a>\n",
    "There are two standard ways to solve for $ \\sigma^* $\n",
    "\n",
    "1. Time iteration (TI) using the Euler equality  \n",
    "1. Value function iteration (VFI)  \n",
    "\n",
    "\n",
    "Let’s look at these in turn."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Time Iteration\n",
    "\n",
    "We can rewrite [(5)](#equation-eqeul0) to make it a statement about functions rather than\n",
    "random variables.\n",
    "\n",
    "In particular, consider the functional equation\n",
    "\n",
    "\n",
    "<a id='equation-eqeul1'></a>\n",
    "$$\n",
    "u' \\circ \\sigma \\, (a, z)\n",
    "= \\max \\left\\{\n",
    "\\gamma \\int u' \\circ \\sigma \\, \\{R a + z - c(a, z), \\, \\acute z\\}\n",
    "\\, \\Pi(z,d \\acute z)\n",
    "\\, , \\;\n",
    "     u'(Ra + z + b)\n",
    "     \\right\\} \\tag{7}\n",
    "$$\n",
    "\n",
    "where $ \\gamma := \\beta R $ and $ u' \\circ c(s) := u'(c(s)) $.\n",
    "\n",
    "Equation [(7)](#equation-eqeul1) is a functional equation in $ \\sigma $.\n",
    "\n",
    "In order to identify a solution, let $ \\mathscr{C} $ be the set of\n",
    "candidate consumption functions $ \\sigma \\colon S \\to \\mathbb R $ such that\n",
    "\n",
    "- each $ \\sigma \\in \\mathscr{C} $ is continuous and (weakly) increasing  \n",
    "- $ \\min Z \\leq c(a,z) \\leq Ra + z + b $ for all $ (a,z) \\in S $  \n",
    "\n",
    "\n",
    "In addition, let $ K \\colon \\mathscr{C} \\to \\mathscr{C} $ be defined as follows.\n",
    "\n",
    "For given $ \\sigma \\in \\mathscr{C} $, the value $ K \\sigma (a,z) $ is the unique $ t \\in J(a,z) $ that solves\n",
    "\n",
    "\n",
    "<a id='equation-eqsifc'></a>\n",
    "$$\n",
    "u'(t)\n",
    "= \\max \\left\\{\n",
    "\\gamma \\int u' \\circ \\sigma \\, \\{R a + z - t, \\, \\acute z\\}\n",
    "\\, \\Pi(z,d \\acute z)\n",
    "\\, , \\;\n",
    "     u'(Ra + z + b)\n",
    "     \\right\\} \\tag{8}\n",
    "$$\n",
    "\n",
    "where\n",
    "\n",
    "\n",
    "<a id='equation-eqbos'></a>\n",
    "$$\n",
    "J(a,z) := \\{t \\in \\mathbb{R} \\,:\\, \\min Z \\leq t \\leq Ra+ z + b\\} \\tag{9}\n",
    "$$\n",
    "\n",
    "We refer to $ K $ as Coleman’s policy function operator [[Col90]](https://lectures.quantecon.org/zreferences.html#coleman1990).\n",
    "\n",
    "It is <a href=_static/lecture_specific/ifp/pi2.pdf download>known</a> that\n",
    "\n",
    "- $ K $ is a contraction mapping on $ \\mathscr{C} $ under the metric  \n",
    "\n",
    "\n",
    "$$\n",
    "\\rho(c, d) := \\| \\, u' \\circ \\sigma_1 - u' \\circ \\sigma_2 \\, \\|\n",
    "    := \\sup_{s \\in S} | \\, u'(\\sigma_1(s))  - u'(\\sigma_2(s)) \\, |\n",
    " \\qquad \\quad (\\sigma_1, \\sigma_2 \\in \\mathscr{C})\n",
    "$$\n",
    "\n",
    "- The metric $ \\rho $ is complete on $ \\mathscr{C} $  \n",
    "- Convergence in $ \\rho $ implies uniform convergence on compacts  \n",
    "\n",
    "\n",
    "In consequence, $ K $ has a unique fixed point $ \\sigma^* \\in \\mathscr{C} $\n",
    "and $ K^n c \\to \\sigma^* $ as $ n \\to \\infty $ for any $ \\sigma \\in \\mathscr{C} $.\n",
    "\n",
    "By the definition of $ K $, the fixed points of $ K $ in $ \\mathscr{C} $ coincide with\n",
    "the solutions to [(7)](#equation-eqeul1) in $ \\mathscr{C} $.\n",
    "\n",
    "In particular, it <a href=_static/lecture_specific/ifp/pi2.pdf download>can be shown</a> that the path $ \\{c_t\\} $\n",
    "generated from $ (a_0,z_0) \\in S $ using policy function $ \\sigma^* $ is\n",
    "the unique optimal path from $ (a_0,z_0) \\in S $.\n",
    "\n",
    "**TL;DR** The unique optimal policy can be computed by picking any\n",
    "$ \\sigma \\in \\mathscr{C} $ and iterating with the operator $ K $ defined in [(8)](#equation-eqsifc)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Value Function Iteration\n",
    "\n",
    "The Bellman operator for this problem is given by\n",
    "\n",
    "\n",
    "<a id='equation-eqbop'></a>\n",
    "$$\n",
    "Tv(a, z)\n",
    "= \\max_{0 \\leq \\sigma \\leq Ra + z + b}\n",
    "\\left\\{\n",
    "    u(c) + \\beta \\int v(Ra + z - \\sigma, \\acute z) \\Pi(z, d \\acute z)\n",
    "\\right\\} \\tag{10}\n",
    "$$\n",
    "\n",
    "We have to be careful with VFI (i.e., iterating with\n",
    "$ T $) in this setting because $ u $ is not assumed to be bounded\n",
    "\n",
    "- In fact typically unbounded both above and below — e.g. $ u(c) = \\log c $.  \n",
    "- In which case, the standard DP theory does not apply.  \n",
    "- $ T^n v $ is not guaranteed to converge to the value function for arbitrary continuous bounded $ v $.  \n",
    "\n",
    "\n",
    "Nonetheless, we can always try the popular strategy “iterate and hope”.\n",
    "\n",
    "We can then check the outcome by comparing with that produced by TI.\n",
    "\n",
    "The latter is known to converge, as described above."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Implementation\n",
    "\n",
    "\n",
    "<a id='index-3'></a>\n",
    "First, we build a class called `ConsumerProblem` that stores the model primitives."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "class ConsumerProblem:\n",
    "    \"\"\"\n",
    "    A class that stores primitives for the income fluctuation problem. The\n",
    "    income process is assumed to be a finite state Markov chain.\n",
    "    \"\"\"\n",
    "    def __init__(self,\n",
    "                 r=0.01,                        # Interest rate\n",
    "                 β=0.96,                        # Discount factor\n",
    "                 Π=((0.6, 0.4),\n",
    "                    (0.05, 0.95)),              # Markov matrix for z_t\n",
    "                 z_vals=(0.5, 1.0),             # State space of z_t\n",
    "                 b=0,                           # Borrowing constraint\n",
    "                 grid_max=16,\n",
    "                 grid_size=50,\n",
    "                 u=np.log,                      # Utility function\n",
    "                 du=njit(lambda x: 1/x)):       # Derivative of utility\n",
    "\n",
    "        self.u, self.du = u, du\n",
    "        self.r, self.R = r, 1 + r\n",
    "        self.β, self.b = β, b\n",
    "        self.Π, self.z_vals = np.array(Π), tuple(z_vals)\n",
    "        self.asset_grid = np.linspace(-b, grid_max, grid_size)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The function `operator_factory` returns the operator `K` as specified above"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def operator_factory(cp):\n",
    "    \"\"\"\n",
    "    A function factory for building operator K.\n",
    "\n",
    "    Here cp is an instance of ConsumerProblem.\n",
    "    \"\"\"\n",
    "    # Simplify names, set up arrays\n",
    "    R, Π, β, u, b, du = cp.R, cp.Π, cp.β, cp.u, cp.b, cp.du\n",
    "    asset_grid, z_vals = cp.asset_grid, cp.z_vals\n",
    "    γ = R * β\n",
    "\n",
    "\n",
    "    @njit\n",
    "    def euler_diff(c, a, z, i_z, σ):\n",
    "        \"\"\"\n",
    "        The difference of the left-hand side and the right-hand side\n",
    "        of the Euler Equation.\n",
    "        \"\"\"\n",
    "        lhs = du(c)\n",
    "        expectation = 0\n",
    "        for i in range(len(z_vals)):\n",
    "            expectation += du(interp(asset_grid, σ[:, i], R * a + z - c)) \\\n",
    "                * Π[i_z, i]\n",
    "        rhs = max(γ * expectation, du(R * a + z + b))\n",
    "\n",
    "        return lhs - rhs\n",
    "\n",
    "    @njit\n",
    "    def K(σ):\n",
    "        \"\"\"\n",
    "        The operator K.\n",
    "\n",
    "        Iteration with this operator corresponds to time iteration on the\n",
    "        Euler equation.  Computes and returns the updated consumption policy\n",
    "        σ.  The array σ is replaced with a function cf that implements\n",
    "        univariate linear interpolation over the asset grid for each\n",
    "        possible value of z.\n",
    "        \"\"\"\n",
    "        σ_new = np.empty_like(σ)\n",
    "        for i_a in range(len(asset_grid)):\n",
    "            a = asset_grid[i_a]\n",
    "            for i_z in range(len(z_vals)):\n",
    "                z = z_vals[i_z]\n",
    "                c_star = brentq(euler_diff, 1e-8, R * a + z + b, \\\n",
    "                    args=(a, z, i_z, σ)).root\n",
    "                σ_new[i_a, i_z] = c_star\n",
    "\n",
    "        return σ_new\n",
    "\n",
    "    return K"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`K` uses linear interpolation along the asset grid to approximate the value and consumption functions.\n",
    "\n",
    "To solve for the optimal policy function, we will write a function `solve_model`\n",
    "to iterate and find the optimal $ \\sigma $."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def solve_model(cp,\n",
    "                tol=1e-4,\n",
    "                max_iter=1000,\n",
    "                verbose=True,\n",
    "                print_skip=25):\n",
    "\n",
    "    \"\"\"\n",
    "    Solves for the optimal policy using time iteration\n",
    "\n",
    "    * cp is an instance of ConsumerProblem\n",
    "    \"\"\"\n",
    "\n",
    "    u, β, b, R = cp.u, cp.β, cp.b, cp.R\n",
    "    asset_grid, z_vals = cp.asset_grid, cp.z_vals\n",
    "\n",
    "    # Initial guess of σ\n",
    "    σ = np.empty((len(asset_grid), len(z_vals)))\n",
    "    for i_a, a in enumerate(asset_grid):\n",
    "        for i_z, z in enumerate(z_vals):\n",
    "            c_max = R * a + z + b\n",
    "            σ[i_a, i_z] = c_max\n",
    "\n",
    "    K = operator_factory(cp)\n",
    "\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": [
    "Plotting the result using the default parameters of the `ConsumerProblem` class"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "cp = ConsumerProblem()\n",
    "σ_star = solve_model(cp)\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(10, 6))\n",
    "ax.plot(cp.asset_grid, σ_star[:, 0], label='$\\sigma^*$')\n",
    "ax.set(xlabel='asset level', ylabel='optimal consumption')\n",
    "ax.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The following exercises walk you through several applications where policy functions are computed."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercises"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 1\n",
    "\n",
    "Next, let’s consider how the interest rate affects consumption.\n",
    "\n",
    "Reproduce the following figure, which shows (approximately) optimal consumption policies for different interest rates\n",
    "\n",
    "<img src=\"https://s3-ap-southeast-2.amazonaws.com/lectures.quantecon.org/py/_static/lecture_specific/ifp/ifp_policies.png\" style=\"\">\n",
    "\n",
    "  \n",
    "- Other than `r`, all parameters are at their default values.  \n",
    "- `r` steps through `np.linspace(0, 0.04, 4)`.  \n",
    "- Consumption is plotted against assets for income shock fixed at the smallest value.  \n",
    "\n",
    "\n",
    "The figure shows that higher interest rates boost savings and hence suppress consumption."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 2\n",
    "\n",
    "Now let’s consider the long run asset levels held by households.\n",
    "\n",
    "We’ll take `r = 0.03` and otherwise use default parameters.\n",
    "\n",
    "The following figure is a 45 degree diagram showing the law of motion for assets when consumption is optimal"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "m = ConsumerProblem(r=0.03, grid_max=4)\n",
    "K = operator_factory(m)\n",
    "\n",
    "σ_star = solve_model(m, verbose=False)\n",
    "a = m.asset_grid\n",
    "R, z_vals = m.R, m.z_vals\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(10, 8))\n",
    "ax.plot(a, R * a + z_vals[0] - σ_star[:, 0], label='Low income')\n",
    "ax.plot(a, R * a + z_vals[1] - σ_star[:, 1], label='High income')\n",
    "ax.plot(a, a, 'k--')\n",
    "ax.set(xlabel='Current assets',\n",
    "       ylabel='Next period assets',\n",
    "       xlim=(0, 4), ylim=(0, 4))\n",
    "ax.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The blue line and orange line represent the function\n",
    "\n",
    "$$\n",
    "a' = h(a, z) := R a + z - \\sigma^*(a, z)\n",
    "$$\n",
    "\n",
    "when income $ z $ takes its high and low values respectively.\n",
    "\n",
    "The dashed line is the 45 degree line.\n",
    "\n",
    "We can see from the figure that the dynamics will be stable — assets do not\n",
    "diverge.\n",
    "\n",
    "In fact there is a unique stationary distribution of assets that we can calculate by simulation\n",
    "\n",
    "- Can be proved via theorem 2 of [[HP92]](https://lectures.quantecon.org/zreferences.html#hopenhaynprescott1992).  \n",
    "- Represents the long run dispersion of assets across households when households have idiosyncratic shocks.  \n",
    "\n",
    "\n",
    "Ergodicity is valid here, so stationary probabilities can be calculated by averaging over a single long time series.\n",
    "\n",
    "Hence to approximate the stationary distribution we can simulate a long time series for assets and histogram, as in the following figure\n",
    "\n",
    "<img src=\"https://s3-ap-southeast-2.amazonaws.com/lectures.quantecon.org/py/_static/lecture_specific/ifp/ifp_histogram.png\" style=\"\">\n",
    "\n",
    "  \n",
    "Your task is to replicate the figure\n",
    "\n",
    "- Parameters are as discussed above.  \n",
    "- The histogram in the figure used a single time series $ \\{a_t\\} $ of length 500,000.  \n",
    "- Given the length of this time series, the initial condition $ (a_0, z_0) $ will not matter.  \n",
    "- You might find it helpful to use the `MarkovChain` class from `quantecon`.  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 3\n",
    "\n",
    "Following on from exercises 1 and 2, let’s look at how savings and aggregate asset holdings vary with the interest rate\n",
    "\n",
    "- Note: [[LS18]](https://lectures.quantecon.org/zreferences.html#ljungqvist2012) section 18.6 can be consulted for more background on the topic treated in this exercise.  \n",
    "\n",
    "\n",
    "For a given parameterization of the model, the mean of the stationary distribution can be interpreted as aggregate capital in an economy with a unit mass of *ex-ante* identical households facing idiosyncratic shocks.\n",
    "\n",
    "Let’s look at how this measure of aggregate capital varies with the interest\n",
    "rate and borrowing constraint.\n",
    "\n",
    "The next figure plots aggregate capital against the interest rate for `b in (1, 3)`\n",
    "\n",
    "<img src=\"https://s3-ap-southeast-2.amazonaws.com/lectures.quantecon.org/py/_static/lecture_specific/ifp/ifp_agg_savings.png\" style=\"\">\n",
    "\n",
    "  \n",
    "As is traditional, the price (interest rate) is on the vertical axis.\n",
    "\n",
    "The horizontal axis is aggregate capital computed as the mean of the stationary distribution.\n",
    "\n",
    "Exercise 3 is to replicate the figure, making use of code from previous exercises.\n",
    "\n",
    "Try to explain why the measure of aggregate capital is equal to $ -b $\n",
    "when $ r=0 $ for both cases shown here."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Solutions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "r_vals = np.linspace(0, 0.04, 4)\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(10, 8))\n",
    "for r_val in r_vals:\n",
    "    cp = ConsumerProblem(r=r_val)\n",
    "    σ_star = solve_model(cp, verbose=False)\n",
    "    ax.plot(cp.asset_grid, σ_star[:, 0], label=f'$r = {r_val:.3f}$')\n",
    "\n",
    "ax.set(xlabel='asset level', ylabel='consumption (low income)')\n",
    "ax.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def compute_asset_series(cp, T=500000, verbose=False):\n",
    "    \"\"\"\n",
    "    Simulates a time series of length T for assets, given optimal\n",
    "    savings behavior.\n",
    "\n",
    "    cp is an instance of ConsumerProblem\n",
    "    \"\"\"\n",
    "    Π, z_vals, R = cp.Π, cp.z_vals, cp.R  # Simplify names\n",
    "    mc = MarkovChain(Π)\n",
    "    σ_star = solve_model(cp, verbose=False)\n",
    "    cf = lambda a, i_z: interp(cp.asset_grid, σ_star[:, i_z], a)\n",
    "    a = np.zeros(T+1)\n",
    "    z_seq = mc.simulate(T)\n",
    "    for t in range(T):\n",
    "        i_z = z_seq[t]\n",
    "        a[t+1] = R * a[t] + z_vals[i_z] - cf(a[t], i_z)\n",
    "    return a\n",
    "\n",
    "cp = ConsumerProblem(r=0.03, grid_max=4)\n",
    "a = compute_asset_series(cp)\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(10, 8))\n",
    "ax.hist(a, bins=20, alpha=0.5, density=True)\n",
    "ax.set(xlabel='assets', xlim=(-0.05, 0.75))\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "M = 25\n",
    "r_vals = np.linspace(0, 0.04, M)\n",
    "fig, ax = plt.subplots(figsize=(10, 8))\n",
    "\n",
    "for b in (1, 3):\n",
    "    asset_mean = []\n",
    "    for r_val in r_vals:\n",
    "        cp = ConsumerProblem(r=r_val, b=b)\n",
    "        mean = np.mean(compute_asset_series(cp, T=250000))\n",
    "        asset_mean.append(mean)\n",
    "    ax.plot(asset_mean, r_vals, label=f'$b = {b:d}$')\n",
    "    print(f\"Finished iteration b = {b:d}\")\n",
    "\n",
    "ax.set(xlabel='capital', ylabel='interest rate')\n",
    "ax.grid()\n",
    "ax.legend()\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "filename": "ifp.rst",
  "kernelspec": {
   "display_name": "Python",
   "language": "python3",
   "name": "python3"
  },
  "title": "The Income Fluctuation Problem"
 },
 "nbformat": 4,
 "nbformat_minor": 2
}