day_4/european_option_jax.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "b93734d3",
   "metadata": {},
   "source": [
    "## Pricing a European Call Option (Numba vs JAX)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ee528fd1",
   "metadata": {},
   "source": [
    "#### Written for the CBC QuantEcon Workshop (September 2022)\n",
    "\n",
    "#### Author: [John Stachurski](http://johnstachurski.net/)\n",
    "\n",
    "\n",
    "Previously we computed the value of a European call option via Monte Carlo using Numba-based routines.\n",
    "\n",
    "Here we try the same operations using JAX."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "cfdd9072",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aa8bbfa1",
   "metadata": {},
   "source": [
    "Recall that we want to compute\n",
    "\n",
    "\n",
    "$$ P = \\beta^n \\mathbb E \\max\\{ S_n - K, 0 \\} $$\n",
    "\n",
    "We suppose that"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "f6528adb",
   "metadata": {},
   "outputs": [],
   "source": [
    "n = 20\n",
    "β = 0.99\n",
    "K = 100"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5741765b",
   "metadata": {},
   "source": [
    "The dynamics are\n",
    "\n",
    "$$ \\ln \\frac{S_{t+1}}{S_t} = \\mu + \\sigma_t \\xi_{t+1} $$\n",
    "\n",
    "where \n",
    "\n",
    "$$ \n",
    "    \\sigma_t = \\exp(h_t), \n",
    "    \\quad\n",
    "        h_{t+1} = \\rho h_t + \\nu \\eta_{t+1}\n",
    "$$\n",
    "\n",
    "Here $\\{\\xi_t\\}$ and $\\{\\eta_t\\}$ are IID and standard normal."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5058d5ca",
   "metadata": {},
   "source": [
    "With $s_t := \\ln S_t$, the price dynamics become\n",
    "\n",
    "$$ s_{t+1} = s_t + \\mu + \\exp(h_t) \\xi_{t+1} $$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8aa8da35",
   "metadata": {},
   "source": [
    "We use the following defaults."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "e1247de4",
   "metadata": {},
   "outputs": [],
   "source": [
    "μ  = 0.0001\n",
    "ρ  = 0.1\n",
    "ν  = 0.001\n",
    "S0 = 10\n",
    "h0 = 0"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1e050d4a",
   "metadata": {},
   "source": [
    "(Here `S0` is $S_0$ and `h0` is $h_0$.)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0f9a0b27",
   "metadata": {},
   "source": [
    "We used the following estimate of the price, computed via Monte Carlo and applying Numba and parallelization."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "546888a1",
   "metadata": {},
   "outputs": [],
   "source": [
    "from numba import njit, prange\n",
    "from numpy.random import randn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "29c445a5",
   "metadata": {},
   "outputs": [],
   "source": [
    "M = 10_000_000"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "49ba9d36",
   "metadata": {},
   "outputs": [],
   "source": [
    "@njit(parallel=True)\n",
    "def compute_call_price_parallel(β=β,\n",
    "                                μ=μ,\n",
    "                                S0=S0,\n",
    "                                h0=h0,\n",
    "                                K=K,\n",
    "                                n=n,\n",
    "                                ρ=ρ,\n",
    "                                ν=ν,\n",
    "                                M=M):\n",
    "    current_sum = 0.0\n",
    "    # For each sample path\n",
    "    for m in prange(M):\n",
    "        s = np.log(S0)\n",
    "        h = h0\n",
    "        # Simulate forward in time\n",
    "        for t in range(n):\n",
    "            s = s + μ + np.exp(h) * randn()\n",
    "            h = ρ * h + ν * randn()\n",
    "        # And add the value max{S_n - K, 0} to current_sum\n",
    "        current_sum += np.maximum(np.exp(s) - K, 0)\n",
    "        \n",
    "    return β**n * current_sum / M"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "200bed14",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "6"
      ]
     },
     "execution_count": 40,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from numba import get_num_threads, set_num_threads\n",
    "get_num_threads()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "c1ddc360",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 12 s, sys: 4.07 ms, total: 12 s\n",
      "Wall time: 2.47 s\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "121617.90871201943"
      ]
     },
     "execution_count": 38,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%time\n",
    "compute_call_price_parallel()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "15c645cc",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 14.4 s, sys: 3.96 ms, total: 14.4 s\n",
      "Wall time: 1.96 s\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "124535.64389811848"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%time\n",
    "compute_call_price_parallel()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2c2444dd",
   "metadata": {},
   "source": [
    "### Exercise\n",
    "\n",
    "Try to shift the whole operation to the GPU using JAX and test your speed gain."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c24104d0",
   "metadata": {},
   "source": [
    "### Solution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "8a51bfdc",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Sat Sep 17 05:47:33 2022       \r\n",
      "+-----------------------------------------------------------------------------+\r\n",
      "| NVIDIA-SMI 515.65.01    Driver Version: 515.65.01    CUDA Version: 11.7     |\r\n",
      "|-------------------------------+----------------------+----------------------+\r\n",
      "| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |\r\n",
      "| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |\r\n",
      "|                               |                      |               MIG M. |\r\n",
      "|===============================+======================+======================|\r\n",
      "|   0  NVIDIA GeForce ...  On   | 00000000:3B:00.0 Off |                  N/A |\r\n",
      "| 30%   29C    P8    25W / 320W |   9042MiB / 10240MiB |      0%      Default |\r\n",
      "|                               |                      |                  N/A |\r\n",
      "+-------------------------------+----------------------+----------------------+\r\n",
      "                                                                               \r\n",
      "+-----------------------------------------------------------------------------+\r\n",
      "| Processes:                                                                  |\r\n",
      "|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |\r\n",
      "|        ID   ID                                                   Usage      |\r\n",
      "|=============================================================================|\r\n",
      "|    0   N/A  N/A   1442158      C   ...john/anaconda3/bin/python     9039MiB |\r\n",
      "+-----------------------------------------------------------------------------+\r\n"
     ]
    }
   ],
   "source": [
    "!nvidia-smi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "cb274020",
   "metadata": {},
   "outputs": [],
   "source": [
    "import jax\n",
    "import jax.numpy as jnp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "17428b84",
   "metadata": {},
   "outputs": [],
   "source": [
    "@jax.jit\n",
    "def compute_call_price_jax(β=β,\n",
    "                           μ=μ,\n",
    "                           S0=S0,\n",
    "                           h0=h0,\n",
    "                           K=K,\n",
    "                           n=n,\n",
    "                           ρ=ρ,\n",
    "                           ν=ν,\n",
    "                           M=M,\n",
    "                           key=jax.random.PRNGKey(1)):\n",
    "\n",
    "    s = jnp.full(M, np.log(S0))\n",
    "    h = jnp.full(M, h0)\n",
    "    for t in range(n):\n",
    "        key, subkey = jax.random.split(key)\n",
    "        Z = jax.random.normal(subkey, (2, M))\n",
    "        s = s + μ + jnp.exp(h) * Z[0, :]\n",
    "        h = ρ * h + ν * Z[1, :]\n",
    "    expectation = jnp.mean(jnp.maximum(jnp.exp(s) - K, 0))\n",
    "        \n",
    "    return β**n * expectation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "4009135c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 2.71 s, sys: 55.9 ms, total: 2.77 s\n",
      "Wall time: 4.43 s\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "DeviceArray(152444.25, dtype=float32)"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%time \n",
    "compute_call_price_jax().block_until_ready()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "377fa60d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 1.71 ms, sys: 3 µs, total: 1.71 ms\n",
      "Wall time: 35.2 ms\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "DeviceArray(152444.25, dtype=float32)"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%time \n",
    "compute_call_price_jax().block_until_ready()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2f28bc7f",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}