aiyagari_jax.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "bed2c952",
   "metadata": {},
   "source": [
    "# The Aiyagari Model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "44fabb6b",
   "metadata": {},
   "source": [
    "# GPU\n",
    "\n",
    "This lecture was built using a machine with JAX installed and access to a GPU.\n",
    "\n",
    "To run this lecture on [Google Colab](https://colab.research.google.com/), click on the “play” icon top right, select Colab, and set the runtime environment to include a GPU.\n",
    "\n",
    "To run this lecture on your own machine, you need to install [Google JAX](https://github.com/google/jax)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "35eb1dbd",
   "metadata": {},
   "source": [
    "## Overview\n",
    "\n",
    "In this lecture, we describe the structure of a class of models that build on work by Truman Bewley [[Bew77](https://python.quantecon.org/zreferences.html#id173)].\n",
    "\n",
    "We begin by discussing an example of a Bewley model due to Rao Aiyagari [[Aiy94](https://python.quantecon.org/zreferences.html#id137)].\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 [[Aiy94](https://python.quantecon.org/zreferences.html#id138)]  \n",
    "- risk sharing and asset pricing [[HL96](https://python.quantecon.org/zreferences.html#id130)]  \n",
    "- the shape of the wealth distribution [[BBZ15](https://python.quantecon.org/zreferences.html#id131)]  "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ef15fbe5",
   "metadata": {},
   "source": [
    "### References\n",
    "\n",
    "The primary reference for this lecture is [[Aiy94](https://python.quantecon.org/zreferences.html#id138)].\n",
    "\n",
    "A textbook treatment is available in chapter 18 of [[LS18](https://python.quantecon.org/zreferences.html#id183)].\n",
    "\n",
    "A less sophisticated version of this lecture (without JAX) can be found\n",
    "[here](https://python.quantecon.org/aiyagari.html)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aad09f32",
   "metadata": {},
   "source": [
    "### Preliminaries\n",
    "\n",
    "We use the following imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8104d8e0",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "import time\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import jax\n",
    "import jax.numpy as jnp\n",
    "from collections import namedtuple"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e765c7ee",
   "metadata": {},
   "source": [
    "Let’s check the GPU we are running"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e2e52927",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "!nvidia-smi"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b33e2d48",
   "metadata": {},
   "source": [
    "We will use 64 bit floats with JAX in order to increase the precision."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "322acf1f",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "jax.config.update(\"jax_enable_x64\", True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "86573cd5",
   "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": "57d199c1",
   "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": "90ae8bbb",
   "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": "ec80d4ba",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "Firm = namedtuple('Firm', ('A', 'N', 'α', 'δ'))\n",
    "\n",
    "def create_firm(A=1.0,\n",
    "                N=1.0,\n",
    "                α=0.33,\n",
    "                δ=0.05):\n",
    "    \"\"\"\n",
    "    Create a namedtuple that stores firm data.\n",
    "    \n",
    "    \"\"\"\n",
    "    return Firm(A=A, N=N, α=α, δ=δ)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6e69ddd5",
   "metadata": {},
   "source": [
    "From the first-order condition with respect to capital, the firm’s inverse demand for capital is\n",
    "\n",
    "\n",
    "<a id='equation-equation-aiy-rgk'></a>\n",
    "$$\n",
    "r = A \\alpha  \\left( \\frac{N}{K} \\right)^{1 - \\alpha} - \\delta \\tag{18.1}\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e9ce7371",
   "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": "48c38e45",
   "metadata": {},
   "source": [
    "Using [(18.1)](#equation-equation-aiy-rgk) and the firm’s first-order condition for labor,\n",
    "\n",
    "\n",
    "<a id='equation-equation-aiy-wgr'></a>\n",
    "$$\n",
    "w(r) = A  (1 - \\alpha)  (A \\alpha / (r + \\delta))^{\\alpha / (1 - \\alpha)} \\tag{18.2}\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b72e5463",
   "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": "0483bdac",
   "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",
    "In this simple version of the model, households supply labor  inelastically because they do not value leisure.\n",
    "\n",
    "Below we provide code to solve the household problem, taking $ r $ and $ w $ as fixed."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "20e85937",
   "metadata": {},
   "source": [
    "### Primitives and Operators\n",
    "\n",
    "We will solve the household problem using Howard policy iteration\n",
    "(see Ch 5 of [Dynamic Programming](https://dp.quantecon.org/)).\n",
    "\n",
    "First we set up a namedtuple to store the parameters that define a household asset\n",
    "accumulation problem, as well as the grids used to solve it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d68ed886",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "Household = namedtuple('Household', \n",
    "                       ('β', 'a_grid', 'z_grid', 'Π'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "702ff6f4",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "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=20,       # Asset grid\n",
    "                     a_size=200):\n",
    "    \"\"\"\n",
    "    Create a namedtuple that stores all data needed to solve the household\n",
    "    problem, given prices.\n",
    "\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": "f4c093fe",
   "metadata": {},
   "source": [
    "For now we assume that $ u(c) = \\log(c) $.\n",
    "\n",
    "(CRRA utility is treated in the exercises.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6d18690d",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "u = jnp.log"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a18b9594",
   "metadata": {},
   "source": [
    "Here’s a tuple that stores the wage rate and interest rate, as well as a function that creates a price namedtuple with default values."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0efd1846",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "Prices = namedtuple('Prices', ('r', 'w'))\n",
    "\n",
    "def create_prices(r=0.01,   # Interest rate\n",
    "                  w=1.0):   # Wages\n",
    "    return Prices(r=r, w=w)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ee7a283f",
   "metadata": {},
   "source": [
    "Now we set up a vectorized version of the right-hand side of the Bellman equation\n",
    "(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') Π(z, z')\n",
    "$$\n",
    "\n",
    "for all $ (a, z, a') $."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fb237881",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "@jax.jit\n",
    "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": "a4692863",
   "metadata": {},
   "source": [
    "The next function computes greedy policies."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10c908ec",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "@jax.jit\n",
    "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",
    "    \"\"\"\n",
    "    return jnp.argmax(B(v, household, prices), axis=-1) # argmax over ap"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d502cbbf",
   "metadata": {},
   "source": [
    "The following function computes the array $ r_{\\sigma} $ which gives current\n",
    "rewards given policy $ \\sigma $."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "af09f7fc",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "@jax.jit\n",
    "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",
    "    \"\"\"\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": "549dda72",
   "metadata": {},
   "source": [
    "The value $ v_{\\sigma} $ of a policy $ \\sigma $ is defined as\n",
    "\n",
    "$$\n",
    "v_{\\sigma} = (I - \\beta P_{\\sigma})^{-1} r_{\\sigma}\n",
    "$$\n",
    "\n",
    "(See Ch 5 of [Dynamic Programming](https://dp.quantecon.org/) for notation and background on Howard policy iteration.)\n",
    "\n",
    "To compute this vector, we set up the linear map $ v \\rightarrow R_{\\sigma} v $, where $ R_{\\sigma} := I - \\beta P_{\\sigma} $.\n",
    "\n",
    "This map can be expressed as\n",
    "\n",
    "$$\n",
    "(R_{\\sigma} v)(a, z) = v(a, z) - \\beta \\sum_{z'} v(\\sigma(a, z), z') Π(z, z')\n",
    "$$\n",
    "\n",
    "(Notice that $ R_\\sigma $ is expressed as a linear operator rather than a matrix – this is much easier and cleaner to code, and also exploits sparsity.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fe14cd8c",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "@jax.jit\n",
    "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": "f58319fa",
   "metadata": {},
   "source": [
    "The next function computes the lifetime value of a given policy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "111a3274",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "@jax.jit\n",
    "def get_value(σ, household, prices):\n",
    "    \"\"\"\n",
    "    Get the lifetime value of policy σ by computing\n",
    "\n",
    "        v_σ = R_σ^{-1} r_σ\n",
    "\n",
    "    \"\"\"\n",
    "    r_σ = compute_r_σ(σ, household, prices)\n",
    "    # Reduce R_σ to a function in v\n",
    "    _R_σ = lambda v: R_σ(v, σ, household)\n",
    "    # Compute v_σ = R_σ^{-1} r_σ using an iterative routing.\n",
    "    return jax.scipy.sparse.linalg.bicgstab(_R_σ, r_σ)[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "da748144",
   "metadata": {},
   "source": [
    "Here’s the Howard policy iteration."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1cc1a63d",
   "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",
    "    \"\"\"\n",
    "    β, a_grid, z_grid, Π = household\n",
    "    a_size, z_size = len(a_grid), len(z_grid)\n",
    "    σ = jnp.zeros((a_size, z_size), dtype=int)\n",
    "    \n",
    "    v_σ = get_value(σ, household, prices)\n",
    "    i = 0\n",
    "    error = tol + 1\n",
    "    while error > tol and i < max_iter:\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",
    "        σ = σ_new\n",
    "        v_σ = v_σ_new\n",
    "        i = i + 1\n",
    "        if verbose:\n",
    "            print(f\"Concluded loop {i} with error {error}.\")\n",
    "    return σ"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "634e5685",
   "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": "894dfe2f",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "# Create an instance of Household\n",
    "household = create_household()\n",
    "prices = create_prices()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b75af5a9",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "r, w = prices"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c637509b",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "r, w"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b11b3509",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "%time σ_star = howard_policy_iteration(household, prices, verbose=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2ced6644",
   "metadata": {},
   "source": [
    "The next plot shows asset accumulation policies at different values of the exogenous state."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "759d3a7b",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "β, a_grid, z_grid, Π = household\n",
    "\n",
    "fig, ax = plt.subplots()\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": "4937e77a",
   "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\n",
    "  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": "8c35818f",
   "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": "03e57236",
   "metadata": {},
   "source": [
    "Let’s give this a test run."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "159d0532",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "ψ_a = compute_asset_stationary(σ_star, household)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "40b3280c",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "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": "afd1bc8f",
   "metadata": {},
   "source": [
    "The distribution should sum to one:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "22ba82ef",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "ψ_a.sum()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3e79572f",
   "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": "2ce527bc",
   "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",
    "    \"\"\"\n",
    "    β, a_grid, z_grid, Π = household\n",
    "    ψ_a = compute_asset_stationary(σ, household)\n",
    "    return float(jnp.sum(ψ_a * a_grid))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7550ffa7",
   "metadata": {},
   "source": [
    "## Equilibrium\n",
    "\n",
    "We compute a **stationary rational expectations equilibrium** (SREE) as follows:\n",
    "\n",
    "1. set $ n=0 $, start with 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 the 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": "b86067a7",
   "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",
    "    # Generate a household object with these prices, compute\n",
    "    # aggregate capital.\n",
    "    prices = create_prices(r=r, w=w)\n",
    "    σ_star = howard_policy_iteration(household, prices)\n",
    "    return capital_supply(σ_star, household)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bd2a2a04",
   "metadata": {},
   "source": [
    "### Visual inspection\n",
    "\n",
    "Let’s inspect visually as a first pass."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5f723d76",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "num_points = 50\n",
    "firm = create_firm()\n",
    "household = create_household()\n",
    "k_vals = np.linspace(4, 12, num_points)\n",
    "out = [G(k, firm, household) for k in k_vals]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fb4f719d",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\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": "24295279",
   "metadata": {},
   "source": [
    "### Computing the equilibrium\n",
    "\n",
    "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": "b3f6f061",
   "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": "3be20cc0",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "firm = create_firm()\n",
    "household = create_household()\n",
    "print(\"\\nComputing equilibrium capital stock\")\n",
    "start = time.time()\n",
    "K_star, n = compute_equilibrium(firm, household, K0=6.0, verbose=True)\n",
    "elapsed = time.time() - start\n",
    "print(f\"Computed equilibrium {K_star:.5} in {n} iterations and {elapsed} seconds\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fe3e0b5e",
   "metadata": {},
   "source": [
    "This 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": "efe367e5",
   "metadata": {},
   "source": [
    "## Exercises"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6280ecf9",
   "metadata": {},
   "source": [
    "## Exercise 18.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 that 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": "d1b01bea",
   "metadata": {},
   "source": [
    "## Solution to[ Exercise 18.1](https://jax.quantecon.org/#aiyagari-ex1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c53cdc47",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "from scipy.optimize import bisect"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c9f3f1ff",
   "metadata": {},
   "source": [
    "We use bisection to find the zero of the function $ h(k) = k - G(k) $."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "34f3d35f",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def compute_equilibrium(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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3b1e4df1",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "firm = create_firm()\n",
    "household = create_household()\n",
    "print(\"\\nComputing equilibrium capital stock\")\n",
    "start = time.time()\n",
    "K_star = compute_equilibrium(firm, household)\n",
    "elapsed = time.time() - start\n",
    "print(f\"Computed equilibrium capital stock {K_star:.5} in {elapsed} seconds\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b6d290ff",
   "metadata": {},
   "source": [
    "Bisection seems to be faster than the damped iteration scheme."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2e010fa2",
   "metadata": {},
   "source": [
    "## Exercise 18.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": "bc6f1989",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "β_vals = np.linspace(0.94, 0.98, 20)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bf291b51",
   "metadata": {},
   "source": [
    "## Solution to[ Exercise 18.2](https://jax.quantecon.org/#aiyagari-ex2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5192cce4",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "K_vals = np.empty_like(β_vals)\n",
    "K = 6.0  # initial guess\n",
    "\n",
    "for i, β in enumerate(β_vals):\n",
    "    household = create_household(β=β)\n",
    "    K = compute_equilibrium(firm, household, 0.5 * K, 1.5 * K)\n",
    "    print(f\"Computed equilibrium {K:.4} at β = {β}\")\n",
    "    K_vals[i] = K"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a35ccdef",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "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()"
   ]
  }
 ],
 "metadata": {
  "date": 1729509368.5433939,
  "filename": "aiyagari_jax.md",
  "kernelspec": {
   "display_name": "Python",
   "language": "python3",
   "name": "python3"
  },
  "title": "The Aiyagari Model"
 },
 "nbformat": 4,
 "nbformat_minor": 5
}