numba_intro.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "view-in-github"
   },
   "source": [
    "<a href=\"https://colab.research.google.com/github/evaneschneider/parallel-programming/blob/master/numba_intro.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "cBAVRqT-D8JR"
   },
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import matplotlib.pylab as plt\n",
    "import numpy as np\n",
    "import math\n",
    "from numba import jit, njit, vectorize"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "72QKUcHITZXs"
   },
   "source": [
    "##What is Numba?\n",
    "Numba is a **just-in-time**, **type-specializing**, **function compiler** for accelerating **numerically-focused** Python. That's a long list, so let's break down those terms:\n",
    "\n",
    "+ **function compiler**: Numba compiles Python functions, not entire applications, and not parts of functions. Numba does not replace your Python interpreter, but is just another Python module that can turn a function into a (usually) faster function.\n",
    "+ **type-specializing**: Numba speeds up your function by generating a specialized implementation for the specific data types you are using. Python functions are designed to operate on generic data types, which makes them very flexible, but also very slow. In practice, you only will call a function with a small number of argument types, so Numba will generate a fast implementation for each set of types.\n",
    "+ **just-in-time**: Numba translates functions when they are first called. This ensures the compiler knows what argument types you will be using. This also allows Numba to be used interactively in a Jupyter notebook just as easily as a traditional application.\n",
    "+ **numerically-focused**: Currently, Numba is focused on numerical data types, like int, float, and complex. There is very limited string processing support, and many string use cases are not going to work well on the GPU. To get best results with Numba, you will likely be using NumPy arrays."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "SgFcESP_EaNv"
   },
   "source": [
    "### Problem 1 - A First Numba Function\n",
    "\n",
    "**1a)** To start our exploration of Numba's features, let's write a python function to add two numbers. We'll creatively name it `add`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "yzKiUuSjEeDy"
   },
   "outputs": [],
   "source": [
    "def add(x, y):\n",
    "# add code here"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "2Fnz9vr0FG7k"
   },
   "source": [
    "Now, test the function, first with two scalar integers:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "o13KJoxJFDyD"
   },
   "outputs": [],
   "source": [
    "# add code here"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "7HdCKleiFd6r"
   },
   "source": [
    "**1b)** With Numpy, we can use our function to add not just scalars, but vectors as well. Using your favorite array creation routine, create two integer arrays with ten elements each, called `a` and `b`, and use your `add` function to add them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "EWwdk1dYE5pN"
   },
   "outputs": [],
   "source": [
    "# add code here"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Jwqy1reHM0nA"
   },
   "source": [
    "Okay, so our function can add things. Now, let's use Numba's `jit` function to create a Numba version of our addition function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "SfnuyYZ1MjxR"
   },
   "outputs": [],
   "source": [
    "numba_add = jit(add)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "FW3GvGwKkawy"
   },
   "source": [
    "More commonly, you will use `jit` as a decorator, by adding `@jit` to the line above your function definition, but the above version shows you that at heart, `@jit` is just a python function that takes other functions as its argument!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "X726qBEztHg5"
   },
   "source": [
    "**1c)**  By default, a Numba function saves the original python version of the function in the variable `py_func`. Check that the original python version gives you the same answer as the Numba version."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "XS0-gUOWN12L"
   },
   "outputs": [],
   "source": [
    "# add code here"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "QQx121wWGD_Y"
   },
   "source": [
    "#### Profiling\n",
    "A central feature of parallel programming, Numba, and writing efficient code more generally is **profiling**, or understanding how long various pieces of your program take to run. Profiling tools are becoming ever more sophisticated, but for today we're going to stick with the tried-and-true method of timing things. An easy way to do this in python is using the `%timeit` magic function. Let's try it out on our addition function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "XC1o2Ki8Fx2U"
   },
   "outputs": [],
   "source": [
    "%timeit add(1,2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "krQ0FwkjGw-_"
   },
   "source": [
    "What's going on here? `%timeit` is running our function many times, and then reporting the average time it takes to run. This is generally a better approach than timing a single function execution, because it accounts for random events that may cause any given run to perform poorly."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "0xXjAFh-MXiv"
   },
   "source": [
    "**1d)** Compare the time it takes to run your function with scalar vs array arguments, then your function vs python's add function (the standard ''+'' operator)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "3AgB0xKOHjcS"
   },
   "outputs": [],
   "source": [
    "# add code here\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "YsZ_clWelcVJ"
   },
   "source": [
    "So, scalars are faster than arrays (makes sense), and python's addition function is better than ours (seems reasonable). Now, let's see how fast our pre-compiled Numba addition function is."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "pqNzG9mhN-lp"
   },
   "outputs": [],
   "source": [
    "# add code here"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "pgDvbPLml3Dm"
   },
   "source": [
    "Hold on - our new pre-compiled function is running even slower than the original python version! What's going on here?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "szNrfvJxl_Vj"
   },
   "source": [
    "### Problem 2 - A Better Numba Function\n",
    "\n",
    "(This problem borrowed from [seibert's 2018 gtc numba tutorial](https://github.com/ContinuumIO/gtc2018-numba).)\n",
    "\n",
    "As we saw in the first example, Numba isn't going to speed up everything. Generally, Numba will help you most in circumstances where python's line-by-line interperability and lack of type casting is slowing it down. We can use a slightly more complicated function to demonstrate this. The following is a function to calculate the hypotenuse of two numbers, that has been carefully designed to compensate for the computer's finite precision representation of numbers (check out https://en.wikipedia.org/wiki/Hypot for more info)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "gFmqy_gjvWJX"
   },
   "source": [
    "**2a)** Use the `@jit` decorator to generate a Numba version of this function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "VNDuzJvGlkfg"
   },
   "outputs": [],
   "source": [
    "# add code here\n",
    "def hypotenuse(x, y):\n",
    "  x = abs(x);\n",
    "  y = abs(y);\n",
    "  t = min(x, y);\n",
    "  x = max(x, y);\n",
    "  t = t / x;\n",
    "  return x * math.sqrt(1+t*t)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "52cFaUWboRZO"
   },
   "source": [
    "**2b)** Use the `%timeit` function to determine whether the Numba version of the hyptonenuse function is better than the original Python implementation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "HS7QNmoUmcYG"
   },
   "outputs": [],
   "source": [
    "# add code here"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "0ZQmD-3Ho5oJ"
   },
   "source": [
    "**2c)** Numba functions can call other functions, provided they are also Numba functions. Below is a function that loops through two numpy arrays and puts their sum into an output array. Modify the following function to calculate the hypotenuse instead."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "0gO3Pxqbo4Nz"
   },
   "outputs": [],
   "source": [
    "@njit # this is an alias for @jit(nopython=True)\n",
    "def ex_func(x, y, out):\n",
    "  for i in range(x.shape[0]):\n",
    "    out[i] = x[i] + y[i] # change this line"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "m62ZUF84opUb"
   },
   "outputs": [],
   "source": [
    "in1 = np.arange(10, dtype=np.float64)\n",
    "in2 = 2 * in1 + 1\n",
    "out = np.empty_like(in1)\n",
    "\n",
    "print('in1:', in1)\n",
    "print('in2:', in2)\n",
    "\n",
    "ex_func(in1, in2, out)\n",
    "\n",
    "print('out:', out)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "q1GhkW0LBuV_"
   },
   "outputs": [],
   "source": [
    "# This test will fail until you fix the ex1 function\n",
    "np.testing.assert_almost_equal(out, np.hypot(in1, in2))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "zhO1-4Das9cm"
   },
   "source": [
    "### Problem 3 - Fun with Fractals\n",
    "Now that we've got the basics of the Numba `jit` decorator down, let's have a little fun. A classic example problem in parallel programming is the calculation of a fractal, because a large fraction of the work can be done in parallel. Below is some code that calculates whether a number is a member of the Julia set, and then computes the set on a discrete domain to calculate a fractal."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "tH6JqqMLw5Hq"
   },
   "source": [
    "**3a)** Modify the code below to use Numba and test how much faster it is than the original python implementation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "XormIu6WqCOa"
   },
   "outputs": [],
   "source": [
    "# add code here\n",
    "def julia(x, y, max_iters):\n",
    "    \"\"\"\n",
    "    Given the real and imaginary parts of a complex number,\n",
    "    determine if it is a candidate for membership in the Julia\n",
    "    set given a fixed number of iterations.\n",
    "    \"\"\"\n",
    "    i = 0\n",
    "    c = complex(-0.8, 0.156)\n",
    "    a = complex(x,y)\n",
    "    for i in range(max_iters):\n",
    "        a = a*a + c\n",
    "        if (a.real*a.real + a.imag*a.imag) > 1000:\n",
    "            return 0\n",
    "    return 255"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "LCr_6rlEt6KU"
   },
   "outputs": [],
   "source": [
    "# add code here\n",
    "def create_fractal(min_x, max_x, min_y, max_y, image, iters):\n",
    "    height = image.shape[0]\n",
    "    width = image.shape[1]\n",
    "\n",
    "    pixel_size_x = (max_x - min_x) / width\n",
    "    pixel_size_y = (max_y - min_y) / height\n",
    "    for x in range(width):\n",
    "        real = min_x + x * pixel_size_x\n",
    "        for y in range(height):\n",
    "            imag = min_y + y * pixel_size_y\n",
    "            color = julia(real, imag, iters)\n",
    "            image[y, x] = color\n",
    "\n",
    "    return image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "fZ-U8-tLuDjl"
   },
   "outputs": [],
   "source": [
    "image = np.zeros((500, 750), dtype=np.uint8)\n",
    "%timeit create_fractal(-2.0, 2.0, -1.0, 1.0, image, 200)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "FI8RpWQ4uYTL"
   },
   "source": [
    "Want to see what you made? Run the following code to plot the image. Feel free to pick your favorite matplotlib color map :)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "5QY46v6Oua79"
   },
   "outputs": [],
   "source": [
    "plt.imshow(image)\n",
    "plt.viridis()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "KvUWJRTJG7r7"
   },
   "source": [
    "**3b)** There is more than one type of fractal in the world, however! Below is a function that determines membership in the Mandelbrot set. Modify the function below using Numba, then modify the code above to produce a new pretty picture."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "6Ow0_mRBHIxC"
   },
   "outputs": [],
   "source": [
    "# add code here\n",
    "def mandel(x, y, max_iters):\n",
    "    \"\"\"\n",
    "    Given the real and imaginary parts of a complex number,\n",
    "    determine if it is a candidate for membership in the Mandelbrot\n",
    "    set given a fixed number of iterations.\n",
    "    \"\"\"\n",
    "    i = 0\n",
    "    c = complex(x,y)\n",
    "    z = 0.0j\n",
    "    for i in range(max_iters):\n",
    "        z = z*z + c\n",
    "        if (z.real*z.real + z.imag*z.imag) >= 4:\n",
    "            return i\n",
    "    return 255"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "hHu6VBXYHXxv"
   },
   "outputs": [],
   "source": [
    "# add code here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "dlW4gww-HftJ"
   },
   "outputs": [],
   "source": [
    "image = np.zeros((500, 750), dtype=np.uint8)\n",
    "create_fractal(-2.0, 2.0, -1.0, 1.0, image, 200)\n",
    "plt.imshow(image)\n",
    "plt.viridis()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "GvUi6dVvVeY5"
   },
   "source": [
    "### Problem 4 - Typing\n",
    "\n",
    "Much of the power of Numba comes from its ability to compile a specific version of a python function based on the **data types** of its arguments. The data type describes what kind of variables the function uses, and in Numpy and Numba, pre-defined type names are based on the kind and size of the number in memory. You can find the type of a variable (or array) using numpy's `dtype` object. For example, let's see what type our `a` array is."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "L6LBCgxiuu4Q"
   },
   "outputs": [],
   "source": [
    "a.dtype"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "aF8y1QTW8ZK4"
   },
   "source": [
    "This tells us the array contains integers, and each integer has been assigned 64bits in memory (or equivalently, 8 bytes). Most python functions are defined to work on arbitrary types, so that if you use the + operator, for example, you can add integers, floats, complex numbers, or even strings! However, this flexibility comes at a cost, performance-wise. Numba, on the other hand, compiles each function based on the types of its arguments, and infers the type of the result. You can see this if you run the `inspect_types` function on a numba function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "NJoPUPFPZb3w"
   },
   "outputs": [],
   "source": [
    "numba_add.inspect_types()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "D0mmp24rrKUp"
   },
   "source": [
    "**4a)** Numba has inferred the types for this function based on how we've used it. Try out your `numba_add` function with two floating point numbers, then re-inspect the types of the Numba function. Are they the same?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Ce0oK34-rcBW"
   },
   "outputs": [],
   "source": [
    "#Add code here"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Fnuhv8VGCWFx"
   },
   "source": [
    "So far we have been using what Numba refers to as \"lazy\" (or \"call-time\") decoration. Basically, we've been letting Numba do the work of figuring out how we're using the function and inferring the types for us. Alternatively, if we know how we are going to use a given function, we can use \"eager\" (or \"compile-time\") decoration. To do this, we make use of the `vectorize` decorator. For example, if we want to make an integer-only version of our addition function, we could write:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "j5BI_s1NCVN_"
   },
   "outputs": [],
   "source": [
    "@vectorize(['int64(int64, int64)'], target='cpu')\n",
    "def add_ufunc(x, y):\n",
    "  return x + y"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Dhm4Kf5HDpJY"
   },
   "source": [
    "#### Numpy Universal Functions\n",
    "You'll notice a couple of new things here. In the first set of brackets, we have specified both the argument types of the function (those are inside the parentheses), as well as the return type of the function. This is just making explicit what Numba was previously inferring on our behalf. In second set of brackets you'll see that we have specified a 'target' architechture for the function. The default is `cpu`, which means that Numba is optimizing the function to your specific machine. Other options include `parallel`, which allows you to take advantage of multicore processors, and `cuda`, which we'll be discussing more tomorrow. You'll also notice that we called this a 'ufunc', which is short for [Universal Function](https://docs.scipy.org/doc/numpy/reference/ufuncs.html). In brief, universal functions are numpy functions that operate on ndarrays in element by element fashion. So if we pass two vectors to our `add_func`, they will be added together and a return vector of the same shape will be returned."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "rxiVxBslV-0z"
   },
   "outputs": [],
   "source": [
    "%timeit add_ufunc(a,b)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "oUbAJTzlFTqA"
   },
   "source": [
    "**4b)** Try your ufunc out with a new target, 'parallel'. How does the speed compare? What if the array size is much larger?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "X4A_kvi1GOkm"
   },
   "outputs": [],
   "source": [
    "# add code here"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "KREEf5ZNMKBX"
   },
   "source": [
    "### Problem 5 - Direct Summation\n",
    "(This problem borrowed in its entirety from [gforsyth's scipy17 numba tutorial](https://github.com/gforsyth/numba_tutorial_scipy2017).)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "xMPg_TRvRE8Q"
   },
   "source": [
    "Many physical problems require the evaluation of all pairwise interactions of a large number of particles, so-called N-body problems. These problems arise in molecular dynamics, astrodynamics and electromagnetics among others. \n",
    "\n",
    "Their pairwise interactions can be expressed as:\n",
    "\n",
    "\\begin{equation}\n",
    "f_i = \\sum_{j=1}^n{P \\left(\\boldsymbol{x}_i, \\boldsymbol{x}_j \\right)w_j} \\ \\ \\ \\text{for } i=1,2,...,n \n",
    "\\end{equation}\n",
    "\n",
    "*  where subscripts $i$,  $j$ respectively denote *target* and *source*\n",
    "*  $f_i$ can be a *potential* (or *force*) at target point $i$\n",
    "*  $w_j$ is the *source weight* \n",
    "*  $\\boldsymbol{x}_i, \\boldsymbol{x}_j$ are the *spatial positions* of particles \n",
    "*  $P \\left(\\boldsymbol{x}_i, \\boldsymbol{x}_j \\right)$ is the *interaction kernel*. \n",
    "\n",
    "In order to evalute the potential $f_i$ at a target point $i$, we have to loop over each source particle $j$. Since there are $n$ target points $i$, this 'brute-force' approach costs $\\mathcal{O} \\left(n^2 \\right)$ operations. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "T8Fz6v8cRZe2"
   },
   "source": [
    "One possible approach in this kind of problem is to define a few classes, say `Point` and `Particle` and then loop over the objects and perform the necessary point-to-point calculations.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Kgk_4HeaMNI4"
   },
   "outputs": [],
   "source": [
    "class Point():\n",
    "    \"\"\"    \n",
    "    Arguments:\n",
    "        domain: the domain of random generated coordinates x,y,z, \n",
    "                default=1.0\n",
    "    \n",
    "    Attributes:\n",
    "        x, y, z: coordinates of the point\n",
    "    \"\"\"\n",
    "    def __init__(self, domain=1.0):\n",
    "        self.x = domain * np.random.random()\n",
    "        self.y = domain * np.random.random()\n",
    "        self.z = domain * np.random.random()\n",
    "            \n",
    "    def distance(self, other):\n",
    "        return ((self.x - other.x)**2 + \n",
    "                (self.y - other.y)**2 + \n",
    "                (self.z - other.z)**2)**.5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "XgKlyjf1Rcyl"
   },
   "outputs": [],
   "source": [
    "class Particle(Point):\n",
    "    \"\"\"    \n",
    "    Attributes:\n",
    "        m: mass of the particle\n",
    "        phi: the potential of the particle\n",
    "    \"\"\"\n",
    "    \n",
    "    def __init__(self, domain=1.0, m=1.0):\n",
    "        Point.__init__(self, domain)\n",
    "        self.m = m\n",
    "        self.phi = 0."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "54LlwJUARoQx"
   },
   "source": [
    "Next, we define a function to calculate the particle interaction via direct summation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "eaotXcZzRsI6"
   },
   "outputs": [],
   "source": [
    "def direct_sum(particles):\n",
    "    \"\"\"\n",
    "    Calculate the potential at each particle\n",
    "    using direct summation method.\n",
    "\n",
    "    Arguments:\n",
    "        particles: the list of particles\n",
    "\n",
    "    \"\"\"\n",
    "    for i, target in enumerate(particles):\n",
    "        for source in (particles[:i] + particles[i+1:]):\n",
    "            r = target.distance(source)\n",
    "            target.phi += source.m / r"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "jaUtRaKKS1ZL"
   },
   "source": [
    "All that's left is to create a list of random particles with assigned masses:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "R8JmgOBtRlnV"
   },
   "outputs": [],
   "source": [
    "n = 1000\n",
    "particles = [Particle(m = 1 / n) for i in range(n)]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "tZPmeLQ2SRF-"
   },
   "source": [
    "**5a)** Run the direct summation code and determine how long it takes with 10, 100, 1000 particles. Is there a relationship?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "5BxIb_3YRuja"
   },
   "outputs": [],
   "source": [
    "# add code here"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "OeTfGagnTcEd"
   },
   "source": [
    "Add answer here."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "BkC-j2c2UXbg"
   },
   "source": [
    "**How do we use Numba on this problem?** There is a subtle issue here - Numba doesn't support jitting native Python classes.  There is a `jit_class` structure in Numba but it's still in early development. But we'd like to have attributes for readable programming. The solution is to build NumPy custom dtypes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "pwAbPrw3SG67"
   },
   "outputs": [],
   "source": [
    "particle_dtype = np.dtype({'names':['x','y','z','m','phi'], \n",
    "                             'formats':[np.double, \n",
    "                                        np.double, \n",
    "                                        np.double, \n",
    "                                        np.double, \n",
    "                                        np.double]})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "0YQzXLouUsXC"
   },
   "outputs": [],
   "source": [
    "myarray = np.ones(3, dtype=particle_dtype)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "gWtollTtU05Q"
   },
   "outputs": [],
   "source": [
    "myarray"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "irR1wcMOVPzH"
   },
   "source": [
    "You can access an individual \"attribute\" like this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "w4oyOqCRU4Zi"
   },
   "outputs": [],
   "source": [
    "myarray[0]['x'] = 2.0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "sL-2QZVKVYYd"
   },
   "source": [
    "**5b)** Write a `jit` function `create_n_random_particles` that takes the arguments `n` (number of particles), `m` (mass of every particle) and a domain within which to generate a random number (as in the class above).\n",
    "It should create an array with `n` elements and `dtype=particle_dtype` and then return that array.\n",
    "​\n",
    "For each particle, the mass should be initialized to the value of `m` and the potential `phi` initialized to zero."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "zuPJFbAFXGZV"
   },
   "source": [
    "*Hint: You will probably want to loop over the number of particles within the function to assign attributes.*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "6_NGgEa_VbLf"
   },
   "outputs": [],
   "source": [
    "# add code here"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Vp_EbHgJXl9y"
   },
   "source": [
    "Now we'll create our array of particles using the new function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "0KKyJWWoXkgj"
   },
   "outputs": [],
   "source": [
    "particles = create_n_random_particles(1000, .001, 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "vcaGjOQIXrjU"
   },
   "outputs": [],
   "source": [
    "particles[:3]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Vx32Tp4kWwQ-"
   },
   "source": [
    "We don't have a `distance` method anymore, so we need to write a function to take care of that.\n",
    "\n",
    "**5c)** Write a `jit` function `distance` to calculate the distance between two particles of dtype `particle_dtype`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "1EkEnyJSVp9G"
   },
   "outputs": [],
   "source": [
    "# add code here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "nX8TaDkLXcJK"
   },
   "outputs": [],
   "source": [
    "distance(particles[0], particles[1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "0y18OqcDX-bR"
   },
   "outputs": [],
   "source": [
    "%%timeit\n",
    "distance(particles[0], particles[1])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "JdZRXDJfYN5t"
   },
   "source": [
    "**5d)** Modify the `direct_sum` function above to instead work on a NumPy array of particles.  Loop over each element in the array and calculate its total potential. Time the result and compare it to your previous version of this function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "FhglZ1oVYEvL"
   },
   "outputs": [],
   "source": [
    "# add code here"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "include_colab_link": true,
   "name": "numba_intro.ipynb",
   "provenance": [],
   "version": "0.3.2"
  },
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}