{ "cells": [ { "cell_type": "markdown", "id": "e273d424", "metadata": {}, "source": [ "$$\n", "\\newcommand{\\argmax}{arg\\,max}\n", "\\newcommand{\\argmin}{arg\\,min}\n", "$$" ] }, { "cell_type": "markdown", "id": "caeb8698", "metadata": {}, "source": [ "\n", "\n", "
\n", " \n", " \"QuantEcon\"\n", " \n", "
" ] }, { "cell_type": "markdown", "id": "ee674936", "metadata": {}, "source": [ "# The Aiyagari Model" ] }, { "cell_type": "markdown", "id": "38e8ea0c", "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": "02d537da", "metadata": {}, "source": [ "## Contents\n", "\n", "- [The Aiyagari Model](#The-Aiyagari-Model) \n", " - [Overview](#Overview) \n", " - [The Economy](#The-Economy) \n", " - [Implementation](#Implementation) \n", " - [Exercises](#Exercises) " ] }, { "cell_type": "markdown", "id": "0f607af8", "metadata": {}, "source": [ "In addition to what’s included in base Anaconda, we need to install JAX" ] }, { "cell_type": "code", "execution_count": null, "id": "c0151e56", "metadata": { "hide-output": false }, "outputs": [], "source": [ "!pip install quantecon jax" ] }, { "cell_type": "markdown", "id": "94515edc", "metadata": {}, "source": [ "## Overview\n", "\n", "In this lecture, we describe the structure of a class of models that build on work by Truman Bewley [[Bewley, 1977](https://python.quantecon.org/zreferences.html#id191)].\n", "\n", "We begin by discussing an example of a Bewley model due to Rao Aiyagari [[Aiyagari, 1994](https://python.quantecon.org/zreferences.html#id155)].\n", "\n", "The model features\n", "\n", "- heterogeneous agents \n", "- a single exogenous vehicle for borrowing and lending \n", "- limits on amounts individual agents may borrow \n", "\n", "\n", "The Aiyagari model has been used to investigate many topics, including\n", "\n", "- precautionary savings and the effect of liquidity constraints [[Aiyagari, 1994](https://python.quantecon.org/zreferences.html#id155)] \n", "- risk sharing and asset pricing [[Heaton and Lucas, 1996](https://python.quantecon.org/zreferences.html#id147)] \n", "- the shape of the wealth distribution [[Benhabib *et al.*, 2015](https://python.quantecon.org/zreferences.html#id148)] \n", "- etc., etc., etc. " ] }, { "cell_type": "markdown", "id": "8dcf953a", "metadata": {}, "source": [ "### Preliminaries\n", "\n", "We use the following imports:" ] }, { "cell_type": "code", "execution_count": null, "id": "5f9a6c38", "metadata": { "hide-output": false }, "outputs": [], "source": [ "import quantecon as qe\n", "import matplotlib.pyplot as plt\n", "import jax\n", "import jax.numpy as jnp\n", "from typing import NamedTuple\n", "from scipy.optimize import bisect" ] }, { "cell_type": "markdown", "id": "9ba9ba39", "metadata": {}, "source": [ "We will use 64-bit floats with JAX in order to increase precision." ] }, { "cell_type": "code", "execution_count": null, "id": "0239737d", "metadata": { "hide-output": false }, "outputs": [], "source": [ "jax.config.update(\"jax_enable_x64\", True)" ] }, { "cell_type": "markdown", "id": "4a1129b2", "metadata": {}, "source": [ "We will use the following function to compute stationary distributions of stochastic matrices (for a reference to the algorithm, see p. 88 of [Economic Dynamics](https://johnstachurski.net/edtc))." ] }, { "cell_type": "code", "execution_count": null, "id": "77636ad4", "metadata": { "hide-output": false }, "outputs": [], "source": [ "@jax.jit\n", "def compute_stationary(P):\n", " n = P.shape[0]\n", " I = jnp.identity(n)\n", " O = jnp.ones((n, n))\n", " A = I - jnp.transpose(P) + O\n", " return jnp.linalg.solve(A, jnp.ones(n))" ] }, { "cell_type": "markdown", "id": "5a73ab73", "metadata": {}, "source": [ "### References\n", "\n", "The primary reference for this lecture is [[Aiyagari, 1994](https://python.quantecon.org/zreferences.html#id155)].\n", "\n", "A textbook treatment is available in chapter 18 of [[Ljungqvist and Sargent, 2018](https://python.quantecon.org/zreferences.html#id200)].\n", "\n", "A continuous time version of the model by SeHyoun Ahn and Benjamin Moll can be found [here](https://nbviewer.org/github/QuantEcon/QuantEcon.notebooks/blob/master/aiyagari_continuous_time.ipynb)." ] }, { "cell_type": "markdown", "id": "8933ba46", "metadata": {}, "source": [ "## The Economy" ] }, { "cell_type": "markdown", "id": "3b235927", "metadata": {}, "source": [ "### Households\n", "\n", "Infinitely lived households / consumers face idiosyncratic income shocks.\n", "\n", "A unit interval of *ex-ante* identical households face a common borrowing constraint.\n", "\n", "The savings problem faced by a typical household is\n", "\n", "$$\n", "\\max \\mathbb E \\sum_{t=0}^{\\infty} \\beta^t u(c_t)\n", "$$\n", "\n", "subject to\n", "\n", "$$\n", "a_{t+1} + c_t \\leq w z_t + (1 + r) a_t\n", "\\quad\n", "c_t \\geq 0,\n", "\\quad \\text{and} \\quad\n", "a_t \\geq -B\n", "$$\n", "\n", "where\n", "\n", "- $ c_t $ is current consumption \n", "- $ a_t $ is assets \n", "- $ z_t $ is an exogenous component of labor income capturing stochastic unemployment risk, etc. \n", "- $ w $ is a wage rate \n", "- $ r $ is a net interest rate \n", "- $ B $ is the maximum amount that the agent is allowed to borrow \n", "\n", "\n", "The exogenous process $ \\{z_t\\} $ follows a finite state Markov chain with given stochastic matrix $ P $.\n", "\n", "The wage and interest rate are fixed over time.\n", "\n", "In this simple version of the model, households supply labor inelastically because they do not value leisure." ] }, { "cell_type": "markdown", "id": "2a77b709", "metadata": {}, "source": [ "### Firms\n", "\n", "Firms produce output by hiring capital and labor.\n", "\n", "Firms act competitively and face constant returns to scale.\n", "\n", "Since returns to scale are constant, the number of firms does not matter.\n", "\n", "Hence we can consider a single (but nonetheless competitive) representative firm.\n", "\n", "The firm’s output is\n", "\n", "$$\n", "Y = A K^{\\alpha} N^{1 - \\alpha}\n", "$$\n", "\n", "where\n", "\n", "- $ A $ and $ \\alpha $ are parameters with $ A > 0 $ and $ \\alpha \\in (0, 1) $ \n", "- $ K $ is aggregate capital \n", "- $ N $ is total labor supply (which is constant in this simple version of the model) \n", "\n", "\n", "The firm’s problem is\n", "\n", "$$\n", "\\max_{K, N} \\left\\{ A K^{\\alpha} N^{1 - \\alpha} - (r + \\delta) K - w N \\right\\}\n", "$$\n", "\n", "The parameter $ \\delta $ is the depreciation rate.\n", "\n", "These parameters are stored in the following namedtuple:" ] }, { "cell_type": "code", "execution_count": null, "id": "caf6348d", "metadata": { "hide-output": false }, "outputs": [], "source": [ "class Firm(NamedTuple):\n", " A: float = 1.0 # Total factor productivity\n", " N: float = 1.0 # Total labor supply\n", " α: float = 0.33 # Capital share\n", " δ: float = 0.05 # Depreciation rate" ] }, { "cell_type": "markdown", "id": "e696f46a", "metadata": {}, "source": [ "From the first-order condition with respect to capital, the firm’s inverse demand for capital is\n", "\n", "\n", "\n", "$$\n", "r = A \\alpha \\left( \\frac{N}{K} \\right)^{1 - \\alpha} - \\delta \\tag{79.1}\n", "$$" ] }, { "cell_type": "code", "execution_count": null, "id": "33f6ff2e", "metadata": { "hide-output": false }, "outputs": [], "source": [ "def r_given_k(K, firm):\n", " \"\"\"\n", " Inverse demand curve for capital. The interest rate associated with a\n", " given demand for capital K.\n", " \"\"\"\n", " A, N, α, δ = firm\n", " return A * α * (N / K)**(1 - α) - δ" ] }, { "cell_type": "markdown", "id": "a42c62a7", "metadata": {}, "source": [ "Using this expression and the firm’s first-order condition for labor, we can pin down\n", "the equilibrium wage rate as a function of $ r $ as\n", "\n", "\n", "\n", "$$\n", "w(r) = A (1 - \\alpha) (A \\alpha / (r + \\delta))^{\\alpha / (1 - \\alpha)} \\tag{79.2}\n", "$$" ] }, { "cell_type": "code", "execution_count": null, "id": "8531a313", "metadata": { "hide-output": false }, "outputs": [], "source": [ "def r_to_w(r, firm):\n", " \"\"\"\n", " Equilibrium wages associated with a given interest rate r.\n", " \"\"\"\n", " A, N, α, δ = firm\n", " return A * (1 - α) * (A * α / (r + δ))**(α / (1 - α))" ] }, { "cell_type": "markdown", "id": "0dc4d69c", "metadata": {}, "source": [ "### Equilibrium\n", "\n", "We construct a **stationary rational expectations equilibrium (SREE)**.\n", "\n", "In such an equilibrium\n", "\n", "- prices induce behavior that generates aggregate quantities consistent with the prices \n", "- aggregate quantities and prices are constant over time \n", "\n", "\n", "In more detail, an SREE lists a set of prices, savings and production policies such that\n", "\n", "- households want to choose the specified savings policies taking the prices as given \n", "- firms maximize profits taking the same prices as given \n", "- the resulting aggregate quantities are consistent with the prices; in particular, the demand for capital equals the supply \n", "- aggregate quantities (defined as cross-sectional averages) are constant " ] }, { "cell_type": "markdown", "id": "0636f3b3", "metadata": {}, "source": [ "## Implementation\n", "\n", "Let’s look at how we might compute such an equilibrium in practice.\n", "\n", "Below we provide code to solve the household problem, taking $ r $ and $ w $ as fixed." ] }, { "cell_type": "markdown", "id": "7737a860", "metadata": {}, "source": [ "### Primitives and operators\n", "\n", "We will solve the household problem using value function iteration.\n", "\n", "First we set up a `NamedTuple` to store the parameters that define a household asset accumulation problem, as well as the grids used to solve it" ] }, { "cell_type": "code", "execution_count": null, "id": "1a8578bd", "metadata": { "hide-output": false }, "outputs": [], "source": [ "class Household(NamedTuple):\n", " β: float # Discount factor\n", " a_grid: jnp.ndarray # Asset grid\n", " z_grid: jnp.ndarray # Exogenous states\n", " Π: jnp.ndarray # Transition matrix\n", "\n", "def create_household(β=0.96, # Discount factor\n", " Π=[[0.9, 0.1], [0.1, 0.9]], # Markov chain\n", " z_grid=[0.1, 1.0], # Exogenous states\n", " a_min=1e-10, a_max=12.5, # Asset grid\n", " a_size=100):\n", " \"\"\"\n", " Create a Household namedtuple with custom grids.\n", " \"\"\"\n", " a_grid = jnp.linspace(a_min, a_max, a_size)\n", " z_grid, Π = map(jnp.array, (z_grid, Π))\n", " return Household(β=β, a_grid=a_grid, z_grid=z_grid, Π=Π)" ] }, { "cell_type": "markdown", "id": "d74a36ce", "metadata": {}, "source": [ "For now we assume that $ u(c) = \\log(c) $" ] }, { "cell_type": "code", "execution_count": null, "id": "933fc737", "metadata": { "hide-output": false }, "outputs": [], "source": [ "u = jnp.log" ] }, { "cell_type": "markdown", "id": "fec1f5d8", "metadata": {}, "source": [ "Here’s a namedtuple that stores the wage rate and interest rate with default values" ] }, { "cell_type": "code", "execution_count": null, "id": "5e6b9594", "metadata": { "hide-output": false }, "outputs": [], "source": [ "class Prices(NamedTuple):\n", " r: float = 0.01 # Interest rate\n", " w: float = 1.0 # Wages" ] }, { "cell_type": "markdown", "id": "37c06aad", "metadata": {}, "source": [ "Now we set up a vectorized version of the right-hand side of the Bellman equation (before maximization), which is a 3D array representing\n", "\n", "$$\n", "B(a, z, a') = u(wz + (1+r)a - a') + \\beta \\sum_{z'} v(a', z') \\Pi(z, z')\n", "$$\n", "\n", "for all $ (a, z, a') $." ] }, { "cell_type": "code", "execution_count": null, "id": "c1982633", "metadata": { "hide-output": false }, "outputs": [], "source": [ "def B(v, household, prices):\n", " # Unpack\n", " β, a_grid, z_grid, Π = household\n", " a_size, z_size = len(a_grid), len(z_grid)\n", " r, w = prices\n", "\n", " # Compute current consumption as array c[i, j, ip]\n", " a = jnp.reshape(a_grid, (a_size, 1, 1)) # a[i] -> a[i, j, ip]\n", " z = jnp.reshape(z_grid, (1, z_size, 1)) # z[j] -> z[i, j, ip]\n", " ap = jnp.reshape(a_grid, (1, 1, a_size)) # ap[ip] -> ap[i, j, ip]\n", " c = w * z + (1 + r) * a - ap\n", "\n", " # Calculate continuation rewards at all combinations of (a, z, ap)\n", " v = jnp.reshape(v, (1, 1, a_size, z_size)) # v[ip, jp] -> v[i, j, ip, jp]\n", " Π = jnp.reshape(Π, (1, z_size, 1, z_size)) # Π[j, jp] -> Π[i, j, ip, jp]\n", " EV = jnp.sum(v * Π, axis=-1) # sum over last index jp\n", "\n", " # Compute the right-hand side of the Bellman equation\n", " return jnp.where(c > 0, u(c) + β * EV, -jnp.inf)" ] }, { "cell_type": "markdown", "id": "5941752a", "metadata": {}, "source": [ "The next function computes greedy policies" ] }, { "cell_type": "code", "execution_count": null, "id": "98f2dec5", "metadata": { "hide-output": false }, "outputs": [], "source": [ "def get_greedy(v, household, prices):\n", " \"\"\"\n", " Computes a v-greedy policy σ, returned as a set of indices. If\n", " σ[i, j] equals ip, then a_grid[ip] is the maximizer at i, j.\n", " \"\"\"\n", " # argmax over ap\n", " return jnp.argmax(B(v, household, prices), axis=-1)" ] }, { "cell_type": "markdown", "id": "a225cebc", "metadata": {}, "source": [ "We define the Bellman operator $ T $, which takes a value function $ v $ and returns $ Tv $ as given in the Bellman equation" ] }, { "cell_type": "code", "execution_count": null, "id": "2134922e", "metadata": { "hide-output": false }, "outputs": [], "source": [ "def T(v, household, prices):\n", " \"\"\"\n", " The Bellman operator. Takes a value function v and returns Tv.\n", " \"\"\"\n", " return jnp.max(B(v, household, prices), axis=-1)" ] }, { "cell_type": "markdown", "id": "27c0f886", "metadata": {}, "source": [ "Here’s value function iteration, which repeatedly applies the Bellman operator until convergence" ] }, { "cell_type": "code", "execution_count": null, "id": "2095848f", "metadata": { "hide-output": false }, "outputs": [], "source": [ "@jax.jit\n", "def value_function_iteration(household, prices, tol=1e-4, max_iter=10_000):\n", " \"\"\"\n", " Implements value function iteration using a compiled JAX loop.\n", " \"\"\"\n", " β, a_grid, z_grid, Π = household\n", " a_size, z_size = len(a_grid), len(z_grid)\n", "\n", " def condition_function(loop_state):\n", " i, v, error = loop_state\n", " return jnp.logical_and(error > tol, i < max_iter)\n", "\n", " def update(loop_state):\n", " i, v, error = loop_state\n", " v_new = T(v, household, prices)\n", " error = jnp.max(jnp.abs(v_new - v))\n", " return i + 1, v_new, error\n", "\n", " # Initial loop state\n", " v_init = jnp.zeros((a_size, z_size))\n", " loop_state_init = (0, v_init, tol + 1)\n", "\n", " # Run the fixed point iteration\n", " i, v, error = jax.lax.while_loop(condition_function, update, loop_state_init)\n", "\n", " return get_greedy(v, household, prices)" ] }, { "cell_type": "markdown", "id": "a93271f9", "metadata": {}, "source": [ "As a first example of what we can do, let’s compute and plot an optimal accumulation policy at fixed prices" ] }, { "cell_type": "code", "execution_count": null, "id": "a50fc0a0", "metadata": { "hide-output": false }, "outputs": [], "source": [ "# Create an instance of Household\n", "household = create_household()\n", "prices = Prices()\n", "\n", "r, w = prices\n", "print(f\"Interest rate: {r}, Wage: {w}\")" ] }, { "cell_type": "code", "execution_count": null, "id": "fef25d79", "metadata": { "hide-output": false }, "outputs": [], "source": [ "with qe.Timer():\n", " σ_star = value_function_iteration(household, prices).block_until_ready()" ] }, { "cell_type": "markdown", "id": "4e678d76", "metadata": {}, "source": [ "The next plot shows asset accumulation policies at different values of the exogenous state" ] }, { "cell_type": "code", "execution_count": null, "id": "06603fe4", "metadata": { "hide-output": false }, "outputs": [], "source": [ "β, a_grid, z_grid, Π = household\n", "\n", "fig, ax = plt.subplots(figsize=(9, 9))\n", "ax.plot(a_grid, a_grid, 'k--', label=\"45 degrees\") \n", "for j, z in enumerate(z_grid):\n", " lb = f'$z = {z:.2}$'\n", " policy_vals = a_grid[σ_star[:, j]]\n", " ax.plot(a_grid, policy_vals, lw=2, alpha=0.6, label=lb)\n", " ax.set_xlabel('current assets')\n", " ax.set_ylabel('next period assets')\n", "ax.legend(loc='upper left')\n", "plt.show()" ] }, { "cell_type": "markdown", "id": "a201fc91", "metadata": {}, "source": [ "The plot shows asset accumulation policies at different values of the exogenous state." ] }, { "cell_type": "markdown", "id": "16706cbb", "metadata": {}, "source": [ "### Capital supply\n", "\n", "To start thinking about equilibrium, we need to know how much capital households supply at a given interest rate $ r $.\n", "\n", "This quantity can be calculated by taking the stationary distribution of assets under the optimal policy and computing the mean.\n", "\n", "The next function computes the stationary distribution for a given policy $ \\sigma $ via the following steps:\n", "\n", "- Compute the stationary distribution $ \\psi = (\\psi(a, z)) $ of $ P_{\\sigma} $, which defines the Markov chain of the state $ (a_t, z_t) $ under policy $ \\sigma $. \n", "- Sum out $ z_t $ to get the marginal distribution for $ a_t $. " ] }, { "cell_type": "code", "execution_count": null, "id": "e5b441b5", "metadata": { "hide-output": false }, "outputs": [], "source": [ "@jax.jit\n", "def compute_asset_stationary(σ, household):\n", " # Unpack\n", " β, a_grid, z_grid, Π = household\n", " a_size, z_size = len(a_grid), len(z_grid)\n", "\n", " # Construct P_σ as an array of the form P_σ[i, j, ip, jp]\n", " ap_idx = jnp.arange(a_size)\n", " ap_idx = jnp.reshape(ap_idx, (1, 1, a_size, 1))\n", " σ = jnp.reshape(σ, (a_size, z_size, 1, 1))\n", " A = jnp.where(σ == ap_idx, 1, 0)\n", " Π = jnp.reshape(Π, (1, z_size, 1, z_size))\n", " P_σ = A * Π\n", "\n", " # Reshape P_σ into a matrix\n", " n = a_size * z_size\n", " P_σ = jnp.reshape(P_σ, (n, n))\n", "\n", " # Get stationary distribution and reshape back onto [i, j] grid\n", " ψ = compute_stationary(P_σ)\n", " ψ = jnp.reshape(ψ, (a_size, z_size))\n", "\n", " # Sum along the rows to get the marginal distribution of assets\n", " ψ_a = jnp.sum(ψ, axis=1)\n", " return ψ_a" ] }, { "cell_type": "markdown", "id": "00b6c080", "metadata": {}, "source": [ "Let’s give this a test run." ] }, { "cell_type": "code", "execution_count": null, "id": "c808298e", "metadata": { "hide-output": false }, "outputs": [], "source": [ "ψ_a = compute_asset_stationary(σ_star, household)\n", "\n", "fig, ax = plt.subplots()\n", "ax.bar(household.a_grid, ψ_a)\n", "ax.set_xlabel(\"asset level\")\n", "ax.set_ylabel(\"probability mass\")\n", "plt.show()" ] }, { "cell_type": "markdown", "id": "c4a52bba", "metadata": {}, "source": [ "The distribution should sum to one:" ] }, { "cell_type": "code", "execution_count": null, "id": "cee04ff4", "metadata": { "hide-output": false }, "outputs": [], "source": [ "ψ_a.sum()" ] }, { "cell_type": "markdown", "id": "2965f76e", "metadata": {}, "source": [ "The next function computes aggregate capital supply by households under policy $ \\sigma $, given wages and interest rates" ] }, { "cell_type": "code", "execution_count": null, "id": "d2fc315d", "metadata": { "hide-output": false }, "outputs": [], "source": [ "def capital_supply(σ, household):\n", " \"\"\"\n", " Induced level of capital stock under the policy, taking r and w as given.\n", " \"\"\"\n", " β, a_grid, z_grid, Π = household\n", " ψ_a = compute_asset_stationary(σ, household)\n", " return float(jnp.sum(ψ_a * a_grid))" ] }, { "cell_type": "markdown", "id": "0900f8a0", "metadata": {}, "source": [ "### Equilibrium\n", "\n", "We compute a SREE as follows:\n", "\n", "1. Set $ n=0 $ and start with an initial guess $ K_0 $ for aggregate capital. \n", "1. Determine prices $ r, w $ from the firm decision problem, given $ K_n $. \n", "1. Compute the optimal savings policy of households given these prices. \n", "1. Compute aggregate capital $ K_{n+1} $ as the mean of steady-state capital given this savings policy. \n", "1. If $ K_{n+1} \\approx K_n $, stop; otherwise, go to step 2. \n", "\n", "\n", "We can write the sequence of operations in steps 2-4 as\n", "\n", "$$\n", "K_{n + 1} = G(K_n)\n", "$$\n", "\n", "If $ K_{n+1} $ agrees with $ K_n $ then we have a SREE.\n", "\n", "In other words, our problem is to find the fixed point of the one-dimensional map $ G $.\n", "\n", "Here’s $ G $ expressed as a Python function" ] }, { "cell_type": "code", "execution_count": null, "id": "f98f48cc", "metadata": { "hide-output": false }, "outputs": [], "source": [ "def G(K, firm, household):\n", " # Get prices r, w associated with K\n", " r = r_given_k(K, firm)\n", " w = r_to_w(r, firm)\n", "\n", " # Generate a household object with these prices, compute\n", " # aggregate capital.\n", " prices = Prices(r=r, w=w)\n", " σ_star = value_function_iteration(household, prices)\n", " return capital_supply(σ_star, household)" ] }, { "cell_type": "markdown", "id": "57965552", "metadata": {}, "source": [ "Let’s inspect visually as a first pass" ] }, { "cell_type": "code", "execution_count": null, "id": "873a3c7b", "metadata": { "hide-output": false }, "outputs": [], "source": [ "num_points = 50\n", "firm = Firm()\n", "household = create_household()\n", "k_vals = jnp.linspace(4, 12, num_points)\n", "out = [G(k, firm, household) for k in k_vals]\n", "\n", "fig, ax = plt.subplots(figsize=(11, 8))\n", "ax.plot(k_vals, out, lw=2, alpha=0.6, label='$G$')\n", "ax.plot(k_vals, k_vals, 'k--', label=\"45 degrees\")\n", "ax.set_xlabel('capital')\n", "ax.legend()\n", "plt.show()" ] }, { "cell_type": "markdown", "id": "e9b91560", "metadata": {}, "source": [ "Now let’s compute the equilibrium.\n", "\n", "Looking at the figure above, we see that a simple iteration scheme $ K_{n+1} = G(K_n) $ will cycle from high to low values, leading to slow convergence.\n", "\n", "As a result, we use a damped iteration scheme of the form\n", "\n", "$$\n", "K_{n+1} = \\alpha K_n + (1-\\alpha) G(K_n)\n", "$$" ] }, { "cell_type": "code", "execution_count": null, "id": "aa082fd4", "metadata": { "hide-output": false }, "outputs": [], "source": [ "def compute_equilibrium(firm, household,\n", " K0=6, α=0.99, max_iter=1_000, tol=1e-4, \n", " print_skip=10, verbose=False):\n", " n = 0\n", " K = K0\n", " error = tol + 1\n", " while error > tol and n < max_iter:\n", " new_K = α * K + (1 - α) * G(K, firm, household)\n", " error = abs(new_K - K)\n", " K = new_K\n", " n += 1\n", " if verbose and n % print_skip == 0:\n", " print(f\"At iteration {n} with error {error}\")\n", " return K, n" ] }, { "cell_type": "code", "execution_count": null, "id": "a75e9962", "metadata": { "hide-output": false }, "outputs": [], "source": [ "firm = Firm()\n", "household = create_household()\n", "print(\"\\nComputing equilibrium capital stock\")\n", "with qe.Timer():\n", " K_star, n = compute_equilibrium(firm, household, K0=6.0)\n", "print(f\"Computed equilibrium {K_star:.5} in {n} iterations\")" ] }, { "cell_type": "markdown", "id": "ffb20f11", "metadata": {}, "source": [ "This convergence is not very fast, given how quickly we can solve the household problem.\n", "\n", "You can try varying $ \\alpha $, but usually this parameter is hard to set a priori.\n", "\n", "In the exercises below you will be asked to use bisection instead, which generally performs better." ] }, { "cell_type": "markdown", "id": "ea9f422d", "metadata": {}, "source": [ "### Supply and demand curves\n", "\n", "We can visualize the equilibrium using supply and demand curves.\n", "\n", "The following code draws the aggregate supply and demand curves.\n", "\n", "The intersection gives the equilibrium interest rate and capital" ] }, { "cell_type": "code", "execution_count": null, "id": "fae2330e", "metadata": { "hide-output": false }, "outputs": [], "source": [ "def prices_to_capital_stock(household, r, firm):\n", " \"\"\"\n", " Map prices to the induced level of capital stock.\n", " \"\"\"\n", " w = r_to_w(r, firm)\n", " prices = Prices(r=r, w=w)\n", "\n", " # Compute the optimal policy\n", " σ_star = value_function_iteration(household, prices)\n", "\n", " # Compute capital supply\n", " return capital_supply(σ_star, household)\n", "\n", "# Create a grid of r values to compute demand and supply of capital\n", "num_points = 20\n", "r_vals = jnp.linspace(0.005, 0.04, num_points)\n", "\n", "# Compute supply of capital\n", "k_vals = []\n", "for r in r_vals:\n", " k_vals.append(prices_to_capital_stock(household, r, firm))\n", "\n", "# Plot against demand for capital by firms\n", "fig, ax = plt.subplots(figsize=(11, 8))\n", "ax.plot(k_vals, r_vals, lw=2, alpha=0.6, \n", " label='supply of capital')\n", "ax.plot(k_vals, r_given_k(\n", " jnp.array(k_vals), firm), lw=2, alpha=0.6, \n", " label='demand for capital')\n", "\n", "# Add marker at equilibrium\n", "r_star = r_given_k(K_star, firm)\n", "ax.plot(K_star, r_star, 'o', markersize=10, label='equilibrium')\n", "\n", "ax.set_xlabel('capital')\n", "ax.set_ylabel('interest rate')\n", "ax.legend(loc='upper right')\n", "\n", "plt.show()" ] }, { "cell_type": "markdown", "id": "763fad07", "metadata": {}, "source": [ "## Exercises" ] }, { "cell_type": "markdown", "id": "d234ad4e", "metadata": {}, "source": [ "## Exercise 79.1\n", "\n", "Write a new version of `compute_equilibrium` that uses `bisect` from `scipy.optimize` instead of damped iteration.\n", "\n", "See if you can make it faster than the previous version.\n", "\n", "In `bisect`,\n", "\n", "- you should set `xtol=1e-4` to have the same error tolerance as the previous version. \n", "- for the lower and upper bounds of the bisection routine try `a = 1.0` and `b = 20.0`. " ] }, { "cell_type": "markdown", "id": "adbb26ba", "metadata": {}, "source": [ "## Solution\n", "\n", "We use bisection to find the zero of the function $ h(k) = k - G(k) $" ] }, { "cell_type": "code", "execution_count": null, "id": "63e2f6a7", "metadata": { "hide-output": false }, "outputs": [], "source": [ "def compute_equilibrium_bisect(firm, household, a=1.0, b=20.0):\n", " K = bisect(lambda k: k - G(k, firm, household), a, b, xtol=1e-4)\n", " return K\n", "\n", "firm = Firm()\n", "household = create_household()\n", "print(\"\\nComputing equilibrium capital stock using bisection\")\n", "with qe.Timer():\n", " K_star = compute_equilibrium_bisect(firm, household)\n", "print(f\"Computed equilibrium capital stock {K_star:.5}\")" ] }, { "cell_type": "markdown", "id": "d4d87083", "metadata": {}, "source": [ "The bisection method is faster than the damped iteration scheme." ] }, { "cell_type": "markdown", "id": "f027e3d7", "metadata": {}, "source": [ "## Exercise 79.2\n", "\n", "Show how equilibrium capital stock changes with $ \\beta $.\n", "\n", "Use the following values of $ \\beta $ and plot the relationship you find." ] }, { "cell_type": "code", "execution_count": null, "id": "1c4b699d", "metadata": { "hide-output": false }, "outputs": [], "source": [ "β_vals = jnp.linspace(0.94, 0.98, 20)" ] }, { "cell_type": "markdown", "id": "bb2b70f6", "metadata": {}, "source": [ "## Solution" ] }, { "cell_type": "code", "execution_count": null, "id": "bbb88fe7", "metadata": { "hide-output": false }, "outputs": [], "source": [ "K_vals = []\n", "K = 6.0 # initial guess\n", "\n", "for β in β_vals:\n", " household = create_household(β=β)\n", " K = compute_equilibrium_bisect(firm, household, 0.5 * K, 1.5 * K)\n", " print(f\"Computed equilibrium {K:.4} at β = {β}\")\n", " K_vals.append(K)\n", "\n", "fig, ax = plt.subplots()\n", "ax.plot(β_vals, K_vals, ms=2)\n", "ax.set_xlabel(r'$\\beta$')\n", "ax.set_ylabel('capital')\n", "plt.show()" ] }, { "cell_type": "markdown", "id": "c7b96740", "metadata": {}, "source": [ "## Exercise 79.3\n", "\n", "In this lecture, we used value function iteration to solve the household problem.\n", "\n", "An alternative is Howard policy iteration (HPI), which is discussed in detail in [Dynamic Programming](https://dp.quantecon.org/).\n", "\n", "HPI can be faster than VFI for some problems because it uses fewer but more computationally intensive iterations.\n", "\n", "Your task is to implement Howard policy iteration and compare the results with value function iteration.\n", "\n", "**Key concepts you’ll need:**\n", "\n", "Howard policy iteration requires computing the value $ v_{\\sigma} $ of a policy $ \\sigma $, defined as:\n", "\n", "$$\n", "v_{\\sigma} = (I - \\beta P_{\\sigma})^{-1} r_{\\sigma}\n", "$$\n", "\n", "where $ r_{\\sigma} $ is the reward vector under policy $ \\sigma $, and $ P_{\\sigma} $ is the transition matrix induced by $ \\sigma $.\n", "\n", "To solve this, you’ll need to:\n", "\n", "1. Compute current rewards $ r_{\\sigma}(a, z) = u((1 + r)a + wz - \\sigma(a, z)) $ \n", "1. Set up the linear operator $ R_{\\sigma} $ where $ (R_{\\sigma} v)(a, z) = v(a, z) - \\beta \\sum_{z'} v(\\sigma(a, z), z') \\Pi(z, z') $ \n", "1. Solve $ v_{\\sigma} = R_{\\sigma}^{-1} r_{\\sigma} $ using `jax.scipy.sparse.linalg.bicgstab` \n", "\n", "\n", "You can use the `get_greedy` function that’s already defined in this lecture.\n", "\n", "Implement the following Howard policy iteration routine:" ] }, { "cell_type": "code", "execution_count": null, "id": "402166b6", "metadata": { "hide-output": false }, "outputs": [], "source": [ "def howard_policy_iteration(household, prices,\n", " tol=1e-4, max_iter=10_000, verbose=False):\n", " \"\"\"\n", " Howard policy iteration routine.\n", " \"\"\"\n", " # Your code here\n", " pass" ] }, { "cell_type": "markdown", "id": "59b58173", "metadata": {}, "source": [ "Once implemented, compute the equilibrium capital stock using HPI and verify that it produces approximately the same result as VFI at the default parameter values." ] }, { "cell_type": "markdown", "id": "a31bd995", "metadata": {}, "source": [ "## Solution\n", "\n", "First, we need to implement the helper functions for Howard policy iteration.\n", "\n", "The following function computes the array $ r_{\\sigma} $ which gives current rewards given policy $ \\sigma $:" ] }, { "cell_type": "code", "execution_count": null, "id": "99e93d80", "metadata": { "hide-output": false }, "outputs": [], "source": [ "def compute_r_σ(σ, household, prices):\n", " \"\"\"\n", " Compute current rewards at each i, j under policy σ. In particular,\n", "\n", " r_σ[i, j] = u((1 + r)a[i] + wz[j] - a'[ip])\n", "\n", " when ip = σ[i, j].\n", " \"\"\"\n", " # Unpack\n", " β, a_grid, z_grid, Π = household\n", " a_size, z_size = len(a_grid), len(z_grid)\n", " r, w = prices\n", "\n", " # Compute r_σ[i, j]\n", " a = jnp.reshape(a_grid, (a_size, 1))\n", " z = jnp.reshape(z_grid, (1, z_size))\n", " ap = a_grid[σ]\n", " c = (1 + r) * a + w * z - ap\n", " r_σ = u(c)\n", "\n", " return r_σ" ] }, { "cell_type": "markdown", "id": "71427afe", "metadata": {}, "source": [ "The linear operator $ R_{\\sigma} $ is defined as:" ] }, { "cell_type": "code", "execution_count": null, "id": "a6d8b20c", "metadata": { "hide-output": false }, "outputs": [], "source": [ "def R_σ(v, σ, household):\n", " # Unpack\n", " β, a_grid, z_grid, Π = household\n", " a_size, z_size = len(a_grid), len(z_grid)\n", "\n", " # Set up the array v[σ[i, j], jp]\n", " zp_idx = jnp.arange(z_size)\n", " zp_idx = jnp.reshape(zp_idx, (1, 1, z_size))\n", " σ = jnp.reshape(σ, (a_size, z_size, 1))\n", " V = v[σ, zp_idx]\n", "\n", " # Expand Π[j, jp] to Π[i, j, jp]\n", " Π = jnp.reshape(Π, (1, z_size, z_size))\n", "\n", " # Compute and return v[i, j] - β Σ_jp v[σ[i, j], jp] * Π[j, jp]\n", " return v - β * jnp.sum(V * Π, axis=-1)" ] }, { "cell_type": "markdown", "id": "5907cfa3", "metadata": {}, "source": [ "The next function computes the lifetime value of a given policy:" ] }, { "cell_type": "code", "execution_count": null, "id": "799fb3da", "metadata": { "hide-output": false }, "outputs": [], "source": [ "def get_value(σ, household, prices):\n", " \"\"\"\n", " Get the lifetime value of policy σ by computing\n", "\n", " v_σ = R_σ^{-1} r_σ\n", " \"\"\"\n", " r_σ = compute_r_σ(σ, household, prices)\n", "\n", " # Reduce R_σ to a function in v\n", " _R_σ = lambda v: R_σ(v, σ, household)\n", "\n", " # Compute v_σ = R_σ^{-1} r_σ using an iterative routine.\n", " return jax.scipy.sparse.linalg.bicgstab(_R_σ, r_σ)[0]" ] }, { "cell_type": "markdown", "id": "4871f9d4", "metadata": {}, "source": [ "Now we can implement Howard policy iteration:" ] }, { "cell_type": "code", "execution_count": null, "id": "0f94a24a", "metadata": { "hide-output": false }, "outputs": [], "source": [ "@jax.jit\n", "def howard_policy_iteration(household, prices, tol=1e-4, max_iter=10_000):\n", " \"\"\"\n", " Howard policy iteration routine using a compiled JAX loop.\n", " \"\"\"\n", " β, a_grid, z_grid, Π = household\n", " a_size, z_size = len(a_grid), len(z_grid)\n", "\n", " def condition_function(loop_state):\n", " i, σ, v_σ, error = loop_state\n", " return jnp.logical_and(error > tol, i < max_iter)\n", "\n", " def update(loop_state):\n", " i, σ, v_σ, error = loop_state\n", " σ_new = get_greedy(v_σ, household, prices)\n", " v_σ_new = get_value(σ_new, household, prices)\n", " error = jnp.max(jnp.abs(v_σ_new - v_σ))\n", " return i + 1, σ_new, v_σ_new, error\n", "\n", " # Initial loop state\n", " σ_init = jnp.zeros((a_size, z_size), dtype=int)\n", " v_σ_init = get_value(σ_init, household, prices)\n", " loop_state_init = (0, σ_init, v_σ_init, tol + 1)\n", "\n", " # Run the fixed point iteration\n", " i, σ, v_σ, error = jax.lax.while_loop(condition_function, update, loop_state_init)\n", "\n", " return σ" ] }, { "cell_type": "markdown", "id": "adaae003", "metadata": {}, "source": [ "Now let’s create a modified version of the G function that uses HPI:" ] }, { "cell_type": "code", "execution_count": null, "id": "a1166dce", "metadata": { "hide-output": false }, "outputs": [], "source": [ "def G_hpi(K, firm, household):\n", " # Get prices r, w associated with K\n", " r = r_given_k(K, firm)\n", " w = r_to_w(r, firm)\n", "\n", " # Generate prices and compute aggregate capital using HPI.\n", " prices = Prices(r=r, w=w)\n", " σ_star = howard_policy_iteration(household, prices)\n", " return capital_supply(σ_star, household)" ] }, { "cell_type": "markdown", "id": "6fde8496", "metadata": {}, "source": [ "And compute the equilibrium using HPI:" ] }, { "cell_type": "code", "execution_count": null, "id": "bc69cdf8", "metadata": { "hide-output": false }, "outputs": [], "source": [ "def compute_equilibrium_bisect_hpi(firm, household, a=1.0, b=20.0):\n", " K = bisect(lambda k: k - G_hpi(k, firm, household), a, b, xtol=1e-4)\n", " return K\n", "\n", "firm = Firm()\n", "household = create_household()\n", "print(\"\\nComputing equilibrium capital stock using HPI\")\n", "with qe.Timer():\n", " K_star_hpi = compute_equilibrium_bisect_hpi(firm, household)\n", "print(f\"Computed equilibrium capital stock with HPI: {K_star_hpi:.5}\")\n", "print(f\"Previous equilibrium capital stock with VFI: {K_star:.5}\")\n", "print(f\"Difference: {abs(K_star_hpi - K_star):.6}\")" ] }, { "cell_type": "markdown", "id": "15eb872c", "metadata": {}, "source": [ "The results show that both methods produce approximately the same equilibrium, confirming that HPI is a valid alternative to VFI." ] } ], "metadata": { "date": 1770028415.6623738, "filename": "aiyagari.md", "kernelspec": { "display_name": "Python", "language": "python3", "name": "python3" }, "title": "The Aiyagari Model" }, "nbformat": 4, "nbformat_minor": 5 }