{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Lab 2: Numba\n",
    "\n",
    "Links: \n",
    "- https://github.com/jpivarski-talks/2021-02-03-pyhep-numba-tutorial\n",
    "- https://henryiii.github.io/compclass/week10/2_rk.html\n",
    "- https://carpentries-incubator.github.io/lesson-gpu-programming/02-cupy/index.html\n",
    "- https://docs.google.com/presentation/d/13MmOKh30IbRJo5sZ8ZujQFr1J3rhqDjA-0Qkfg7JZmE/edit#slide=id.p\n",
    "- https://indico.cern.ch/event/1019958/timetable/#43-cuda-and-python-with-numba\n",
    "- https://jakevdp.github.io/blog/2012/08/18/matplotlib-animation-tutorial/"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Simple stuff\n",
    "\n",
    "Even if you've used Numba before, let's start with the basics.\n",
    "\n",
    "In fact, let's start with NumPy itself."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### NumPy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "NumPy accelerates Python code by replacing loops in Python's virtual machine (with type-checks at runtime) with precompiled loops that transform arrays into arrays."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = np.arange(1000000)\n",
    "data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit\n",
    "\n",
    "output = np.empty(len(data))\n",
    "\n",
    "for i, x in enumerate(data):\n",
    "    output[i] = x**2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit\n",
    "\n",
    "output = data**2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "But if you have to compute a complex formula, NumPy would have to _make an array for each intermediate step_.\n",
    "\n",
    "(There are tricks for circumventing this, but we won't get into that.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "energy = np.random.normal(100, 10, 1000000)\n",
    "px = np.random.normal(0, 10, 1000000)\n",
    "py = np.random.normal(0, 10, 1000000)\n",
    "pz = np.random.normal(0, 10, 1000000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit\n",
    "\n",
    "mass = np.sqrt(energy**2 - px**2 - py**2 - pz**2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The above is equivalent to"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit\n",
    "\n",
    "tmp1 = energy**2\n",
    "tmp2 = px**2\n",
    "tmp3 = py**2\n",
    "tmp4 = pz**2\n",
    "tmp5 = tmp1 - tmp2\n",
    "tmp6 = tmp5 - tmp3\n",
    "tmp7 = tmp6 - tmp4\n",
    "mass = np.sqrt(tmp7)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Numba"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numba as nb"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Numba lets us compile a function to compute a whole formula in one step.\n",
    "\n",
    "```python\n",
    "@nb.jit\n",
    "```\n",
    "\n",
    "means \"JIT-compile\" and\n",
    "\n",
    "```python\n",
    "@nb.njit\n",
    "```\n",
    "\n",
    "means \"really JIT-compile\" because the original has a fallback mode that's getting deprecated. If we're using Numba at all, we don't want it to fall back to ordinary Python."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@nb.njit\n",
    "def compute_mass(energy, px, py, pz):\n",
    "    mass = np.empty(len(energy))\n",
    "    for i in range(len(energy)):\n",
    "        mass[i] = np.sqrt(energy[i]**2 - px[i]**2 - py[i]**2 - pz[i]**2)\n",
    "    return mass"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `compute_mass` function is now \"ready to be compiled.\" It will be compiled when we give it arguments, so that it can propagate types."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "compute_mass(energy, px, py, pz)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit\n",
    "\n",
    "mass = compute_mass(energy, px, py, pz)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Type considerations"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Dynamically typed programs that are not statically typed\n",
    "\n",
    "Much of Python's flexibility comes from the fact that it does not need to know the types of all variables before it starts to run.\n",
    "\n",
    "That dynamism makes it easier to express complex logic (what I call \"bookkeeping\"), but it is a hurdle for speed."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Avoid lists and dicts\n",
    "\n",
    "Although Numba developers are doing a lot of work on supporting Python `list` and `dict` (of identically typed contents), I find it to be too easy to run into unsupported cases. The main problem is that their contents _must_ be typed, but the Python language didn't have ways to express types.\n",
    "\n",
    "(Yes, Python has type annotations now, but they're not fully integrated into Numba yet.)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Runge–Kutta example\n",
    "\n",
    "Let's use Numba to accelerate RK4\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x_max = 1  # Size of x max\n",
    "v_0 = 0\n",
    "koverm = 1  # k / m\n",
    "\n",
    "\n",
    "def f(t, y):\n",
    "    \"Y has two elements, x and v\"\n",
    "    return np.array([-koverm * y[1], y[0]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def rk4_ivp(f, init_y, t):\n",
    "    steps = len(t)\n",
    "    order = len(init_y)\n",
    "\n",
    "    y = np.empty((steps, order))\n",
    "    y[0] = init_y\n",
    "\n",
    "    for n in range(steps - 1):\n",
    "        h = t[n + 1] - t[n]\n",
    "\n",
    "        k1 = h * f(t[n], y[n])  # 2.1\n",
    "        k2 = h * f(t[n] + h / 2, y[n] + k1 / 2)  # 2.2\n",
    "        k3 = h * f(t[n] + h / 2, y[n] + k2 / 2)  # 2.3\n",
    "        k4 = h * f(t[n] + h, y[n] + k3)  # 2.4\n",
    "\n",
    "        y[n + 1] = y[n] + 1 / 6 * (k1 + 2 * k2 + 2 * k3 + k4)  # 2.0\n",
    "\n",
    "    return y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ts = np.linspace(0, 40, 100 + 1)\n",
    "y = rk4_ivp(f, [x_max, v_0], ts)\n",
    "plt.plot(ts, np.cos(ts))\n",
    "plt.plot(ts, y[:, 0], \"--\")\n",
    "plt.xlabel(\"t\")\n",
    "plt.ylabel(\"x\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit\n",
    "ts = np.linspace(0, 40, 1000 + 1)\n",
    "y = rk4_ivp(f, [x_max, v_0], ts)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Animation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib\n",
    "from matplotlib.animation import FuncAnimation\n",
    "\n",
    "# First set up the figure, the axis, and the plot element we want to animate\n",
    "fig = plt.figure()\n",
    "ax = plt.axes(xlim=(-2, 2), ylim=(-2, 2), xlabel=\"x\", ylabel=\"v\")\n",
    "line, = ax.plot([], [], lw=2, marker=\".\")\n",
    "\n",
    "# initialization function: plot the background of each frame\n",
    "def init():\n",
    "    line.set_data([], [])\n",
    "    return line,\n",
    "\n",
    "# animation function.  This is called sequentially\n",
    "# in this case, we will plot the solution from t=0 to the current time step\n",
    "def animate(i):\n",
    "    line.set_data(y[0:i+1, 0], y[0:i+1, 1])\n",
    "    return line,\n",
    "\n",
    "# call the animator.  blit=True means only re-draw the parts that have changed.\n",
    "anim = FuncAnimation(fig, animate, init_func=init,\n",
    "                     frames=len(y), interval=20, blit=True)\n",
    "\n",
    "# save the animation as an mp4.  This requires ffmpeg or mencoder to be\n",
    "# installed.  The extra_args ensure that the x264 codec is used, so that\n",
    "# the video can be embedded in html5.  You may need to adjust this for\n",
    "# your system: for more information, see\n",
    "# http://matplotlib.sourceforge.net/api/animation_api.html\n",
    "anim.save('basic_animation.mp4', fps=30, extra_args=['-vcodec', 'libx264'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Display .mp4 video\n",
    "from IPython.display import Video\n",
    "Video(data=\"basic_animation.mp4\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Note we can also save it as a .gif\n",
    "from IPython.display import Image\n",
    "anim.save(\"basic_animation.gif\", fps=30, writer=\"imagemagick\")\n",
    "Image(open(\"basic_animation.gif\", \"rb\").read())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.5"
  },
  "toc-autonumbering": false,
  "toc-showcode": false,
  "toc-showmarkdowntxt": false
 },
 "nbformat": 4,
 "nbformat_minor": 4
}