opt_savings_1.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0a4181eb",
   "metadata": {},
   "source": [
    "# Optimal Savings I: Value Function Iteration"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e65d60fa",
   "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).\n",
    "\n",
    "In addition to JAX and Anaconda, this lecture will need the following libraries:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "75173810",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "!pip install quantecon"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "39620bd9",
   "metadata": {},
   "source": [
    "We will use the following imports:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b15f4aef",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "import quantecon as qe\n",
    "import numpy as np\n",
    "import jax\n",
    "import jax.numpy as jnp\n",
    "import matplotlib.pyplot as plt\n",
    "from collections import namedtuple\n",
    "from time import time"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "43ab3c35",
   "metadata": {},
   "source": [
    "Let’s check the GPU we are running"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7f132f2c",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "!nvidia-smi"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fed269fb",
   "metadata": {},
   "source": [
    "We’ll use 64 bit floats to gain extra precision."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "60b7625c",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "jax.config.update(\"jax_enable_x64\", True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5a9460cb",
   "metadata": {},
   "source": [
    "## Overview\n",
    "\n",
    "We consider an optimal savings problem with CRRA utility and budget constraint\n",
    "\n",
    "$$\n",
    "W_{t+1} + C_t \\leq R W_t + Y_t\n",
    "$$\n",
    "\n",
    "where\n",
    "\n",
    "- $ C_t $ is consumption and $ C_t \\geq 0 $,  \n",
    "- $ W_t $ is wealth and $ W_t \\geq 0 $,  \n",
    "- $ R > 0 $ is a gross rate of return, and  \n",
    "- $ (Y_t) $ is labor income.  \n",
    "\n",
    "\n",
    "We assume below that labor income is a discretized AR(1) process.\n",
    "\n",
    "The Bellman equation is\n",
    "\n",
    "$$\n",
    "v(w) = \\max_{0 \\leq w' \\leq Rw + y}\n",
    "    \\left\\{\n",
    "        u(Rw + y - w') + β \\sum_{y'} v(w', y') Q(y, y') \n",
    "    \\right\\}\n",
    "$$\n",
    "\n",
    "where\n",
    "\n",
    "$$\n",
    "u(c) = \\frac{c^{1-\\gamma}}{1-\\gamma}\n",
    "$$\n",
    "\n",
    "In the code we use the function\n",
    "\n",
    "$$\n",
    "B((w, y), w', v) = u(Rw + y - w') + β \\sum_{y'} v(w', y') Q(y, y').\n",
    "$$\n",
    "\n",
    "the encapsulate the right hand side of the Bellman equation."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c563d8d0",
   "metadata": {},
   "source": [
    "## Starting with NumPy\n",
    "\n",
    "Let’s start with a standard NumPy version running on the CPU.\n",
    "\n",
    "Starting with this traditional approach will allow us to record the speed gain\n",
    "associated with switching to JAX.\n",
    "\n",
    "(NumPy operations are similar to MATLAB operations, so this also serves as a\n",
    "rough comparison with MATLAB.)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6ae8e524",
   "metadata": {},
   "source": [
    "### Functions and operators\n",
    "\n",
    "The following function contains default parameters and returns tuples that\n",
    "contain the key computational components of the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "35f3b9bf",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def create_consumption_model(R=1.01,                    # Gross interest rate\n",
    "                             β=0.98,                    # Discount factor\n",
    "                             γ=2,                       # CRRA parameter\n",
    "                             w_min=0.01,                # Min wealth\n",
    "                             w_max=5.0,                 # Max wealth\n",
    "                             w_size=150,                # Grid side\n",
    "                             ρ=0.9, ν=0.1, y_size=100): # Income parameters\n",
    "    \"\"\"\n",
    "    A function that takes in parameters and returns parameters and grids \n",
    "    for the optimal savings problem.\n",
    "    \"\"\"\n",
    "    # Build grids and transition probabilities\n",
    "    w_grid = np.linspace(w_min, w_max, w_size)\n",
    "    mc = qe.tauchen(n=y_size, rho=ρ, sigma=ν)\n",
    "    y_grid, Q = np.exp(mc.state_values), mc.P\n",
    "    # Pack and return\n",
    "    params = β, R, γ\n",
    "    sizes = w_size, y_size\n",
    "    arrays = w_grid, y_grid, Q\n",
    "    return params, sizes, arrays"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f9d78837",
   "metadata": {},
   "source": [
    "(The function returns sizes of arrays because we use them later to help\n",
    "compile functions in JAX.)\n",
    "\n",
    "To produce efficient NumPy code, we will use a vectorized approach.\n",
    "\n",
    "The first step is to create the right hand side of the Bellman equation as a\n",
    "multi-dimensional array with dimensions over all states and controls."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e3be7dff",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def B(v, params, sizes, arrays):\n",
    "    \"\"\"\n",
    "    A vectorized version of the right-hand side of the Bellman equation\n",
    "    (before maximization), which is a 3D array representing\n",
    "\n",
    "        B(w, y, w′) = u(Rw + y - w′) + β Σ_y′ v(w′, y′) Q(y, y′)\n",
    "\n",
    "    for all (w, y, w′).\n",
    "    \"\"\"\n",
    "\n",
    "    # Unpack\n",
    "    β, R, γ = params\n",
    "    w_size, y_size = sizes\n",
    "    w_grid, y_grid, Q = arrays\n",
    "\n",
    "    # Compute current rewards r(w, y, wp) as array r[i, j, ip]\n",
    "    w  = np.reshape(w_grid, (w_size, 1, 1))    # w[i]   ->  w[i, j, ip]\n",
    "    y  = np.reshape(y_grid, (1, y_size, 1))    # z[j]   ->  z[i, j, ip]\n",
    "    wp = np.reshape(w_grid, (1, 1, w_size))    # wp[ip] -> wp[i, j, ip]\n",
    "    c = R * w + y - wp\n",
    "\n",
    "    # Calculate continuation rewards at all combinations of (w, y, wp)\n",
    "    v = np.reshape(v, (1, 1, w_size, y_size))  # v[ip, jp] -> v[i, j, ip, jp]\n",
    "    Q = np.reshape(Q, (1, y_size, 1, y_size))  # Q[j, jp]  -> Q[i, j, ip, jp]\n",
    "    EV = np.sum(v * Q, axis=3)                 # sum over last index jp\n",
    "\n",
    "    # Compute the right-hand side of the Bellman equation\n",
    "    return np.where(c > 0, c**(1-γ)/(1-γ) + β * EV, -np.inf)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "54d22d21",
   "metadata": {},
   "source": [
    "Here are two functions we need for value function iteration.\n",
    "\n",
    "The first is the Bellman operator.\n",
    "\n",
    "The second computes a $ v $-greedy policy given $ v $ (i.e., the policy that\n",
    "maximizes the right-hand side of the Bellman equation.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "82c2562f",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def T(v, params, sizes, arrays):\n",
    "    \"The Bellman operator.\"\n",
    "    return np.max(B(v, params, sizes, arrays), axis=2)\n",
    "\n",
    "def get_greedy(v, params, sizes, arrays):\n",
    "    \"Computes a v-greedy policy, returned as a set of indices.\"\n",
    "    return np.argmax(B(v, params, sizes, arrays), axis=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4595c044",
   "metadata": {},
   "source": [
    "### Value function iteration\n",
    "\n",
    "Here’s a routine that performs value function iteration."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bd22d0f0",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def value_function_iteration(model, max_iter=10_000, tol=1e-5):\n",
    "    params, sizes, arrays = model\n",
    "    v = np.zeros(sizes)\n",
    "    i, error = 0, tol + 1\n",
    "    while error > tol and i < max_iter:\n",
    "        v_new = T(v, params, sizes, arrays)\n",
    "        error = np.max(np.abs(v_new - v))\n",
    "        i += 1\n",
    "        v = v_new\n",
    "    return v, get_greedy(v, params, sizes, arrays)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c9d5cd54",
   "metadata": {},
   "source": [
    "Now we create an instance, unpack it, and test how long it takes to solve the\n",
    "model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f87c4512",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "model = create_consumption_model()\n",
    "# Unpack\n",
    "params, sizes, arrays = model\n",
    "β, R, γ = params\n",
    "w_size, y_size = sizes\n",
    "w_grid, y_grid, Q = arrays"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6e8d798f",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "print(\"Starting VFI.\")\n",
    "start = time()\n",
    "v_star, σ_star = value_function_iteration(model)\n",
    "numpy_with_compile = time() - start\n",
    "print(f\"VFI completed in {numpy_with_compile} seconds.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a3fb214e",
   "metadata": {},
   "source": [
    "Let’s run it again to eliminate compile time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f825c220",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "start = time()\n",
    "v_star, σ_star = value_function_iteration(model)\n",
    "numpy_without_compile = time() - start\n",
    "print(f\"VFI completed in {numpy_without_compile} seconds.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "412f3f58",
   "metadata": {},
   "source": [
    "Here’s a plot of the policy function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0899f61f",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "ax.plot(w_grid, w_grid, \"k--\", label=\"45\")\n",
    "ax.plot(w_grid, w_grid[σ_star[:, 1]], label=\"$\\\\sigma^*(\\cdot, y_1)$\")\n",
    "ax.plot(w_grid, w_grid[σ_star[:, -1]], label=\"$\\\\sigma^*(\\cdot, y_N)$\")\n",
    "ax.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "038d7d13",
   "metadata": {},
   "source": [
    "## Switching to JAX\n",
    "\n",
    "To switch over to JAX, we change `np` to `jnp` throughout and add some\n",
    "`jax.jit` requests."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3725dc2a",
   "metadata": {},
   "source": [
    "### Functions and operators\n",
    "\n",
    "We redefine `create_consumption_model` to produce JAX arrays.\n",
    "\n",
    "\n",
    "<a id='prgm-create-consumption-model'></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e27f7df0",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def create_consumption_model(R=1.01,                    # Gross interest rate\n",
    "                             β=0.98,                    # Discount factor\n",
    "                             γ=2,                       # CRRA parameter\n",
    "                             w_min=0.01,                # Min wealth\n",
    "                             w_max=5.0,                 # Max wealth\n",
    "                             w_size=150,                # Grid side\n",
    "                             ρ=0.9, ν=0.1, y_size=100): # Income parameters\n",
    "    \"\"\"\n",
    "    A function that takes in parameters and returns parameters and grids \n",
    "    for the optimal savings problem.\n",
    "    \"\"\"\n",
    "    w_grid = jnp.linspace(w_min, w_max, w_size)\n",
    "    mc = qe.tauchen(n=y_size, rho=ρ, sigma=ν)\n",
    "    y_grid, Q = jnp.exp(mc.state_values), jax.device_put(mc.P)\n",
    "    sizes = w_size, y_size\n",
    "    return (β, R, γ), sizes, (w_grid, y_grid, Q)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9271aa27",
   "metadata": {},
   "source": [
    "The right hand side of the Bellman equation is the same as the NumPy version\n",
    "after switching `np` to `jnp`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c67e8781",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def B(v, params, sizes, arrays):\n",
    "    \"\"\"\n",
    "    A vectorized version of the right-hand side of the Bellman equation\n",
    "    (before maximization), which is a 3D array representing\n",
    "\n",
    "        B(w, y, w′) = u(Rw + y - w′) + β Σ_y′ v(w′, y′) Q(y, y′)\n",
    "\n",
    "    for all (w, y, w′).\n",
    "    \"\"\"\n",
    "\n",
    "    # Unpack\n",
    "    β, R, γ = params\n",
    "    w_size, y_size = sizes\n",
    "    w_grid, y_grid, Q = arrays\n",
    "\n",
    "    # Compute current rewards r(w, y, wp) as array r[i, j, ip]\n",
    "    w  = jnp.reshape(w_grid, (w_size, 1, 1))    # w[i]   ->  w[i, j, ip]\n",
    "    y  = jnp.reshape(y_grid, (1, y_size, 1))    # z[j]   ->  z[i, j, ip]\n",
    "    wp = jnp.reshape(w_grid, (1, 1, w_size))    # wp[ip] -> wp[i, j, ip]\n",
    "    c = R * w + y - wp\n",
    "\n",
    "    # Calculate continuation rewards at all combinations of (w, y, wp)\n",
    "    v = jnp.reshape(v, (1, 1, w_size, y_size))  # v[ip, jp] -> v[i, j, ip, jp]\n",
    "    Q = jnp.reshape(Q, (1, y_size, 1, y_size))  # Q[j, jp]  -> Q[i, j, ip, jp]\n",
    "    EV = jnp.sum(v * Q, axis=3)                 # sum over last index jp\n",
    "\n",
    "    # Compute the right-hand side of the Bellman equation\n",
    "    return jnp.where(c > 0, c**(1-γ)/(1-γ) + β * EV, -jnp.inf)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "649994b3",
   "metadata": {},
   "source": [
    "Some readers might be concerned that we are creating high dimensional arrays,\n",
    "leading to inefficiency.\n",
    "\n",
    "Could they be avoided by more careful vectorization?\n",
    "\n",
    "In fact this is not necessary: this function will be JIT-compiled by JAX, and\n",
    "the JIT compiler will optimize compiled code to minimize memory use."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d3f1305b",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "B = jax.jit(B, static_argnums=(2,))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5968158c",
   "metadata": {},
   "source": [
    "In the call above, we indicate to the compiler that `sizes` is static, so the\n",
    "compiler can parallelize optimally while taking array sizes as fixed.\n",
    "\n",
    "The Bellman operator $ T $ can be implemented by"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f12b7ecc",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def T(v, params, sizes, arrays):\n",
    "    \"The Bellman operator.\"\n",
    "    return jnp.max(B(v, params, sizes, arrays), axis=2)\n",
    "\n",
    "T = jax.jit(T, static_argnums=(2,))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a8fdcec7",
   "metadata": {},
   "source": [
    "The next function computes a $ v $-greedy policy given $ v $ (i.e., the policy that\n",
    "maximizes the right-hand side of the Bellman equation.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dcf76011",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def get_greedy(v, params, sizes, arrays):\n",
    "    \"Computes a v-greedy policy, returned as a set of indices.\"\n",
    "    return jnp.argmax(B(v, params, sizes, arrays), axis=2)\n",
    "\n",
    "get_greedy = jax.jit(get_greedy, static_argnums=(2,))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "892a6a38",
   "metadata": {},
   "source": [
    "### Successive approximation\n",
    "\n",
    "Now we define a solver that implements VFI.\n",
    "\n",
    "We could use the one we built for NumPy above, after changing `np` to `jnp`.\n",
    "\n",
    "Alternatively, we can push a bit harder and write a compiled version using\n",
    "`jax.lax.while_loop`.\n",
    "\n",
    "This will give us just a bit more speed.\n",
    "\n",
    "The first step is to write a compiled successive approximation routine that\n",
    "performs fixed point iteration on some given function `T`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c2f7d25e",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def successive_approx_jax(T,                     # Operator (callable)\n",
    "                          x_0,                   # Initial condition                \n",
    "                          tolerance=1e-6,        # Error tolerance\n",
    "                          max_iter=10_000):      # Max iteration bound\n",
    "    def body_fun(k_x_err):\n",
    "        k, x, error = k_x_err\n",
    "        x_new = T(x)\n",
    "        error = jnp.max(jnp.abs(x_new - x))\n",
    "        return k + 1, x_new, error\n",
    "\n",
    "    def cond_fun(k_x_err):\n",
    "        k, x, error = k_x_err\n",
    "        return jnp.logical_and(error > tolerance, k < max_iter)\n",
    "\n",
    "    k, x, error = jax.lax.while_loop(cond_fun, body_fun, \n",
    "                                    (1, x_0, tolerance + 1))\n",
    "    return x\n",
    "\n",
    "successive_approx_jax = \\\n",
    "    jax.jit(successive_approx_jax, static_argnums=(0,))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "778f5cee",
   "metadata": {},
   "source": [
    "Our value function iteration routine calls `successive_approx_jax` while passing\n",
    "in the Bellman operator."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "143bee36",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def value_function_iteration(model, tol=1e-5):\n",
    "    params, sizes, arrays = model\n",
    "    vz = jnp.zeros(sizes)\n",
    "    _T = lambda v: T(v, params, sizes, arrays)\n",
    "    v_star = successive_approx_jax(_T, vz, tolerance=tol)\n",
    "    return v_star, get_greedy(v_star, params, sizes, arrays)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cefc3fa1",
   "metadata": {},
   "source": [
    "### Timing\n",
    "\n",
    "Let’s create an instance and unpack it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e84a3fc0",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "model = create_consumption_model()\n",
    "# Unpack\n",
    "params, sizes, arrays = model\n",
    "β, R, γ = params\n",
    "w_size, y_size = sizes\n",
    "w_grid, y_grid, Q = arrays"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "52e39ae9",
   "metadata": {},
   "source": [
    "Let’s see how long it takes to solve this model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2fba84a6",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "print(\"Starting VFI using vectorization.\")\n",
    "start = time()\n",
    "v_star_jax, σ_star_jax = value_function_iteration(model)\n",
    "jax_with_compile = time() - start\n",
    "print(f\"VFI completed in {jax_with_compile} seconds.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "82603304",
   "metadata": {},
   "source": [
    "Let’s run it again to eliminate compile time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10ed04ab",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "start = time()\n",
    "v_star_jax, σ_star_jax = value_function_iteration(model)\n",
    "jax_without_compile = time() - start\n",
    "print(f\"VFI completed in {jax_without_compile} seconds.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f51ff3a6",
   "metadata": {},
   "source": [
    "The relative speed gain is"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e627690d",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "print(f\"Relative speed gain = {numpy_without_compile / jax_without_compile}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9cdaab34",
   "metadata": {},
   "source": [
    "This is an impressive speed up and in fact we can do better still by switching\n",
    "to alternative algorithms that are better suited to parallelization.\n",
    "\n",
    "These algorithms are discussed in a [separate lecture](https://jax.quantecon.org/opt_savings_2.html)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0f761e11",
   "metadata": {},
   "source": [
    "## Switching to vmap\n",
    "\n",
    "Before we discuss alternative algorithms, let’s take another look at value\n",
    "function iteration.\n",
    "\n",
    "For this simple optimal savings problem, direct vectorization is relatively easy.\n",
    "\n",
    "In particular, it’s straightforward to express the right hand side of the\n",
    "Bellman equation as an array that stores evaluations of the function at every\n",
    "state and control.\n",
    "\n",
    "For more complex models direct vectorization can be much harder.\n",
    "\n",
    "For this reason, it helps to have another approach to fast JAX implementations\n",
    "up our sleeves.\n",
    "\n",
    "Here’s a version that\n",
    "\n",
    "1. writes the right hand side of the Bellman operator as a function of individual states and controls, and  \n",
    "1. applies `jax.vmap` on the outside to achieve a parallelized solution.  \n",
    "\n",
    "\n",
    "First let’s rewrite `B`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ff03b5e7",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def B(v, params, arrays, i, j, ip):\n",
    "    \"\"\"\n",
    "    The right-hand side of the Bellman equation before maximization, which takes\n",
    "    the form\n",
    "\n",
    "        B(w, y, w′) = u(Rw + y - w′) + β Σ_y′ v(w′, y′) Q(y, y′)\n",
    "\n",
    "    The indices are (i, j, ip) -> (w, y, w′).\n",
    "    \"\"\"\n",
    "    β, R, γ = params\n",
    "    w_grid, y_grid, Q = arrays\n",
    "    w, y, wp  = w_grid[i], y_grid[j], w_grid[ip]\n",
    "    c = R * w + y - wp\n",
    "    EV = jnp.sum(v[ip, :] * Q[j, :]) \n",
    "    return jnp.where(c > 0, c**(1-γ)/(1-γ) + β * EV, -jnp.inf)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b7e7634a",
   "metadata": {},
   "source": [
    "Now we successively apply `vmap` to simulate nested loops."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c1401578",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "B_1    = jax.vmap(B,   in_axes=(None, None, None, None, None, 0))\n",
    "B_2    = jax.vmap(B_1, in_axes=(None, None, None, None, 0,    None))\n",
    "B_vmap = jax.vmap(B_2, in_axes=(None, None, None, 0,    None, None))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e329aa68",
   "metadata": {},
   "source": [
    "Here’s the Bellman operator and the `get_greedy` functions for the `vmap` case."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a5352abd",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def T_vmap(v, params, sizes, arrays):\n",
    "    \"The Bellman operator.\"\n",
    "    w_size, y_size = sizes\n",
    "    w_indices, y_indices = jnp.arange(w_size), jnp.arange(y_size)\n",
    "    B_values = B_vmap(v, params, arrays, w_indices, y_indices, w_indices)\n",
    "    return jnp.max(B_values, axis=-1)\n",
    "\n",
    "T_vmap = jax.jit(T_vmap, static_argnums=(2,))\n",
    "\n",
    "def get_greedy_vmap(v, params, sizes, arrays):\n",
    "    \"Computes a v-greedy policy, returned as a set of indices.\"\n",
    "    w_size, y_size = sizes\n",
    "    w_indices, y_indices = jnp.arange(w_size), jnp.arange(y_size)\n",
    "    B_values = B_vmap(v, params, arrays, w_indices, y_indices, w_indices)\n",
    "    return jnp.argmax(B_values, axis=-1)\n",
    "\n",
    "get_greedy_vmap = jax.jit(get_greedy_vmap, static_argnums=(2,))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "68f9482e",
   "metadata": {},
   "source": [
    "Here’s the iteration routine."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cfbeed8d",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def value_iteration_vmap(model, tol=1e-5):\n",
    "    params, sizes, arrays = model\n",
    "    vz = jnp.zeros(sizes)\n",
    "    _T = lambda v: T_vmap(v, params, sizes, arrays)\n",
    "    v_star = successive_approx_jax(_T, vz, tolerance=tol)\n",
    "    return v_star, get_greedy(v_star, params, sizes, arrays)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "501fd3e0",
   "metadata": {},
   "source": [
    "Let’s see how long it takes to solve the model using the `vmap` method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2e1261e5",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "print(\"Starting VFI using vmap.\")\n",
    "start = time()\n",
    "v_star_vmap, σ_star_vmap = value_iteration_vmap(model)\n",
    "jax_vmap_with_compile = time() - start\n",
    "print(f\"VFI completed in {jax_vmap_with_compile} seconds.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6b04aa1e",
   "metadata": {},
   "source": [
    "Let’s run it again to get rid of compile time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "50a3eadf",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "start = time()\n",
    "v_star_vmap, σ_star_vmap = value_iteration_vmap(model)\n",
    "jax_vmap_without_compile = time() - start\n",
    "print(f\"VFI completed in {jax_vmap_without_compile} seconds.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "06d0a1a2",
   "metadata": {},
   "source": [
    "We need to make sure that we got the same result."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "516db0ac",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "print(jnp.allclose(v_star_vmap, v_star_jax))\n",
    "print(jnp.allclose(σ_star_vmap, σ_star_jax))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0d247b0c",
   "metadata": {},
   "source": [
    "Here’s the speed gain associated with switching from the NumPy version to JAX with `vmap`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cee90c91",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "print(f\"Relative speed = {numpy_without_compile/jax_vmap_without_compile}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1202027c",
   "metadata": {},
   "source": [
    "And here’s the comparison with the first JAX implementation (which used direct vectorization)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7c8fab4d",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "print(f\"Relative speed = {jax_without_compile / jax_vmap_without_compile}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "21cc6e2b",
   "metadata": {},
   "source": [
    "The execution times for the two JAX versions are relatively similar.\n",
    "\n",
    "However, as emphasized above, having a second method up our sleeves (i.e, the\n",
    "`vmap` approach) will be helpful when confronting dynamic programs with more\n",
    "sophisticated Bellman equations."
   ]
  }
 ],
 "metadata": {
  "date": 1729509369.27677,
  "filename": "opt_savings_1.md",
  "kernelspec": {
   "display_name": "Python",
   "language": "python3",
   "name": "python3"
  },
  "title": "Optimal Savings I: Value Function Iteration"
 },
 "nbformat": 4,
 "nbformat_minor": 5
}