kesten_processes.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "eef9a741",
   "metadata": {},
   "source": [
    "# Kesten Processes and Firm Dynamics\n",
    "\n",
    "\n",
    "<a id='index-0'></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6c9bd872",
   "metadata": {},
   "source": [
    "# GPU\n",
    "\n",
    "This lecture was built using a machine with JAX installed and access to a GPU.\n",
    "\n",
    "To run this lecture on [Google Colab](https://colab.research.google.com/), click on the “play” icon top right, select Colab, and set the runtime environment to include a GPU.\n",
    "\n",
    "To run this lecture on your own machine, you need to install [Google JAX](https://github.com/google/jax).\n",
    "\n",
    "In addition to JAX and Anaconda, this lecture will need the following libraries:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e5f7dfe1",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "!pip install quantecon"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0282e193",
   "metadata": {},
   "source": [
    "## Overview\n",
    "\n",
    "This lecture describes Kesten processes, which are an important class of\n",
    "stochastic processes, and an application of firm dynamics.\n",
    "\n",
    "The lecture draws on [an earlier QuantEcon\n",
    "lecture](https://python.quantecon.org/kesten_processes.html), which uses Numba\n",
    "to accelerate the computations.\n",
    "\n",
    "In that earlier lecture you can find a more detailed discussion of the concepts involved.\n",
    "\n",
    "This lecture focuses on implementing the same computations in JAX.\n",
    "\n",
    "Let’s start with some imports:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1cf2058e",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import quantecon as qe\n",
    "import jax\n",
    "import jax.numpy as jnp\n",
    "from jax import random\n",
    "from jax import lax"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "07b78c41",
   "metadata": {},
   "source": [
    "Let’s check the GPU we are running"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10b5de65",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "!nvidia-smi"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c64fcf6f",
   "metadata": {},
   "source": [
    "## Kesten processes\n",
    "\n",
    "\n",
    "<a id='index-1'></a>\n",
    "A **Kesten process** is a stochastic process of the form\n",
    "\n",
    "\n",
    "<a id='equation-kesproc'></a>\n",
    "$$\n",
    "X_{t+1} = a_{t+1} X_t + \\eta_{t+1} \\tag{6.1}\n",
    "$$\n",
    "\n",
    "where $ \\{a_t\\}_{t \\geq 1} $ and $ \\{\\eta_t\\}_{t \\geq 1} $ are IID\n",
    "sequences.\n",
    "\n",
    "We are interested in the dynamics of $ \\{X_t\\}_{t \\geq 0} $ when $ X_0 $ is given.\n",
    "\n",
    "We will focus on the nonnegative scalar case, where $ X_t $ takes values in $ \\mathbb R_+ $.\n",
    "\n",
    "In particular, we will assume that\n",
    "\n",
    "- the initial condition $ X_0 $ is nonnegative,  \n",
    "- $ \\{a_t\\}_{t \\geq 1} $ is a nonnegative IID stochastic process and  \n",
    "- $ \\{\\eta_t\\}_{t \\geq 1} $ is another nonnegative IID stochastic process, independent of the first.  "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "20c25847",
   "metadata": {},
   "source": [
    "### Application: firm dynamics\n",
    "\n",
    "In this section we apply Kesten process theory to the study of firm dynamics."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "498730fd",
   "metadata": {},
   "source": [
    "#### Gibrat’s law\n",
    "\n",
    "It was postulated many years ago by Robert Gibrat that firm size evolves\n",
    "according to a simple rule whereby size next period is proportional to current\n",
    "size.\n",
    "\n",
    "This is now known as [Gibrat’s law of proportional growth](https://en.wikipedia.org/wiki/Gibrat%27s_law).\n",
    "\n",
    "We can express this idea by stating that a suitably defined measure\n",
    "$ s_t $ of firm size obeys\n",
    "\n",
    "\n",
    "<a id='equation-firm-dynam-gb'></a>\n",
    "$$\n",
    "\\frac{s_{t+1}}{s_t} = a_{t+1} \\tag{6.2}\n",
    "$$\n",
    "\n",
    "for some positive IID sequence $ \\{a_t\\} $.\n",
    "\n",
    "Subsequent empirical research has shown that this specification is not accurate,\n",
    "particularly for small firms.\n",
    "\n",
    "However, we can get close to the data by modifying [(6.2)](#equation-firm-dynam-gb) to\n",
    "\n",
    "\n",
    "<a id='equation-firm-dynam'></a>\n",
    "$$\n",
    "s_{t+1} = a_{t+1} s_t + b_{t+1} \\tag{6.3}\n",
    "$$\n",
    "\n",
    "where $ \\{a_t\\} $ and $ \\{b_t\\} $ are both IID and independent of each\n",
    "other.\n",
    "\n",
    "We now study the implications of this specification."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0db552dd",
   "metadata": {},
   "source": [
    "#### Heavy tails\n",
    "\n",
    "If the conditions of the [Kesten–Goldie\n",
    "Theorem](https://python.quantecon.org/kesten_processes.html#the-kestengoldie-theorem)\n",
    "are satisfied, then [(6.3)](#equation-firm-dynam) implies that the firm size distribution\n",
    "will have Pareto tails.\n",
    "\n",
    "This matches empirical findings across many data sets.\n",
    "\n",
    "But there is another unrealistic aspect of the firm dynamics specified in [(6.3)](#equation-firm-dynam) that we need to address: it ignores entry and exit.\n",
    "\n",
    "In any given period and in any given market, we observe significant numbers of\n",
    "firms entering and exiting the market.\n",
    "\n",
    "In this setting, firm dynamics can be expressed as\n",
    "\n",
    "\n",
    "<a id='equation-firm-dynam-ee'></a>\n",
    "$$\n",
    "s_{t+1} = e_{t+1} \\mathbb{1}\\{s_t < \\bar s\\} +\n",
    "    (a_{t+1} s_t + b_{t+1}) \\mathbb{1}\\{s_t \\geq \\bar s\\} \\tag{6.4}\n",
    "$$\n",
    "\n",
    "The motivation behind and interpretation of [(6.4)](#equation-firm-dynam-ee) can be found in\n",
    "[our earlier Kesten process lecture](https://python.quantecon.org/kesten_processes.html).\n",
    "\n",
    "What can we say about dynamics?\n",
    "\n",
    "Although [(6.4)](#equation-firm-dynam-ee) is not a Kesten process, it does update in the\n",
    "same way as a Kesten process when $ s_t $ is large.\n",
    "\n",
    "So perhaps its stationary distribution still has Pareto tails?\n",
    "\n",
    "We can investigate this question via simulation and rank-size plots.\n",
    "\n",
    "The approach will be to\n",
    "\n",
    "1. generate $ M $ draws of $ s_T $ when $ M $ and $ T $ are\n",
    "  large and  \n",
    "1. plot the largest 1,000 of the resulting draws in a rank-size plot.  \n",
    "\n",
    "\n",
    "(The distribution of $ s_T $ will be close to the stationary distribution\n",
    "when $ T $ is large.)\n",
    "\n",
    "In the simulation, we assume that each of $ a_t, b_t $ and $ e_t $ is lognormal.\n",
    "\n",
    "Here’s code to update a cross-section of firms according to the dynamics in\n",
    "[(6.4)](#equation-firm-dynam-ee)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "472d8fc7",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "@jax.jit\n",
    "def update_s(s, s_bar, a_random, b_random, e_random):\n",
    "    exp_a = jnp.exp(a_random)\n",
    "    exp_b = jnp.exp(b_random)\n",
    "    exp_e = jnp.exp(e_random)\n",
    "\n",
    "    s = jnp.where(s < s_bar,\n",
    "                  exp_e,\n",
    "                  exp_a * s + exp_b)\n",
    "\n",
    "    return s"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7a95520d",
   "metadata": {},
   "source": [
    "Now we write a for loop that repeatedly calls this function, to push a\n",
    "cross-section of firms forward in time.\n",
    "\n",
    "For sufficiently large `T`, the cross-section it returns (the cross-section at\n",
    "time `T`) corresponds to firm size distribution in (approximate) equilibrium."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ed05ba60",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def generate_draws(M=1_000_000,\n",
    "                   μ_a=-0.5,\n",
    "                   σ_a=0.1,\n",
    "                   μ_b=0.0,\n",
    "                   σ_b=0.5,\n",
    "                   μ_e=0.0,\n",
    "                   σ_e=0.5,\n",
    "                   s_bar=1.0,\n",
    "                   T=500,\n",
    "                   s_init=1.0,\n",
    "                   seed=123):\n",
    "\n",
    "    key = random.PRNGKey(seed)\n",
    "\n",
    "    # Initialize the array of s values with the initial value\n",
    "    s = jnp.full((M, ), s_init)\n",
    "\n",
    "    # Perform updates on s for time t\n",
    "    for t in range(T):\n",
    "        keys = random.split(key, 3)\n",
    "        a_random = μ_a + σ_a * random.normal(keys[0], (M, ))\n",
    "        b_random = μ_b + σ_b * random.normal(keys[1], (M, ))\n",
    "        e_random = μ_e + σ_e * random.normal(keys[2], (M, ))\n",
    "\n",
    "        s = update_s(s, s_bar, a_random, b_random, e_random)\n",
    "        \n",
    "        # Generate new key for the next iteration\n",
    "        key = random.fold_in(key, t)\n",
    "\n",
    "    return s\n",
    "\n",
    "%time data = generate_draws().block_until_ready()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9721ba60",
   "metadata": {},
   "source": [
    "Running the above function again so we can see the speed with and without compile time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "492719d0",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "%time data = generate_draws().block_until_ready()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "91efa9fb",
   "metadata": {},
   "source": [
    "Notice that we do not JIT-compile the `for` loops, since\n",
    "\n",
    "1. acceleration of the outer loop makes little difference terms of compute\n",
    "  time and  \n",
    "1. compiling the outer loop is often very slow.  \n",
    "\n",
    "\n",
    "Let’s produce the rank-size plot and check the distribution:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "41e58792",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "\n",
    "rank_data, size_data = qe.rank_size(data, c=0.01)\n",
    "ax.loglog(rank_data, size_data, 'o', markersize=3.0, alpha=0.5)\n",
    "ax.set_xlabel(\"log rank\")\n",
    "ax.set_ylabel(\"log size\")\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3610fc69",
   "metadata": {},
   "source": [
    "The plot produces a straight line, consistent with a Pareto tail."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a65aa6f3",
   "metadata": {},
   "source": [
    "#### Alternative implementation with `lax.fori_loop`\n",
    "\n",
    "If the time horizon is not too large, we can try to further accelerate our code\n",
    "by replacing the `for` loop with\n",
    "[`lax.fori_loop`](https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.fori_loop.html).\n",
    "\n",
    "Note, however, that\n",
    "\n",
    "1. as mentioned above, there is not much speed gain in accelerating outer loops,  \n",
    "1. `lax.fori_loop` has a more complicated syntax, and, most importantly,  \n",
    "1. the `lax.fori_loop` implementation consumes far more memory, as we need to have to\n",
    "  store large matrices of random draws  \n",
    "\n",
    "\n",
    "Hence the code below will fail due to out-of-memory errors when `T` and `M` are large.\n",
    "\n",
    "Here is the `lax.fori_loop` version:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e0b8ca7a",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "@jax.jit\n",
    "def generate_draws_lax(μ_a=-0.5,\n",
    "                       σ_a=0.1,\n",
    "                       μ_b=0.0,\n",
    "                       σ_b=0.5,\n",
    "                       μ_e=0.0,\n",
    "                       σ_e=0.5,\n",
    "                       s_bar=1.0,\n",
    "                       T=500,\n",
    "                       M=500_000,\n",
    "                       s_init=1.0,\n",
    "                       seed=123):\n",
    "\n",
    "    key = random.PRNGKey(seed)\n",
    "    keys = random.split(key, 3)\n",
    "    \n",
    "    # Generate random draws and initial values\n",
    "    a_random = μ_a + σ_a * random.normal(keys[0], (T, M))\n",
    "    b_random = μ_b + σ_b * random.normal(keys[1], (T, M))\n",
    "    e_random = μ_e + σ_e * random.normal(keys[2], (T, M))\n",
    "    s = jnp.full((M, ), s_init)\n",
    "\n",
    "    # Define the function for each update\n",
    "    def update_s(i, s):\n",
    "        a, b, e = a_random[i], b_random[i], e_random[i]\n",
    "        s = jnp.where(s < s_bar,\n",
    "                      jnp.exp(e),\n",
    "                      jnp.exp(a) * s + jnp.exp(b))\n",
    "        return s\n",
    "\n",
    "    # Use lax.scan to perform the calculations on all states\n",
    "    s_final = lax.fori_loop(0, T, update_s, s)\n",
    "    return s_final\n",
    "\n",
    "%time data = generate_draws_lax().block_until_ready()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e6ccb633",
   "metadata": {},
   "source": [
    "In this case, `M` is small enough for the code to run and\n",
    "we see some speed gain over the for loop implementation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fb001731",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "%time data = generate_draws_lax().block_until_ready()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bd4080f5",
   "metadata": {},
   "source": [
    "Here we produce the same rank-size plot:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d35ce445",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "\n",
    "rank_data, size_data = qe.rank_size(data, c=0.01)\n",
    "ax.loglog(rank_data, size_data, 'o', markersize=3.0, alpha=0.5)\n",
    "ax.set_xlabel(\"log rank\")\n",
    "ax.set_ylabel(\"log size\")\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "85bbbd2b",
   "metadata": {},
   "source": [
    "Let’s rerun the `for` loop version on smaller `M` to compare the speed"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "478cbe58",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "%time generate_draws(M=500_000).block_until_ready()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ec659b54",
   "metadata": {},
   "source": [
    "Let’s run it again to get rid of the compile time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "379bdb80",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "%time generate_draws(M=500_000).block_until_ready()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fcdf24b6",
   "metadata": {},
   "source": [
    "We see that the `lax.fori_loop` version is faster than the `for` loop version\n",
    "when memory is not an issue."
   ]
  }
 ],
 "metadata": {
  "date": 1729509368.9640515,
  "filename": "kesten_processes.md",
  "kernelspec": {
   "display_name": "Python",
   "language": "python3",
   "name": "python3"
  },
  "title": "Kesten Processes and Firm Dynamics"
 },
 "nbformat": 4,
 "nbformat_minor": 5
}