opt_savings_2.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "13da8b74",
   "metadata": {},
   "source": [
    "# Optimal Savings II: Alternative Algorithms"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c14aac0c",
   "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 [Optimal Savings I: Value Function Iteration](https://jax.quantecon.org/opt_savings_1.html) we solved a simple version of the household optimal\n",
    "savings problem via value function iteration (VFI) using JAX.\n",
    "\n",
    "In this lecture we tackle exactly the same problem while adding in two\n",
    "alternative algorithms:\n",
    "\n",
    "- optimistic policy iteration (OPI) and  \n",
    "- Howard policy iteration (HPI).  \n",
    "\n",
    "\n",
    "We will see that both of these algorithms outperform traditional VFI.\n",
    "\n",
    "One reason for this is that the algorithms have good convergence properties.\n",
    "\n",
    "Another is that one of them, HPI, is particularly well suited to pairing with\n",
    "JAX.\n",
    "\n",
    "The reason is that HPI uses a relatively small number of computationally expensive steps,\n",
    "whereas VFI uses a longer sequence of small steps.\n",
    "\n",
    "In other words, VFI is inherently more sequential than HPI, and sequential\n",
    "routines are hard to parallelize.\n",
    "\n",
    "By comparison, HPI is less sequential – the small number of computationally\n",
    "intensive steps can be effectively parallelized by JAX.\n",
    "\n",
    "This is particularly valuable when the underlying hardware includes a GPU.\n",
    "\n",
    "Details on VFI, HPI and OPI can be found in [this book](https://dp.quantecon.org), for which a PDF is freely available.\n",
    "\n",
    "Here we assume readers have some knowledge of the algorithms and focus on\n",
    "computation.\n",
    "\n",
    "For the details of the savings model, readers can refer to [Optimal Savings I: Value Function Iteration](https://jax.quantecon.org/opt_savings_1.html).\n",
    "\n",
    "In addition to what’s in Anaconda, this lecture will need the following libraries:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4b7cd97d",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "!pip install quantecon"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d297f30f",
   "metadata": {},
   "source": [
    "We will use the following imports:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6b71e066",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "import quantecon as qe\n",
    "import jax\n",
    "import jax.numpy as jnp\n",
    "import matplotlib.pyplot as plt\n",
    "from time import time"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dbb9b1ad",
   "metadata": {},
   "source": [
    "Let’s check the GPU we are running."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6085264c",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "!nvidia-smi"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a59a3825",
   "metadata": {},
   "source": [
    "We’ll use 64 bit floats to gain extra precision."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8f62fddb",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "jax.config.update(\"jax_enable_x64\", True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "374dbe4d",
   "metadata": {},
   "source": [
    "## Model primitives\n",
    "\n",
    "First we define a model that stores parameters and grids.\n",
    "\n",
    "The [following code](https://jax.quantecon.org/opt_savings_1.html#prgm-create-consumption-model) is repeated from [Optimal Savings I: Value Function Iteration](https://jax.quantecon.org/opt_savings_1.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8bc73079",
   "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 = 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), mc.P\n",
    "    # Pack and return\n",
    "    params = β, R, γ\n",
    "    sizes = w_size, y_size\n",
    "    arrays = w_grid, y_grid, jnp.array(Q)\n",
    "    return params, sizes, arrays"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4af35601",
   "metadata": {},
   "source": [
    "Here’s the right hand side of the Bellman equation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13845249",
   "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": "0e3514ea",
   "metadata": {},
   "source": [
    "Now we successively apply `vmap` to vectorize $ B $ by simulating nested loops."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3354429f",
   "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": "df92e0e0",
   "metadata": {},
   "source": [
    "Here’s a fully vectorized version of $ B $."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8cd93e03",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def B(v, params, sizes, arrays):\n",
    "    w_size, y_size = sizes\n",
    "    w_indices, y_indices = jnp.arange(w_size), jnp.arange(y_size)\n",
    "    return B_vmap(v, params, arrays, w_indices, y_indices, w_indices)\n",
    "\n",
    "B = jax.jit(B, static_argnums=(2,))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fa044e2b",
   "metadata": {},
   "source": [
    "## Operators\n",
    "\n",
    "Here’s the Bellman operator $ T $"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "76e3a878",
   "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=-1)\n",
    "\n",
    "T = jax.jit(T, static_argnums=(2,))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "059e0b49",
   "metadata": {},
   "source": [
    "The next function computes a $ v $-greedy policy given $ v $"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "57f916e6",
   "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=-1)\n",
    "\n",
    "get_greedy = jax.jit(get_greedy, static_argnums=(2,))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f749934d",
   "metadata": {},
   "source": [
    "We define a function to compute the current rewards $ r_\\sigma $ given policy $ \\sigma $,\n",
    "which is defined as the vector\n",
    "\n",
    "$$\n",
    "r_\\sigma(w, y) := r(w, y, \\sigma(w, y))\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e6440c11",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def _compute_r_σ(σ, params, arrays, i, j):\n",
    "    \"\"\"\n",
    "    With indices (i, j) -> (w, y) and wp = σ[i, j], compute \n",
    "        \n",
    "        r_σ[i, j] = u(Rw + y - wp)   \n",
    "\n",
    "    which gives current rewards under policy σ.\n",
    "    \"\"\"\n",
    "\n",
    "    # Unpack model\n",
    "    β, R, γ = params\n",
    "    w_grid, y_grid, Q = arrays\n",
    "    # Compute r_σ[i, j]\n",
    "    w, y, wp = w_grid[i], y_grid[j], w_grid[σ[i, j]]\n",
    "    c = R * w + y - wp\n",
    "    r_σ = c**(1-γ)/(1-γ)\n",
    "\n",
    "    return r_σ"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cabf4a77",
   "metadata": {},
   "source": [
    "Now we successively apply `vmap` to simulate nested loops."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1aae2d23",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "r_1 = jax.vmap(_compute_r_σ,  in_axes=(None, None, None, None, 0))\n",
    "r_σ_vmap = jax.vmap(r_1,      in_axes=(None, None, None, 0,    None))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4caaa105",
   "metadata": {},
   "source": [
    "Here’s a fully vectorized version of $ r_\\sigma $."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "425dc930",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def compute_r_σ(σ, params, sizes, arrays):\n",
    "    w_size, y_size = sizes\n",
    "    w_indices, y_indices = jnp.arange(w_size), jnp.arange(y_size)\n",
    "    return r_σ_vmap(σ, params, arrays, w_indices, y_indices)\n",
    "\n",
    "compute_r_σ = jax.jit(compute_r_σ, static_argnums=(2,))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "025b54c9",
   "metadata": {},
   "source": [
    "Now we define the policy operator $ T_\\sigma $ going through similar steps"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "45d4c53e",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def _T_σ(v, σ, params, arrays, i, j):\n",
    "    \"The σ-policy operator.\"\n",
    "\n",
    "    # Unpack model\n",
    "    β, R, γ = params\n",
    "    w_grid, y_grid, Q = arrays\n",
    "\n",
    "    r_σ  = _compute_r_σ(σ, params, arrays, i, j)\n",
    "    # Calculate the expected sum Σ_jp v[σ[i, j], jp] * Q[i, j, jp]\n",
    "    EV = jnp.sum(v[σ[i, j], :] * Q[j, :])\n",
    "\n",
    "    return r_σ + β * EV\n",
    "\n",
    "\n",
    "T_1 = jax.vmap(_T_σ,      in_axes=(None, None, None, None, None, 0))\n",
    "T_σ_vmap = jax.vmap(T_1,  in_axes=(None, None, None, None, 0,    None))\n",
    "\n",
    "def T_σ(v, σ, params, sizes, arrays):\n",
    "    w_size, y_size = sizes\n",
    "    w_indices, y_indices = jnp.arange(w_size), jnp.arange(y_size)\n",
    "    return T_σ_vmap(v, σ, params, arrays, w_indices, y_indices)\n",
    "\n",
    "T_σ = jax.jit(T_σ, static_argnums=(3,))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "83c8f802",
   "metadata": {},
   "source": [
    "The function below computes the value $ v_\\sigma $ of following policy $ \\sigma $.\n",
    "\n",
    "This lifetime value is a function $ v_\\sigma $ that satisfies\n",
    "\n",
    "$$\n",
    "v_\\sigma(w, y) = r_\\sigma(w, y) + \\beta \\sum_{y'} v_\\sigma(\\sigma(w, y), y') Q(y, y')\n",
    "$$\n",
    "\n",
    "We wish to solve this equation for $ v_\\sigma $.\n",
    "\n",
    "Suppose we define the linear operator $ L_\\sigma $ by\n",
    "\n",
    "$$\n",
    "(L_\\sigma v)(w, y) = v(w, y) - \\beta \\sum_{y'} v(\\sigma(w, y), y') Q(y, y')\n",
    "$$\n",
    "\n",
    "With this notation, the problem is to solve for $ v $ via\n",
    "\n",
    "$$\n",
    "(L_{\\sigma} v)(w, y) = r_\\sigma(w, y)\n",
    "$$\n",
    "\n",
    "In vector for this is $ L_\\sigma v = r_\\sigma $, which tells us that the function\n",
    "we seek is\n",
    "\n",
    "$$\n",
    "v_\\sigma = L_\\sigma^{-1} r_\\sigma\n",
    "$$\n",
    "\n",
    "JAX allows us to solve linear systems defined in terms of operators; the first\n",
    "step is to define the function $ L_{\\sigma} $."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9d8b0982",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def _L_σ(v, σ, params, arrays, i, j):\n",
    "    \"\"\"\n",
    "    Here we set up the linear map v -> L_σ v, where \n",
    "\n",
    "        (L_σ v)(w, y) = v(w, y) - β Σ_y′ v(σ(w, y), y′) Q(y, y′)\n",
    "\n",
    "    \"\"\"\n",
    "    # Unpack\n",
    "    β, R, γ = params\n",
    "    w_grid, y_grid, Q = arrays\n",
    "    # Compute and return v[i, j] - β Σ_jp v[σ[i, j], jp] * Q[j, jp]\n",
    "    return v[i, j]  - β * jnp.sum(v[σ[i, j], :] * Q[j, :])\n",
    "\n",
    "L_1 = jax.vmap(_L_σ,      in_axes=(None, None, None, None, None, 0))\n",
    "L_σ_vmap = jax.vmap(L_1,  in_axes=(None, None, None, None, 0,    None))\n",
    "\n",
    "def L_σ(v, σ, params, sizes, arrays):\n",
    "    w_size, y_size = sizes\n",
    "    w_indices, y_indices = jnp.arange(w_size), jnp.arange(y_size)\n",
    "    return L_σ_vmap(v, σ, params, arrays, w_indices, y_indices)\n",
    "\n",
    "L_σ = jax.jit(L_σ, static_argnums=(3,))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f9d57b48",
   "metadata": {},
   "source": [
    "Now we can define a function to compute $ v_{\\sigma} $"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cbddb758",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def get_value(σ, params, sizes, arrays):\n",
    "    \"Get the value v_σ of policy σ by inverting the linear map L_σ.\"\n",
    "\n",
    "    # Unpack\n",
    "    β, R, γ = params\n",
    "    w_size, y_size = sizes\n",
    "    w_grid, y_grid, Q = arrays\n",
    "\n",
    "    r_σ = compute_r_σ(σ, params, sizes, arrays)\n",
    "\n",
    "    # Reduce L_σ to a function in v\n",
    "    partial_L_σ = lambda v: L_σ(v, σ, params, sizes, arrays)\n",
    "\n",
    "    return jax.scipy.sparse.linalg.bicgstab(partial_L_σ, r_σ)[0]\n",
    "\n",
    "get_value = jax.jit(get_value, static_argnums=(2,))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e47caddc",
   "metadata": {},
   "source": [
    "## Iteration\n",
    "\n",
    "We use successive approximation for VFI."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "56bc8559",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def successive_approx_jax(T,                     # Operator (callable)\n",
    "                          x_0,                   # Initial condition                \n",
    "                          tol=1e-6,              # Error tolerance\n",
    "                          max_iter=10_000):      # Max iteration bound\n",
    "    def update(inputs):\n",
    "        k, x, error = inputs\n",
    "        x_new = T(x)\n",
    "        error = jnp.max(jnp.abs(x_new - x))\n",
    "        return k + 1, x_new, error\n",
    "\n",
    "    def condition_function(inputs):\n",
    "        k, x, error = inputs\n",
    "        return jnp.logical_and(error > tol, k < max_iter)\n",
    "\n",
    "    k, x, error = jax.lax.while_loop(condition_function, \n",
    "                                     update, \n",
    "                                     (1, x_0, tol + 1))\n",
    "    return x\n",
    "\n",
    "successive_approx_jax = jax.jit(successive_approx_jax, static_argnums=(0,))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "080be192",
   "metadata": {},
   "source": [
    "For OPI we’ll add a compiled routine that computes $ T_σ^m v $."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "da807bc3",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def iterate_policy_operator(σ, v, m, params, sizes, arrays):\n",
    "\n",
    "    def update(i, v):\n",
    "        v = T_σ(v, σ, params, sizes, arrays)\n",
    "        return v\n",
    "    \n",
    "    v = jax.lax.fori_loop(0, m, update, v)\n",
    "    return v\n",
    "\n",
    "iterate_policy_operator = jax.jit(iterate_policy_operator,\n",
    "                                  static_argnums=(4,))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3b390fd7",
   "metadata": {},
   "source": [
    "## Solvers\n",
    "\n",
    "Now we define the solvers, which implement VFI, HPI and OPI.\n",
    "\n",
    "Here’s VFI."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e86ba022",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def value_function_iteration(model, tol=1e-5):\n",
    "    \"\"\"\n",
    "    Implements value function iteration.\n",
    "    \"\"\"\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, tol=tol)\n",
    "    return get_greedy(v_star, params, sizes, arrays)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6edec4e2",
   "metadata": {},
   "source": [
    "For OPI we will use a compiled JAX `lax.while_loop` operation to speed execution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "865b5b9f",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def opi_loop(params, sizes, arrays, m, tol, max_iter):\n",
    "    \"\"\"\n",
    "    Implements optimistic policy iteration (see dp.quantecon.org) with \n",
    "    step size m.\n",
    "\n",
    "    \"\"\"\n",
    "    v_init = jnp.zeros(sizes)\n",
    "\n",
    "    def condition_function(inputs):\n",
    "        i, v, error = inputs\n",
    "        return jnp.logical_and(error > tol, i < max_iter)\n",
    "\n",
    "    def update(inputs):\n",
    "        i, v, error = inputs\n",
    "        last_v = v\n",
    "        σ = get_greedy(v, params, sizes, arrays)\n",
    "        v = iterate_policy_operator(σ, v, m, params, sizes, arrays)\n",
    "        error = jnp.max(jnp.abs(v - last_v))\n",
    "        i += 1\n",
    "        return i, v, error\n",
    "\n",
    "    num_iter, v, error = jax.lax.while_loop(condition_function,\n",
    "                                            update,\n",
    "                                            (0, v_init, tol + 1))\n",
    "\n",
    "    return get_greedy(v, params, sizes, arrays)\n",
    "\n",
    "opi_loop = jax.jit(opi_loop, static_argnums=(1,))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "15064c26",
   "metadata": {},
   "source": [
    "Here’s a friendly interface to OPI"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9c925a01",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def optimistic_policy_iteration(model, m=10, tol=1e-5, max_iter=10_000):\n",
    "    params, sizes, arrays = model\n",
    "    σ_star = opi_loop(params, sizes, arrays, m, tol, max_iter)\n",
    "    return σ_star"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4bb16ca4",
   "metadata": {},
   "source": [
    "Here’s HPI."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "07cc5c6d",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def howard_policy_iteration(model, maxiter=250):\n",
    "    \"\"\"\n",
    "    Implements Howard policy iteration (see dp.quantecon.org)\n",
    "    \"\"\"\n",
    "    params, sizes, arrays = model\n",
    "    σ = jnp.zeros(sizes, dtype=int)\n",
    "    i, error = 0, 1.0\n",
    "    while error > 0 and i < maxiter:\n",
    "        v_σ = get_value(σ, params, sizes, arrays)\n",
    "        σ_new = get_greedy(v_σ, params, sizes, arrays)\n",
    "        error = jnp.max(jnp.abs(σ_new - σ))\n",
    "        σ = σ_new\n",
    "        i = i + 1\n",
    "        print(f\"Concluded loop {i} with error {error}.\")\n",
    "    return σ"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ac2f94f3",
   "metadata": {},
   "source": [
    "## Tests\n",
    "\n",
    "Let’s create a model for consumption, and plot the resulting optimal policy function using all the three algorithms and also check the time taken by each solver."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7d4da6c2",
   "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": "b1fa1e02",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "print(\"Starting HPI.\")\n",
    "start = time()\n",
    "σ_star_hpi = howard_policy_iteration(model).block_until_ready()\n",
    "hpi_with_compile = time() - start\n",
    "print(f\"HPI completed in {hpi_with_compile} seconds.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a7cb37d7",
   "metadata": {},
   "source": [
    "We run it again to get rid of compile time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "59b6b34b",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "start = time()\n",
    "σ_star_hpi = howard_policy_iteration(model).block_until_ready()\n",
    "hpi_without_compile = time() - start\n",
    "print(f\"HPI completed in {hpi_without_compile} seconds.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3e04a150",
   "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_hpi[:, 1]], label=\"$\\\\sigma^{*}_{HPI}(\\cdot, y_1)$\")\n",
    "ax.plot(w_grid, w_grid[σ_star_hpi[:, -1]], label=\"$\\\\sigma^{*}_{HPI}(\\cdot, y_N)$\")\n",
    "ax.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b214de60",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "print(\"Starting VFI.\")\n",
    "start = time()\n",
    "σ_star_vfi = value_function_iteration(model).block_until_ready()\n",
    "vfi_with_compile = time() - start\n",
    "print(f\"VFI completed in {vfi_with_compile} seconds.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "089dd922",
   "metadata": {},
   "source": [
    "We run it again to eliminate compile time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0ea7e62c",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "start = time()\n",
    "σ_star_vfi = value_function_iteration(model).block_until_ready()\n",
    "vfi_without_compile = time() - start\n",
    "print(f\"VFI completed in {vfi_without_compile} seconds.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fd06587c",
   "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_vfi[:, 1]], label=\"$\\\\sigma^{*}_{VFI}(\\cdot, y_1)$\")\n",
    "ax.plot(w_grid, w_grid[σ_star_vfi[:, -1]], label=\"$\\\\sigma^{*}_{VFI}(\\cdot, y_N)$\")\n",
    "ax.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b46ed4c9",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "print(\"Starting OPI.\")\n",
    "start = time()\n",
    "σ_star_opi = optimistic_policy_iteration(model, m=100).block_until_ready()\n",
    "opi_with_compile = time() - start\n",
    "print(f\"OPI completed in {opi_with_compile} seconds.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e45d949a",
   "metadata": {},
   "source": [
    "Let’s run it again to get rid of compile time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0a5a36ba",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "start = time()\n",
    "σ_star_opi = optimistic_policy_iteration(model, m=100).block_until_ready()\n",
    "opi_without_compile = time() - start\n",
    "print(f\"OPI completed in {opi_without_compile} seconds.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dc4c8ac9",
   "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_opi[:, 1]], label=\"$\\\\sigma^{*}_{OPI}(\\cdot, y_1)$\")\n",
    "ax.plot(w_grid, w_grid[σ_star_opi[:, -1]], label=\"$\\\\sigma^{*}_{OPI}(\\cdot, y_N)$\")\n",
    "ax.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1595b7cb",
   "metadata": {},
   "source": [
    "We observe that all the solvers produce the same output from the above three plots.\n",
    "\n",
    "Now, let’s create a plot to visualize the time differences among these algorithms."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7d0b421e",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def run_algorithm(algorithm, model, **kwargs):\n",
    "    result = algorithm(model, **kwargs)\n",
    "\n",
    "    # Now time it without compile time\n",
    "    start = time()\n",
    "    result = algorithm(model, **kwargs).block_until_ready()\n",
    "    algorithm_without_compile = time() - start\n",
    "    print(f\"{algorithm.__name__} completed in {algorithm_without_compile:.2f} seconds.\")\n",
    "    return result, algorithm_without_compile"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1dd446a7",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "σ_pi, pi_time = run_algorithm(howard_policy_iteration, model)\n",
    "σ_vfi, vfi_time = run_algorithm(value_function_iteration, model, tol=1e-5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b45f78de",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "m_vals = range(5, 600, 40)\n",
    "opi_times = []\n",
    "for m in m_vals:\n",
    "    σ_opi, opi_time = run_algorithm(optimistic_policy_iteration, \n",
    "                                    model, m=m, tol=1e-5)\n",
    "    opi_times.append(opi_time)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "086f2251",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "ax.plot(m_vals, \n",
    "        jnp.full(len(m_vals), hpi_without_compile), \n",
    "        lw=2, label=\"Howard policy iteration\")\n",
    "ax.plot(m_vals, \n",
    "        jnp.full(len(m_vals), vfi_without_compile), \n",
    "        lw=2, label=\"value function iteration\")\n",
    "ax.plot(m_vals, opi_times, \n",
    "        lw=2, label=\"optimistic policy iteration\")\n",
    "ax.legend(frameon=False)\n",
    "ax.set_xlabel(\"$m$\")\n",
    "ax.set_ylabel(\"time\")\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "date": 1732064318.2627592,
  "filename": "opt_savings_2.md",
  "kernelspec": {
   "display_name": "Python",
   "language": "python3",
   "name": "python3"
  },
  "title": "Optimal Savings II: Alternative Algorithms"
 },
 "nbformat": 4,
 "nbformat_minor": 5
}