nbody_demo.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Optimization and performance in Python\n",
    "\n",
    "Python code has a reputation for being slower than other programming languages, and that is often well-deserved. But we will see here that there are ways to make your python code run orders-of-magnitude faster with relatively little extra work."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "%load_ext line_profiler"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For our problem let's assume we have a collection of point-masses, each with its own 3D coordinates, and we wish to know the gravitational acceleration exerted on each of the masses. This is requied e.g. for performing N-body simulations of planetary systems, star clusters, galaxies, and dark matter structure."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Initialization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_particles = 10**3 # number of particles\n",
    "masses = np.random.rand(num_particles) # particle masses chosen at random on [0,1)\n",
    "coordinates = np.random.normal(size=(num_particles,3)) # particle positions chosen at random in a 3D Gaussian blob"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's code a typical Numpy function that will compute the N-body gravitational acceleration. We will start as simple as possible while working within the numpy paradigm for arrays and docstrings, and explicitly writing out the loop operations.\n",
    "\n",
    "## Naïve implementation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def nbody_accel(masses, coordinates, G=1.):\n",
    "    \"\"\"\n",
    "    Computes the Newtontian gravitational acceleration exerted on each of a set\n",
    "    of point masses.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    masses: array_like\n",
    "        Shape (N,) array of masses\n",
    "    coordinates: array_like\n",
    "        Shape (N,3) array of coordinates\n",
    "    G: float, optional\n",
    "        Value of Newton's gravitational constant (default to 1)\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    accel: ndarray\n",
    "        Shape (N,3) array containing the gravitational acceleration experienced\n",
    "        by each point-mass\n",
    "    \"\"\"\n",
    "\n",
    "    # first declare the array that will store the acceleration\n",
    "    accel = np.zeros_like(coordinates) # array of zeros shaped like coordinate (N,3)\n",
    "    N = coordinates.shape[0] # number of particles\n",
    "\n",
    "    for i in range(N):\n",
    "        for j in range(N):\n",
    "            if i==j: continue # self-force is 0\n",
    "            # first need to calculate the distance between i and j\n",
    "            distance = 0.\n",
    "            for k in range(3):\n",
    "                dx = coordinates[j,k] - coordinates[i,k]\n",
    "                distance += dx*dx\n",
    "            if distance == 0: continue # just skip if points lie on top of each other\n",
    "            distance = np.sqrt(distance)\n",
    "\n",
    "            # now compute the acceleration\n",
    "            for k in range(3):\n",
    "                dx = coordinates[j,k] - coordinates[i,k]\n",
    "                accel[i,k] += G * masses[j] * dx / distance**3\n",
    "\n",
    "    return accel"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Performance test\n",
    "\n",
    "One way to test the performance of a function is to use the `%timeit` magic, which will run the function repeatedly and give you an average of how long it took."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "7.97 s ± 205 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "%timeit nbody_accel(masses,coordinates)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Oh shit that took forever didn't it? And that's just for 1000 masses - imagine if we had to simulate a Milky Way of ~10^11 masses! It would take (10^11/10^3)^2 = 10^16 times longer! Who's got time for that?\n",
    "\n",
    "What's taking so long? One way to get a breakdown of which lines in your function are taking the most time is with the `%lprun` magic provided by the `line_profiler` package"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Timer unit: 1e-06 s\n",
      "\n",
      "Total time: 24.1356 s\n",
      "File: <ipython-input-3-1a5f7a885e4f>\n",
      "Function: nbody_accel at line 1\n",
      "\n",
      "Line #      Hits         Time  Per Hit   % Time  Line Contents\n",
      "==============================================================\n",
      "     1                                           def nbody_accel(masses, coordinates, G=1.):\n",
      "     2                                               \"\"\"\n",
      "     3                                               Computes the Newtontian gravitational acceleration exerted on each of a set\n",
      "     4                                               of point masses.\n",
      "     5                                           \n",
      "     6                                               Parameters\n",
      "     7                                               ----------\n",
      "     8                                               masses: array_like\n",
      "     9                                                   Shape (N,) array of masses\n",
      "    10                                               coordinates: array_like\n",
      "    11                                                   Shape (N,3) array of coordinates\n",
      "    12                                               G: float, optional\n",
      "    13                                                   Value of Newton's gravitational constant (default to 1)\n",
      "    14                                           \n",
      "    15                                               Returns\n",
      "    16                                               -------\n",
      "    17                                               accel: ndarray\n",
      "    18                                                   Shape (N,3) array containing the gravitational acceleration experienced\n",
      "    19                                                   by each point-mass\n",
      "    20                                               \"\"\"\n",
      "    21                                           \n",
      "    22                                               # first declare the array that will store the acceleration\n",
      "    23         1        183.0    183.0      0.0      accel = np.zeros_like(coordinates) # array of zeros shaped like coordinate (N,3)\n",
      "    24         1          3.0      3.0      0.0      N = coordinates.shape[0] # number of particles\n",
      "    25                                           \n",
      "    26      1001        721.0      0.7      0.0      for i in range(N):\n",
      "    27   1001000     533501.0      0.5      2.2          for j in range(N):\n",
      "    28   1000000     555480.0      0.6      2.3              if i==j: continue # self-force is 0\n",
      "    29                                                       # first need to calculate the distance between i and j\n",
      "    30    999000     518899.0      0.5      2.1              distance = 0.\n",
      "    31   3996000    2554522.0      0.6     10.6              for k in range(3):\n",
      "    32   2997000    3466373.0      1.2     14.4                  dx = coordinates[j,k] - coordinates[i,k]\n",
      "    33   2997000    2305378.0      0.8      9.6                  distance += dx*dx\n",
      "    34    999000     659161.0      0.7      2.7              if distance == 0: continue # just skip if points lie on top of each other\n",
      "    35    999000    2183010.0      2.2      9.0              distance = np.sqrt(distance)\n",
      "    36                                           \n",
      "    37                                                       # now compute the acceleration\n",
      "    38   3996000    2678010.0      0.7     11.1              for k in range(3):\n",
      "    39   2997000    3421611.0      1.1     14.2                  dx = coordinates[j,k] - coordinates[i,k]\n",
      "    40   2997000    5258707.0      1.8     21.8                  accel[i,k] += G * masses[j] * dx / distance**3\n",
      "    41                                           \n",
      "    42         1          1.0      1.0      0.0      return accel"
     ]
    }
   ],
   "source": [
    "%lprun -f nbody_accel nbody_accel(masses, coordinates)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We see that there is no single line that takes the vast majority of the time, so optimizing this line-by-line could be time-consuming.\n",
    "\n",
    "**Exercise**: How much can you optimize the above function just by rearranging the way the loop is structured and the floating-point calculations are carried out? (there should be a factor of ~2 on the table). \n",
    "\n",
    "But the reason for the low performance is more fundamental: explicit indexed loop operations like this have a huge amount of overhead when running natively in the Python interpreter. Much of python's flexibility comes with a trade-off for performance, so pure-python numerical code like this will almost always be outperformed by a compiled language like C++ or Fortran.\n",
    "\n",
    "Is there a way to get around this?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Numba\n",
    "\n",
    "Numba performs JIT (just-in-time) compilation of your numerical Python code so that it is possible in theory to obtain performance equal to compiled languages. The simplest way to use it is just to put a `@jit` decordator on your function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/mgrudic/opt/anaconda3/lib/python3.8/site-packages/scipy/__init__.py:146: UserWarning: A NumPy version >=1.16.5 and <1.23.0 is required for this version of SciPy (detected version 1.24.4\n",
      "  warnings.warn(f\"A NumPy version >={np_minversion} and <{np_maxversion}\"\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "9.14 ms ± 382 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "from numba import jit\n",
    "\n",
    "@jit(nopython=True)\n",
    "def nbody_accel_numba(masses, coordinates, G=1.):\n",
    "    \"\"\"\n",
    "    Computes the Newtontian gravitational acceleration exerted on each of a set\n",
    "    of point masses.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    masses: array_like\n",
    "        Shape (N,) array of masses\n",
    "    coordinates: array_like\n",
    "        Shape (N,3) array of coordinates\n",
    "    G: float, optional\n",
    "        Value of Newton's gravitational constant (default to 1)\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    accel: ndarray\n",
    "        Shape (N,3) array containing the gravitational acceleration experienced\n",
    "        by each point-mass\n",
    "    \"\"\"\n",
    "\n",
    "    # first declare the array that will store the acceleration\n",
    "    accel = np.zeros_like(coordinates) # array of zeros shaped like coordinate (N,3)\n",
    "    N = coordinates.shape[0] # number of particles\n",
    "\n",
    "    for i in range(N):\n",
    "        for j in range(N):\n",
    "            if i==j: continue # self-force is 0\n",
    "            # first need to calculate the distance between i and j\n",
    "            distance = 0.\n",
    "            for k in range(3):\n",
    "                dx = coordinates[j,k] - coordinates[i,k]\n",
    "                distance += dx*dx\n",
    "            if distance == 0: continue # just skip if points lie on top of each other\n",
    "            distance = np.sqrt(distance)\n",
    "\n",
    "            # now compute the acceleration\n",
    "            for k in range(3):\n",
    "                dx = coordinates[j,k] - coordinates[i,k]\n",
    "                accel[i,k] += G * masses[j] * dx / distance**3\n",
    "\n",
    "    return accel\n",
    "\n",
    "%timeit nbody_accel_numba(masses,coordinates)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note the factor of ~1000 speedup, obtained by fundamentally changing the way the code gets transformed into instructions! \n",
    "\n",
    "We can get a little more performance by adding the `fastmath=True` argument, which relaxes the requirement that floating-point operations be performed according to the standard spec. This is usually fine, but **always test your function's accuracy to make sure it calculates the result to desired accuracy**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "<ipython-input-7-008862e9faf2>:1: NumbaDeprecationWarning: \u001b[1mThe 'nopython' keyword argument was not supplied to the 'numba.jit' decorator. The implicit default value for this argument is currently False, but it will be changed to True in Numba 0.59.0. See https://numba.readthedocs.io/en/stable/reference/deprecation.html#deprecation-of-object-mode-fall-back-behaviour-when-using-jit for details.\u001b[0m\n",
      "  @jit(fastmath=True)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "4.98 ms ± 276 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "@jit(fastmath=True)\n",
    "def nbody_accel_numba_fastmath(masses, coordinates, G=1.):\n",
    "    \"\"\"\n",
    "    Computes the Newtontian gravitational acceleration exerted on each of a set\n",
    "    of point masses.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    masses: array_like\n",
    "        Shape (N,) array of masses\n",
    "    coordinates: array_like\n",
    "        Shape (N,3) array of coordinates\n",
    "    G: float, optional\n",
    "        Value of Newton's gravitational constant (default to 1)\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    accel: ndarray\n",
    "        Shape (N,3) array containing the gravitational acceleration experienced\n",
    "        by each point-mass\n",
    "    \"\"\"\n",
    "\n",
    "    # first declare the array that will store the acceleration\n",
    "    accel = np.zeros_like(coordinates) # array of zeros shaped like coordinate (N,3)\n",
    "    N = coordinates.shape[0] # number of particles\n",
    "\n",
    "    for i in range(N):\n",
    "        for j in range(N):\n",
    "            if i==j: continue # self-force is 0\n",
    "            # first need to calculate the distance between i and j\n",
    "            distance = 0.\n",
    "            for k in range(3):\n",
    "                dx = coordinates[j,k] - coordinates[i,k]\n",
    "                distance += dx*dx\n",
    "            if distance == 0: continue # just skip if points lie on top of each other\n",
    "            distance = np.sqrt(distance)\n",
    "\n",
    "            # now compute the acceleration\n",
    "            for k in range(3):\n",
    "                dx = coordinates[j,k] - coordinates[i,k]\n",
    "                accel[i,k] += G * masses[j] * dx / distance**3\n",
    "\n",
    "    return accel\n",
    "\n",
    "%timeit nbody_accel_numba_fastmath(masses, coordinates)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Parallelism\n",
    "There is even more performance on the table here: your computer has **multiple** processors, and so far we have just been using only one. The easiest way to parallelize this is to run the outer loop in parallel using `prange`, and tell numba to use it with `parallel=True`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "<ipython-input-8-dd9c992efc50>:3: NumbaDeprecationWarning: \u001b[1mThe 'nopython' keyword argument was not supplied to the 'numba.jit' decorator. The implicit default value for this argument is currently False, but it will be changed to True in Numba 0.59.0. See https://numba.readthedocs.io/en/stable/reference/deprecation.html#deprecation-of-object-mode-fall-back-behaviour-when-using-jit for details.\u001b[0m\n",
      "  @jit(fastmath=True,parallel=True)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.35 ms ± 200 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/mgrudic/opt/anaconda3/lib/python3.8/site-packages/numba/np/ufunc/parallel.py:371: NumbaWarning: \u001b[1mThe TBB threading layer requires TBB version 2021 update 6 or later i.e., TBB_INTERFACE_VERSION >= 12060. Found TBB_INTERFACE_VERSION = 12040. The TBB threading layer is disabled.\u001b[0m\n",
      "  warnings.warn(problem)\n"
     ]
    }
   ],
   "source": [
    "from numba import prange\n",
    "\n",
    "@jit(fastmath=True,parallel=True)\n",
    "def nbody_accel_numba_fastmath_parallel(masses, coordinates, G=1.):\n",
    "    \"\"\"\n",
    "    Computes the Newtontian gravitational acceleration exerted on each of a set\n",
    "    of point masses.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    masses: array_like\n",
    "        Shape (N,) array of masses\n",
    "    coordinates: array_like\n",
    "        Shape (N,3) array of coordinates\n",
    "    G: float, optional\n",
    "        Value of Newton's gravitational constant (default to 1)\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    accel: ndarray\n",
    "        Shape (N,3) array containing the gravitational acceleration experienced\n",
    "        by each point-mass\n",
    "    \"\"\"\n",
    "\n",
    "    # first declare the array that will store the acceleration\n",
    "    accel = np.zeros_like(coordinates) # array of zeros shaped like coordinate (N,3)\n",
    "    N = coordinates.shape[0] # number of particles\n",
    "\n",
    "    for i in prange(N):\n",
    "        for j in range(N):\n",
    "            if i==j: continue # self-force is 0\n",
    "            # first need to calculate the distance between i and j\n",
    "            distance = 0.\n",
    "            for k in range(3):\n",
    "                dx = coordinates[j,k] - coordinates[i,k]\n",
    "                distance += dx*dx\n",
    "            if distance == 0: continue # just skip if points lie on top of each other\n",
    "            distance = np.sqrt(distance)\n",
    "\n",
    "            # now compute the acceleration\n",
    "            for k in range(3):\n",
    "                dx = coordinates[j,k] - coordinates[i,k]\n",
    "                accel[i,k] += G * masses[j] * dx / distance**3\n",
    "\n",
    "    return accel\n",
    "\n",
    "%timeit nbody_accel_numba_fastmath_parallel(masses, coordinates, G=1.)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Other options for optimization\n",
    "\n",
    "There are many other ways to write performant python code than just the numba parallel CPU coding we have done here. Any python code can in principle be parallelized using python's native [multiprocessing](https://docs.python.org/3/library/multiprocessing.html). [joblib](https://joblib.readthedocs.io/en/stable/) offers functionality for parallelism with a different implementation.  [numba also supports running code on the GPU](https://numba.readthedocs.io/en/stable/cuda/index.html). [JAX](https://jax.readthedocs.io/en/latest/quickstart.html) is best known for its use in machine learning but is also more broadly useful for coding GPU-portable python code consisting of function compositions and array operations. [mpi4py](https://mpi4py.readthedocs.io/en/stable/mpi4py.html) implements MPI (Message Passing Interface) for programs that may need to run distributed across multiple computers (i.e. supercomputers). The best choice will depend on your particular problem and requirements.\n",
    "\n",
    "Python may still have a reputation for being slow, but the reality is that in many cases the time-to-solution (coding+computation) for a give project can be shorter, just by taking advantage of Python's extensive community library support while optimizing the most numerically-intensive parts as we have here."
   ]
  }
 ],
 "metadata": {
  "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.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}