03-hep.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "84697070-5cc7-4494-9611-bd3a3bb6a206",
   "metadata": {},
   "source": [
    "# A HEP analysis using dask\n",
    "\n",
    "In this notebook, we utilize dask alongside [uproot](https://uproot.readthedocs.io/en/latest/), [awkward](https://awkward-array.readthedocs.io/en/latest/), and [hist](https://github.com/scikit-hep/hist/) to analyze some CMS Run I open data.\n",
    "\n",
    "This analysis touches on a lot of the previous material, using dask delayed, bag, and dataframe objects in unison.\n",
    "The complicated task graph that we eventually build and run will allow us to discuss some of the visualization and monitoring tools to inspect workflow progress."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "974d7382-fd05-4c56-9171-735ee6f41073",
   "metadata": {},
   "outputs": [],
   "source": [
    "import distributed\n",
    "\n",
    "client = distributed.Client()\n",
    "client"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "04bb1911-0fb9-4d65-abd7-fb64cd7da482",
   "metadata": {},
   "source": [
    "We need to teach dask about our HEP-specific data types, so it can make intelligent decisions about when to cache and when to reproduce intermediate results. The `patch.py` module in this directory does this for us by registering the appropriate object size functions with the `dask.sizeof` utility."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "43dc9d3c-bbc3-4197-a8b9-61ea4af62526",
   "metadata": {},
   "outputs": [],
   "source": [
    "import patch\n",
    "client.upload_file(\"patch.py\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9f5df098-b56d-4c5a-9362-76d2819dcbfd",
   "metadata": {},
   "source": [
    "Our inputs consist of 20 open data files from CMS data collection in 2012. We set up a delayed routine to fetch the `Events` tree out of each file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "79466eec-e151-4b40-8318-c83310034cfa",
   "metadata": {},
   "outputs": [],
   "source": [
    "import dask\n",
    "from dask import delayed\n",
    "import uproot\n",
    "\n",
    "\n",
    "@delayed(pure=True)\n",
    "def get_tree(url):    \n",
    "    return uproot.open(url)[\"Events\"]\n",
    "\n",
    "urls = [\n",
    "    f\"root://eospublic.cern.ch//eos/root-eos/benchmark/CMSOpenDataDimuon/Run2012BC_DoubleMuParked_Muons_{i}.root\"\n",
    "    for i in range(1, 21)\n",
    "]\n",
    "inputs = [get_tree(url) for url in urls]\n",
    "inputs[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12d3f519-612b-4c70-8bc0-dd1529c8a999",
   "metadata": {},
   "source": [
    "We can compute one input and bring it back to our client to inspect the available columns in our data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c78d68f2-fa21-4ab1-81be-27ccde9fbf42",
   "metadata": {},
   "outputs": [],
   "source": [
    "inputs[0].compute().show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4d67e38a-8842-4d37-a046-9567106dc375",
   "metadata": {},
   "source": [
    "With this information in hand, we can create an awkward array structure representing these muons, as Lorentz vectors with a charge attribute."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0c399f47-a9b8-4d17-a95e-1619bf7a7710",
   "metadata": {},
   "outputs": [],
   "source": [
    "import awkward as ak\n",
    "# NB: https://github.com/scikit-hep/vector is now in beta\n",
    "from coffea.nanoevents.methods import vector\n",
    "\n",
    "\n",
    "@delayed(pure=True)\n",
    "def muon_struct(tree, entry_start, entry_stop):\n",
    "    def get(name):\n",
    "        return tree[name].array(entry_start=entry_start, entry_stop=entry_stop)\n",
    "\n",
    "    return ak.zip(\n",
    "        {\n",
    "            \"pt\": get(\"Muon_pt\"),\n",
    "            \"eta\": get(\"Muon_eta\"),\n",
    "            \"phi\": get(\"Muon_phi\"),\n",
    "            \"mass\": get(\"Muon_mass\"),\n",
    "            \"charge\": get(\"Muon_charge\"),\n",
    "        },\n",
    "        with_name=\"PtEtaPhiMLorentzVector\",\n",
    "        behavior=vector.behavior,\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "41005b5f-a524-4281-8c0c-1f6c477dd1d6",
   "metadata": {},
   "source": [
    "When handling large arrays, we need to be careful to make sure that the chunks each have appropriate data volumes, namely somewhere on the order of 1-100 megabytes. We need to chunk the files (each being over 2 GB) and a convenient way to do that with uproot is to set entry (collision event) ranges. But first, we need to know how many events are in each file:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8dcfff81-262e-411d-90dd-f0d8c4b64578",
   "metadata": {},
   "outputs": [],
   "source": [
    "nevents, = dask.compute([t.num_entries for t in inputs])\n",
    "nevents"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eb668fe3-3f86-47c7-ace7-e9b1f559fce7",
   "metadata": {},
   "source": [
    "Now we can declare a list of delayed objects representing chunks of event-muons:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f0a74a84-b7e6-4e27-ab84-c848b072de5a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "def chunks(n, target_size):\n",
    "    edges = np.linspace(0, n, n // target_size, dtype=int)\n",
    "    return zip(edges[:-1], edges[1:])\n",
    "\n",
    "chunksize = 400_000\n",
    "muons = [\n",
    "    muon_struct(tree, start, stop)\n",
    "    for tree, nev in zip(inputs, nevents)\n",
    "    for start, stop in chunks(nev, chunksize)\n",
    "]\n",
    "len(muons), muons[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1778c4d3-69bc-4462-80e6-86e58926cb68",
   "metadata": {},
   "source": [
    "[persist](https://docs.dask.org/en/latest/api.html?highlight=persist#dask.persist) turns lazy Dask collections into Dask collections with the same metadata, but now with their results fully computed or actively computing in the background.\n",
    "\n",
    "Let's persist the first chunk in the cluster so that we can quickly use it in subsequent testing. We'll also explicitly pull its value back to the client to do some preliminary inspection with matplotlib, plotting the number of muons per event as well as the invariant mass of dimuon pairs from events with exactly two muons."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b630c2e8-a745-4ebb-8fbd-7cc4981519cb",
   "metadata": {},
   "outputs": [],
   "source": [
    "muons[0] = muons[0].persist()\n",
    "muons_chunk = muons[0].compute()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e1cd3b74-8718-4600-bd9e-72f6df9870c0",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "plt.hist(ak.num(muons_chunk), bins=range(6));"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b9c5c15e-789a-4843-8670-99c8393f8dca",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.hist(muons_chunk[ak.num(muons_chunk)==2].sum().mass, bins=np.geomspace(0.2, 200, 200));\n",
    "plt.gca().set_xscale(\"log\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cc2b71d1-e7ca-4a10-9166-6024586270d7",
   "metadata": {},
   "source": [
    "These plots look nice, but we'd rather be filling histograms in our distributed cluster, aggregating them, and returning the total back to the client. Let's see how we can do that with dask:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f357963c-1c1c-47a7-9c0a-00648c322467",
   "metadata": {},
   "outputs": [],
   "source": [
    "import hist\n",
    "\n",
    "@delayed(pure=True)\n",
    "def nmuons_plot(muons_chunk):\n",
    "    return (\n",
    "        hist.Hist.new\n",
    "        .Reg(6, 0, 6, name=\"nMuons\")\n",
    "        .Double()\n",
    "        .fill(ak.num(muons_chunk))\n",
    "    )\n",
    "\n",
    "@delayed(pure=True)\n",
    "def filter_muons(muons_chunk):\n",
    "    return muons_chunk[\n",
    "        (ak.num(muons_chunk)==2)\n",
    "        & (ak.sum(muons_chunk.charge, axis=1) == 0)\n",
    "    ]\n",
    "\n",
    "@delayed(pure=True)\n",
    "def dimuon_cand(muons_chunk):\n",
    "    return muons_chunk.sum()\n",
    "\n",
    "@delayed(pure=True)\n",
    "def mass_plot(cand_chunk):\n",
    "    return (\n",
    "        hist.Hist.new\n",
    "        .Log(1000, 0.2, 200, name=\"mass\", label=\"Di-muon invariant mass\")\n",
    "        .Double()\n",
    "        .fill(cand_chunk.mass)\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2cd73597-9abc-4031-af50-22888db75776",
   "metadata": {},
   "source": [
    "We can use `dask.compute` to compute several delayed objects at once instead of calling the `.compute()` method on each in turn. This allows the scheduler to recognize shared components in the computation graph and ensure they are not executed more than necessary."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a5713d68-e4fa-44e1-888c-2ae6eeec5feb",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_nmu, test_mass = dask.compute(\n",
    "    nmuons_plot(muons[0]),\n",
    "    mass_plot(dimuon_cand(filter_muons(muons[0])))\n",
    ")\n",
    "display(test_nmu, test_mass)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8f6a2c97-b38e-499d-9a50-2ec55f420f8a",
   "metadata": {},
   "source": [
    "Suppose, in addition to the plots above, we also want to save a reduced dataset of dimuon candidate events, with each constituent muon's kinematic attributes corresponding to a column in a table. We can construct a pandas dataframe from our filtered muons and composite candidate:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d0a6f348-4bca-4784-9770-ca9088f4fc77",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "\n",
    "@delayed(pure=True)\n",
    "def cand_table(muons_chunk, cand_chunk):\n",
    "    mupos = ak.firsts(muons_chunk[muons_chunk.charge == 1])\n",
    "    muneg = ak.firsts(muons_chunk[muons_chunk.charge == -1])\n",
    "    return pd.DataFrame({\n",
    "        \"mass\": cand_chunk.mass,\n",
    "        \"pt\": cand_chunk.pt,\n",
    "        \"mu+_pt\": mupos.pt,\n",
    "        \"mu+_eta\": mupos.eta,\n",
    "        \"mu+_phi\": mupos.phi,\n",
    "        \"mu-_pt\": muneg.pt,\n",
    "        \"mu-_eta\": muneg.eta,\n",
    "        \"mu-_phi\": muneg.phi,\n",
    "    })\n",
    "\n",
    "\n",
    "cand_table(\n",
    "    filter_muons(muons[0]),\n",
    "    dimuon_cand(filter_muons(muons[0])),\n",
    ").compute()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b3ef31a2-0564-4cad-8391-07eaf6560909",
   "metadata": {},
   "source": [
    "As mentioned above, when calling `dask.compute` on multiple delayed objects, the scheduler first optimizes the task graph. Let's prove this. The `results` method creates the three delayed results we are interested in without any attempt to re-use shared components, as opposed to the `results_opt` method. Yet, if we visualize the task graph for either method applied to the first chunk of muon-events, we see no difference."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4b002392-5545-45e3-b7e8-296659653970",
   "metadata": {},
   "outputs": [],
   "source": [
    "def results(muons_chunk):\n",
    "    return (\n",
    "        nmuons_plot(muons_chunk),\n",
    "        mass_plot(dimuon_cand(filter_muons(muons_chunk))),\n",
    "        cand_table(\n",
    "            filter_muons(muons_chunk),\n",
    "            dimuon_cand(filter_muons(muons_chunk)),\n",
    "        ),\n",
    "    )\n",
    "\n",
    "def results_opt(muons_chunk):\n",
    "    filtered = filter_muons(muons_chunk)\n",
    "    cand = dimuon_cand(filtered)\n",
    "    return (\n",
    "        nmuons_plot(muons_chunk),\n",
    "        mass_plot(cand),\n",
    "        cand_table(filtered, cand),\n",
    "    )\n",
    "\n",
    "dask.visualize(results(muons[0]), optimize_graph=True)  # dask.compute always optimizes the graph"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3ec39065-baaf-4552-a29c-ac5a64a57d75",
   "metadata": {},
   "source": [
    "Now lets reduce the histogram results by using the dask bag [fold](https://docs.dask.org/en/latest/bag-api.html#dask.bag.Bag.fold) method and save our reduced dataset to parquet files using the dask dataframe [to_parquet](https://docs.dask.org/en/latest/dataframe-api.html?#dask.dataframe.to_parquet) after some additional filtering to only consider Z boson candidates (i.e. in a mass window 60-120 GeV).\n",
    "\n",
    "We do have to work a bit harder to construct a dask bag from delayed objects since it expects lists of items, and our results are `hist.Hist` objects. By adding a simple intermediate function `to_list` to group results together, we can then use [from_delayed](https://docs.dask.org/en/latest/bag-api.html#dask.bag.from_delayed). It may seem that `from_sequence` might do what we want, but be warned this will not properly treat sequences of delayed objects!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0cb435d9-2692-4a88-bdce-81978454a637",
   "metadata": {},
   "outputs": [],
   "source": [
    "import dask.dataframe as dd\n",
    "import dask.bag as db\n",
    "from operator import add\n",
    "\n",
    "@delayed(pure=True)\n",
    "def to_list(*args):\n",
    "    return list(args)\n",
    "\n",
    "\n",
    "result_set = np.array([results(chunk) for chunk in muons[:10]])\n",
    "nmuons_final = db.from_delayed(\n",
    "    to_list(*result_set[start:stop, 0])\n",
    "    for start, stop in chunks(len(result_set), 4)\n",
    ").fold(add)\n",
    "mass_final = db.from_delayed(\n",
    "    to_list(*result_set[start:stop, 1])\n",
    "    for start, stop in chunks(len(result_set), 4)\n",
    ").fold(add)\n",
    "table_full = dd.from_delayed(result_set[:, 2])\n",
    "table_skim = (\n",
    "    table_full[abs(table_full.mass - 90.0) < 30.0]\n",
    "    .repartition(2)\n",
    "    .to_parquet(\"zmmtable\", compute=False)\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "caa30762-1d75-4c77-b209-47480ed29833",
   "metadata": {},
   "outputs": [],
   "source": [
    "dask.visualize((nmuons_final, mass_final, table_skim), optimize_graph=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ef76aff1-e750-4bca-bb3a-9113c3efaf23",
   "metadata": {},
   "source": [
    "While this computation is running, we'll explore the dask dashboard (URL listed in `client` repr at the top of the notebook)\n",
    "This can also be seen from the jupyter lab by pasting the URL into the labextension pane."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ad1835fe-8881-4741-9689-5ba9f2e8f8a1",
   "metadata": {},
   "outputs": [],
   "source": [
    "res, = dask.compute((nmuons_final, mass_final, table_skim))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e80a34f5-ab26-4b28-a90a-6b51cfa69bce",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "res[0].plot(ax=ax)\n",
    "ax.set_xlabel(\"Number of muons\")\n",
    "ax.set_ylabel(\"Event counts\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6bbb915d-dc87-42be-841a-06687a2506a6",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "res[1].plot(ax=ax)\n",
    "ax.set_xscale(\"log\")\n",
    "ax.set_xlabel(\"Di-muon invariant mass [GeV]\")\n",
    "ax.set_ylabel(\"Event counts\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "841f6fcc-50c5-4011-8c30-57094437fe7f",
   "metadata": {},
   "outputs": [],
   "source": [
    "pd.read_parquet(\"zmmtable\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "91cb0193-d468-4f72-a843-e8da34431074",
   "metadata": {},
   "source": [
    "Before closing our client, take a look at the memory usage.\n",
    "We only hold a persistent reference to `muons[0]` at this point, yet the memory usage is high. After restarting, it drops.\n",
    "The main culprit in this instance is not leaks but rather the way that memory allocations work: memory is requested from the operating system by worker processes to satisfy peak usage, and is not relinquished immediately in all cases."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f91b749d-7986-4655-be4a-916b6e71c18c",
   "metadata": {},
   "outputs": [],
   "source": [
    "client.restart()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2f08016e-6b25-4f25-ace3-472e05fc37e3",
   "metadata": {},
   "outputs": [],
   "source": [
    "client.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e31b4a3b-cf75-439d-8a00-15c44d3b2e7b",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}