{ "cells": [ { "cell_type": "markdown", "id": "2c81f590", "metadata": {}, "source": [ "$$\n", "\\newcommand{\\argmax}{arg\\,max}\n", "\\newcommand{\\argmin}{arg\\,min}\n", "$$" ] }, { "cell_type": "markdown", "id": "e07f04c8", "metadata": {}, "source": [ "
\n", " \n", " \"QuantEcon\"\n", " \n", "
" ] }, { "cell_type": "markdown", "id": "df9fa656", "metadata": {}, "source": [ "# The Income Fluctuation Problem V: Stochastic Returns on Assets" ] }, { "cell_type": "markdown", "id": "f3997080", "metadata": {}, "source": [ "# GPU\n", "\n", "This lecture was built using a machine with access to a GPU — although it will also run without one.\n", "\n", "[Google Colab](https://colab.research.google.com/) has a free tier with GPUs\n", "that you can access as follows:\n", "\n", "1. Click on the “play” icon top right \n", "1. Select Colab \n", "1. Set the runtime environment to include a GPU " ] }, { "cell_type": "markdown", "id": "2edd9133", "metadata": {}, "source": [ "## Contents\n", "\n", "- [The Income Fluctuation Problem V: Stochastic Returns on Assets](#The-Income-Fluctuation-Problem-V:-Stochastic-Returns-on-Assets) \n", " - [Overview](#Overview) \n", " - [The Model](#The-Model) \n", " - [Solution Algorithm](#Solution-Algorithm) \n", " - [Implementation](#Implementation) \n", " - [Simulation](#Simulation) \n", " - [Wealth Inequality](#Wealth-Inequality) \n", " - [Exercises](#Exercises) " ] }, { "cell_type": "markdown", "id": "cb41d1cb", "metadata": {}, "source": [ "In addition to what’s in Anaconda, this lecture will need the following libraries:" ] }, { "cell_type": "code", "execution_count": null, "id": "6c87a68c", "metadata": { "hide-output": false }, "outputs": [], "source": [ "!pip install quantecon" ] }, { "cell_type": "markdown", "id": "6685cb27", "metadata": {}, "source": [ "## Overview\n", "\n", "In this lecture, we continue our study of the income fluctuation problem described in [The Income Fluctuation Problem III: The Endogenous Grid Method](https://python.quantecon.org/ifp_egm.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#id148)] or [[Stachurski and Toda, 2019](https://python.quantecon.org/zreferences.html#id271)]).\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#id270)].\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": "83b91a4d", "metadata": { "hide-output": false }, "outputs": [], "source": [ "import matplotlib.pyplot as plt\n", "import numpy as np\n", "import quantecon as qe\n", "import jax\n", "import jax.numpy as jnp\n", "from jax import vmap\n", "from typing import NamedTuple\n", "from functools import partial" ] }, { "cell_type": "markdown", "id": "266efab7", "metadata": {}, "source": [ "## The Model\n", "\n", "In this section we review the household problem and optimality results." ] }, { "cell_type": "markdown", "id": "0e0f03c7", "metadata": {}, "source": [ "### Set Up\n", "\n", "A household chooses a consumption-asset path $ \\{(c_t, a_t)\\} $ to\n", "maximize\n", "\n", "\n", "\n", "$$\n", "\\mathbb E \\left\\{ \\sum_{t=0}^\\infty \\beta^t u(c_t) \\right\\} \\tag{61.1}\n", "$$\n", "\n", "subject to\n", "\n", "\n", "\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{61.2}\n", "$$\n", "\n", "with initial condition $ (a_0, Z_0)=(a,z) $ treated as given.\n", "\n", "The only difference from [The Income Fluctuation Problem IV: Transient Income Shocks](https://python.quantecon.org/ifp_egm_transient_shocks.html) is that $ \\{R_t\\}_{t \\geq 1} $, the gross rate of return on wealth, is allowed to be stochastic.\n", "\n", "In particular, we assume that\n", "\n", "\n", "\n", "$$\n", "R_t = R(Z_t, \\zeta_t)\n", " \\quad \\text{and} \\quad\n", " Y_t = Y(Z_t, \\eta_t), \\tag{61.3}\n", "$$\n", "\n", "where\n", "\n", "- $ 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 a 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", "In what follows, $ \\mathbb E_z \\hat X $ means expectation of next period value $ \\hat X $ given current value $ Z = z $." ] }, { "cell_type": "markdown", "id": "703945ad", "metadata": {}, "source": [ "### Assumptions\n", "\n", "We need restrictions to ensure that the objective [(61.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", "Since it is now stochastic, we require (see [[Ma *et al.*, 2020](https://python.quantecon.org/zreferences.html#id270)]) that\n", "\n", "\n", "\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{61.4}\n", "$$\n", "\n", "The value $ G_R $ can be thought of as the long run (geometric) average\n", "gross rate of return.\n", "\n", "To simplify this lecture, we will *assume that the interest rate process is\n", "IID*.\n", "\n", "In that case, it is clear from the definition of $ G_R $ that $ G_R $ is just $ \\mathbb E R_t $.\n", "\n", "We test the condition $ \\beta \\mathbb E R_t < 1 $ in the code 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#id148)]." ] }, { "cell_type": "markdown", "id": "8a80160e", "metadata": {}, "source": [ "### Optimality\n", "\n", "Let the class of candidate consumption policies $ \\mathscr C $ be defined as in [The Income Fluctuation Problem III: The Endogenous Grid Method](https://python.quantecon.org/ifp_egm.html).\n", "\n", "In [[Ma *et al.*, 2020](https://python.quantecon.org/zreferences.html#id270)] 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", "\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{61.5}\n", "$$\n", "\n", "(Intuition and derivation are similar to [The Income Fluctuation Problem III: The Endogenous Grid Method](https://python.quantecon.org/ifp_egm.html).)\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", "[(61.5)](#equation-ifpa-euler)." ] }, { "cell_type": "markdown", "id": "8454cf22", "metadata": {}, "source": [ "## Solution Algorithm\n", "\n", "\n", "" ] }, { "cell_type": "markdown", "id": "cd5de651", "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 [The Income Fluctuation Problem III: The Endogenous Grid Method](https://python.quantecon.org/ifp_egm.html).\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", "\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{61.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#id270)] for more details)." ] }, { "cell_type": "markdown", "id": "978f2558", "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": "ee3c035e", "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/os_egm.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": "e4e52353", "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 [(61.5)](#equation-ifpa-euler) drops out, and we solve for\n", "\n", "\n", "\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{61.7}\n", "$$\n", "\n", "at each $ s_i $." ] }, { "cell_type": "markdown", "id": "b9fe9206", "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": "cf7d4314", "metadata": {}, "source": [ "## Implementation\n", "\n", "Here’s the model as a `NamedTuple`." ] }, { "cell_type": "code", "execution_count": null, "id": "4b0fdfd2", "metadata": { "hide-output": false }, "outputs": [], "source": [ "class IFP(NamedTuple):\n", " \"\"\"\n", " A NamedTuple that stores primitives for the income fluctuation\n", " problem, using JAX.\n", " \"\"\"\n", " γ: float\n", " β: float\n", " P: jnp.ndarray\n", " a_r: float\n", " b_r: float\n", " a_y: float\n", " b_y: float\n", " s_grid: jnp.ndarray\n", " η_draws: jnp.ndarray\n", " ζ_draws: jnp.ndarray\n", "\n", "\n", "def create_ifp(\n", " γ=1.5, # Utility parameter\n", " β=0.96, # Discount factor\n", " P=jnp.array([(0.9, 0.1), # Default Markov chain for Z\n", " (0.1, 0.9)]),\n", " a_r=0.16, # Volatility term in R shock\n", " b_r=0.0, # Mean shift R shock\n", " a_y=0.2, # Volatility term in Y shock\n", " b_y=0.5, # Mean shift Y shock\n", " shock_draw_size=100, # For Monte Carlo\n", " grid_max=100, # Exogenous grid max\n", " grid_size=100, # Exogenous grid size\n", " seed=1234 # Random seed\n", " ):\n", " \"\"\"\n", " Create an instance of IFP with the given parameters.\n", "\n", " \"\"\"\n", " # Test stability assuming {R_t} is IID and ln R ~ N(b_r, a_r)\n", " ER = np.exp(b_r + a_r**2 / 2)\n", " assert β * ER < 1, \"Stability condition failed.\"\n", "\n", " # Generate random draws using JAX\n", " key = jax.random.PRNGKey(seed)\n", " subkey1, subkey2 = jax.random.split(key)\n", " η_draws = jax.random.normal(subkey1, (shock_draw_size,))\n", " ζ_draws = jax.random.normal(subkey2, (shock_draw_size,))\n", " s_grid = jnp.linspace(0, grid_max, grid_size)\n", "\n", " return IFP(\n", " γ, β, P, a_r, b_r, a_y, b_y, s_grid, η_draws, ζ_draws\n", " )\n", "\n", "\n", "def u_prime(c, γ):\n", " \"\"\"Marginal utility\"\"\"\n", " return c**(-γ)\n", "\n", "def u_prime_inv(c, γ):\n", " \"\"\"Inverse of marginal utility\"\"\"\n", " return c**(-1/γ)\n", "\n", "def R(z, ζ, a_r, b_r):\n", " \"\"\"Gross return on assets\"\"\"\n", " return jnp.exp(a_r * ζ + b_r)\n", "\n", "def Y(z, η, a_y, b_y):\n", " \"\"\"Labor income\"\"\"\n", " return jnp.exp(a_y * η + (z * b_y))" ] }, { "cell_type": "markdown", "id": "81a073c4", "metadata": {}, "source": [ "Here’s the Coleman-Reffett operator using JAX:" ] }, { "cell_type": "code", "execution_count": null, "id": "2ad4596f", "metadata": { "hide-output": false }, "outputs": [], "source": [ "def K(\n", " a_in: jnp.array, # a_in[i, z] is an asset grid\n", " c_in: jnp.array, # c_in[i, z] = consumption at a_in[i, z]\n", " ifp: IFP\n", " ):\n", " \"\"\"\n", " The Coleman--Reffett operator for the income fluctuation problem,\n", " using the endogenous grid method with JAX.\n", "\n", " \"\"\"\n", "\n", " # Extract parameters from ifp\n", " γ, β, P, a_r, b_r, a_y, b_y, s_grid, η_draws, ζ_draws = ifp\n", " n = len(P)\n", "\n", " def compute_expectation(s, z):\n", " def inner_expectation(z_hat):\n", " def compute_term(η, ζ):\n", " R_hat = R(z_hat, ζ, a_r, b_r)\n", " Y_hat = Y(z_hat, η, a_y, b_y)\n", " a_val = R_hat * s + Y_hat\n", " # Interpolate consumption\n", " c_interp = jnp.interp(a_val, a_in[:, z_hat], c_in[:, z_hat])\n", " mu = u_prime(c_interp, γ)\n", " return R_hat * mu\n", " # Vectorize over all shock combinations\n", " η_grid, ζ_grid = jnp.meshgrid(η_draws, ζ_draws, indexing='ij')\n", " terms = vmap(vmap(compute_term))(η_grid, ζ_grid)\n", " return P[z, z_hat] * jnp.mean(terms)\n", " # Sum over z_hat states\n", " Ez = jnp.sum(vmap(inner_expectation)(jnp.arange(n)))\n", " return u_prime_inv(β * Ez, γ)\n", "\n", " # Vectorize over s_grid and z\n", " compute_exp_v1 = vmap(compute_expectation, in_axes=(None, 0))\n", " compute_exp_v2 = vmap(compute_exp_v1, in_axes=(0, None))\n", " c_out = compute_exp_v2(s_grid, jnp.arange(n))\n", " # Calculate endogenous asset grid\n", " a_out = s_grid[:, None] + c_out\n", " # Fix consumption-asset pair at (0, 0) \n", " c_out = c_out.at[0, :].set(0)\n", " a_out = a_out.at[0, :].set(0)\n", "\n", " return a_out, c_out" ] }, { "cell_type": "markdown", "id": "9a6d5812", "metadata": {}, "source": [ "The next function solves for an approximation of the optimal consumption policy\n", "via time iteration using JAX:" ] }, { "cell_type": "code", "execution_count": null, "id": "dc827135", "metadata": { "hide-output": false }, "outputs": [], "source": [ "@jax.jit\n", "def solve_model(\n", " ifp: IFP,\n", " c_init: jnp.ndarray, # Initial guess of σ on grid endogenous grid\n", " a_init: jnp.ndarray, # Initial endogenous grid\n", " tol: float = 1e-5,\n", " max_iter: int = 1000\n", " ) -> jnp.ndarray:\n", " \" Solve the model using time iteration with EGM. \"\n", "\n", " def condition(loop_state):\n", " c_in, a_in, i, error = loop_state\n", " return (error > tol) & (i < max_iter)\n", "\n", " def body(loop_state):\n", " c_in, a_in, i, error = loop_state\n", " c_out, a_out = K(c_in, a_in, ifp)\n", " error = jnp.max(jnp.abs(c_out - c_in))\n", " i += 1\n", " return c_out, a_out, i, error\n", "\n", " i, error = 0, tol + 1\n", " initial_state = (c_init, a_init, i, error)\n", " final_loop_state = jax.lax.while_loop(condition, body, initial_state)\n", " c_out, a_out, i, error = final_loop_state\n", "\n", " return c_out, a_out" ] }, { "cell_type": "markdown", "id": "536b9966", "metadata": {}, "source": [ "Now we can create an instance and solve the model using JAX:" ] }, { "cell_type": "code", "execution_count": null, "id": "0b624e28", "metadata": { "hide-output": false }, "outputs": [], "source": [ "ifp = create_ifp()" ] }, { "cell_type": "markdown", "id": "2463e1cb", "metadata": {}, "source": [ "Set up the initial condition:" ] }, { "cell_type": "code", "execution_count": null, "id": "c4f43bbe", "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 = jnp.empty((k, n))\n", "for z in range(n):\n", " σ_init = σ_init.at[:, z].set(ifp.s_grid)\n", "a_init = σ_init.copy()" ] }, { "cell_type": "markdown", "id": "0dc80aed", "metadata": {}, "source": [ "Let’s generate an approximation solution with JAX:" ] }, { "cell_type": "code", "execution_count": null, "id": "1e435b32", "metadata": { "hide-output": false }, "outputs": [], "source": [ "a_star, σ_star = solve_model(ifp, a_init, σ_init)" ] }, { "cell_type": "markdown", "id": "e3d01810", "metadata": {}, "source": [ "Let’s try it again with a timer." ] }, { "cell_type": "code", "execution_count": null, "id": "c7ff66f9", "metadata": { "hide-output": false }, "outputs": [], "source": [ "with qe.Timer(precision=8):\n", " a_star, σ_star = solve_model(ifp, a_init, σ_init)\n", " a_star.block_until_ready()" ] }, { "cell_type": "markdown", "id": "d7da23d1", "metadata": {}, "source": [ "## Simulation\n", "\n", "Let’s return to the default model and study the stationary distribution of assets.\n", "\n", "Our plan is to run a large number of households forward for $ T $ periods and then\n", "histogram the cross-sectional distribution of assets.\n", "\n", "Set `num_households=50_000, T=500`.\n", "\n", "First we write a function to run a single household forward in time and record\n", "the final value of assets.\n", "\n", "The function takes a solution pair `c_vec` and `a_vec`, understanding them\n", "as representing an optimal policy associated with a given model `ifp`" ] }, { "cell_type": "code", "execution_count": null, "id": "dbcffcae", "metadata": { "hide-output": false }, "outputs": [], "source": [ "def simulate_household(\n", " key, a_0, z_idx_0, c_vec, a_vec, ifp, T\n", " ):\n", " \"\"\"\n", " Simulates a single household for T periods to approximate the stationary\n", " distribution of assets.\n", "\n", " - key is the state of the random number generator\n", " - ifp is an instance of IFP\n", " - c_vec, a_vec are the optimal consumption policy, endogenous grid for ifp\n", "\n", " \"\"\"\n", " # Extract parameters from ifp\n", " γ, β, P, a_r, b_r, a_y, b_y, s_grid, η_draws, ζ_draws = ifp\n", " n_z = len(P)\n", "\n", " # Create interpolation function for consumption policy\n", " σ = lambda a, z_idx: jnp.interp(a, a_vec[:, z_idx], c_vec[:, z_idx])\n", "\n", " # Simulate forward T periods\n", " def update(t, state):\n", " a, z_idx = state\n", " # Draw next shock z' from P[z, z']\n", " current_key = jax.random.fold_in(key, 3*t)\n", " z_next_idx = jax.random.choice(current_key, n_z, p=P[z_idx]).astype(jnp.int32)\n", " # Draw η shock for income\n", " η_key = jax.random.fold_in(key, 3*t + 1)\n", " η = jax.random.normal(η_key)\n", " # Draw ζ shock for return\n", " ζ_key = jax.random.fold_in(key, 3*t + 2)\n", " ζ = jax.random.normal(ζ_key)\n", " # Compute stochastic return\n", " R_next = R(z_next_idx, ζ, a_r, b_r)\n", " # Compute income\n", " Y_next = Y(z_next_idx, η, a_y, b_y)\n", " # Update assets: a' = R' * (a - c) + Y'\n", " a_next = R_next * (a - σ(a, z_idx)) + Y_next\n", " # Return updated state\n", " return a_next, z_next_idx\n", "\n", " initial_state = a_0, z_idx_0\n", " final_state = jax.lax.fori_loop(0, T, update, initial_state)\n", " a_final, _ = final_state\n", " return a_final" ] }, { "cell_type": "markdown", "id": "d330f226", "metadata": {}, "source": [ "Now we write a function to simulate many households in parallel." ] }, { "cell_type": "code", "execution_count": null, "id": "f13c5bc1", "metadata": { "hide-output": false }, "outputs": [], "source": [ "@partial(jax.jit, static_argnums=(3, 4, 5))\n", "def compute_asset_stationary(\n", " c_vec, a_vec, ifp, num_households=50_000, T=500, seed=1234\n", " ):\n", " \"\"\"\n", " Simulates num_households households for T periods to approximate\n", " the stationary distribution of assets.\n", "\n", " Returns the final cross-section of asset holdings.\n", "\n", " - ifp is an instance of IFP\n", " - c_vec, a_vec are the optimal consumption policy and endogenous grid.\n", "\n", " \"\"\"\n", " # Extract parameters from ifp\n", " γ, β, P, a_r, b_r, a_y, b_y, s_grid, η_draws, ζ_draws = ifp\n", "\n", " # Start with assets = savings_grid_max / 2\n", " a_0_vector = jnp.full(num_households, s_grid[-1] / 2)\n", " # Initialize the exogenous state of each household\n", " z_idx_0_vector = jnp.zeros(num_households).astype(jnp.int32)\n", "\n", " # Vectorize over many households\n", " key = jax.random.PRNGKey(seed)\n", " keys = jax.random.split(key, num_households)\n", " # Vectorize simulate_household in (key, a_0, z_idx_0)\n", " sim_all_households = jax.vmap(\n", " simulate_household, in_axes=(0, 0, 0, None, None, None, None)\n", " )\n", " assets = sim_all_households(keys, a_0_vector, z_idx_0_vector, c_vec, a_vec, ifp, T)\n", "\n", " return jnp.array(assets)" ] }, { "cell_type": "markdown", "id": "27913ad4", "metadata": {}, "source": [ "We’ll need some inequality measures for visualization, so let’s define them first:" ] }, { "cell_type": "code", "execution_count": null, "id": "81069960", "metadata": { "hide-output": false }, "outputs": [], "source": [ "def gini_coefficient(x):\n", " \"\"\"\n", " Compute the Gini coefficient for array x.\n", "\n", " \"\"\"\n", " x = jnp.asarray(x)\n", " n = len(x)\n", " x_sorted = jnp.sort(x)\n", " # Compute Gini coefficient\n", " cumsum = jnp.cumsum(x_sorted)\n", " a = (2 * jnp.sum((jnp.arange(1, n+1)) * x_sorted)) / (n * cumsum[-1])\n", " return a - (n + 1) / n\n", "\n", "\n", "def top_share(\n", " x: jnp.array, # array of wealth values\n", " p: float=0.01 # fraction of top households (default 0.01 for top 1%)\n", " ):\n", " \"\"\"\n", " Compute the share of total wealth held by the top p fraction of households.\n", "\n", " \"\"\"\n", " x = jnp.asarray(x)\n", " x_sorted = jnp.sort(x)\n", " # Number of households in top p%\n", " n_top = int(jnp.ceil(len(x) * p))\n", " # Wealth held by top p%\n", " wealth_top = jnp.sum(x_sorted[-n_top:])\n", " # Total wealth\n", " wealth_total = jnp.sum(x_sorted)\n", " return wealth_top / wealth_total" ] }, { "cell_type": "markdown", "id": "79a4b3ce", "metadata": {}, "source": [ "Now we call the function, generate the asset distribution and visualize it:" ] }, { "cell_type": "code", "execution_count": null, "id": "5c528e13", "metadata": { "hide-output": false }, "outputs": [], "source": [ "ifp = create_ifp()\n", "# Extract parameters for initialization\n", "s_grid = ifp.s_grid\n", "n_z = len(ifp.P)\n", "a_init = s_grid[:, None] * jnp.ones(n_z)\n", "c_init = a_init\n", "a_vec, c_vec = solve_model(ifp, a_init, c_init)\n", "assets = compute_asset_stationary(c_vec, a_vec, ifp, num_households=200_000)\n", "\n", "# Compute Gini coefficient for the plot\n", "gini_plot = gini_coefficient(assets)\n", "\n", "# Plot histogram of log wealth\n", "fig, ax = plt.subplots(figsize=(10, 6))\n", "ax.hist(jnp.log(assets), bins=40, alpha=0.5, density=True)\n", "ax.set(xlabel='log assets', ylabel='density', title=\"Wealth Distribution\")\n", "plt.tight_layout()\n", "plt.show()" ] }, { "cell_type": "markdown", "id": "2cac7860", "metadata": {}, "source": [ "The histogram shows the distribution of log wealth.\n", "\n", "Bearing in mind that we are looking at log values, the histogram suggests\n", "a long right tail of the distribution.\n", "\n", "Below we examine this in more detail." ] }, { "cell_type": "markdown", "id": "cd6e6fcc", "metadata": {}, "source": [ "## Wealth Inequality\n", "\n", "Lets’ look at wealth inequality by computing some standard measures of this phenomenon.\n", "\n", "We will also examine how inequality varies with the interest rate." ] }, { "cell_type": "markdown", "id": "5f3fecea", "metadata": {}, "source": [ "### Measuring Inequality\n", "\n", "Let’s print the Gini coefficient and the top 1% wealth share from our simulation:" ] }, { "cell_type": "code", "execution_count": null, "id": "37d1dea5", "metadata": { "hide-output": false }, "outputs": [], "source": [ "gini = gini_coefficient(assets)\n", "top1 = top_share(assets, p=0.01)\n", "\n", "print(f\"Gini coefficient: {gini:.4f}\")\n", "print(f\"Top 1% wealth share: {top1:.4f}\")" ] }, { "cell_type": "markdown", "id": "5ea404f9", "metadata": {}, "source": [ "Recent numbers suggest that\n", "\n", "- the Gini coefficient for wealth in the US is around 0.8 \n", "- the top 1% wealth share is over 0.3 \n", "\n", "\n", "Our model with stochastic returns generates a Gini coefficient close to the\n", "empirical value, demonstrating that capital income risk is an important factor\n", "in wealth inequality.\n", "\n", "The top 1% wealth share is, however, too large.\n", "\n", "Our model needs proper calibration and additional work – we set these tasks aside for now." ] }, { "cell_type": "markdown", "id": "36ab65f8", "metadata": {}, "source": [ "## Exercises" ] }, { "cell_type": "markdown", "id": "59ec7f37", "metadata": {}, "source": [ "## Exercise 61.1\n", "\n", "Plot how the Gini coefficient varies with the volatility of returns on assets.\n", "\n", "Specifically, compute the Gini coefficient for values of `a_r` ranging from 0.10 to 0.16 (use at least 5 different values) and plot the results.\n", "\n", "What does this tell you about the relationship between capital income risk and wealth inequality?" ] }, { "cell_type": "markdown", "id": "2ca4d411", "metadata": {}, "source": [ "## Solution\n", "\n", "We loop over different values of `a_r`, solve the model for each, simulate the wealth distribution, and compute the Gini coefficient." ] }, { "cell_type": "code", "execution_count": null, "id": "f149c626", "metadata": { "hide-output": false }, "outputs": [], "source": [ "# Range of a_r values to explore\n", "a_r_vals = np.linspace(0.10, 0.16, 5)\n", "gini_vals = []\n", "\n", "print(\"Computing Gini coefficients for different return volatilities...\\n\")\n", "\n", "for a_r in a_r_vals:\n", " print(f\"a_r = {a_r:.3f}...\", end=\" \")\n", "\n", " # Create model with this a_r value\n", " ifp_temp = create_ifp(a_r=a_r, grid_max=100)\n", "\n", " # Solve the model\n", " s_grid_temp = ifp_temp.s_grid\n", " n_z_temp = len(ifp_temp.P)\n", " a_init_temp = s_grid_temp[:, None] * jnp.ones(n_z_temp)\n", " c_init_temp = a_init_temp\n", " a_vec_temp, c_vec_temp = solve_model(\n", " ifp_temp, a_init_temp, c_init_temp\n", " )\n", "\n", " # Simulate households\n", " assets_temp = compute_asset_stationary(\n", " c_vec_temp, a_vec_temp, ifp_temp, num_households=200_000\n", " )\n", "\n", " # Compute Gini coefficient\n", " gini_temp = gini_coefficient(assets_temp)\n", " gini_vals.append(gini_temp)\n", " print(f\"Gini = {gini_temp:.4f}\")\n", "\n", "# Plot the results\n", "fig, ax = plt.subplots(figsize=(10, 6))\n", "ax.plot(a_r_vals, gini_vals, 'o-', linewidth=2, markersize=8)\n", "ax.set(xlabel='Return volatility (a_r)',\n", " ylabel='Gini coefficient',\n", " title='Wealth Inequality vs Return Volatility')\n", "ax.axhline(y=0.8, color='k', linestyle='--', linewidth=1,\n", " label='Empirical US Gini (~0.8)')\n", "ax.legend()\n", "plt.tight_layout()\n", "plt.show()" ] }, { "cell_type": "markdown", "id": "bfbb0cd2", "metadata": {}, "source": [ "The plot shows that wealth inequality (measured by the Gini coefficient) increases with return volatility.\n", "\n", "This demonstrates that capital income risk is a key driver of wealth inequality.\n", "\n", "When returns are more volatile, lucky households who experience sequences of\n", "high returns accumulate substantially more wealth than unlucky households,\n", "leading to greater inequality in the wealth distribution." ] }, { "cell_type": "markdown", "id": "0dd56fa0", "metadata": {}, "source": [ "## Exercise 61.2\n", "\n", "Plot how the Gini coefficient varies with the volatility of labor income.\n", "\n", "Specifically, compute the Gini coefficient for values of `a_y` ranging from\n", "0.125 to 0.20 and plot the results. Set `a_r=0.10` for this exercise.\n", "\n", "What does this tell you about the relationship between labor income risk and\n", "wealth inequality? Can we achieve the same rise in inequality by varying labor\n", "income volatility as we can by varying return volatility?" ] }, { "cell_type": "markdown", "id": "6e04963c", "metadata": {}, "source": [ "## Solution\n", "\n", "We loop over different values of `a_y`, solve the model for each, simulate the wealth distribution, and compute the Gini coefficient." ] }, { "cell_type": "code", "execution_count": null, "id": "54e89cad", "metadata": { "hide-output": false }, "outputs": [], "source": [ "# Range of a_y values to explore\n", "a_y_vals = np.linspace(0.125, 0.20, 5)\n", "gini_vals_y = []\n", "\n", "print(\"Computing Gini coefficients for different labor income volatilities...\\n\")\n", "\n", "for a_y in a_y_vals:\n", " print(f\"a_y = {a_y:.3f}...\", end=\" \")\n", "\n", " # Create model with this a_y value and a_r=0.10\n", " ifp_temp = create_ifp(a_y=a_y, a_r=0.10, grid_max=100)\n", "\n", " # Solve the model\n", " s_grid_temp = ifp_temp.s_grid\n", " n_z_temp = len(ifp_temp.P)\n", " a_init_temp = s_grid_temp[:, None] * jnp.ones(n_z_temp)\n", " c_init_temp = a_init_temp\n", " a_vec_temp, c_vec_temp = solve_model(\n", " ifp_temp, a_init_temp, c_init_temp\n", " )\n", "\n", " # Simulate households\n", " assets_temp = compute_asset_stationary(\n", " c_vec_temp, a_vec_temp, ifp_temp, num_households=200_000\n", " )\n", "\n", " # Compute Gini coefficient\n", " gini_temp = gini_coefficient(assets_temp)\n", " gini_vals_y.append(gini_temp)\n", " print(f\"Gini = {gini_temp:.4f}\")\n", "\n", "# Plot the results\n", "fig, ax = plt.subplots(figsize=(10, 6))\n", "ax.plot(a_y_vals, gini_vals_y, 'o-', linewidth=2, markersize=8, color='green')\n", "ax.set(xlabel='Labor income volatility (a_y)',\n", " ylabel='Gini coefficient',\n", " title='Wealth Inequality vs Labor Income Volatility')\n", "ax.axhline(y=0.8, color='k', linestyle='--', linewidth=1,\n", " label='Empirical US Gini (~0.8)')\n", "ax.legend()\n", "plt.tight_layout()\n", "plt.show()" ] }, { "cell_type": "markdown", "id": "5228b1a1", "metadata": {}, "source": [ "The plot shows that wealth inequality increases with labor income volatility, but the effect is much weaker than the effect of return volatility.\n", "\n", "Comparing the two exercises:\n", "\n", "- When return volatility (`a_r`) varies from 0.10 to 0.16, the Gini coefficient rises dramatically from around 0.20 to 0.79 \n", "- When labor income volatility (`a_y`) varies from 0.125 to 0.20, a similar amount in percentage terms, the Gini coefficient increases but by a much smaller amount \n", "\n", "\n", "This suggests that capital income risk is a more important driver of wealth inequality than labor income risk.\n", "\n", "The intuition is that wealth accumulation compounds over time: households who\n", "experience favorable returns on their assets can reinvest those returns, leading\n", "to exponential growth.\n", "\n", "In contrast, labor income shocks, while they affect current consumption and\n", "savings, do not have the same compounding effect on wealth accumulation." ] } ], "metadata": { "date": 1770028417.9397175, "filename": "ifp_advanced.md", "kernelspec": { "display_name": "Python", "language": "python3", "name": "python3" }, "title": "The Income Fluctuation Problem V: Stochastic Returns on Assets" }, "nbformat": 4, "nbformat_minor": 5 }