numba.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "<a id='speed'></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": [
    "# Numba"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Contents\n",
    "\n",
    "- [Numba](#Numba)  \n",
    "  - [Overview](#Overview)  \n",
    "  - [Where are the Bottlenecks?](#Where-are-the-Bottlenecks?)  \n",
    "  - [Vectorization](#Vectorization)  \n",
    "  - [Numba](#Numba)  "
   ]
  },
  {
   "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",
    "In our lecture on [NumPy](https://lectures.quantecon.org/numpy.html), we learned one method to improve speed and efficiency in numerical work.\n",
    "\n",
    "That method, called *vectorization*, involved sending array processing operations in batch to efficient low-level code.\n",
    "\n",
    "This clever idea dates back to Matlab, which uses it extensively.\n",
    "\n",
    "Unfortunately, vectorization is limited and has several weaknesses.\n",
    "\n",
    "One weakness is that it is highly memory-intensive.\n",
    "\n",
    "Another problem is that only some algorithms can be vectorized.\n",
    "\n",
    "In the last few years, a new Python library called [Numba](http://numba.pydata.org/) has appeared that\n",
    "solves many of these problems.\n",
    "\n",
    "It does so through something called **just in time (JIT) compilation**.\n",
    "\n",
    "JIT compilation is effective in many numerical settings and can generate extremely fast, efficient code.\n",
    "\n",
    "It can also do other tricks such as facilitate multithreading (a form of parallelization well suited to numerical work)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### The Need for Speed\n",
    "\n",
    "To understand what Numba does and why, we need some background knowledge.\n",
    "\n",
    "Let’s start by thinking about higher-level languages, such as Python.\n",
    "\n",
    "These languages are optimized for humans.\n",
    "\n",
    "This means that the programmer can leave many details to the runtime environment\n",
    "\n",
    "- specifying variable types  \n",
    "- memory allocation/deallocation, etc.  \n",
    "\n",
    "\n",
    "The upside is that, compared to low-level languages, Python is typically faster to write, less error-prone and  easier to debug.\n",
    "\n",
    "The downside is that Python is harder to optimize — that is, turn into fast machine code — than languages like C or Fortran.\n",
    "\n",
    "Indeed, the standard implementation of Python (called CPython) cannot match the speed of compiled languages such as C or Fortran.\n",
    "\n",
    "Does that mean that we should just switch to C or Fortran for everything?\n",
    "\n",
    "The answer is no, no and one hundred times no.\n",
    "\n",
    "High productivity languages should be chosen over high-speed languages for the great majority of scientific computing tasks.\n",
    "\n",
    "This is because\n",
    "\n",
    "1. Of any given program, relatively few lines are ever going to be time-critical.  \n",
    "1. For those lines of code that *are* time-critical, we can achieve C-like speed using a combination of NumPy and Numba.  \n",
    "\n",
    "\n",
    "This lecture provides a guide."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Where are the Bottlenecks?\n",
    "\n",
    "Let’s start by trying to understand why high-level languages like Python are slower than compiled code."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Dynamic Typing\n",
    "\n",
    "\n",
    "<a id='index-0'></a>\n",
    "Consider this Python operation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "a, b = 10, 10\n",
    "a + b"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Even for this simple operation, the Python interpreter has a fair bit of work to do.\n",
    "\n",
    "For example, in the statement `a + b`, the interpreter has to know which\n",
    "operation to invoke.\n",
    "\n",
    "If `a` and `b` are strings, then `a + b` requires string concatenation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "a, b = 'foo', 'bar'\n",
    "a + b"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If `a` and `b` are lists, then `a + b` requires list concatenation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "a, b = ['foo'], ['bar']\n",
    "a + b"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "(We say that the operator `+` is *overloaded* — its action depends on the\n",
    "type of the objects on which it acts)\n",
    "\n",
    "As a result, Python must check the type of the objects and then call the correct operation.\n",
    "\n",
    "This involves substantial overheads."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Static Types\n",
    "\n",
    "\n",
    "<a id='index-1'></a>\n",
    "Compiled languages avoid these overheads with explicit, static types.\n",
    "\n",
    "For example, consider the following C code, which sums the integers from 1 to 10"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hide-output": false
   },
   "source": [
    "```c\n",
    "#include <stdio.h>\n",
    "\n",
    "int main(void) {\n",
    "    int i;\n",
    "    int sum = 0;\n",
    "    for (i = 1; i <= 10; i++) {\n",
    "        sum = sum + i;\n",
    "    }\n",
    "    printf(\"sum = %d\\n\", sum);\n",
    "    return 0;\n",
    "}\n",
    "```\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The variables `i` and `sum` are explicitly declared to be integers.\n",
    "\n",
    "Hence, the meaning of addition here is completely unambiguous."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Data Access\n",
    "\n",
    "Another drag on speed for high-level languages is data access.\n",
    "\n",
    "To illustrate, let’s consider the problem of summing some data — say, a collection of integers."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Summing with Compiled Code\n",
    "\n",
    "In C or Fortran, these integers would typically be stored in an array, which\n",
    "is a simple data structure for storing homogeneous data.\n",
    "\n",
    "Such an array is stored in a single contiguous block of memory\n",
    "\n",
    "- In modern computers, memory addresses are allocated to each byte (one byte = 8 bits).  \n",
    "- For example, a 64 bit integer is stored in 8 bytes of memory.  \n",
    "- An array of $ n $ such integers occupies $ 8n $ **consecutive** memory slots.  \n",
    "\n",
    "\n",
    "Moreover, the compiler is made aware of the data type by the programmer.\n",
    "\n",
    "- In this case 64 bit integers  \n",
    "\n",
    "\n",
    "Hence, each successive data point can be accessed by shifting forward in memory\n",
    "space by a known and fixed amount.\n",
    "\n",
    "- In this case 8 bytes  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Summing in Pure Python\n",
    "\n",
    "Python tries to replicate these ideas to some degree.\n",
    "\n",
    "For example, in the standard Python implementation (CPython), list elements are placed in memory locations that are in a sense contiguous.\n",
    "\n",
    "However, these list elements are more like pointers to data rather than actual data.\n",
    "\n",
    "Hence, there is still overhead involved in accessing the data values themselves.\n",
    "\n",
    "This is a considerable drag on speed.\n",
    "\n",
    "In fact, it’s generally true that memory traffic is a major culprit when it comes to slow execution.\n",
    "\n",
    "Let’s look at some ways around these problems."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Vectorization\n",
    "\n",
    "\n",
    "<a id='index-3'></a>\n",
    "Vectorization is about sending batches of related operations to native machine code.\n",
    "\n",
    "- The machine code itself is typically compiled from carefully optimized C or Fortran.  \n",
    "\n",
    "\n",
    "This can greatly accelerate many (but not all) numerical computations."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Operations on Arrays\n",
    "\n",
    "\n",
    "<a id='index-4'></a>\n",
    "First, let’s run some imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "import random\n",
    "import numpy as np\n",
    "import quantecon as qe"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let’s try this non-vectorized code"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "qe.util.tic()   # Start timing\n",
    "n = 1_000_000\n",
    "sum = 0\n",
    "for i in range(n):\n",
    "    x = random.uniform(0, 1)\n",
    "    sum += x**2\n",
    "qe.util.toc()   # End timing"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now compare this vectorized code"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "qe.util.tic()\n",
    "n = 1_000_000\n",
    "x = np.random.uniform(0, 1, n)\n",
    "np.sum(x**2)\n",
    "qe.util.toc()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The second code block — which achieves the same thing as the first — runs\n",
    "much faster.\n",
    "\n",
    "The reason is that in the second implementation we have broken the loop down into three basic operations\n",
    "\n",
    "1. draw `n` uniforms  \n",
    "1. square them  \n",
    "1. sum them  \n",
    "\n",
    "\n",
    "These are sent as batch operators to optimized machine code.\n",
    "\n",
    "Apart from minor overheads associated with sending data back and forth, the result is C or Fortran-like speed.\n",
    "\n",
    "When we run batch operations on arrays like this, we say that the code is *vectorized*.\n",
    "\n",
    "Vectorized code is typically fast and efficient.\n",
    "\n",
    "It is also surprisingly flexible, in the sense that many operations can be vectorized.\n",
    "\n",
    "The next section illustrates this point.\n",
    "\n",
    "\n",
    "<a id='ufuncs'></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Universal Functions\n",
    "\n",
    "\n",
    "<a id='index-5'></a>\n",
    "Many functions provided by NumPy are so-called *universal functions* — also called [ufuncs](https://docs.scipy.org/doc/numpy/reference/ufuncs.html).\n",
    "\n",
    "This means that they\n",
    "\n",
    "- map scalars into scalars, as expected  \n",
    "- map arrays into arrays, acting element-wise  \n",
    "\n",
    "\n",
    "For example, `np.cos` is a ufunc:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "np.cos(1.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "np.cos(np.linspace(0, 1, 3))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "By exploiting ufuncs, many operations can be vectorized.\n",
    "\n",
    "For example, consider the problem of maximizing a function $ f $ of two\n",
    "variables $ (x,y) $ over the square $ [-a, a] \\times [-a, a] $.\n",
    "\n",
    "For $ f $ and $ a $ let’s choose\n",
    "\n",
    "$$\n",
    "f(x,y) = \\frac{\\cos(x^2 + y^2)}{1 + x^2 + y^2}\n",
    "\\quad \\text{and} \\quad\n",
    "a = 3\n",
    "$$\n",
    "\n",
    "Here’s a plot of $ f $"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "from mpl_toolkits.mplot3d.axes3d import Axes3D\n",
    "from matplotlib import cm\n",
    "\n",
    "def f(x, y):\n",
    "    return np.cos(x**2 + y**2) / (1 + x**2 + y**2)\n",
    "\n",
    "xgrid = np.linspace(-3, 3, 50)\n",
    "ygrid = xgrid\n",
    "x, y = np.meshgrid(xgrid, ygrid)\n",
    "\n",
    "fig = plt.figure(figsize=(8, 6))\n",
    "ax = fig.add_subplot(111, projection='3d')\n",
    "ax.plot_surface(x,\n",
    "                y,\n",
    "                f(x, y),\n",
    "                rstride=2, cstride=2,\n",
    "                cmap=cm.jet,\n",
    "                alpha=0.7,\n",
    "                linewidth=0.25)\n",
    "ax.set_zlim(-0.5, 1.0)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To maximize it, we’re going to use a naive grid search:\n",
    "\n",
    "1. Evaluate $ f $ for all $ (x,y) $ in a grid on the square.  \n",
    "1. Return the maximum of observed values.  \n",
    "\n",
    "\n",
    "Here’s a non-vectorized version that uses Python loops"
   ]
  },
  {
   "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, 1000)\n",
    "m = -np.inf\n",
    "\n",
    "qe.tic()\n",
    "for x in grid:\n",
    "    for y in grid:\n",
    "        z = f(x, y)\n",
    "        if z > m:\n",
    "            m = z\n",
    "\n",
    "qe.toc()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And here’s a vectorized version"
   ]
  },
  {
   "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, 1000)\n",
    "x, y = np.meshgrid(grid, grid)\n",
    "\n",
    "qe.tic()\n",
    "np.max(f(x, y))\n",
    "qe.toc()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the vectorized version, all the looping takes place in compiled code.\n",
    "\n",
    "As you can see, the second version is **much** faster.\n",
    "\n",
    "(We’ll make it even faster again below when we discuss Numba)\n",
    "\n",
    "\n",
    "<a id='numba-p-c-vectorization'></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Pros and Cons of Vectorization\n",
    "\n",
    "At its best, vectorization yields fast, simple code.\n",
    "\n",
    "However, it’s not without disadvantages.\n",
    "\n",
    "One issue is that it can be highly memory-intensive.\n",
    "\n",
    "For example, the vectorized maximization routine above is far more memory\n",
    "intensive than the non-vectorized version that preceded it.\n",
    "\n",
    "Another issue is that not all algorithms can be vectorized.\n",
    "\n",
    "In these kinds of settings, we need to go back to loops.\n",
    "\n",
    "Fortunately, there are nice ways to speed up Python loops.\n",
    "\n",
    "\n",
    "<a id='numba-link'></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Numba\n",
    "\n",
    "\n",
    "<a id='index-7'></a>\n",
    "One exciting development in this direction is [Numba](http://numba.pydata.org/).\n",
    "\n",
    "Numba aims to automatically compile functions to native machine code instructions on the fly.\n",
    "\n",
    "The process isn’t flawless, since Numba needs to infer type information on all variables to generate pure machine instructions.\n",
    "\n",
    "Such inference isn’t possible in every setting.\n",
    "\n",
    "But for simple routines, Numba infers types very well.\n",
    "\n",
    "Moreover, the “hot loops” at the heart of our code that we need to speed up are often such simple routines."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Prerequisites\n",
    "\n",
    "If you [followed our set up instructions](https://lectures.quantecon.org/getting_started.html), then Numba should be installed.\n",
    "\n",
    "Make sure you have the latest version of Anaconda by running `conda update anaconda` from a terminal (Mac, Linux) / Anaconda command prompt (Windows).\n",
    "\n",
    "\n",
    "<a id='quad-map-eg'></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### An Example\n",
    "\n",
    "Let’s consider some problems that are difficult to vectorize.\n",
    "\n",
    "One is generating the trajectory of a difference equation given an initial\n",
    "condition.\n",
    "\n",
    "Let’s take the difference equation to be the quadratic map\n",
    "\n",
    "$$\n",
    "x_{t+1} = 4 x_t (1 - x_t)\n",
    "$$\n",
    "\n",
    "Here’s the plot of a typical trajectory, starting from $ x_0 = 0.1 $, with $ t $ on the x-axis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def qm(x0, n):\n",
    "    x = np.empty(n+1)\n",
    "    x[0] = x0\n",
    "    for t in range(n):\n",
    "        x[t+1] = 4 * x[t] * (1 - x[t])\n",
    "    return x\n",
    "\n",
    "x = qm(0.1, 250)\n",
    "fig, ax = plt.subplots(figsize=(10, 6))\n",
    "ax.plot(x, 'b-', lw=2, alpha=0.8)\n",
    "ax.set_xlabel('time', fontsize=16)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To speed this up using Numba is trivial using Numba’s `jit` function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "from numba import jit\n",
    "\n",
    "qm_numba = jit(qm)  # qm_numba is now a 'compiled' version of qm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let’s time and compare identical function calls across these two versions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "qe.util.tic()\n",
    "qm(0.1, int(10**5))\n",
    "time1 = qe.util.toc()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "qe.util.tic()\n",
    "qm_numba(0.1, int(10**5))\n",
    "time2 = qe.util.toc()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The first execution is relatively slow because of JIT compilation (see below).\n",
    "\n",
    "Next time and all subsequent times it runs much faster:\n",
    "\n",
    "\n",
    "<a id='qm-numba-result'></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "qe.util.tic()\n",
    "qm_numba(0.1, int(10**5))\n",
    "time2 = qe.util.toc()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "time1 / time2  # Calculate speed gain"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "That’s a speed increase of two orders of magnitude!\n",
    "\n",
    "Your mileage will of course vary depending on hardware and so on.\n",
    "\n",
    "Nonetheless, two orders of magnitude is huge relative to how simple and clear the implementation is."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Decorator Notation\n",
    "\n",
    "If you don’t need a separate name for the “numbafied” version of `qm`,\n",
    "you can just put `@jit` before the function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "@jit\n",
    "def qm(x0, n):\n",
    "    x = np.empty(n+1)\n",
    "    x[0] = x0\n",
    "    for t in range(n):\n",
    "        x[t+1] = 4 * x[t] * (1 - x[t])\n",
    "    return x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is equivalent to `qm = jit(qm)`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### How and When it Works\n",
    "\n",
    "Numba attempts to generate fast machine code using the infrastructure provided by the [LLVM Project](http://llvm.org/).\n",
    "\n",
    "It does this by inferring type information on the fly.\n",
    "\n",
    "As you can imagine, this is easier for simple Python objects (simple scalar data types, such as floats, integers, etc.).\n",
    "\n",
    "Numba also plays well with NumPy arrays, which it treats as typed memory regions.\n",
    "\n",
    "In an ideal setting, Numba can infer all necessary type information.\n",
    "\n",
    "This allows it to generate native machine code, without having to call the Python runtime environment.\n",
    "\n",
    "In such a setting, Numba will be on par with machine code from low-level languages.\n",
    "\n",
    "When Numba cannot infer all type information, some Python objects are given generic `object` status, and some code is generated using the Python runtime.\n",
    "\n",
    "In this second setting, Numba typically provides only minor speed gains — or none at all.\n",
    "\n",
    "Hence, it’s prudent when using Numba to focus on speeding up small, time-critical snippets of code.\n",
    "\n",
    "This will give you much better performance than blanketing your Python programs with `@jit` statements."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### A Gotcha: Global Variables\n",
    "\n",
    "Consider the following example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "a = 1\n",
    "\n",
    "@jit\n",
    "def add_x(x):\n",
    "    return a + x\n",
    "\n",
    "print(add_x(10))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "a = 2\n",
    "\n",
    "print(add_x(10))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notice that changing the global had no effect on the value returned by the\n",
    "function.\n",
    "\n",
    "When Numba compiles machine code for functions, it treats global variables as constants to ensure type stability."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Numba for Vectorization\n",
    "\n",
    "Numba can also be used to create custom [ufuncs](#ufuncs) with the [@vectorize](http://numba.pydata.org/numba-doc/dev/user/vectorize.html) decorator.\n",
    "\n",
    "To illustrate the advantage of using Numba to vectorize a function, we\n",
    "return to a maximization problem [discussed above](#ufuncs)"
   ]
  },
  {
   "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",
    "grid = np.linspace(-3, 3, 1000)\n",
    "x, y = np.meshgrid(grid, grid)\n",
    "\n",
    "np.max(f_vec(x, y))  # Run once to compile\n",
    "\n",
    "qe.tic()\n",
    "np.max(f_vec(x, y))\n",
    "qe.toc()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is faster than our vectorized version using NumPy’s ufuncs.\n",
    "\n",
    "Why should that be?  After all, anything vectorized with NumPy will be running in fast C or Fortran code.\n",
    "\n",
    "The reason is that it’s much less memory-intensive.\n",
    "\n",
    "For example, when NumPy computes `np.cos(x**2 + y**2)` it first creates the\n",
    "intermediate arrays `x**2` and `y**2`, then it creates the array `np.cos(x**2 + y**2)`.\n",
    "\n",
    "In our `@vectorize` version using Numba, the entire operator is reduced to a\n",
    "single vectorized process and none of these intermediate arrays are created.\n",
    "\n",
    "We can gain further speed improvements using Numba’s automatic parallelization\n",
    "feature by specifying `target='parallel'`.\n",
    "\n",
    "In this case, we need to specify the types of our inputs and outputs"
   ]
  },
  {
   "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\n",
    "\n",
    "qe.tic()\n",
    "np.max(f_vec(x, y))\n",
    "qe.toc()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is a striking speed up with very little effort."
   ]
  }
 ],
 "metadata": {
  "filename": "numba.rst",
  "kernelspec": {
   "display_name": "Python",
   "language": "python3",
   "name": "python3"
  },
  "title": "Numba"
 },
 "nbformat": 4,
 "nbformat_minor": 2
}