ifp_advanced.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "787c3f26",
   "metadata": {},
   "source": [
    "# The Income Fluctuation Problem II: Stochastic Returns on Assets"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "16928940",
   "metadata": {},
   "source": [
    "## Contents\n",
    "\n",
    "- [The Income Fluctuation Problem II: Stochastic Returns on Assets](#The-Income-Fluctuation-Problem-II:-Stochastic-Returns-on-Assets)  \n",
    "  - [Overview](#Overview)  \n",
    "  - [The Savings Problem](#The-Savings-Problem)  \n",
    "  - [Solution Algorithm](#Solution-Algorithm)  \n",
    "  - [Implementation](#Implementation)  \n",
    "  - [Exercises](#Exercises)  "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f596b751",
   "metadata": {},
   "source": [
    "In addition to what’s in Anaconda, this lecture will need the following libraries:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5d67a1c5",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "!pip install quantecon"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b777d6f9",
   "metadata": {},
   "source": [
    "## Overview\n",
    "\n",
    "In this lecture, we continue our study of the [income fluctuation problem](https://python.quantecon.org/ifp.html).\n",
    "\n",
    "While the interest rate was previously taken to be fixed, we now allow\n",
    "returns on assets to be state-dependent.\n",
    "\n",
    "This matches the fact that most households with a positive level of assets\n",
    "face some capital income risk.\n",
    "\n",
    "It has been argued that modeling capital income risk is essential for\n",
    "understanding the joint distribution of income and wealth (see, e.g.,\n",
    "[[Benhabib *et al.*, 2015](https://python.quantecon.org/zreferences.html#id134)] or [[Stachurski and Toda, 2019](https://python.quantecon.org/zreferences.html#id257)]).\n",
    "\n",
    "Theoretical properties of the household savings model presented here are\n",
    "analyzed in detail in [[Ma *et al.*, 2020](https://python.quantecon.org/zreferences.html#id256)].\n",
    "\n",
    "In terms of computation, we use a combination of time iteration and the\n",
    "endogenous grid method to solve the model quickly and accurately.\n",
    "\n",
    "We require the following imports:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2312a71e",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from numba import jit, float64\n",
    "from numba.experimental import jitclass\n",
    "from quantecon import MarkovChain"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "41472c20",
   "metadata": {},
   "source": [
    "## The Savings Problem\n",
    "\n",
    "In this section we review the household problem and optimality results."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "169eb400",
   "metadata": {},
   "source": [
    "### Set Up\n",
    "\n",
    "A household chooses a consumption-asset path $ \\{(c_t, a_t)\\} $ to\n",
    "maximize\n",
    "\n",
    "\n",
    "<a id='equation-trans-at'></a>\n",
    "$$\n",
    "\\mathbb E \\left\\{ \\sum_{t=0}^\\infty \\beta^t u(c_t) \\right\\} \\tag{44.1}\n",
    "$$\n",
    "\n",
    "subject to\n",
    "\n",
    "\n",
    "<a id='equation-trans-at2'></a>\n",
    "$$\n",
    "a_{t+1} = R_{t+1} (a_t - c_t) + Y_{t+1}\n",
    "\\; \\text{ and } \\;\n",
    "0 \\leq c_t \\leq a_t, \\tag{44.2}\n",
    "$$\n",
    "\n",
    "with initial condition $ (a_0, Z_0)=(a,z) $ treated as given.\n",
    "\n",
    "Note that $ \\{R_t\\}_{t \\geq 1} $, the gross rate of return on wealth, is allowed to be stochastic.\n",
    "\n",
    "The sequence $ \\{Y_t \\}_{t \\geq 1} $ is non-financial income.\n",
    "\n",
    "The stochastic components of the problem obey\n",
    "\n",
    "\n",
    "<a id='equation-eq-ry-func'></a>\n",
    "$$\n",
    "R_t = R(Z_t, \\zeta_t)\n",
    "  \\quad \\text{and} \\quad\n",
    "Y_t = Y(Z_t, \\eta_t), \\tag{44.3}\n",
    "$$\n",
    "\n",
    "where\n",
    "\n",
    "- the maps $ R $ and $ Y $ are time-invariant nonnegative functions,  \n",
    "- the innovation processes $ \\{\\zeta_t\\} $ and\n",
    "  $ \\{\\eta_t\\} $ are IID and independent of each other, and  \n",
    "- $ \\{Z_t\\}_{t \\geq 0} $ is an irreducible time-homogeneous Markov chain on a finite set $ \\mathsf Z $  \n",
    "\n",
    "\n",
    "Let $ P $ represent the Markov matrix for the chain $ \\{Z_t\\}_{t \\geq 0} $.\n",
    "\n",
    "Our assumptions on preferences are the same as our [previous lecture](https://python.quantecon.org/ifp.html) on the income fluctuation problem.\n",
    "\n",
    "As before, $ \\mathbb E_z \\hat X $ means expectation of next period value\n",
    "$ \\hat X $ given current value $ Z = z $."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "135c4584",
   "metadata": {},
   "source": [
    "### Assumptions\n",
    "\n",
    "We need restrictions to ensure that the objective [(44.1)](#equation-trans-at) is finite and\n",
    "the solution methods described below converge.\n",
    "\n",
    "We also need to ensure that the present discounted value of wealth\n",
    "does not grow too quickly.\n",
    "\n",
    "When $ \\{R_t\\} $ was constant we required that $ \\beta R < 1 $.\n",
    "\n",
    "Now it is stochastic, we require that\n",
    "\n",
    "\n",
    "<a id='equation-fpbc2'></a>\n",
    "$$\n",
    "\\beta G_R < 1,\n",
    "\\quad \\text{where} \\quad\n",
    "G_R := \\lim_{n \\to \\infty}\n",
    "\\left(\\mathbb E \\prod_{t=1}^n R_t \\right)^{1/n} \\tag{44.4}\n",
    "$$\n",
    "\n",
    "Notice that, when $ \\{R_t\\} $ takes some constant value $ R $, this\n",
    "reduces to the previous restriction $ \\beta R < 1 $\n",
    "\n",
    "The value $ G_R $ can be thought of as the long run (geometric) average\n",
    "gross rate of return.\n",
    "\n",
    "More intuition behind [(44.4)](#equation-fpbc2) is provided in [[Ma *et al.*, 2020](https://python.quantecon.org/zreferences.html#id256)].\n",
    "\n",
    "Discussion on how to check it is given below.\n",
    "\n",
    "Finally, we impose some routine technical restrictions on non-financial income.\n",
    "\n",
    "$$\n",
    "\\mathbb E \\, Y_t < \\infty \\text{ and } \\mathbb E \\, u'(Y_t) < \\infty\n",
    "\\label{a:y0}\n",
    "$$\n",
    "\n",
    "One relatively simple setting where all these restrictions are satisfied is\n",
    "the IID and CRRA environment of [[Benhabib *et al.*, 2015](https://python.quantecon.org/zreferences.html#id134)]."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d0538fdb",
   "metadata": {},
   "source": [
    "### Optimality\n",
    "\n",
    "Let the class of candidate consumption policies $ \\mathscr C $ be defined\n",
    "[as before](https://python.quantecon.org/ifp.html).\n",
    "\n",
    "In [[Ma *et al.*, 2020](https://python.quantecon.org/zreferences.html#id256)] it is shown that, under the stated assumptions,\n",
    "\n",
    "- any $ \\sigma \\in \\mathscr C $ satisfying the Euler equation is an\n",
    "  optimal policy and  \n",
    "- exactly one such policy exists in $ \\mathscr C $.  \n",
    "\n",
    "\n",
    "In the present setting, the Euler equation takes the form\n",
    "\n",
    "\n",
    "<a id='equation-ifpa-euler'></a>\n",
    "$$\n",
    "(u' \\circ \\sigma) (a, z) =\n",
    "\\max \\left\\{\n",
    "           \\beta \\, \\mathbb E_z \\,\\hat{R} \\,\n",
    "             (u' \\circ \\sigma)[\\hat{R}(a - \\sigma(a, z)) + \\hat{Y}, \\, \\hat{Z}],\n",
    "          \\, u'(a)\n",
    "       \\right\\} \\tag{44.5}\n",
    "$$\n",
    "\n",
    "(Intuition and derivation are similar to our [earlier lecture](https://python.quantecon.org/ifp.html) on\n",
    "the income fluctuation problem.)\n",
    "\n",
    "We again solve the Euler equation using time iteration, iterating with a\n",
    "Coleman–Reffett operator $ K $ defined to match the Euler equation\n",
    "[(44.5)](#equation-ifpa-euler)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7300fdc6",
   "metadata": {},
   "source": [
    "## Solution Algorithm\n",
    "\n",
    "\n",
    "<a id='index-1'></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "430b8f8d",
   "metadata": {},
   "source": [
    "### A Time Iteration Operator\n",
    "\n",
    "Our definition of the candidate class $ \\sigma \\in \\mathscr C $ of consumption\n",
    "policies is the same as in our [earlier lecture](https://python.quantecon.org/ifp.html) on the income\n",
    "fluctuation problem.\n",
    "\n",
    "For fixed $ \\sigma \\in \\mathscr C $ and $ (a,z) \\in \\mathbf S $, the value\n",
    "$ K\\sigma(a,z) $ of the function $ K\\sigma $ at $ (a,z) $ is defined as the\n",
    "$ \\xi \\in (0,a] $ that solves\n",
    "\n",
    "\n",
    "<a id='equation-k-opr'></a>\n",
    "$$\n",
    "u'(\\xi) =\n",
    "\\max \\left\\{\n",
    "          \\beta \\, \\mathbb E_z \\, \\hat{R} \\,\n",
    "             (u' \\circ \\sigma)[\\hat{R}(a - \\xi) + \\hat{Y}, \\, \\hat{Z}],\n",
    "          \\, u'(a)\n",
    "       \\right\\} \\tag{44.6}\n",
    "$$\n",
    "\n",
    "The idea behind $ K $ is that, as can be seen from the definitions,\n",
    "$ \\sigma \\in \\mathscr C $ satisfies the Euler equation\n",
    "if and only if $ K\\sigma(a, z) = \\sigma(a, z) $ for all $ (a, z) \\in\n",
    "\\mathbf S $.\n",
    "\n",
    "This means that fixed points of $ K $ in $ \\mathscr C $ and optimal\n",
    "consumption policies exactly coincide (see [[Ma *et al.*, 2020](https://python.quantecon.org/zreferences.html#id256)] for more details)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1c0cda79",
   "metadata": {},
   "source": [
    "### Convergence Properties\n",
    "\n",
    "As before, we pair $ \\mathscr C $ with the distance\n",
    "\n",
    "$$\n",
    "\\rho(c,d)\n",
    ":= \\sup_{(a,z) \\in \\mathbf S}\n",
    "          \\left|\n",
    "              \\left(u' \\circ c \\right)(a,z) -\n",
    "              \\left(u' \\circ d \\right)(a,z)\n",
    "          \\right|,\n",
    "$$\n",
    "\n",
    "It can be shown that\n",
    "\n",
    "1. $ (\\mathscr C, \\rho) $ is a complete metric space,  \n",
    "1. there exists an integer $ n $ such that $ K^n $ is a contraction\n",
    "  mapping on $ (\\mathscr C, \\rho) $, and  \n",
    "1. The unique fixed point of $ K $ in $ \\mathscr C $ is\n",
    "  the unique optimal policy in $ \\mathscr C $.  \n",
    "\n",
    "\n",
    "We now have a clear path to successfully approximating the optimal policy:\n",
    "choose some $ \\sigma \\in \\mathscr C $ and then iterate with $ K $ until\n",
    "convergence (as measured by the distance $ \\rho $)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "394f624c",
   "metadata": {},
   "source": [
    "### Using an Endogenous Grid\n",
    "\n",
    "In the study of that model we found that it was possible to further\n",
    "accelerate time iteration via the [endogenous grid method](https://python.quantecon.org/egm_policy_iter.html).\n",
    "\n",
    "We will use the same method here.\n",
    "\n",
    "The methodology is the same as it was for the optimal growth model, with the\n",
    "minor exception that we need to remember that consumption is not always\n",
    "interior.\n",
    "\n",
    "In particular, optimal consumption can be equal to assets when the level of\n",
    "assets is low."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1a11f2e0",
   "metadata": {},
   "source": [
    "#### Finding Optimal Consumption\n",
    "\n",
    "The endogenous grid method (EGM) calls for us to take a grid of *savings*\n",
    "values $ s_i $, where each such $ s $ is interpreted as $ s = a -\n",
    "c $.\n",
    "\n",
    "For the lowest grid point we take $ s_0 = 0 $.\n",
    "\n",
    "For the corresponding $ a_0, c_0 $ pair we have $ a_0 = c_0 $.\n",
    "\n",
    "This happens close to the origin, where assets are low and the household\n",
    "consumes all that it can.\n",
    "\n",
    "Although there are many solutions, the one we take is $ a_0 = c_0 = 0 $,\n",
    "which pins down the policy at the origin, aiding interpolation.\n",
    "\n",
    "For $ s > 0 $, we have, by definition, $ c < a $, and hence\n",
    "consumption is interior.\n",
    "\n",
    "Hence the max component of [(44.5)](#equation-ifpa-euler) drops out, and we solve for\n",
    "\n",
    "\n",
    "<a id='equation-eqsifc2'></a>\n",
    "$$\n",
    "c_i =\n",
    "(u')^{-1}\n",
    "\\left\\{\n",
    "    \\beta \\, \\mathbb E_z\n",
    "    \\hat R\n",
    "    (u' \\circ \\sigma) \\, [\\hat R s_i + \\hat Y, \\, \\hat Z]\n",
    "\\right\\} \\tag{44.7}\n",
    "$$\n",
    "\n",
    "at each $ s_i $."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aa4159b0",
   "metadata": {},
   "source": [
    "#### Iterating\n",
    "\n",
    "Once we have the pairs $ \\{s_i, c_i\\} $, the endogenous asset grid is\n",
    "obtained by $ a_i = c_i + s_i $.\n",
    "\n",
    "Also, we held $ z \\in \\mathsf Z $ in the discussion above so we can pair\n",
    "it with $ a_i $.\n",
    "\n",
    "An approximation of the policy $ (a, z) \\mapsto \\sigma(a, z) $ can be\n",
    "obtained by interpolating $ \\{a_i, c_i\\} $ at each $ z $.\n",
    "\n",
    "In what follows, we use linear interpolation."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "76e6069e",
   "metadata": {},
   "source": [
    "### Testing the Assumptions\n",
    "\n",
    "Convergence of time iteration is dependent on the condition $ \\beta G_R < 1 $ being satisfied.\n",
    "\n",
    "One can check this using the fact that $ G_R $ is equal to the spectral\n",
    "radius of the matrix $ L $ defined by\n",
    "\n",
    "$$\n",
    "L(z, \\hat z) := P(z, \\hat z) \\int R(\\hat z, x) \\phi(x) dx\n",
    "$$\n",
    "\n",
    "This identity is proved in [[Ma *et al.*, 2020](https://python.quantecon.org/zreferences.html#id256)], where $ \\phi $ is the\n",
    "density of the innovation $ \\zeta_t $ to returns on assets.\n",
    "\n",
    "(Remember that $ \\mathsf Z $ is a finite set, so this expression defines a matrix.)\n",
    "\n",
    "Checking the condition is even easier when $ \\{R_t\\} $ is IID.\n",
    "\n",
    "In that case, it is clear from the definition of $ G_R $ that $ G_R $\n",
    "is just $ \\mathbb E R_t $.\n",
    "\n",
    "We test the condition $ \\beta \\mathbb E R_t < 1 $ in the code below."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "22966ce8",
   "metadata": {},
   "source": [
    "## Implementation\n",
    "\n",
    "We will assume that $ R_t = \\exp(a_r \\zeta_t + b_r) $ where $ a_r, b_r $\n",
    "are constants and $ \\{ \\zeta_t\\} $ is IID standard normal.\n",
    "\n",
    "We allow labor income to be correlated, with\n",
    "\n",
    "$$\n",
    "Y_t = \\exp(a_y \\eta_t + Z_t b_y)\n",
    "$$\n",
    "\n",
    "where $ \\{ \\eta_t\\} $ is also IID standard normal\n",
    "and $ \\{ Z_t\\} $ is a Markov chain taking values in $ \\{0, 1\\} $."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "305f4700",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "ifp_data = [\n",
    "    ('γ', float64),              # utility parameter\n",
    "    ('β', float64),              # discount factor\n",
    "    ('P', float64[:, :]),        # transition probs for z_t\n",
    "    ('a_r', float64),            # scale parameter for R_t\n",
    "    ('b_r', float64),            # additive parameter for R_t\n",
    "    ('a_y', float64),            # scale parameter for Y_t\n",
    "    ('b_y', float64),            # additive parameter for Y_t\n",
    "    ('s_grid', float64[:]),      # Grid over savings\n",
    "    ('η_draws', float64[:]),     # Draws of innovation η for MC\n",
    "    ('ζ_draws', float64[:])      # Draws of innovation ζ for MC\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9a3edf6c",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "@jitclass(ifp_data)\n",
    "class IFP:\n",
    "    \"\"\"\n",
    "    A class that stores primitives for the income fluctuation\n",
    "    problem.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self,\n",
    "                 γ=1.5,\n",
    "                 β=0.96,\n",
    "                 P=np.array([(0.9, 0.1),\n",
    "                             (0.1, 0.9)]),\n",
    "                 a_r=0.1,\n",
    "                 b_r=0.0,\n",
    "                 a_y=0.2,\n",
    "                 b_y=0.5,\n",
    "                 shock_draw_size=50,\n",
    "                 grid_max=10,\n",
    "                 grid_size=100,\n",
    "                 seed=1234):\n",
    "\n",
    "        np.random.seed(seed)  # arbitrary seed\n",
    "\n",
    "        self.P, self.γ, self.β = P, γ, β\n",
    "        self.a_r, self.b_r, self.a_y, self.b_y = a_r, b_r, a_y, b_y\n",
    "        self.η_draws = np.random.randn(shock_draw_size)\n",
    "        self.ζ_draws = np.random.randn(shock_draw_size)\n",
    "        self.s_grid = np.linspace(0, grid_max, grid_size)\n",
    "\n",
    "        # Test stability assuming {R_t} is IID and adopts the lognormal\n",
    "        # specification given below.  The test is then β E R_t < 1.\n",
    "        ER = np.exp(b_r + a_r**2 / 2)\n",
    "        assert β * ER < 1, \"Stability condition failed.\"\n",
    "\n",
    "    # Marginal utility\n",
    "    def u_prime(self, c):\n",
    "        return c**(-self.γ)\n",
    "\n",
    "    # Inverse of marginal utility\n",
    "    def u_prime_inv(self, c):\n",
    "        return c**(-1/self.γ)\n",
    "\n",
    "    def R(self, z, ζ):\n",
    "        return np.exp(self.a_r * ζ + self.b_r)\n",
    "\n",
    "    def Y(self, z, η):\n",
    "        return np.exp(self.a_y * η + (z * self.b_y))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "17cbeca6",
   "metadata": {},
   "source": [
    "Here’s the Coleman-Reffett operator based on EGM:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2bab02d4",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "@jit\n",
    "def K(a_in, σ_in, ifp):\n",
    "    \"\"\"\n",
    "    The Coleman--Reffett operator for the income fluctuation problem,\n",
    "    using the endogenous grid method.\n",
    "\n",
    "        * ifp is an instance of IFP\n",
    "        * a_in[i, z] is an asset grid\n",
    "        * σ_in[i, z] is consumption at a_in[i, z]\n",
    "    \"\"\"\n",
    "\n",
    "    # Simplify names\n",
    "    u_prime, u_prime_inv = ifp.u_prime, ifp.u_prime_inv\n",
    "    R, Y, P, β = ifp.R, ifp.Y, ifp.P, ifp.β\n",
    "    s_grid, η_draws, ζ_draws = ifp.s_grid, ifp.η_draws, ifp.ζ_draws\n",
    "    n = len(P)\n",
    "\n",
    "    # Create consumption function by linear interpolation\n",
    "    σ = lambda a, z: np.interp(a, a_in[:, z], σ_in[:, z])\n",
    "\n",
    "    # Allocate memory\n",
    "    σ_out = np.empty_like(σ_in)\n",
    "\n",
    "    # Obtain c_i at each s_i, z, store in σ_out[i, z], computing\n",
    "    # the expectation term by Monte Carlo\n",
    "    for i, s in enumerate(s_grid):\n",
    "        for z in range(n):\n",
    "            # Compute expectation\n",
    "            Ez = 0.0\n",
    "            for z_hat in range(n):\n",
    "                for η in ifp.η_draws:\n",
    "                    for ζ in ifp.ζ_draws:\n",
    "                        R_hat = R(z_hat, ζ)\n",
    "                        Y_hat = Y(z_hat, η)\n",
    "                        U = u_prime(σ(R_hat * s + Y_hat, z_hat))\n",
    "                        Ez += R_hat * U * P[z, z_hat]\n",
    "            Ez = Ez / (len(η_draws) * len(ζ_draws))\n",
    "            σ_out[i, z] =  u_prime_inv(β * Ez)\n",
    "\n",
    "    # Calculate endogenous asset grid\n",
    "    a_out = np.empty_like(σ_out)\n",
    "    for z in range(n):\n",
    "        a_out[:, z] = s_grid + σ_out[:, z]\n",
    "\n",
    "    # Fixing a consumption-asset pair at (0, 0) improves interpolation\n",
    "    σ_out[0, :] = 0\n",
    "    a_out[0, :] = 0\n",
    "\n",
    "    return a_out, σ_out"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bec44cef",
   "metadata": {},
   "source": [
    "The next function solves for an approximation of the optimal consumption policy via time iteration."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ab1c3d72",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def solve_model_time_iter(model,        # Class with model information\n",
    "                          a_vec,        # Initial condition for assets\n",
    "                          σ_vec,        # Initial condition for consumption\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",
    "        a_new, σ_new = K(a_vec, σ_vec, model)\n",
    "        error = np.max(np.abs(σ_vec - σ_new))\n",
    "        i += 1\n",
    "        if verbose and i % print_skip == 0:\n",
    "            print(f\"Error at iteration {i} is {error}.\")\n",
    "        a_vec, σ_vec = np.copy(a_new), np.copy(σ_new)\n",
    "\n",
    "    if error > tol:\n",
    "        print(\"Failed to converge!\")\n",
    "    elif verbose:\n",
    "        print(f\"\\nConverged in {i} iterations.\")\n",
    "\n",
    "    return a_new, σ_new"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5ea3f6f6",
   "metadata": {},
   "source": [
    "Now we are ready to create an instance at the default parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "297b8f55",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "ifp = IFP()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7bf28618",
   "metadata": {},
   "source": [
    "Next we set up an initial condition, which corresponds to consuming all\n",
    "assets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6254d343",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "# Initial guess of σ = consume all assets\n",
    "k = len(ifp.s_grid)\n",
    "n = len(ifp.P)\n",
    "σ_init = np.empty((k, n))\n",
    "for z in range(n):\n",
    "    σ_init[:, z] = ifp.s_grid\n",
    "a_init = np.copy(σ_init)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f0891173",
   "metadata": {},
   "source": [
    "Let’s generate an approximation solution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18fc8b87",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "a_star, σ_star = solve_model_time_iter(ifp, a_init, σ_init, print_skip=5)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c1841c7b",
   "metadata": {},
   "source": [
    "Here’s a plot of the resulting consumption policy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "462acfdf",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "for z in range(len(ifp.P)):\n",
    "    ax.plot(a_star[:, z], σ_star[:, z], label=f\"consumption when $z={z}$\")\n",
    "\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2b7dcb17",
   "metadata": {},
   "source": [
    "Notice that we consume all assets in the lower range of the asset space.\n",
    "\n",
    "This is because we anticipate income $ Y_{t+1} $ tomorrow, which makes the need to save less urgent.\n",
    "\n",
    "Can you explain why consuming all assets ends earlier (for lower values of\n",
    "assets) when $ z=0 $?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2db937c4",
   "metadata": {},
   "source": [
    "### Law of Motion\n",
    "\n",
    "Let’s try to get some idea of what will happen to assets over the long run\n",
    "under this consumption policy.\n",
    "\n",
    "As with our [earlier lecture](https://python.quantecon.org/ifp.html) on the income fluctuation problem, we\n",
    "begin by producing a 45 degree diagram showing the law of motion for assets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ff57bfb3",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "# Good and bad state mean labor income\n",
    "Y_mean = [np.mean(ifp.Y(z, ifp.η_draws)) for z in (0, 1)]\n",
    "# Mean returns\n",
    "R_mean = np.mean(ifp.R(z, ifp.ζ_draws))\n",
    "\n",
    "a = a_star\n",
    "fig, ax = plt.subplots()\n",
    "for z, lb in zip((0, 1), ('bad state', 'good state')):\n",
    "    ax.plot(a[:, z], R_mean * (a[:, z] - σ_star[:, z]) + Y_mean[z] , label=lb)\n",
    "\n",
    "ax.plot(a[:, 0], a[:, 0], 'k--')\n",
    "ax.set(xlabel='current assets', ylabel='next period assets')\n",
    "\n",
    "ax.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9358ba79",
   "metadata": {},
   "source": [
    "The unbroken lines represent, for each $ z $, an average update function\n",
    "for assets, given by\n",
    "\n",
    "$$\n",
    "a \\mapsto \\bar R (a - \\sigma^*(a, z)) + \\bar Y(z)\n",
    "$$\n",
    "\n",
    "Here\n",
    "\n",
    "- $ \\bar R = \\mathbb E R_t $, which is mean returns and  \n",
    "- $ \\bar Y(z) = \\mathbb E_z Y(z, \\eta_t) $, which is mean labor income in state $ z $.  \n",
    "\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 even in the highest state."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "beac4811",
   "metadata": {},
   "source": [
    "## Exercises"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6fb4c087",
   "metadata": {},
   "source": [
    "## Exercise 44.1\n",
    "\n",
    "Let’s repeat our [earlier exercise](https://python.quantecon.org/ifp.html#ifp_ex2) on the long-run\n",
    "cross sectional distribution of assets.\n",
    "\n",
    "In that exercise, we used a relatively simple income fluctuation model.\n",
    "\n",
    "In the solution, we found the shape of the asset distribution to be unrealistic.\n",
    "\n",
    "In particular, we failed to match the long right tail of the wealth distribution.\n",
    "\n",
    "Your task is to try again, repeating the exercise, but now with our more sophisticated model.\n",
    "\n",
    "Use the default parameters."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "24d804a1",
   "metadata": {},
   "source": [
    "## Solution to[ Exercise 44.1](https://python.quantecon.org/#ifpa_ex1)\n",
    "\n",
    "First we write a function to compute a long asset series.\n",
    "\n",
    "Because we want to JIT-compile the function, we code the solution in a way\n",
    "that breaks some rules on good programming style.\n",
    "\n",
    "For example, we will pass in the solutions `a_star, σ_star` along with\n",
    "`ifp`, even though it would be more natural to just pass in `ifp` and then\n",
    "solve inside the function.\n",
    "\n",
    "The reason we do this is that `solve_model_time_iter` is not\n",
    "JIT-compiled."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1ea43ed5",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "@jit\n",
    "def compute_asset_series(ifp, a_star, σ_star, z_seq, T=500_000):\n",
    "    \"\"\"\n",
    "    Simulates a time series of length T for assets, given optimal\n",
    "    savings behavior.\n",
    "\n",
    "        * ifp is an instance of IFP\n",
    "        * a_star is the endogenous grid solution\n",
    "        * σ_star is optimal consumption on the grid\n",
    "        * z_seq is a time path for {Z_t}\n",
    "\n",
    "    \"\"\"\n",
    "\n",
    "    # Create consumption function by linear interpolation\n",
    "    σ = lambda a, z: np.interp(a, a_star[:, z], σ_star[:, z])\n",
    "\n",
    "    # Simulate the asset path\n",
    "    a = np.zeros(T+1)\n",
    "    for t in range(T):\n",
    "        z = z_seq[t]\n",
    "        ζ, η = np.random.randn(), np.random.randn()\n",
    "        R = ifp.R(z, ζ)\n",
    "        Y = ifp.Y(z, η)\n",
    "        a[t+1] = R * (a[t] - σ(a[t], z)) + Y\n",
    "    return a"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "811280c3",
   "metadata": {},
   "source": [
    "Now we call the function, generate the series and then histogram it, using the\n",
    "solutions computed above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "62861cbc",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "T = 1_000_000\n",
    "mc = MarkovChain(ifp.P)\n",
    "z_seq = mc.simulate(T, random_state=1234)\n",
    "\n",
    "a = compute_asset_series(ifp, a_star, σ_star, z_seq, T=T)\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "ax.hist(a, bins=40, alpha=0.5, density=True)\n",
    "ax.set(xlabel='assets')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "57b48531",
   "metadata": {},
   "source": [
    "Now we have managed to successfully replicate the long right tail of the\n",
    "wealth distribution.\n",
    "\n",
    "Here’s another view of this using a horizontal violin plot."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "603cd2e8",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "ax.violinplot(a, vert=False, showmedians=True)\n",
    "ax.set(xlabel='assets')\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "date": 1731921899.4091794,
  "filename": "ifp_advanced.md",
  "kernelspec": {
   "display_name": "Python",
   "language": "python3",
   "name": "python3"
  },
  "title": "The Income Fluctuation Problem II: Stochastic Returns on Assets"
 },
 "nbformat": 4,
 "nbformat_minor": 5
}