02-columnar-analysis-awkward-array.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Columnar analysis with Awkward Array"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## How this works as a hands-on tutorial\n",
    "\n",
    "Even though I don't have formal exercises scattered throughout these notebooks, this session can still be interactive.\n",
    "\n",
    "   * **You** should open each notebook in Binder (see [GitHub README](https://github.com/jpivarski/2020-06-08-uproot-awkward-columnar-hats)) and evaluate cells, following along with me.\n",
    "   * **I** should pause frequently and stay open to questions. I'll be monitoring the videoconference chat.\n",
    "   * **We** should feel free to step off the path and try to answer \"What if?\" questions in real time.\n",
    "\n",
    "Not all digressions will lead to an answer—I often realize, \"That's why it didn't work!\" long after the tutorial is over—but tinkering is how we learn.\n",
    "\n",
    "Consider this a tour and I'm your guide. The planned route is a suggestion to get things started, but your questions and wayfaring are more important.\n",
    "\n",
    "<br><br><br>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Array-based programming\n",
    "\n",
    "One of the first programming languages, named **APL** (\"A Programming Language\") was array-based. It started as a notation for _describing_ hand-written machine code and was later made interactive.\n",
    "\n",
    "**Nial** was also theoretically motivated, and the two of these inspired a generation of direct descendants (green).\n",
    "\n",
    "Meanwhile, the **S** language for statistics borrowed many of these ideas while being focused on a particular domain. Its descendent, **R**, is still widely used.\n",
    "\n",
    "**IDL** was invented for the sciences and gained a lot of traction as an alternative to writing custom Fortran, again using vectorization as a first-class concept.\n",
    "\n",
    "**MATLAB** was similarly gained traction in the sciences as a commercial product.\n",
    "\n",
    "**PDL** (Perl Data Language) and **NumPy** introduced the same concepts as libraries within an established language (Perl and Python). **Julia** has some vector-like interfaces, though its focus is on just-in-time compiling imperative code."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![](img/apl-timeline.png)\n",
    "\n",
    "<br><br><br>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Common features of array-based languages:\n",
    "\n",
    "   * Arrays are the central data type with most operations applying to arrays. (By contrast, C requires explicit iteration over the arrays: it's imperative.)\n",
    "   * They are _all_ interactive languages. The array-at-a-time logic makes it possible to define precompiled routines that run in response to user commands.\n",
    "   * They are primarily data analysis languages, highly targeted to the sciences and statistics.\n",
    "\n",
    "In retrospect, it sounds like a perfect fit to what we want to do.\n",
    "\n",
    "<br><br><br>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this plot of the \"astronomical\" rise of Python, note that 2 of the 3 languages it's displacing are array languages.\n",
    "\n",
    "![](img/mentions-of-programming-languages.png)\n",
    "\n",
    "<br><br><br>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Why not for particle physics?\n",
    "\n",
    "Because **data structures**. Particle physicists have _always_ needed to deal with complex data structures, so much so that we invented packages to add them to Fortran.\n",
    "\n",
    "The following is from [_Initiation to HYDRA_ by R.K. Böck (1976)](https://cds.cern.ch/record/864527?ln=en) as part of an explanation of what a \"data structure\" is, at a time before Fortran had `FOR` loops. (HYDRA was merged into ZEBRA, which became the basis for ROOT I/O.)\n",
    "\n",
    "We would draw similar diagrams today.\n",
    "\n",
    "![](img/hydra-2.png)\n",
    "\n",
    "<br><br><br>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "But the modify-compile-rerun cycle of C++ is too long for interactive data analysis. That's why ROOT invented CINT and then Cling.\n",
    "\n",
    "But C++ is too complex of a language for data-focused tasks. That's why I was thinking a lot about [extending query languages (like SQL) to data structures](https://stackoverflow.com/questions/38831961/what-declarative-language-is-good-at-analysis-of-tree-like-data).\n",
    "\n",
    "But I was surprised by how useful the simple JaggedArray class in Uproot turned out to be. My conclusion was that you don't need a new language, just some data types and operations.\n",
    "\n",
    "![](img/uproot-awkward-timeline.png)\n",
    "\n",
    "<br><br><br>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<font size=\"15\">That's what </font><img src=\"img/awkward-logo-300px.png\" style=\"vertical-align:middle\"><font size=\"15\"> is.</font>\n",
    "\n",
    "\n",
    "Just arrays and functions on arrays, but with awkward shapes.\n",
    "\n",
    "![](img/cartoon-schematic.png)\n",
    "\n",
    "<br><br><br>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Let's start with a non-physics example\n",
    "\n",
    "To get a feel for what this means, let's look at _something completely different_ from a Z peak: [Chicago bike paths](https://github.com/Chicago/osd-bike-routes/blob/master/data/Bikeroutes.geojson)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "bikeroutes_json = open(\"data/Bikeroutes.geojson\").read()\n",
    "bikeroutes_pyobj = json.loads(bikeroutes_json)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# First thousand bytes...\n",
    "print(bikeroutes_json[:1000])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It has a lot of structure—metadata and street names mixed in with the longitude, latitude points. But we can read it in as arrays."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import awkward1 as ak\n",
    "\n",
    "bikeroutes = ak.Record(bikeroutes_pyobj)\n",
    "bikeroutes"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The analysis-relevant information about the data's structure is contained in its `type`. (Note: this is [Datashape](https://datashape.readthedocs.io/en/latest/) notation, which was invented for this sort of thing.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ak.type(bikeroutes)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Everything is a `string` or an `option[string]` (i.e. can be `null`) except the coordinates, which are `var * var * var * float64` (triply jagged array).\n",
    "\n",
    "Let's go straight for the coordinates."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bikeroutes[\"features\", \"geometry\", \"coordinates\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "or"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bikeroutes.features.geometry.coordinates"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ak.to_list(bikeroutes.features.geometry.coordinates[:3])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The third axis happens to have length 2 in all cases, but it came from JSON, which can't guarantee that lists have a certain length, so Awkward identifies it as `var` (variable-width, i.e. \"jagged\").\n",
    "\n",
    "We _could_ enforce this by applying a slice that has a fixed length."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ak.type(bikeroutes.features.geometry.coordinates)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ak.type(bikeroutes.features.geometry.coordinates[:, :, :, [0, 1]])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This distinction between fixed-size and in-principle-variable size is important in general, though not very important for this example.\n",
    "\n",
    "A more important question is, what do these levels of jaggedness represent?\n",
    "\n",
    "Let's pick one item and print it out in full detail."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ak.to_list(bikeroutes.features[751])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The hint is \"MultiLineString\": a bike route consists of disconnected lines. (I guess you have to pick up your bike and walk it.)\n",
    "\n",
    "Most routes are a single connected line; I found this extreme using [ak.num](https://awkward-array.readthedocs.io/en/latest/_auto/ak.num.html), a function for jagged multiplicity."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ak.num(bikeroutes.features.geometry.coordinates)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ak.argmax(ak.num(bikeroutes.features.geometry.coordinates))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[ak.max](https://awkward-array.readthedocs.io/en/latest/_auto/ak.max.html) is just like [np.max](https://numpy.org/doc/1.18/reference/generated/numpy.amax.html) from NumPy except that it recognizes Awkward Arrays.\n",
    "\n",
    "By contrast, [ak.num](https://awkward-array.readthedocs.io/en/latest/_auto/ak.num.html) could not have a NumPy equivalent because it provides information that would always be trivial with NumPy's rectangular arrays.\n",
    "\n",
    "Functions that overlap NumPy functions, like [ak.max](https://awkward-array.readthedocs.io/en/latest/_auto/ak.max.html) does for [np.max](https://numpy.org/doc/1.18/reference/generated/numpy.amax.html), have exactly the same interface and defaults. If you have NumPy 1.17 or above, they're actually interchangeable (NumPy recognizes that it's looking at a non-NumPy arrays and defers to our implementation)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "np.argmax(ak.num(bikeroutes.features.geometry.coordinates))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To fill out the pattern set by NumPy, most of the functions have an `axis` parameter indicating the depth of nestedness where you want the function to apply.\n",
    "\n",
    "This can be particularly useful for [ak.num](https://awkward-array.readthedocs.io/en/latest/_auto/ak.num.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# most routes have a single contiguous path\n",
    "ak.num(bikeroutes.features.geometry.coordinates, axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# paths can have many or few longitude, latitude points\n",
    "ak.num(bikeroutes.features.geometry.coordinates, axis=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# all of the longitude, latitude points have exactly two numbers\n",
    "ak.num(bikeroutes.features.geometry.coordinates, axis=3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notice the data type: the number of entries in deeply nested data is itself a nested structure.\n",
    "\n",
    "We can verify that the longitude, latitude points really do have two values like this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "num = ak.num(bikeroutes.features.geometry.coordinates, axis=3)\n",
    "num == 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ak.all(num == 2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[ak.all](https://awkward-array.readthedocs.io/en/latest/_auto/ak.all.html) is a reducer (like [np.all](https://numpy.org/doc/1.18/reference/generated/numpy.all.html)), which turns arrays into scalars.\n",
    "\n",
    "Its default `axis` is `None`, meaning \"reduce everything.\" We can also partially reduce.\n",
    "\n",
    "`axis=-1` means \"deepest axis,\" which is the most-often useful axis, apart from `None`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ak.all(num == 2, axis=-1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So now let's do something useful: how about computing the length of each bike route?\n",
    "\n",
    "First, get the longitude and latitude separately."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "longitude = bikeroutes.features.geometry.coordinates[..., 0]\n",
    "latitude = bikeroutes.features.geometry.coordinates[..., 1]\n",
    "longitude, latitude"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The ellipsis (`...`) saved me from having to type `coordinates[:, :, :, 0]`, having to know the exact depth when I wanted the deepest. That's also a NumPy thing."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "At our longtidue and latitude, one degree of longitude corresponds to 82.7 km and one degree of latitude corresponds to 111.1 km (I looked that up elsewhere).\n",
    "\n",
    "Functions like [ak.mean](https://awkward-array.readthedocs.io/en/latest/_auto/ak.mean.html)/[np.mean](https://numpy.org/doc/1.18/reference/generated/numpy.mean.html) have the same interface as reducers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "km_east = (longitude - np.mean(longitude)) * 82.7\n",
    "km_north = (latitude - np.mean(latitude)) * 111.1\n",
    "km_east, km_north"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So now all of the paths are in distance units (km), relative to the center of Chicago (the center of all the points at least; we only needed a convenient origin).\n",
    "\n",
    "Think, for a moment, about what that transformation would have required to do it in \"for\" loops. Even if speed were not an issue, it would be a lot of typing."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To compute lengths, we need distances _between_ points. So we want to match pairs of points along each path.\n",
    "\n",
    "The way you'd do this with a NumPy array is with slices:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "path = np.array([1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8, 9.9])\n",
    "path[1:] - path[:-1]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "where `path[1:]` drops the first element and `path[:-1]` drops the last element, which is exactly what you want to subtract to get all the distances in between.\n",
    "\n",
    "We can do the same thing with our Awkward Arrays, even though they have different lengths at the deepest level of jaggedness."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "km_east[0, 0, :-1], km_east[0, 0, 1:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "km_east[0, 0, :-1] - km_east[0, 0, 1:]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Doing it for all paths at once is no more difficult than doing it for the first path."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "km_east[:, :, :-1] - km_east[:, :, 1:]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Since we're all familiar with using $\\sqrt{(x_i - x_{i + 1})^2 + (y_i - y_{i + 1})^2}$ as a distance formula, let's jump to the answer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "segment_length = np.sqrt((km_east[:, :, 1:] - km_east[:, :, :-1])**2 +\n",
    "                         (km_north[:, :, 1:] - km_north[:, :, :-1])**2)\n",
    "segment_length"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So now we have replaced paths of length _n_ (where _n_ is variable) with segment distances of length _n ‒ 1_.\n",
    "\n",
    "We probably want the length of each path, so... reducer! This one is [ak.sum](https://awkward-array.readthedocs.io/en/latest/_auto/ak.sum.html)/[np.sum](https://numpy.org/doc/1.18/reference/generated/numpy.sum.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "path_length = np.sum(segment_length, axis=-1)\n",
    "path_length"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Okay, but some routes have multiple paths (though most have exactly one). These `path_lengths` have the same multiplicity as the paths."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ak.num(path_length), ak.num(bikeroutes.features.geometry.coordinates)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ak.num(path_length) == ak.num(bikeroutes.features.geometry.coordinates)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ak.all(ak.num(path_length) == ak.num(bikeroutes.features.geometry.coordinates))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So let's reduce again."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "route_length = np.sum(path_length, axis=-1)\n",
    "route_length"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "There you have it. We can also put this new derived column into the original array, if that's useful for anything.\n",
    "\n",
    "(Note: you have to assign with square brackets and strings, not attributes, because attribute-assignment would lead to confusion about assigning to temporary copies. Pandas has the same problem, and they're deprecating attribute-assignemnt because of it.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bikeroutes[\"features\", \"route_length\"] = route_length"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ak.to_list(bikeroutes.features[751])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now every record has a `route_length` field."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bikeroutes.features.route_length"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "That's how these calculations go. If we were to do the same thing with Python for loops, it would be a lot more verbose and slower."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit\n",
    "\n",
    "total_length = []\n",
    "for route in bikeroutes_pyobj[\"features\"]:\n",
    "    route_length = []\n",
    "    for polyline in route[\"geometry\"][\"coordinates\"]:\n",
    "        segment_length = []\n",
    "        last = None\n",
    "        for lng, lat in polyline:\n",
    "            km_east = lng * 82.7\n",
    "            km_north = lat * 111.1\n",
    "            if last is not None:\n",
    "                dx2 = (km_east - last[0])**2\n",
    "                dy2 = (km_north - last[1])**2\n",
    "                segment_length.append(np.sqrt(dx2 + dy2))\n",
    "            last = (km_east, km_north)\n",
    "\n",
    "        route_length.append(sum(segment_length))\n",
    "    total_length.append(sum(route_length))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit\n",
    "\n",
    "km_east = bikeroutes.features.geometry.coordinates[..., 0] * 82.7\n",
    "km_north = bikeroutes.features.geometry.coordinates[..., 1] * 111.1\n",
    "\n",
    "segment_length = np.sqrt((km_east[:, :, 1:] - km_east[:, :, :-1])**2 +\n",
    "                         (km_north[:, :, 1:] - km_north[:, :, :-1])**2)\n",
    "\n",
    "route_length = np.sum(segment_length, axis=-1)\n",
    "total_length = np.sum(route_length, axis=-1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the limit, it's a factor of 8 faster. This isn't even a _great_ speedup; other examples easily reach a factor of 100.\n",
    "\n",
    "The point is that we're in the \"fast math\" regieme, where typical examples are an order of magnitude or two faster than \"slow Python.\"\n",
    "\n",
    "![](img/bikeroutes-scaling.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Particle physics data\n",
    "\n",
    "Now, of course, we want to see what this can do for particle physics data. This public CMS file has a million events and is often used to make dimuon spectra.\n",
    "\n",
    "The recommended way to get Awkward 1 arrays is to read with Uproot 3 (the current version) and then call [ak.from_awkward0](https://awkward-array.readthedocs.io/en/latest/_auto/ak.from_awkward0.html).\n",
    "\n",
    "(Uproot 4 will produce Awkward 1 arrays natively.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import uproot\n",
    "cms_dict = uproot.open(\"data/cms_opendata_2012_nanoaod_DoubleMuParked.root\")[\"Events\"].arrays()\n",
    "cms_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cms_dict_ak1 = {name.decode(): ak.from_awkward0(array) for name, array in cms_dict.items()}\n",
    "cms_dict_ak1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You might have noticed that the conversion was faster than the reading—actually the \"conversion\" only replaces old metadata with new metadata—they're the same array buffers in memory.\n",
    "\n",
    "Switching Awkward 0 ↔ Awkward 1 is an inexpensive operation, in terms of time and memory."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Speaking of memory, I loaded these million events into memory because they're only 250 MB of RAM.\n",
    "\n",
    "As Python objects, they're 1307 MB, 5.2× larger."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data as columns AND objects\n",
    "\n",
    "NanoAOD was invented at around the same time as Uproot—the success of NanoAOD is part of what's convincing me that directly engaging in array-based programming is working for physicists.\n",
    "\n",
    "Instead of letting C++ objects be split as an implementation detail in ROOT (see Uproot session), the data are \"pre-split\" into branches like `Muon_pt`, `Muon_eta`, and `Muon_phi`.\n",
    "\n",
    "But now we can have it both ways. The branches of a NanoAOD file can be the columns of an Awkward Array with event/particle structure."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cms_events = ak.zip({\n",
    "    \"run\": cms_dict_ak1[\"run\"],\n",
    "    \"luminosityBlock\": cms_dict_ak1[\"luminosityBlock\"],\n",
    "    \"event\": cms_dict_ak1[\"event\"],\n",
    "    \"PV\": ak.zip({\n",
    "        \"x\": cms_dict_ak1[\"PV_x\"],\n",
    "        \"y\": cms_dict_ak1[\"PV_y\"],\n",
    "        \"z\": cms_dict_ak1[\"PV_z\"],\n",
    "    }),\n",
    "    \"muons\": ak.zip({\n",
    "        \"pt\": cms_dict_ak1[\"Muon_pt\"],\n",
    "        \"eta\": cms_dict_ak1[\"Muon_eta\"],\n",
    "        \"phi\": cms_dict_ak1[\"Muon_phi\"],\n",
    "        \"mass\": cms_dict_ak1[\"Muon_mass\"],\n",
    "        \"charge\": cms_dict_ak1[\"Muon_charge\"],\n",
    "        \"pfRelIso04_all\": cms_dict_ak1[\"Muon_pfRelIso04_all\"],\n",
    "        \"tightId\": cms_dict_ak1[\"Muon_tightId\"]\n",
    "    })\n",
    "}, depth_limit=1)\n",
    "\n",
    "cms_events"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# this is the first event\n",
    "ak.to_list(cms_events[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# this is the first muon in the first event\n",
    "ak.to_list(cms_events[0, \"muons\", 0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The above _manually_ links NanoAOD branches to record array structures, using [ak.zip](https://awkward-array.readthedocs.io/en/latest/_auto/ak.zip.html), but Nick has written a more general (and lazy) converter called [NanoEvents](https://github.com/CoffeaTeam/coffea/blob/9a29fe47fc690051be50773d262ee74e805a2f60/binder/nanoevents.ipynb) (look it up in the Coffea tools).\n",
    "\n",
    "That is what I would recommend for analysis of NanoAOD."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Why making and breaking record structures is okay\n",
    "\n",
    "Zipping and projecting are essentially \"free\" from a computational point of view because we're just moving around metadata (i.e. _O(1)_, the computation time does not scale with the number of events).\n",
    "\n",
    "This is easier to see with a simple example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = ak.Array([[0.0, 1.1, 2.2], [], [3.3, 4.4]])\n",
    "y = ak.Array([[[], [1], [1, 2]], [], [[1, 2, 3], [1, 2, 3, 4]]])\n",
    "\n",
    "columnar = ak.zip({\"x\": x, \"y\": y}, depth_limit=2)\n",
    "columnar"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ak.to_list(columnar)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ak.type(columnar)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We've \"zipped\" `x` and `y` together to make records with fields named `x` and `y`.\n",
    "\n",
    "The `depth_limit=2` ensured that the outermost level of jaggedness is shared in the output but not the innermost. Here's what would happen if `depth_limit` were `1` or `3`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ak.to_list(ak.zip({\"x\": x, \"y\": y}, depth_limit=1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ak.to_list(ak.zip({\"x\": x, \"y\": y}, depth_limit=2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ak.to_list(ak.zip({\"x\": x, \"y\": y}, depth_limit=3))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "(The duplication of `x` fields to match `y` is called \"broadcasting.\" The same idea [exists in NumPy](https://numpy.org/doc/stable/user/basics.broadcasting.html), but we've extended it for jaggedness.)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "I've been using [ak.to_list](https://awkward-array.readthedocs.io/en/latest/_auto/ak.to_list.html) to show the array in its entirety by converting it to Python objects and letting Jupyter format them nicely.\n",
    "\n",
    "We should be thinking about the arrays in these terms: as nested lists and records. Printing out a very small sample with [ak.to_list](https://awkward-array.readthedocs.io/en/latest/_auto/ak.to_list.html) is a great debugging technique.\n",
    "\n",
    "But this is not how the data are laid out in memory. If you want to get into the details of how it works, look into the `layout`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "columnar.layout"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This hierarchical structure of the layout mirrors the structure of the [ak.to_list](https://awkward-array.readthedocs.io/en/latest/_auto/ak.to_list.html) output with one important difference: the number of nodes scales with the complexity of the data structure, not the number of elements in the arrays. This example has 5 nodes, but if it had a million elements of the same data _type_, it will still just have 5 nodes.\n",
    "\n",
    "All the slow, bookkeeping-type code that has to have dynamic types to address different types without recompilation only has to walk over these 5 nodes.\n",
    "\n",
    "All the fast, vectorized code is precompiled to run over the arrays within those 5 nodes.\n",
    "\n",
    "![](img/example-hierarchy.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For our `cms_events`, this means that 17 nodes represent a million events."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cms_events.layout"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Another consequence of this is that it's easy to \"project\" fields out of nested records into non-record arrays. (Again, just rearranging metadata.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ak.to_list(columnar)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ak.to_list(columnar.x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ak.to_list(columnar.y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cms_events.PV.x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cms_events.muons.pt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In your analysis, feel free to zip and project as much as you find useful. Records are _fluid_."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Imperative programming\n",
    "\n",
    "At this point, I'd usually show how slicing and combinatorics functions like [ak.cartesian](https://awkward-array.readthedocs.io/en/latest/_auto/ak.cartesian.html) and [ak.combinations](https://awkward-array.readthedocs.io/en/latest/_auto/ak.combinations.html) can be used to develop an analysis.\n",
    "\n",
    "But that gives the impression that columnar operations are the _only_ way to work. Although some operations are more convenient in this form, others are not.\n",
    "\n",
    "Therefore, I've been putting a lot of effort into making these arrays usable from [Numba](http://numba.pydata.org/), a just-in-time compiler for Python, which lets you write standard for-loop style (imperative) code, compile it on the fly, and run it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numba as nb\n",
    "\n",
    "@nb.jit\n",
    "def Zmass(events):\n",
    "    # How big of an output array should we allocate?\n",
    "    numZ = 0\n",
    "    for event in events:\n",
    "        numZ += max(0, len(event.muons) * (len(event.muons) - 1) // 2)\n",
    "    \n",
    "    # Allocate the output array.\n",
    "    out = np.empty(numZ, np.float64)\n",
    "\n",
    "    # Fill the output array.\n",
    "    numZ = 0\n",
    "    for event in events:\n",
    "        for i in range(len(event.muons)):\n",
    "            for j in range(i + 1, len(event.muons)):\n",
    "                m1 = event.muons[i]\n",
    "                m2 = event.muons[j]\n",
    "                out[numZ] = np.sqrt(2*m1.pt*m2.pt*(np.cosh(m1.eta - m2.eta) - np.cos(m1.phi - m2.phi)))\n",
    "                numZ += 1\n",
    "\n",
    "    return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Zmass(cms_events)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "plt.hist(Zmass(cms_events), bins=100, range=(60, 120));"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.hist(Zmass(cms_events), bins=300, range=(0, 12))\n",
    "plt.yscale(\"log\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now the precise way the events were zipped _is_ important. With a flatter structure, the code would have to be entirely different:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cms_events_flat = ak.zip(cms_dict_ak1, depth_limit=1)\n",
    "ak.type(cms_events_flat)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@nb.jit\n",
    "def Zmass_flat(events):\n",
    "    numZ = 0\n",
    "    for event in events:\n",
    "        numZ += max(0, event.nMuon * (event.nMuon - 1) // 2)\n",
    "    out = np.empty(numZ, np.float64)\n",
    "\n",
    "    numZ = 0\n",
    "    for event in events:\n",
    "        for i in range(event.nMuon):\n",
    "            for j in range(i + 1, event.nMuon):\n",
    "                pt1 = event.Muon_pt[i]    # muons are not objects containing pt, eta, phi\n",
    "                pt2 = event.Muon_pt[j]    # the Muon_pt, Muon_eta, Muon_phi arrays are separate\n",
    "                eta1 = event.Muon_eta[i]\n",
    "                eta2 = event.Muon_eta[j]\n",
    "                phi1 = event.Muon_phi[i]\n",
    "                phi2 = event.Muon_phi[j]\n",
    "                out[numZ] = np.sqrt(2*pt1*pt2*(np.cosh(eta1 - eta2) - np.cos(phi1 - phi2)))\n",
    "                numZ += 1\n",
    "    return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with np.errstate(divide=\"ignore\"):\n",
    "    plt.hist(np.log(Zmass(cms_events)), bins=1000, range=(0, 5))\n",
    "plt.yscale(\"log\")\n",
    "plt.xlabel(\"log(mass)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you'd rather write conventional loop-style analysis code, then Numba may be the best way to do it.\n",
    "\n",
    "The advantage **Numba** brings is that the compilation is fast, doesn't take you out of the interactive environment, and is identical to Python code that you can test on small samples.\n",
    "\n",
    "The advantage **Awkward Array** brings is efficiently getting data from ROOT files into Numba compiled functions. Numba's builtin interface acceps NumPy arrays, which are unstructured, and Python objects, which are slow to \"unbox\" into statically typed values. Awkward keeps the columnar ROOT data columnar with static types for compilation.\n",
    "\n",
    "There are some **limitations**, however:\n",
    "\n",
    "   * Numba publishes a list of [Python language features](https://numba.pydata.org/numba-doc/dev/reference/pysupported.html) and [NumPy functions](https://numba.pydata.org/numba-doc/dev/reference/numpysupported.html) that it supports. You have to learn the \"sublanguage\" of Python that can be statically typed and compiled.\n",
    "   * Array-at-a-time functions in the Awkward Array library are not supported: in a compiled function, you _must_ use imperative code.\n",
    "   * Awkward Arrays are read-only. Therefore, you can't create new arrays inside a compiled function."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Creating arrays in a Numba-compiled function\n",
    "\n",
    "In the above example, the output was a NumPy array because Numba supports the creation and filling of NumPy arrays.\n",
    "\n",
    "If the derived quantities you want to create are \"flat\" in that sense, outputting NumPy is a good option. Note that this includes boolean arrays that you might use to cut an Awkward Array outside of the compiled function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@nb.jit\n",
    "def event_has_z(event):\n",
    "    for i in range(len(event.muons)):\n",
    "        for j in range(i + 1, len(event.muons)):\n",
    "            m1 = event.muons[i]\n",
    "            m2 = event.muons[j]\n",
    "            mass = np.sqrt(2*m1.pt*m2.pt*(np.cosh(m1.eta - m2.eta) - np.cos(m1.phi - m2.phi)))\n",
    "            if 60 < mass and mass < 120:\n",
    "                return True\n",
    "    return False\n",
    "\n",
    "@nb.jit\n",
    "def has_z(events):\n",
    "    i = 0\n",
    "    out = np.empty(len(events), np.bool_)\n",
    "    for event in events:\n",
    "        out[i] = event_has_z(event)\n",
    "        i += 1\n",
    "    return out\n",
    "\n",
    "mask = has_z(cms_events)\n",
    "mask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cms_events[mask]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this case, what we wanted and got in the end is a data structure, but Numba only had to return a flat array of booleans.\n",
    "\n",
    "However, there will be cases when you want to get a data structure out. [ak.ArrayBuilder](https://awkward-array.readthedocs.io/en/latest/_auto/ak.ArrayBuilder.html) is intended for that."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@nb.jit\n",
    "def Zcandidates(events, builder):\n",
    "    for event in events:\n",
    "\n",
    "        # Begin a list in the output, like printing \"[\"\n",
    "        builder.begin_list()\n",
    "\n",
    "        for i in range(len(event.muons)):\n",
    "            for j in range(i + 1, len(event.muons)):\n",
    "                m1 = event.muons[i]\n",
    "                m2 = event.muons[j]\n",
    "                mass = np.sqrt(2*m1.pt*m2.pt*(np.cosh(m1.eta - m2.eta) - np.cos(m1.phi - m2.phi)))\n",
    "                if 60 < mass and mass < 120:\n",
    "\n",
    "                    # Begin a record in the output, like printing \"{\"\n",
    "                    builder.begin_record()\n",
    "\n",
    "                    builder.field(\"mass\")\n",
    "                    builder.append(mass)\n",
    "\n",
    "                    builder.field(\"muon1\")\n",
    "                    builder.append(m1)\n",
    "\n",
    "                    builder.field(\"muon2\")\n",
    "                    builder.append(m2)\n",
    "\n",
    "                    # End a record in the output, like printing \"}\"\n",
    "                    builder.end_record()\n",
    "\n",
    "        # End a list in the output, like printing \"]\"\n",
    "        builder.end_list()\n",
    "\n",
    "    return builder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "builder = Zcandidates(cms_events, ak.ArrayBuilder())\n",
    "builder"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This ak.ArrayBuilder is not an array, but it can quickly (i.e. zero-copy) become one by calling `snapshot`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "builder.snapshot()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This array has a lot of structure: each element is a list of records with a `mass` (float), `muon1` (another record), and `muon2` (another record)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ak.type(builder.snapshot())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ak.to_list(builder.snapshot()[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "An ArrayBuilder makes arrays whose type is determined by the _order in which its methods are called_:\n",
    "\n",
    "```\n",
    "b = ak.ArrayBuilder()\n",
    "\n",
    "# fill commands   # as JSON   # current array type\n",
    "##########################################################################################\n",
    "b.begin_list()    # [         # 0 * var * unknown     (initially, the type is unknown)\n",
    "b.integer(1)      #   1,      # 0 * var * int64\n",
    "b.integer(2)      #   2,      # 0 * var * int64\n",
    "b.real(3)         #   3.0     # 0 * var * float64     (all the integers have become floats)\n",
    "b.end_list()      # ],        # 1 * var * float64\n",
    "b.begin_list()    # [         # 1 * var * float64\n",
    "b.end_list()      # ],        # 2 * var * float64\n",
    "b.begin_list()    # [         # 2 * var * float64\n",
    "b.integer(4)      #   4,      # 2 * var * float64\n",
    "b.null()          #   null,   # 2 * var * ?float64    (now the floats are nullable)\n",
    "b.integer(5)      #   5       # 2 * var * ?float64\n",
    "b.end_list()      # ],        # 3 * var * ?float64\n",
    "b.begin_list()    # [         # 3 * var * ?float64\n",
    "b.begin_record()  #   {       # 3 * var * ?union[float64, {}]\n",
    "b.field(\"x\")      #     \"x\":  # 3 * var * ?union[float64, {\"x\": unknown}]\n",
    "b.integer(1)      #      1,   # 3 * var * ?union[float64, {\"x\": int64}]\n",
    "b.field(\"y\")      #      \"y\": # 3 * var * ?union[float64, {\"x\": int64, \"y\": unknown}]\n",
    "b.begin_list()    #      [    # 3 * var * ?union[float64, {\"x\": int64, \"y\": var * unknown}]\n",
    "b.integer(2)      #        2, # 3 * var * ?union[float64, {\"x\": int64, \"y\": var * int64}]\n",
    "b.integer(3)      #        3  # 3 * var * ?union[float64, {\"x\": int64, \"y\": var * int64}]\n",
    "b.end_list()      #      ]    # 3 * var * ?union[float64, {\"x\": int64, \"y\": var * int64}]\n",
    "b.end_record()    #   }       # 3 * var * ?union[float64, {\"x\": int64, \"y\": var * int64}]\n",
    "b.end_list()      # ]         # 4 * var * ?union[float64, {\"x\": int64, \"y\": var * int64}]\n",
    "```\n",
    "\n",
    "It is a bit like writing JSON from a script, except that it fills columnar data structures, rather than writing text.\n",
    "\n",
    "Thus, it is very important to close each `begin_X` with an `end_X` (at the right place) to avoid making enormous, unclosed data structures! Try your scripts on small samples (outside of Numba) before scaling up.\n",
    "\n",
    "The `append` method is a synonym for `null`, `boolean`, `integer`, `real`, etc., determining its type from its argument. In Numba, that type-determination happens at compile-time (fast). But it can also insert whole data structures, like `muon1` and `muon2` in the Z candidates above. The muons are inserted _by reference_ (not copied)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "One more caveat: since ArrayBuilder is dynamically typed and its length scales dynamically (like `std::vector`), it's slower than filling a NumPy array:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@nb.jit\n",
    "def fill_with_numpy():\n",
    "    out = np.empty(10000000, np.int64)\n",
    "    for i in range(10000000):\n",
    "        out[i] = i\n",
    "    return out\n",
    "\n",
    "@nb.jit\n",
    "def fill_with_builder(builder):\n",
    "    for i in range(10000000):\n",
    "        builder.integer(i)\n",
    "    return builder\n",
    "\n",
    "# run them each once to not count the compilation step in the timing below\n",
    "fill_with_numpy()\n",
    "fill_with_builder(ak.ArrayBuilder());"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit\n",
    "\n",
    "fill_with_numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit\n",
    "\n",
    "fill_with_builder(ak.ArrayBuilder())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "I had to make the workload trivial (just filling with increasing numbers) to see such a dramatic difference, but there's a factor of 10 right there.\n",
    "\n",
    "If what you're doing is moderately complex, then the distinction is hidden by the rest of the workload (e.g. only a factor of 2 for filling with `np.random.normal(0, 1)` and no visible difference when filling Z candidates).\n",
    "\n",
    "So if it's just as easy to fill NumPy arrays and [ak.zip](https://awkward-array.readthedocs.io/en/latest/_auto/ak.zip.html) them or use them as slices later, then do so."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## And now, columnar analysis\n",
    "\n",
    "We can analyze physics data like the way we analyzed the Chicago bike routes: by slicing, projecting, and zipping nested data structures.\n",
    "\n",
    "We could start with a Z peak, the \"hello world\" of particle physics because it has minimal, non-trivial combinatorics.\n",
    "\n",
    "It's actually easier than the imperative code in the Numba examples if you have functions like [ak.cartesian](https://awkward-array.readthedocs.io/en/latest/_auto/ak.cartesian.html) and [ak.combinations](https://awkward-array.readthedocs.io/en/latest/_auto/ak.combinations.html)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Cartesian (cross) products\n",
    "\n",
    "[ak.cartesian](https://awkward-array.readthedocs.io/en/latest/_auto/ak.cartesian.html) makes all possible pairs (or triples, etc.) from two (or three, etc.) collections.\n",
    "\n",
    "If `axis=1` (the default), it does so for each list in arrays of lists.\n",
    "\n",
    "![](img/cartoon-cartesian.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ak.cartesian(([[  1,   2, 3], [   ], [4       ]],\n",
    "              [[\"a\", \"b\"   ], [\"c\"], [\"d\", \"e\"]])).tolist()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Combinations\n",
    "\n",
    "[ak.combinations](https://awkward-array.readthedocs.io/en/latest/_auto/ak.combinations.html) makes all possible pairs (or triples, etc.) of a collection with itself, without duplicates.\n",
    "\n",
    "If `axis=1` (the default), it does so for each list in an array of lists.\n",
    "\n",
    "![](img/cartoon-combinations.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ak.combinations([[1, 2, 3, 4], [], [5, 6]], 2).tolist()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Applying this to the Z peak\n",
    "\n",
    "To make a Z peak, we must consider all possible pairs of muons, without duplicates.\n",
    "\n",
    "In the columnar analysis way of thinking, we make those pairs before we've fully decided what we're going to do with them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pairs = ak.combinations(cms_events.muons, 2)\n",
    "pairs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ak.type(pairs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It would be more convenient if the \"left muon\" and \"right muon\" in each pair were in separate arrays.\n",
    "\n",
    "That's an easy operation with projection/zipping/etc."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "m1, m2 = ak.unzip(pairs)\n",
    "m1, m2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can just apply the formula, vectorially."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mass = np.sqrt(2*m1.pt*m2.pt*(np.cosh(m1.eta - m2.eta) - np.cos(m1.phi - m2.phi)))\n",
    "mass"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This array still has event structure (some events have zero pairs, others have several).\n",
    "\n",
    "We can plot it now to see what the data look like, but we have to flatten it. (Note: the default [ak.flatten](https://awkward-array.readthedocs.io/en/latest/_auto/ak.flatten.html) axis is 1.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.hist(ak.flatten(mass), bins=100, range=(60, 120));"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The next step would probably be to filter non-Z dimuons. We do that by making an array of booleans and slicing the array with it, as in NumPy.\n",
    "\n",
    "Note that `and` and `or` don't work for (vectorized) logic on arrays. Just as in NumPy, we're forced to use `&` and `|`, which has a different operator precedence, so I put expressions in parentheses for safety."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mask = (60 < mass) & (mass < 120)\n",
    "mask"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Unlike NumPy slicing, this is a _jagged_ array of booleans, but that's fine for masking a jagged array, as long as the lengths of all sublists match."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mass[mask]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can make a Z record structure, just like the one we built with an ak.ArrayBuilder in Numba."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "z_records = ak.zip({\"mass\": mass[mask], \"muon1\": m1[mask], \"muon2\": m2[mask]})\n",
    "z_records"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ak.type(z_records)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ak.to_list(z_records[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Cuts on particles, cuts on events\n",
    "\n",
    "The jagged array of booleans is an Awkward extension of NumPy's slicing syntax.\n",
    "\n",
    "We use it to make a distinction between filtering particles:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cms_events.muons[cms_events.muons.pt > 50]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "and filtering events:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cms_events[ak.any(cms_events.muons.pt > 50, axis=1)]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To filter particles, the array must be a jagged array of booleans:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cms_events.muons.pt > 50"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "and to filter particles, it must be a flat array of booleans:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ak.any(cms_events.muons.pt > 50, axis=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here, we used a reducer, [ak.any](https://awkward-array.readthedocs.io/en/latest/_auto/ak.any.html), to turn a list of `True`/`False` for each muon into a single `True`/`False` for the event.\n",
    "\n",
    "You can read it as, \"Select CMS events in which any CMS events' muons' $p_T$ is greater than 50.\"\n",
    "\n",
    "The `axis=1` is to apply to the first level of lists-within-the-array."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Cutting vs masking\n",
    "\n",
    "Cuts eliminate events (or particles), which may have two different intentions:\n",
    "\n",
    "   1. to reduce the size of the dataset so that it's more manageable on disk or in memory;\n",
    "   2. to improve statistical significance by eliminating backgrounds.\n",
    "\n",
    "For rare signals, (1) implies (2), though they're usually broken into stages: coarse cuts and fine cuts.\n",
    "\n",
    "Assuming that we're concerned with (2) only, it can be a _problem_ that cutting eliminates data because most array-at-a-time functions assume that their arguments align: element `i` of the first array corresponds to element `i` of the second array.\n",
    "\n",
    "Cuts break that relationship."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# applying the same cut to all arguments is okay\n",
    "ak.zip({\"mass\": mass[mask], \"muon1\": m1[mask], \"muon2\": m2[mask]})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# applying the cut after the operation is okay\n",
    "ak.zip({\"mass\": mass, \"muon1\": m1, \"muon2\": m2})[mask]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "raises-exception"
    ]
   },
   "outputs": [],
   "source": [
    "# but applying the cut to two out of three arguments is bad\n",
    "ak.zip({\"mass\": mass[mask], \"muon1\": m1[mask], \"muon2\": m2})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Worse still, applying _nearly the same, but not quite_ cuts is terrible.\n",
    "\n",
    "(It would be terrible in imperative code, too, but without an error message.)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Masking\n",
    "\n",
    "One of the data types that Awkward Array provides allows values to be `None` (the \"option\" type, represented by `?` or `option[XYZ]`).\n",
    "\n",
    "If we `mask`, rather than cut, arrays and lists within arrays maintain their lengths, but `None` values are inserted where bad/background values were."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# cutting particles\n",
    "cms_events.muons[cms_events.muons.pt > 50]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# masking particles\n",
    "cms_events.muons.mask[cms_events.muons.pt > 50]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# cutting events\n",
    "cms_events[ak.any(cms_events.muons.pt > 50, axis=1)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# masking events\n",
    "cms_events.mask[ak.any(cms_events.muons.pt > 50, axis=1)]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "All of the Awkward operations are aware of option types and have rules for how to handle them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "masked_muons = cms_events.muons.mask[cms_events.muons.pt > 50]\n",
    "\n",
    "masked_muons[\"pz\"] = masked_muons.pt * np.sinh(masked_muons.eta)\n",
    "masked_muons.pz"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Including Numba."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@nb.njit\n",
    "def find_first_high_pt(events, threshold):\n",
    "    for event in events:\n",
    "        for muon in event:\n",
    "            if muon is None:\n",
    "                pass\n",
    "            elif muon.pt > threshold:\n",
    "                return muon\n",
    "    return None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ak.to_list(find_first_high_pt(masked_muons, 1000))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Pandas!\n",
    "\n",
    "Awkward Arrays can be columns in a Pandas DataFrame. (It was possible in Awkward 0, but made more general now.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "\n",
    "df = pd.DataFrame({\"data\": columnar})\n",
    "df"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "But it's not clear that this is the most useful form, since Pandas doesn't have any operations that know what to do with Awkward Arrays.\n",
    "\n",
    "NumPy functions, at least, pass through."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df + 100"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It is perhaps more useful to \"explode\" an Awkward Array into a DataFrame, such that the DataFrame cells are simple numbers. Then, at least, Pandas operations will know what to do with the data.\n",
    "\n",
    "This is the [ak.pandas.df](https://awkward-array.readthedocs.io/en/latest/ak.pandas.df.html) function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ak.pandas.df(cms_events)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now the event numbers (\"entry\") are distinguished from the particle numbers (\"subentry\"), and the nesting of records becomes hierarchical columns.\n",
    "\n",
    "However, this isn't an exact representation:\n",
    "\n",
    "   * event-level quantities (like PV) are duplicated for each particle;\n",
    "   * events with zero particles drop event-level quantities entirely;\n",
    "   * two particle types with different jaggedness (i.e. different numbers of jets and muons in each event) don't have a natural mapping onto the same \"subentry\" numbers, though you can force it with the `how` parameter (tells Pandas how to JOIN different multiplicities, as in the previous Uproot tutorial)."
   ]
  }
 ],
 "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.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}