parallelization.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "<a id='parallel'></a>\n",
    "<div id=\"qe-notebook-header\" align=\"right\" style=\"text-align:right;\">\n",
    "        <a href=\"https://quantecon.org/\" title=\"quantecon.org\">\n",
    "                <img style=\"width:250px;display:inline;\" width=\"250px\" src=\"https://assets.quantecon.org/img/qe-menubar-logo.svg\" alt=\"QuantEcon\">\n",
    "        </a>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Parallelization"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Contents\n",
    "\n",
    "- [Parallelization](#Parallelization)  \n",
    "  - [Overview](#Overview)  \n",
    "  - [Types of Parallelization](#Types-of-Parallelization)  \n",
    "  - [Implicit Multithreading in NumPy](#Implicit-Multithreading-in-NumPy)  \n",
    "  - [Multithreaded Loops in Numba](#Multithreaded-Loops-in-Numba)  \n",
    "  - [Exercises](#Exercises)  \n",
    "  - [Solutions](#Solutions)  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In addition to what’s in Anaconda, this lecture will need the following libraries:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": true
   },
   "outputs": [],
   "source": [
    "!pip install --upgrade quantecon"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Overview\n",
    "\n",
    "The growth of CPU clock speed (i.e., the speed at which a single chain of logic can\n",
    "be run) has slowed dramatically in recent years.\n",
    "\n",
    "This is unlikely to change in the near future, due to inherent physical\n",
    "limitations on the construction of chips and circuit boards.\n",
    "\n",
    "Chip designers and computer programmers have responded to the slowdown by\n",
    "seeking a different path to fast execution: parallelization.\n",
    "\n",
    "Hardware makers have increased the number of cores (physical CPUs) embedded in each machine.\n",
    "\n",
    "For programmers, the challenge has been to exploit these multiple CPUs by running many processes in parallel (i.e., simultaneously).\n",
    "\n",
    "This is particularly important in scientific programming, which requires handling\n",
    "\n",
    "- large amounts of data and  \n",
    "- CPU intensive simulations and other calculations.  \n",
    "\n",
    "\n",
    "In this lecture we discuss parallelization for scientific computing, with a focus on\n",
    "\n",
    "1. the best tools for parallelization in Python and  \n",
    "1. how these tools can be applied to quantitative economic problems.  \n",
    "\n",
    "\n",
    "Let’s start with some imports:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import quantecon as qe\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Types of Parallelization\n",
    "\n",
    "Large textbooks have been written on different approaches to parallelization but we will keep a tight focus on what’s most useful to us.\n",
    "\n",
    "We will briefly review the two main kinds of parallelization commonly used in\n",
    "scientific computing and discuss their pros and cons."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Multiprocessing\n",
    "\n",
    "Multiprocessing means concurrent execution of multiple processes using more than one processor.\n",
    "\n",
    "In this context, a **process** is a chain of instructions (i.e., a program).\n",
    "\n",
    "Multiprocessing can be carried out on one machine with multiple CPUs or on a\n",
    "collection of machines connected by a network.\n",
    "\n",
    "In the latter case, the collection of machines is usually called a\n",
    "**cluster**.\n",
    "\n",
    "With multiprocessing, each process has its own memory space, although the\n",
    "physical memory chip might be shared."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Multithreading\n",
    "\n",
    "Multithreading is similar to multiprocessing, except that, during execution, the threads all share the same memory space.\n",
    "\n",
    "Native Python struggles to implement multithreading due to some [legacy design\n",
    "features](https://wiki.python.org/moin/GlobalInterpreterLock).\n",
    "\n",
    "But this is not a restriction for scientific libraries like NumPy and Numba.\n",
    "\n",
    "Functions imported from these libraries and JIT-compiled code run in low level\n",
    "execution environments where Python’s legacy restrictions don’t apply."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Advantages and Disadvantages\n",
    "\n",
    "Multithreading is more lightweight because most system and memory resources\n",
    "are shared by the threads.\n",
    "\n",
    "In addition, the fact that multiple threads all access a shared pool of memory\n",
    "is extremely convenient for numerical programming.\n",
    "\n",
    "On the other hand, multiprocessing is more flexible and can be distributed\n",
    "across clusters.\n",
    "\n",
    "For the great majority of what we do in these lectures, multithreading will\n",
    "suffice."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Implicit Multithreading in NumPy\n",
    "\n",
    "Actually, you have already been using multithreading in your Python code,\n",
    "although you might not have realized it.\n",
    "\n",
    "(We are, as usual, assuming that you are running the latest version of\n",
    "Anaconda Python.)\n",
    "\n",
    "This is because NumPy cleverly implements multithreading in a lot of its\n",
    "compiled code.\n",
    "\n",
    "Let’s look at some examples to see this in action."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### A Matrix Operation\n",
    "\n",
    "The next piece of code computes the eigenvalues of a large number of randomly\n",
    "generated matrices.\n",
    "\n",
    "It takes a few seconds to run."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "n = 20\n",
    "m = 1000\n",
    "for i in range(n):\n",
    "    X = np.random.randn(m, m)\n",
    "    λ = np.linalg.eigvals(X)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, let’s look at the output of the htop system monitor on our machine while\n",
    "this code is running:\n",
    "\n",
    "<img src=\"https://s3-ap-southeast-2.amazonaws.com/python.quantecon.org/_static/lecture_specific/parallelization/htop_parallel_npmat.png\" style=\"\">\n",
    "\n",
    "  \n",
    "We can see that 4 of the 8 CPUs are running at full speed.\n",
    "\n",
    "This is because NumPy’s `eigvals` routine neatly splits up the tasks and\n",
    "distributes them to different threads."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### A Multithreaded Ufunc\n",
    "\n",
    "Over the last few years, NumPy has managed to push this kind of multithreading\n",
    "out to more and more operations.\n",
    "\n",
    "For example, let’s return to a maximization problem [discussed previously](https://python.quantecon.org/need_for_speed.html#ufuncs):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def f(x, y):\n",
    "    return np.cos(x**2 + y**2) / (1 + x**2 + y**2)\n",
    "\n",
    "grid = np.linspace(-3, 3, 5000)\n",
    "x, y = np.meshgrid(grid, grid)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "%timeit np.max(f(x, y))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you have a system monitor such as htop (Linux/Mac) or perfmon\n",
    "(Windows), then try running this and then observing the load on your CPUs.\n",
    "\n",
    "(You will probably need to bump up the grid size to see large effects.)\n",
    "\n",
    "At least on our machine, the output shows that the operation is successfully\n",
    "distributed across multiple threads.\n",
    "\n",
    "This is one of the reasons why the vectorized code above is fast."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### A Comparison with Numba\n",
    "\n",
    "To get some basis for comparison for the last example, let’s try the same\n",
    "thing with Numba.\n",
    "\n",
    "In fact there is an easy way to do this, since Numba can also be used to\n",
    "create custom [ufuncs](https://python.quantecon.org/need_for_speed.html#ufuncs) with the [@vectorize](http://numba.pydata.org/numba-doc/dev/user/vectorize.html) decorator."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "from numba import vectorize\n",
    "\n",
    "@vectorize\n",
    "def f_vec(x, y):\n",
    "    return np.cos(x**2 + y**2) / (1 + x**2 + y**2)\n",
    "\n",
    "np.max(f_vec(x, y))  # Run once to compile"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "%timeit np.max(f_vec(x, y))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "At least on our machine, the difference in the speed between the\n",
    "Numba version and the vectorized NumPy version shown above is not large.\n",
    "\n",
    "But there’s quite a bit going on here so let’s try to break down what is\n",
    "happening.\n",
    "\n",
    "Both Numba and NumPy use efficient machine code that’s specialized to these\n",
    "floating point operations.\n",
    "\n",
    "However, the code NumPy uses is, in some ways, less efficient.\n",
    "\n",
    "The reason is that, in NumPy, the operation `np.cos(x**2 + y**2) / (1 +\n",
    "x**2 + y**2)` generates several intermediate arrays.\n",
    "\n",
    "For example, a new array is created when `x**2` is calculated.\n",
    "\n",
    "The same is true when `y**2` is calculated, and then `x**2 + y**2` and so on.\n",
    "\n",
    "Numba avoids creating all these intermediate arrays by compiling one\n",
    "function that is specialized to the entire operation.\n",
    "\n",
    "But if this is true, then why isn’t the Numba code faster?\n",
    "\n",
    "The reason is that NumPy makes up for its disadvantages with implicit\n",
    "multithreading, as we’ve just discussed."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Multithreading a Numba Ufunc\n",
    "\n",
    "Can we get both of these advantages at once?\n",
    "\n",
    "In other words, can we pair\n",
    "\n",
    "- the efficiency of Numba’s highly specialized JIT compiled function and  \n",
    "- the speed gains from parallelization obtained by NumPy’s implicit\n",
    "  multithreading?  \n",
    "\n",
    "\n",
    "It turns out that we can, by adding some type information plus `target='parallel'`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "@vectorize('float64(float64, float64)', target='parallel')\n",
    "def f_vec(x, y):\n",
    "    return np.cos(x**2 + y**2) / (1 + x**2 + y**2)\n",
    "\n",
    "np.max(f_vec(x, y))  # Run once to compile"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "%timeit np.max(f_vec(x, y))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now our code runs significantly faster than the NumPy version."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Multithreaded Loops in Numba\n",
    "\n",
    "We just saw one approach to parallelization in Numba, using the `parallel`\n",
    "flag in `@vectorize`.\n",
    "\n",
    "This is neat but, it turns out, not well suited to many problems we consider.\n",
    "\n",
    "Fortunately, Numba provides another approach to multithreading that will work\n",
    "for us almost everywhere parallelization is possible.\n",
    "\n",
    "To illustrate, let’s look first at a simple, single-threaded (i.e., non-parallelized) piece of code.\n",
    "\n",
    "The code simulates updating the wealth $ w_t $ of a household via the rule\n",
    "\n",
    "$$\n",
    "w_{t+1} = R_{t+1} s w_t + y_{t+1}\n",
    "$$\n",
    "\n",
    "Here\n",
    "\n",
    "- $ R $ is the gross rate of return on assets  \n",
    "- $ s $ is the savings rate of the household and  \n",
    "- $ y $ is labor income.  \n",
    "\n",
    "\n",
    "We model both $ R $ and $ y $ as independent draws from a lognormal\n",
    "distribution.\n",
    "\n",
    "Here’s the code:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "from numpy.random import randn\n",
    "from numba import njit\n",
    "\n",
    "@njit\n",
    "def h(w, r=0.1, s=0.3, v1=0.1, v2=1.0):\n",
    "    \"\"\"\n",
    "    Updates household wealth.\n",
    "    \"\"\"\n",
    "\n",
    "    # Draw shocks\n",
    "    R = np.exp(v1 * randn()) * (1 + r)\n",
    "    y = np.exp(v2 * randn())\n",
    "\n",
    "    # Update wealth\n",
    "    w = R * s * w + y\n",
    "    return w"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let’s have a look at how wealth evolves under this rule."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "\n",
    "T = 100\n",
    "w = np.empty(T)\n",
    "w[0] = 5\n",
    "for t in range(T-1):\n",
    "    w[t+1] = h(w[t])\n",
    "\n",
    "ax.plot(w)\n",
    "ax.set_xlabel('$t$', fontsize=12)\n",
    "ax.set_ylabel('$w_{t}$', fontsize=12)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let’s suppose that we have a large population of households and we want to\n",
    "know what median wealth will be.\n",
    "\n",
    "This is not easy to solve with pencil and paper, so we will use simulation\n",
    "instead.\n",
    "\n",
    "In particular, we will simulate a large number of households and then\n",
    "calculate median wealth for this group.\n",
    "\n",
    "Suppose we are interested in the long-run average of this median over time.\n",
    "\n",
    "It turns out that, for the specification that we’ve chosen above, we can\n",
    "calculate this by taking a one-period snapshot of what has happened to median\n",
    "wealth of the group at the end of a long simulation.\n",
    "\n",
    "Moreover, provided the simulation period is long enough, initial conditions\n",
    "don’t matter.\n",
    "\n",
    "- This is due to something called ergodicity, which we will discuss [later on](https://python.quantecon.org/finite_markov.html#Ergodicity).  \n",
    "\n",
    "\n",
    "So, in summary, we are going to simulate 50,000 households by\n",
    "\n",
    "1. arbitrarily setting initial wealth to 1 and  \n",
    "1. simulating forward in time for 1,000 periods.  \n",
    "\n",
    "\n",
    "Then we’ll calculate median wealth at the end period.\n",
    "\n",
    "Here’s the code:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "@njit\n",
    "def compute_long_run_median(w0=1, T=1000, num_reps=50_000):\n",
    "\n",
    "    obs = np.empty(num_reps)\n",
    "    for i in range(num_reps):\n",
    "        w = w0\n",
    "        for t in range(T):\n",
    "            w = h(w)\n",
    "        obs[i] = w\n",
    "\n",
    "    return np.median(obs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let’s see how fast this runs:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "compute_long_run_median()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To speed this up, we’re going to parallelize it via multithreading.\n",
    "\n",
    "To do so, we add the `parallel=True` flag and change `range` to `prange`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "from numba import prange\n",
    "\n",
    "@njit(parallel=True)\n",
    "def compute_long_run_median_parallel(w0=1, T=1000, num_reps=50_000):\n",
    "\n",
    "    obs = np.empty(num_reps)\n",
    "    for i in prange(num_reps):\n",
    "        w = w0\n",
    "        for t in range(T):\n",
    "            w = h(w)\n",
    "        obs[i] = w\n",
    "\n",
    "    return np.median(obs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let’s look at the timing:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "compute_long_run_median_parallel()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The speed-up is significant."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### A Warning\n",
    "\n",
    "Parallelization works well in the outer loop of the last example because the individual tasks inside the loop are independent of each other.\n",
    "\n",
    "If this independence fails then parallelization is often problematic.\n",
    "\n",
    "For example, each step inside the inner loop depends on the last step, so\n",
    "independence fails, and this is why we use ordinary `range` instead of `prange`.\n",
    "\n",
    "When you see us using `prange` in later lectures, it is because the\n",
    "independence of tasks holds true.\n",
    "\n",
    "When you see us using ordinary `range` in a jitted function, it is either because the speed gain from parallelization is small or because independence fails."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercises"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 1\n",
    "\n",
    "In [an earlier exercise](https://python.quantecon.org/numba.html#speed-ex1), we used Numba to accelerate an\n",
    "effort to compute the constant $ \\pi $ by Monte Carlo.\n",
    "\n",
    "Now try adding parallelization and see if you get further speed gains.\n",
    "\n",
    "You should not expect huge gains here because, while there are many\n",
    "independent tasks (draw point and test if in circle), each one has low\n",
    "execution time.\n",
    "\n",
    "Generally speaking, parallelization is less effective when the individual\n",
    "tasks to be parallelized are very small relative to total execution time.\n",
    "\n",
    "This is due to overheads associated with spreading all of these small tasks across multiple CPUs.\n",
    "\n",
    "Nevertheless, with suitable hardware, it is possible to get nontrivial speed gains in this exercise.\n",
    "\n",
    "For the size of the Monte Carlo simulation, use something substantial, such as\n",
    "`n = 100_000_000`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Solutions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 1\n",
    "\n",
    "Here is one solution:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "from random import uniform\n",
    "\n",
    "@njit(parallel=True)\n",
    "def calculate_pi(n=1_000_000):\n",
    "    count = 0\n",
    "    for i in prange(n):\n",
    "        u, v = uniform(0, 1), uniform(0, 1)\n",
    "        d = np.sqrt((u - 0.5)**2 + (v - 0.5)**2)\n",
    "        if d < 0.5:\n",
    "            count += 1\n",
    "\n",
    "    area_estimate = count / n\n",
    "    return area_estimate * 4  # dividing by radius**2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let’s see how fast it runs:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "%time calculate_pi()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "%time calculate_pi()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "By switching parallelization on and off (selecting `True` or\n",
    "`False` in the `@njit` annotation), we can test the speed gain that\n",
    "multithreading provides on top of JIT compilation.\n",
    "\n",
    "On our workstation, we find that parallelization increases execution speed by\n",
    "a factor of 2 or 3.\n",
    "\n",
    "(If you are executing locally, you will get different numbers, depending mainly\n",
    "on the number of CPUs on your machine.)"
   ]
  }
 ],
 "metadata": {
  "date": 1584334751.6544552,
  "filename": "parallelization.rst",
  "kernelspec": {
   "display_name": "Python",
   "language": "python3",
   "name": "python3"
  },
  "title": "Parallelization"
 },
 "nbformat": 4,
 "nbformat_minor": 2
}