tutorial.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Welcome to PyHEP 2020!\n",
    "\n",
    "<br><br><br>\n",
    "\n",
    "Before writing this tutorial, I took a look at the survey..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas\n",
    "df = pandas.read_csv(\"survey-results.csv\")\n",
    "df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df[\"Professional life: What best describes your occupation?\"].value_counts(ascending=True).plot.barh();"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df[\"Professional life: What best describes the stage of your professional career?\"].value_counts(ascending=True).plot.barh();"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "languages = [\n",
    "    \"C or C++\",\n",
    "    \"Python\",\n",
    "    \"Matlab\",\n",
    "    \"Javascript or other browser-based (e.g. TypeScript, CoffeeScript)\",\n",
    "    \"Verilog, VHDL, or other hardware description language\",\n",
    "    \"R\",\n",
    "    \"Java or other JVM-based (e.g. Kotlin, Scala, Clojure)\",\n",
    "    \"Perl\",\n",
    "    \"PHP\",\n",
    "    \"C#\",\n",
    "    \"Julia\",\n",
    "    \"Go\",\n",
    "    \"Swift\",\n",
    "    \"Rust\",\n",
    "    \"Ruby\",\n",
    "    \"Haskell\",\n",
    "    \"Raw assembly or machine code\",\n",
    "    \"Other, not listed above\",\n",
    "]\n",
    "def explode(responses):\n",
    "    responses = [response.strip() for response in responses.split(\";\")]\n",
    "    return [1.0 if language in responses else 0.0 for language in languages]\n",
    "exploded = df[[\"Computing and programming: Which of the following languages do you use regularly (i.e. more than 10% of your work)?\"]].fillna(\"\").applymap(explode)\n",
    "indicator = pandas.DataFrame(exploded.iloc[:, 0].tolist(), columns=languages)\n",
    "indicator.div(indicator.sum(axis=1), axis=0).sum(axis=0).iloc[::-1].plot.barh(figsize=(5, 7));"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df[[\n",
    "    \"Computing and programming: Do you *expect* to use Python more or less in the future (as a fraction of your programming time)?\",\n",
    "    \"Computing and programming: Do you *want* to use Python more or less in the future (as a fraction of your programming time)?\"\n",
    "]].apply(pandas.Series.value_counts).loc[[\"Less\", \"About the same\", \"More\", \"Don't know\"]].plot.bar(rot=0).legend(bbox_to_anchor=(1.2, 0.5));"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cols = {x: x.split(\":\")[1].strip() for x in df.columns if x.startswith(\"Python ecosystem:\") and \"?\" not in x}\n",
    "order = ((df[list(cols)] == \"Don't know what it is\") | (df[list(cols)] == \"Never\")).sum(axis=0).sort_values(ascending=False).index.tolist()\n",
    "pkgs = df[order].rename(columns=cols).apply(pandas.Series.value_counts).T[[\n",
    "    \"Don't know what it is\", \"Never\", \"Through dependencies only\", \"Regularly\", \"All the time\"\n",
    "]].fillna(0)\n",
    "pkgs.insert(0, \"No selection\", pkgs.sum(axis=1).max() - pkgs.sum(axis=1))\n",
    "pkgs.plot.barh(stacked=True, width=0.9, figsize=(20, 20), color=[\"#5e79e0\", \"#798bd1\", \"#992cc7\", \"#f5f518\", \"#ffa640\", \"#ff5a30\"]).legend(bbox_to_anchor=(1.2, 0.5));"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cols = {x: x.split(\":\")[1].strip() for x in df.columns if x.startswith(\"Particle physics ecosystem:\") and \"?\" not in x}\n",
    "order = ((df[list(cols)] == \"Don't know what it is\") | (df[list(cols)] == \"Never\")).sum(axis=0).sort_values(ascending=False).index.tolist()\n",
    "pkgs = df[order].rename(columns=cols).apply(pandas.Series.value_counts).T[[\n",
    "    \"Don't know what it is\", \"Never\", \"Through dependencies only\", \"Regularly\", \"All the time\"\n",
    "]].fillna(0)\n",
    "pkgs.insert(0, \"No selection\", pkgs.sum(axis=1).max() - pkgs.sum(axis=1))\n",
    "pkgs.plot.barh(stacked=True, width=0.9, figsize=(20, 20), color=[\"#5e79e0\", \"#798bd1\", \"#992cc7\", \"#f5f518\", \"#ffa640\", \"#ff5a30\"]).legend(bbox_to_anchor=(1.2, 0.5));"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hopes = [\n",
    "    \"Particle physics analysis tools (other than ROOT)\",\n",
    "    \"General-purpose data analysis toolkits\",\n",
    "    \"Machine learning/deep learning toolkits\",\n",
    "    \"Software engineering skills (beyond the fundamentals)\",\n",
    "    \"ROOT and PyROOT\",\n",
    "    \"Python fundamentals (how to program in Python)\",\n",
    "    \"Collaboration-specific topics\",\n",
    "    \"Other\",\n",
    "]\n",
    "def explode(responses):\n",
    "    responses = [response.strip() for response in responses.split(\";\")]\n",
    "    return [1.0 if hope in responses else 0.0 for hope in hopes]\n",
    "exploded = df[[\"PyHEP feedback: What are you hoping to learn from this workshop?\"]].fillna(\"\").applymap(explode)\n",
    "indicator = pandas.DataFrame(exploded.iloc[:, 0].tolist(), columns=hopes)\n",
    "indicator.div(indicator.sum(axis=1), axis=0).sum(axis=0).iloc[::-1].plot.barh(figsize=(5, 7));"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br><br><br><br><br>\n",
    "\n",
    "## Conclusions:\n",
    "\n",
    "   1. You are mostly grad students and postdocs engaged in physics research.\n",
    "   2. You use Python and C++ about equally, but want to use Python more.\n",
    "   3. You are familiar with the major libraries of the Python world: NumPy, Matplotlib, machine learning.\n",
    "   4. You are less familiar with Python libraries intended for physics analysis.\n",
    "   5. But you want to learn.\n",
    "\n",
    "So let's get started!\n",
    "\n",
    "<br><br><br><br><br>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<img src=\"img/uproot-logo-300px.png\" alt=\"Uproot\" width=\"300px\" style=\"margin-bottom: -50px; margin-right: 20px\"><font size=\"5\"> is a pure-Python implementation of ROOT I/O.</font>\n",
    "\n",
    "<br><br>\n",
    "\n",
    "<img src=\"img/abstraction-layers.png\" width=\"900px\">"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<img src=\"img/awkward-logo-600px.png\" alt=\"Uproot\" width=\"350px\" style=\"margin-bottom: -29px; margin-right: 20px\"><font size=\"5\"> is a generalization of NumPy to data structures (such as jagged arrays).</font>\n",
    "\n",
    "<br><br>\n",
    "\n",
    "<img src=\"img/cartoon-schematic.png\" width=\"1000px\">"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br><br><br>\n",
    "\n",
    "# Interesting times!\n",
    "\n",
    "<font size=\"4\">We're in in the middle of a transition from <b>Uproot 3.x → Uproot 4.x</b> and <b>Awkward 0.x → Awkward 1.x</b>.</font>\n",
    "\n",
    "<img src=\"img/uproot-awkward-timeline.png\" width=\"900px\">\n",
    "\n",
    "<font size=\"4\">You can use both! Old and new versions are independently installable/importable.</font>\n",
    "\n",
    "<table style=\"font-size: 1.5em; font-weight: bold; margin-left: 0px\">\n",
    "    <tr style=\"background: white\"><td></td><td style=\"color: gray\">Now</td><td style=\"color: gray\">Later this year</td></tr>\n",
    "    <tr style=\"background: white\"><td style=\"color: gray\">Old versions</td><td style=\"color: blue\">uproot, awkward</td><td>uproot3, awkward0</td></tr>\n",
    "    <tr style=\"background: white\"><td style=\"color: gray\">New versions</td><td>uproot4, awkward1</td><td style=\"color: blue\">uproot, awkward</td></tr>\n",
    "</table>\n",
    "\n",
    "<img src=\"img/Raiders-of-the-Lost-Ark-Chamber.jpg\" width=\"800px\">\n",
    "\n",
    "<br><br><br><br><br>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# What will this tutorial use?\n",
    "\n",
    "New versions of both: <b>Uproot 4</b> and <b>Awkward 1</b>. <font color=\"red\">This tutorial is bleeding edge.</font>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import uproot4\n",
    "import awkward1 as ak\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# In a nutshell\n",
    "\n",
    "Uproot provides a short path from ROOT files to arrays."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.array(uproot4.open(\"data/opendata_muons.root:Events/nMuon\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's break that down."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Exploring a file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "root = uproot4.open(\"data/opendata_muons.root\")\n",
    "root"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "When you open a file, you get its root directory, which has the properties of a Python dict.\n",
    "\n",
    "You can list its keys."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "root.keys()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can get an item from it using square brackets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "root[\"Events\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "(The `;1` wasn't necesssary—it's a \"cycle number,\" which ROOT uses to distinguish objects in the same directory with the same name. If unspecified, you get the highest cycle number.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections.abc import Mapping\n",
    "\n",
    "isinstance(root, Mapping)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can also get listings of objects by type."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "root.classnames()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Perhaps this file is more interesting:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nesteddirs = uproot4.open(\"data/nesteddirs.root\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nesteddirs.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nesteddirs[\"one/two\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nesteddirs[\"one/two\"].keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nesteddirs.classnames()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "At all levels, you can filter by name or type."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nesteddirs.keys(filter_classname=\"TTree\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nesteddirs.classnames(filter_name=\"*t*\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Histograms!\n",
    "\n",
    "Uproot can read histograms (as well as most other ROOT objects), but it doesn't deal directly with them. The first thing that you do when extracting a histogram is to convert it to another library."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "histograms = uproot4.open(\"data/hepdata-example.root\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "histograms.classnames()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "histograms[\"hpx\"].to_boost()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is a [boost-histogram](https://github.com/scikit-hep/boost-histogram), a clean design of N-dimensional histograms in the [Boost](https://www.boost.org/doc/libs/release/libs/histogram/doc/html/index.html) C++ library (with Python bindings). Boost-histogram focuses just on **filling and manipulating (e.g. slicing)** histograms.\n",
    "\n",
    "But we want to plot it, right? There's another library, [mplhep](https://github.com/scikit-hep/mplhep), which focuses just on **plotting** histograms in Matplotlib.\n",
    "\n",
    "<table style=\"margin-left: 0px\">\n",
    "    <tr style=\"background: white\">\n",
    "        <td><img src=\"img/BoostHistogramCppLogo.png\" width=\"300px\" style=\"margin-right: 20px\"></td>\n",
    "        <td><img src=\"img/BoostHistogramPythonLogo.png\" width=\"300px\" style=\"margin-right: 20px\"></td>\n",
    "        <td><img src=\"img/mplhep.png\" width=\"300px\"></td>\n",
    "    </tr>\n",
    "</table>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import mplhep\n",
    "\n",
    "mplhep.histplot(histograms[\"hpx\"].to_boost())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.style.use(mplhep.style.CMS)\n",
    "mplhep.histplot(histograms[\"hpx\"].to_boost())\n",
    "mplhep.cms.label()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.style.use(mplhep.style.ATLAS)\n",
    "mplhep.histplot(histograms[\"hpx\"].to_boost())\n",
    "mplhep.atlas.label()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mplhep.hist2dplot(histograms[\"hpxpy\"].to_boost());"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## pip install hist\n",
    "\n",
    "The [hist](https://github.com/scikit-hep/hist) project is being developed to pull all of this into a common interface.\n",
    "\n",
    "(Note: currently, you have to `pip install 'hist>=2.0.0a2'` because it has only been released as an alpha version.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "histograms[\"hpx\"].to_hist().plot();"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "histograms[\"hpxpy\"].to_hist().plot();"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "histograms[\"hpxpy\"].to_hist().plot2d_full();"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from uncertainties import unumpy as unp\n",
    "\n",
    "def pdf(x, a=1/np.sqrt(2*np.pi), x0=0, sigma=1, offset=0):\n",
    "    exp = unp.exp if a.dtype == np.dtype(\"O\") else np.exp\n",
    "    return a * exp(-(x-x0)**2/(2*sigma**2)) + offset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "histograms[\"hpx\"].to_hist().plot_pull(pdf);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "See the [hist examples](https://github.com/scikit-hep/hist/tree/master/notebooks) on how to do this more beautifully/elegantly."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Exploring a TTree\n",
    "\n",
    "<img src=\"img/terminology.png\" width=\"1000px\">"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It's generally useful to first look at a TTree with `show`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tree = root[\"Events\"]\n",
    "tree.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "These are all the branches of the TTree with the type name of the branch (if Uproot can determine it) and its interpretation as an array (if possible).\n",
    "\n",
    "TTrees also have a dict-like interface."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tree.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tree.items()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tree.typenames()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "isinstance(tree, Mapping)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tree.keys(filter_name=\"Muon_*\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tree.keys(filter_typename=\"float[]\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tree.keys(filter_branch=lambda branch: not isinstance(branch.interpretation, uproot4.AsJagged))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Turning branches into arrays\n",
    "\n",
    "If a branch has a known interpretation, you can call `array` on it to get an array."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tree[\"Muon_pt\"].array()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First thing to notice: this is not a NumPy array. It's because the data have different numbers of values in each element (a jagged array)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tree[\"Muon_pt\"].array()[:20].tolist()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can (in Uproot 4) _force_ it to be a NumPy array, but it isn't pretty:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tree[\"Muon_pt\"].array(library=\"np\")[:20]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The data type (`dtype`) of this NumPy array is `object`, meaning that each element it contains is a Python object, namely another NumPy array.\n",
    "\n",
    "The default is for all arrays to be Awkward arrays, but you can override this by specifying `library`.\n",
    "\n",
    "The difference is that Awkward arrays interpret nested lists as a second dimension, whereas NumPy object arrays do not:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "awkward_array = tree[\"Muon_pt\"].array(library=\"ak\")\n",
    "numpy_array = tree[\"Muon_pt\"].array(library=\"np\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# from the first 20 events, get the first item\n",
    "awkward_array[:20, 0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "raises-exception"
    ]
   },
   "outputs": [],
   "source": [
    "# doesn't work with NumPy object arrays because contents are not guaranteed to be arrays\n",
    "numpy_array[:20, 0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Another valid library is Pandas. Pandas has its own way of describing variable length structures (`MultiIndex`)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tree[\"Muon_pt\"].array(library=\"pd\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the original one-liner, we used colon (`:`) to separate a file path/URL from an object path to get to the branch:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "uproot4.open(\"data/opendata_muons.root:Events/nMuon\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And \"cast\" the branch as a NumPy array, which is the same as calling `array` with `library=\"np\"`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.array(uproot4.open(\"data/opendata_muons.root:Events/nMuon\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This can be useful if you're passing the branch to a library that expects an array. _(Warning: only do this with non-jagged arrays!)_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "plt.hist(uproot4.open(\"data/opendata_muons.root:Events/nMuon\"), bins=11, range=(-0.5, 10.5));"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Pluralization\n",
    "\n",
    "If you look carefully, you'll notice that there's an `array` function and an `arrays` function. The latter gets multiple arrays."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# NumPy arrays in a dict\n",
    "pv_numpy = tree.arrays(filter_name=\"PV_*\", library=\"np\")\n",
    "pv_numpy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Awkward record-array\n",
    "pv_awkward = tree.arrays(filter_name=\"PV_*\", library=\"ak\")\n",
    "pv_awkward"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Pandas DataFrame (as opposed to Series for a single array)\n",
    "pv_pandas = tree.arrays(filter_name=\"PV_*\", library=\"pd\")\n",
    "pv_pandas"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "All of these are \"packages\" of arrays that you might use in your analysis."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pv_numpy[\"PV_x\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pv_awkward[\"PV_x\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pv_awkward.PV_x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pv_pandas[\"PV_x\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pv_pandas.PV_x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Above, we used `filter_name` to select branches that match a pattern. We can also request specific branches:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tree.arrays([\"PV_x\", \"PV_y\", \"PV_z\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Or do calculations. (This feature exists for TTree aliases, which can be formulas.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tree.arrays(\"PV\", aliases={\"PV\": \"sqrt(PV_x**2 + PV_y**2 + PV_z**2)\"})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tree.arrays(\"PV\", cut=\"sqrt(PV_x**2 + PV_y**2) < 0.1\", aliases={\"PV\": \"sqrt(PV_x**2 + PV_y**2 + PV_z**2)\"})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Multiple files\n",
    "\n",
    "Typically, you'll have a lot of files with similar contents.\n",
    "\n",
    "## Concatenation\n",
    "\n",
    "The simplest way to deal with this is to read a selection of branches entirely into memory, concatenating them.\n",
    "\n",
    "If you have enough memory, go for it!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "all_in_memory = uproot4.concatenate(\"data/uproot-sample-*.root:sample\", [\"i4\", \"ai8\", \"Af8\", \"str\"])\n",
    "all_in_memory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "all_in_memory.i4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "all_in_memory.ai8"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "all_in_memory.Af8"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "all_in_memory.str"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "But, often enough, you don't have enough memory. What then?\n",
    "\n",
    "## Laziness\n",
    "\n",
    "One option is to open them as lazy arrays, which opens the files (to get the number of events in each), but doesn't read the data until you use it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "not_in_memory_yet = uproot4.lazy(\"data/uproot-sample-*.root:sample\")\n",
    "not_in_memory_yet"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\"If it's not in memory, how can I see the values?\"\n",
    "\n",
    "Only the parts of the files (branches and batches of events) that are visible are read. In the above, `n` and `b` from the first file and `str` from the last file must have been read.\n",
    "\n",
    "Let's get the `Af8` field:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "not_in_memory_yet.Af8"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Again, this only read from the first and last files to show the first and last values.\n",
    "\n",
    "A mathematical operation would cause them all to be read in."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "not_in_memory_yet.Af8 + 100"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Controlling the lazy cache\n",
    "\n",
    "After (part of) a lazy array has been read, how long does it stay in memory? Is it constantly being re-read every time we do a calculation?\n",
    "\n",
    "By default, a 100 MB cache is associated with the lazy array, but we can provide our own if we want a bigger or smaller one."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cache = uproot4.LRUArrayCache(\"1 GB\")\n",
    "\n",
    "not_in_memory_yet = uproot4.lazy(\"data/uproot-sample-*.root:sample\", array_cache=cache)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cache"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "not_in_memory_yet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cache"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "not_in_memory_yet + 100"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cache"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "What's more, we can clear it when we need to."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cache.clear()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cache"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Iteration\n",
    "\n",
    "But often, that's still not enough control.\n",
    "\n",
    "We don't read arrays into memory for the fun of it, we do it to perform calculations, and lazy arrays don't control which parts of which arrays are in memory during a calculation.\n",
    "\n",
    "If you're worried about memory, the safest thing to do is to iterate over the data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for arrays in uproot4.iterate(\"data/uproot-sample-*.root:sample\", [\"i4\", \"Af8\"]):\n",
    "    print(arrays[\"i4\"] + arrays[\"Af8\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This iteration is _not_ one event at a time. (This set of TTrees has 420 entries.) It's a _chunk of events_ at a time.\n",
    "\n",
    "In each step, a chunk of events for all specified arrays (`[\"i4\", \"Af8\"]`) is read. You do your calculation, move on to the next step, and all the previous arrays are dropped. (Only TBasket data carries over if event steps don't line up with TBasket boundaries—a low-level detail.)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Weird objects\n",
    "\n",
    "Although Uproot is primarily for TTrees with simple types, it is possible to read some complex types."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "weird = uproot4.open(\"data/uproot-stl_containers.root:tree\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "weird.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.array(weird[\"map_string_vector_string\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "HZZ = uproot4.open(\"data/uproot-HZZ-objects.root:events\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "HZZ.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "HZZ[\"muonp4\"].array(library=\"pd\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Sometimes, errors that people report as \"Uproot can't read type X\" don't have anything to do with type X.\n",
    "\n",
    "To help zero in on actual deserialization errors, a lot more diagnostic information has been included. In Uproot 4, this is what a deserialization error looks like (encountered last Thursday):\n",
    "\n",
    "```\n",
    "    TH1F version 2 as <dynamic>.Model_TH1F_v2 (939 bytes)\n",
    "        TH1 version 7 as <dynamic>.Model_TH1_v7 (893 bytes)\n",
    "            (base): <TNamed 'cutflow' title='dijethad' at 0x7fafb4505f90>\n",
    "            (base): <TAttLine (version 2) at 0x7fafb4506350>\n",
    "            (base): <TAttFill (version 2) at 0x7fafb4506390>\n",
    "            (base): <TAttMarker (version 2) at 0x7fafb4506310>\n",
    "            fNcells: 9\n",
    "            TAxis version 9 as <dynamic>.Model_TAxis_v9 (417 bytes)\n",
    "                (base): <TNamed 'xaxis' at 0x7fafb4506890>\n",
    "                (base): <TAttAxis (version 4) at 0x7fafb4506950>\n",
    "                fNbins: 7\n",
    "                fXmin: 0.0\n",
    "                fXmax: 7.0\n",
    "                fXbins: <TArrayD [] at 0x7fafb4506910>\n",
    "                fFirst: 0\n",
    "                fLast: 0\n",
    "                fBits2: 4\n",
    "                fTimeDisplay: False\n",
    "                fTimeFormat: <TString '' at 0x7fafb44dc1d0>\n",
    "                THashList version 5 as <dynamic>.Model_THashList_v0 (294 bytes)\n",
    "                    TList version 1 as uproot4.models.TList.Model_TList (? bytes)\n",
    "                        (base): <TObject None None at 0x7fafb4506fd0>\n",
    "                        fName: ''\n",
    "                        fSize: 475136\n",
    "\n",
    "attempting to get bytes 1851028560:1851028561\n",
    "outside expected range 0:939 for this Chunk\n",
    "in file /home/pivarski/miniconda3/lib/python3.7/site-packages/skhep_testdata/data/uproot-issue33.root\n",
    "in object /cutflow\n",
    "```\n",
    "\n",
    "The values get wonky in the THashList (`fSize: 475136`), which led directly to the issue."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# To split or not to split\n",
    "\n",
    "\"Splitting branches\" is a technical detail in ROOT, but very important in Uproot, since Uproot views branches as arrays.\n",
    "\n",
    "Here is an \"unsplit\" file:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "unsplit = uproot4.open(\"data/uproot-small-evnt-tree-nosplit.root:tree\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "unsplit.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And here is a \"split\" file:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "split = uproot4.open(\"data/uproot-small-evnt-tree-fullsplit.root:tree\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "split.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It's the same data, and they look the same in ROOT. But since a branch is a unit of what can be read into an array, we have to read whole objects or nothing."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "unsplit[\"evt\"].array()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "But with the split file, we can choose which parts to read."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "split.arrays([\"P3.Px\", \"P3.Py\", \"P3.Pz\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If there are any types in an unsplit object that Uproot can't deserialize, Uproot won't be able to read _any_ of its data.\n",
    "\n",
    "If there are any types in a split object that Uproot can't deserialize, Uproot will be able to read everything else.\n",
    "\n",
    "Also, split objects are usually faster to read, even if you read everything.\n",
    "\n",
    "**Split your data whenever possible!**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Writing to ROOT files\n",
    "\n",
    "Uproot 4 cannot write to ROOT files yet: see Uproot 3 documentation.\n",
    "\n",
    "Some caveats, though:\n",
    "\n",
    "   * Writing ROOT files with Uproot will always be minimal: histograms and only basic types in TTrees.\n",
    "   * You won't be able to update an existing file, only make new files.\n",
    "   * It won't be as fast as writing with ROOT.\n",
    "\n",
    "The issues involved in _writing_ an established format are considerably different from _reading_. If anyone thinks they can do better, they're welcome to try!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br><br><br>\n",
    "\n",
    "<img src=\"img/awkward-logo-600px.png\" alt=\"Uproot\" width=\"350px\">\n",
    "\n",
    "<br>\n",
    "\n",
    "Why do we need our own array library? It's a long story, but if you're interested, I presented it to non-physicists here:\n",
    "\n",
    "<a href=\"https://youtu.be/WlnUF3LRBj4\"><img src=\"img/my-video.png\"></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's take another look at the muon data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tree = uproot4.open(\"data/opendata_muons.root:Events\")\n",
    "tree.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We've already seen that it's \"awkward\" to deal with the jagged arrays in NumPy. However, they look and feel like records if \"zipped\" into an Awkward array."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "events = tree.arrays(library=\"ak\", how=\"zip\")\n",
    "events"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The data type encapsulates the structure of the events."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ak.type(events)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`1000000 *` means that there are a million events, `\"Muon\": var *` means that the contents of the `\"Muon\"` field are jagged: there's a variable number of them per event.\n",
    "\n",
    "We could look at a few of these as Python lists and dicts."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ak.to_list(events[:3])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "But the data are not actually arranged as objects in memory; each field (`\"pt\"`, `\"eta\"`, `\"phi\"`, etc.) is in an array by itself.\n",
    "\n",
    "This means that structure-changing things like pulling out the kinematics are not expensive computations. (That is, they do not scale with the size of the dataset.)\n",
    "\n",
    "You can project out and remix them without penalty."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "events[\"Muon\", [\"pt\", \"eta\", \"phi\"]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ak.type(events[\"Muon\", [\"pt\", \"eta\", \"phi\"]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "events[\"Muon\", \"pz\"] = events[\"Muon\", \"pt\"] * np.sinh(events[\"Muon\", \"eta\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ak.type(events.Muon)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "events.Muon.pz"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Since nearly all of you are familiar with NumPy, slicing arrays with boolean arrays is probably familiar to you.\n",
    "\n",
    "What's new is that the boolean arrays can now be jagged to slice jagged arrays (i.e. cut particles)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "events.Muon.pt > 20"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "events.Muon[events.Muon.pt > 20]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To cut events, make the jagged array of booleans a one-dimensional array of booleans. You can do this with a reducer (such as [ak.sum](https://awkward-array.readthedocs.io/en/latest/_auto/ak.sum.html) or [ak.max](https://awkward-array.readthedocs.io/en/latest/_auto/ak.max.html), but most likely [ak.any](https://awkward-array.readthedocs.io/en/latest/_auto/ak.any.html) and [ak.all](https://awkward-array.readthedocs.io/en/latest/_auto/ak.all.html) for booleans)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ak.any(events.Muon.pt > 20, axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "events.Muon[ak.any(events.Muon.pt > 20, axis=1)]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Awkward analysis\n",
    "\n",
    "Several new operations are needed when arrays can have arbitrary data structures, so the Awkward Array library is best seen as a collection of functions acting on [ak.Array](https://awkward-array.readthedocs.io/en/latest/_auto/ak.Array.html) (the array type).\n",
    "\n",
    "Probably the most important of these is [ak.num](https://awkward-array.readthedocs.io/en/latest/_auto/ak.num.html), which tells us the number of elements in each nested list."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ak.num(events.Muon)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "That's how many muons there are in each event.\n",
    "\n",
    "This becomes necessary if we ever try to select the first (and second, etc.) element in each event. Some events might not have any."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "raises-exception"
    ]
   },
   "outputs": [],
   "source": [
    "events.Muon[:, 0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So we use [ak.num](https://awkward-array.readthedocs.io/en/latest/_auto/ak.num.html) to slice the first dimension."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "events.Muon[ak.num(events.Muon) > 0, 0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Masking vs cutting\n",
    "\n",
    "In the nearly two years that physicists have been doing analyses with Awkward Arrays, they've found that cuts are difficult to accumulate. If the first cut changes the length of the array from 1000000 to 969031, boolean arrays that could be applied to the first array can't be applied to the second array.\n",
    "\n",
    "In practice, they've taken a logical-and of all cuts and apply them at the end, but we can do better: we can mask, rather than cut."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "events.Muon.mask[events.Muon.pt > 20]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "One of the new types that Awkward Array introduces is the \"option type,\" which allows some values to be `None`. (It's a `?` in the type specification.)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Making regular arrays\n",
    "\n",
    "If you're feeding your data into a machine learning pipeline, you might need the jaggedness to go away. There are functions for padding jagged arrays (with `None`) so that they reach a desired length (and replacing `None` with a preferred value): [ak.pad_none](https://awkward-array.readthedocs.io/en/latest/_auto/ak.pad_none.html) (and [ak.fill_none](https://awkward-array.readthedocs.io/en/latest/_auto/ak.fill_none.html))."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ak.pad_none(events.Muon.pt, 3, clip=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ak.fill_none(ak.pad_none(events.Muon.pt, 3, clip=True), 0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "When you're plotting something, you usually want to flatten the jaggedness."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.hist(ak.flatten(events.Muon.pt), bins=120, range=(0, 60));"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Awkward combinatorics\n",
    "\n",
    "We often want to find pairs of particles with some invariant mass. To do that, we need combinatoric 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).\n",
    "\n",
    "<table style=\"margin-left: 0px\">\n",
    "    <tr style=\"background: white\"><td style=\"font-size: 1.75em; font-weight: bold; text-align: center\">Cartesian product (per event)</td><td style=\"font-size: 1.75em; font-weight: bold; text-align: center\">n-choose-k combinations (per event)</td></tr>\n",
    "    <tr style=\"background: white\"><td><img src=\"img/cartoon-cartesian.png\"></td><td><img src=\"img/cartoon-combinations.png\"></td></tr>\n",
    "</table>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "muon_pairs = ak.combinations(events.Muon, 2)\n",
    "muon_pairs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "m1, m2 = ak.unzip(muon_pairs)\n",
    "m1, m2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "masses = np.sqrt(2*m1.pt*m2.pt*(np.cosh(m1.eta - m2.eta) - np.cos(m1.phi - m2.phi)))\n",
    "masses"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.hist(ak.flatten(masses), bins=80, range=(70, 110));"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.hist(ak.flatten(masses), bins=240, range=(0, 12))\n",
    "plt.yscale(\"log\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Numba\n",
    "\n",
    "The most important Python library that most of you aren't aware of is Numba:\n",
    "\n",
    "<a href=\"https://numba.pydata.org/\"><img src=\"img/numba.png\"></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The NumPy-like expressions for combinatorics look weird. It's fine for simple things when you get the hang of it, but really complex or performance-critical array expressions can be hard to write.\n",
    "\n",
    "Numba lets you write Python for loops, but then it compiles them as though they were C functions.\n",
    "\n",
    "Only a subset of Python and NumPy are accepted by its compiler, but Awkward Arrays are included."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numba as nb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@nb.jit\n",
    "def find_energetic_muon(events):\n",
    "    for event in events:\n",
    "        for muon in event.Muon:\n",
    "            if muon.pt > 100000:\n",
    "                return muon\n",
    "    return None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "find_energetic_muon(events).tolist()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Or better yet,"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@nb.jit\n",
    "def find_energetic_muons_event(events):\n",
    "    for i in range(len(events)):\n",
    "        for muon in events[i].Muon:\n",
    "            if muon.pt > 100000:\n",
    "                return i\n",
    "    return -1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "find_energetic_muons_event(events)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "events[7400].tolist()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The invariant mass calculation looks more conventional this way, though it is more verbose, too. You also need two passes to allocate an array of the right size."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@nb.jit\n",
    "def invariant_mass(events):\n",
    "    num_pairs = 0\n",
    "    for event in events:\n",
    "        num_pairs += max(0, len(event.Muon) * (len(event.Muon) - 1) // 2)\n",
    "    out = np.empty(num_pairs, np.float64)\n",
    "    \n",
    "    num_pairs = 0\n",
    "    for event in events:\n",
    "        for i in range(len(event.Muon)):\n",
    "            for j in range(i + 1, len(event.Muon)):\n",
    "                m1 = event.Muon[i]\n",
    "                m2 = event.Muon[j]\n",
    "                out[num_pairs] = np.sqrt(2*m1.pt*m2.pt*(np.cosh(m1.eta - m2.eta) - np.cos(m1.phi - m2.phi)))\n",
    "                num_pairs += 1\n",
    "    \n",
    "    return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "masses = invariant_mass(events)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.hist(masses, bins=80, range=(70, 110));"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.hist(masses, bins=240, range=(0, 12))\n",
    "plt.yscale(\"log\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The important point is that _you don't have to choose_ between Awkward combinatorics and Numba. You can pass the same data structures through `ak.this` and `ak.that` functions as you can pass into Numba-compiled functions.\n",
    "\n",
    "   * Awkward's array-at-a-time functions are best for interactive exploration. It can be hard to express deeply nested loops with [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).\n",
    "   * Numba's just-in-time compilation is best for tricky algorithms and high performance. It can be hard to satisfy Numba's type constraints and stay within the supported [Python subset](https://numba.pydata.org/numba-doc/dev/reference/pysupported.html) and [NumPy subset](https://numba.pydata.org/numba-doc/dev/reference/numpysupported.html).\n",
    "\n",
    "Use them both in your data analysis! Use Numba to make indexes to slice Awkward arrays; use Awkward to prepare structures for Numba.\n",
    "\n",
    "(I haven't even mentioned the fact that Awkward Arrays are [accessible from C++](https://github.com/scikit-hep/awkward-1.0/tree/master/dependent-project) and will [soon run on GPUs](https://github.com/scikit-hep/awkward-1.0/projects/1).)"
   ]
  }
 ],
 "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
}