inventory_ssd.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "7d31df8d",
   "metadata": {},
   "source": [
    "# Inventory Management Model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8bae2abf",
   "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 provides a JAX implementation of a model in [Dynamic Programming](https://dp.quantecon.org/).\n",
    "\n",
    "In addition to JAX and Anaconda, this lecture will need the following libraries:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "93c23a73",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "!pip install --upgrade quantecon"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "08a8d182",
   "metadata": {},
   "source": [
    "## A model with constant discounting\n",
    "\n",
    "We study a firm where a manager tries to maximize shareholder value.\n",
    "\n",
    "To simplify the problem, we  assume that the firm only sells one product.\n",
    "\n",
    "Letting $ \\pi_t $ be profits at time $ t $ and $ r > 0 $ be the interest rate, the value of the firm is\n",
    "\n",
    "$$\n",
    "V_0 = \\sum_{t \\geq 0} \\beta^t \\pi_t\n",
    "    \\qquad\n",
    "    \\text{ where }\n",
    "    \\quad \\beta := \\frac{1}{1+r}.\n",
    "$$\n",
    "\n",
    "Suppose the firm faces exogenous demand process $ (D_t)_{t \\geq 0} $.\n",
    "\n",
    "We assume $ (D_t)_{t \\geq 0} $ is IID with common distribution $ \\phi \\in (Z_+) $.\n",
    "\n",
    "Inventory $ (X_t)_{t \\geq 0} $ of the product obeys\n",
    "\n",
    "$$\n",
    "X_{t+1} = f(X_t, D_{t+1}, A_t)\n",
    "    \\qquad\n",
    "    \\text{where}\n",
    "    \\quad\n",
    "    f(x,a,d) := (x - d)\\vee 0 + a.\n",
    "$$\n",
    "\n",
    "The term $ A_t $ is units of stock ordered this period, which take one period to\n",
    "arrive.\n",
    "\n",
    "We assume that the firm can store at most $ K $ items at one time.\n",
    "\n",
    "Profits are given by\n",
    "\n",
    "$$\n",
    "\\pi_t := X_t \\wedge D_{t+1} - c A_t - \\kappa 1\\{A_t > 0\\}.\n",
    "$$\n",
    "\n",
    "We take the minimum of current stock and demand because orders in excess of\n",
    "inventory are assumed to be lost rather than back-filled.\n",
    "\n",
    "Here $ c $ is unit product cost and $ \\kappa $ is a fixed cost of ordering inventory.\n",
    "\n",
    "We can map our inventory problem into a dynamic program with state space\n",
    "$ X := \\{0, \\ldots, K\\} $ and action space $ A := X $.\n",
    "\n",
    "The feasible correspondence $ \\Gamma $ is\n",
    "\n",
    "$$\n",
    "\\Gamma(x) := \\{0, \\ldots, K - x\\},\n",
    "$$\n",
    "\n",
    "which represents the set of feasible orders when the current inventory\n",
    "state is $ x $.\n",
    "\n",
    "The reward function is expected current profits, or\n",
    "\n",
    "$$\n",
    "r(x, a)  := \\sum_{d \\geq 0} (x \\wedge d) \\phi(d)\n",
    "        - c a - \\kappa 1\\{a > 0\\}.\n",
    "$$\n",
    "\n",
    "The stochastic kernel (i.e., state-transition probabilities) from the set of feasible state-action pairs is\n",
    "\n",
    "$$\n",
    "P(x, a, x') := P\\{ f(x, a, D) = x' \\}\n",
    "    \\qquad \\text{when} \\quad\n",
    "    D \\sim \\phi.\n",
    "$$\n",
    "\n",
    "When discounting is constant, the Bellman equation takes the form\n",
    "\n",
    "\n",
    "<a id='equation-inventory-ssd-v1'></a>\n",
    "$$\n",
    "v(x)\n",
    "    = \\max_{a \\in \\Gamma(x)} \\left\\{\n",
    "        r(x, a)\n",
    "        + \\beta\n",
    "        \\sum_{d \\geq 0} v(f(x, a, d)) \\phi(d)\n",
    "    \\right\\} \\tag{15.1}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b1806bc5",
   "metadata": {},
   "source": [
    "## Time varing discount rates\n",
    "\n",
    "We wish to consider a more sophisticated model with time-varying discounting.\n",
    "\n",
    "This time variation accommodates non-constant interest rates.\n",
    "\n",
    "To this end, we replace the constant $ \\beta $ in\n",
    "[(15.1)](#equation-inventory-ssd-v1) with a stochastic process $ (\\beta_t) $ where\n",
    "\n",
    "- $ \\beta_t = 1/(1+r_t) $ and  \n",
    "- $ r_t $ is the interest rate at time $ t $  \n",
    "\n",
    "\n",
    "We suppose that the dynamics can be expressed as $ \\beta_t = \\beta(Z_t) $, where the exogenous process $ (Z_t)_{t \\geq 0} $ is a Markov chain\n",
    "on $ Z $ with Markov matrix $ Q $.\n",
    "\n",
    "After relabeling inventory $ X_t $ as $ Y_t $ and $ x $ as $ y $, the Bellman equation becomes\n",
    "\n",
    "$$\n",
    "v(y, z) = \\max_{a \\in \\Gamma(x)} B((y, z), a, v)\n",
    "$$\n",
    "\n",
    "where\n",
    "\n",
    "\n",
    "<a id='equation-inventory-ssd-b1'></a>\n",
    "$$\n",
    "B((y, z), a, v)\n",
    "    =\n",
    "        r(y, a)\n",
    "        + \\beta(z)\n",
    "        \\sum_{d, \\, z'} v(f(y, a, d), z') \\phi(d) Q(z, z'). \\tag{15.2}\n",
    "$$\n",
    "\n",
    "We set\n",
    "\n",
    "$$\n",
    "R(y, a, y')\n",
    "        := P\\{f(y, a, d) = y'\\} \\quad \\text{when} \\quad D \\sim \\phi,\n",
    "$$\n",
    "\n",
    "Now $ R(y, a, y') $ is the probability of realizing next period inventory level\n",
    "$ y' $ when the current level is $ y $ and the action is $ a $.\n",
    "\n",
    "Hence we can rewrite [(15.2)](#equation-inventory-ssd-b1) as\n",
    "\n",
    "$$\n",
    "B((y, z), a, v)\n",
    "    =  r(y, a)\n",
    "        + \\beta(z)\n",
    "        \\sum_{y', z'} v(y', z') Q(z, z') R(y, a, y') .\n",
    "$$\n",
    "\n",
    "Let’s begin with the following imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d6fc6890",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "import quantecon as qe\n",
    "import jax\n",
    "import jax.numpy as jnp\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from collections import namedtuple\n",
    "import numba\n",
    "from numba import prange\n",
    "from time import time"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "31baf70c",
   "metadata": {},
   "source": [
    "Let’s check the GPU we are running"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c06c71b3",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "!nvidia-smi"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "71e630d4",
   "metadata": {},
   "source": [
    "We will use 64 bit floats with JAX in order to increase the precision."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8cfe1ae8",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "jax.config.update(\"jax_enable_x64\", True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "299e8d99",
   "metadata": {},
   "source": [
    "Let’s define a model to represent the inventory management."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e9d2d4d9",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "# NamedTuple Model\n",
    "Model = namedtuple(\"Model\", (\"c\", \"κ\", \"p\", \"z_vals\", \"Q\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f4e4727c",
   "metadata": {},
   "source": [
    "We need the following successive approximation function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9b8215f7",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def successive_approx(T,                     # Operator (callable)\n",
    "                      x_0,                   # Initial condition\n",
    "                      tolerance=1e-6,        # Error tolerance\n",
    "                      max_iter=10_000,       # Max iteration bound\n",
    "                      print_step=25,         # Print at multiples\n",
    "                      verbose=False):\n",
    "    x = x_0\n",
    "    error = tolerance + 1\n",
    "    k = 1\n",
    "    while error > tolerance and k <= max_iter:\n",
    "        x_new = T(x)\n",
    "        error = jnp.max(jnp.abs(x_new - x))\n",
    "        if verbose and k % print_step == 0:\n",
    "            print(f\"Completed iteration {k} with error {error}.\")\n",
    "        x = x_new\n",
    "        k += 1\n",
    "    if error > tolerance:\n",
    "        print(f\"Warning: Iteration hit upper bound {max_iter}.\")\n",
    "    elif verbose:\n",
    "        print(f\"Terminated successfully in {k} iterations.\")\n",
    "    return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9dbe2c57",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "@jax.jit\n",
    "def demand_pdf(p, d):\n",
    "    return (1 - p)**d * p"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cae543d3",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "K = 100\n",
    "D_MAX = 101"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8889a054",
   "metadata": {},
   "source": [
    "Let’s define a function to create an inventory model using the given parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "014506d0",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def create_sdd_inventory_model(\n",
    "        ρ=0.98, ν=0.002, n_z=100, b=0.97,          # Z state parameters\n",
    "        c=0.2, κ=0.8, p=0.6,                       # firm and demand parameters\n",
    "        use_jax=True):\n",
    "    mc = qe.tauchen(n_z, ρ, ν)\n",
    "    z_vals, Q = mc.state_values + b, mc.P\n",
    "    if use_jax:\n",
    "        z_vals, Q = map(jnp.array, (z_vals, Q))\n",
    "    return Model(c=c, κ=κ, p=p, z_vals=z_vals, Q=Q)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3b0724a8",
   "metadata": {},
   "source": [
    "Here’s the function `B` on the right-hand side of the Bellman equation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "569c577d",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "@jax.jit\n",
    "def B(x, i_z, a, v, model):\n",
    "    \"\"\"\n",
    "    The function B(x, z, a, v) = r(x, a) + β(z) Σ_x′ v(x′) P(x, a, x′).\n",
    "    \"\"\"\n",
    "    c, κ, p, z_vals, Q = model\n",
    "    z = z_vals[i_z]\n",
    "    d_vals = jnp.arange(D_MAX)\n",
    "    ϕ_vals = demand_pdf(p, d_vals)\n",
    "    revenue = jnp.sum(jnp.minimum(x, d_vals)*ϕ_vals)\n",
    "    profit = revenue - c * a - κ * (a > 0)\n",
    "    v_R = jnp.sum(v[jnp.maximum(x - d_vals, 0) + a].T * ϕ_vals, axis=1)\n",
    "    cv = jnp.sum(v_R*Q[i_z])\n",
    "    return profit + z * cv"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e6f6dd06",
   "metadata": {},
   "source": [
    "We need to vectorize this function so that we can use it efficiently in JAX.\n",
    "\n",
    "We apply a sequence of `vmap` operations to vectorize appropriately in each\n",
    "argument."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "97714e0a",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "B_vec_a = jax.vmap(B, in_axes=(None, None, 0, None, None))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6d0866a0",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "@jax.jit\n",
    "def B2(x, i_z, v, model):\n",
    "    \"\"\"\n",
    "    The function B(x, z, a, v) = r(x, a) + β(z) Σ_x′ v(x′) P(x, a, x′).\n",
    "    \"\"\"\n",
    "    c, κ, p, z_vals, Q = model\n",
    "    a_vals = jnp.arange(K)\n",
    "    res = B_vec_a(x, i_z, a_vals, v, model)\n",
    "    return jnp.where(a_vals < K - x + 1, res, -jnp.inf)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c83927af",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "B2_vec_z = jax.vmap(B2, in_axes=(None, 0, None, None))\n",
    "B2_vec_z_x = jax.vmap(B2_vec_z, in_axes=(0, None, None, None))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "54e2f6bc",
   "metadata": {},
   "source": [
    "Next we define the Bellman operator."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5921ad32",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "@jax.jit\n",
    "def T(v, model):\n",
    "    \"\"\"The Bellman operator.\"\"\"\n",
    "    c, κ, p, z_vals, Q = model\n",
    "    i_z_range = jnp.arange(len(z_vals))\n",
    "    x_range = jnp.arange(K + 1)\n",
    "    res = B2_vec_z_x(x_range, i_z_range, v, model)\n",
    "    return jnp.max(res, axis=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3beed5ae",
   "metadata": {},
   "source": [
    "The following function computes a v-greedy policy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ceb266e8",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "@jax.jit\n",
    "def get_greedy(v, model):\n",
    "    \"\"\"Get a v-greedy policy.  Returns a zero-based array.\"\"\"\n",
    "    c, κ, p, z_vals, Q  = model\n",
    "    i_z_range = jnp.arange(len(z_vals))\n",
    "    x_range = jnp.arange(K + 1)\n",
    "    res = B2_vec_z_x(x_range, i_z_range, v, model)\n",
    "    return jnp.argmax(res, axis=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "67e4d2c9",
   "metadata": {},
   "source": [
    "Here’s code to solve the model using value function iteration."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7499e853",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def solve_inventory_model(v_init, model):\n",
    "    \"\"\"Use successive_approx to get v_star and then compute greedy.\"\"\"\n",
    "    v_star = successive_approx(lambda v: T(v, model), v_init, verbose=True)\n",
    "    σ_star = get_greedy(v_star, model)\n",
    "    return v_star, σ_star"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4aab797d",
   "metadata": {},
   "source": [
    "Now let’s create an instance and solve it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "71f3fbe9",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "model = create_sdd_inventory_model()\n",
    "c, κ, p, z_vals, Q = model\n",
    "n_z = len(z_vals)\n",
    "v_init = jnp.zeros((K + 1, n_z), dtype=float)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8da1d521",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "start = time()\n",
    "v_star, σ_star = solve_inventory_model(v_init, model)\n",
    "jax_time_with_compile = time() - start\n",
    "print(\"Jax compile plus execution time = \", jax_time_with_compile)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "59a4a33e",
   "metadata": {},
   "source": [
    "Let’s run again to get rid of the compile time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15772273",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "start = time()\n",
    "v_star, σ_star = solve_inventory_model(v_init, model)\n",
    "jax_time_without_compile = time() - start\n",
    "print(\"Jax execution time = \", jax_time_without_compile)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d8dc2230",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "z_mc = qe.MarkovChain(Q, z_vals)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "38be9cc0",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def sim_inventories(ts_length, X_init=0):\n",
    "    \"\"\"Simulate given the optimal policy.\"\"\"\n",
    "    global p, z_mc\n",
    "    i_z = z_mc.simulate_indices(ts_length, init=1)\n",
    "    X = np.zeros(ts_length, dtype=np.int32)\n",
    "    X[0] = X_init\n",
    "    rand = np.random.default_rng().geometric(p=p, size=ts_length-1) - 1\n",
    "    for t in range(ts_length-1):\n",
    "        X[t+1] = np.maximum(X[t] - rand[t], 0) + σ_star[X[t], i_z[t]]\n",
    "    return X, z_vals[i_z]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "29bfbf7c",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def plot_ts(ts_length=400, fontsize=10):\n",
    "    X, Z = sim_inventories(ts_length)\n",
    "    fig, axes = plt.subplots(2, 1, figsize=(9, 5.5))\n",
    "\n",
    "    ax = axes[0]\n",
    "    ax.plot(X, label=r\"$X_t$\", alpha=0.7)\n",
    "    ax.set_xlabel(r\"$t$\", fontsize=fontsize)\n",
    "    ax.set_ylabel(\"inventory\", fontsize=fontsize)\n",
    "    ax.legend(fontsize=fontsize, frameon=False)\n",
    "    ax.set_ylim(0, np.max(X)+3)\n",
    "\n",
    "    # calculate interest rate from discount factors\n",
    "    r = (1 / Z) - 1\n",
    "\n",
    "    ax = axes[1]\n",
    "    ax.plot(r, label=r\"$r_t$\", alpha=0.7)\n",
    "    ax.set_xlabel(r\"$t$\", fontsize=fontsize)\n",
    "    ax.set_ylabel(\"interest rate\", fontsize=fontsize)\n",
    "    ax.legend(fontsize=fontsize, frameon=False)\n",
    "\n",
    "    plt.tight_layout()\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6d2e12cd",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "plot_ts()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7869e564",
   "metadata": {},
   "source": [
    "## Numba implementation\n",
    "\n",
    "Let’s try the same operations in Numba in order to compare the speed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "35bb3468",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "@numba.njit\n",
    "def demand_pdf_numba(p, d):\n",
    "    return (1 - p)**d * p\n",
    "\n",
    "@numba.njit\n",
    "def B_numba(x, i_z, a, v, model):\n",
    "    \"\"\"\n",
    "    The function B(x, z, a, v) = r(x, a) + β(z) Σ_x′ v(x′) P(x, a, x′).\n",
    "    \"\"\"\n",
    "    c, κ, p, z_vals, Q = model\n",
    "    z = z_vals[i_z]\n",
    "    d_vals = np.arange(D_MAX)\n",
    "    ϕ_vals = demand_pdf_numba(p, d_vals)\n",
    "    revenue = np.sum(np.minimum(x, d_vals)*ϕ_vals)\n",
    "    profit = revenue - c * a - κ * (a > 0)\n",
    "    v_R = np.sum(v[np.maximum(x - d_vals, 0) + a].T * ϕ_vals, axis=1)\n",
    "    cv = np.sum(v_R*Q[i_z])\n",
    "    return profit + z * cv\n",
    "\n",
    "\n",
    "@numba.njit(parallel=True)\n",
    "def T_numba(v, model):\n",
    "    \"\"\"The Bellman operator.\"\"\"\n",
    "    c, κ, p, z_vals, Q = model\n",
    "    new_v = np.empty_like(v)\n",
    "    for i_z in prange(len(z_vals)):\n",
    "        for x in prange(K+1):\n",
    "            v_1 = np.array([B_numba(x, i_z, a, v, model)\n",
    "                             for a in range(K-x+1)])\n",
    "            new_v[x, i_z] = np.max(v_1)\n",
    "    return new_v\n",
    "\n",
    "\n",
    "@numba.njit(parallel=True)\n",
    "def get_greedy_numba(v, model):\n",
    "    \"\"\"Get a v-greedy policy.  Returns a zero-based array.\"\"\"\n",
    "    c, κ, p, z_vals, Q = model\n",
    "    n_z = len(z_vals)\n",
    "    σ_star = np.zeros((K+1, n_z), dtype=np.int32)\n",
    "    for i_z in prange(n_z):\n",
    "        for x in range(K+1):\n",
    "            v_1 = np.array([B_numba(x, i_z, a, v, model)\n",
    "                             for a in range(K-x+1)])\n",
    "            σ_star[x, i_z] = np.argmax(v_1)\n",
    "    return σ_star\n",
    "\n",
    "\n",
    "\n",
    "def solve_inventory_model_numba(v_init, model):\n",
    "    \"\"\"Use successive_approx to get v_star and then compute greedy.\"\"\"\n",
    "    v_star = successive_approx(lambda v: T_numba(v, model), v_init, verbose=True)\n",
    "    σ_star = get_greedy_numba(v_star, model)\n",
    "    return v_star, σ_star"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0ce7513d",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "model = create_sdd_inventory_model(use_jax=False)\n",
    "c, κ, p, z_vals, Q  = model\n",
    "n_z = len(z_vals)\n",
    "v_init = np.zeros((K + 1, n_z), dtype=float)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "744d921d",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "start = time()\n",
    "v_star_numba, σ_star_numba = solve_inventory_model_numba(v_init, model)\n",
    "numba_time_with_compile = time() - start\n",
    "print(\"Numba compile plus execution time = \", numba_time_with_compile)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dded0c9a",
   "metadata": {},
   "source": [
    "Let’s run again to eliminate the compile time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f636e1b5",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "start = time()\n",
    "v_star_numba, σ_star_numba = solve_inventory_model_numba(v_init, model)\n",
    "numba_time_without_compile = time() - start\n",
    "print(\"Numba execution time = \", numba_time_without_compile)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0fb5f269",
   "metadata": {},
   "source": [
    "Let’s verify that the Numba and JAX implementations converge to the same solution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "76d31f20",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "np.allclose(v_star_numba, v_star)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "87265b97",
   "metadata": {},
   "source": [
    "Here’s the speed comparison."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fcaa3d6e",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "print(\"JAX vectorized implementation is \"\n",
    "      f\"{numba_time_without_compile/jax_time_without_compile} faster \"\n",
    "       \"than Numba's parallel implementation\")"
   ]
  }
 ],
 "metadata": {
  "date": 1732064317.7839913,
  "filename": "inventory_ssd.md",
  "kernelspec": {
   "display_name": "Python",
   "language": "python3",
   "name": "python3"
  },
  "title": "Inventory Management Model"
 },
 "nbformat": 4,
 "nbformat_minor": 5
}