{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Coffea-Casa Benchmark Example 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "%matplotlib inline\n",
    "import hist\n",
    "import coffea.processor as processor\n",
    "import awkward as ak\n",
    "from coffea.nanoevents import schemas"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# This program plots a per-event array (in this case, Jet pT). In Coffea, this is not very dissimilar from the event-level process.\n",
    "\n",
    "class Processor(processor.ProcessorABC):\n",
    "    def __init__(self):\n",
    "        dataset_axis = hist.axis.StrCategory(name=\"dataset\", label=\"\", categories=[], growth=True)\n",
    "        Jet_axis = hist.axis.Regular(name=\"Jet_pt\", label=\"Jet_pt [GeV]\", bins=50, start=15, stop=60)\n",
    "        \n",
    "        self.output = processor.dict_accumulator({\n",
    "            'Jet_pt': hist.Hist(dataset_axis, Jet_axis),\n",
    "            'cutflow': processor.defaultdict_accumulator(int)\n",
    "        })\n",
    "    \n",
    "    def process(self, events):\n",
    "        dataset = events.metadata['dataset']\n",
    "        Jet_pt = events.Jet.pt\n",
    "        # As before, we can get the number of events by checking the size of the array. To get the number of jets, which varies per event, though, we need to count up the number in each event, and then sum those counts (count subarray sizes, sum them).\n",
    "        self.output['cutflow']['all events'] += ak.size(Jet_pt, axis=0)\n",
    "        self.output['cutflow']['all jets'] += ak.sum(ak.count(Jet_pt, axis=1))\n",
    "        \n",
    "        # .flatten() removes jaggedness; plotting jagged data is meaningless, we just want to plot flat jets.\n",
    "        self.output['Jet_pt'].fill(dataset=dataset, Jet_pt=ak.flatten(Jet_pt))\n",
    "        \n",
    "        return self.output\n",
    "    \n",
    "    def postprocess(self, accumulator):\n",
    "        pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dask.distributed import Client\n",
    "\n",
    "client = Client(\"tls://localhost:8786\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[########################################] | 100% Completed | 47.6s\r"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'bytesread': 1612443666,\n",
       " 'columns': ['Jet_pt', 'nJet'],\n",
       " 'entries': 53446198,\n",
       " 'processtime': 432.0077931880951,\n",
       " 'chunks': 534}"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "fileset = {'SingleMu' : [\"root://eospublic.cern.ch//eos/root-eos/benchmark/Run2012B_SingleMu.root\"]}\n",
    "\n",
    "executor = processor.DaskExecutor(client=client)\n",
    "\n",
    "run = processor.Runner(executor=executor,\n",
    "                        schema=schemas.NanoAODSchema,\n",
    "                        savemetrics=True\n",
    "                      )\n",
    "\n",
    "output, metrics = run(fileset, \"Events\", processor_instance=Processor())\n",
    "\n",
    "metrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[StairsArtists(stairs=<matplotlib.patches.StepPatch object at 0x7f719cfab9d0>, errorbar=None, legend_artist=None)]"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAXQAAAESCAYAAADtzi4UAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjUuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/YYfK9AAAACXBIWXMAAAsTAAALEwEAmpwYAAASUUlEQVR4nO3de6ykd13H8fenFy7abgrsIt3etpANuKBtyWZLhZCKioUUmxhIlygoXtZiMZDgBdHQYCRiTFCk0naFilUoINeFlJugaTGldHdpaXsqugLaza52odKlQsDC1z/mWTs9e86ZOWdnzsz5zfuVTM48z/ObZ779pfk8v/3NM79JVSFJWvuOm3QBkqTRMNAlqREGuiQ1wkCXpEYY6JLUCANdkhox0UBPcm2Se5PcOUTbP01yW/f4lyTfWIUSJWnNyCTvQ0/ybOAB4LqqetoyXvcbwHlV9UtjK06S1piJjtCr6kbgvv59SZ6U5ONJ9iS5KclTFnjpi4HrV6VISVojTph0AQvYCVxWVf+a5HzgrcBzjhxMchZwNvCZCdUnSVNpqgI9yUnAjwF/l+TI7kfOa7YdeF9VfW81a5OkaTdVgU5vCugbVXXuEm22A5evTjmStHZM1W2LVXUY+EqSFwGk55wjx5M8GXgMcPOESpSkqTXp2xavpxfOT06yP8kvAz8H/HKS24G7gEv6XvJi4N3lEpGSdJSJ3rYoSRqdqZpykSSt3MQ+FF2/fn1t2rRpUm8vSWvSnj17vlZVGxY6NrFA37RpE7t3757U20vSmpTk3xc75pSLJDXCQJekRhjoktQIA12SGmGgS1IjDHRJaoSBLkmNMNAlqRHTtnzuQK//yF3MHTi86PEtG9dxxQueuooVSdJ0GDhCT3JGkn9IcneSu5K8coE2Fya5v+9HnF83nnJh7sBh5g4uHOhzBw8vGfaS1LJhRugPAq+uqr1JTgb2JPlUVc3Na3dTVV08+hKPtuXUdbzn1y44av+l17hMuqTZNXCEXlUHq2pv9/ybwN3AaeMuTJK0PMv6UDTJJuA84JYFDl+Q5PYkH0uy4CR2kh1JdifZfejQoeVXK0la1NCB3v2A8/uBV3U/FddvL3BWVZ0DvAX40ELnqKqdVbW1qrZu2LDg6o+SpBUaKtCTnEgvzN9ZVR+Yf7yqDlfVA93zG4ATk6wfaaWSpCUNc5dLgLcDd1fVmxZp84SuHUm2def9+igLlSQtbZi7XJ4JvAS4I8lt3b7XAmcCVNXVwAuBlyd5EPg2sN0fcpak1TUw0Kvqs0AGtLkSuHJURUmSls+v/ktSIwx0SWqEgS5JjTDQJakRBrokNcJAl6RGGOiS1AgDXZIaYaBLUiMMdElqhIEuSY0w0CWpEQa6JDXCQJekRhjoktQIA12SGmGgS1IjDHRJaoSBLkmNMNAlqREGuiQ1wkCXpEYY6JLUCANdkhphoEtSIwx0SWqEgS5JjTDQJakRBrokNcJAl6RGGOiS1IgTBjVIcgZwHfAE4PvAzqp687w2Ad4MPB/4FvCLVbV39OUONnfwMJdec/NR+7dsXMcVL3jqBCqSpNUxMNCBB4FXV9XeJCcDe5J8qqrm+to8D9jcPc4Hrur+rqotG9ctuH/u4OFVrkSSVt/AQK+qg8DB7vk3k9wNnAb0B/olwHVVVcDnkpyS5NTutatmsRH4QiN2SWrNsubQk2wCzgNumXfoNOCevu393T5J0ioZOtCTnAS8H3hVVc2fw8gCL6kFzrEjye4kuw8dOrS8SiVJSxoq0JOcSC/M31lVH1igyX7gjL7t04ED8xtV1c6q2lpVWzds2LCSeiVJixgY6N0dLG8H7q6qNy3SbBfw0vQ8A7h/tefPJWnWDXOXyzOBlwB3JLmt2/da4EyAqroauIHeLYv76N22+LKRVypJWtIwd7l8loXnyPvbFHD5qIqSJC2f3xSVpEYY6JLUCANdkhphoEtSIwx0SWqEgS5JjTDQJakRBrokNcJAl6RGGOiS1AgDXZIaYaBLUiMMdElqhIEuSY0w0CWpEQa6JDXCQJekRhjoktQIA12SGmGgS1IjDHRJaoSBLkmNMNAlqREnTLqA1TJ38DCXXnPzUfu3bFzHFS946gQqkqTRmolA37Jx3YL75w4eXuVKJGl8ZiLQFxuBLzRil6S1yjl0SWqEgS5JjTDQJakRBrokNcJAl6RGGOiS1IiBgZ7k2iT3JrlzkeMXJrk/yW3d43WjL1OSNMgw96G/A7gSuG6JNjdV1cUjqUiStCIDR+hVdSNw3yrUIkk6BqOaQ78gye1JPpZk0YVRkuxIsjvJ7kOHDo3orSVJMJpA3wucVVXnAG8BPrRYw6raWVVbq2rrhg0bRvDWkqQjjjnQq+pwVT3QPb8BODHJ+mOuTJK0LMcc6EmekCTd823dOb9+rOeVJC3PwLtcklwPXAisT7IfuAI4EaCqrgZeCLw8yYPAt4HtVVVjq1iStKCBgV5VLx5w/Ep6tzVKkibIb4pKUiMMdElqhIEuSY0w0CWpEQa6JDXCQJekRhjoktQIA12SGjHMeuhNmzt4mEuvufmo/Vs2ruOKFyy6cKQkTZ2ZDvQtG9ctuH/u4OFVrkSSjt1MB/piI/CFRuySNO2cQ5ekRhjoktQIA12SGmGgS1IjDHRJaoSBLkmNMNAlqREGuiQ1wkCXpEYY6JLUCANdkhphoEtSIwx0SWqEgS5JjTDQJakRBrokNcJAl6RGGOiS1IiZ/gm6pfjj0ZLWGgN9Af54tKS1yEBfgD8eLWktGjiHnuTaJPcmuXOR40ny50n2JflikqePvkxJ0iDDfCj6DuCiJY4/D9jcPXYAVx17WZKk5RoY6FV1I3DfEk0uAa6rns8BpyQ5dVQFSpKGM4rbFk8D7unb3t/tO0qSHUl2J9l96NChEby1JOmIUQR6FthXCzWsqp1VtbWqtm7YsGEEby1JOmIUgb4fOKNv+3TgwAjOK0lahlEE+i7gpd3dLs8A7q+qgyM4ryRpGQbeh57keuBCYH2S/cAVwIkAVXU1cAPwfGAf8C3gZeMqVpK0uIGBXlUvHnC8gMtHVpEkaUVcnEuSGmGgS1IjDHRJaoSBLkmNcLXFZXKddEnTykBfBtdJlzTNDPRlcJ10SdPMOXRJaoSBLkmNMNAlqREGuiQ1wkCXpEYY6JLUCANdkhphoEtSIwx0SWqEgS5JjfCr/yPiol2SJs1AHwEX7ZI0DQz0EXDRLknTwDl0SWqEgS5JjTDQJakRBrokNcJAl6RGGOiS1AgDXZIa4X3oY+Y3SCWtFgN9jPwGqaTVZKCPkd8glbSanEOXpEYY6JLUiKECPclFSb6UZF+S1yxw/MIk9ye5rXu8bvSlSpKWMnAOPcnxwF8APwXsB25Nsquq5uY1vamqLh5DjZKkIQwzQt8G7KuqL1fVd4F3A5eMtyxJ0nINc5fLacA9fdv7gfMXaHdBktuBA8BvVtVd8xsk2QHsADjzzDOXX21DvD9d0qgNE+hZYF/N294LnFVVDyR5PvAhYPNRL6raCewE2Lp16/xzzAzvT5c0DsME+n7gjL7t0+mNwv9fVR3ue35DkrcmWV9VXxtNmW3x/nRJ4zDMHPqtwOYkZyd5BLAd2NXfIMkTkqR7vq0779dHXawkaXEDR+hV9WCSVwCfAI4Hrq2qu5Jc1h2/Gngh8PIkDwLfBrZX1cxOqUjSJAz11f+qugG4Yd6+q/ueXwlcOdrSJEnL4VouU8a7XyStlIE+Rbz7RdKxMNCniHe/SDoWLs4lSY0w0CWpEU65rBF+WCppEAN9DfDDUknDMNDXAD8slTQM59AlqRGO0Nc459YlHWGgr2HOrUvqZ6CvYc6tS+pnoDdqsakYcDpGapWB3qDFpmLA6RipZQZ6g5YafV96zc1+kCo1ykCfMX6QKrXLQJ8xS32Q6shdWtsMdAGO3KUWGOgCHLlLLTDQtaTFRu63fOU+bvnKfcwdOHoEb9BLk2Gga0mLBfPrP3LXgmFu0EuTY6BrRQx6afoY6Bopg16aHANdq8Kgl8bPQNdEjTLoF+MFQLPCQNdUWm7QL2bQBcCwV0sMdK0pyw3fpS4AKxntL8YLg6aBga6mLRWyyx3tL8YLg6aFga6ZNarg9MKgaWGgS8eo5QvDSngxmRwDXZoS03ZhWIlJX0yWo8ULj4EuNWaSITXJi8lyTPrCM66LyVCBnuQi4M3A8cDbquqN846nO/584FvAL1bV3hHXKmnKrZUR7yQvPONcknpgoCc5HvgL4KeA/cCtSXZV1Vxfs+cBm7vH+cBV3V9JmjqTvPAs9uPto3DcEG22Afuq6stV9V3g3cAl89pcAlxXPZ8DTkly6ohrlSQtYZgpl9OAe/q293P06HuhNqcBB/sbJdkB7Og2H0jypWVV2+e9l630lSO1HvjapIuYIvbHw9kfD7EvHm79ey9bcX+ctdiBYQI9C+yrFbShqnYCO4d4zzUhye6q2jrpOqaF/fFw9sdD7IuHG1d/DDPlsh84o2/7dODACtpIksZomEC/Fdic5OwkjwC2A7vmtdkFvDQ9zwDur6qD808kSRqfgVMuVfVgklcAn6B32+K1VXVXksu641cDN9C7ZXEfvdsWXza+kqdKM9NHI2J/PJz98RD74uHG0h+pOmqqW5K0Bg0z5SJJWgMMdElqhIE+pCRnJPmHJHcnuSvJK7v9j03yqST/2v19zKRrHbckj0ry+SS3d33x+m7/zPVFvyTHJ/lCko922zPbH0m+muSOJLcl2d3tm8n+SHJKkvcl+ecuPy4YV18Y6MN7EHh1Vf0w8Azg8iRbgNcAn66qzcCnu+3WfQd4TlWdA5wLXNTd3TSLfdHvlcDdfduz3h8/XlXn9t1vPav98Wbg41X1FOAcev+PjKcvqsrHCh7Ah+mtb/Ml4NRu36nAlyZd2yr3ww8Ae+l9e3hm+4Ledy8+DTwH+Gi3b5b746vA+nn7Zq4/gHXAV+huQBl3XzhCX4Ekm4DzgFuAH6runvvu7+MnWNqq6aYXbgPuBT5VVTPbF50/A34b+H7fvlnujwI+mWRPt+QHzGZ/PBE4BPxVNx33tiQ/yJj6wkBfpiQnAe8HXlVV07/w85hU1feq6lx6I9NtSZ424ZImJsnFwL1VtWfStUyRZ1bV0+mtxHp5kmdPuqAJOQF4OnBVVZ0H/A9jnGoy0JchyYn0wvydVfWBbvd/HVlZsvt776Tqm4Sq+gbwj8BFzG5fPBP4mSRfpbca6XOS/C2z2x9U1YHu773AB+mt2jqL/bEf2N/9CxbgffQCfix9YaAPqfsRj7cDd1fVm/oO7QJ+oXv+C/Tm1puWZEOSU7rnjwZ+EvhnZrAvAKrqd6vq9KraRG9pjM9U1c8zo/2R5AeTnHzkOfBc4E5msD+q6j+Be5I8udv1E8AcY+oLvyk6pCTPAm4C7uChedLX0ptHfy9wJvAfwIuq6r6JFLlKkvwo8Nf0loI4DnhvVf1BkscxY30xX5ILgd+sqotntT+SPJHeqBx6Uw7vqqo3zHB/nAu8DXgE8GV6S6Mcxxj6wkCXpEY45SJJjTDQJakRBrokNcJAl6RGGOiS1AgDXZIaYaBr6iV5YIljpyT59WM492uXOPa9bvnXjd32SUmuSvJv3boce5L86oDz/2OSn56371VJ3prkSd35F/3vk5bDQNdadwqw4kCn9+WwxXy7esu/Hui23wb8N7C5W5fjIuCxA85/Pb1vj/bbDlxfVf/WrYcjjYSBrjUjyW8luTXJF4/8qAbwRuDISPdPFnndhUluTPLBJHNJrk5yXJI3Ao/uXvvOAe/9JHrrkfx+VX0foKoOVdUfD6jvfcDFSR7ZtdkEbAQ+u/KekBZ2wqQLkIaR5LnAZnqhGmBXt4Lfa4CnDTHS3QZsAf4d+Djws1X1miSvGHKU/FTg9iNhPmx9VXVjks/TG81/mN7o/D3lV7Q1Bo7QtVY8t3t8gd4PajyFXoAO6/NV9eWq+h69aZBnHUsxSX6vG9kfmY5Zqr7+aZft3bY0co7QtVYE+KOquuZhO3tTGMOYPyJe7gh5DjgnyXFV9f2qegPwhr4PNBesr/Mh4E1Jng48uqr2LvO9paE4Qtda8Qngl7ofGCHJaUkeD3wTOHmI129LcnaS44BLeWgO+3+7de6XVFX7gN3AHyY5vqvhUfSCfKn6qKoH6K0Zfy2OzjVGBrqmWpITgO9U1SeBdwE3J7mD3oeNJ1fV14F/SnLnYh+Kdm6m9wHqnfR+4/HI8q47gS8O+lC08yvA44B9SfYAfw/8DsBi9fW99np6PxD87iHeR1oRl8/VVEtyDvCXVbXtGM5xId0a5ct83QNVddJK33fa3kftc4SuqZXkMnoj29+fUAmH+79YNGpHvlgE/Nc4zq/Z4whdzUjyI8DfzNv9nao6fxL1SKvNQJekRjjlIkmNMNAlqREGuiQ1wkCXpEb8H1FMBwOj8dyEAAAAAElFTkSuQmCC\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "output['Jet_pt'].plot1d()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "all events 53446198\n",
      "all jets 170952895\n"
     ]
    }
   ],
   "source": [
    "for key, value in output['cutflow'].items():\n",
    "    print(key, value)"
   ]
  }
 ],
 "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.8.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}