{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Coffea-Casa Benchmark Example 3"
   ]
  },
  {
   "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 (jet_pt) that has been masked to meet certain conditions (in this case, abs(jet eta) < 1).\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",
    "        jets = events.Jet\n",
    "\n",
    "        self.output['cutflow']['all events'] += ak.size(jets, axis=0)\n",
    "        self.output['cutflow']['number of jets'] += ak.sum(ak.num(jets))\n",
    "        \n",
    "        # We want jets with an abs(eta) < 1. Conditionals act on every value in an array in Coffea, so this is easy.\n",
    "        eta_max = (np.absolute(jets.eta) < 1)\n",
    "        # eta_max is a Boolean array, with True in the place of values where the condition is met, and False otherwise. We want to sum up all the Trues (=1) in each sublist, then sum up all the sublists to get the number of jets with pt > 20.\n",
    "        self.output['cutflow']['abs(eta) < 1'] += ak.sum(ak.sum(eta_max, axis=1))\n",
    "            \n",
    "        # We define good_jets as the actual jets we want to graph. We mask it with the jets that have abs(eta) < 1.\n",
    "        good_jets = jets[eta_max]\n",
    "        # good_jets is no longer a Boolean array, so we can't just sum up the True's. We count the amount of jets and sum that.\n",
    "        self.output['cutflow']['final good jets'] += ak.sum(ak.num(good_jets, axis=1))\n",
    "        \n",
    "        self.output['Jet_pt'].fill(dataset=dataset, Jet_pt=ak.flatten(good_jets.pt))\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 | 53.2s\r"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'bytesread': 2727541868,\n",
       " 'columns': ['nJet', 'Jet_pt', 'Jet_eta'],\n",
       " 'entries': 53446198,\n",
       " 'processtime': 994.5796327590942,\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 0x7f6d6428eee0>, errorbar=None, legend_artist=None)]"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAXQAAAESCAYAAADtzi4UAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjUuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/YYfK9AAAACXBIWXMAAAsTAAALEwEAmpwYAAATyklEQVR4nO3df7BcZ33f8fdHthJobI+SSBRZthChSqig9Y9qZLtkMg5NqO0x9UyH1KZtSGka1dR0YAbSujSDh06ZptMZpqEGyyo4iVswcYE4DqNAKIWxydjGluJfkvBUMaRWpWJjwEKFoTF8+8ce4dXV7r17r3bv7n3u+zWzc/ec8+zu1894PvvoOc85m6pCkrTyrZl2AZKk8TDQJakRBrokNcJAl6RGGOiS1AgDXZIaMdVAT3JbkqeTPD5i+7+X5ECS/Uk+Oun6JGklyTTXoSf5OeA4cHtVvXqBtluBO4HXVtU3k7ykqp5ejjolaSWY6gi9qu4BvtG/L8krknw6yd4k9yZ5ZXfo14APVNU3u9ca5pLUZxbn0HcD/7yq/gbwTuCD3f6fBn46yZ8kuT/JFVOrUJJm0JnTLqBfkrOAvwn8tyQndv9o9/dMYCtwOXAecG+SV1fVt5a5TEmaSTMV6PT+xfCtqrpwwLHDwP1V9RfAV5I8QS/gH1zG+iRpZs3UlEtVHaMX1r8EkJ4LusN3AT/f7V9PbwrmyWnUKUmzaNrLFu8A7gN+JsnhJL8K/APgV5M8AuwHrumafwZ4NskB4PPAr1fVs9OoW5Jm0YLLFpO8CLiH3lz2mcDHq+qmOW0C/BZwFfAd4B9V1b6JVCxJGmiUOfTv0Vv7fTzJWuCLSf6oqu7va3MlvfnsrcAlwC3dX0nSMlkw0Ks3hD/eba7tHnOH9dfQuziogPuTrEuysaqODnvf9evX15YtW5ZWtSStUnv37v16VW0YdGykVS5JzgD2An+F3sU9D8xpsgl4qm/7cLfvpEBPshPYCbB582Yeeuihkf4DJEk9Sf582LGRTopW1fe7pYTnATuSzL1MP6e+6pRRPFW1u6q2V9X2DRsGfsFIkpZoUatcuot4vgDMvUrzMHB+3/Z5wJHTKUyStDgLBnqSDUnWdc9fDPwC8OU5ze4G3tStG78UeG6++XNJ0viNMoe+Efjdbh59DXBnVX0qyfUAVbUL2ENvyeIhessW3zyheiVJQ4yyyuVR4KIB+3f1PS/ghvGWJklajJm69F+StHQGuiQ1wkCXpEbM2u1zF/SeP9zPgSPHhh7fdu453PT6Vy1jRZI0G1ZcoB84cowDR4+xbeM5px47OjzoJal1Ky7QAbZtPIff+6eXnbL/2lvvm0I1kjQbnEOXpEYY6JLUCANdkhphoEtSIwx0SWqEgS5JjTDQJakRBrokNcJAl6RGGOiS1AgDXZIaYaBLUiMMdElqhIEuSY0w0CWpEQa6JDXCQJekRhjoktQIA12SGmGgS1IjDHRJasSCgZ7k/CSfT3Iwyf4kbxvQ5vIkzyV5uHu8ezLlSpKGOXOENs8D76iqfUnOBvYm+WxVHZjT7t6qunr8JUqSRrHgCL2qjlbVvu75t4GDwKZJFyZJWpxFzaEn2QJcBDww4PBlSR5J8kdJXjXk9TuTPJTkoWeeeWbx1UqShho50JOcBXwCeHtVHZtzeB/wsqq6APhPwF2D3qOqdlfV9qravmHDhiWWLEkaZKRAT7KWXph/pKo+Ofd4VR2rquPd8z3A2iTrx1qpJGleo6xyCfBh4GBVvW9Im5d27Uiyo3vfZ8dZqCRpfqOscnkN8MvAY0ke7va9C9gMUFW7gDcAb0nyPPBd4LqqqvGXK0kaZsFAr6ovAlmgzc3AzeMqSpK0eF4pKkmNGGXKZUU5cPQY19563yn7t517Dje9fuBqSklqQlOBvu3ccwbuP3B07ipLSWpPU4E+bAQ+aMQuSa1xDl2SGmGgS1IjDHRJaoSBLkmNMNAlqREGuiQ1wkCXpEYY6JLUCANdkhphoEtSIwx0SWqEgS5JjTDQJakRBrokNcJAl6RGGOiS1AgDXZIaYaBLUiMMdElqhIEuSY0w0CWpEQa6JDXCQJekRiwY6EnOT/L5JAeT7E/ytgFtkuT9SQ4leTTJxZMpV5I0zJkjtHkeeEdV7UtyNrA3yWer6kBfmyuBrd3jEuCW7u/MOHD0GNfeet8p+7edew43vf5VU6hIksZrwUCvqqPA0e75t5McBDYB/YF+DXB7VRVwf5J1STZ2r526beeeM3D/gaPHlrkSSZqcUUboP5RkC3AR8MCcQ5uAp/q2D3f7Tgr0JDuBnQCbN29eZKlLN2wEPmjELkkr1cgnRZOcBXwCeHtVzR3aZsBL6pQdVburantVbd+wYcPiKpUkzWukQE+yll6Yf6SqPjmgyWHg/L7t84Ajp1+eJGlUo6xyCfBh4GBVvW9Is7uBN3WrXS4FnpuV+XNJWi1GmUN/DfDLwGNJHu72vQvYDFBVu4A9wFXAIeA7wJvHXqkkaV6jrHL5IoPnyPvbFHDDuIqSJC2eV4pKUiMMdElqhIEuSY0w0CWpEQa6JDXCQJekRhjoktQIA12SGmGgS1IjDHRJaoSBLkmNMNAlqRGL+sWiFvlbo5JasaoD3d8aldSSVR3o/taopJY4hy5JjTDQJakRBrokNcJAl6RGGOiS1AgDXZIaYaBLUiMMdElqhIEuSY0w0CWpEQa6JDXCQJekRiwY6EluS/J0kseHHL88yXNJHu4e7x5/mZKkhYxyt8XfAW4Gbp+nzb1VdfVYKpoR3idd0kqzYKBX1T1JtixDLTPD+6RLWonGdT/0y5I8AhwB3llV+8f0vlPhfdIlrUTjCPR9wMuq6niSq4C7gK2DGibZCewE2Lx58xg+WpJ0wmmvcqmqY1V1vHu+B1ibZP2QtrurantVbd+wYcPpfrQkqc9pB3qSlyZJ93xH957Pnu77SpIWZ8EplyR3AJcD65McBm4C1gJU1S7gDcBbkjwPfBe4rqpqYhVLkgYaZZXLGxc4fjO9ZY2SpCnySlFJaoSBLkmNMNAlqREGuiQ1YlxXiq4a3uNF0qwy0BfBe7xImmUG+iJ4jxdJs8w5dElqhIEuSY0w0CWpEQa6JDXCQJekRhjoktQIA12SGuE69DHxClJJ02agj4FXkEqaBQb6GHgFqaRZ4By6JDXCQJekRhjoktQIA12SGmGgS1IjXOUyYa5Pl7RcDPQJcn26pOVkoE+Q69MlLSfn0CWpEQa6JDViwUBPcluSp5M8PuR4krw/yaEkjya5ePxlSpIWMsoI/XeAK+Y5fiWwtXvsBG45/bIkSYu14EnRqronyZZ5mlwD3F5VBdyfZF2SjVV1dFxFtsjljJLGbRyrXDYBT/VtH+72GehDuJxR0iSMI9AzYF8NbJjspDctw+bNm8fw0SuTyxklTcI4VrkcBs7v2z4PODKoYVXtrqrtVbV9w4YNY/hoSdIJ4wj0u4E3datdLgWec/5ckpbfglMuSe4ALgfWJzkM3ASsBaiqXcAe4CrgEPAd4M2TKnY18GSppKUaZZXLGxc4XsANY6toFfNkqaTT4b1cZognSyWdDgN9hXAqRtJCDPQVwKkYSaMw0FcAp2IkjcJAX+GcipF0goG+gjkVI6mfgb6CORUjqZ+B3qhhUzHgdIzUKgO9QcOmYsDpGKllBnqD5ht9X3vrfZ5IlRploK8ynkiV2mWgrzLznUh15C6tbAa6AEfuUgsMdAEugZRaYKBrQU7FSCuDga55ORUjrRwGuublSVRp5TDQtSSO3KXZY6BrSRy5S7PHQNdYOXKXpsdA11i5/FGaHgNdy8apGGmyDHQtC6dipMkz0LUsPIkqTZ6Brqly5C6Nj4GuqXLkLo2Pga6Z5MhdWjwDXTPJkbu0eGtGaZTkiiRPJDmU5MYBxy9P8lySh7vHu8dfqtQL7W0bTx29Hzh6jANHHL1rdVtwhJ7kDOADwC8Ch4EHk9xdVQfmNL23qq6eQI3SDzlyl4YbZcplB3Coqp4ESPIx4BpgbqBLU+OcuzRaoG8CnurbPgxcMqDdZUkeAY4A76yq/XMbJNkJ7ATYvHnz4quVhljKyB0cvastowR6BuyrOdv7gJdV1fEkVwF3AVtPeVHVbmA3wPbt2+e+hzR2w0buAA985Rs88JVvDJx7N+i1Eo0S6IeB8/u2z6M3Cv+hqjrW93xPkg8mWV9VXx9PmdLSzBfK7/nD/QPD3GkarVSjBPqDwNYkLwf+N3Ad8Pf7GyR5KfC1qqokO+itnnl23MVK4+QJVrVmwUCvqueTvBX4DHAGcFtV7U9yfXd8F/AG4C1Jnge+C1xXVU6paEUaNk3jFI1m3UgXFlXVHmDPnH27+p7fDNw83tKk6RgWzMOmaAx6zQqvFJVGtNigdy5ey81Al06Tc/GaFQa6NCHOxWu5GejShDgXr+VmoEvLzKDXpBjo0oww6HW6DHRpxhn0GpWBLq1QBr3mMtClxowz6MGwX0kMdGmVWGzQw8JhP4hfANNjoEur3FLuSDmM0zrTZaBLGmqxAbzUaZ1B/AJYPANd0tgsZVpnEEf6S2OgS5o4R/rLw0CXNHOWY6Q/zEr+AjDQJa0Y4xrpD7OULwCYnS8BA11Ssyb9BQCz9a8AA12SOksJ2cV+CUzyh08MdEk6DYv9Ehj0gyfjsmZi7yxJWlYGuiQ1wkCXpEYY6JLUCANdkhphoEtSIwx0SWqEgS5JjRgp0JNckeSJJIeS3DjgeJK8vzv+aJKLx1+qJGk+CwZ6kjOADwBXAtuANybZNqfZlcDW7rETuGXMdUqSFjDKCH0HcKiqnqyq/wd8DLhmTptrgNur535gXZKNY65VkjSPUe7lsgl4qm/7MHDJCG02AUf7GyXZSW8ED3A8yROLqrbPndcv9ZVjtR74+rSLmCH2x8nsjxfYFydbf+f1S+6Plw07MEqgZ8C+WkIbqmo3sHuEz1wRkjxUVdunXcessD9OZn+8wL442aT6Y5Qpl8PA+X3b5wFHltBGkjRBowT6g8DWJC9P8iPAdcDdc9rcDbypW+1yKfBcVR2d+0aSpMlZcMqlqp5P8lbgM8AZwG1VtT/J9d3xXcAe4CrgEPAd4M2TK3mmNDN9NCb2x8nsjxfYFyebSH+k6pSpbknSCuSVopLUCANdkhphoI8oyflJPp/kYJL9Sd7W7f+JJJ9N8j+7vz8+7VonLcmLknwpySNdX7yn27/q+qJfkjOS/GmST3Xbq7Y/knw1yWNJHk7yULdvVfZHknVJPp7ky11+XDapvjDQR/c88I6q+qvApcAN3S0QbgQ+V1Vbgc912637HvDaqroAuBC4olvdtBr7ot/bgIN926u9P36+qi7sW2+9Wvvjt4BPV9UrgQvo/T8ymb6oKh9LeAB/APwi8ASwsdu3EXhi2rUtcz/8JWAfvauHV21f0Lv24nPAa4FPdftWc398FVg/Z9+q6w/gHOArdAtQJt0XjtCXIMkW4CLgAeAvV7fmvvv7kimWtmy66YWHgaeBz1bVqu2Lzn8E/gXwg759q7k/CvjjJHu7W37A6uyPnwKeAX67m477UJIfY0J9YaAvUpKzgE8Ab6+qY9OuZ1qq6vtVdSG9kemOJK+ecklTk+Rq4Omq2jvtWmbIa6rqYnp3Yr0hyc9Nu6ApORO4GLilqi4C/i8TnGoy0BchyVp6Yf6Rqvpkt/trJ+4s2f19elr1TUNVfQv4AnAFq7cvXgP8nSRfpXc30tcm+a+s3v6gqo50f58Gfp/eXVtXY38cBg53/4IF+Di9gJ9IXxjoI0oS4MPAwap6X9+hu4Ff6Z7/Cr259aYl2ZBkXff8xcAvAF9mFfYFQFX9q6o6r6q20Ls1xv+oqn/IKu2PJD+W5OwTz4HXAY+zCvujqv4P8FSSn+l2/S3gABPqC68UHVGSnwXuBR7jhXnSd9GbR78T2Az8L+CXquobUylymST568Dv0rsVxBrgzqr6N0l+klXWF3MluRx4Z1VdvVr7I8lP0RuVQ2/K4aNV9d5V3B8XAh8CfgR4kt6tUdYwgb4w0CWpEU65SFIjDHRJaoSBLkmNMNAlqREGuiQ1wkCXpEYY6Jp5SY7Pc2xdkn92Gu/9rnmOfb+7/eu53fZZSW5J8mfdfTn2Jvm1Bd7/C0n+9px9b0/ywSSv6N5/6H+ftBgGula6dcCSA53exWHDfLd6t3890m1/CPgmsLW7L8cVwE8s8P530Lt6tN91wB1V9Wfd/XCksTDQtWIk+fUkDyZ59MSPagC/CZwY6f6HIa+7PMk9SX4/yYEku5KsSfKbwIu7135kgc9+Bb37kfxGVf0AoKqeqap/v0B9HweuTvKjXZstwLnAF5feE9JgZ067AGkUSV4HbKUXqgHu7u7gdyPw6hFGujuAbcCfA58G/m5V3ZjkrSOOkl8FPHIizEetr6ruSfIleqP5P6A3Ov+98hJtTYAjdK0Ur+sef0rvBzVeSS9AR/Wlqnqyqr5PbxrkZ0+nmCT/uhvZn5iOma++/mmX67ptaewcoWulCPDvqurWk3b2pjBGMXdEvNgR8gHggiRrquoHVfVe4L19JzQH1te5C3hfkouBF1fVvkV+tjQSR+haKT4D/OPuB0ZIsinJS4BvA2eP8PodSV6eZA1wLS/MYf9Fd5/7eVXVIeAh4N8mOaOr4UX0gny++qiq4/TuGX8bjs41QQa6ZlqSM4HvVdUfAx8F7kvyGL2TjWdX1bPAnyR5fNhJ0c599E6gPk7vNx5P3N51N/DoQidFO/8E+EngUJK9wH8H/iXAsPr6XnsHvR8I/tgInyMtibfP1UxLcgHwn6tqx2m8x+V09yhf5OuOV9VZS/3cWfsctc8RumZWkuvpjWx/Y0olHOu/sGjcTlxYBHxtEu+v1ccRupqR5K8B/2XO7u9V1SXTqEdabga6JDXCKRdJaoSBLkmNMNAlqREGuiQ14v8D/A8wF5enXNIAAAAASUVORK5CYII=\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",
      "number of jets 170952895\n",
      "abs(eta) < 1 39563735\n",
      "final good jets 39563735\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
}