inventory_dynamics.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "19618159",
   "metadata": {},
   "source": [
    "# Inventory Dynamics"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5c09d3e9",
   "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",
    "\n",
    "<a id='index-0'></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fe6363f3",
   "metadata": {},
   "source": [
    "## Overview\n",
    "\n",
    "This lecture explores the inventory dynamics of a firm using so-called s-S inventory control.\n",
    "\n",
    "Loosely speaking, this means that the firm\n",
    "\n",
    "- waits until inventory falls below some value $ s $  \n",
    "- and then restocks with a bulk order of $ S $ units (or, in some models, restocks up to level $ S $).  \n",
    "\n",
    "\n",
    "We will be interested in the distribution of the associated Markov process,\n",
    "which can be thought of as cross-sectional distributions of inventory levels\n",
    "across a large number of firms, all of which\n",
    "\n",
    "1. evolve independently and  \n",
    "1. have the same dynamics.  \n",
    "\n",
    "\n",
    "Note that we also studied this model in a [separate\n",
    "lecture](https://python.quantecon.org/inventory_dynamics.html), using Numba.\n",
    "\n",
    "Here we study the same problem using JAX.\n",
    "\n",
    "We will use the following imports:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b1132bca",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import jax\n",
    "import jax.numpy as jnp\n",
    "from jax import random, lax\n",
    "from collections import namedtuple"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ac930dd0",
   "metadata": {},
   "source": [
    "Here’s a description of our GPU:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "92b2f8bc",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "!nvidia-smi"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fb05d20d",
   "metadata": {},
   "source": [
    "## Sample paths\n",
    "\n",
    "Consider a firm with inventory $ X_t $.\n",
    "\n",
    "The firm waits until $ X_t \\leq s $ and then restocks up to $ S $ units.\n",
    "\n",
    "It faces stochastic demand $ \\{ D_t \\} $, which we assume is IID across time and\n",
    "firms.\n",
    "\n",
    "With notation $ a^+ := \\max\\{a, 0\\} $, inventory dynamics can be written\n",
    "as\n",
    "\n",
    "$$\n",
    "X_{t+1} =\n",
    "    \\begin{cases}\n",
    "      ( S - D_{t+1})^+ & \\quad \\text{if } X_t \\leq s \\\\\n",
    "      ( X_t - D_{t+1} )^+ &  \\quad \\text{if } X_t > s\n",
    "    \\end{cases}\n",
    "$$\n",
    "\n",
    "In what follows, we will assume that each $ D_t $ is lognormal, so that\n",
    "\n",
    "$$\n",
    "D_t = \\exp(\\mu + \\sigma Z_t)\n",
    "$$\n",
    "\n",
    "where $ \\mu $ and $ \\sigma $ are parameters and $ \\{Z_t\\} $ is IID\n",
    "and standard normal.\n",
    "\n",
    "Here’s a `namedtuple` that stores parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f29de941",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "Parameters = namedtuple('Parameters', ['s', 'S', 'μ', 'σ'])\n",
    "\n",
    "# Create a default instance\n",
    "params = Parameters(s=10, S=100, μ=1.0, σ=0.5)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "abe8fe9a",
   "metadata": {},
   "source": [
    "## Cross-sectional distributions\n",
    "\n",
    "Now let’s look at the marginal distribution $ \\psi_T $ of $ X_T $ for some fixed $ T $.\n",
    "\n",
    "The probability distribution $ \\psi_T $ is the time $ T $ distribution of firm\n",
    "inventory levels implied by the model.\n",
    "\n",
    "We will approximate this distribution by\n",
    "\n",
    "1. fixing $ n $ to be some large number, indicating the number of firms in the\n",
    "  simulation,  \n",
    "1. fixing $ T $, the time period we are interested in,  \n",
    "1. generating $ n $ independent draws from some fixed distribution $ \\psi_0 $ that gives the\n",
    "  initial cross-section of inventories for the $ n $ firms, and  \n",
    "1. shifting this distribution forward in time $ T $ periods, updating each firm\n",
    "  $ T $ times via the dynamics described above (independent of other firms).  \n",
    "\n",
    "\n",
    "We will then visualize $ \\psi_T $ by histogramming the cross-section.\n",
    "\n",
    "We will use the following code to update the cross-section of firms by one period."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a013a54e",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "@jax.jit\n",
    "def update_cross_section(params, X_vec, D):\n",
    "    \"\"\"\n",
    "    Update by one period a cross-section of firms with inventory levels given by\n",
    "    X_vec, given the vector of demand shocks in D.\n",
    "\n",
    "       * D[i] is the demand shock for firm i with current inventory X_vec[i]\n",
    "\n",
    "    \"\"\"\n",
    "    # Unpack\n",
    "    s, S = params.s, params.S\n",
    "    # Restock if the inventory is below the threshold\n",
    "    X_new = jnp.where(X_vec <= s, \n",
    "                      jnp.maximum(S - D, 0), jnp.maximum(X_vec - D, 0))\n",
    "    return X_new"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cbfa8558",
   "metadata": {},
   "source": [
    "### For loop version\n",
    "\n",
    "Now we provide code to compute the cross-sectional distribution $ \\psi_T $ given some\n",
    "initial distribution $ \\psi_0 $ and a positive integer $ T $.\n",
    "\n",
    "In this code we use an ordinary Python `for` loop to step forward through time\n",
    "\n",
    "While Python loops are slow, this approach is reasonable here because\n",
    "efficiency of outer loops has far less influence on runtime than efficiency of inner loops.\n",
    "\n",
    "(Below we will squeeze out more speed by compiling the outer loop as well as the\n",
    "update rule.)\n",
    "\n",
    "In the code below, the initial distribution $ \\psi_0 $ takes all firms to have\n",
    "initial inventory `x_init`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e8252646",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def compute_cross_section(params, x_init, T, key, num_firms=50_000):\n",
    "    # Set up initial distribution\n",
    "    X_vec = jnp.full((num_firms, ), x_init)\n",
    "    # Loop\n",
    "    for i in range(T):\n",
    "        Z = random.normal(key, shape=(num_firms, ))\n",
    "        D = jnp.exp(params.μ + params.σ * Z)\n",
    "\n",
    "        X_vec = update_cross_section(params, X_vec, D)\n",
    "        _, key = random.split(key)\n",
    "\n",
    "    return X_vec"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14a7c3ea",
   "metadata": {},
   "source": [
    "We’ll use the following specification"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ae165921",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "x_init = 50\n",
    "T = 500\n",
    "# Initialize random number generator\n",
    "key = random.PRNGKey(10)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7503d025",
   "metadata": {},
   "source": [
    "Let’s look at the timing."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8bd571fb",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "%time X_vec = compute_cross_section(params, \\\n",
    "        x_init, T, key).block_until_ready()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "03e3c101",
   "metadata": {},
   "source": [
    "Let’s run again to eliminate compile time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9dcba0d1",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "%time X_vec = compute_cross_section(params, \\\n",
    "        x_init, T, key).block_until_ready()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "01b8d54b",
   "metadata": {},
   "source": [
    "Here’s a histogram of inventory levels at time $ T $."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "097049cc",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "ax.hist(X_vec, bins=50, \n",
    "        density=True, \n",
    "        histtype='step', \n",
    "        label=f'cross-section when $t = {T}$')\n",
    "ax.set_xlabel('inventory')\n",
    "ax.set_ylabel('probability')\n",
    "ax.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11f82e8f",
   "metadata": {},
   "source": [
    "### Compiling the outer loop\n",
    "\n",
    "Now let’s see if we can gain some speed by compiling the outer loop, which steps\n",
    "through the time dimension.\n",
    "\n",
    "We will do this using `jax.jit` and a `fori_loop`, which is a compiler-ready version of a `for` loop provided by JAX."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a936491f",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def compute_cross_section_fori(params, x_init, T, key, num_firms=50_000):\n",
    "\n",
    "    s, S, μ, σ = params.s, params.S, params.μ, params.σ\n",
    "    X = jnp.full((num_firms, ), x_init)\n",
    "\n",
    "    # Define the function for each update\n",
    "    def fori_update(t, inputs):\n",
    "        # Unpack\n",
    "        X, key = inputs\n",
    "        # Draw shocks using key\n",
    "        Z = random.normal(key, shape=(num_firms,))\n",
    "        D = jnp.exp(μ + σ * Z)\n",
    "        # Update X\n",
    "        X = jnp.where(X <= s,\n",
    "                  jnp.maximum(S - D, 0),\n",
    "                  jnp.maximum(X - D, 0))\n",
    "        # Refresh the key\n",
    "        key, subkey = random.split(key)\n",
    "        return X, subkey\n",
    "\n",
    "    # Loop t from 0 to T, applying fori_update each time.\n",
    "    # The initial condition for fori_update is (X, key).\n",
    "    X, key = lax.fori_loop(0, T, fori_update, (X, key))\n",
    "\n",
    "    return X\n",
    "\n",
    "# Compile taking T and num_firms as static (changes trigger recompile)\n",
    "compute_cross_section_fori = jax.jit(\n",
    "    compute_cross_section_fori, static_argnums=(2, 4))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "908a1c6c",
   "metadata": {},
   "source": [
    "Let’s see how fast this runs with compile time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "02ba0b27",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "%time X_vec = compute_cross_section_fori(params, \\\n",
    "                x_init, T, key).block_until_ready()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e618537c",
   "metadata": {},
   "source": [
    "And let’s see how fast it runs without compile time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0503f4f7",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "%time X_vec = compute_cross_section_fori(params, \\\n",
    "                x_init, T, key).block_until_ready()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7ed1fa7e",
   "metadata": {},
   "source": [
    "Compared to the original version with a pure Python outer loop, we have\n",
    "produced a nontrivial speed gain.\n",
    "\n",
    "This is due to the fact that we have compiled the whole operation."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c15ed76c",
   "metadata": {},
   "source": [
    "### Further vectorization\n",
    "\n",
    "For relatively small problems, we can make this code run even faster by generating\n",
    "all random variables at once.\n",
    "\n",
    "This improves efficiency because we are taking more operations out of the loop."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e976c94c",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def compute_cross_section_fori(params, x_init, T, key, num_firms=50_000):\n",
    "\n",
    "    s, S, μ, σ = params.s, params.S, params.μ, params.σ\n",
    "    X = jnp.full((num_firms, ), x_init)\n",
    "    Z = random.normal(key, shape=(T, num_firms))\n",
    "    D = jnp.exp(μ + σ * Z)\n",
    "\n",
    "    def update_cross_section(i, X):\n",
    "        X = jnp.where(X <= s,\n",
    "                  jnp.maximum(S - D[i, :], 0),\n",
    "                  jnp.maximum(X - D[i, :], 0))\n",
    "        return X\n",
    "\n",
    "    X = lax.fori_loop(0, T, update_cross_section, X)\n",
    "\n",
    "    return X\n",
    "\n",
    "# Compile taking T and num_firms as static (changes trigger recompile)\n",
    "compute_cross_section_fori = jax.jit(\n",
    "    compute_cross_section_fori, static_argnums=(2, 4))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "89d1566f",
   "metadata": {},
   "source": [
    "Let’s test it with compile time included."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1bbf962b",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "%time X_vec = compute_cross_section_fori(params, \\\n",
    "                x_init, T, key).block_until_ready()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5c895a0a",
   "metadata": {},
   "source": [
    "Let’s run again to eliminate compile time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "546bbc85",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "%time X_vec = compute_cross_section_fori(params, \\\n",
    "                x_init, T, key).block_until_ready()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1ec0df53",
   "metadata": {},
   "source": [
    "On one hand, this version is faster than the previous one, where random variables were\n",
    "generated inside the loop.\n",
    "\n",
    "On the other hand, this implementation consumes far more memory, as we need to\n",
    "store large arrays of random draws.\n",
    "\n",
    "The high memory consumption becomes problematic for large problems."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c88fde07",
   "metadata": {},
   "source": [
    "## Distribution dynamics\n",
    "\n",
    "Next let’s take a look at how the distribution sequence evolves over time.\n",
    "\n",
    "We will go back to using ordinary Python `for` loops.\n",
    "\n",
    "Here is code that repeatedly shifts the cross-section forward while\n",
    "recording the cross-section at the dates in `sample_dates`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aa8351ea",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def shift_forward_and_sample(x_init, params, sample_dates,\n",
    "                        key, num_firms=50_000, sim_length=750):\n",
    "\n",
    "    X = res = jnp.full((num_firms, ), x_init)\n",
    "\n",
    "    # Use for loop to update X and collect samples\n",
    "    for i in range(sim_length):\n",
    "        Z = random.normal(key, shape=(num_firms, ))\n",
    "        D = jnp.exp(params.μ + params.σ * Z)\n",
    "\n",
    "        X = update_cross_section(params, X, D)\n",
    "        _, key = random.split(key)\n",
    "\n",
    "        # draw a sample at the sample dates\n",
    "        if (i+1 in sample_dates):\n",
    "          res = jnp.vstack((res, X))\n",
    "\n",
    "    return res[1:]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8755cba8",
   "metadata": {},
   "source": [
    "Let’s test it"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6b57d28c",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "x_init = 50\n",
    "num_firms = 10_000\n",
    "sample_dates = 10, 50, 250, 500, 750\n",
    "key = random.PRNGKey(10)\n",
    "\n",
    "\n",
    "%time X = shift_forward_and_sample(x_init, params, \\\n",
    "                    sample_dates, key).block_until_ready()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "21a247ed",
   "metadata": {},
   "source": [
    "We run the code again to eliminate compile time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1fcf44f1",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "%time X = shift_forward_and_sample(x_init, params, \\\n",
    "                    sample_dates, key).block_until_ready()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7f36a388",
   "metadata": {},
   "source": [
    "Let’s plot the output."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4e423c44",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "\n",
    "for i, date in enumerate(sample_dates):\n",
    "    ax.hist(X[i, :], bins=50, \n",
    "            density=True, \n",
    "            histtype='step',\n",
    "            label=f'cross-section when $t = {date}$')\n",
    "\n",
    "ax.set_xlabel('inventory')\n",
    "ax.set_ylabel('probability')\n",
    "ax.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3bbb820c",
   "metadata": {},
   "source": [
    "This model for inventory dynamics is asymptotically stationary, with a unique\n",
    "stationary distribution.\n",
    "\n",
    "In particular, the sequence of marginal distributions $ \\{\\psi_t\\} $\n",
    "converges to a unique limiting distribution that does not depend on\n",
    "initial conditions.\n",
    "\n",
    "Although we will not prove this here, we can see it in the simulation above.\n",
    "\n",
    "By $ t=500 $ or $ t=750 $ the distributions are barely changing.\n",
    "\n",
    "If you test a few different initial conditions, you will see that they do not affect long-run outcomes."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "93f549d6",
   "metadata": {},
   "source": [
    "## Restock frequency\n",
    "\n",
    "As an exercise, let’s study the probability that firms need to restock over a given time period.\n",
    "\n",
    "In the exercise, we will\n",
    "\n",
    "- set the starting stock level to $ X_0 = 70 $ and  \n",
    "- calculate the proportion of firms that need to order twice or more in the first 50 periods.  \n",
    "\n",
    "\n",
    "This proportion approximates the probability of the event when the sample size\n",
    "is large."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5300f2d3",
   "metadata": {},
   "source": [
    "### For loop version\n",
    "\n",
    "We start with an easier `for` loop implementation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e796e216",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "# Define a jitted function for each update\n",
    "@jax.jit\n",
    "def update_stock(n_restock, X, params, D):\n",
    "    n_restock = jnp.where(X <= params.s,\n",
    "                          n_restock + 1,\n",
    "                          n_restock)\n",
    "    X = jnp.where(X <= params.s,\n",
    "                  jnp.maximum(params.S - D, 0),\n",
    "                  jnp.maximum(X - D, 0))\n",
    "    return n_restock, X, key\n",
    "\n",
    "def compute_freq(params, key,\n",
    "                 x_init=70,\n",
    "                 sim_length=50,\n",
    "                 num_firms=1_000_000):\n",
    "\n",
    "    # Prepare initial arrays\n",
    "    X = jnp.full((num_firms, ), x_init)\n",
    "\n",
    "    # Stack the restock counter on top of the inventory\n",
    "    n_restock = jnp.zeros((num_firms, ))\n",
    "\n",
    "    # Use a for loop to perform the calculations on all states\n",
    "    for i in range(sim_length):\n",
    "        Z = random.normal(key, shape=(num_firms, ))\n",
    "        D = jnp.exp(params.μ + params.σ * Z)\n",
    "        n_restock, X, key = update_stock(\n",
    "            n_restock, X, params, D)\n",
    "        key = random.fold_in(key, i)\n",
    "\n",
    "    return jnp.mean(n_restock > 1, axis=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5d8f6ddb",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "key = random.PRNGKey(27)\n",
    "\n",
    "%time freq = compute_freq(params, key).block_until_ready()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5f6b4b3b",
   "metadata": {},
   "source": [
    "We run the code again to get rid of compile time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6a58f6be",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "%time freq = compute_freq(params, key).block_until_ready()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d1039f16",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "print(f\"Frequency of at least two stock outs = {freq}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e1009588",
   "metadata": {},
   "source": [
    "### Exercise 5.1\n",
    "\n",
    "Write a `fori_loop` version of the last function.  See if you can increase the\n",
    "speed while generating a similar answer."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f50a0881",
   "metadata": {},
   "source": [
    "### Solution to[ Exercise 5.1](https://jax.quantecon.org/#inventory_dynamics_ex1)\n",
    "\n",
    "Here is a `lax.fori_loop` version that JIT compiles the whole function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b0192c10",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "@jax.jit\n",
    "def compute_freq(params, key,\n",
    "                 x_init=70,\n",
    "                 sim_length=50,\n",
    "                 num_firms=1_000_000):\n",
    "\n",
    "    s, S, μ, σ = params.s, params.S, params.μ, params.σ\n",
    "    # Prepare initial arrays\n",
    "    X = jnp.full((num_firms, ), x_init)\n",
    "    Z = random.normal(key, shape=(sim_length, num_firms))\n",
    "    D = jnp.exp(μ + σ * Z)\n",
    "\n",
    "    # Stack the restock counter on top of the inventory\n",
    "    restock_count = jnp.zeros((num_firms, ))\n",
    "    Xs = (X, restock_count)\n",
    "\n",
    "    # Define the function for each update\n",
    "    def update_cross_section(i, Xs):\n",
    "        # Separate the inventory and restock counter\n",
    "        x, restock_count = Xs[0], Xs[1]\n",
    "        restock_count = jnp.where(x <= s,\n",
    "                                restock_count + 1,\n",
    "                                restock_count)\n",
    "        x = jnp.where(x <= s,\n",
    "                      jnp.maximum(S - D[i], 0),\n",
    "                      jnp.maximum(x - D[i], 0))\n",
    "\n",
    "        Xs = (x, restock_count)\n",
    "        return Xs\n",
    "\n",
    "    # Use lax.fori_loop to perform the calculations on all states\n",
    "    X_final = lax.fori_loop(0, sim_length, update_cross_section, Xs)\n",
    "\n",
    "    return jnp.mean(X_final[1] > 1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c1b7a163",
   "metadata": {},
   "source": [
    "Note the time the routine takes to run, as well as the output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f74ba10b",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "%time freq = compute_freq(params, key).block_until_ready()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eafd1300",
   "metadata": {},
   "source": [
    "We run the code again to eliminate the compile time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "013d521a",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "%time freq = compute_freq(params, key).block_until_ready()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d1f5f5a6",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "print(f\"Frequency of at least two stock outs = {freq}\")"
   ]
  }
 ],
 "metadata": {
  "date": 1732064317.7557838,
  "filename": "inventory_dynamics.md",
  "kernelspec": {
   "display_name": "Python",
   "language": "python3",
   "name": "python3"
  },
  "title": "Inventory Dynamics"
 },
 "nbformat": 4,
 "nbformat_minor": 5
}