cake_eating_numerical.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0355c95e",
   "metadata": {},
   "source": [
    "# Cake Eating: Numerical Methods"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "518c125c",
   "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",
    "This lecture is the extended JAX implementation of [this lecture](https://python.quantecon.org/cake_eating_numerical.html).\n",
    "\n",
    "Please refer that lecture for all background and notation.\n",
    "\n",
    "In addition to JAX and Anaconda, this lecture will need the following libraries:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "82c19e76",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "!pip install quantecon"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "73461742",
   "metadata": {},
   "source": [
    "We will use the following imports."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2772d468",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "import jax\n",
    "import jax.numpy as jnp\n",
    "import matplotlib.pyplot as plt\n",
    "from collections import namedtuple\n",
    "import time"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9d4caca9",
   "metadata": {},
   "source": [
    "Let’s check the GPU we are running"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5ebff432",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "!nvidia-smi"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "484b052f",
   "metadata": {},
   "source": [
    "## Reviewing the Model\n",
    "\n",
    "Recall in particular that the Bellman equation is\n",
    "\n",
    "\n",
    "<a id='equation-bellman-cen'></a>\n",
    "$$\n",
    "v(x) = \\max_{0\\leq c \\leq x} \\{u(c) + \\beta v(x-c)\\}\n",
    "\\quad \\text{for all } x \\geq 0. \\tag{16.1}\n",
    "$$\n",
    "\n",
    "where $ u $ is the CRRA utility function."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cff7a4e5",
   "metadata": {},
   "source": [
    "## Implementation using JAX\n",
    "\n",
    "The analytical solutions for the value function and optimal policy were found\n",
    "to be as follows."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2a182152",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "@jax.jit\n",
    "def c_star(x, β, γ):\n",
    "    return (1 - β ** (1/γ)) * x\n",
    "\n",
    "@jax.jit\n",
    "def v_star(x, β, γ):\n",
    "    return (1 - β**(1 / γ))**(-γ) * (x**(1-γ) / (1-γ))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e79f2516",
   "metadata": {},
   "source": [
    "Let’s define a model to represent the Cake Eating Problem."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "829658ad",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "CEM = namedtuple('CakeEatingModel',\n",
    "                    ('β', 'γ', 'x_grid', 'c_grid'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "befb17aa",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def create_cake_eating_model(β=0.96,           # discount factor\n",
    "                             γ=1.5,            # degree of relative risk aversion\n",
    "                             x_grid_min=1e-3,  # exclude zero for numerical stability\n",
    "                             x_grid_max=2.5,   # size of cake\n",
    "                             x_grid_size=200):\n",
    "    x_grid = jnp.linspace(x_grid_min, x_grid_max, x_grid_size)\n",
    "\n",
    "    # c_grid used for finding maximize function values using brute force\n",
    "    c_grid = jnp.linspace(x_grid_min, x_grid_max, 100*x_grid_size)\n",
    "    return CEM(β=β, γ=γ, x_grid=x_grid, c_grid=c_grid)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d33a5d72",
   "metadata": {},
   "source": [
    "Now let’s define the CRRA utility function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "810800a0",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "# Utility function\n",
    "@jax.jit\n",
    "def u(c, cem):\n",
    "    return (c ** (1 - cem.γ)) / (1 - cem.γ)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "51c9fd98",
   "metadata": {},
   "source": [
    "### The Bellman Operator\n",
    "\n",
    "We introduce the **Bellman operator** $ T $ that takes a function v as an\n",
    "argument and returns a new function $ Tv $ defined by\n",
    "\n",
    "$$\n",
    "Tv(x) = \\max_{0 \\leq c \\leq x} \\{u(c) + \\beta v(x - c)\\}\n",
    "$$\n",
    "\n",
    "From $ v $ we get $ Tv $, and applying $ T $ to this yields\n",
    "$ T^2 v := T (Tv) $ and so on.\n",
    "\n",
    "This is called **iterating with the Bellman operator** from initial guess\n",
    "$ v $."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e501164c",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "@jax.jit\n",
    "def state_action_value(x, c, v_array, ce):\n",
    "    \"\"\"\n",
    "    Right hand side of the Bellman equation given x and c.\n",
    "    * x: scalar element `x`\n",
    "    * c: c_grid, 1-D array\n",
    "    * v_array: value function array guess, 1-D array\n",
    "    * ce: Cake Eating Model instance\n",
    "    \"\"\"\n",
    "\n",
    "    return jnp.where(c <= x,\n",
    "                     u(c, ce) + ce.β * jnp.interp(x - c, ce.x_grid, v_array),\n",
    "                     -jnp.inf)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "485be622",
   "metadata": {},
   "source": [
    "In order to create a vectorized function using `state_action_value`, we use [jax.vmap](https://jax.readthedocs.io/en/latest/_autosummary/jax.vmap.html).\n",
    "This function returns a new vectorized version of the above function which is vectorized on the argument `x`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5c2396c2",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "state_action_value_vec = jax.vmap(state_action_value, (0, None, None, None))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ef15c222",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "@jax.jit\n",
    "def T(v, ce):\n",
    "    \"\"\"\n",
    "    The Bellman operator. Updates the guess of the value function.\n",
    "\n",
    "    * ce: Cake Eating Model instance\n",
    "    * v: value function array guess, 1-D array\n",
    "\n",
    "    \"\"\"\n",
    "    return jnp.max(state_action_value_vec(ce.x_grid, ce.c_grid, v, ce), axis=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "929b1623",
   "metadata": {},
   "source": [
    "Let’s start by creating a Cake Eating Model instance using the default parameterization."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f0229093",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "ce = create_cake_eating_model()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c40d7b39",
   "metadata": {},
   "source": [
    "Now let’s see the iteration of the value function in action.\n",
    "\n",
    "We start from guess $ v $ given by $ v(x) = u(x) $ for every\n",
    "$ x $ grid point."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5aecbae0",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "x_grid = ce.x_grid\n",
    "v = u(x_grid, ce)       # Initial guess\n",
    "n = 12                 # Number of iterations\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "\n",
    "ax.plot(x_grid, v, color=plt.cm.jet(0),\n",
    "        lw=2, alpha=0.6, label='Initial guess')\n",
    "\n",
    "for i in range(n):\n",
    "    v = T(v, ce)  # Apply the Bellman operator\n",
    "    ax.plot(x_grid, v, color=plt.cm.jet(i / n), lw=2, alpha=0.6)\n",
    "\n",
    "ax.legend()\n",
    "ax.set_ylabel('value', fontsize=12)\n",
    "ax.set_xlabel('cake size $x$', fontsize=12)\n",
    "ax.set_title('Value function iterations')\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "75d7f9da",
   "metadata": {},
   "source": [
    "Let’s introduce a wrapper function called `compute_value_function`\n",
    "that iterates until some convergence conditions are satisfied."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "acfbb8a7",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def compute_value_function(ce,\n",
    "                           tol=1e-4,\n",
    "                           max_iter=1000,\n",
    "                           verbose=True,\n",
    "                           print_skip=25):\n",
    "\n",
    "    # Set up loop\n",
    "    v = jnp.zeros(len(ce.x_grid)) # Initial guess\n",
    "    i = 0\n",
    "    error = tol + 1\n",
    "\n",
    "    while i < max_iter and error > tol:\n",
    "        v_new = T(v, ce)\n",
    "\n",
    "        error = jnp.max(jnp.abs(v - v_new))\n",
    "        i += 1\n",
    "\n",
    "        if verbose and i % print_skip == 0:\n",
    "            print(f\"Error at iteration {i} is {error}.\")\n",
    "\n",
    "        v = v_new\n",
    "\n",
    "    if error > tol:\n",
    "        print(\"Failed to converge!\")\n",
    "    elif verbose:\n",
    "        print(f\"\\nConverged in {i} iterations.\")\n",
    "\n",
    "    return v_new"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8e13e06b",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "in_time = time.time()\n",
    "v_jax = compute_value_function(ce)\n",
    "jax_time = time.time() - in_time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5851cbba",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "\n",
    "ax.plot(x_grid, v_jax, label='Approximate value function')\n",
    "ax.set_ylabel('$V(x)$', fontsize=12)\n",
    "ax.set_xlabel('$x$', fontsize=12)\n",
    "ax.set_title('Value function')\n",
    "ax.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ad382e7f",
   "metadata": {},
   "source": [
    "Next let’s compare it to the analytical solution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b7030243",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "v_analytical = v_star(ce.x_grid, ce.β, ce.γ)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3d3f60c8",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "\n",
    "ax.plot(x_grid, v_analytical, label='analytical solution')\n",
    "ax.plot(x_grid, v_jax, label='numerical solution')\n",
    "ax.set_ylabel('$V(x)$', fontsize=12)\n",
    "ax.set_xlabel('$x$', fontsize=12)\n",
    "ax.legend()\n",
    "ax.set_title('Comparison between analytical and numerical value functions')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d73a7a7c",
   "metadata": {},
   "source": [
    "### Policy Function\n",
    "\n",
    "Recall that the optimal consumption policy was shown to be\n",
    "\n",
    "$$\n",
    "\\sigma^*(x) = \\left(1-\\beta^{1/\\gamma} \\right) x\n",
    "$$\n",
    "\n",
    "Let’s see if our numerical results lead to something similar.\n",
    "\n",
    "Our numerical strategy will be to compute\n",
    "\n",
    "$$\n",
    "\\sigma(x) = \\arg \\max_{0 \\leq c \\leq x} \\{u(c) + \\beta v(x - c)\\}\n",
    "$$\n",
    "\n",
    "on a grid of $ x $ points and then interpolate.\n",
    "\n",
    "For $ v $ we will use the approximation of the value function we obtained\n",
    "above.\n",
    "\n",
    "Here’s the function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "57c0d4fd",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "@jax.jit\n",
    "def σ(ce, v):\n",
    "    \"\"\"\n",
    "    The optimal policy function. Given the value function,\n",
    "    it finds optimal consumption in each state.\n",
    "\n",
    "    * ce: Cake Eating Model instance\n",
    "    * v: value function array guess, 1-D array\n",
    "\n",
    "    \"\"\"\n",
    "    i_cs =  jnp.argmax(state_action_value_vec(ce.x_grid, ce.c_grid, v, ce), axis=1)\n",
    "    return ce.c_grid[i_cs]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7bbd3db5",
   "metadata": {},
   "source": [
    "Now let’s pass the approximate value function and compute optimal consumption:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5930f970",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "c = σ(ce, v_jax)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "508f4c0a",
   "metadata": {},
   "source": [
    "Let’s plot this next to the true analytical solution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0a2aadef",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "c_analytical = c_star(ce.x_grid, ce.β, ce.γ)\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "\n",
    "ax.plot(ce.x_grid, c_analytical, label='analytical')\n",
    "ax.plot(ce.x_grid, c, label='numerical')\n",
    "ax.set_ylabel(r'$\\sigma(x)$')\n",
    "ax.set_xlabel('$x$')\n",
    "ax.legend()\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4c502de0",
   "metadata": {},
   "source": [
    "## Numba implementation\n",
    "\n",
    "This section of the lecture is directly adapted from [this lecture](https://python.quantecon.org/cake_eating_numerical.html)\n",
    "for the purpose of comparing the results of JAX implementation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6a9482c1",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from numba import prange, njit\n",
    "from quantecon.optimize import brent_max"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9a80ca6c",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "CEMN = namedtuple('CakeEatingModelNumba',\n",
    "                    ('β', 'γ', 'x_grid'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f535dabb",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def create_cake_eating_model_numba(β=0.96,           # discount factor\n",
    "                                   γ=1.5,            # degree of relative risk aversion\n",
    "                                   x_grid_min=1e-3,  # exclude zero for numerical stability\n",
    "                                   x_grid_max=2.5,   # size of cake\n",
    "                                   x_grid_size=200):\n",
    "    x_grid = np.linspace(x_grid_min, x_grid_max, x_grid_size)\n",
    "    return CEMN(β=β, γ=γ, x_grid=x_grid)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3687d69c",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "# Utility function\n",
    "@njit\n",
    "def u_numba(c, cem):\n",
    "    return (c ** (1 - cem.γ)) / (1 - cem.γ)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "636bc343",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "@njit\n",
    "def state_action_value_numba(c, x, v_array, cem):\n",
    "    \"\"\"\n",
    "    Right hand side of the Bellman equation given x and c.\n",
    "    * x: scalar element `x`\n",
    "    * c: consumption\n",
    "    * v_array: value function array guess, 1-D array\n",
    "    * cem: Cake Eating Numba Model instance\n",
    "    \"\"\"\n",
    "    return u_numba(c, cem) + cem.β * np.interp(x - c, cem.x_grid, v_array)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4c7dac6b",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "@njit\n",
    "def T_numba(v, ce):\n",
    "    \"\"\"\n",
    "    The Bellman operator.  Updates the guess of the value function.\n",
    "\n",
    "    * ce is an instance of CakeEatingNumba Model\n",
    "    * v is an array representing a guess of the value function\n",
    "\n",
    "    \"\"\"\n",
    "    v_new = np.empty_like(v)\n",
    "\n",
    "    for i in prange(len(ce.x_grid)):\n",
    "        # Maximize RHS of Bellman equation at state x\n",
    "        v_new[i] = brent_max(state_action_value_numba, 1e-10, ce.x_grid[i],\n",
    "                             args=(ce.x_grid[i], v, ce))[1]\n",
    "    return v_new"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c9a98455",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def compute_value_function_numba(ce,\n",
    "                           tol=1e-4,\n",
    "                           max_iter=1000,\n",
    "                           verbose=True,\n",
    "                           print_skip=25):\n",
    "\n",
    "    # Set up loop\n",
    "    v = np.zeros(len(ce.x_grid)) # Initial guess\n",
    "    i = 0\n",
    "    error = tol + 1\n",
    "\n",
    "    while i < max_iter and error > tol:\n",
    "        v_new = T_numba(v, ce)\n",
    "\n",
    "        error = np.max(np.abs(v - v_new))\n",
    "        i += 1\n",
    "\n",
    "        if verbose and i % print_skip == 0:\n",
    "            print(f\"Error at iteration {i} is {error}.\")\n",
    "\n",
    "        v = v_new\n",
    "\n",
    "    if error > tol:\n",
    "        print(\"Failed to converge!\")\n",
    "    elif verbose:\n",
    "        print(f\"\\nConverged in {i} iterations.\")\n",
    "\n",
    "    return v_new"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c6512d1a",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "cen = create_cake_eating_model_numba()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d6ead6a1",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "in_time = time.time()\n",
    "v_np = compute_value_function_numba(cen)\n",
    "numba_time = time.time() - in_time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f29046d8",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "ratio = numba_time/jax_time\n",
    "print(f\"JAX implementation is {ratio} times faster than NumPy.\")\n",
    "print(f\"JAX time: {jax_time}\")\n",
    "print(f\"Numba time: {numba_time}\")"
   ]
  }
 ],
 "metadata": {
  "date": 1715241420.1465895,
  "filename": "cake_eating_numerical.md",
  "kernelspec": {
   "display_name": "Python",
   "language": "python3",
   "name": "python3"
  },
  "title": "Cake Eating: Numerical Methods"
 },
 "nbformat": 4,
 "nbformat_minor": 5
}