SROptimizationAnalysisFull.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Import modules"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "import time\n",
    "import json\n",
    "import pickle\n",
    "\n",
    "%matplotlib inline\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import uproot\n",
    "import concurrent.futures\n",
    "\n",
    "from XRootD import client\n",
    "from XRootD.client.flags import DirListFlags, StatInfoFlags, OpenFlags, MkDirFlags, QueryCode"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Import local classes from utils"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/javascript": [
       "if (!(\"Notification\" in window)) {\n",
       "    alert(\"This browser does not support desktop notifications, so the %%notify magic will not work.\");\n",
       "} else if (Notification.permission !== 'granted' && Notification.permission !== 'denied') {\n",
       "    Notification.requestPermission(function (permission) {\n",
       "        if(!('permission' in Notification)) {\n",
       "            Notification.permission = permission;\n",
       "        }\n",
       "    })\n",
       "}\n"
      ],
      "text/plain": [
       "<IPython.core.display.Javascript object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "%load_ext jupyternotify\n",
    "%autonotify -a 10\n",
    "\n",
    "%load_ext autoreload\n",
    "%autoreload 1\n",
    "%aimport utils.ObjectExtractor\n",
    "%aimport utils.PlotMaker\n",
    "%aimport utils.HistogramContainer\n",
    "%aimport utils.HistogramCalculator\n",
    "OE = utils.ObjectExtractor\n",
    "PM = utils.PlotMaker\n",
    "HCont = utils.HistogramContainer\n",
    "HCalc = utils.HistogramCalculator"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Global initialization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "sys.version_info(major=3, minor=6, micro=4, releaselevel='final', serial=0)\n"
     ]
    }
   ],
   "source": [
    "print(sys.version_info)\n",
    "# num_cores = multiprocessing.cpu_count()\n",
    "# print(num_cores)\n",
    "\n",
    "\n",
    "executor = concurrent.futures.ThreadPoolExecutor(48)\n",
    "\n",
    "plt.style.use('default')\n",
    "plt.rcParams['grid.linestyle'] = ':'\n",
    "plt.rcParams.update({'font.size': 10})\n",
    "\n",
    "numCuts = np.arange(0,16)\n",
    "\n",
    "branch_path = 'SREffi_gbm'\n",
    "\n",
    "labels = [ f'cut{cut}' for cut in numCuts ]\n",
    "cut_descriptions = [\n",
    "    'cut1: MET/MHT trigger fired (120 GeV)',\n",
    "    'cut2: j1 pT > 120 GeV, <= 2j w/ pT > 30 GeV',\n",
    "    'cut3: mu1 pT > 5 GeV, 0.1 < |dxy| < 700 cm',\n",
    "    'cut4: mu2 pT > 5 GeV, 0.1 < |dxy| < 700 cm',\n",
    "    r'cut5: $|\\Delta\\Phi$(MET, mu pair)| < 0.4',\n",
    "    'cut6: ',\n",
    "    'cut7: ',\n",
    "    'cut8: ',\n",
    "    'cut9: ',\n",
    "    'cut10:',\n",
    "    'cut11:',\n",
    "    'cut12:',\n",
    "    'cut13:',\n",
    "    'cut14:',\n",
    "    'cut15:',\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "all_plot_vars = ['metpt', 'jetpt','metjetphi', 'metmuphi', 'leadingmupt', 'subleadingmupt','recodr', 'recovertex']\n",
    "#all_plot_vars = ['reco_PF_MET_pt', 'reco_PF_jet_pt','metjetphi', 'metmuphi', 'leadingmupt', 'subleadingmupt', 'recodr', 'recovertex']\n",
    "plot_vars_metjet = all_plot_vars[0:4] #['metpt', 'jetpt', 'metjetphi', 'metmuphi']\n",
    "plot_vars_muons = all_plot_vars[4:8] #['leadingmupt', 'subleadingmupt', 'recodr', 'recovertex']\n",
    "cutflow_vars = ['cutflow_incl', 'cutflow_excl']\n",
    "all_plot_xlabels = [\n",
    "    'MET [GeV]', 'Leading jet pT [GeV]', '$\\Delta\\Phi$(MET, jet)', '$\\Delta\\Phi$(MET, di-muon)',\n",
    "    'Leading muon pT [GeV]', 'Subleading muon pT [GeV]', 'dR(muons)', 'Di-muon vertex [cm]']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "histos = {}\n",
    "all_bins = {}\n",
    "for plot_var in all_plot_vars:\n",
    "    histos[plot_var] = {}\n",
    "    all_bins[plot_var] = 60\n",
    "histos['cutflow_incl'] = {}\n",
    "histos['cutflow_excl'] = {}\n",
    "histos['sumgenwgt'] = {}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Process signal"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "masses = [('60p0','20p0'),('6p0','2p0'),('52p5','5p0'),('5p25','0p5')]\n",
    "def print_masses(mass):\n",
    "    return f'({float(mass[0].replace(\"p\",\".\"))-float(mass[1].replace(\"p\",\".\"))/2}, ' + \\\n",
    "           f'{float(mass[0].replace(\"p\",\".\"))+float(mass[1].replace(\"p\",\".\"))/2}) GeV'\n",
    "    \n",
    "mchis = dict([(mass[0], print_masses(mass)) for mass in masses])\n",
    "ctaus = [10]#, 10, 100, 1000]\n",
    "\n",
    "labels = [ f'cut{cut}' for cut in numCuts ]\n",
    "cut_descriptions = [\n",
    "    'cut1: MET/MHT trigger fired (120 GeV)',\n",
    "    'cut2: j1 pT > 120 GeV, <= 2j w/ pT > 30 GeV',\n",
    "    'cut3: mu1 pT > 5 GeV, 0.1 < |dxy| < 700 cm',\n",
    "    'cut4: mu2 pT > 5 GeV, 0.1 < |dxy| < 700 cm',\n",
    "    r'cut5: $|\\Delta\\Phi$(MET, mu pair)| < 0.4',\n",
    "]\n",
    "\n",
    "\n",
    "base_dir = '../iDM-analysis-skimming/washAOD/SROptimization/'\n",
    "def filename(Mchi, dMchi, ctau): \n",
    "    return base_dir + f'Mchi-{Mchi}_dMchi-{dMchi}_ctau-{ctau}.root'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "trees_gbm = dict()\n",
    "gen_info_gbm = dict()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1 of 4: ../iDM-analysis-skimming/washAOD/SROptimization/Mchi-60p0_dMchi-20p0_ctau-10.root\n",
      "2 of 4: ../iDM-analysis-skimming/washAOD/SROptimization/Mchi-6p0_dMchi-2p0_ctau-10.root\n",
      "3 of 4: ../iDM-analysis-skimming/washAOD/SROptimization/Mchi-52p5_dMchi-5p0_ctau-10.root\n",
      "4 of 4: ../iDM-analysis-skimming/washAOD/SROptimization/Mchi-5p25_dMchi-0p5_ctau-10.root\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[('Mchi-6p0_dMchi-2p0_ctau-1', 1),\n",
       " ('Mchi-6p0_dMchi-2p0_ctau-10', 1),\n",
       " ('Mchi-6p0_dMchi-2p0_ctau-100', 1),\n",
       " ('Mchi-6p0_dMchi-2p0_ctau-1000', 1),\n",
       " ('Mchi-60p0_dMchi-20p0_ctau-1', 1),\n",
       " ('Mchi-60p0_dMchi-20p0_ctau-10', 1),\n",
       " ('Mchi-60p0_dMchi-20p0_ctau-100', 1),\n",
       " ('Mchi-60p0_dMchi-20p0_ctau-1000', 1),\n",
       " ('Mchi-52p5_dMchi-5p0_ctau-1', 1),\n",
       " ('Mchi-52p5_dMchi-5p0_ctau-10', 1),\n",
       " ('Mchi-52p5_dMchi-5p0_ctau-100', 1),\n",
       " ('Mchi-52p5_dMchi-5p0_ctau-1000', 1),\n",
       " ('Mchi-5p25_dMchi-0p5_ctau-1', 1),\n",
       " ('Mchi-5p25_dMchi-0p5_ctau-10', 1),\n",
       " ('Mchi-5p25_dMchi-0p5_ctau-100', 1),\n",
       " ('Mchi-5p25_dMchi-0p5_ctau-1000', 1)]"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "## new signal input\n",
    "with open('config/sig.json') as sigs_json_file:\n",
    "    sigs = json.load(sigs_json_file)\n",
    "\n",
    "xrdfs = client.FileSystem(\"root://cmseos.fnal.gov/\")\n",
    "\n",
    "redirector = 'root://cmsxrootd.fnal.gov'\n",
    "sig_base_dir = '/store/group/lpcmetx/iDM/Ntuples/2018/signal/'\n",
    "files = {}\n",
    "\n",
    "for sig, properties in sigs.items():\n",
    "    files[sig] = []\n",
    "    status, listing = xrdfs.dirlist(f'{sig_base_dir}/{properties[\"dir\"]}', DirListFlags.STAT)\n",
    "    for file in listing:\n",
    "        if '.root' in file.name:\n",
    "            files[sig].append(f'{redirector}/{sig_base_dir}/{properties[\"dir\"]}/{file.name}')\n",
    "num_files_total = np.sum(np.array([len(files[i]) for i in files]))\n",
    "print(num_files_total)\n",
    "[(i, len(files[i])) for i in files]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processing signal Mchi-6p0_dMchi-2p0_ctau-1 (1/16)\n",
      "Reading file 1 of 1, global 1 of 16 (0.00%)\n",
      "Sample \"\" does not have either pileup or weight information\n",
      "Processing signal Mchi-6p0_dMchi-2p0_ctau-10 (2/16)\n",
      "Reading file 1 of 1, global 2 of 16 (6.25%)\n",
      "Sample \"\" does not have either pileup or weight information\n",
      "Processing signal Mchi-6p0_dMchi-2p0_ctau-100 (3/16)\n",
      "Reading file 1 of 1, global 3 of 16 (12.50%)\n",
      "Sample \"\" does not have either pileup or weight information\n",
      "Processing signal Mchi-6p0_dMchi-2p0_ctau-1000 (4/16)\n",
      "Reading file 1 of 1, global 4 of 16 (18.75%)\n",
      "Sample \"\" does not have either pileup or weight information\n",
      "Processing signal Mchi-60p0_dMchi-20p0_ctau-1 (5/16)\n",
      "Reading file 1 of 1, global 5 of 16 (25.00%)\n",
      "Sample \"\" does not have either pileup or weight information\n",
      "Processing signal Mchi-60p0_dMchi-20p0_ctau-10 (6/16)\n",
      "Reading file 1 of 1, global 6 of 16 (31.25%)\n",
      "Sample \"\" does not have either pileup or weight information\n",
      "Processing signal Mchi-60p0_dMchi-20p0_ctau-100 (7/16)\n",
      "Reading file 1 of 1, global 7 of 16 (37.50%)\n",
      "Sample \"\" does not have either pileup or weight information\n",
      "Processing signal Mchi-60p0_dMchi-20p0_ctau-1000 (8/16)\n",
      "Reading file 1 of 1, global 8 of 16 (43.75%)\n",
      "Sample \"\" does not have either pileup or weight information\n",
      "Processing signal Mchi-52p5_dMchi-5p0_ctau-1 (9/16)\n",
      "Reading file 1 of 1, global 9 of 16 (50.00%)\n",
      "Sample \"\" does not have either pileup or weight information\n",
      "Processing signal Mchi-52p5_dMchi-5p0_ctau-10 (10/16)\n",
      "Reading file 1 of 1, global 10 of 16 (56.25%)\n",
      "Sample \"\" does not have either pileup or weight information\n",
      "Processing signal Mchi-52p5_dMchi-5p0_ctau-100 (11/16)\n",
      "Reading file 1 of 1, global 11 of 16 (62.50%)\n",
      "Sample \"\" does not have either pileup or weight information\n",
      "Processing signal Mchi-52p5_dMchi-5p0_ctau-1000 (12/16)\n",
      "Reading file 1 of 1, global 12 of 16 (68.75%)\n",
      "Sample \"\" does not have either pileup or weight information\n",
      "Processing signal Mchi-5p25_dMchi-0p5_ctau-1 (13/16)\n",
      "Reading file 1 of 1, global 13 of 16 (75.00%)\n",
      "Sample \"\" does not have either pileup or weight information\n",
      "Processing signal Mchi-5p25_dMchi-0p5_ctau-10 (14/16)\n",
      "Reading file 1 of 1, global 14 of 16 (81.25%)\n",
      "Sample \"\" does not have either pileup or weight information\n",
      "Processing signal Mchi-5p25_dMchi-0p5_ctau-100 (15/16)\n",
      "Reading file 1 of 1, global 15 of 16 (87.50%)\n",
      "Sample \"\" does not have either pileup or weight information\n",
      "Processing signal Mchi-5p25_dMchi-0p5_ctau-1000 (16/16)\n",
      "Reading file 1 of 1, global 16 of 16 (93.75%)\n",
      "Sample \"\" does not have either pileup or weight information\n",
      "CPU times: user 3min 44s, sys: 43.5 s, total: 4min 27s\n",
      "Wall time: 6min 7s\n"
     ]
    },
    {
     "data": {
      "application/javascript": [
       "$(document).ready(\n",
       "    function() {\n",
       "        function appendUniqueDiv(){\n",
       "            // append a div with our uuid so we can check that it's already\n",
       "            // been sent and avoid duplicates on page reload\n",
       "            var notifiedDiv = document.createElement(\"div\")\n",
       "            notifiedDiv.id = \"63583018-1218-4f39-a79c-4cbadf93f7b3\"\n",
       "            element.append(notifiedDiv)\n",
       "        }\n",
       "\n",
       "        // only send notifications if the pageload is complete; this will\n",
       "        // help stop extra notifications when a saved notebook is loaded,\n",
       "        // which during testing gives us state \"interactive\", not \"complete\"\n",
       "        if (document.readyState === 'complete') {\n",
       "            // check for the div that signifies that the notification\n",
       "            // was already sent\n",
       "            if (document.getElementById(\"63583018-1218-4f39-a79c-4cbadf93f7b3\") === null) {\n",
       "                var notificationPayload = {\"requireInteraction\": false, \"icon\": \"/static/base/images/favicon.ico\", \"body\": \"Cell Execution Has Finished!!\", \"autonotify_after\": \"10\", \"autonotify_output\": false};\n",
       "                if (Notification.permission !== 'denied') {\n",
       "                    if (Notification.permission !== 'granted') { \n",
       "                        Notification.requestPermission(function (permission) {\n",
       "                            if(!('permission' in Notification)) {\n",
       "                                Notification.permission = permission\n",
       "                            }\n",
       "                        })\n",
       "                    }\n",
       "                    if (Notification.permission === 'granted') {\n",
       "                    var notification = new Notification(\"Jupyter Notebook\", notificationPayload)\n",
       "                    appendUniqueDiv()\n",
       "                    notification.onclick = function () {\n",
       "                        window.focus();\n",
       "                        this.close();\n",
       "                        };\n",
       "                    } \n",
       "                }     \n",
       "            }\n",
       "        }\n",
       "    }\n",
       ")\n"
      ],
      "text/plain": [
       "<IPython.core.display.Javascript object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "MAX_FILES=None # To load all possible files\n",
    "# MAX_FILES=1 # For testing\n",
    "\n",
    "### Initialize empty dicts of histograms \n",
    "# histos = {}\n",
    "# all_bins = {}\n",
    "# for plot_var in all_plot_vars:\n",
    "#     histos[plot_var] = {}\n",
    "#     all_bins[plot_var] = 60\n",
    "# histos['cutflow_incl'] = {}\n",
    "# histos['cutflow_excl'] = {}\n",
    "# histos['sumgenwgt'] = {}\n",
    "\n",
    "global_file_counter = 1\n",
    "\n",
    "for sig in sigs:\n",
    "    \n",
    "    print(f'Processing signal {sig} ({(list(sigs.keys())).index(sig)+1}/{len(sigs)})')\n",
    "    \n",
    "    ### Initialize histograms as empty HistogramContainers\n",
    "    for plot_var in all_plot_vars:\n",
    "        histos[plot_var][sig] = HCont.HistogramContainer(all_bins[plot_var])\n",
    "    histos['cutflow_incl'][sig] = np.zeros(len(numCuts))\n",
    "    histos['cutflow_excl'][sig] = np.zeros(len(numCuts))\n",
    "    histos['sumgenwgt'][sig] = 0.0\n",
    "    \n",
    "    ### Load data\n",
    "    file_counter = 1\n",
    "    for file in files[sig][slice(0,MAX_FILES)]:\n",
    "        \n",
    "        if file_counter % 10 == 1:\n",
    "            print(f'Reading file {file_counter} of {len(files[sig])},'\n",
    "                  f' global {global_file_counter} of {num_files_total}'\n",
    "                  f' ({100*(global_file_counter-1)/num_files_total:.2f}%)')\n",
    "            with open('histos_temp.dat', 'wb') as histos_file:\n",
    "                pickle.dump(histos, histos_file)\n",
    "        file_counter += 1\n",
    "        global_file_counter += 1\n",
    "        \n",
    "        ### Open ROOT file and get tree\n",
    "        tree = uproot.open(file)[branch_path + '/reco']\n",
    "        \n",
    "        ### Make pandas dataframes and create all objects that will be passed to histo functions\n",
    "        obj_extractor = OE.ObjectExtractor(tree)\n",
    "        objects = obj_extractor.get_all()\n",
    "            \n",
    "        ## Add to sum of genwgts\n",
    "        histos['sumgenwgt'][sig] += np.sum(objects['gen_wgt'])\n",
    "        \n",
    "        ### Calculate histograms and cutflows\n",
    "        histo_maker = HCalc.HistogramCalculator(objects, sig)\n",
    "            \n",
    "        ### Cutflows\n",
    "        incl, excl = histo_maker.cutflows()\n",
    "        histos['cutflow_incl'][sig] += incl\n",
    "        histos['cutflow_excl'][sig] += excl\n",
    "        \n",
    "        ### Histograms\n",
    "        for plot_var in all_plot_vars:\n",
    "            new_hist = eval(f'histo_maker.{plot_var}()')\n",
    "            histos[plot_var][sig] += new_hist\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mchi-6p0_dMchi-2p0_ctau-1 7053.7400881057265 45400.0\n",
      "Mchi-6p0_dMchi-2p0_ctau-10 7057.159857199524 45378.0\n",
      "Mchi-6p0_dMchi-2p0_ctau-100 7053.7400881057265 45400.0\n",
      "Mchi-6p0_dMchi-2p0_ctau-1000 7053.7400881057265 45400.0\n",
      "Mchi-60p0_dMchi-20p0_ctau-1 2279.205722216291 140505.0\n",
      "Mchi-60p0_dMchi-20p0_ctau-10 2278.4759871931697 140550.0\n",
      "Mchi-60p0_dMchi-20p0_ctau-100 2278.897554866073 140524.0\n",
      "Mchi-60p0_dMchi-20p0_ctau-1000 2276.5646771120655 140668.0\n",
      "Mchi-52p5_dMchi-5p0_ctau-1 2031.5401499676466 157634.0\n",
      "Mchi-52p5_dMchi-5p0_ctau-10 2033.2556618688136 157501.0\n",
      "Mchi-52p5_dMchi-5p0_ctau-100 2040.497763504989 156942.0\n",
      "Mchi-52p5_dMchi-5p0_ctau-1000 2036.4363613239643 157255.0\n",
      "Mchi-5p25_dMchi-0p5_ctau-1 5487.034593835136 58363.0\n",
      "Mchi-5p25_dMchi-0p5_ctau-10 5491.174402853272 58319.0\n",
      "Mchi-5p25_dMchi-0p5_ctau-100 5485.436793422405 58380.0\n",
      "Mchi-5p25_dMchi-0p5_ctau-1000 5485.436793422405 58380.0\n"
     ]
    }
   ],
   "source": [
    "luminosity = 59.97 * 1000 # 1/pb\n",
    "for sig, properties in sigs.items():\n",
    "    properties['weight'] = luminosity * properties['xsec'] / histos['sumgenwgt'][sig]\n",
    "#     except KeyError:\n",
    "#         properties['weight'] = 1\n",
    "for sig, properties in sigs.items():\n",
    "    try:\n",
    "        print(sig, luminosity * properties['xsec'] / histos['sumgenwgt'][sig], histos['sumgenwgt'][sig])\n",
    "    except KeyError: pass\n",
    "    \n",
    "with open('histos_signal_objects_gbm.dat', 'wb') as histos_file:\n",
    "    pickle.dump(histos, histos_file)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Process backgrounds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('config/bkgs2.json') as bkgs_json_file:\n",
    "    bkgs = json.load(bkgs_json_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "xrdfs = client.FileSystem(\"root://cmseos.fnal.gov/\")\n",
    "\n",
    "redirector = 'root://cmsxrootd.fnal.gov'\n",
    "bkg_base_dir = '/store/group/lpcmetx/iDM/Ntuples/2018/backgrounds'\n",
    "files = {}\n",
    "\n",
    "for bkg, properties in bkgs.items():\n",
    "    files[bkg] = []\n",
    "    status, listing = xrdfs.dirlist(f'{bkg_base_dir}/{properties[\"dir\"]}', DirListFlags.STAT)\n",
    "    for file in listing:\n",
    "        if '.root' in file.name:\n",
    "            files[bkg].append(f'{redirector}/{bkg_base_dir}/{properties[\"dir\"]}/{file.name}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1364\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[('DYJetsToLL', 490),\n",
       " ('QCD_bEnriched_HT100to200', 82),\n",
       " ('QCD_bEnriched_HT200to300', 63),\n",
       " ('QCD_bEnriched_HT300to500', 13),\n",
       " ('QCD_bEnriched_HT500to700', 24),\n",
       " ('QCD_bEnriched_HT700to1000', 11),\n",
       " ('QCD_bEnriched_HT1000to1500', 1),\n",
       " ('QCD_bEnriched_HT1500to2000', 1),\n",
       " ('QCD_bEnriched_HT2000toINF', 1),\n",
       " ('TTTo2L2Nu', 172),\n",
       " ('TTJets', 30),\n",
       " ('TT_diLept', 5),\n",
       " ('WJetsToLNu_HT-70To100', 168),\n",
       " ('WJetsToLNu_HT-100To200', 66),\n",
       " ('WJetsToLNu_HT-200To400', 37),\n",
       " ('WJetsToLNu_HT-400To600', 7),\n",
       " ('ZJetsToNuNu_HT-100To200', 54),\n",
       " ('ZJetsToNuNu_HT-200To400', 41),\n",
       " ('ZJetsToNuNu_HT-400To600', 36),\n",
       " ('ZJetsToNuNu_HT-600To800', 12),\n",
       " ('ZJetsToNuNu_HT-800To1200', 10),\n",
       " ('ZJetsToNuNu_HT-1200To2500', 1),\n",
       " ('ZJetsToNuNu_HT-2500ToInf', 1),\n",
       " ('WWJJToLNuLNu', 1),\n",
       " ('WWTo2L2Nu', 20),\n",
       " ('WZTo3LNu', 1),\n",
       " ('ZZTo2L2Nu', 16)]"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "num_files_total = np.sum(np.array([len(files[i]) for i in files]))\n",
    "print(num_files_total)\n",
    "[(i, len(files[i])) for i in files]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processing background DYJetsToLL (1/27)\n",
      "Reading file 1 of 490, global 1 of 1364 (0.00%)\n",
      "Processing background QCD_bEnriched_HT100to200 (2/27)\n",
      "Reading file 1 of 82, global 3 of 1364 (0.15%)\n",
      "Processing background QCD_bEnriched_HT200to300 (3/27)\n",
      "Reading file 1 of 63, global 5 of 1364 (0.29%)\n",
      "Processing background QCD_bEnriched_HT300to500 (4/27)\n",
      "Reading file 1 of 13, global 7 of 1364 (0.44%)\n",
      "Processing background QCD_bEnriched_HT500to700 (5/27)\n",
      "Reading file 1 of 24, global 9 of 1364 (0.59%)\n",
      "Processing background QCD_bEnriched_HT700to1000 (6/27)\n",
      "Reading file 1 of 11, global 11 of 1364 (0.73%)\n",
      "Processing background QCD_bEnriched_HT1000to1500 (7/27)\n",
      "Reading file 1 of 1, global 13 of 1364 (0.88%)\n",
      "Processing background QCD_bEnriched_HT1500to2000 (8/27)\n",
      "Reading file 1 of 1, global 14 of 1364 (0.95%)\n",
      "Processing background QCD_bEnriched_HT2000toINF (9/27)\n",
      "Reading file 1 of 1, global 15 of 1364 (1.03%)\n",
      "Processing background TTTo2L2Nu (10/27)\n",
      "Reading file 1 of 172, global 16 of 1364 (1.10%)\n",
      "Processing background TTJets (11/27)\n",
      "Reading file 1 of 30, global 18 of 1364 (1.25%)\n",
      "Processing background TT_diLept (12/27)\n",
      "Reading file 1 of 5, global 20 of 1364 (1.39%)\n",
      "Processing background WJetsToLNu_HT-70To100 (13/27)\n",
      "Reading file 1 of 168, global 22 of 1364 (1.54%)\n",
      "Processing background WJetsToLNu_HT-100To200 (14/27)\n",
      "Reading file 1 of 66, global 24 of 1364 (1.69%)\n",
      "Processing background WJetsToLNu_HT-200To400 (15/27)\n",
      "Reading file 1 of 37, global 26 of 1364 (1.83%)\n",
      "Processing background WJetsToLNu_HT-400To600 (16/27)\n",
      "Reading file 1 of 7, global 28 of 1364 (1.98%)\n",
      "Processing background ZJetsToNuNu_HT-100To200 (17/27)\n",
      "Reading file 1 of 54, global 30 of 1364 (2.13%)\n",
      "Processing background ZJetsToNuNu_HT-200To400 (18/27)\n",
      "Reading file 1 of 41, global 32 of 1364 (2.27%)\n",
      "Processing background ZJetsToNuNu_HT-400To600 (19/27)\n",
      "Reading file 1 of 36, global 34 of 1364 (2.42%)\n",
      "Processing background ZJetsToNuNu_HT-600To800 (20/27)\n",
      "Reading file 1 of 12, global 36 of 1364 (2.57%)\n",
      "Processing background ZJetsToNuNu_HT-800To1200 (21/27)\n",
      "Reading file 1 of 10, global 38 of 1364 (2.71%)\n",
      "Processing background ZJetsToNuNu_HT-1200To2500 (22/27)\n",
      "Reading file 1 of 1, global 40 of 1364 (2.86%)\n",
      "Processing background ZJetsToNuNu_HT-2500ToInf (23/27)\n",
      "Reading file 1 of 1, global 41 of 1364 (2.93%)\n",
      "Processing background WWJJToLNuLNu (24/27)\n",
      "Reading file 1 of 1, global 42 of 1364 (3.01%)\n",
      "Processing background WWTo2L2Nu (25/27)\n",
      "Reading file 1 of 20, global 43 of 1364 (3.08%)\n",
      "Processing background WZTo3LNu (26/27)\n",
      "Reading file 1 of 1, global 45 of 1364 (3.23%)\n",
      "Processing background ZZTo2L2Nu (27/27)\n",
      "Reading file 1 of 16, global 46 of 1364 (3.30%)\n",
      "CPU times: user 6min 40s, sys: 28.2 s, total: 7min 8s\n",
      "Wall time: 11min 12s\n"
     ]
    },
    {
     "data": {
      "application/javascript": [
       "$(document).ready(\n",
       "    function() {\n",
       "        function appendUniqueDiv(){\n",
       "            // append a div with our uuid so we can check that it's already\n",
       "            // been sent and avoid duplicates on page reload\n",
       "            var notifiedDiv = document.createElement(\"div\")\n",
       "            notifiedDiv.id = \"0cd58700-6a38-472d-8d4c-0a63207af637\"\n",
       "            element.append(notifiedDiv)\n",
       "        }\n",
       "\n",
       "        // only send notifications if the pageload is complete; this will\n",
       "        // help stop extra notifications when a saved notebook is loaded,\n",
       "        // which during testing gives us state \"interactive\", not \"complete\"\n",
       "        if (document.readyState === 'complete') {\n",
       "            // check for the div that signifies that the notification\n",
       "            // was already sent\n",
       "            if (document.getElementById(\"0cd58700-6a38-472d-8d4c-0a63207af637\") === null) {\n",
       "                var notificationPayload = {\"requireInteraction\": false, \"icon\": \"/static/base/images/favicon.ico\", \"body\": \"Cell Execution Has Finished!!\", \"autonotify_after\": \"10\", \"autonotify_output\": false};\n",
       "                if (Notification.permission !== 'denied') {\n",
       "                    if (Notification.permission !== 'granted') { \n",
       "                        Notification.requestPermission(function (permission) {\n",
       "                            if(!('permission' in Notification)) {\n",
       "                                Notification.permission = permission\n",
       "                            }\n",
       "                        })\n",
       "                    }\n",
       "                    if (Notification.permission === 'granted') {\n",
       "                    var notification = new Notification(\"Jupyter Notebook\", notificationPayload)\n",
       "                    appendUniqueDiv()\n",
       "                    notification.onclick = function () {\n",
       "                        window.focus();\n",
       "                        this.close();\n",
       "                        };\n",
       "                    } \n",
       "                }     \n",
       "            }\n",
       "        }\n",
       "    }\n",
       ")\n"
      ],
      "text/plain": [
       "<IPython.core.display.Javascript object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "# MAX_FILES=None # To load all possible files\n",
    "MAX_FILES=2 # For testing\n",
    "\n",
    "### Initialize empty dicts of histograms \n",
    "# histos = {}\n",
    "# all_bins = {}\n",
    "# for plot_var in all_plot_vars:\n",
    "#     histos[plot_var] = {}\n",
    "#     all_bins[plot_var] = 60\n",
    "# histos['cutflow_incl'] = {}\n",
    "# histos['cutflow_excl'] = {}\n",
    "# histos['sumgenwgt'] = {}\n",
    "\n",
    "global_file_counter = 1\n",
    "\n",
    "for bkg in bkgs:\n",
    "    \n",
    "    print(f'Processing background {bkg} ({(list(bkgs.keys())).index(bkg)+1}/{len(bkgs)})')\n",
    "    \n",
    "    ### Initialize histograms as empty HistogramContainers\n",
    "    for plot_var in all_plot_vars:\n",
    "        histos[plot_var][bkg] = HCont.HistogramContainer(all_bins[plot_var])\n",
    "    histos['cutflow_incl'][bkg] = np.zeros(len(numCuts))\n",
    "    histos['cutflow_excl'][bkg] = np.zeros(len(numCuts))\n",
    "    histos['sumgenwgt'][bkg] = 0.0\n",
    "    \n",
    "    ### Load data\n",
    "    file_counter = 1\n",
    "    for file in files[bkg][slice(0,MAX_FILES)]:\n",
    "        \n",
    "        if file_counter % 10 == 1:\n",
    "            print(f'Reading file {file_counter} of {len(files[bkg])},'\n",
    "                  f' global {global_file_counter} of {num_files_total}'\n",
    "                  f' ({100*(global_file_counter-1)/num_files_total:.2f}%)')\n",
    "            with open('histos_temp.dat', 'wb') as histos_file:\n",
    "                pickle.dump(histos, histos_file)\n",
    "        file_counter += 1\n",
    "        global_file_counter += 1\n",
    "        \n",
    "        ### Open ROOT file and get tree\n",
    "        tree = uproot.open(file)[branch_path + '/cutsTree']\n",
    "        \n",
    "        ### Make pandas dataframes and create all objects that will be passed to histo functions\n",
    "        obj_extractor = OE.ObjectExtractor(tree)\n",
    "        objects = obj_extractor.get_all()\n",
    "            \n",
    "        ## Add to sum of genwgts\n",
    "        histos['sumgenwgt'][bkg] += np.sum(objects['genwgt'])\n",
    "        \n",
    "        ### Calculate histograms and cutflows\n",
    "        histo_maker = HCalc.HistogramCalculator(objects, bkg)\n",
    "            \n",
    "        ### Cutflows\n",
    "        incl, excl = histo_maker.cutflows()\n",
    "        histos['cutflow_incl'][bkg] += incl\n",
    "        histos['cutflow_excl'][bkg] += excl\n",
    "        \n",
    "        ### Histograms\n",
    "        for plot_var in all_plot_vars:\n",
    "            new_hist = eval(f'histo_maker.{plot_var}()')\n",
    "            histos[plot_var][bkg] += new_hist"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('histos_temp.dat', 'wb') as histos_file:\n",
    "    pickle.dump(histos, histos_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "luminosity = 59.97 * 1000 # 1/pb\n",
    "for bkg, properties in bkgs.items():\n",
    "    properties['weight'] = luminosity * properties['xsec'] / histos['sumgenwgt'][bkg]\n",
    "#     except KeyError:\n",
    "#         properties['weight'] = 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "DYJetsToLL 1532.446448558506 208972.9140625\n",
      "QCD_bEnriched_HT100to200 121663.67032035017 553052.03125\n",
      "QCD_bEnriched_HT200to300 12747.545378030525 377154.5625\n",
      "QCD_bEnriched_HT300to500 939.2816048504089 1065600.875\n",
      "QCD_bEnriched_HT500to700 234.88200260330424 377873.140625\n",
      "QCD_bEnriched_HT700to1000 72.78128544965914 244473.546875\n",
      "QCD_bEnriched_HT1000to1500 5.668163836746107 491976.78125\n",
      "QCD_bEnriched_HT1500to2000 0.5451332217184982 409126.46875\n",
      "QCD_bEnriched_HT2000toINF 0.0979757226515774 396573.375\n",
      "TTTo2L2Nu 1.8448089475505214 22335856.0\n",
      "TTJets 122.50301775396018 240559.4453125\n",
      "TT_diLept 0.03809437870998642 114101496.0\n",
      "WJetsToLNu_HT-70To100 380.05058913085185 203239.6796875\n",
      "WJetsToLNu_HT-100To200 311.53985586101516 268146.140625\n",
      "WJetsToLNu_HT-200To400 61.5533895601258 395361.265625\n",
      "WJetsToLNu_HT-400To600 1.8279181987180197 1895963.5625\n",
      "ZJetsToNuNu_HT-100To200 92.10092407487133 197944.59375\n",
      "ZJetsToNuNu_HT-200To400 8.996377229024846 610673.78125\n",
      "ZJetsToNuNu_HT-400To600 3.554660664233398 221007.59375\n",
      "ZJetsToNuNu_HT-600To800 0.9479853540262179 206165.875\n",
      "ZJetsToNuNu_HT-800To1200 0.30041589979118105 298636.390625\n",
      "ZJetsToNuNu_HT-1200To2500 0.06038058799838905 340171.0\n",
      "ZJetsToNuNu_HT-2500ToInf 0.0009013426397961769 350168.84375\n",
      "WWJJToLNuLNu 0.2627504828990591 493681.71875\n",
      "WWTo2L2Nu 0.45607818158356417 1456916.0\n",
      "WZTo3LNu 4.7573253800545885 58717.921875\n",
      "ZZTo2L2Nu 0.08136568285220586 442815.3828125\n"
     ]
    }
   ],
   "source": [
    "for bkg, properties in bkgs.items():\n",
    "    try:\n",
    "        print(bkg, luminosity * properties['xsec'] / histos['sumgenwgt'][bkg], histos['sumgenwgt'][bkg])\n",
    "    except KeyError: pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# with open('histos_bkgs_objects_gbm.dat', 'wb') as histos_file:\n",
    "#     pickle.dump(histos, histos_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "cutFlowInclGrp = {}\n",
    "for grp in bkg_grps:\n",
    "    if '60p0' in grp or '5p25' in grp or '52p5' in grp or '6p0' in grp: continue\n",
    "    for bkg in bkg_grps[grp]:\n",
    "        if grp in cutFlowInclGrp.keys():\n",
    "            cutFlowInclGrp[grp] += histos['cutflow_incl'][bkg].astype(int)\n",
    "        else:\n",
    "            cutFlowInclGrp[grp] = histos['cutflow_incl'][bkg].astype(int)\n",
    "\n",
    "pd.DataFrame.from_dict(cutFlowInclGrp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cutFlowInclGrp2 = {}\n",
    "for grp in bkg_grps:\n",
    "#     if '60p0' in grp or '5p25' in grp or '52p5' in grp or '6p0' in grp: continue\n",
    "    for bkg in bkg_grps[grp]:\n",
    "        if grp in cutFlowInclGrp2.keys():\n",
    "            try:\n",
    "                cutFlowInclGrp2[grp] += (histos['cutflow_incl'][bkg]*bkgs[bkg]['weight']).astype(int)\n",
    "            except KeyError: pass\n",
    "        else:\n",
    "            try:\n",
    "                cutFlowInclGrp2[grp] = (histos['cutflow_incl'][bkg]*bkgs[bkg]['weight']).astype(int)\n",
    "            except KeyError: pass\n",
    "\n",
    "\n",
    "pd.DataFrame.from_dict(cutFlowInclGrp2)"
   ]
  },
  {
   "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.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}