{
 "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.\n",
    "This also shows the analysis object syntax implemented by FCAT `JaggedCandidateArray`, and the usage of an accumulator class provided by FCAT."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "import uproot\n",
    "import awkward\n",
    "\n",
    "%matplotlib inline\n",
    "from coffea import hist\n",
    "from coffea.analysis_objects import JaggedCandidateArray\n",
    "from coffea.processor import defaultdict_accumulator"
   ]
  },
  {
   "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",
    "cutflow = defaultdict_accumulator(lambda: 0)\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",
    "    muons = JaggedCandidateArray.candidatesfromcounts(chunk['nMuon'],\n",
    "                                            pt=chunk['Muon_pt'].content,\n",
    "                                            eta=chunk['Muon_eta'].content,\n",
    "                                            phi=chunk['Muon_phi'].content,\n",
    "                                            mass=chunk['Muon_mass'].content,\n",
    "                                            charge=chunk['Muon_charge'].content,\n",
    "                                           )\n",
    "    \n",
    "    cutflow['all events'] += muons.size\n",
    "    twomuons = (muons.counts == 2)\n",
    "    cutflow['two muons'] += twomuons.sum()\n",
    "    opposite_charge = twomuons & (muons['charge'].prod() == -1)\n",
    "    cutflow['opposite charge'] += opposite_charge.sum()\n",
    "    dimuons = muons[opposite_charge].distincts()\n",
    "    masshist.fill(mass=dimuons.mass.flatten())\n",
    "    \n",
    "elapsed = time.time() - tstart\n",
    "print(dict(cutflow))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "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:\", cutflow['all events']/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.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}