ifp_egm.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "620e927d",
   "metadata": {},
   "source": [
    "# Endogenous Grid Method"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0ad08acc",
   "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": "cdb1f501",
   "metadata": {},
   "source": [
    "## Overview\n",
    "\n",
    "In this lecture we use the endogenous grid method (EGM) to solve a basic income\n",
    "fluctuation (optimal savings) problem.\n",
    "\n",
    "Background on the endogenous grid method can be found in [an earlier\n",
    "QuantEcon lecture](https://python.quantecon.org/egm_policy_iter.html).\n",
    "\n",
    "Here we focus on providing an efficient JAX implementation.\n",
    "\n",
    "In addition to JAX and Anaconda, this lecture will need the following libraries:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "59da51a8",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "!pip install --upgrade quantecon"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "03832f88",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "import quantecon as qe\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import jax\n",
    "import jax.numpy as jnp\n",
    "import numba\n",
    "from time import time"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "43b1c6f8",
   "metadata": {},
   "source": [
    "Let’s check the GPU we are running"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ad8b0ad6",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "!nvidia-smi"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "de508fe7",
   "metadata": {},
   "source": [
    "We use 64 bit floating point numbers for extra precision."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5a19e7a4",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "jax.config.update(\"jax_enable_x64\", True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3087b186",
   "metadata": {},
   "source": [
    "## Setup\n",
    "\n",
    "We consider a household that chooses a state-contingent consumption plan $ \\{c_t\\}_{t \\geq 0} $ to maximize\n",
    "\n",
    "$$\n",
    "\\mathbb{E} \\, \\sum_{t=0}^{\\infty} \\beta^t u(c_t)\n",
    "$$\n",
    "\n",
    "subject to\n",
    "\n",
    "$$\n",
    "a_{t+1} \\leq  R(a_t - c_t)  + Y_{t+1},\n",
    "    \\quad c_t \\geq 0,\n",
    "    \\quad a_t \\geq 0\n",
    "    \\quad t = 0, 1, \\ldots\n",
    "$$\n",
    "\n",
    "Here $ R = 1 + r $ where $ r $ is the interest rate.\n",
    "\n",
    "The income process $ \\{Y_t\\} $ is a [Markov chain](https://python.quantecon.org/finite_markov.html) generated by stochastic matrix $ P $.\n",
    "\n",
    "The matrix $ P $ and the grid of values taken by $ Y_t $ are obtained by\n",
    "discretizing the AR(1) process\n",
    "\n",
    "$$\n",
    "Y_{t+1} = \\rho Y_t + \\nu \\epsilon_{t+1}\n",
    "$$\n",
    "\n",
    "where $ \\{\\epsilon_t\\} $ is IID and standard normal.\n",
    "\n",
    "Utility has the CRRA specification\n",
    "\n",
    "$$\n",
    "u(c) = \\frac{c^{1 - \\gamma}} {1 - \\gamma}\n",
    "$$\n",
    "\n",
    "The following function stores default parameter values for the income\n",
    "fluctuation problem and creates suitable arrays."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "133c8e8a",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def ifp(R=1.01,             # gross interest rate\n",
    "        β=0.99,             # discount factor\n",
    "        γ=1.5,              # CRRA preference parameter\n",
    "        s_max=16,           # savings grid max\n",
    "        s_size=200,         # savings grid size\n",
    "        ρ=0.99,             # income persistence\n",
    "        ν=0.02,             # income volatility\n",
    "        y_size=25):         # income grid size\n",
    "  \n",
    "    # require R β < 1 for convergence\n",
    "    assert R * β < 1, \"Stability condition failed.\"\n",
    "    # Create income Markov chain\n",
    "    mc = qe.tauchen(y_size, ρ, ν)\n",
    "    y_grid, P = jnp.exp(mc.state_values), mc.P\n",
    "    # Shift to JAX arrays\n",
    "    P, y_grid = jax.device_put((P, y_grid))\n",
    "    s_grid = jnp.linspace(0, s_max, s_size)\n",
    "    # Pack and return\n",
    "    constants = β, R, γ\n",
    "    sizes = s_size, y_size\n",
    "    arrays = s_grid, y_grid, P\n",
    "    return constants, sizes, arrays"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "63451fef",
   "metadata": {},
   "source": [
    "## Solution method\n",
    "\n",
    "Let $ S = \\mathbb R_+ \\times \\mathsf Y $ be the set of possible values for the\n",
    "state $ (a_t, Y_t) $.\n",
    "\n",
    "We aim to compute an optimal consumption policy $ \\sigma^* \\colon S \\to \\mathbb\n",
    "R $, under which dynamics are given by\n",
    "\n",
    "$$\n",
    "c_t = \\sigma^*(a_t, Y_t)\n",
    "    \\quad \\text{and} \\quad\n",
    "    a_{t+1} = R (a_t - c_t) + Y_{t+1}\n",
    "$$\n",
    "\n",
    "In this section we discuss how we intend to solve for this policy."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a340174b",
   "metadata": {},
   "source": [
    "### Euler equation\n",
    "\n",
    "The Euler equation for the optimization problem is\n",
    "\n",
    "$$\n",
    "u' (c_t)\n",
    "    = \\max \\left\\{\n",
    "        \\beta R \\,  \\mathbb{E}_t  u'(c_{t+1})  \\,,\\;  u'(a_t)\n",
    "    \\right\\}\n",
    "$$\n",
    "\n",
    "An explanation for this expression can be found [here](https://python.quantecon.org/ifp.html#value-function-and-euler-equation).\n",
    "\n",
    "We rewrite the Euler equation in functional form\n",
    "\n",
    "$$\n",
    "(u' \\circ \\sigma)  (a, y)\n",
    "    = \\max \\left\\{\n",
    "    \\beta R \\, \\mathbb E_y (u' \\circ \\sigma)\n",
    "        [R (a - \\sigma(a, y)) + \\hat Y, \\, \\hat Y]\n",
    "    \\, , \\;\n",
    "         u'(a)\n",
    "         \\right\\}\n",
    "$$\n",
    "\n",
    "where $ (u' \\circ \\sigma)(a, y) := u'(\\sigma(a, y)) $ and $ \\sigma $ is a consumption\n",
    "policy.\n",
    "\n",
    "For given consumption policy $ \\sigma $, we define $ (K \\sigma) (a,y) $ as the unique $ c \\in [0, a] $ that solves\n",
    "\n",
    "\n",
    "<a id='equation-eq-kaper'></a>\n",
    "$$\n",
    "u'(c)\n",
    "= \\max \\left\\{\n",
    "           \\beta R \\, \\mathbb E_y (u' \\circ \\sigma) \\,\n",
    "           [R (a - c) + \\hat Y, \\, \\hat Y]\n",
    "           \\, , \\;\n",
    "           u'(a)\n",
    "     \\right\\} \\tag{16.1}\n",
    "$$\n",
    "\n",
    "It [can be shown that](https://python.quantecon.org/ifp.html)\n",
    "\n",
    "1. iterating with $ K $ computes an optimal policy and  \n",
    "1. if $ \\sigma $ is increasing in its first argument, then so is $ K\\sigma $  \n",
    "\n",
    "\n",
    "Hence below we always assume that $ \\sigma $ is increasing in its first argument.\n",
    "\n",
    "The EGM is a technique for computing the update $ K\\sigma $ given $ \\sigma $ along a grid of asset values.\n",
    "\n",
    "Notice that, since $ u'(a) \\to \\infty $ as $ a \\downarrow 0 $, the second term in\n",
    "the max in [(16.1)](#equation-eq-kaper) dominates for sufficiently small $ a $.\n",
    "\n",
    "Also, again using [(16.1)](#equation-eq-kaper), we have $ c=a $ for all such $ a $.\n",
    "\n",
    "Hence, for sufficiently small $ a $,\n",
    "\n",
    "$$\n",
    "u'(a) \\geq\n",
    "   \\beta R \\, \\mathbb E_y (u' \\circ \\sigma) \\,\n",
    "           [\\hat Y, \\, \\hat Y]\n",
    "$$\n",
    "\n",
    "Equality holds at $ \\bar a(y) $ given by\n",
    "\n",
    "$$\n",
    "\\bar a (y) =\n",
    "   (u')^{-1}\n",
    "   \\left\\{\n",
    "       \\beta R \\, \\mathbb E_y (u' \\circ \\sigma) \\,\n",
    "               [\\hat Y, \\, \\hat Y]\n",
    "   \\right\\}\n",
    "$$\n",
    "\n",
    "We can now write\n",
    "\n",
    "$$\n",
    "u'(c)\n",
    "    = \\begin{cases}\n",
    "        \\beta R \\, \\mathbb E_y (u' \\circ \\sigma) \\,\n",
    "               [R (a - c) + \\hat Y, \\, \\hat Y]\n",
    "               & \\text{if } a > \\bar a (y) \\\\\n",
    "        u'(a)  & \\text{if } a \\leq \\bar a (y)\n",
    "    \\end{cases}\n",
    "$$\n",
    "\n",
    "Equivalently, we can state that the $ c $ satisfying $ c = (K\\sigma)(a, y) $ obeys\n",
    "\n",
    "\n",
    "<a id='equation-eq-oro'></a>\n",
    "$$\n",
    "c = \\begin{cases}\n",
    "        (u')^{-1}\n",
    "        \\left\\{\n",
    "            \\beta R \\, \\mathbb E_y (u' \\circ \\sigma) \\,\n",
    "               [R (a - c) + \\hat Y, \\, \\hat Y]\n",
    "        \\right\\}\n",
    "               & \\text{if } a > \\bar a (y) \\\\\n",
    "            a  & \\text{if } a \\leq \\bar a (y)\n",
    "    \\end{cases} \\tag{16.2}\n",
    "$$\n",
    "\n",
    "We begin with an *exogenous* grid of saving values $ 0 = s_0 < \\ldots < s_{N-1} $\n",
    "\n",
    "Using the exogenous savings grid, and a fixed value of $ y $, we create an *endogenous* asset grid\n",
    "$ a_0, \\ldots, a_{N-1} $ and a consumption grid $ c_0, \\ldots, c_{N-1} $ as follows.\n",
    "\n",
    "First we set $ a_0 = c_0 = 0 $, since zero consumption is an optimal (in fact the only) choice when $ a=0 $.\n",
    "\n",
    "Then, for $ i > 0 $, we compute\n",
    "\n",
    "\n",
    "<a id='equation-eq-kaperc'></a>\n",
    "$$\n",
    "c_i\n",
    "    = (u')^{-1}\n",
    "    \\left\\{ \n",
    "        \\beta R \\, \\mathbb E_y (u' \\circ \\sigma) \\,\n",
    "               [R s_i + \\hat Y, \\, \\hat Y]\n",
    "     \\right\\}\n",
    "     \\quad \\text{for all } i \\tag{16.3}\n",
    "$$\n",
    "\n",
    "and we set\n",
    "\n",
    "$$\n",
    "a_i = s_i + c_i\n",
    "$$\n",
    "\n",
    "We claim that each pair $ a_i, c_i $ obeys [(16.2)](#equation-eq-oro).\n",
    "\n",
    "Indeed, since $ s_i > 0 $, choosing $ c_i $ according to [(16.3)](#equation-eq-kaperc) gives\n",
    "\n",
    "$$\n",
    "c_i\n",
    "    = (u')^{-1}\n",
    "    \\left\\{ \n",
    "        \\beta R \\, \\mathbb E_y (u' \\circ \\sigma) \\,\n",
    "               [R s_i + \\hat Y, \\, \\hat Y]\n",
    "     \\right\\}\n",
    "     \\geq \\bar a(y)\n",
    "$$\n",
    "\n",
    "where the inequality uses the fact that $ \\sigma $ is increasing in its first argument.\n",
    "\n",
    "If we now take $ a_i = s_i + c_i $ we get $ a_i > \\bar a(y) $, so the pair $ (a_i, c_i) $ satisfies\n",
    "\n",
    "$$\n",
    "c_i\n",
    "    = (u')^{-1}\n",
    "    \\left\\{ \n",
    "        \\beta R \\, \\mathbb E_y (u' \\circ \\sigma) \\,\n",
    "               [R (a_i - c_i) + \\hat Y, \\, \\hat Y]\n",
    "     \\right\\}\n",
    "     \\quad \\text{and} \\quad a_i > \\bar a(y)\n",
    "$$\n",
    "\n",
    "Hence [(16.2)](#equation-eq-oro) holds.\n",
    "\n",
    "We are now ready to iterate with $ K $."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b8d56689",
   "metadata": {},
   "source": [
    "### JAX version\n",
    "\n",
    "First we define a vectorized operator $ K $ based on the EGM.\n",
    "\n",
    "Notice in the code below that\n",
    "\n",
    "- we avoid all loops and any mutation of arrays  \n",
    "- the function is pure (no globals, no mutation of inputs)  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6e4fb66e",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def K_egm(a_in, σ_in, constants, sizes, arrays):\n",
    "    \"\"\"\n",
    "    The vectorized operator K using EGM.\n",
    "\n",
    "    \"\"\"\n",
    "    \n",
    "    # Unpack\n",
    "    β, R, γ = constants\n",
    "    s_size, y_size = sizes\n",
    "    s_grid, y_grid, P = arrays\n",
    "    \n",
    "    def u_prime(c):\n",
    "        return c**(-γ)\n",
    "\n",
    "    def u_prime_inv(u):\n",
    "            return u**(-1/γ)\n",
    "\n",
    "    # Linearly interpolate σ(a, y)\n",
    "    def σ(a, y):\n",
    "        return jnp.interp(a, a_in[:, y], σ_in[:, y])\n",
    "    σ_vec = jnp.vectorize(σ)\n",
    "\n",
    "    # Broadcast and vectorize\n",
    "    y_hat = jnp.reshape(y_grid, (1, 1, y_size))\n",
    "    y_hat_idx = jnp.reshape(jnp.arange(y_size), (1, 1, y_size))\n",
    "    s = jnp.reshape(s_grid, (s_size, 1, 1))\n",
    "    P = jnp.reshape(P, (1, y_size, y_size))\n",
    "    \n",
    "    # Evaluate consumption choice\n",
    "    a_next = R * s + y_hat\n",
    "    σ_next = σ_vec(a_next, y_hat_idx)\n",
    "    up = u_prime(σ_next)\n",
    "    E = jnp.sum(up * P, axis=-1)\n",
    "    c = u_prime_inv(β * R * E)\n",
    "\n",
    "    # Set up a column vector with zero in the first row and ones elsewhere\n",
    "    e_0 = jnp.ones(s_size) - jnp.identity(s_size)[:, 0]\n",
    "    e_0 = jnp.reshape(e_0, (s_size, 1))\n",
    "\n",
    "    # The policy is computed consumption with the first row set to zero\n",
    "    σ_out = c * e_0\n",
    "\n",
    "    # Compute a_out by a = s + c\n",
    "    a_out = np.reshape(s_grid, (s_size, 1)) + σ_out\n",
    "    \n",
    "    return a_out, σ_out"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5938cd24",
   "metadata": {},
   "source": [
    "Then we use `jax.jit` to compile $ K $.\n",
    "\n",
    "We use `static_argnums` to allow a recompile whenever `sizes` changes, since the compiler likes to specialize on shapes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "620f7986",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "K_egm_jax = jax.jit(K_egm, static_argnums=(3,))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "de746674",
   "metadata": {},
   "source": [
    "Next we define a successive approximator that repeatedly applies $ K $."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a835a32b",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def successive_approx_jax(model,        \n",
    "            tol=1e-5,\n",
    "            max_iter=100_000,\n",
    "            verbose=True,\n",
    "            print_skip=25):\n",
    "\n",
    "    # Unpack\n",
    "    constants, sizes, arrays = model\n",
    "    β, R, γ = constants\n",
    "    s_size, y_size = sizes\n",
    "    s_grid, y_grid, P = arrays\n",
    "    \n",
    "    # Initial condition is to consume all in every state\n",
    "    σ_init = jnp.repeat(s_grid, y_size)\n",
    "    σ_init = jnp.reshape(σ_init, (s_size, y_size))\n",
    "    a_init = jnp.copy(σ_init)\n",
    "    a_vec, σ_vec = a_init, σ_init\n",
    "    \n",
    "    i = 0\n",
    "    error = tol + 1\n",
    "\n",
    "    while i < max_iter and error > tol:\n",
    "        a_new, σ_new = K_egm_jax(a_vec, σ_vec, constants, sizes, arrays)    \n",
    "        error = jnp.max(jnp.abs(σ_vec - σ_new))\n",
    "        i += 1\n",
    "        if verbose and i % print_skip == 0:\n",
    "            print(f\"Error at iteration {i} is {error}.\")\n",
    "        a_vec, σ_vec = jnp.copy(a_new), jnp.copy(σ_new)\n",
    "\n",
    "    if error > tol:\n",
    "        print(\"Failed to converge!\")\n",
    "    else:\n",
    "        print(f\"\\nConverged in {i} iterations.\")\n",
    "\n",
    "    return a_new, σ_new"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1b0df3df",
   "metadata": {},
   "source": [
    "### Numba version\n",
    "\n",
    "Below we provide a second set of code, which solves the same model with Numba.\n",
    "\n",
    "The purpose of this code is to cross-check our results from the JAX version, as\n",
    "well as to do a runtime comparison.\n",
    "\n",
    "Most readers will want to skip ahead to the next section, where we solve the\n",
    "model and run the cross-check."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "711ccfda",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "@numba.jit\n",
    "def K_egm_nb(a_in, σ_in, constants, sizes, arrays):\n",
    "    \"\"\"\n",
    "    The operator K using Numba.\n",
    "\n",
    "    \"\"\"\n",
    "    \n",
    "    # Simplify names\n",
    "    β, R, γ = constants\n",
    "    s_size, y_size = sizes\n",
    "    s_grid, y_grid, P = arrays\n",
    "\n",
    "    def u_prime(c):\n",
    "        return c**(-γ)\n",
    "\n",
    "    def u_prime_inv(u):\n",
    "        return u**(-1/γ)\n",
    "\n",
    "    # Linear interpolation of policy using endogenous grid\n",
    "    def σ(a, z):\n",
    "        return np.interp(a, a_in[:, z], σ_in[:, z])\n",
    "    \n",
    "    # Allocate memory for new consumption array\n",
    "    σ_out = np.zeros_like(σ_in)\n",
    "    a_out = np.zeros_like(σ_out)\n",
    "    \n",
    "    for i, s in enumerate(s_grid[1:]):\n",
    "        i += 1\n",
    "        for z in range(y_size):\n",
    "            expect = 0.0\n",
    "            for z_hat in range(y_size):\n",
    "                expect += u_prime(σ(R * s + y_grid[z_hat], z_hat)) * \\\n",
    "                            P[z, z_hat]\n",
    "            c = u_prime_inv(β * R * expect)\n",
    "            σ_out[i, z] = c\n",
    "            a_out[i, z] = s + c\n",
    "    \n",
    "    return a_out, σ_out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c6009e04",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def successive_approx_numba(model,        # Class with model information\n",
    "                              tol=1e-5,\n",
    "                              max_iter=100_000,\n",
    "                              verbose=True,\n",
    "                              print_skip=25):\n",
    "\n",
    "    # Unpack\n",
    "    constants, sizes, arrays = model\n",
    "    s_size, y_size = sizes\n",
    "    # make NumPy versions of arrays\n",
    "    arrays = tuple(map(np.array, arrays))\n",
    "    s_grid, y_grid, P = arrays\n",
    "    \n",
    "    σ_init = np.repeat(s_grid, y_size)\n",
    "    σ_init = np.reshape(σ_init, (s_size, y_size))\n",
    "    a_init = np.copy(σ_init)\n",
    "    a_vec, σ_vec = a_init, σ_init\n",
    "    \n",
    "    # Set up loop\n",
    "    i = 0\n",
    "    error = tol + 1\n",
    "\n",
    "    while i < max_iter and error > tol:\n",
    "        a_new, σ_new = K_egm_nb(a_vec, σ_vec, constants, sizes, arrays)\n",
    "        error = np.max(np.abs(σ_vec - σ_new))\n",
    "        i += 1\n",
    "        if verbose and i % print_skip == 0:\n",
    "            print(f\"Error at iteration {i} is {error}.\")\n",
    "        a_vec, σ_vec = np.copy(a_new), np.copy(σ_new)\n",
    "\n",
    "    if error > tol:\n",
    "        print(\"Failed to converge!\")\n",
    "    else:\n",
    "        print(f\"\\nConverged in {i} iterations.\")\n",
    "\n",
    "    return a_new, σ_new"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2a98ff83",
   "metadata": {},
   "source": [
    "## Solutions\n",
    "\n",
    "Here we solve the IFP with JAX and Numba.\n",
    "\n",
    "We will compare both the outputs and the execution time."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4328d581",
   "metadata": {},
   "source": [
    "### Outputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9f7bcd81",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "model = ifp()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1b0ede14",
   "metadata": {},
   "source": [
    "Here’s a first run of the JAX code."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "494d0359",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "a_star_egm_jax, σ_star_egm_jax = successive_approx_jax(model,\n",
    "                                        print_skip=1000)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f10bb323",
   "metadata": {},
   "source": [
    "Next let’s solve the same IFP with Numba."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11e7b64c",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "a_star_egm_nb, σ_star_egm_nb = successive_approx_numba(model,\n",
    "                                    print_skip=1000)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7e46a0f5",
   "metadata": {},
   "source": [
    "Now let’s check the outputs in a plot to make sure they are the same."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a17dcd20",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "constants, sizes, arrays = model\n",
    "β, R, γ = constants\n",
    "s_size, y_size = sizes\n",
    "s_grid, y_grid, P = arrays\n",
    "\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "\n",
    "for z in (0, y_size-1):\n",
    "    ax.plot(a_star_egm_nb[:, z], \n",
    "            σ_star_egm_nb[:, z], \n",
    "            '--', lw=2,\n",
    "            label=f\"Numba EGM: consumption when $z={z}$\")\n",
    "    ax.plot(a_star_egm_jax[:, z], \n",
    "            σ_star_egm_jax[:, z], \n",
    "            label=f\"JAX EGM: consumption when $z={z}$\")\n",
    "\n",
    "ax.set_xlabel('asset')\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d112fb17",
   "metadata": {},
   "source": [
    "### Timing\n",
    "\n",
    "Now let’s compare execution time of the two methods."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "59b100af",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "start = time()\n",
    "a_star_egm_jax, σ_star_egm_jax = successive_approx_jax(model,\n",
    "                                         print_skip=1000)\n",
    "jax_time_without_compile = time() - start\n",
    "print(\"Jax execution time = \", jax_time_without_compile)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "07de715b",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "start = time()\n",
    "a_star_egm_nb, σ_star_egm_nb = successive_approx_numba(model,\n",
    "                                         print_skip=1000)\n",
    "numba_time_without_compile = time() - start\n",
    "print(\"Numba execution time = \", numba_time_without_compile)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eaff79de",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "jax_time_without_compile / numba_time_without_compile"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7551dab1",
   "metadata": {},
   "source": [
    "The JAX code is significantly faster, as expected.\n",
    "\n",
    "This difference will increase when more features (and state variables) are added\n",
    "to the model."
   ]
  }
 ],
 "metadata": {
  "date": 1729509368.7936976,
  "filename": "ifp_egm.md",
  "kernelspec": {
   "display_name": "Python",
   "language": "python3",
   "name": "python3"
  },
  "title": "Endogenous Grid Method"
 },
 "nbformat": 4,
 "nbformat_minor": 5
}