{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Learning the awkward arrays syntax\n",
    "\n",
    "The core of `coffea` leverages a completely new syntax for expressing analysis\n",
    "computations: `awkward arrays` and its index-based notation. For people coming\n",
    "from a more traditional loop-based programming syntax, the syntax will take\n",
    "some getting use to, but this tutorial can hopefully help you understand how to\n",
    "understand the syntax and how to understand the various method.\n",
    "\n",
    "Let use begin by first understanding what you need to explore the contents of a\n",
    "typical n-tuple file using `coffea` related tools. For the examples below we\n",
    "will be using the [CMS open data](http://opendata.cern.ch/record/12354#).\n",
    "Either download the file to your working directory to follow along.\n",
    "\n",
    "First import the relevant `coffea` objects:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from coffea.nanoevents import NanoEventsFactory\n",
    "from coffea.nanoevents.schemas import  NanoAODSchema\n",
    "import numpy as np\n",
    "import awkward1 as ak"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can create the event list as an awkward array using `coffea` tools like:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "events = NanoEventsFactory.from_root(\n",
    "    'file:TTbar.root',    # The file, notice the prefix \"file:\" for local file operation\n",
    "    'Events',             # Name of the tree object to open\n",
    "    entry_stop=50,        # Limit the number of events to process, nice for small scale debugging\n",
    "    schemaclass=NanoAODSchema\n",
    ").events()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The last `schemaclass` argument will be let unexplained for now, see the schema\n",
    "tutorials to learn more about what this is. Here we have created the events as\n",
    "an awkward array. To see that is stored in the array we can use:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['Muon', 'luminosityBlock', 'Jet', 'HLT', 'PV', 'Tau', 'GenPart', 'MET', 'run', 'event']\n"
     ]
    }
   ],
   "source": [
    "print(events.fields)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Indicating the collections that are stored in the awkward array. To see how\n",
    "many events exists in our file we can use the typical python method:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "50\n"
     ]
    }
   ],
   "source": [
    "print(len(events))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The 50 here in corresponds correctly to the `entry_stop` used in to open the\n",
    "file. Next we can, of course, start to explore the contents of the various\n",
    "object collections. One can access the fields of the event as if it was a\n",
    "regular data member"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['pt', 'eta', 'phi', 'mass', 'charge', 'pfRelIso03_all', 'pfRelIso04_all', 'tightId', 'softId', 'dxy', 'dxyErr', 'dz', 'dzErr', 'jetIdx', 'genPartIdx', 'jetIdxG', 'genPartIdxG']\n",
      "['pt', 'eta', 'phi', 'mass', 'puId', 'btag']\n"
     ]
    }
   ],
   "source": [
    "print(events.Muon.fields)\n",
    "print(events.Jet.fields)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Ah-ha! Now, we are starting to see numbers we can play around with. Notice that\n",
    "`coffea` was written with high energy physics analysis in mind, so even if the\n",
    "electron energy doesn't look like it is stored from the output of the fields,\n",
    "we can still access methods that we will typically associate with 4-vectors. In\n",
    "particular, notice that we can call the `energy` field of the electron\n",
    "collection, even though the energy field isn't explicitly defined. `coffea` is\n",
    "designed with 4 vectors in mind, so the energy collection is calculated on the\n",
    "fly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[23.6], [], [], [57.3, 29.3, 7.7], [], [], ... [], [], [57.2], [], [19.7], [], []]\n",
      "[[52.3], [], [], [95.6, 30.2, 8.01], [], [], ... [], [], [309], [], [19.7], [], []]\n"
     ]
    }
   ],
   "source": [
    "print(events.Muon.pt)\n",
    "print(events.Muon.energy)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, looking at the output, we can begin to get a grasp of what awkward arrays\n",
    "are: the variable `events.Muon.pt` variable represents an $N$ events times $A$\n",
    "objects array of floating point, the `events.Muon` variable represents and $N$\n",
    "events times $A$ *objects*, acting as a *collection* of floating point arrays.\n",
    "\n",
    "The \"awkward\" part of awkward array refers to two parts:\n",
    "- The value of `A` is different for each event and for each collection. In this\n",
    "  demonstration, our 0-th event has 1 muon, the 1-st and 2-rd event has no\n",
    "  muons, and the 3-rd event has 3 muons and so on.\n",
    "- Each collection can have a different number of fields. In a sense, the\n",
    "  `events`, `Muon` and `pt` variables are just an easy way for representing the\n",
    "  final `NxA` array that we might be interested in for the analysis.\n",
    "\n",
    "In our case the `N` number of events is what will refer to as the outermost\n",
    "**dimension** or axis of the various objects, `A` is the one inner dimension of\n",
    "the of array. One thing to keep in mind is that we typically refer to the\n",
    "dimension of `events.Muon` and `events.Muon.pt` as both being `NxA`, except\n",
    "`events.Muon` is a `NxA` array of objects (or collection), while the\n",
    "`events.Electron.pt` being a `NxA` array of a concrete data type (in this case\n",
    "a float).\n",
    "\n",
    "We can use the usual index notation to look at a particular object of interest.\n",
    "For instance if we want to look at the 0-th Muon of the 3-rd event in our event\n",
    "list, we can write:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "57.327919006347656 1.0991891622543335\n"
     ]
    }
   ],
   "source": [
    "print(events[3].Muon[0].pt, events[3].Muon[0].eta)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "But the real power of using awkward arrays is for using awkward arrays comes in\n",
    "when you don't explicitly use a concrete index, and instead call calculation all\n",
    "in an abstract form"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Basic object and event selection\n",
    "\n",
    "Let us start with the most basic example of event selection. Say we want to\n",
    "select event with electrons that have $p_T > 30$ GeV and $|\\eta| < 2.0$. The\n",
    "awkward array allows us to write something like:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[False], [], [], [True, False, False], [], ... [], [True], [], [False], [], []]\n",
      "[[True], [], [], [True, True, True], [], [], ... [], [], [False], [], [True], [], []]\n",
      "[[False], [], [], [True, False, False], [], ... [], [False], [], [False], [], []]\n"
     ]
    }
   ],
   "source": [
    "mask_pt = events.Muon.pt > 50\n",
    "mask_eta = np.abs(events.Muon.eta) < 2.0\n",
    "muon_mask = mask_pt & mask_eta\n",
    "print(mask_pt)\n",
    "print(mask_eta)\n",
    "print(muon_mask)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can see that the usual logic comparison operators generate a `NxA` boolean\n",
    "array telling use which electron (or more specifically which electron.pt and\n",
    "electron etas) pass this particular selection criteria. This particular boolean\n",
    "array generated from a logic operation on usual arrays is typically call a\n",
    "`mask`. We can use the typical boolean operation `&` operation to get the\n",
    "intersect of multiple masks, or maybe the `|` operator for the union. Now the\n",
    "problem is where can we use this mask? The answer is any array that has a `NxA`\n",
    "structure and can receive these masks to create a reduced array!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[], [], [], [57.3], [], [], [], [], [], ... [], [], [], [], [], [], [], [], []]\n",
      "[[], [], [], [1.1], [], [], [], [], [], ... [], [], [], [], [], [], [], [], []]\n",
      "[[], [], [], [57.3], [], [], [], [], [], ... [], [], [], [], [], [], [], [], []]\n"
     ]
    }
   ],
   "source": [
    "print(events.Muon.pt[muon_mask])\n",
    "print(events.Muon.eta[muon_mask])\n",
    "selectedMuons = events.Muon[muon_mask]\n",
    "print(selectedMuons.pt)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Probably the most important place to put the mask is directly in the\n",
    "`events.Muon` index, this generates a new collection of electrons that\n",
    "preserves the `NxA` structure, but have various collection instances filtered\n",
    "out. If you are familiar with `numpy`, this sort of index-based array filtering\n",
    "look familiar. The difference is that because awkward arrays accept arrays of\n",
    "varying inner dimensions, it can truly preserve the structure of such\n",
    "selection, rather than having everything be flattened out."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2 4 6 8 2]\n",
      "[2 4 6 8 2]\n",
      "[[2, 4], [6, 8], [2]]\n"
     ]
    }
   ],
   "source": [
    "x = np.array([1,2,3,4,5,6,7,8,1,1,1,2])\n",
    "print( x[x% 2 == 0])\n",
    "y = np.array([[1,2,3,4],[5,6,7,8],[1,1,1,2]])\n",
    "print( y[y%2==0])\n",
    "z = ak.Array([[1,2,3,4],[5,6,7,8],[1,1,1,2]])\n",
    "print(z[z%2==0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now suppose we only want events that have at least 1 electron selected event.\n",
    "What we need are a set of functions that can reduces this `NxA'` array to\n",
    "something of just dimension `N`. Formally this is called **reduction**\n",
    "operations, and the awkward package has a large set of functions that can\n",
    "reduce the dimension of arrays. In our case, what we want is:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[False, False, False, True, False, False, ... False, False, False, False, False]\n"
     ]
    }
   ],
   "source": [
    "muon_count = ak.count(selectedMuons.pt, axis=-1)\n",
    "event_mask = muon_count >= 1\n",
    "print(event_mask)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To break this down, `ak.count`, as the method name suggests \"counts\" the number\n",
    "of elements along a certain axis, in our case, what we are interested is the\n",
    "innermost dimension/axis, hence the typical python notation of `axis=-1`.\n",
    "Using this we can run the event selection using the usual masking notation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[False, False, False, True, False, False, ... False, False, False, False, False]\n",
      "[[23.6], [], [], [57.3, 29.3, 7.7], [], [], ... [], [], [57.2], [], [19.7], [], []]\n",
      "[[57.3, 29.3, 7.7], [52.2]]\n",
      "2\n"
     ]
    }
   ],
   "source": [
    "selectedEvents = events[event_mask]\n",
    "print(event_mask)\n",
    "print(events.Muon.pt)\n",
    "print(selectedEvents.Muon.pt)\n",
    "print(len(selectedEvents))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here we can confirm that the first event to pass the event selection is the\n",
    "1-st event in the event list, and the 0-th instance in the\n",
    "`selectedEvents.Muon.pt` result of the `selectedEvents` indeed corresponds to\n",
    "the values stored in the 1-st event of the original event list."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Object storage and collection creation\n",
    "\n",
    "Having completed the selection, we might be rather annoyed that we didn't just\n",
    "store the selected Electron, since these are the objects that we are likely\n",
    "going to use for further calculation. Following from the code above, what we\n",
    "can do is add the additional selection to the `selectedMuon` collections. This\n",
    "is valid since the `N` dimensional event mask \"makes sense\" performed on the\n",
    "`NxA'` dimensional `selectedMuon` object.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[57.3], [52.2]]\n",
      "2\n"
     ]
    }
   ],
   "source": [
    "our_selectedMuons = selectedMuons[event_mask]\n",
    "print(our_selectedMuons.pt)\n",
    "print(len(our_selectedMuons))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "However, this is rather undesirable, since now we have some a bunch of detected\n",
    "collections, and event lists that we need to take care of: `selectedMuon`,\n",
    "`selectedEvents`, `our_selectedMuons`. And this is with just one toy object\n",
    "selection. One can imagine if there isn't some sort of way to store collections\n",
    "into events, the analysis code will get out of hands very quick. This also ties\n",
    "into the topic that there might be certain physics quantities that are specific\n",
    "to a certain analysis that might be used for the analysis object selection and\n",
    "would be nice to add to the electron collection if it isn't a standard variable\n",
    "that is maintained by the NanoAOD development team. Here we are going to add a\n",
    "very artificial example of calculating the inverse of the electron pt, then\n",
    "selecting on the inverse pt. This very simple example will demonstrate the\n",
    "typical syntax used for storing variables as well as exposing one of the\n",
    "particuliar quirks of awkward arrays:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "First attempt at adding extended variables to events\n",
      "['Muon', 'luminosityBlock', 'Jet', 'HLT', 'PV', 'Tau', 'GenPart', 'MET', 'run', 'event', 'selectedMuon_1']\n",
      "['charge', 'dxy', 'dxyErr', 'dz', 'dzErr', 'eta', 'genPartIdx', 'genPartIdxG', 'jetIdx', 'jetIdxG', 'mass', 'pfRelIso03_all', 'pfRelIso04_all', 'phi', 'pt', 'softId', 'tightId']\n",
      "['charge', 'dxy', 'dxyErr', 'dz', 'dzErr', 'eta', 'genPartIdx', 'genPartIdxG', 'jetIdx', 'jetIdxG', 'mass', 'pfRelIso03_all', 'pfRelIso04_all', 'phi', 'pt', 'softId', 'tightId']\n",
      "\n",
      "\n",
      "Second attemp at adding extended variables to events\n",
      "['Muon', 'luminosityBlock', 'Jet', 'HLT', 'PV', 'Tau', 'GenPart', 'MET', 'run', 'event', 'selectedMuon_1', 'myMuon', 'selectedMuon_2']\n",
      "['charge', 'dxy', 'dxyErr', 'dz', 'dzErr', 'eta', 'genPartIdx', 'genPartIdxG', 'jetIdx', 'jetIdxG', 'mass', 'pfRelIso03_all', 'pfRelIso04_all', 'phi', 'pt', 'softId', 'tightId']\n",
      "['charge', 'dxy', 'dxyErr', 'dz', 'dzErr', 'eta', 'genPartIdx', 'genPartIdxG', 'jetIdx', 'jetIdxG', 'mass', 'pfRelIso03_all', 'pfRelIso04_all', 'phi', 'pt', 'softId', 'tightId']\n",
      "\n",
      "\n",
      "Third attemp at adding extended variables to events\n",
      "['Muon', 'luminosityBlock', 'Jet', 'HLT', 'PV', 'Tau', 'GenPart', 'MET', 'run', 'event', 'selectedMuon_1', 'myMuon', 'selectedMuon_2', 'selectedMuon_3']\n",
      "['charge', 'dxy', 'dxyErr', 'dz', 'dzErr', 'eta', 'genPartIdx', 'genPartIdxG', 'jetIdx', 'jetIdxG', 'mass', 'pfRelIso03_all', 'pfRelIso04_all', 'phi', 'pt', 'softId', 'tightId', 'invpt']\n",
      "['charge', 'dxy', 'dxyErr', 'dz', 'dzErr', 'eta', 'genPartIdx', 'genPartIdxG', 'jetIdx', 'jetIdxG', 'mass', 'pfRelIso03_all', 'pfRelIso04_all', 'phi', 'pt', 'softId', 'tightId', 'invpt']\n"
     ]
    }
   ],
   "source": [
    "print('First attempt at adding extended variables to events')\n",
    "events.Muon['invpt'] = 1/events.Muon.pt\n",
    "events['selectedMuon_1'] = events.Muon[events.Muon.pt > 50]\n",
    "\n",
    "print(events.fields)\n",
    "print(events.Muon.fields)\n",
    "print(events.selectedMuon_1.fields)\n",
    "\n",
    "print('\\n\\nSecond attemp at adding extended variables to events')\n",
    "events['myMuon'] = events.Muon[:]\n",
    "events.myMuon['invpt'] = 1/events.myMuon.pt\n",
    "events['selectedMuon_2'] = events.myMuon[events.myMuon.pt > 50]\n",
    "\n",
    "print(events.fields)\n",
    "print(events.myMuon.fields)\n",
    "print(events.selectedMuon_2.fields)\n",
    "\n",
    "print('\\n\\nThird attemp at adding extended variables to events')\n",
    "myMuon = events.Muon[:]\n",
    "myMuon['invpt'] = 1/myMuon.pt\n",
    "events['selectedMuon_3'] = myMuon[myMuon.pt > 50]\n",
    "\n",
    "print(events.fields)\n",
    "print(myMuon.fields)\n",
    "print(events.selectedMuon_3.fields)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's get the straightforward part of the code clear up. The addition of\n",
    "collections looks very straight forward, one can think of the `events` as\n",
    "something that looks like a \"dictionary of collection with a common outer\n",
    "dimension\", so the addition of the two electron collections to the event has a\n",
    "very distionary-esque notation. What is strange is the persistence of the\n",
    "extended collection for the electrons. Logically, the operation looks\n",
    "identical, but the first attempt to add the new variable `invpt` directly to\n",
    "`events.Muon` fails to persist, and thus all direct extensions of `events.Muon`\n",
    "doesn't include the new `invpt` field.\n",
    "\n",
    "The reason for this is rather technical regarding the mutability of objects in\n",
    "python and awkward. The rule-of-thumb is that collections that are directly\n",
    "generated from the file, (a.k.a. the collections directly obtained listed the\n",
    "`events.fields` immediate after opening a file) can **never** be altered, and\n",
    "therefore cannot have extended variables added to them. To create an extended\n",
    "variable to some collection, we will need to make some sort of copy of the\n",
    "original either by some trivial kinematic selection (ex. `myMuons =\n",
    "events.Muon[events.Muon.pt > 0]`) or some trivial splicing (`myMuon =\n",
    "events.Muons[:]`). Another feature of mutability is that once the collection is\n",
    "added to the event collection, it becomes immutable. That is why the third\n",
    "attempt is the one that adds the both the electron extended variable and the\n",
    "extended electron collection to the event.\n",
    "\n",
    "Because of these quirks, it would typically be worth it to wrap the object\n",
    "selection into a function if the object selection is typical within an\n",
    "analysis, and it also helps with code readability"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['Muon', 'luminosityBlock', 'Jet', 'HLT', 'PV', 'Tau', 'GenPart', 'MET', 'run', 'event', 'selectedMuon_1', 'myMuon', 'selectedMuon_2', 'selectedMuon_3', 'selectedMuon_f']\n",
      "['charge', 'dxy', 'dxyErr', 'dz', 'dzErr', 'eta', 'genPartIdx', 'genPartIdxG', 'jetIdx', 'jetIdxG', 'mass', 'pfRelIso03_all', 'pfRelIso04_all', 'phi', 'pt', 'softId', 'tightId', 'invpt']\n"
     ]
    }
   ],
   "source": [
    "def SelectMuons(muon):\n",
    "    muon = muon[muon.pt > 50]\n",
    "    muon['invpt']  = 1.0 / muon.pt\n",
    "    return muon\n",
    "\n",
    "events['selectedMuon_f'] = SelectMuons(events.Muon)\n",
    "print(events.fields)\n",
    "print(events.selectedMuon_f.fields)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Once the new object collection has been added to the event collection, they\n",
    "will persist to arbitrary levels of event selection:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['Muon', 'luminosityBlock', 'Jet', 'HLT', 'PV', 'Tau', 'GenPart', 'MET', 'run', 'event', 'selectedMuon_1', 'myMuon', 'selectedMuon_2', 'selectedMuon_3', 'selectedMuon_f']\n",
      "['charge', 'dxy', 'dxyErr', 'dz', 'dzErr', 'eta', 'genPartIdx', 'genPartIdxG', 'jetIdx', 'jetIdxG', 'mass', 'pfRelIso03_all', 'pfRelIso04_all', 'phi', 'pt', 'softId', 'tightId', 'invpt']\n",
      "['Muon', 'luminosityBlock', 'Jet', 'HLT', 'PV', 'Tau', 'GenPart', 'MET', 'run', 'event', 'selectedMuon_1', 'myMuon', 'selectedMuon_2', 'selectedMuon_3', 'selectedMuon_f']\n",
      "['charge', 'dxy', 'dxyErr', 'dz', 'dzErr', 'eta', 'genPartIdx', 'genPartIdxG', 'jetIdx', 'jetIdxG', 'mass', 'pfRelIso03_all', 'pfRelIso04_all', 'phi', 'pt', 'softId', 'tightId', 'invpt']\n",
      "['Muon', 'luminosityBlock', 'Jet', 'HLT', 'PV', 'Tau', 'GenPart', 'MET', 'run', 'event', 'selectedMuon_1', 'myMuon', 'selectedMuon_2', 'selectedMuon_3', 'selectedMuon_f']\n",
      "['charge', 'dxy', 'dxyErr', 'dz', 'dzErr', 'eta', 'genPartIdx', 'genPartIdxG', 'jetIdx', 'jetIdxG', 'mass', 'pfRelIso03_all', 'pfRelIso04_all', 'phi', 'pt', 'softId', 'tightId', 'invpt']\n"
     ]
    }
   ],
   "source": [
    "myevents = events[ak.count(events.selectedMuon_f.pt,axis=-1) > 0 ]\n",
    "\n",
    "print(myevents.fields)\n",
    "print(myevents.selectedMuon_f.fields)\n",
    "\n",
    "myevents = events[ak.count(events.selectedMuon_f.pt,axis=-1) > 1 ]\n",
    "\n",
    "print(myevents.fields)\n",
    "print(myevents.selectedMuon_f.fields)\n",
    "myevents = events[ak.count(events.selectedMuon_f.pt,axis=-1) > 2 ]\n",
    "\n",
    "print(myevents.fields)\n",
    "print(myevents.selectedMuon_f.fields)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Comparing array based computations and loop based computation\n",
    "\n",
    "So to put this together into a single code block, suppose our analysis\n",
    "consisted of selecting events that have at least 1 electron with $p_{T} >\n",
    "50GeV$, $|\\eta| < 2.4$, and we want to calculate the average of all such\n",
    "electron's inverse $p_{T}$ within the selected events. Our awkward array code\n",
    "would look something like:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.018032373239596684\n"
     ]
    }
   ],
   "source": [
    "events = NanoEventsFactory.from_root( 'file:TTbar.root',\n",
    "                                      'Events',\n",
    "                                      entry_stop=50,\n",
    "                                      schemaclass=NanoAODSchema).events()\n",
    "\n",
    "## Object selection\n",
    "selectedMuon = events.Muon[ (events.Muon.pt > 50) &\n",
    "                            (np.abs(events.Muon.eta)<2.4) ]\n",
    "selectedMuon['invpt'] = 1/selectedMuon.pt\n",
    "events['selectedMuon'] = selectedMuon\n",
    "\n",
    "# Event selection\n",
    "events = events[ak.count(events.selectedMuon.pt,axis=-1) >= 1]\n",
    "\n",
    "# Calculating the total average\n",
    "print(ak.sum(events.selectedMuon.invpt)/ak.count(events.selectedMuon.invpt))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In total, this is 4 statements (not counting the file reading step) used to\n",
    "make this analysis. Compare that with the loop based notation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.01803237348586199\n"
     ]
    }
   ],
   "source": [
    "events = NanoEventsFactory.from_root( 'file:TTbar.root',\n",
    "                                      'Events',\n",
    "                                      entry_stop=50,\n",
    "                                      schemaclass=NanoAODSchema).events()\n",
    "\n",
    "count = 0\n",
    "suminv = 0\n",
    "for i in range(len(events)):\n",
    "    is_good = []\n",
    "    for j in range(len(events[i].Muon)):\n",
    "        if events[i].Muon[j].pt > 50 and np.abs(events[i].Muon[j].eta) < 2.4:\n",
    "            is_good.append(j)\n",
    "    if len(is_good) >= 1:\n",
    "        for j in is_good:\n",
    "            count = count +1\n",
    "            suminv += 1.0/ events[i].Muon[j].pt\n",
    "\n",
    "print(suminv/count)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notice the results are only difference because the 32-bit to 64-bit float\n",
    "conversion is happening at different places. For awkward arrays, this is\n",
    "happening only after the sum has been performed. For the loop based approach\n",
    "this happening every time the `+=` operator is called.\n",
    "\n",
    "In the loop based analysis, notice for such a simple analysis, many lines of\n",
    "code are dedicated to just bookkeeping stuff: number of electrons passing\n",
    "criteria, adding a counter variable and sum variable... etc., instead of actually\n",
    "analysis computation. The array based notation for expressing the analysis is\n",
    "much cleaner, if rather more unfamiliar to typical users.\n",
    "\n",
    "Of course, this isn't the end. Physics analysis are typically more involved\n",
    "than just basic selection and counting. In the next session, we will talk about\n",
    "how to perform more involved calculations with awkward arrays that use\n",
    "multiple collections within an event collection."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## More on index notation\n",
    "\n",
    "Let us dive in a bit more regarding the index notation, with a fresh set of\n",
    "events."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "events = NanoEventsFactory.from_root( 'file:TTbar.root',\n",
    "                                      'Events',\n",
    "                                      entry_stop=1000,\n",
    "                                      schemaclass=NanoAODSchema).events()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The index notation largely follows in index notation of `numpy` arrays, which\n",
    "means we can use some of the following\n",
    "\n",
    "- `array[ -n ]`: to the objects in the array in reverse order\n",
    "- `array[ start:end:step ]`: to slice an array given a certain range and step size,\n",
    "- `array[ [list, of, indexes] ]`: create a new array composed of the given indices\n",
    "\n",
    "This is very simply when the array is just `N` dimensional (a.k.a. contains a\n",
    "single index). Let us use the most simple example, the `events.MET.pt` object\n",
    "is just a list floats, with the length matching the number of events. Using\n",
    "this we can get:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The entire MET list:                   [18.3, 50.1, 6.06, 38.8, 4.99, 68, 51.1, ... 8.04, 30.6, 107, 20.8, 31.5, 85.6, 40.6]\n",
      "The MET from the second of last event: 85.59586334228516\n",
      "The MET of the first 5 events:         [18.3, 50.1, 6.06, 38.8, 4.99]\n",
      "The MET of every other event:          [18.3, 6.06, 4.99, 51.1, 28.5, 12.7, 16.8, ... 56.2, 87.7, 133, 30.6, 20.8, 85.6]\n",
      "The MET of events 1 3 and 10:          [50.1, 38.8, 12.7]\n"
     ]
    }
   ],
   "source": [
    "print(\"The entire MET list:                  \", events.MET.pt)\n",
    "print(\"The MET from the second of last event:\", events.MET.pt[-2])\n",
    "print(\"The MET of the first 5 events:        \", events.MET.pt[0:5])\n",
    "print(\"The MET of every other event:         \", events.MET.pt[0::2])\n",
    "print(\"The MET of events 1 3 and 10:         \", events.MET.pt[[1,3,10]])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Very rarely will a typical analysis want to directly operator on the event\n",
    "index. What we can do is operate on object indices using the notation:\n",
    "\n",
    "```\n",
    "object[event_index, object_index]\n",
    "```\n",
    "\n",
    "There the comma in the square braces is used to separate the first index\n",
    "argument that operation on the outermost dimension (in our case the event\n",
    "index), while the second argument is for the inner dimension index (in which\n",
    "our case the object index within an event). For instance, if we want to get the\n",
    "leading jet's $p_{T}$ for every event, we can use the following\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[147, 76.9, 84.4, 52.8, 105, 94.3, 34.8, ... 41.7, 133, 210, 91.3, 76.3, 96.2, 139]\n"
     ]
    }
   ],
   "source": [
    "print(events.Jet.pt[:,0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The first argument `:` is a trivial slice operation, indicating that we want\n",
    "every event. The second operator indicates that we want the leading (a.k.a.\n",
    "0-th) jet in each event. Notice that the user is responsible to make sure a\n",
    "corresponding object exists. For instance, if we run the same notation using\n",
    "the muon collection instead:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "ename": "ValueError",
     "evalue": "in ListArray64 attempting to get 0, index out of range\n\n(https://github.com/scikit-hep/awkward-1.0/blob/0.4.5/src/cpu-kernels/awkward_NumpyArray_getitem_next_at.cpp#L21)",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-22-af27855176d3>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mevents\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mMuon\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpt\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;32m~/VirtualENV/coffeaenv/lib/python3.8/site-packages/awkward1/highlevel.py\u001b[0m in \u001b[0;36m__getitem__\u001b[0;34m(self, where)\u001b[0m\n\u001b[1;32m    962\u001b[0m         \u001b[0mhave\u001b[0m \u001b[0mthe\u001b[0m \u001b[0msame\u001b[0m \u001b[0mdimension\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0mthe\u001b[0m \u001b[0marray\u001b[0m \u001b[0mbeing\u001b[0m \u001b[0mindexed\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    963\u001b[0m         \"\"\"\n\u001b[0;32m--> 964\u001b[0;31m         \u001b[0;32mreturn\u001b[0m \u001b[0mawkward1\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_util\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mwrap\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_layout\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mwhere\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_behavior\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    965\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    966\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0m__setitem__\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mwhere\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mwhat\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mValueError\u001b[0m: in ListArray64 attempting to get 0, index out of range\n\n(https://github.com/scikit-hep/awkward-1.0/blob/0.4.5/src/cpu-kernels/awkward_NumpyArray_getitem_next_at.cpp#L21)"
     ]
    }
   ],
   "source": [
    "print(events.Muon.pt[:,0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We get an out of range error, since not every event has a muon. In this case,\n",
    "we would first need to filter on events to ensure events contain muons, and\n",
    "then run same notation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[23.6, 57.3, 25.6, 3.13, 12.8, 17.9, 8.44, ... 3.42, 7.64, 8.26, 3.29, 27, 14.6, 3.3]\n"
     ]
    }
   ],
   "source": [
    "events_with_muons = events[ak.count(events.Muon.pt,axis=-1) > 0 ]\n",
    "print(events_with_muons.Muon.pt[:,0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The second argument doesn't need to be singular either, for instance, if our\n",
    "analysis is only interested in the 2 leading jets within an event, what we can\n",
    "do is:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[147, 72.7], [76.9, 70.5], [84.4, 60.5], ... [76.3, 72.7], [96.2, 77.5], [139, 106]]\n"
     ]
    }
   ],
   "source": [
    "events_with_2jets = events[ak.count(events.Jet.pt,axis=-1) >= 2 ]\n",
    "events_with_2jets['selectedJetsPt'] = events_with_2jets.Jet[:,0:2]\n",
    "print(events_with_2jets.selectedJetsPt.pt)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Where we have now created the collection which contains just the 2 leading jets\n",
    "in the event using a slicing notation. The list of index notation is also\n",
    "helpful if we want to reorder objects. For example if we want to find the 2\n",
    "jets that are the most forward in the detector, we can use awkward1 methods to\n",
    "generate the required index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[-2.1, -1.46, -1.73, -1.57, 0.0878], ... 0.94, -1.85, 3.86, 4.29, -1.51, -2.49]]\n",
      "[[0, 2, 3, 1, 4], [1, 4, 0, 2, 3], [5, ... 2, 0, 5, 1, 3], [5, 4, 7, 3, 1, 6, 2, 0]]\n",
      "[[-2.1, -1.73, -1.57, -1.46, 0.0878], ... -2.49, -1.85, -1.75, -1.51, 0.94, 0.0518]]\n"
     ]
    }
   ],
   "source": [
    "order = ak.argsort(np.abs(events.Jet.eta), axis=-1 , ascending=False)\n",
    "print(events.Jet.eta)\n",
    "print(order)\n",
    "events['JetByEta'] = events.Jet[order]\n",
    "print(events.JetByEta.eta)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now this example is a bit more involved:\n",
    "\n",
    "- First the method\n",
    "  [`ak.argsort`](https://awkward-array.readthedocs.io/en/latest/_auto/ak.argsort.html),\n",
    "  generates the index by which the object should be called along the last\n",
    "  axis(`axis=-1`) in descending order: For example in the first event, the jet\n",
    "  with the largest $|\\eta|$ is the 0-th on, followed by the 2-nd one and so on.\n",
    "  The `order` variable would be a `NxA` array of integers.\n",
    "- Next, when placed in the square parenthesis, since the outer dimension of the\n",
    "  `order` matches the outer dimension of `events.Jet` it automatically knows to\n",
    "  apply the index by event to each inner dimension.\n",
    "\n",
    "Notice because of how the second line is phrased, the resulting\n",
    "`events.JetByEta` has an inner dimension that is the same as the `order`\n",
    "variable and *not* the `events.Jet` variable. For instance, if we are only\n",
    "interested in the 2 least forward jets we can use:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[-1.46, 0.0878], [-0.27, -0.178], [0.738, ... [-0.399, -0.239], [0.94, 0.0518]]\n"
     ]
    }
   ],
   "source": [
    "new_order = order[:, -2:]\n",
    "events['JetByEta'] = events.Jet[new_order]\n",
    "print(events.JetByEta.eta)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Another handy application of this is the index notation is the old process of\n",
    "using object indices in the place of pointers. In this NanoAOD event, muons\n",
    "have a `genPartIdx` field, which indicates which `events.GenPart` collection it\n",
    "should be matched to. Notice that we will need to some additional parsing to\n",
    "ensure, that the muons of interest have correct `genPartIdx` values (>= 0)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[23.6], [], [], [57.3, 29.3, 7.7, ... [3.3, 4.25, 3.28, 19.6, 17.3, 6.62, 3.72]]\n",
      "[[24], [], [], [56.2, 29.4, 29.4], ... [14.7], [3.31, 27, 27, 27, 27, 27, 3.31]]\n"
     ]
    }
   ],
   "source": [
    "muon_withgen = events.Muon[events.Muon.genPartIdx >= 0]\n",
    "muon_genpart = events.GenPart[muon_withgen.genPartIdx]\n",
    "print(muon_withgen.pt)\n",
    "print(muon_genpart.pt)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now since the `muon_genpart` has the same dimension as that of the\n",
    "`events.MuonWithGen` collection, we can add the `muon_genpart` to the\n",
    "`MuonWithGen` collection. This makes additional codes much more useful:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[-13], [], [], [13, 13, 13], [], ... -13], [-13], [-13, 11, 11, 11, 11, 11, -13]]\n",
      "[[1], [], [], [1, 1, 1], [], [], [2], ... [], [1, 1], [1], [1, 1, 1, 1, 1, 1, 1]]\n"
     ]
    }
   ],
   "source": [
    "muon_withgen['GenPart'] = muon_genpart\n",
    "print(muon_withgen.GenPart.pdgId)\n",
    "print(muon_withgen.GenPart.status)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now the code mimics quite a bit with what you would normally get with object\n",
    "based code, which makes awkward arrays rather more useful than just the flat\n",
    "arrays one typically get in C++ n-tuple codes."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Calculation with objects\n",
    "\n",
    "Up til now, the calculations have all removed around single collection. Even in\n",
    "the case of the cell involving two collection, one is just used to generate\n",
    "index dimension. The first instance of calculation, where we use a mathematic\n",
    "operation to generate a new array is in the case of mask generation and the\n",
    "calculation of the inverse of muon pt, but that apparently only involved 1\n",
    "array. In this section, we dive a bit deeper into what mathematical operators\n",
    "actually imply in awkward.\n",
    "\n",
    "The mathematical operations of awkward arrays is member-wise: Given 2 arrays of\n",
    "identical dimension ($a$ and $b$) and a generic binary operation $\\otimes$, the\n",
    "statement $a\\otimes b$, will generate a new array with the same dimensions as\n",
    "$a$, with output being a member-wise results of $o_{ijk\\ldots} = a_{ijk\\ldots}\n",
    "\\otimes b{ijk\\ldots}$. Now what if $a$ and $b$ are different dimension, say in\n",
    "our case $a$ is a scalar 1, and $b$ is the muon $p_{T}$ array. What happens is\n",
    "formally known as\n",
    "[**broadcasting**](https://numpy.org/doc/stable/user/basics.broadcasting.html):\n",
    "you can think that awkward automatically expands the scalar $a$ into a `NxA`\n",
    "array and performs a typical member-wise arithmetics. Broadcasting is a rather\n",
    "complicated subject to fully discuss, but a  rule-of-thumb is that the\n",
    "follow operations would be considered valid:\n",
    "\n",
    "- scalar $\\otimes$ `N`: scalar is expanded to a `N`\n",
    "- scalar $\\otimes$ `NxA`: scalar is expanded to a `NxA` array\n",
    "- `NxA`  $\\otimes$ `NxA`: typical member-wise calculation\n",
    "- `N` $\\otimes$ `N`: expansion of the first `N` array to a `NxA` array\n",
    "\n",
    "Since in awkward array, the `A` dimension is a variable one, we will never\n",
    "encounter the case of a `A`$\\otimes$`NxA` array. To use some concrete\n",
    "examples:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "scalar @ N   : [92.7, 145, 99.4, 121, 101, 35.7, 146, ... 108, 99.9, 207, 120, 76.7, 14.7, 139]\n",
      "scalar @ NxA : [[89.8], [], [], [148, 116, 104], ... [88.3], [103, 104, 103, 119, 117, 107, 104]]\n",
      "NxA    @ NxA : [[11.1], [], [], [16, -7.89, ... [3.99, 4.12, 3.91, 23.5, 21.4, 6.87, 4.26]]\n",
      "N      @ NxA : [[-17.5], [], [], [68.7, 37.4, 25, ... [42.3, 43.3, 42.3, 58.2, 55.6, 45.7, 42.7]]\n"
     ]
    }
   ],
   "source": [
    "print('scalar @ N   :', 100 +  events.MET.px)\n",
    "print('scalar @ NxA :', 100 + events.Muon.px)\n",
    "print('NxA    @ NxA :', events.Muon.px + events.Muon.py)\n",
    "print('N      @ NxA :', events.MET.px + events.Muon.px)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Pretty much all calculations can be broken down into binary operations like\n",
    "above, `coffea` even provides some common HEP function such as `delta_r`, so\n",
    "you don't have to re-write common routines from scratch. The difficulty comes\n",
    "in mangling the array into the correct dimension, so that broadcasting can take\n",
    "over. For instance, the calculation of the $\\Delta R$ of the first and second\n",
    "leading jets in an event is simple, this is just a reduction of the 2 `NxA`\n",
    "array to 2 `N` array via index selection, then calling an in-built binary\n",
    "operation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[3.18, 4.35, 5.69, 0.947, 2.7, 1.8, 1.65, ... 2.9, 3.19, 2.89, 3.04, 2.85, 3.61]\n"
     ]
    }
   ],
   "source": [
    "evt = events[ak.count(events.Jet.pt,axis=-1) >= 2 ]\n",
    "first_jet = evt.Jet[:,0]\n",
    "second_jet = evt.Jet[:,1]\n",
    "print(first_jet.delta_r(second_jet))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To calculated of $\\Delta R$ the leading jet with all muon in the event is also\n",
    "simple (notice how it takes on the dimension of the Muon collection)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[3.16], [], [], [2.66, 3.49, 3.56, ... [3.5, 3.47, 3.58, 3.55, 3.48, 3.55, 3.53]]\n"
     ]
    }
   ],
   "source": [
    "print(first_jet.delta_r(evt.Muon))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now suppose we want to perform a jet-cone matching of the gen particles\n",
    "collection with the **all** jets in the event. How do we perform broadcasting\n",
    "of a `NxA` array with a `NxB` array? The answer is to expand one of the array\n",
    "into a `NxAxB` array for the broadcast for higher dimensions can be performed.\n",
    "Details of how this method works can be found the awkward documentation for the\n",
    "[`cartesian`](https://awkward-array.readthedocs.io/en/latest/_auto/ak.cartesian.html)\n",
    "and\n",
    "[`unzip`](https://awkward-array.readthedocs.io/en/latest/_auto/ak.unzip.html)\n",
    "method, but here is skeleton code performing object x object computations:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[[], [24, 24, 24, 24], [], [], []], ... 27], [], [], [], [], [0.545, 0.545], []]]\n"
     ]
    }
   ],
   "source": [
    "jets = evt.Jet[:]\n",
    "_, genpart = ak.unzip(ak.cartesian([evt.Jet,evt.GenPart],nested=True))\n",
    "jets['genpart'] = genpart[jets.delta_r(genpart) < 0.4]\n",
    "print(jets.genpart.pt)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Breaking down the code above:\n",
    "\n",
    "The `ak.unzip` and `ak.cartesian` method expands the `NxA` `evt.Jet` array and\n",
    "the `NxA` `evt.GenPart` array in to a `NxBxA` arrays (named `_` since we\n",
    "wouldn't be using this) and `NxAxB` array `genpart`. The broadcast is the\n",
    "performed on the `NxA` jets array and `NxAxB` `genpart` array to generate a\n",
    "delta R mask of dimension `NxAxB`, which is then used to filter the `genpart`.\n",
    "Now we can store the `NxAxB` array into the `NxA` jet collection, which is\n",
    "used to find all gen particles associated with a jet via delta-R matching!\n",
    "\n",
    "This is probably the most advanced that calculation that a typical analysis\n",
    "will need in computation: the computation of arrays of mismatch dimensions.\n",
    "Hopefully this introduction is enough to help people start to translate their\n",
    "loop based analysis calculation into array-based syntax."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## When all else fails - numba for loops\n",
    "\n",
    "TO BE WRITTEN: assigned -Yi-Mu Chen"
   ]
  }
 ],
 "metadata": {
  "jupytext": {
   "main_language": "python"
  },
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.1-final"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}