{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Dimuon spectrum\n",
    "\n",
    "This code is a columnar adaptation of [a ROOT tutorial](https://root.cern.ch/doc/master/df102__NanoAODDimuonAnalysis_8py.html) showcasing the awkward array toolset, and utilizing FCAT histograms."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "import uproot\n",
    "import uproot_methods\n",
    "import awkward\n",
    "\n",
    "%matplotlib inline\n",
    "from coffea import hist"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# uproot supports xrootd, but its nicer to have them local (about 7 GB)\n",
    "!mkdir -p data\n",
    "!xrdcp root://eospublic.cern.ch//eos/root-eos/cms_opendata_2012_nanoaod/Run2012B_DoubleMuParked.root data/\n",
    "!xrdcp root://eospublic.cern.ch//eos/root-eos/cms_opendata_2012_nanoaod/Run2012C_DoubleMuParked.root data/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tstart = time.time()\n",
    "\n",
    "files = [\n",
    "    'data/Run2012B_DoubleMuParked.root',\n",
    "    'data/Run2012C_DoubleMuParked.root',\n",
    "]\n",
    "\n",
    "masshist = hist.Hist(\"Counts\", hist.Bin(\"mass\", r\"$m_{\\mu\\mu}$ [GeV]\", 30000, 0.25, 300))\n",
    "\n",
    "branches = ['nMuon', 'Muon_pt', 'Muon_eta', 'Muon_phi', 'Muon_mass', 'Muon_charge']\n",
    "for chunk in uproot.iterate(files, 'Events', branches=branches, entrysteps=500000, namedecode='ascii'):\n",
    "    p4 = uproot_methods.TLorentzVectorArray.from_ptetaphim(\n",
    "        chunk.pop('Muon_pt'),\n",
    "        chunk.pop('Muon_eta'),\n",
    "        chunk.pop('Muon_phi'),\n",
    "        chunk.pop('Muon_mass'),\n",
    "    )\n",
    "    muons = awkward.JaggedArray.zip(p4=p4, charge=chunk['Muon_charge'])\n",
    "\n",
    "    twomuons = (muons.counts == 2)\n",
    "    opposite_charge = (muons['charge'].prod() == -1)\n",
    "    dimuons = muons[twomuons & opposite_charge].distincts()\n",
    "    dimuon_mass = (dimuons.i0['p4'] + dimuons.i1['p4']).mass\n",
    "    masshist.fill(mass=dimuon_mass.flatten())\n",
    "    \n",
    "elapsed = time.time() - tstart"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax, _ = hist.plot1d(masshist)\n",
    "ax.set_xscale('log')\n",
    "ax.set_yscale('log')\n",
    "ax.set_ylim(0.1, 1e6)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Events/s:\", masshist.values()[()].sum()/elapsed)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}