pyhep2022-combinatorics.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "4187e701-98e8-4fe3-9078-c760baf51000",
   "metadata": {},
   "source": [
    "# Uproot, Awkward Array, hist, Vector: from basics to combinatorics"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6dc5ce9e-61be-44aa-87e2-0768f54eeee1",
   "metadata": {},
   "source": [
    "This is an introduction to doing particle physics analysis with Scikit-HEP tools: Uproot, Awkward Array, hist, and Vector.\n",
    "\n",
    "It starts at a basic level—exploring files, making plots—and ramps up to resolving e⁺e⁻e⁺e⁻, μ⁺μ⁻μ⁺μ⁻, and e⁺e⁻μ⁺μ⁻ final states in Higgs decays.\n",
    "\n",
    "Due to time constraints, this is not a tutorial with exercises for you to try. But I've given 10 such tutorials in the modern era (Awkward Array version ≥ 1), which you can find online, and 5 are on YouTube.\n",
    "\n",
    "| <!-- --> | <!-- --> | <!-- --> | <!-- --> | <!-- --> |\n",
    "|:--|:--|:--|:--|:--|\n",
    "| 2020 | EIC collaboration tutorial | | [notebooks](https://github.com/jpivarski-talks/2020-04-08-eic-jlab#readme) | [video](https://www.youtube.com/watch?v=FoxNS6nlbD0) |\n",
    "| 2020 | CMS HATS tutorial | | [notebooks](https://github.com/jpivarski/2020-06-08-uproot-awkward-columnar-hats#readme) | |\n",
    "| 2020 | SciPy talk | non-HEP | [notebooks](https://github.com/jpivarski-talks/2020-07-06-scipy2020/blob/master/bikeroutes-execution.ipynb) | [video](https://youtu.be/WlnUF3LRBj4) |\n",
    "| 2020 | PyHEP tutorial | | [notebooks](https://github.com/jpivarski/2020-07-13-pyhep2020-tutorial#readme) | [video](https://youtu.be/ea-zYLQBS4U) |\n",
    "| 2021 | CMS HATS tutorial | | [notebooks](https://github.com/jpivarski-talks/2021-06-14-uproot-awkward-columnar-hats#readme) | |\n",
    "| 2021 | PyHEP tutorial | | [notebooks](https://github.com/jpivarski-talks/2021-07-06-pyhep-uproot-awkward-tutorial#readme) | [video](https://youtu.be/s47Nz0h0vcg) |\n",
    "| 2021 | STAR collaboration tutorial | | [notebooks](https://github.com/jpivarski-talks/2021-09-13-star-uproot-awkward-tutorial#readme) | |\n",
    "| 2022 | PyHEP topical meeting | | [notebooks](https://indico.cern.ch/event/1140031/) | |\n",
    "| 2022 | SciPy tutorial | non-HEP | [notebooks](https://github.com/jpivarski-talks/2022-07-11-scipy-loopy-tutorial#readme) | [video](https://youtu.be/Dovyd72eD70) |\n",
    "| 2022 | CoDaS-HEP tutorial | | [notebooks](https://github.com/jpivarski-talks/2022-08-03-codas-hep-columnar-tutorial#readme) | |\n",
    "\n",
    "Also, there are [HSF training events](https://hepsoftwarefoundation.org/Schools/events.html), featuring an [Uproot module](https://hsf-training.github.io/hsf-training-uproot-webpage/) and a [Scikit-HEP module](https://hsf-training.github.io/hsf-training-scikit-hep-webpage/).\n",
    "\n",
    "And, of course, each of these libraries has online documentation:\n",
    "\n",
    "- Uproot documentation at [https://uproot.readthedocs.io/](https://uproot.readthedocs.io/)\n",
    "- Awkward Array documentation at [https://awkward-array.org/](https://awkward-array.org/)\n",
    "- hist documentation at [https://hist.readthedocs.io/](https://hist.readthedocs.io/)\n",
    "- Vector documentation at [https://vector.readthedocs.io/](https://vector.readthedocs.io/)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2b396a73-d137-40ca-8d13-56d7b03def49",
   "metadata": {},
   "source": [
    "<br><br><br><br><br>\n",
    "\n",
    "This presentation uses the current versions of Uproot, Awkward Array, hist, and Vector that you would get from pip:\n",
    "\n",
    "```bash\n",
    "pip install numpy matplotlib uproot awkward hist[plot] vector\n",
    "```\n",
    "\n",
    "Note that new major versions are in development, to be released in December 2022 (Uproot 5.0, Awkward Array 2.0, and Vector 1.0), but the code in this notebook will still work in the new versions. (The outputs will look a little different—better!)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "da7b13a2-8d30-4408-a57a-abadca5154e6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'1.22.4'"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import numpy as np\n",
    "np.__version__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "5f19b3d1-096e-4fde-ae98-8e6913c227ef",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'3.5.3'"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "matplotlib.__version__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "e593041c-b85f-4fcf-a1e8-3d9f0266fc5e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'4.3.5'"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import uproot\n",
    "uproot.__version__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "e25cfd69-de7a-49f5-8428-953854db3456",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'1.9.0'"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import awkward as ak\n",
    "ak.__version__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "ee934a30-6804-44e4-ad0c-132b1d68ed5f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'2.6.1'"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import hist\n",
    "hist.__version__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "ab366034-7984-41ca-bf1b-e14fa70e745e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'0.10.0'"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import vector\n",
    "vector.__version__"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "33241fa2-e22c-47cc-9757-0ac86e7ae521",
   "metadata": {},
   "source": [
    "<br><br><br><br><br>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "54c4a955-f443-4b12-8153-ec6ffe5e1006",
   "metadata": {},
   "source": [
    "## Overview"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e2d3f297-0501-4c9c-8888-ea3c5c0cf9ad",
   "metadata": {},
   "source": [
    "Just as in the broader Python world, functionality in Scikit-HEP is provided by small, specialized, interoperating libraries that build on each other."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dcc9b413-1a4d-4227-8721-360e705f2169",
   "metadata": {},
   "source": [
    "<table width=\"1200px\"><tr style=\"background: white\">\n",
    "    <td align=\"center\" width=\"50%\" style=\"font-size: 24px; font-weight: bold\">Scientific Python ecosystem</td>\n",
    "    <td align=\"center\" width=\"50%\" style=\"font-size: 24px; font-weight: bold\">Scikit-HEP ecosystem</td>\n",
    "</tr><tr style=\"background: white\">\n",
    "    <td align=\"center\" width=\"50%\"><img src=\"img/shells-border.png\" width=\"95%\"></td>\n",
    "    <td align=\"center\" width=\"50%\"><img src=\"img/shells-hep.svg\" width=\"95%\"></td>\n",
    "</tr></table>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bfab4fd2-cbcf-47a9-ad73-c99def888088",
   "metadata": {},
   "source": [
    "<br><br>\n",
    "\n",
    "<img src=\"img/uproot-logo.png\" width=\"350px\">\n",
    "\n",
    "<p style=\"font-size: 14pt\">Reads ROOT data as <span style=\"background: yellow\">arrays</span>.</p>\n",
    "\n",
    "<br><br>\n",
    "\n",
    "<img src=\"img/awkward-logo.png\" width=\"350px\">\n",
    "\n",
    "<p style=\"font-size: 14pt\">Manipulates <span style=\"background: yellow\">arrays</span> of arbitrary data structures.</p>\n",
    "\n",
    "<br><br>\n",
    "\n",
    "<img src=\"img/hist-logo.png\" width=\"300px\">\n",
    "\n",
    "<p style=\"font-size: 14pt\">Fills histograms from <span style=\"background: yellow\">arrays</span> of data.</p>\n",
    "\n",
    "<br><br>\n",
    "\n",
    "<img src=\"img/vector-logo.png\" width=\"350px\">\n",
    "\n",
    "<p style=\"font-size: 14pt\">Transforms <span style=\"background: yellow\">arrays</span> of 2D, 3D, and Lorentz vectors.</p>\n",
    "\n",
    "<br><br>\n",
    "\n",
    "_Using them effectively will require thinking about whole arrays at a time._\n",
    "\n",
    "<br><br><br>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "39b4384b-49fb-431a-add8-7a44e46fc35c",
   "metadata": {},
   "source": [
    "## Array-at-a-time thinking and interactivity"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "31200110-52a2-4113-b90e-e94102bdfa14",
   "metadata": {},
   "source": [
    "There's a long history of \"array-oriented\" languages, in which the fundamental data type is an array and most functions operate on arrays."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "94983d0a-ab0a-4892-9a25-b5e709f0e933",
   "metadata": {},
   "source": [
    "<img src=\"img/apl-timeline.svg\" width=\"800\">\n",
    "\n",
    "<br><br><br>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a21f49c2-7069-41f1-9e6b-9b16ddf6f0da",
   "metadata": {},
   "source": [
    "Interestingly enough, all of these (except Fortran 90 and Dex) are _interactive_, command-line or notebook-driven languages.\n",
    "\n",
    "<br>\n",
    "\n",
    "They work especially well for studying _distributions of datasets_ and debugging intermediate values of _whole datasets_, not _individual entries_."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6bb4da3e-1e33-4352-8e53-de17b35339c1",
   "metadata": {},
   "source": [
    "<br><br>\n",
    "\n",
    "For example, suppose we start with a Gaussian-distributed dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2fabf753-9a41-4989-8104-6d417e6ab950",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = np.random.normal(0, 1, 1000000)  # one MILLION data points"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bec3b4c7-91a0-4445-ac82-fde52b7e470c",
   "metadata": {},
   "source": [
    "<br><br>\n",
    "\n",
    "You ask, \"How are they distributed?\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8cb71688-9e50-4ef5-93a5-9c2762f917fd",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.hist(dataset, bins=100);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ac34fe07-1010-4658-9d8a-a3b33a1a20d2",
   "metadata": {},
   "source": [
    "<br><br>\n",
    "\n",
    "\"How is the square of each point distributed?\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d650e80e-4c4e-4205-85c7-ab7ceb74f8f4",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.hist(dataset**2, bins=100);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eb96d11b-2c17-4509-8658-15737d952ffa",
   "metadata": {},
   "source": [
    "<br><br>\n",
    "\n",
    "\"What about this weirdo function? $\\sin\\left(\\frac{1}{x^2}\\right)$\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e8843bee-d3aa-40ae-a5df-0ab2a915aca0",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.hist(np.sin(1/dataset**2), bins=100);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a5e6380b-f5c7-4f9d-bbd2-d1e2bd516fd8",
   "metadata": {},
   "source": [
    "<br><br>\n",
    "\n",
    "The shape of the last distribution would have been hard to guess, and as long as we're in \"compute operation, plot, compute operation, plot\" mode, setting up `for` loops would be a hindrance.\n",
    "\n",
    "<br><br>\n",
    "\n",
    "Array-oriented computations and interactivity go together. If you don't like Jupyter, consider IPython, the basic Python prompt, or IDE features that let you interact with an ongoing session. Don't build up a long-running script by running the whole thing over and over until you've already figured out the details of your project.\n",
    "\n",
    "<br><br><br><br><br>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d3a59814-6966-4c09-af07-d5eb93284b0e",
   "metadata": {},
   "source": [
    "## Exploring a file with Uproot"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "db3d2276-5fe4-4237-852c-ffc9c49a1bf2",
   "metadata": {},
   "source": [
    "Uproot is a reimplementation of ROOT I/O in Python, with no dependence on C++ ROOT."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fef73c2d-283c-4ebd-9200-5dd7bd998ac8",
   "metadata": {},
   "source": [
    "<img src=\"img/abstraction-layers.svg\" width=\"800\">\n",
    "\n",
    "<br><br><br>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "539dfd85-97f2-4f98-bf58-b51e3db69c67",
   "metadata": {},
   "source": [
    "\"Where's my file?\"\n",
    "\n",
    "(Executing shell commands with `!` is a Jupyter/IPython thing.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d9abe970-c356-4015-b745-884812204ec2",
   "metadata": {},
   "outputs": [],
   "source": [
    "!ls -l"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f640b2ec-f390-44f7-8cf9-2af9abc1790f",
   "metadata": {},
   "outputs": [],
   "source": [
    "!ls -l data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "30f7c6fe-f52f-404a-bdbe-f6da1af2e5f0",
   "metadata": {},
   "source": [
    "<br><br>\n",
    "\n",
    "\"Open it in Uproot.\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3aaa37de-4dfd-43dc-8ae8-bcc71195dea3",
   "metadata": {},
   "outputs": [],
   "source": [
    "file = uproot.open(\"data/SMHiggsToZZTo4L.root\")\n",
    "file"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9f492326-970a-4377-8b98-97dce984fb46",
   "metadata": {},
   "source": [
    "<br><br>\n",
    "\n",
    "\"What's in the file?\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b40fdb35-3203-4da1-a415-6943c5b91f90",
   "metadata": {},
   "outputs": [],
   "source": [
    "file.keys()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "caeef5d6-c5eb-4478-9fd0-409a7db9e382",
   "metadata": {},
   "source": [
    "<br><br>\n",
    "\n",
    "\"What kind of object is 'Events'?\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5b141aa9-6d0f-4ef9-af75-f16589341a51",
   "metadata": {},
   "outputs": [],
   "source": [
    "file.classnames()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "54574f8d-b1ea-4652-a02d-3ba7f2db7916",
   "metadata": {},
   "source": [
    "<br><br>\n",
    "\n",
    "\"Get the TTree metadata (not the whole arrays; not yet).\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d3743938-0917-4c7f-b0fd-5d6e779f1b96",
   "metadata": {},
   "outputs": [],
   "source": [
    "tree = file[\"Events\"]\n",
    "tree"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a4a199fe-a559-4ce4-b9cb-5bbc9ee4d01a",
   "metadata": {},
   "outputs": [],
   "source": [
    "tree.num_entries"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f1db9801-4e5a-4ce0-9198-370d7d7f0c11",
   "metadata": {},
   "source": [
    "<br><br>\n",
    "\n",
    "\"What kinds of data are in the TTree?\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eeb3cdd7-ed75-4aab-ab32-106c23afaa08",
   "metadata": {},
   "outputs": [],
   "source": [
    "tree.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "24d7549d-f64f-4230-b293-bca21f934cad",
   "metadata": {},
   "source": [
    "<br><br>\n",
    "\n",
    "\"Can I get that information programmatically?\"\n",
    "\n",
    "(Yes.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "af16bb0c-66a7-4308-ae4e-206f57c4cff6",
   "metadata": {},
   "outputs": [],
   "source": [
    "{key: branch.typename for key, branch in tree.items()}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8dd57d94-9a40-4be6-95ac-325dadd73ee6",
   "metadata": {},
   "source": [
    "<br><br>\n",
    "\n",
    "\"Read the muon $p_T$, $\\eta$, $\\phi$, mass, and charge into arrays (and only those columns/fields).\"\n",
    "\n",
    "This is the first _expensive_ operation in our tutorial. If this were a large file, you'd have to wait."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "143437f7-dd0d-4cb2-a967-5cc7da5107af",
   "metadata": {},
   "outputs": [],
   "source": [
    "muon_kinematics = tree.arrays([\"Muon_pt\", \"Muon_eta\", \"Muon_phi\", \"Muon_mass\", \"Muon_charge\"])\n",
    "muon_kinematics"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c11a1be1-369f-4524-ab4f-953a64f0d44f",
   "metadata": {},
   "source": [
    "<br><br>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5f78f135-67d4-4b31-97d0-346a0f59581b",
   "metadata": {},
   "source": [
    "Note that there's a way to limit how much you read/how long you have to wait:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "95837846-7d20-4b07-a952-fbcc78c924eb",
   "metadata": {},
   "outputs": [],
   "source": [
    "tree.arrays(entry_start=0, entry_stop=100)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d5408df3-b26a-4395-b57d-e3caa3b7822b",
   "metadata": {},
   "source": [
    "<br><br>\n",
    "\n",
    "This is an Awkward Array. NumPy can't represent this dataset because there's a different number of muons in each event, and NumPy deals with rectilinear arrays (tensors) only."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5fecf60d-3c86-4cc1-bdef-6de2b8822f16",
   "metadata": {},
   "outputs": [],
   "source": [
    "muon_kinematics[:4].to_list()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fc6a2547-2bc4-4fb2-ad56-cffe11cd7aba",
   "metadata": {},
   "source": [
    "<br><br>\n",
    "\n",
    "\"Extract columns/fields from the already-read data.\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "084b30ba-e3f9-45a5-9489-bcbbb1c559db",
   "metadata": {},
   "outputs": [],
   "source": [
    "muon_kinematics.fields"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a9aa2817-a737-4bdd-82df-f944326f4526",
   "metadata": {},
   "outputs": [],
   "source": [
    "muon_kinematics.Muon_pt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ff3dfc47-9d26-4dba-888a-375e555eb541",
   "metadata": {},
   "outputs": [],
   "source": [
    "muon_kinematics.Muon_eta"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b66d2bee-0e64-4733-82ba-6326990cece4",
   "metadata": {},
   "outputs": [],
   "source": [
    "muon_kinematics.Muon_phi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "98bae042-270d-499d-bf46-9c34dcb0280e",
   "metadata": {},
   "outputs": [],
   "source": [
    "muon_kinematics.Muon_mass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e00ce610-cb49-40fc-8e85-d1135acd4a2b",
   "metadata": {},
   "outputs": [],
   "source": [
    "muon_kinematics.Muon_charge"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "509bcda0-7336-4777-897c-d57cc45021f0",
   "metadata": {
    "tags": []
   },
   "source": [
    "<br><br>\n",
    "\n",
    "\"What is the data type?\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8469f5cf-8032-4497-937e-411c4c749d72",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(muon_kinematics.type)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "28153fbb-ed15-4bcd-9955-f20454fcd924",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(muon_kinematics.Muon_pt.type)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "67809494-df34-450e-81ed-44577ebd6659",
   "metadata": {},
   "source": [
    "<br><br>\n",
    "\n",
    "Back to Uproot:\n",
    "\n",
    "\"Which columns in the TTree have anything to do with muons or electrons?\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a986b5ad-f81d-4bfe-b121-c814fc145dac",
   "metadata": {},
   "outputs": [],
   "source": [
    "tree.keys(filter_name=[\"Muon_*\", \"Electron_*\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "156e0da4-f7a9-4e36-9f3b-7d92d6fb25fe",
   "metadata": {},
   "source": [
    "<br><br>\n",
    "\n",
    "\"Read all of those columns into arrays, using the same `filter_name`.\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "54de7383-cee5-4ac1-878d-8818db562bdf",
   "metadata": {},
   "outputs": [],
   "source": [
    "muons_and_electrons = tree.arrays(filter_name=[\"Muon_*\", \"Electron_*\"])\n",
    "muons_and_electrons"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f6589eb9-423b-4874-b83b-1bfc94c2acf1",
   "metadata": {},
   "source": [
    "(Test a `filter_name` with `tree.keys`, then use it with `tree.arrays`.)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8cf00e93-ffb4-40b9-92b1-03b479f37352",
   "metadata": {},
   "source": [
    "<br><br>\n",
    "\n",
    "\"Show me the type of this larger dataset.\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "07a6ca94-cebb-4290-8898-231e38be1b3c",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(muons_and_electrons.type)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "96c268c1-49d9-4855-9b07-cde04d1f342b",
   "metadata": {},
   "source": [
    "<br><br><br><br><br>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "76ec487b-bab2-40c4-bc8b-1b4509b45228",
   "metadata": {},
   "source": [
    "## Examining distributions with hist"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "28f46bfa-b0e3-4888-8cd0-637fef5252dc",
   "metadata": {},
   "source": [
    "NumPy and Matplotlib both have `histogram` (or `hist`) functions, but these are _display_ options. In particle physics, we almost always want a histogram object that we can fill in parallel, redraw without refilling, etc. Even weighted histograms with error bars require more than what NumPy and Matplotlib give us.\n",
    "\n",
    "So although"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9407d390-b84f-4c5e-9ee5-077b6c3d8e09",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.hist(ak.flatten(muon_kinematics.Muon_pt), bins=100, range=(0, 150));"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a9dbf34d-6ba8-412f-9022-fab42512bbf1",
   "metadata": {},
   "source": [
    "works, you probably want to use the hist library:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "98c8e0e0-9b09-4301-b1d4-a02b9eb499db",
   "metadata": {},
   "outputs": [],
   "source": [
    "h = hist.Hist.new.Regular(100, 0, 150, name=\"muon pT\").Double()\n",
    "\n",
    "h.fill(ak.flatten(muon_kinematics.Muon_pt))\n",
    "\n",
    "h.plot();"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7970a058-1772-47b4-86ce-9b6e40aeb0ee",
   "metadata": {},
   "source": [
    "<br><br><br>\n",
    "\n",
    "Or even"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b3e8d144-31b6-4e8c-8f9e-5c917ae50cdc",
   "metadata": {},
   "outputs": [],
   "source": [
    "(\n",
    "    hist.Hist.new.Regular(100, 0, 150, name=\"muon pT\").Weight()\n",
    "        .fill(\n",
    "            ak.flatten(muon_kinematics.Muon_pt),\n",
    "            weight=ak.flatten(muon_kinematics.Muon_pt),\n",
    "        )\n",
    "        .plot()\n",
    ");"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8a0ab4cd-5e01-4d5b-9b94-68022915a402",
   "metadata": {},
   "source": [
    "<br><br><br>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d1d51748-55cd-401e-88a7-56f6abe2f30d",
   "metadata": {},
   "source": [
    "### Why `ak.flatten`?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e9f62c7d-88a5-438f-af24-cfb4aef5681a",
   "metadata": {},
   "source": [
    "In both of the above examples, we didn't just plot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1038e219-3c00-4295-a140-126365166a3f",
   "metadata": {},
   "outputs": [],
   "source": [
    "muon_kinematics.Muon_pt"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f785a128-30e3-4756-81ac-1edf1e8382a2",
   "metadata": {},
   "source": [
    "but"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "57b1f7f3-28fe-4222-8454-ca12010d2e37",
   "metadata": {},
   "outputs": [],
   "source": [
    "ak.flatten(muon_kinematics.Muon_pt)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7f79508f-0412-426a-868c-561f41f0c04a",
   "metadata": {},
   "source": [
    "<br><br>\n",
    "\n",
    "The first array is an array of _lists_. To fill a histogram, you need an array of _numbers_. `ak.flatten` removes the list structure.\n",
    "\n",
    "<br><br>\n",
    "\n",
    "There are other ways of removing the list structure, depending on what you're trying to do:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "57ff969f-ce45-44a1-b523-dae0e02f2242",
   "metadata": {},
   "outputs": [],
   "source": [
    "ak.sum(muon_kinematics.Muon_pt, axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fa86366d-bcaa-4a33-afae-a4282fed72f3",
   "metadata": {},
   "outputs": [],
   "source": [
    "(\n",
    "    hist.Hist.new.Regular(100, 0, 150, name=\"muon pT\").Double()\n",
    "        .fill(ak.sum(muon_kinematics.Muon_pt, axis=1))\n",
    "        .plot()\n",
    ");\n",
    "plt.yscale(\"log\");"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "43baf6c4-e8d1-4a80-90b8-86591df03049",
   "metadata": {},
   "source": [
    "<br><br>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c48592d7-10aa-4ef7-94e3-432f56358a13",
   "metadata": {},
   "outputs": [],
   "source": [
    "ak.max(muon_kinematics.Muon_pt, axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e4321689-b6e9-4a49-9839-8ce32533b2d3",
   "metadata": {},
   "outputs": [],
   "source": [
    "ak.fill_none(ak.max(muon_kinematics.Muon_pt, axis=1), 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b3ec201d-2f4f-4f12-ae04-c22e66f9b7e3",
   "metadata": {},
   "outputs": [],
   "source": [
    "(\n",
    "    hist.Hist.new.Regular(100, 0, 150, name=\"muon pT\").Double()\n",
    "        .fill(ak.fill_none(ak.max(muon_kinematics.Muon_pt, axis=1), 0))\n",
    "        .plot()\n",
    ");\n",
    "plt.yscale(\"log\");"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bdbccf91-dd04-411d-95a8-496f2222a500",
   "metadata": {},
   "source": [
    "<br><br><br><br><br>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9b894b49-9e4b-4c97-8f2d-1bd1ee595963",
   "metadata": {},
   "source": [
    "## Computing vector coordinates"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fbb59bf2-f942-4a77-adda-4e6d61bb82ac",
   "metadata": {},
   "source": [
    "We can compute coordinate transformations and invariant masses manually, using NumPy's array-at-a-time functions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1cb60002-7b64-43c0-82fb-2417e887a036",
   "metadata": {},
   "outputs": [],
   "source": [
    "px = muon_kinematics.Muon_pt * np.cos(muon_kinematics.Muon_phi)\n",
    "py = muon_kinematics.Muon_pt * np.sin(muon_kinematics.Muon_phi)\n",
    "pz = muon_kinematics.Muon_pt * np.sinh(muon_kinematics.Muon_eta)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "51fd95b7-fc3c-4605-baab-deeff49bf53a",
   "metadata": {},
   "outputs": [],
   "source": [
    "px"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4a7afefe-133f-4434-a232-38e2ba80abf2",
   "metadata": {},
   "outputs": [],
   "source": [
    "py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17228761-b0f5-44ac-aba7-55d5fd0b1fb5",
   "metadata": {},
   "outputs": [],
   "source": [
    "pz"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ea6c9c50-0a3f-45d6-ae74-47f09c460b6c",
   "metadata": {},
   "outputs": [],
   "source": [
    "(\n",
    "    hist.Hist.new.Regular(100, -25, 25, name=\"px\").Regular(100, -25, 25, name=\"py\").Double()\n",
    "        .fill(ak.flatten(px), ak.flatten(py))\n",
    "        .plot()\n",
    ");"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2f64b9df-af72-4209-bfe7-3288356e7eab",
   "metadata": {},
   "outputs": [],
   "source": [
    "(\n",
    "    hist.Hist.new.Regular(100, -50, 50, name=\"pz\").Regular(100, 0, 15, name=\"pt\").Double()\n",
    "        .fill(ak.flatten(pz), ak.flatten(np.sqrt(px**2 + py**2)))\n",
    "        .plot2d_full()\n",
    ");"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "77c53c30-46e1-470c-8b0d-caddab9635a9",
   "metadata": {},
   "source": [
    "<br><br><br>\n",
    "\n",
    "But I bet you'd rather use the Vector library."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b38f512e-e232-48c8-a573-7b6d14216b70",
   "metadata": {},
   "outputs": [],
   "source": [
    "muons = vector.arr(\n",
    "    ak.zip({\n",
    "        \"pt\": muon_kinematics.Muon_pt,\n",
    "        \"eta\": muon_kinematics.Muon_eta,\n",
    "        \"phi\": muon_kinematics.Muon_phi,\n",
    "        \"mass\": muon_kinematics.Muon_mass,\n",
    "        \"charge\": muon_kinematics.Muon_charge,  # not necessary for vector\n",
    "    })\n",
    ")\n",
    "muons"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7e33b9da-51ea-469d-bd2b-6206025c5395",
   "metadata": {},
   "source": [
    "<br><br><br>\n",
    "\n",
    "Note that each `Momentum4D` vector has a single numerical value for each coordinate:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4e713376-9e9f-481c-8b53-90b9263472c3",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(muons[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cd465cad-2eac-4367-b910-4bb6ec4fba9a",
   "metadata": {},
   "source": [
    "whereas the original `muon_kinematics` had _lists_ of each coordinate:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b2b0e602-1e5e-4761-a786-6111ac85e7eb",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(muon_kinematics[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "36de85d6-48f9-4c08-a358-45ecd6bf3c7d",
   "metadata": {},
   "source": [
    "<br><br><br>\n",
    "\n",
    "That's what the `ak.zip` function does: it combines nested lists of the same lengths into lists of records."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "92775652-61aa-4762-9804-3abffcb81f36",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(ak.Array({\"a\": [[1, 2], [], [3]], \"b\": [[1.1, 2.2], [], [3.3]]}))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2bfb29c9-56f7-4a3b-9494-1de45685eaca",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(ak.zip({\"a\": [[1, 2], [], [3]], \"b\": [[1.1, 2.2], [], [3.3]]}))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "988293c6-d604-4a62-b263-034e9d98c771",
   "metadata": {},
   "source": [
    "<br><br><br>\n",
    "\n",
    "The other function, `vector.arr`, interprets records with specially named fields as coordinates of vectors."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c023b4e0-9cf4-4c5a-ba04-1ca10dfc66b5",
   "metadata": {},
   "outputs": [],
   "source": [
    "vector.arr([[{\"x\": 1, \"y\": 1.1}, {\"x\": 2, \"y\": 2.2}], [], [{\"x\": 3, \"y\": 3.3}]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df452814-179a-4ba2-89b4-605d0c045382",
   "metadata": {},
   "outputs": [],
   "source": [
    "vector.arr([[{\"px\": 1, \"py\": 1.1}, {\"px\": 2, \"py\": 2.2}], [], [{\"px\": 3, \"py\": 3.3}]])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "90e75629-05e6-4283-a1f7-72c5ba0382fe",
   "metadata": {},
   "source": [
    "<br><br><br>\n",
    "\n",
    "Arrays of vectors have methods and properties to convert between coordinates."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e39cc280-7fa9-428b-b34f-dbee51ee9517",
   "metadata": {},
   "outputs": [],
   "source": [
    "muons.px"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "604d8ff3-f001-4b01-a5d8-477784807896",
   "metadata": {},
   "outputs": [],
   "source": [
    "muons.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ba2b513d-fc74-4c21-aedb-824fcb06f9f8",
   "metadata": {},
   "outputs": [],
   "source": [
    "muons.pz"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f8b7c473-0059-4689-b5ca-c50f97f29ce2",
   "metadata": {},
   "outputs": [],
   "source": [
    "muons.energy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bd7cb131-53a7-40fd-97c0-ad2320cbba74",
   "metadata": {},
   "outputs": [],
   "source": [
    "(\n",
    "    hist.Hist.new.Regular(100, -50, 50, name=\"pz\").Regular(100, 0, 15, name=\"pt\").Double()\n",
    "        .fill(ak.flatten(muons.pz), ak.flatten(muons.pt))\n",
    "        .plot2d_full()\n",
    ");"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ec4d6246-dafd-4e23-b070-8ba30cf7d7cb",
   "metadata": {},
   "source": [
    "<br><br><br><br><br>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0ca17b50-9f02-4533-b51c-18bf0e7f68b9",
   "metadata": {},
   "source": [
    "## Reconstructing the Z boson mass"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "085ddf3c-95d6-4ccb-ab3b-815d90769b9e",
   "metadata": {},
   "source": [
    "Now let's do some analysis-like stuff.\n",
    "\n",
    "If two muons are the decay products of a Z boson, the mass of the sum of those two muons is the Z boson mass."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b960eb2a-6402-424a-bf7c-ee5f66a705a4",
   "metadata": {},
   "outputs": [],
   "source": [
    "muons[0, 0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b80dea34-91b5-42a6-a22b-58214e3345d3",
   "metadata": {},
   "outputs": [],
   "source": [
    "muons[0, 1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6bd6cd76-ce6c-437c-a109-04f19e3919d9",
   "metadata": {},
   "outputs": [],
   "source": [
    "(muons[0, 0] + muons[0, 1]).mass"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "71bb0a10-8614-4888-a5b6-12d33d3e2a44",
   "metadata": {},
   "source": [
    "<br><br><br>\n",
    "\n",
    "We can vectorize this calculation by selecting all events (first dimension) and only muon number `0` or `1` (second dimension)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "360da8d0-443d-4a71-8d06-89e9712eb196",
   "metadata": {
    "tags": [
     "raises-exception"
    ]
   },
   "outputs": [],
   "source": [
    "muons[:, 0], muons[:, 1]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "16120ac5-6fd7-4bc6-a37e-2b82bc5da49e",
   "metadata": {},
   "source": [
    "<br><br><br>\n",
    "\n",
    "It didn't work because some events don't have two—or even one—muon.\n",
    "\n",
    "<br><br><br>\n",
    "\n",
    "\"How many muons does each event have?\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d09c0864-06f0-428c-8016-ac4e87b59b2c",
   "metadata": {},
   "outputs": [],
   "source": [
    "ak.num(muons)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c7e611ed-06e6-4fb7-8c43-0ddec173f20a",
   "metadata": {},
   "source": [
    "<br><br><br>\n",
    "\n",
    "\"Which events have at least two muons?\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0d3aacf1-e4ce-40a2-ad44-f6ac0ec01f33",
   "metadata": {},
   "outputs": [],
   "source": [
    "ak.num(muons) >= 2"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f6e2e050-d0c3-424b-a184-ef2b57f368b8",
   "metadata": {},
   "source": [
    "<br><br><br>\n",
    "\n",
    "\"Select events with at least two muons.\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "425c0456-ff78-49b0-8d55-1ae41c26ce07",
   "metadata": {},
   "outputs": [],
   "source": [
    "at_least_2 = muons[ak.num(muons) >= 2]\n",
    "\n",
    "ak.num(at_least_2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "59381a63-41b9-44d4-a0f9-937db56faeaa",
   "metadata": {},
   "source": [
    "(Notice that there are fewer of these events.)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "123b7001-03b8-4e68-8599-28728177ca64",
   "metadata": {},
   "source": [
    "<br><br><br>\n",
    "\n",
    "Now we can get the first and second muons of these events."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6588e052-e869-43c6-8752-c88e832bb347",
   "metadata": {},
   "outputs": [],
   "source": [
    "mu1 = at_least_2[:, 0]\n",
    "mu2 = at_least_2[:, 1]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9dded1e0-0f05-457f-b62f-c6734b319993",
   "metadata": {},
   "source": [
    "<br><br><br>\n",
    "\n",
    "And compute the invariant mass."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3a847427-84af-4b79-a93d-1bafa2122114",
   "metadata": {},
   "outputs": [],
   "source": [
    "(mu1 + mu2).mass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21e462b9-0188-438d-b75b-03a1a4118b49",
   "metadata": {},
   "outputs": [],
   "source": [
    "(\n",
    "    hist.Hist.new.Regular(100, 0, 150, name=\"muon pT\").Double()\n",
    "        .fill((mu1 + mu2).mass)\n",
    "        .plot()\n",
    ");"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b89b3508-eda3-4ef5-b7af-b834f0e632b3",
   "metadata": {},
   "source": [
    "Question to ponder: why did the above not need an `ak.flatten`?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "85e9de73-6b19-4c4b-936b-d60489b42715",
   "metadata": {},
   "source": [
    "<br><br><br><br><br>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "427e86d4-df30-4a4c-82f2-72f18524aa48",
   "metadata": {},
   "source": [
    "### Physics cuts: opposite sign muons"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d08c041c-1646-4093-8c8d-f2c50d3df5d5",
   "metadata": {},
   "source": [
    "In the above, we included all first and second muons, regardless of whether they are μ⁺μ⁻, μ⁺μ⁺, or μ⁻μ⁻. Only μ⁺μ⁻ can be decay products of a Z boson.\n",
    "\n",
    "<br><br><br>\n",
    "\n",
    "\"Which of `mu1` and `mu2` have opposite charges?\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f48ed0a4-c2e4-41b6-8cd4-6537787ad9ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "opposite_charges = mu1.charge != mu2.charge\n",
    "opposite_charges"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "572cf715-4540-42a6-8c94-3345ecbd904d",
   "metadata": {},
   "source": [
    "<br><br><br>\n",
    "\n",
    "\"What does the distribution look like if we only use the pairs with opposite charges?\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "73306f48-3be1-45e6-82d7-c5a1b1d79fe2",
   "metadata": {},
   "outputs": [],
   "source": [
    "(mu1[opposite_charges] + mu2[opposite_charges]).mass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "39583827-6a97-435a-a438-8f3d1df6128b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# without a cut (blue)\n",
    "(\n",
    "    hist.Hist.new.Regular(100, 0, 150, name=\"muon pT\").Double()\n",
    "        .fill((mu1 + mu2).mass)\n",
    "        .plot()\n",
    ");\n",
    "# with a cut (orange)\n",
    "(\n",
    "    hist.Hist.new.Regular(100, 0, 150, name=\"muon pT\").Double()\n",
    "        .fill((mu1[opposite_charges] + mu2[opposite_charges]).mass)\n",
    "        .plot()\n",
    ");"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ccfa82a3-c6e8-4058-b950-4401c689ac7a",
   "metadata": {},
   "source": [
    "<br><br><br><br><br>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d77fa7c9-6636-48b3-b3bd-97372d960511",
   "metadata": {},
   "source": [
    "## Particle combinatorics"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "409131c0-d646-4689-8cb6-067829c06902",
   "metadata": {},
   "source": [
    "In the above, we selected the first and second muons from each event, though many events have more than two good muons.\n",
    "\n",
    "What if the Z boson decays into muons `1` and `3`, instead of `0` and `1`? Or any other pair from {`0`, `1`, `2`, `3`}?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dadebe44-cc3d-4adf-b8f9-051b6458e4f9",
   "metadata": {},
   "source": [
    "<br><br><br>\n",
    "\n",
    "Awkward Array has two combinatorial primitives:\n",
    "\n",
    "<table style=\"margin-left: 0px\">\n",
    "    <tr style=\"background: white\"><td style=\"font-size: 1.75em; font-weight: bold; text-align: center\"><a href=\"https://awkward-array.readthedocs.io/en/latest/_auto/ak.cartesian.html\">ak.cartesian</a></td><td style=\"font-size: 1.75em; font-weight: bold; text-align: center\"><a href=\"https://awkward-array.readthedocs.io/en/latest/_auto/ak.combinations.html\">ak.combinations</a></td></tr>\n",
    "    <tr style=\"background: white\"><td><img src=\"img/cartoon-cartesian.svg\" width=\"300\"></td><td><img src=\"img/cartoon-combinations.svg\" width=\"300\"></td></tr>\n",
    "</table>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c494d140-862f-4672-958a-f3e7c70074ce",
   "metadata": {},
   "outputs": [],
   "source": [
    "cartesian_example = ak.cartesian(([[1, 2, 3], [], [4]], [[\"a\", \"b\"], [\"c\"], [\"d\", \"e\"]]))\n",
    "cartesian_example.tolist()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "256e5ca5-50ec-47c6-9fb7-e3173279e3d8",
   "metadata": {},
   "outputs": [],
   "source": [
    "combinations_example = ak.combinations([[1.1, 2.2, 3.3, 4.4], [], [5.5, 6.6]], 2)\n",
    "combinations_example.tolist()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f8fce59b-8e44-4ff0-b053-611526661417",
   "metadata": {},
   "source": [
    "<br><br><br>\n",
    "\n",
    "Note that these are creating new 2-tuple structures (parentheses), not lists (square brackets).\n",
    "\n",
    "Tuples can be separated, array-wise, with [ak.unzip](https://awkward-array.readthedocs.io/en/latest/_auto/ak.unzip.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "73bba511-1b21-4c9a-a974-6a2f44de0b46",
   "metadata": {},
   "outputs": [],
   "source": [
    "lefts, rights = ak.unzip(cartesian_example)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fe652bdc-6d36-472a-83b3-ffd7312d1ff2",
   "metadata": {},
   "outputs": [],
   "source": [
    "lefts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3da17ddb-60c5-4e63-9dcd-efc2bf1f76ff",
   "metadata": {},
   "outputs": [],
   "source": [
    "rights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "59fde816-b058-4fbf-9a74-24611cc4f09d",
   "metadata": {},
   "outputs": [],
   "source": [
    "ak.num(lefts), ak.num(rights)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e8ba9bc6-6b30-4862-9747-70d593b123be",
   "metadata": {},
   "source": [
    "<br><br>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d59f8aea-178d-45ac-83a5-ef9c4578baab",
   "metadata": {},
   "outputs": [],
   "source": [
    "lefts, rights = ak.unzip(combinations_example)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e0b99a32-fdfc-49a3-93a7-4f7b014cdf7f",
   "metadata": {},
   "outputs": [],
   "source": [
    "lefts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "00982afc-5722-47b9-8384-620cddbf67a4",
   "metadata": {},
   "outputs": [],
   "source": [
    "rights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "86706217-2a09-491f-90e0-cda6100161cb",
   "metadata": {},
   "outputs": [],
   "source": [
    "ak.num(lefts), ak.num(rights)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "964419b8-89b7-4c34-a491-2a21aaefb35d",
   "metadata": {},
   "source": [
    "<br><br><br>\n",
    "\n",
    "Applying [ak.combinations](https://awkward-array.readthedocs.io/en/latest/_auto/ak.combinations.html) to `muons`, we get every muon pair from each event.\n",
    "\n",
    "We don't even need to restrict ourselves to events with at least two muons."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "87200c56-ea3a-4a80-bd9a-d24ba8ea9618",
   "metadata": {},
   "outputs": [],
   "source": [
    "mu1, mu2 = ak.unzip(ak.combinations(muons, 2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a3c3cf71-e862-4012-865b-8aa9c0de8621",
   "metadata": {},
   "outputs": [],
   "source": [
    "mu1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5dc179f8-0b62-4f8e-ad88-338991155b7a",
   "metadata": {},
   "outputs": [],
   "source": [
    "mu2"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f7326693-baac-474a-a1e2-716eea8c55f8",
   "metadata": {},
   "source": [
    "<br><br>\n",
    "\n",
    "Invariant mass and plot..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "993cb92d-7ae9-4cfe-a93c-f57e51ef459d",
   "metadata": {},
   "outputs": [],
   "source": [
    "(mu1 + mu2).mass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6a0a8bdd-2753-488b-b6ff-ae57350c4dff",
   "metadata": {},
   "outputs": [],
   "source": [
    "(\n",
    "    hist.Hist.new.Regular(100, 0, 150, name=\"muon pT\").Double()\n",
    "        .fill(ak.flatten((mu1 + mu2).mass))\n",
    "        .plot()\n",
    ");"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0f7727b4-4e39-4fc5-80d2-d925ebcbedb8",
   "metadata": {},
   "source": [
    "<br><br>\n",
    "\n",
    "Now for the opposite charge cut..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a51f2e72-b7b9-40b5-ac26-b670dda163bf",
   "metadata": {},
   "outputs": [],
   "source": [
    "opposite_charges = mu1.charge != mu2.charge"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8f1193f7-ec7d-4c94-9109-89c122103e8a",
   "metadata": {},
   "outputs": [],
   "source": [
    "(\n",
    "    hist.Hist.new.Regular(100, 0, 150, name=\"muon pT\").Double()\n",
    "        .fill(ak.flatten((mu1 + mu2).mass))\n",
    "        .plot()\n",
    ");\n",
    "(\n",
    "    hist.Hist.new.Regular(100, 0, 150, name=\"muon pT\").Double()\n",
    "        .fill(ak.flatten((mu1[opposite_charges] + mu2[opposite_charges]).mass))\n",
    "        .plot()\n",
    ");"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1279385d-954d-4490-bd25-cd9f93ad44e1",
   "metadata": {},
   "source": [
    "In Higgs → ZZ, only one of the two Z bosons has a mass of 91 GeV.\n",
    "\n",
    "The other is \"off-shell,\" a \"virtual\" boson."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2c171b13-e478-4b1b-8675-b4da6fbcce68",
   "metadata": {},
   "source": [
    "<br><br><br><br><br>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8a8b35d9-266c-40c5-b19c-d5a508da9b02",
   "metadata": {},
   "source": [
    "## Higgs mass in μ⁺μ⁻μ⁺μ⁻"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1ca3f4e0-7478-4a71-9337-a2aadba5ec1e",
   "metadata": {},
   "source": [
    "When a Higgs decays into 4 muons, we don't need to know which Z boson is which.\n",
    "\n",
    "(If you want to try that exercise, see the [CoDaS-HEP 2022 tutorial](https://github.com/jpivarski-talks/2022-08-03-codas-hep-columnar-tutorial#readme).)\n",
    "\n",
    "The Higgs mass is formed from a combination of 4 muons."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c49597e8-78b8-4c4f-b7f2-34daa045eb71",
   "metadata": {},
   "outputs": [],
   "source": [
    "mu1, mu2, mu3, mu4 = ak.unzip(ak.combinations(muons, 4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9d367670-b268-4372-9588-5ffb3cbefb3c",
   "metadata": {},
   "outputs": [],
   "source": [
    "mu1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6c5f3263-3667-469b-a95b-afe7c0693406",
   "metadata": {},
   "outputs": [],
   "source": [
    "mu2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "72a3a507-18ee-488b-8ac4-73840e27c1a5",
   "metadata": {},
   "outputs": [],
   "source": [
    "mu3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8987d3ff-98f9-4e7d-8da4-5d0a314231ca",
   "metadata": {},
   "outputs": [],
   "source": [
    "mu4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b338982f-9f9d-4a22-95ee-022cce637e6f",
   "metadata": {},
   "outputs": [],
   "source": [
    "(mu1 + mu2 + mu3 + mu4).mass"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "05f30002-e962-44d0-abf7-aad9f7036971",
   "metadata": {},
   "source": [
    "<br><br><br>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14f33981-60df-4578-8865-d4224a54f0f7",
   "metadata": {},
   "outputs": [],
   "source": [
    "correct_charge = (mu1.charge + mu2.charge + mu3.charge + mu4.charge) == 0\n",
    "correct_charge"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a17a5521-0895-49fe-82ac-bca3c0eda400",
   "metadata": {},
   "outputs": [],
   "source": [
    "(\n",
    "    hist.Hist.new.Regular(100, 0, 150, name=\"muon pT\").Double()\n",
    "        .fill(ak.flatten((mu1 + mu2 + mu3 + mu4).mass))\n",
    "        .plot()\n",
    ");\n",
    "(\n",
    "    hist.Hist.new.Regular(100, 0, 150, name=\"muon pT\").Double()\n",
    "        .fill(ak.flatten((mu1[correct_charge] + mu2[correct_charge] + mu3[correct_charge] + mu4[correct_charge]).mass))\n",
    "        .plot()\n",
    ");"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d743155c-b9f0-4910-933e-5b3d03ba4e88",
   "metadata": {},
   "source": [
    "<br><br><br><br><br>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b577dcbb-ca58-4bc0-b53f-bd784e3d51cf",
   "metadata": {},
   "source": [
    "## Higgs mass in e⁺e⁻e⁺e⁻"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a0c6d3aa-e730-4042-9ec9-20d631a628b7",
   "metadata": {},
   "source": [
    "The 4 electron case is just like the 4 muon case, but we need to go back to the original file to extract electron quantities.\n",
    "\n",
    "Do you remember how?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "02cbf65c-dc2f-4a5d-800c-45e10780847e",
   "metadata": {},
   "outputs": [],
   "source": [
    "electron_kinematics = tree.arrays(filter_name=\"Electron_*\")\n",
    "electron_kinematics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "47078145-c6d2-40cf-9581-908603372c1e",
   "metadata": {},
   "outputs": [],
   "source": [
    "electrons = vector.arr(\n",
    "    ak.zip({\n",
    "        \"pt\": electron_kinematics.Electron_pt,\n",
    "        \"eta\": electron_kinematics.Electron_eta,\n",
    "        \"phi\": electron_kinematics.Electron_phi,\n",
    "        \"mass\": electron_kinematics.Electron_mass,\n",
    "        \"charge\": electron_kinematics.Electron_charge,  # not necessary for vector\n",
    "    })\n",
    ")\n",
    "electrons"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4c9b0231-85f1-4907-a65d-bb84dd010053",
   "metadata": {},
   "outputs": [],
   "source": [
    "e1, e2, e3, e4 = ak.unzip(ak.combinations(electrons, 4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0bee8b3c-843d-4c9b-be0d-a45769c586be",
   "metadata": {},
   "outputs": [],
   "source": [
    "correct_charge = (e1.charge + e2.charge + e3.charge + e4.charge) == 0\n",
    "correct_charge"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1aa47088-f458-429b-adda-ebff6decee39",
   "metadata": {},
   "outputs": [],
   "source": [
    "(\n",
    "    hist.Hist.new.Regular(100, 0, 150, name=\"muon pT\").Double()\n",
    "        .fill(ak.flatten((e1 + e2 + e3 + e4).mass))\n",
    "        .plot()\n",
    ");\n",
    "(\n",
    "    hist.Hist.new.Regular(100, 0, 150, name=\"muon pT\").Double()\n",
    "        .fill(ak.flatten((e1[correct_charge] + e2[correct_charge] + e3[correct_charge] + e4[correct_charge]).mass))\n",
    "        .plot()\n",
    ");"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3672f2c4-b1d7-4852-b7ec-c2414f6328f4",
   "metadata": {},
   "source": [
    "<br><br><br><br><br>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0f95b492-5b02-4f4c-9963-d9f5f1347e36",
   "metadata": {},
   "source": [
    "## Higgs mass in e⁺e⁻μ⁺μ⁻"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "83a64931-5aff-4773-90ab-bf6cf3c60e99",
   "metadata": {},
   "source": [
    "It may seem like the mixed case is harder, but it's not (especially if we want to identify intermediate Z bosons).\n",
    "\n",
    "Each of the 4 particles in the final state is distinguishable.\n",
    "\n",
    "Electrons and muons often come in different collections, while positive and negative charges often come in the same collection, but there's no reason not to make separate collections for both flavor _and_ charge."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "563c2a5e-faca-4939-ab44-1ae6910e2c9a",
   "metadata": {},
   "outputs": [],
   "source": [
    "positive_muons = muons[muons.charge > 0]\n",
    "negative_muons = muons[muons.charge < 0]\n",
    "positive_electrons = electrons[electrons.charge > 0]\n",
    "negative_electrons = electrons[electrons.charge < 0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2d76878f-2a90-499c-8f7e-bb2adce936ee",
   "metadata": {},
   "source": [
    "<br><br><br>\n",
    "\n",
    "Now we can construct the Higgs mass from a Cartesian product of these four collections."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d68274a5-00a9-4cfb-b212-dafeca607dfb",
   "metadata": {},
   "outputs": [],
   "source": [
    "mu1, mu2, e1, e2 = ak.unzip(\n",
    "    ak.cartesian([positive_muons, negative_muons, positive_electrons, negative_electrons])\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4f0efff5-6f46-4d7e-9604-22d60d6329f8",
   "metadata": {},
   "source": [
    "<br><br><br>\n",
    "\n",
    "There's no charge cut to apply."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19debe92-33cb-4b25-9235-c0f9f499143b",
   "metadata": {},
   "outputs": [],
   "source": [
    "(\n",
    "    hist.Hist.new.Regular(100, 0, 150, name=\"muon pT\").Double()\n",
    "        .fill(ak.flatten((mu1 + mu2 + e1 + e2).mass))\n",
    "        .plot()\n",
    ");\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e1f6fa56-2439-45b9-a9a1-0605ad3e5006",
   "metadata": {},
   "source": [
    "<br><br><br><br><br>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7c566437-ac80-4720-a4c7-e3f21c29baf3",
   "metadata": {},
   "source": [
    "## Questions?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6c9a00fb-d125-4fad-8563-60b237c49988",
   "metadata": {},
   "source": [
    "I hope you've enjoyed the tour.\n",
    "\n",
    "Now—time for questions!"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}