bacon_gghbb_uproot.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#analysis boilerplate imports\n",
    "import fnal_column_analysis_tools\n",
    "from fnal_column_analysis_tools.lookup_tools import extractor,evaluator\n",
    "from fnal_column_analysis_tools.analysis_objects import JaggedCandidateArray\n",
    "from fnal_column_analysis_tools.striped import PhysicalColumnGroup,ColumnGroup,jaggedFromColumnGroup\n",
    "\n",
    "import numpy as np\n",
    "from awkward import JaggedArray,Table\n",
    "import uproot\n",
    "import time\n",
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#histogram display boilerplate\n",
    "from histbook import Hist, beside, groupby, below\n",
    "from histbook import bin as hbin\n",
    "from collections import OrderedDict\n",
    "from striped.job import SinglePointStripedSession as Session, IPythonDisplay\n",
    "import histbook as hb\n",
    "import math\n",
    "\n",
    "all_hists = OrderedDict()\n",
    "cat = hb.groupby(\"category\", keeporder=True)\n",
    "\n",
    "def add_1d(*args):\n",
    "    h = hb.Hist(hb.bin(*args), cat)\n",
    "    all_hists[args[0]] = h\n",
    "    return h\n",
    "\n",
    "display_opts = {'width': 200, 'height': 200}\n",
    "def buildDisplay(hists, cols=3):\n",
    "    rows = []\n",
    "    for iRow in range(int(math.ceil(len(hists)/cols))+1):\n",
    "        hrowname = hists.keys()[iRow*cols:(iRow+1)*cols]\n",
    "        hrow = hists.values()[iRow*cols:(iRow+1)*cols]\n",
    "        rows.append(hb.beside(*tuple(h.overlay(\"category\").step(n, **display_opts) for h,n in zip(hrow, hrowname))))\n",
    "    return IPythonDisplay(\n",
    "        hb.below(*tuple(rows))\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#make and fill a histogram of the masses\n",
    "add_1d(\"m_J\", 150,50,200)#, weight=\"weight\")\n",
    "add_1d(\"pT_J\", 300,200,800)#, weight=\"weight\")\n",
    "add_1d(\"deepDoubleB\",50,0,1)\n",
    "add_1d(\"tau21\",50,0,1)\n",
    "add_1d(\"tau32\",50,0,1)\n",
    "\n",
    "display = buildDisplay(all_hists,cols=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "extract = extractor()\n",
    "correction_data = open(\"dazsle_corr_files.txt\").readlines()\n",
    "extract.add_weight_sets(correction_data)\n",
    "#define the weights we want to keep\n",
    "keep_weights = [\"Rho2D\",\"ZJetsNLO\",\"WJetsNLO\",\n",
    "                \"PUWeight2016_central\",\"PUWeight2016_up\",\"PUWeight2016_down\",\n",
    "                \"PUWeight2017_central\",\"PUWeight2017_up\",\"PUWeight2017_down\",\n",
    "                \"data_obs_muCR4_denominator\",\"data_obs_muCR4_numerator\",\n",
    "                \"Mu50_PtEtaBins/efficienciesDATA/pt_abseta_DATA\",\n",
    "                \"NUM_SoftID_DEN_genTracks/abseta_pt_value\",\"NUM_SoftID_DEN_genTracks/abseta_pt_error\",\n",
    "                \"NUM_LooseRelIso_DEN_LooseID/abseta_pt_value\",\"NUM_LooseRelIso_DEN_LooseID/abseta_pt_error\",\n",
    "                \"electron_id_EffArea\",\n",
    "                \"photon_id_EA_CHad\",\"photon_id_EA_NHad\",\"photon_id_EA_Pho\",\n",
    "                #MC JEC/JUNC/JRES\n",
    "                'Spring16_25nsV10_MC_PtResolution_AK4PFPuppi',\n",
    "                'Spring16_25nsV10_MC_SF_AK4PFPuppi',\n",
    "                'Summer16_23Sep2016V3_MC_L1FastJet_AK4PFPuppi',\n",
    "                'Summer16_23Sep2016V3_MC_L2Relative_AK4PFPuppi',\n",
    "                'Summer16_23Sep2016V3_MC_L3Absolute_AK4PFPuppi',\n",
    "                #Data BCD JEC/JUNC\n",
    "                'Summer16_23Sep2016BCDV3_DATA_L1FastJet_AK4PFPuppi',\n",
    "                'Summer16_23Sep2016BCDV3_DATA_L2L3Residual_AK4PFPuppi',\n",
    "                'Summer16_23Sep2016BCDV3_DATA_L2Relative_AK4PFPuppi',\n",
    "                'Summer16_23Sep2016BCDV3_DATA_L3Absolute_AK4PFPuppi',\n",
    "                'Summer16_23Sep2016BCDV3_DATA_Uncertainty_AK4PFPuppi',\n",
    "                #Data EF JEC/JUNC\n",
    "                'Summer16_23Sep2016EFV3_DATA_L1FastJet_AK4PFPuppi',\n",
    "                'Summer16_23Sep2016EFV3_DATA_L2L3Residual_AK4PFPuppi',\n",
    "                'Summer16_23Sep2016EFV3_DATA_L2Relative_AK4PFPuppi',\n",
    "                'Summer16_23Sep2016EFV3_DATA_L3Absolute_AK4PFPuppi',\n",
    "                'Summer16_23Sep2016EFV3_DATA_Uncertainty_AK4PFPuppi',\n",
    "                #Data G JEC/JUNC\n",
    "                'Summer16_23Sep2016GV3_DATA_L1FastJet_AK4PFPuppi',\n",
    "                'Summer16_23Sep2016GV3_DATA_L2L3Residual_AK4PFPuppi',\n",
    "                'Summer16_23Sep2016GV3_DATA_L2Relative_AK4PFPuppi',\n",
    "                'Summer16_23Sep2016GV3_DATA_L3Absolute_AK4PFPuppi',\n",
    "                'Summer16_23Sep2016GV3_DATA_Uncertainty_AK4PFPuppi',\n",
    "                #Data H JEC/JUNC\n",
    "                'Summer16_23Sep2016HV3_DATA_L1FastJet_AK4PFPuppi',\n",
    "                'Summer16_23Sep2016HV3_DATA_L2L3Residual_AK4PFPuppi',\n",
    "                'Summer16_23Sep2016HV3_DATA_L2Relative_AK4PFPuppi',\n",
    "                'Summer16_23Sep2016HV3_DATA_L3Absolute_AK4PFPuppi',\n",
    "                'Summer16_23Sep2016HV3_DATA_Uncertainty_AK4PFPuppi',\n",
    "                ]\n",
    "extract.finalize(reduce_list=keep_weights)\n",
    "\n",
    "#make the evaluator\n",
    "evaluate = extract.make_evaluator()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#build a string of \"OR\" for list of triggers from json bits file\n",
    "import json\n",
    "def selectTriggers(imap,triggerMap):\n",
    "    if not imap =={}:\n",
    "        version     = imap['version']\n",
    "        hltNames    = imap['names']\n",
    "        branchName  = imap['branchName'] #triggerBits or moreTriggerBits\n",
    "        #validate input\n",
    "        if version in triggerMap:\n",
    "            NameToBitMap = triggerMap[version]\n",
    "        else:\n",
    "            print \"ERROR! Cannot find triggerbit map of the requested bit version =%s. Possible versions are: %s\"%(version, \",\".join(triggerMap.keys()))\n",
    "        tCuts = []\n",
    "        for hltName in hltNames:\n",
    "            print hltName\n",
    "            if hltName in NameToBitMap:\n",
    "                bitValue = int(2**float(NameToBitMap[hltName]))\n",
    "                tCuts.append(\"%s & %s\"%(branchName,bitValue))\n",
    "            else:\n",
    "                print \"ERROR! Cannot find the TriggerBit for %s\"%hltName \n",
    "        print \"Using trigger bits = \",\"||\".join(tCuts) \n",
    "        return \"||\".join(tCuts) \n",
    "    else:\n",
    "        return \"1\"\n",
    "\n",
    "triggerBitMaps = json.load(open(\"ggH/TriggerBitMap.json\"))\n",
    "\n",
    "triggerNames = {} #???\n",
    "triggerCut = selectTriggers(triggerNames,triggerBitMaps)\n",
    "if not triggerNames=={}:\n",
    "    print \"List of OR Triggers : \",triggerNames['names']\n",
    "    print \"Using trigger cuts  : \",triggerCut"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#trigger eff scale factors (todo: )\n",
    "def VtrigSF(ak8msd,ak8pt):\n",
    "    msd = ak8msd\n",
    "    pt = ak8pt\n",
    "    if isinstance(ak8msd,JaggedArray):\n",
    "        assert (ak8msd.offsets==ak8pt.offsets).all()\n",
    "        msd = ak8msd.flatten()\n",
    "        pt = ak8pt.flatten()\n",
    "    else:\n",
    "        assert ak8msd.size == ak8pt.size\n",
    "    msd = np.clip(msd,0,300)\n",
    "    pt = np.clip(pt,200,1000)\n",
    "    trig_num = evaluate['data_obs_muCR4_numerator'](msd,pt)\n",
    "    trig_den = evaluate['data_obs_muCR4_denominator'](msd,pt)\n",
    "    eff = trig_num/trig_den\n",
    "    if isinstance(ak8msd,JaggedArray):\n",
    "        eff = JaggedArray.fromoffsets(ak8pt.offsets,eff)\n",
    "    return eff"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#WJetsQQ and ZJetsQQ kfactors\n",
    "def VJetsQQ_kFactor(sampleName,lookup,genVPt):\n",
    "    kfactor = genVPt.ones_like()\n",
    "    if 'ZJetsToQQ_' in sampleName:\n",
    "        kfactor = lookup['ZJetsNLO'](np.clip(genVPt,250.,1200.))\n",
    "    elif 'WJetsToQQ_' in sampleName:\n",
    "        kfactor = lookup['WJetsNLO'](np.clip(genVPt,250.,1200.))\n",
    "    return kfactor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#PUPPI soft-drop mass correction\n",
    "def corrGEN(pt):\n",
    "    #self.corrGEN = ROOT.TF1(\"corrGEN\", \"[0]+[1]*pow(x*[2],-[3])\", 200, 3500)\n",
    "    #self.corrGEN.SetParameter(0, 1.00626)\n",
    "    #self.corrGEN.SetParameter(1, -1.06161)\n",
    "    #self.corrGEN.SetParameter(2, 0.0799900)\n",
    "    #self.corrGEN.SetParameter(3, 1.20454)\n",
    "    x = np.clip(pt,200,3500)\n",
    "    return 1.00626 - 1.06161*pow(x*0.0799900,-1.20454)\n",
    "\n",
    "def corrRECO_cen(pt):    \n",
    "    x = np.clip(pt,200,3500)\n",
    "    return 1.09302 + x*(-0.000150068+x*(3.44866e-07+x*(-2.68100e-10+x*(8.67440e-14+ x*(-1.00114e-17)))))\n",
    "\n",
    "def corrRECO_for(pt):  \n",
    "    x = np.clip(pt,200,3500)\n",
    "    return 1.27212 + x*(-0.000571640+x*(8.37289e-07+x*(-5.20433e-10+x*(1.45375e-13+ x*(-1.50389e-17)))))\n",
    "\n",
    "def PUPPIweight(ak8pt,ak8eta):\n",
    "    pt = ak8pt\n",
    "    eta = ak8eta\n",
    "    if isinstance(ak8pt,JaggedArray):\n",
    "        assert (ak8pt.offsets==ak8eta.offsets).all()\n",
    "        pt = ak8pt.flatten()\n",
    "        eta = ak8eta.flatten()\n",
    "    else:\n",
    "        assert ak8pt.size == ak8eta.size\n",
    "        \n",
    "    genCorr = corrGEN(pt)\n",
    "    recoCorr = np.ones_like(eta)\n",
    "    \n",
    "    etaCut = np.abs(eta) < 1.3\n",
    "    recoCorr[etaCut] = corrRECO_cen(pt[etaCut])\n",
    "    recoCorr[~etaCut] = corrRECO_for(pt[~etaCut])\n",
    "    \n",
    "    total = genCorr*recoCorr\n",
    "    \n",
    "    if isinstance(ak8pt,JaggedArray):\n",
    "        total = JaggedArray.fromoffsets(ak8pt.offsets,total)\n",
    "    \n",
    "    return total"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#muon IDs\n",
    "kPOGLooseMuon  =  1 << 0\n",
    "kPOGMediumMuon =  1 << 1\n",
    "kPOGTightMuon  =  1 << 2\n",
    "kPOGSoftMuon   =  1 << 3\n",
    "kPOGHighPtMuon =  1 << 4\n",
    "\n",
    "def calcMuonDeltaBetaIso(muon):\n",
    "    return muon.chHadIso + np.maximum(muon.neuHadIso + muon.gammaIso - 0.5*(muon.puIso), 0.0)\n",
    "\n",
    "def passMuonLooseSel(muon):\n",
    "    return (((muons.pogIDBits & kPOGLooseMuon) > 0) & (muon.deltaBetaIso < 0.25*muon.pt))\n",
    "\n",
    "def passMuonMediumSel(muon):\n",
    "    return (((muons.pogIDBits & kPOGMediumMuon) > 0) & (muon.deltaBetaIso < 0.25*muon.pt))\n",
    "\n",
    "def passMuonTightSel(muon):\n",
    "    return (((muons.pogIDBits & kPOGTightMuon) > 0) & (muon.deltaBetaIso < 0.25*muon.pt))\n",
    "\n",
    "def passHighPtMuonSel(muon):\n",
    "    return (((muons.pogIDBits & kPOGHighPtMuon) > 0) & (muon.deltaBetaIso < 0.25*muon.pt))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#electron IDs\n",
    "kEcalDriven    = 1 << 0\n",
    "kTrackerDriven = 1 << 1\n",
    "\n",
    "def calcElectronEffAreaIso(electron,rho):\n",
    "    return electron.chHadIso + np.maximum(0.0, (electron.gammaIso + electron.neuHadIso - rho.content*electron.EffArea) )\n",
    "    \n",
    "def passEleVetoSel(electron):\n",
    "    outs = (electron.isConv.content == 0)\n",
    "    etaCut = (np.abs(electron.scEta.content) < 1.479)    \n",
    "    notEta = ~etaCut\n",
    "    #barrel\n",
    "    outs[etaCut] &= ( ( electron.effAreaIso.content[etaCut] < \n",
    "                                      0.175*electron.pt.content[etaCut]      ) &\n",
    "                      ( electron.sieie.content[etaCut] < 0.01150             ) & \n",
    "                      ( np.abs(electron.dEtaInSeed.content[etaCut]) < 0.00749) &\n",
    "                      ( np.abs(electron.dPhiIn.content[etaCut]) < 0.22800    ) &\n",
    "                      ( electron.hovere.content[etaCut] < 0.35600            ) &\n",
    "                      ( np.abs(1.0 - electron.eoverp.content[etaCut]) < \n",
    "                               0.29900*electron.ecalEnergy.content[etaCut]   ) &\n",
    "                      ( electron.nMissingHits.content[etaCut] <= 2           )   )\n",
    "    #endcap\n",
    "    outs[notEta] &= ( ( electron.effAreaIso.content[notEta] < \n",
    "                                      0.159*electron.pt.content[notEta]      ) &\n",
    "                      ( electron.sieie.content[notEta] < 0.03700             ) & \n",
    "                      ( np.abs(electron.dEtaInSeed.content[notEta]) < 0.00895) &\n",
    "                      ( np.abs(electron.dPhiIn.content[notEta]) < 0.21300    ) &\n",
    "                      ( electron.hovere.content[notEta] < 0.21100            ) &\n",
    "                      ( np.abs(1.0 - electron.eoverp.content[notEta]) < \n",
    "                               0.15000*electron.ecalEnergy.content[notEta]   ) &\n",
    "                      ( electron.nMissingHits.content[notEta] <= 3           )   )\n",
    "    outs = JaggedArray.fromoffsets(electron.pt.offsets,outs)    \n",
    "    return outs\n",
    "\n",
    "def passEleLooseSel(electron):\n",
    "    outs = (electron.isConv.content == 0)\n",
    "    etaCut = (np.abs(electron.scEta.content) < 1.479)    \n",
    "    notEta = ~etaCut\n",
    "    #barrel\n",
    "    outs[etaCut] &= ( ( electron.effAreaIso.content[etaCut] < \n",
    "                                      0.0994*electron.pt.content[etaCut]     ) &\n",
    "                      ( electron.sieie.content[etaCut] < 0.01100             ) & \n",
    "                      ( np.abs(electron.dEtaInSeed.content[etaCut]) < 0.00477) &\n",
    "                      ( np.abs(electron.dPhiIn.content[etaCut]) < 0.22200    ) &\n",
    "                      ( electron.hovere.content[etaCut] < 0.29800            ) &\n",
    "                      ( np.abs(1.0 - electron.eoverp.content[etaCut]) < \n",
    "                               0.24100*electron.ecalEnergy.content[etaCut]   ) &\n",
    "                      ( electron.nMissingHits.content[etaCut] <= 1           )   )\n",
    "    #endcap\n",
    "    outs[notEta] &= ( ( electron.effAreaIso.content[notEta] < \n",
    "                                      0.107*electron.pt.content[notEta]      ) &\n",
    "                      ( electron.sieie.content[notEta] < 0.03700             ) & \n",
    "                      ( np.abs(electron.dEtaInSeed.content[notEta]) < 0.00895) &\n",
    "                      ( np.abs(electron.dPhiIn.content[notEta]) < 0.21300    ) &\n",
    "                      ( electron.hovere.content[notEta] < 0.10100            ) &\n",
    "                      ( np.abs(1.0 - electron.eoverp.content[notEta]) < \n",
    "                               0.14000*electron.ecalEnergy.content[notEta]   ) &\n",
    "                      ( electron.nMissingHits.content[notEta] <= 1           )   )\n",
    "    outs = JaggedArray.fromoffsets(electron.pt.offsets,outs)\n",
    "    return outs\n",
    "\n",
    "def passEleMediumSel(electron):\n",
    "    outs = (electron.isConv.content == 0)\n",
    "    etaCut = (np.abs(electron.scEta.content) < 1.479)    \n",
    "    notEta = ~etaCut\n",
    "    #barrel\n",
    "    outs[etaCut] &= ( ( electron.effAreaIso.content[etaCut] < \n",
    "                                      0.0695*electron.pt.content[etaCut]     ) &\n",
    "                      ( electron.sieie.content[etaCut] < 0.00998             ) & \n",
    "                      ( np.abs(electron.dEtaInSeed.content[etaCut]) < 0.00311) &\n",
    "                      ( np.abs(electron.dPhiIn.content[etaCut]) < 0.10300    ) &\n",
    "                      ( electron.hovere.content[etaCut] < 0.25300            ) &\n",
    "                      ( np.abs(1.0 - electron.eoverp.content[etaCut]) < \n",
    "                               0.13400*electron.ecalEnergy.content[etaCut]   ) &\n",
    "                      ( electron.nMissingHits.content[etaCut] <= 1           )   )\n",
    "    #endcap\n",
    "    outs[notEta] &= ( ( electron.effAreaIso.content[notEta] < \n",
    "                                      0.0821*electron.pt.content[notEta]     ) &\n",
    "                      ( electron.sieie.content[notEta] < 0.02980             ) & \n",
    "                      ( np.abs(electron.dEtaInSeed.content[notEta]) < 0.00609) &\n",
    "                      ( np.abs(electron.dPhiIn.content[notEta]) < 0.04500    ) &\n",
    "                      ( electron.hovere.content[notEta] < 0.08780            ) &\n",
    "                      ( np.abs(1.0 - electron.eoverp.content[notEta]) < \n",
    "                               0.13000*electron.ecalEnergy.content[notEta]   ) &\n",
    "                      ( electron.nMissingHits.content[notEta] <= 1           )   )\n",
    "    outs = JaggedArray.fromoffsets(electron.pt.offsets,outs)\n",
    "    return outs\n",
    "\n",
    "def passEleTightSel(electron):\n",
    "    outs = (electron.isConv.content == 0)\n",
    "    etaCut = (np.abs(electron.scEta.content) < 1.479)    \n",
    "    notEta = ~etaCut\n",
    "    #barrel\n",
    "    outs[etaCut] &= ( ( electron.effAreaIso.content[etaCut] < \n",
    "                                      0.0588*electron.pt.content[etaCut]     ) &\n",
    "                      ( electron.sieie.content[etaCut] < 0.00998             ) & \n",
    "                      ( np.abs(electron.dEtaInSeed.content[etaCut]) < 0.00308) &\n",
    "                      ( np.abs(electron.dPhiIn.content[etaCut]) < 0.08160    ) &\n",
    "                      ( electron.hovere.content[etaCut] < 0.04140            ) &\n",
    "                      ( np.abs(1.0 - electron.eoverp.content[etaCut]) < \n",
    "                               0.12900*electron.ecalEnergy.content[etaCut]   ) &\n",
    "                      ( electron.nMissingHits.content[etaCut] <= 1           )   )\n",
    "    #endcap\n",
    "    outs[notEta] &= ( ( electron.effAreaIso.content[notEta] < \n",
    "                                      0.0571*electron.pt.content[notEta]     ) &\n",
    "                      ( electron.sieie.content[notEta] < 0.02920             ) & \n",
    "                      ( np.abs(electron.dEtaInSeed.content[notEta]) < 0.00605) &\n",
    "                      ( np.abs(electron.dPhiIn.content[notEta]) < 0.03940    ) &\n",
    "                      ( electron.hovere.content[notEta] < 0.06410            ) &\n",
    "                      ( np.abs(1.0 - electron.eoverp.content[notEta]) < \n",
    "                               0.12900*electron.ecalEnergy.content[notEta]   ) &\n",
    "                      ( electron.nMissingHits.content[notEta] <= 1           )   )\n",
    "    outs = JaggedArray.fromoffsets(electron.pt.offsets,outs)\n",
    "    return outs\n",
    "\n",
    "def passEleHEEPSel(electron,rho,met):\n",
    "    outs = ( (electron.typeBits.content & kEcalDriven) > 0 )\n",
    "    outs &= (met <= 35)\n",
    "    absSCEta = np.abs(electron.scEta.content)\n",
    "    etaBar = (absSCEta < 1.442)    \n",
    "    etaEnd = ( (absSCEta > 1.566) & (absSCEta < 2.5) )\n",
    "    notEta = ~(etaBar | etaEnd)\n",
    "    outs[notEta] = False #anything in the gap dies\n",
    "    # setup all the hcalDepth1Iso cuts    \n",
    "    ones = electron.typeBits.ones_like()\n",
    "    hcalDepth1Cuts = np.zeros_like(ones.content)\n",
    "    mets = (met.content*ones).content\n",
    "    rhos = (rho.content*ones).content\n",
    "    hcalDepth1Cuts[etaBar] = 2.0+0.03*mets[etaBar]+0.28*rhos[etaBar]\n",
    "    hcalDepth1Cuts[etaEnd & (mets < 50)] = 2.5+0.28*rhos[etaEnd & (mets < 50)]\n",
    "    hcalDepth1Cuts[etaEnd & (mets >= 50)] = ( 2.5 + 0.03*(mets[etaEnd & (mets >= 50)]-50) + \n",
    "                                              0.28*rhos[etaEnd & (mets >= 50)] )\n",
    "    \n",
    "    #barrel\n",
    "    outs[etaBar] &= ( ( np.abs(electron.dEtaInSeed.content[etaBar]) < 0.00600) &\n",
    "                      ( np.abs(electron.dPhiIn.content[etaBar]) < 0.06000    ) &\n",
    "                      ( electron.hovere.content[etaBar] < \n",
    "                           (1.0/(electron.ecalEnergy.content[etaBar])+0.05)  ) &\n",
    "                      (electron.e2x5.content[etaBar]/\n",
    "                            electron.e5x5.content[etaBar] > 0.94             ) &\n",
    "                      (electron.e1x5.content[etaBar]/\n",
    "                            electron.e5x5.content[etaBar] > 0.83             ) &\n",
    "                      ( electron.hcalDepth1Iso.content[etaBar] < \n",
    "                                     hcalDepth1Cuts[etaBar]                  ) & \n",
    "                      ( electron.trkIso.content[etaBar] < 5                  ) &\n",
    "                      ( electron.nMissingHits.content[etaBar] <= 1           ) &\n",
    "                      ( electron.d0.content[etaBar] <= 0.02                  )   )\n",
    "    #endcap\n",
    "    outs[etaEnd] &= ( ( np.abs(electron.dEtaInSeed.content[etaEnd]) < 0.00400) &\n",
    "                      ( np.abs(electron.dPhiIn.content[etaEnd]) < 0.06000    ) &\n",
    "                      ( electron.hovere.content[etaEnd] < \n",
    "                           (5.0/(electron.ecalEnergy.content[etaEnd])+0.05)  ) &\n",
    "                      (electron.sieie.content[etaEnd] < 0.03                 ) &                      \n",
    "                      ( electron.hcalDepth1Iso.content[etaEnd] < \n",
    "                                     hcalDepth1Cuts[etaEnd]                  ) & \n",
    "                      ( electron.trkIso.content[etaEnd] < 5                  ) &\n",
    "                      ( electron.nMissingHits.content[etaEnd] <= 1           ) &\n",
    "                      ( electron.d0.content[etaEnd] <= 0.05                  )   )\n",
    "    outs = JaggedArray.fromoffsets(electron.pt.offsets,outs)\n",
    "    return outs\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#tau IDs\n",
    "# descriptions from https://twiki.cern.ch/twiki/bin/view/CMSPublic/SWGuidePFTauID#Discriminators\n",
    "kByDecayModeFinding                        = 1 << 0   # You will always want to use this (see AN-10-82) \n",
    "\n",
    "kByVLooseIsolation                         = 1 << 1   # isolation cone of 0.3, no PF Charged Candidates with pT > 1.5 GeV/c and no PF Gamma candidates with ET > 2.0 GeV\n",
    "kByLooseIsolation                          = 1 << 2   # (description N/A)\n",
    "kByMediumIsolation                         = 1 << 3   # (description N/A)\n",
    "kByTightIsolation                          = 1 << 4   # (description N/A)\n",
    "kByVLooseIsolationDBSumPtCorr              = 1 << 5   # (description N/A)\n",
    "kByLooseIsolationDBSumPtCorr               = 1 << 6   # (description N/A)\n",
    "kByMediumIsolationDBSumPtCorr              = 1 << 7   # (description N/A)\n",
    "kByTightIsolationDBSumPtCorr               = 1 << 8   # (description N/A)\n",
    "kByVLooseCombinedIsolationDBSumPtCorr      = 1 << 9   # isolation cone of 0.3, Delta Beta corrected sum pT of PF charged and PF gamma isolation candidates (pT > 0.5 GeV) less than 3 GeV\n",
    "kByLooseCombinedIsolationDBSumPtCorr       = 1 << 10  # isolation cone of 0.5, Delta Beta corrected sum pT of PF charged and PF gamma isolation candidates (pT > 0.5 GeV) less than 2 GeV\n",
    "kByMediumCombinedIsolationDBSumPtCorr      = 1 << 11  # isolation cone of 0.5, Delta Beta corrected sum pT of PF charged and PF gamma isolation candidates (pT > 0.5 GeV) less than 1 GeV\n",
    "kByTightCombinedIsolationDBSumPtCorr       = 1 << 12  # isolation cone of 0.5, Delta Beta corrected sum pT of PF charged and PF gamma isolation candidates (pT > 0.5 GeV) less than 0.8 GeV\n",
    "kByLooseCombinedIsolationDBSumPtCorr3Hits  = 1 << 13  # same as ByLooseCombinedIsolationDBSumPtCorr but requiring 3 hits (instead of 8) on track of isolation candidates\n",
    "kByMediumCombinedIsolationDBSumPtCorr3Hits = 1 << 14  # same as ByMediumCombinedIsolationDBSumPtCorr but requiring 3 hits (instead of 8) on track of isolation candidates\n",
    "kByTightCombinedIsolationDBSumPtCorr3Hits  = 1 << 15  # same as ByTightCombinedIsolationDBSumPtCorr but requiring 3 hits (instead of 8) on track of isolation candidates\n",
    "kByLooseIsolationMVA                       = 1 << 16  # BDT based selection using isolation in rings around tau direction and shower shape variables\n",
    "kByMediumIsolationMVA                      = 1 << 17  #  \"\n",
    "kByTightIsolationMVA                       = 1 << 18  #  \"\n",
    "kByLooseIsolationMVA2                      = 1 << 19  # same as ByLooseIsolationMVA with new training and improved performance\n",
    "kByMediumIsolationMVA2                     = 1 << 20  # same as ByMediumIsolationMVA with new training and improved performance\n",
    "kByTightIsolationMVA2                      = 1 << 21  # same as ByTightIsolationMVA with new training and improved performance\n",
    "\n",
    "kByLooseElectronRejection                  = 1 << 22  # electron pion MVA discriminator < 0.6\n",
    "kByMediumElectronRejection                 = 1 << 23  # electron pion MVA discriminator < -0.1 and not 1.4442 < |eta| < 1.566\n",
    "kByTightElectronRejection                  = 1 << 24  # electron pion MVA discriminator < -0.1 and not 1.4442 < |eta| < 1.566 and Brem pattern cuts (see AN-10-387)\n",
    "kByMVA3LooseElectronRejection              = 1 << 25  # anti-electron MVA discriminator with improved training\n",
    "kByMVA3MediumElectronRejection             = 1 << 26  #  \"\n",
    "kByMVA3TightElectronRejection              = 1 << 27  #  \"\n",
    "kByMVA3VTightElectronRejection             = 1 << 28  #  \"\n",
    "    \n",
    "kByLooseMuonRejection                      = 1 << 29  # Tau Lead Track not matched to chamber hits\n",
    "kByMediumMuonRejection                     = 1 << 30  # Tau Lead Track not matched to global/tracker muon\n",
    "kByTightMuonRejection                      = 1 << 31  # Tau Lead Track not matched to global/tracker muon and large enough energy deposit in ECAL+HCAL exceeding 0.2 times Lead Track momentum\n",
    "kByLooseMuonRejection2                     = 1 << 32  # Same as LooseMuonRejection\n",
    "kByMediumMuonRejection2                    = 1 << 33  # LooseMuonRejection2 && no DT, CSC or RPC Hits in last 2 Stations\n",
    "kByTightMuonRejection2                     = 1 << 34  # MediumMuonRejection2 && large enough energy deposit in ECAL+HCAL in 1 prong + 0 strip decay mode (SUM(ECAL+HCAL)>0.2*pT)\n",
    "kByLooseMuonRejection3                     = 1 << 35  # Same as LooseMuonRejection\n",
    "kByTightMuonRejection3                     = 1 << 36  # MediumMuonRejection2 && large enough energy deposit in ECAL+HCAL in 1 prong + 0 strip decay mode (SUM(ECAL+HCAL)>0.2*pT)\n",
    "\n",
    "def passTauSel(tau):\n",
    "    return ( ((tau.hpsDisc & kByDecayModeFinding) > 0) & (tau.rawIso3Hits < 5) )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#photon IDs\n",
    "\n",
    "def calcPhotonEffAreaIso(photon,rho):\n",
    "    chHadIso  = np.maximum(photon.chHadIso  - rho.content*photon.EA_CHad,0.);\n",
    "    neuHadIso = np.maximum(photon.neuHadIso - rho.content*photon.EA_NHad,0.);\n",
    "    phoIso    = np.maximum(photon.gammaIso  - rho.content*photon.EA_Pho, 0.);\n",
    "\n",
    "    maxNeuHadLooseIso = np.zeros_like(photon.pt.content)\n",
    "    maxNeuHadMediumIso = np.zeros_like(photon.pt.content)\n",
    "    maxNeuHadTightIso = np.zeros_like(photon.pt.content)\n",
    "    etaCut = (photon.absSCEta.content < 1.479)\n",
    "    notEta = ~etaCut\n",
    "    #common bits\n",
    "    ptshape = np.zeros_like(photon.pt.content)\n",
    "    ptshape[etaCut] = photon.pt.content[etaCut]*(0.0148+0.000017*photon.pt.content[etaCut])\n",
    "    ptshape[notEta] = photon.pt.content[notEta]*(0.0163+0.000014*photon.pt.content[notEta])\n",
    "    \n",
    "    #loose\n",
    "    maxNeuHadLooseIso[etaCut] = 10.910 + ptshape[etaCut]\n",
    "    maxNeuHadLooseIso[notEta] = 5.931  + ptshape[notEta]\n",
    "    maxNeuHadLooseIso = JaggedArray.fromoffsets(photon.pt.offsets,maxNeuHadLooseIso)\n",
    "    \n",
    "    #medium\n",
    "    maxNeuHadMediumIso[etaCut] = 2.725 + ptshape[etaCut]\n",
    "    maxNeuHadMediumIso[notEta] = 1.715 + ptshape[notEta]\n",
    "    maxNeuHadMediumIso = JaggedArray.fromoffsets(photon.pt.offsets,maxNeuHadMediumIso)\n",
    "    \n",
    "    #tight\n",
    "    maxNeuHadTightIso[etaCut] = 0.264 + ptshape[etaCut]\n",
    "    maxNeuHadTightIso[notEta] = 0.586 + ptshape[notEta]   \n",
    "    maxNeuHadTightIso = JaggedArray.fromoffsets(photon.pt.offsets,maxNeuHadTightIso)\n",
    "    \n",
    "    return chHadIso,neuHadIso,phoIso,maxNeuHadLooseIso,maxNeuHadMediumIso,maxNeuHadTightIso\n",
    "\n",
    "def passPhoLooseSel(photon):\n",
    "    outs = np.ones_like(photon.absSCEta.content,dtype=np.bool)\n",
    "    etaCut = ( photon.absSCEta.content < 1.479 )\n",
    "    notEta = ~etaCut\n",
    "    #barrel\n",
    "    outs[etaCut] &= ( (photon.hovere.content[etaCut] <= 0.0597      ) &\n",
    "                      (photon.sieie.content[etaCut] <= 0.01031      ) &\n",
    "                      (photon.chIsoCorr.content[etaCut] <= 1.295    ) &\n",
    "                      (photon.nhIsoCorr.content[etaCut] <= \n",
    "                                photon.neuIsoMax.content[etaCut]    ) &\n",
    "                      (photon.gamIsoCorr.content[etaCut] <= 3.630 + \n",
    "                                   0.0047*photon.pt.content[etaCut] )   )\n",
    "    #endcap\n",
    "    outs[notEta] &= ( (photon.hovere.content[notEta] <= 0.0481      ) &\n",
    "                      (photon.sieie.content[notEta] <= 0.03013      ) &\n",
    "                      (photon.chIsoCorr.content[notEta] <= 1.011    ) &\n",
    "                      (photon.nhIsoCorr.content[notEta] <= \n",
    "                            photon.neuIsoLooseMax.content[notEta]   ) &\n",
    "                      (photon.gamIsoCorr.content[notEta] <= 6.641 + \n",
    "                                   0.0034*photon.pt.content[notEta] )   )\n",
    "    outs = JaggedArray.fromoffsets(photon.pt.offsets,outs)\n",
    "    return outs\n",
    "\n",
    "def passPhoMediumSel(photon):\n",
    "    outs = np.ones_like(photon.absSCEta.content,dtype=np.bool)\n",
    "    etaCut = ( photon.absSCEta.content < 1.479 )\n",
    "    notEta = ~etaCut\n",
    "    #barrel\n",
    "    outs[etaCut] &= ( (photon.hovere.content[etaCut] <= 0.0396      ) &\n",
    "                      (photon.sieie.content[etaCut] <= 0.01022      ) &\n",
    "                      (photon.chIsoCorr.content[etaCut] <= 0.441    ) &\n",
    "                      (photon.nhIsoCorr.content[etaCut] <= \n",
    "                                photon.neuIsoMax.content[etaCut]    ) &\n",
    "                      (photon.gamIsoCorr.content[etaCut] <= 2.571 + \n",
    "                                   0.0047*photon.pt.content[etaCut] )   )\n",
    "    #endcap\n",
    "    outs[notEta] &= ( (photon.hovere.content[notEta] <= 0.0219      ) &\n",
    "                      (photon.sieie.content[notEta] <= 0.03001      ) &\n",
    "                      (photon.chIsoCorr.content[notEta] <= 0.442    ) &\n",
    "                      (photon.nhIsoCorr.content[notEta] <= \n",
    "                            photon.neuIsoMediumMax.content[notEta]  ) &\n",
    "                      (photon.gamIsoCorr.content[notEta] <= 3.863 + \n",
    "                                   0.0034*photon.pt.content[notEta] )   )\n",
    "    outs = JaggedArray.fromoffsets(photon.pt.offsets,outs)\n",
    "    return outs\n",
    "\n",
    "def passPhoTightSel(photon):\n",
    "    outs = np.ones_like(photon.absSCEta.content,dtype=np.bool)\n",
    "    etaCut = ( photon.absSCEta.content < 1.479 )\n",
    "    notEta = ~etaCut\n",
    "    #barrel\n",
    "    outs[etaCut] &= ( (photon.hovere.content[etaCut] <= 0.0269      ) &\n",
    "                      (photon.sieie.content[etaCut] <= 0.00994      ) &\n",
    "                      (photon.chIsoCorr.content[etaCut] <= 0.202    ) &\n",
    "                      (photon.nhIsoCorr.content[etaCut] <= \n",
    "                                photon.neuIsoMax.content[etaCut]    ) &\n",
    "                      (photon.gamIsoCorr.content[etaCut] <= 2.362 + \n",
    "                                   0.0047*photon.pt.content[etaCut] )   )\n",
    "    #endcap\n",
    "    outs[notEta] &= ( (photon.hovere.content[notEta] <= 0.0213      ) &\n",
    "                      (photon.sieie.content[notEta] <= 0.03000      ) &\n",
    "                      (photon.chIsoCorr.content[notEta] <= 0.034    ) &\n",
    "                      (photon.nhIsoCorr.content[notEta] <= \n",
    "                            photon.neuIsoTightMax.content[notEta]   ) &\n",
    "                      (photon.gamIsoCorr.content[notEta] <= 2.617 + \n",
    "                                   0.0034*photon.pt.content[notEta] )   )\n",
    "    outs = JaggedArray.fromoffsets(photon.pt.offsets,outs)\n",
    "    return outs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#btag working points\n",
    "\n",
    "CSVv2_L_AK4 = 0.5426 # CSVv2 WP\n",
    "CSVv2_M_AK4 = 0.8484\n",
    "CSVv2_T_AK4 = 0.9535\n",
    "\n",
    "CSVv2SubJet_L_AK8 = 0.5426; # CSVv2SubJet WP \n",
    "CSVv2SubJet_M_AK8 = 0.8484;\n",
    "\n",
    "def CSV_Medium(jet,ptCut=50,absEtaCut=2.5):\n",
    "    return ( (jet.pt > ptCut) & (np.abs(jet.eta) < absEtaCut) & (jet.csv > CSVv2_M_AK4) )\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#jet selections\n",
    "def passJetTightLepVetoSel(jet):\n",
    "    outs = np.ones_like(jet.pt.content,dtype=np.bool)\n",
    "    absEta = np.abs(jet.eta.content)\n",
    "    etaFor = (absEta <= 2.7)\n",
    "    etaCen = (absEta <= 2.4)\n",
    "    #forward jets\n",
    "    outs[etaFor] &= ( (jet.neuHadFrac.content[etaFor] < 0.90) &\n",
    "                      (jet.neuEmFrac.content[etaFor]  < 0.90) &\n",
    "                      (jet.nParticles.content[etaFor] > 1   ) &\n",
    "                      (jet.muonFrac.content[etaFor]   < 0.8 )  )\n",
    "    #central jets\n",
    "    outs[etaCen] &= ( (jet.chHadFrac.content[etaCen]  > 0.0 ) &\n",
    "                      (jet.nCharged.content[etaCen]   > 0   ) &\n",
    "                      (jet.chEmFrac.content[etaCen]   < 0.9 )  )\n",
    "    outs = JaggedArray.fromoffsets(jet.pt.offsets,outs)\n",
    "    return outs\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#matching functions\n",
    "def argMatchByDPhi(combs,deltaPhiCut=2./3.*math.pi):\n",
    "    deltaPhis = np.abs(combs.i0.delta_phi(combs.i1))\n",
    "    indexOfMin = deltaPhis.argmin()\n",
    "    indexOfMinOutShape = indexOfMin.flatten(axis=1)\n",
    "    passesCut = (deltaPhis[indexOfMin] <= deltaPhiCut)\n",
    "    passesCutOutShape = passesCut.flatten(axis=1)\n",
    "    flatPass = passesCutOutShape.flatten()\n",
    "    flatIdxMin = indexOfMinOutShape.flatten()\n",
    "    flatIdxMin[~flatPass] = -1\n",
    "    return JaggedArray.fromoffsets(passesCutOutShape.offsets,flatIdxMin)\n",
    "\n",
    "def antiMatchByDR(combs,deltaRCut=0.8):\n",
    "    mask = (combs.i0.delta_r(combs.i1) > deltaRCut)\n",
    "    return mask.any()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#analysis cut defaults\n",
    "PTCUT = 300.#450.\n",
    "PTCUTMUCR = 250.#400.\n",
    "#DBTAGCUT = 0.9\n",
    "#DBTAGCUT = 0.86\n",
    "#DBTAGCUT = 0.94\n",
    "T21DDTCUT = 0.55\n",
    "MUONPTCUT = 55\n",
    "METCUT = 140\n",
    "MASSCUT = 40\n",
    "NJETCUT = 100\n",
    "\n",
    "#self._cutFormula = ROOT.TTreeFormula(\"cutFormula\",\n",
    "#                                     \"(\" + cutFormula + \")&&(\"+self._triggerCut+\")&&(AK8Puppijet0_pt>%f||AK8Puppijet0_pt_JESDown>%f||AK8Puppijet0_pt_JESUp>%f||AK8Puppijet0_pt_JERUp>%f||AK8Puppijet0_pt_JERDown>%f)\" % (\n",
    "#                                     PTCUTMUCR, PTCUTMUCR, PTCUTMUCR, PTCUTMUCR, PTCUTMUCR), self._tt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyinstrument import Profiler\n",
    "\n",
    "profiler = Profiler()\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import concurrent.futures\n",
    "import multiprocessing\n",
    "thread_mult = 4\n",
    "thread_pool = concurrent.futures.ThreadPoolExecutor(thread_mult*multiprocessing.cpu_count())\n",
    "\n",
    "outer_tic = time.time()\n",
    "\n",
    "display.init()\n",
    "_ = map(lambda h: h.clear(), all_hists.values())\n",
    "\n",
    "filestr = \"/Users/lagray/coffea/bacon_prod/WJetsToQQ_HT400to600_qc19_3j_TuneCP5_13TeV/test.root\"\n",
    "\n",
    "evtColumns = {'pfmet':'pfMETC','pfmetphi':'pfMETCphi','puppet':'puppETC','puppetphi':'puppETCphi',\n",
    "              'npu':'nPUmean','rhoIso':'rhoIso','rhoJet':'rhoJet'}\n",
    "\n",
    "genPcolumns = {'pt':'GenParticle.pt','eta':'GenParticle.eta','phi':'GenParticle.phi','mass':'GenParticle.mass', \n",
    "               'parent':'GenParticle.parent','pdgId':'GenParticle.pdgId','status':'GenParticle.status',                \n",
    "               'rapidity':'GenParticle.y'}\n",
    "\n",
    "ak8columns = {'pt':'AK8Puppi.pt','eta':'AK8Puppi.eta','phi':'AK8Puppi.phi','mass':'AK8Puppi.mass',\n",
    "              'deepcsvb':'AddAK8Puppi.doublecsv','msd':'AddAK8Puppi.mass_sd0','deepdoubleb':'AddAK8Puppi.deepdoubleb',\n",
    "              'deepdoubleb_nomasssculptpen':'AddAK8Puppi.deepdoubleb_nomasssculptpen','tau1':'AddAK8Puppi.tau1',\n",
    "              'tau2':'AddAK8Puppi.tau2','tau3':'AddAK8Puppi.tau3','e3_v2_sdb1':'AddAK8Puppi.e3_v2_sdb1',\n",
    "              'e2_sdb1':'AddAK8Puppi.e2_sdb1',\n",
    "              'genpt':'AK8Puppi.genpt','geneta':'AK8Puppi.geneta','genphi':'AK8Puppi.genphi',\n",
    "              'genmass':'AK8Puppi.genm','genmsd':'AK8Puppi.genmsd',\n",
    "              'neuHadFrac':'AK8Puppi.neuHadFrac','neuEmFrac':'AK8Puppi.neuEmFrac','nParticles':'AK8Puppi.nParticles',\n",
    "              'muonFrac':'AK8Puppi.muonFrac','chHadFrac':'AK8Puppi.chHadFrac','nCharged':'AK8Puppi.nCharged',\n",
    "              'chEmFrac':'AK8Puppi.chEmFrac'}\n",
    "\n",
    "ak4columns = {'pt':'AK4Puppi.pt','eta':'AK4Puppi.eta','phi':'AK4Puppi.phi','mass':'AK4Puppi.mass',\n",
    "              'qgid':'AK4Puppi.qgid','csv':'AK4Puppi.csv',\n",
    "              'genpt':'AK4Puppi.genpt','geneta':'AK4Puppi.geneta','genphi':'AK4Puppi.genphi',\n",
    "              'genmass':'AK4Puppi.genm','genmsd':'AK4Puppi.genmsd',\n",
    "              'neuHadFrac':'AK4Puppi.neuHadFrac','neuEmFrac':'AK4Puppi.neuEmFrac','nParticles':'AK4Puppi.nParticles',\n",
    "              'muonFrac':'AK4Puppi.muonFrac','chHadFrac':'AK4Puppi.chHadFrac','nCharged':'AK4Puppi.nCharged',\n",
    "              'chEmFrac':'AK4Puppi.chEmFrac'}\n",
    "\n",
    "muoncolumns = {'pt':'Muon.pt','eta':'Muon.eta','phi':'Muon.phi','pogIDBits':'Muon.pogIDBits',\n",
    "               'chHadIso':'Muon.chHadIso','neuHadIso':'Muon.neuHadIso','gammaIso':'Muon.gammaIso',\n",
    "               'puIso':'Muon.puIso'}\n",
    "\n",
    "electroncolumns = {'pt':'Electron.pt','eta':'Electron.eta','phi':'Electron.phi',\n",
    "                   'typeBits':'Electron.typeBits','chHadIso':'Electron.chHadIso',\n",
    "                   'neuHadIso':'Electron.neuHadIso','gammaIso':'Electron.gammaIso',\n",
    "                   'sieie':'Electron.sieie','hovere':'Electron.hovere','eoverp':'Electron.eoverp',\n",
    "                   'nMissingHits':'Electron.nMissingHits','ecalEnergy':'Electron.ecalEnergy',\n",
    "                   'dEtaInSeed':'Electron.dEtaInSeed','dPhiIn':'Electron.dPhiIn','d0':'Electron.d0',\n",
    "                   'trkIso':'Electron.trkIso','hcalDepth1Iso':'Electron.hcalDepth1Iso',\n",
    "                   'e2x5':'Electron.e2x5','e1x5':'Electron.e1x5','e5x5':'Electron.e5x5',\n",
    "                   'isConv':'Electron.isConv','scEta':'Electron.scEta'}\n",
    "\n",
    "taucolumns = {'pt':'Tau.pt','eta':'Tau.eta','phi':'Tau.phi','mass':'Tau.m',\n",
    "              'hpsDisc':'Tau.hpsDisc','rawIso3Hits':'Tau.rawIso3Hits'}\n",
    "\n",
    "photoncolumns = {'pt':'Photon.pt','eta':'Photon.eta','phi':'Photon.phi','scEta':'Photon.scEta',\n",
    "                 'chHadIso':'Photon.chHadIso','neuHadIso':'Photon.neuHadIso','gammaIso':'Photon.gammaIso',\n",
    "                 'hovere':'Photon.hovere','sieie':'Photon.sieie'}\n",
    "\n",
    "all_columns = [evtColumns,genPcolumns,ak8columns,ak4columns,\n",
    "               muoncolumns,electroncolumns,taucolumns,photoncolumns]\n",
    "\n",
    "columns = []\n",
    "\n",
    "profiler.start()\n",
    "for cols in all_columns: columns.extend(list(cols.values()))\n",
    "\n",
    "print 'extracting',len(columns),'columns'\n",
    "    \n",
    "tot_events = 0\n",
    "tot_time = 0.0\n",
    "for arrays in uproot.iterate(filestr,'Events',columns,entrysteps=100000,executor=thread_pool):\n",
    "    tic = time.time()\n",
    "    #event-level info (run / evt / met / etc)\n",
    "    eventInfo = JaggedArray.fromcounts(np.ones_like(arrays[evtColumns['pfmet']],dtype=np.int),\n",
    "                                       Table({key:arrays[val] for key,val in evtColumns.items()}))\n",
    "    \n",
    "    #gen particle info (TODO)\n",
    "    genParticles = JaggedCandidateArray.candidatesfromcounts(arrays[genPcolumns['pt']].counts,\n",
    "                                                             **{key:arrays[val].content for key,val in genPcolumns.items()})\n",
    "    \n",
    "    #k-factors for jets\n",
    "    #vjetsKF = genParticles.pt.ones_like()\n",
    "    #genParticles.add_attributes(vjetsKF=vjetsKF)\n",
    "    \n",
    "    #setup muon candidates (gotta spoof the mass :-/)\n",
    "    muonarrays = {key:arrays[val].content for key,val in muoncolumns.items()}\n",
    "    muonmasses = np.full_like(muonarrays['pt'],0.1056583745)\n",
    "    muonarrays['mass'] = muonmasses\n",
    "    muons = JaggedCandidateArray.candidatesfromcounts(arrays[muoncolumns['pt']].counts,\n",
    "                                                      **muonarrays)\n",
    "    muons.add_attributes(deltaBetaIso = calcMuonDeltaBetaIso(muons))\n",
    "    \n",
    "    #various types of muon we'll need\n",
    "    passLooseMuonId = passMuonLooseSel(muons)\n",
    "    looseMuons = muons[passLooseMuonId]\n",
    "    leadingLooseMuon = looseMuons[looseMuons.pt.argmax()]\n",
    "    \n",
    "    passTightMuonId = passMuonTightSel(muons)\n",
    "    tightMuons = muons[passTightMuonId]\n",
    "    \n",
    "    #setup electrons candidates (gotta spoof the mass :-/)\n",
    "    electronarrays = {key:arrays[val].content for key,val in electroncolumns.items()}\n",
    "    electronmasses = np.full_like(electronarrays['pt'],0.000511)\n",
    "    electronarrays['mass'] = electronmasses\n",
    "    electrons = JaggedCandidateArray.candidatesfromcounts(arrays[electroncolumns['pt']].counts,\n",
    "                                                          **electronarrays)\n",
    "    el_abseta = np.abs(electrons.eta)\n",
    "    electrons.add_attributes(absEta  = el_abseta,\n",
    "                             EffArea = evaluate['electron_id_EffArea'](el_abseta))\n",
    "    electrons.add_attributes(effAreaIso = calcElectronEffAreaIso(electrons,eventInfo['rhoIso']))\n",
    "    \n",
    "    #various types of electrons we'll need\n",
    "    passLooseElectronId = passEleLooseSel(electrons)\n",
    "    looseElectrons = electrons[passLooseElectronId]\n",
    "    \n",
    "    #setup tau candidates    \n",
    "    taus = JaggedCandidateArray.candidatesfromcounts(arrays[taucolumns['pt']].counts,\n",
    "                                                     **{key:arrays[val].content for key,val in taucolumns.items()})\n",
    "    #various types of taus we'll need\n",
    "    passLooseTauId = passTauSel(taus)\n",
    "    looseTaus = taus[passLooseTauId]\n",
    "    \n",
    "    \n",
    "    #setup photon candidates (gotta spoof the mass :-/)\n",
    "    photonarrays = {key:arrays[val].content for key,val in photoncolumns.items()}\n",
    "    photonmasses = np.full_like(photonarrays['pt'],0.0)\n",
    "    photonarrays['mass'] = photonmasses\n",
    "    photons = JaggedCandidateArray.candidatesfromcounts(arrays[photoncolumns['pt']].counts,\n",
    "                                                        **photonarrays)\n",
    "    ph_absSCeta  = np.abs(photons.scEta)\n",
    "    EA_CHad = evaluate['photon_id_EA_CHad'](ph_absSCeta)\n",
    "    EA_NHad = evaluate['photon_id_EA_NHad'](ph_absSCeta)\n",
    "    EA_Pho  = evaluate['photon_id_EA_Pho'](ph_absSCeta)    \n",
    "    photons.add_attributes(absSCEta  = ph_absSCeta,\n",
    "                           EA_CHad = EA_CHad,\n",
    "                           EA_NHad = EA_NHad,\n",
    "                           EA_Pho  = EA_Pho)\n",
    "    photonIsoInfo = calcPhotonEffAreaIso(photons,eventInfo['rhoIso'])\n",
    "    photons.add_attributes(chIsoCorr  = photonIsoInfo[0],\n",
    "                           nhIsoCorr  = photonIsoInfo[1],\n",
    "                           gamIsoCorr = photonIsoInfo[2],\n",
    "                           neuIsoLooseMax  = photonIsoInfo[3],\n",
    "                           neuIsoMediumMax  = photonIsoInfo[4],\n",
    "                           neuIsoTightMax  = photonIsoInfo[5])\n",
    "        \n",
    "    #setup ak4 jet candidates\n",
    "    ak4puppijet = JaggedCandidateArray.candidatesfromcounts(arrays[ak4columns['pt']].counts,\n",
    "                                                            **{key:arrays[val].content for key,val in ak4columns.items()})\n",
    "    ak4_hasbtag = CSV_Medium(ak4puppijet)\n",
    "    ak4_btags = ak4puppijet[ak4_hasbtag]\n",
    "\n",
    "    #setup ak8 jet candidates\n",
    "    ak8puppijet = JaggedCandidateArray.candidatesfromcounts(arrays[ak8columns['pt']].counts,\n",
    "                                                            **{key:arrays[val].content for key,val in ak8columns.items()})\n",
    "    msd_corr_8 = ak8puppijet.msd*PUPPIweight(ak8puppijet.pt,ak8puppijet.eta)\n",
    "    clamped_msd_corr_8 = np.maximum(msd_corr_8,1e-8) #need to bound corrected jet mass from below???\n",
    "    #need JER/JEC uncertainties up/down added here too!!\n",
    "    rh_8 = 2.0*(np.log(clamped_msd_corr_8) - np.log(ak8puppijet.pt))  # tocheck here\n",
    "    rhP_8 = 2.0*np.log(clamped_msd_corr_8) - np.log(ak8puppijet.pt)\n",
    "    jtau21_8 = ak8puppijet.tau2/(ak8puppijet.tau1 + 1e-8)\n",
    "    jt21P_8 = jtau21_8 + 0.063 * rhP_8\n",
    "    N2sdb1=ak8puppijet.e3_v2_sdb1/(ak8puppijet.e2_sdb1**2 + 1e-8)\n",
    "    jtN2b1sdddt_8 = N2sdb1 - evaluate['Rho2D'](rh_8,ak8puppijet.pt)\n",
    "    #trigger weights for the Vs (TODO: up and down)    \n",
    "    trigweight = VtrigSF(ak8puppijet.msd,ak8puppijet.pt)\n",
    "    #store all our calculated stuff\n",
    "    ak8puppijet.add_attributes(msd_corr_8=msd_corr_8,\n",
    "                               rh_8=rh_8,\n",
    "                               rhP_8=rhP_8,\n",
    "                               tau21=jtau21_8,\n",
    "                               tau32=ak8puppijet.tau3/(ak8puppijet.tau2 + 1e-8),\n",
    "                               jt21P_8=jt21P_8,\n",
    "                               N2sdb1=N2sdb1,\n",
    "                               jtN2b1sdddt_8=jtN2b1sdddt_8,\n",
    "                               trigweight=trigweight)\n",
    "    leadingak8jet = ak8puppijet[ak8puppijet.pt.argmax()]\n",
    "    \n",
    "    #todo -> trigger selection\n",
    "    presel = ( (leadingak8jet.pt > PTCUTMUCR).sum() > 0 # | \n",
    "              # AK8Puppijet0_pt_JESUp > PTCUTMUCR | \n",
    "              # AK8Puppijet0_pt_JESDown > PTCUTMUCR | \n",
    "              # AK8Puppijet0_pt_JERUp > PTCUTMURC |\n",
    "              # AK8Puppijet0_pt_JERDown > PTCUTMURC \n",
    "             )\n",
    "    #presel = ( ( ak8puppijet.counts > 0 ) & ( ak8puppijet.tau21 < T21DDTCUT) )\n",
    "    \n",
    "    #single muon control regions\n",
    "    #if (jpt_8 > PTCUTMUCR and jmsd_8 > MASSCUT and \n",
    "    #    nmuLoose == 1 and neleLoose == 0 and ntau == 0 and vmuoLoose0_pt > MUONPTCUT and \n",
    "    #    abs(vmuoLoose0_eta) < 2.1 and isTightVJet and abs(vmuoLoose0_phi - jphi_8) > 2. * ROOT.TMath.Pi() / 3. and \n",
    "    #    n_MdR0p8_4 >= 1):\n",
    "    \n",
    "    #matching criteria for control regions\n",
    "    AK8jet_muon_matches = leadingak8jet.argmatch(leadingLooseMuon,argmatchfunc=argMatchByDPhi)    \n",
    "    AK8jet_AK4bjet_antimatches = leadingak8jet.match(ak4_btags,matchfunc=antiMatchByDR)\n",
    "    singlemuon_cr = ( #jet selection\n",
    "                      ((leadingak8jet.pt > PTCUTMUCR).sum() > 0) &\n",
    "                      ((leadingak8jet.msd_corr_8 > MASSCUT).sum() > 0) &\n",
    "                      (passJetTightLepVetoSel(leadingak8jet).sum() > 0) &\n",
    "                      #muon selection\n",
    "                      ((np.abs(leadingLooseMuon.eta) < 2.1).sum() > 0) &\n",
    "                      ((leadingLooseMuon.pt > MUONPTCUT).sum() > 0) &\n",
    "                      #matches and cross cleaning\n",
    "                      ((AK8jet_muon_matches != -1).counts == 0) &                      \n",
    "                      (AK8jet_AK4bjet_antimatches.sum() > 0) &\n",
    "                      #lepton vetos\n",
    "                      (looseMuons.counts == 1) &\n",
    "                      (looseElectrons.counts == 0) &\n",
    "                      (looseTaus.counts == 0)\n",
    "                    )            \n",
    "    \n",
    "    #fill histograms\n",
    "    all_hists['m_J'].fill(category='WQQ',m_J=ak8puppijet[presel].msd_corr_8.max())\n",
    "    all_hists['pT_J'].fill(category='WQQ',pT_J=ak8puppijet[presel].pt.max())\n",
    "    all_hists['deepDoubleB'].fill(category='WQQ',deepDoubleB=ak8puppijet[presel].deepdoubleb.flatten())\n",
    "    all_hists['tau21'].fill(category='WQQ',tau21=ak8puppijet[presel].tau21.flatten())\n",
    "    all_hists['tau32'].fill(category='WQQ',tau32=ak8puppijet[presel].tau32.flatten())\n",
    "    toc = time.time()\n",
    "    \n",
    "    tot_time += toc-tic\n",
    "    tot_events += arrays[evtColumns['pfmet']].size\n",
    "    display.update()\n",
    "\n",
    "profiler.stop()\n",
    "outer_toc = time.time()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "print tot_events\n",
    "print (outer_toc-outer_tic)/tot_events\n",
    "print tot_time/tot_events\n",
    "print \n",
    "print(profiler.output_text(unicode=True, color=True))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}