notebook.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# Numpy: manipulating arrays"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "Numpy's speedup is no joke."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "**Normal Python:**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "data = []\n",
    "for i in range(1000000):\n",
    "    data.append(random.gauss(0, 1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit\n",
    "data2 = []\n",
    "for x in data:\n",
    "    data2.append(x**2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "**Numpy:**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [],
   "source": [
    "import numpy\n",
    "data = numpy.random.normal(0, 1, 1000000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit\n",
    "data2 = data**2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "**A Numpy array is everything normal Python data is not:**\n",
    "\n",
    "   * Loop performed in native bytecode (compiled C)\n",
    "   * Type-checking performed once before loop\n",
    "   * Data are packed in contiguous bytes\n",
    "   * Python's Global Interpreter Lock (GIL) is released during loop"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "**Bonus:**\n",
    "\n",
    "   * Most methods also benefit from hardware vectorization"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "But you have to write your algorithms \"sideways.\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "px = numpy.random.normal(0, 30, 100000)\n",
    "py = numpy.random.normal(0, 30, 100000)\n",
    "pz = numpy.random.normal(0, 300, 100000)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "**Computing one event at a time (\"normal\"):**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit\n",
    "p = numpy.empty(100000)\n",
    "for i in range(len(p)):                                   # for each px[i], py[i], pz[i]\n",
    "    p[i] = numpy.sqrt(px[i]**2 + py[i]**2 + pz[i]**2)     # compute p[i]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "**Computing one column at a time (\"sideways\"):**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit\n",
    "p = numpy.sqrt(px**2 + py**2 + pz**2)       # compute all px**2, then all py**2, then all pz**2, then sum all, then sqrt all"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "Normal math functions are *scalar* (e.g. binary operators like `+` or functions from `import math`). They perform one operation for each appearance in Python source code.\n",
    "\n",
    "Numpy math functions are *vectorized.* Given equal-length arrays as input, they return the same length array as output, performing all loops in compiled C code, possibly doing 4 or 8 at a time in the CPU."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "small_array = numpy.array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "small_array**2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "numpy.sqrt(small_array)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# this won't work because math.sqrt wants a scalar number\n",
    "import math\n",
    "math.sqrt(small_array)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "Numpy arrays are contiguous blocks of bytes in memory, just like C arrays."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "small_array.view(numpy.uint8)        # view (reinterpret_cast) the 64-bit integers as unsigned 8-bit integers"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "They can be multidimensional."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "twod = small_array.reshape(2, 5)     # view as 2 arrays of 5 elements each (a constant-time reinterpretation)\n",
    "twod"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "They can even be arrays of _structs!_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "table = numpy.array([(0, -9, 0.0), (1, -7, 1.1), (2, -5, 2.2), (3, -3, 3.3), (4, -1, 4.4), (5, 1, 5.5), (6, 3, 6.6), (7, 5, 7.7), (8, 7, 8.8), (9, 9, 9.9)],\n",
    "                    dtype=[(\"one\", numpy.uint8), (\"two\", numpy.int64), (\"three\", numpy.double)])\n",
    "table"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "Besides being a common, agreed-upon language for sharing arrays with C/C++ and Fortran code, Numpy has powerful index selection:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "small_array[-1:1:-2]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "twod[1, ::2]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "table[\"three\"][5:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "table[\"three\"][[False, False, False, False, False, True, False, True, False, True]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "table[\"three\"][[7, 5, 5, 3, 2, 7]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "table[\"three\"][table[\"two\"]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "table[\"three\"][numpy.tile([7, 2], 5)]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "The same rules apply to *assigning* to arrays."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "small_array[[False, False, False, False, False, True, False, True, False, True]] = 5000, 7000, 9000"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "small_array"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "Be careful changing things in place!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "view = small_array[5:]\n",
    "view"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "view[0] = 999"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "small_array"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "view.base is small_array"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "small_array.base is None"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "**Exercise:**\n",
    "\n",
    "Suppose you're given a zillion `(px, py, pz, E)` 4-vectors and you want `(E, px, py, pz)` 4-vectors. Do it *fast!*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ZILLION = 1000000\n",
    "fourvectors = numpy.empty((ZILLION, 4))\n",
    "fourvectors[:, 0] = numpy.random.normal(0, 1, ZILLION)\n",
    "fourvectors[:, 1] = numpy.random.normal(0, 1, ZILLION)\n",
    "fourvectors[:, 2] = numpy.random.normal(0, 10, ZILLION)\n",
    "fourvectors[:, 3] = numpy.random.normal(0, 10, ZILLION)**2\n",
    "fourvectors[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit\n",
    "reordered = ???"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# uproot: ROOT → Numpy"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "There are several ways to get ROOT data into Numpy arrays:\n",
    "\n",
    "   * Pythonic iteration in PyROOT (super slow)\n",
    "   * use ROOT's new `TTree::AsMatrix` (for flat tables)\n",
    "   * through a custom C++ function\n",
    "   * with `root_numpy` (compiles against a ROOT version; can segfault with version mismatch)\n",
    "   * uproot* is a pure-Python, minimal dependency reimplementation of ROOT I/O\n",
    "\n",
    "<i>(*disclosure: I'm the author)</i>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "!pip install uproot --user"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [],
   "source": [
    "import uproot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "events = uproot.open(\"http://scikit-hep.org/uproot/examples/HZZ.root\")[\"events\"]\n",
    "events.keys()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "One-per-event ROOT branches become Numpy arrays:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [],
   "source": [
    "events.array(\"MET_px\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "Multi-per-event ROOT branches become so-called \"jagged arrays\":"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "events.array(\"Jet_Px\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "A jagged array represents a list (events) of lists (particles)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "px_events, py_events, pz_events = events.arrays([\"Jet_Px\", \"Jet_Py\", \"Jet_Pz\"], outputtype=tuple)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit\n",
    "p = numpy.empty(2773)\n",
    "j = 0\n",
    "for i in range(events.numentries):\n",
    "    for px, py, pz in zip(px_events[i], py_events[i], pz_events[i]):\n",
    "        p[j] = math.sqrt(px**2 + py**2 + pz**2)\n",
    "        j += 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "But it consists of three contiguous arrays."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(px_events.starts)                # content index where each event starts (inclusive)\n",
    "print(px_events.stops)                 # content index where each event stops (exclusive)\n",
    "print(px_events.content)               # content data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit\n",
    "p = numpy.sqrt(px_events.content**2 + py_events.content**2 + pz_events.content**2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "Numpy is designed for flat arrays, so X of all particles → Y of all particles is easy. What about particle attributes *per event?*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "p = numpy.full(events.numentries, numpy.nan)    # initialize with nan\n",
    "for i in range(events.numentries):\n",
    "    for px, py, pz in zip(px_events[i], py_events[i], pz_events[i]):\n",
    "        # only fill first in each event\n",
    "        if numpy.isnan(p[i]):\n",
    "            p[i] = math.sqrt(px**2 + py**2 + pz**2)\n",
    "p"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "**Step 1:** Which events have at least one jet?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hasajet = (px_events.stops - px_events.starts != 0)\n",
    "hasajet"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "**Step 2:** Get indexes of first jets in events with at least one jet."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "firsts = px_events.starts[hasajet]\n",
    "firsts"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "**Step 3:** Gather values from per-particle arrays into per-event arrays."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "px, py, pz = px_events.content[firsts], py_events.content[firsts], pz_events.content[firsts]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "**Step 4:** Vectorized calculation, assigned to events with at least one jet."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "p = numpy.full(events.numentries, numpy.nan)       # initialize with nan\n",
    "p[hasajet] = numpy.sqrt(px**2 + py**2 + pz**2)     # assign to events through the mask\n",
    "p"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "**Exercise:**\n",
    "\n",
    "Compute the invariant mass of the first two muons in each event that has at least two muons.\n",
    "\n",
    "$m = \\sqrt{(E_1 + E_2)^2 - (px_1 + px_2)^2 - (py_1 + py_2)^2 - (pz_1 + pz_2)^2}$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "q_events = events.array(\"Muon_Charge\")\n",
    "starts, stops = q_events.starts, q_events.stops\n",
    "energy = events.array(\"Muon_E\").content\n",
    "px = events.array(\"Muon_Px\").content\n",
    "py = events.array(\"Muon_Py\").content\n",
    "pz = events.array(\"Muon_Pz\").content"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [],
   "source": [
    "mass = ???"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "!pip install histbook --user"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from histbook import Hist, bin\n",
    "# from vega import VegaLite                                             # if Jupyter Notebook, not JupyterLab"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Hist(bin(\"mass\", 200, 0, 200), fill=mass).step(width=400, height=400)   # .to(VegaLite)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![](src/visualization.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# Pandas: data analysis on indexed tables"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "Pandas is a library for manipulating tabular data, but as we've just seen, so is Numpy.\n",
    "\n",
    "Pandas is based on Numpy. So what's its \"value added?\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "## Rich indexing!!!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "Although the elements of a Numpy array can be accessed in fancy ways, each item is conceptually addressed by a single integer: the entry number. Indexes in Pandas can be typed (e.g. string keys like a `dict`), structured like the two-component index for jagged data below, and sparse (as below)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true,
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [],
   "source": [
    "df = events.pandas.df([\"MET*\", \"Muon*\"])\n",
    "df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "df.loc[2]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "Indexes can be fluidly manipulated. For instance, convert this structured index into integer-indexed columns:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true,
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [],
   "source": [
    "df.reset_index()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "Flatten the sparse subentries into dense columns: a column for each number of muons:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true,
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [],
   "source": [
    "df2 = df.pivot_table(index=\"entry\", columns=\"subentry\")\n",
    "df2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true,
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "df2[b\"Muon_E\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "Numpy concepts and functions apply to Pandas."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [],
   "source": [
    "# make a mask of events with at least two muons (second muon index is not nan)\n",
    "has2muons = numpy.logical_not(numpy.isnan(df2[b\"Muon_E\", 1]))\n",
    "df3 = df2[has2muons]\n",
    "\n",
    "# select particle attributes for the first and second muons through filter\n",
    "e1,  e2  = df3[b\"Muon_E\",  0], df3[b\"Muon_E\",  1]\n",
    "px1, px2 = df3[b\"Muon_Px\", 0], df3[b\"Muon_Px\", 1]\n",
    "py1, py2 = df3[b\"Muon_Py\", 0], df3[b\"Muon_Py\", 1]\n",
    "pz1, pz2 = df3[b\"Muon_Pz\", 0], df3[b\"Muon_Pz\", 1]\n",
    "\n",
    "# vectorized calculation like Numpy\n",
    "mass = numpy.sqrt((e1 + e2)**2 - (px1 + px2)**2 - (py1 + py2)**2 - (pz1 + pz2)**2)\n",
    "mass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "mass.hist()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# Faster math: NumExpr, Numba, and Cython"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "The Python development strategy is to write slow Python until it gets to be a problem and then accelerate the problem spots."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# some require actual Numpy arrays, not Pandas series\n",
    "e1, e2, px1, px2, py1, py2, pz1, pz2 = (\n",
    "    numpy.array(e1), numpy.array(e2), numpy.array(px1), numpy.array(px2),\n",
    "    numpy.array(py1), numpy.array(py2), numpy.array(pz1), numpy.array(pz2))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## NumExpr\n",
    "\n",
    "Mini-compiler to process Numpy arrays in one pass, rather than one per column."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [],
   "source": [
    "import numexpr\n",
    "numexpr.evaluate(\"sqrt((e1 + e2)**2 - (px1 + px2)**2 - (py1 + py2)**2 - (pz1 + pz2)**2)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## Numba\n",
    "\n",
    "More general: uses LLVM to just-in-time compile Python code to native bytecode."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install numba --user"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numba\n",
    "\n",
    "@numba.jit\n",
    "def f(e1, e2, px1, px2, py1, py2, pz1, pz2):\n",
    "    out = numpy.empty(len(e1))\n",
    "    for i in range(len(e1)):\n",
    "        out[i] = math.sqrt((e1[i] + e2[i])**2 - (px1[i] + px2[i])**2 - (py1[i] + py2[i])**2 - (pz1[i] + pz2[i])**2)\n",
    "    return out\n",
    "\n",
    "f(e1, e2, px1, px2, py1, py2, pz1, pz2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## Cython\n",
    "\n",
    "Cython is a half-way language between C++ and Python for static compilation and C++ bindings."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext Cython"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*(continued next page)*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true,
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "%%cython -a --cplus -c-std=c++11 -c-O3\n",
    "\n",
    "import numpy\n",
    "\n",
    "cdef extern from *:\n",
    "    \"\"\"\n",
    "    void doit(int N, float* out, float* e1, float* e2, float* px1, float* px2,\n",
    "                                 float* py1, float* py2, float* pz1, float* pz2) {\n",
    "        for (int i = 0;  i < N;  i++) {\n",
    "            out[i] = sqrt((e1[i] + e2[i])*(e1[i] + e2[i]) - (px1[i] + px2[i])*(px1[i] + px2[i]) -\n",
    "                          (py1[i] + py2[i])*(py1[i] + py2[i]) - (pz1[i] + pz2[i])*(pz1[i] + pz2[i]));\n",
    "        }\n",
    "    }\n",
    "    \"\"\"\n",
    "    void doit(int N, float* out, float* e1, float* e2, float* px1, float* px2, float* py1, float* py2, float* pz1, float* pz2)\n",
    "\n",
    "def g(e1, e2, px1, px2, py1, py2, pz1, pz2):\n",
    "    if (not isinstance(e1, numpy.ndarray) or not isinstance(e2, numpy.ndarray) or\n",
    "        not isinstance(px1, numpy.ndarray) or not isinstance(px2, numpy.ndarray) or\n",
    "        not isinstance(py1, numpy.ndarray) or not isinstance(py2, numpy.ndarray) or\n",
    "        not isinstance(pz1, numpy.ndarray) or not isinstance(pz2, numpy.ndarray)):\n",
    "        raise TypeError(\"these are not arrays\")\n",
    "\n",
    "    out = numpy.empty(len(e1), dtype=numpy.float32)    \n",
    "    doit(len(e1), <float*>(<size_t>out.ctypes.data),\n",
    "         <float*>(<size_t>e1.ctypes.data), <float*>(<size_t>e2.ctypes.data),\n",
    "         <float*>(<size_t>px1.ctypes.data), <float*>(<size_t>px2.ctypes.data),\n",
    "         <float*>(<size_t>py1.ctypes.data), <float*>(<size_t>py2.ctypes.data),\n",
    "         <float*>(<size_t>pz1.ctypes.data), <float*>(<size_t>pz2.ctypes.data))\n",
    "    return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "g(e1, e2, px1, px2, py1, py2, pz1, pz2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# Dask: delayed and distributed computing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "!pip install \"dask[complete]\" --user"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [],
   "source": [
    "lazyarrays = uproot.daskarrays(\"http://scikit-hep.org/uproot/examples/HZZ.root\", \"events\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "lazyarrays"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pt = numpy.sqrt(lazyarrays[b\"MET_px\"]**2 + lazyarrays[b\"MET_py\"]**2)\n",
    "pt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "`pt` is an *instruction* for a calculation that can be distributed across a cluster. (We don't have a cluster.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pt.compute()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# NumPythia, PyJet, and PyPDT: for quick generator-level studies"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## NumPythia"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [],
   "source": [
    "!pip install numpythia --user      # it takes a long time to compile Pythia..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpythia\n",
    "import numpythia.testcmnd\n",
    "pythia = numpythia.Pythia(numpythia.testcmnd.get_cmnd(\"w\"), random_state=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [],
   "source": [
    "import pandas\n",
    "events = list(pythia(events=1))    # pythia is a generator; evaluate it\n",
    "pandas.DataFrame(events[0].all())  # event.all(SELECTION) returns a Numpy record array, so of course we Pandas it"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## pyjet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [],
   "source": [
    "!pip install pyjet --user"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import pyjet\n",
    "\n",
    "fourvectors = events[0].all()[[\"E\", \"px\", \"py\", \"pz\"]]\n",
    "for jet in pyjet.cluster(fourvectors, R=1.0, p=-1, ep=True).inclusive_jets():\n",
    "    print(jet)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## PyPDT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install pypdt --user"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import pypdt\n",
    "\n",
    "for pdgid in set(events[0].all()[\"pdgid\"]):\n",
    "    p = pypdt.get(pdgid)\n",
    "    if p is not None:\n",
    "        print(\"{}: {} mass: {} width {} lifetime {} ctau {}\".format(\n",
    "            pdgid, p.name, p.mass, p.width, p.lifetime, p.ctau))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# The End!"
   ]
  }
 ],
 "metadata": {
  "celltoolbar": "Slideshow",
  "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.5.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}