preselection_UL.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "4acdc2e9-391a-44cc-829a-773dcee8eb8c",
   "metadata": {},
   "source": [
    "# Preselection \n",
    "This notebook features a simplified version of the preselection step of VBS ssWW with one lepton and one handronic tau in the final state analysis. Events are filtered, and corrections on jet, PU and prefiring are computed, along with the columns necessary for systematic variations. The result of the computation is the snapshot of a skimmed and slimmed dataset."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e9ab02f2-7635-450e-86ef-a679a04ab73e",
   "metadata": {},
   "source": [
    "Import necessary libraries, setup remote reading and configure local/distributed execution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "reported-lunch",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture\n",
    "\n",
    "import ROOT\n",
    "import os\n",
    "from dask.distributed import Client, PipInstall, WorkerPlugin\n",
    "import json\n",
    "from samplesUL import *\n",
    "import sys \n",
    "from dask_jobqueue import HTCondorCluster\n",
    "from distributed.diagnostics.plugin import UploadFile\n",
    "\n",
    "distributed = False\n",
    "MT = False\n",
    "redirector = \"root://t2-xrdcms.lnl.infn.it:7070/\"\n",
    "maxNfilespersample = 99999  #lower this number just for debugging purposes\n",
    "sched_port = 20197  #used only in distributed mode\n",
    "nPartitions = 92*3  #used only in distributed mode\n",
    "\n",
    "os.environ['X509_USER_PROXY'] = '/opt/workspace/persistent-storage/proxy'\n",
    "os.environ['X509_CERT_DIR'] = '/cvmfs/cms.cern.ch/grid/etc/grid-security/certificates/'\n",
    "\n",
    "if distributed != True and MT == True:\n",
    "    ROOT.ROOT.EnableImplicitMT()\n",
    "    \n",
    "def set_proxy(dask_worker):\n",
    "    import os\n",
    "    import shutil\n",
    "    working_dir = dask_worker.local_directory\n",
    "    print(working_dir)\n",
    "    os.environ['X509_USER_PROXY'] = working_dir + '/proxy'\n",
    "    os.environ['X509_CERT_DIR']=\"/cvmfs/grid.cern.ch/etc/grid-security/certificates/\"\n",
    "    shutil.copyfile(working_dir + '/proxy', working_dir + '/../../../proxy')\n",
    "    os.environ['EXTRA_CLING_ARGS'] = \"-O2\"\n",
    "    \n",
    "    return os.environ.get(\"X509_USER_PROXY\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b02a1b7b-af20-455e-86a9-00284cd67c02",
   "metadata": {},
   "source": [
    "Declare to the ROOT interpreter the header file which contains custom defined functions: this is done inside an inizialization function, which is necessary in the case of distributed execution (it is executed at the beginning of each task)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "0a510143-d7bf-4f0f-aa10-2d0d3eaaef40",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.8/dist-packages/pkg_resources/__init__.py:123: PkgResourcesDeprecationWarning: 0.23ubuntu1 is an invalid version and will not be supported in a future release\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['nominal', 'jerup', 'jerdown', 'jesTotalup', 'jesTotaldown', 'unclustEnup', 'unclustEndown']\n",
      "Loading histogram NUM_TightID_DEN_TrackerMuons_abseta_pt from file https://vbs-pg-support.web.cern.ch/nanoAOD-tools/python/postprocessing/data/leptonSF/MuID_Tight_UL2017.root... \n",
      "Loading histogram NUM_IsoMu27_DEN_CutBasedIdTight_and_PFIsoTight_abseta_pt from file https://vbs-pg-support.web.cern.ch/nanoAOD-tools/python/postprocessing/data/leptonSF/MuTRIG_UL2017.root... \n",
      "Loading histogram NUM_TightRelIso_DEN_TightIDandIPCut_abseta_pt from file https://vbs-pg-support.web.cern.ch/nanoAOD-tools/python/postprocessing/data/leptonSF/MuISO_Tight_UL2017.root... \n",
      "Loading histogram NUM_TrackerMuons_DEN_genTracks from file https://vbs-pg-support.web.cern.ch/nanoAOD-tools/python/postprocessing/data/leptonSF/MuRECO_UL2017.root... \n",
      "Loading histogram EGamma_SF2D from file https://vbs-pg-support.web.cern.ch/nanoAOD-tools/python/postprocessing/data/leptonSF/EleRECO_UL2017_EGM2D.root... \n",
      "Loading histogram EGamma_SF2D from file https://vbs-pg-support.web.cern.ch/nanoAOD-tools/python/postprocessing/data/leptonSF/EleID_WP90Iso_UL2017_EGM2D.root... \n"
     ]
    }
   ],
   "source": [
    "#from my_initialization_function_UL import *\n",
    "text_file = open(\"preselection_UL.h\", \"r\")\n",
    "data = text_file.read()\n",
    "\n",
    "text_file = open(\"preselection_part2_UL.h\", \"r\")\n",
    "data_2 = text_file.read()\n",
    "\n",
    "def my_initialization_function():\n",
    "    \n",
    "    import ROOT\n",
    "    import sys\n",
    "    ROOT.EnableThreadSafety()\n",
    "    ROOT.gInterpreter.AddIncludePath(\"/cvmfs/cms.dodas.infn.it/boost_1_77_0\")\n",
    "    sys.path.append(\"/lib/python3.8/site-packages\")\n",
    "    \n",
    "    jec_prefix_MC = \"Summer19UL17_V6_MC\"\n",
    "    jer_prefix_MC = \"Summer19UL17_JRV3_MC\"\n",
    "    \n",
    "    from urllib import request \n",
    "    request.urlretrieve(\"https://vbs-pg-support.web.cern.ch/nanoAOD-tools/python/postprocessing/data/roccor.Run2.v5/RoccoR2017UL.txt\", \"RoccoR2017UL.txt\")\n",
    "    request.urlretrieve(\"https://vbs-pg-support.web.cern.ch/nanoAOD-tools/python/postprocessing/data/roccor.Run2.v5/RoccoR.cc\", \"RoccoR.cc\")\n",
    "    request.urlretrieve(\"https://vbs-pg-support.web.cern.ch/nanoAOD-tools/python/postprocessing/data/roccor.Run2.v5/RoccoR.h\", \"RoccoR.h\")\n",
    "    request.urlretrieve(\"https://vbs-pg-support.web.cern.ch/nanoAOD-tools/data/btagSF/DeepJet_106XUL17_v3_new.csv\", \"DeepJet_106XUL17_v3_new.csv\") \n",
    "    \n",
    "    \n",
    "    ROOT.gInterpreter.Declare(\n",
    "        '''\n",
    "        #ifndef ROCCOR\n",
    "        #define ROCCOR\n",
    "        #include \"RoccoR.cc\"\n",
    "        #endif\n",
    "        ''')\n",
    "\n",
    "    request.urlretrieve(\"https://vbs-pg-support.web.cern.ch/jme_data/{}_L1FastJet_AK4PFchs.txt\".format(jec_prefix_MC), \"{}_L1FastJet_AK4PFchs.txt\".format(jec_prefix_MC))\n",
    "    request.urlretrieve(\"https://vbs-pg-support.web.cern.ch/jme_data/{}_L2Relative_AK4PFchs.txt\".format(jec_prefix_MC), \"{}_L2Relative_AK4PFchs.txt\".format(jec_prefix_MC)) \n",
    "    request.urlretrieve(\"https://vbs-pg-support.web.cern.ch/jme_data/{}_L3Absolute_AK4PFchs.txt\".format(jec_prefix_MC), \"{}_L3Absolute_AK4PFchs.txt\".format(jec_prefix_MC)) \n",
    "    request.urlretrieve(\"https://vbs-pg-support.web.cern.ch/jme_data/{}_L2L3Residual_AK4PFchs.txt\".format(jec_prefix_MC), \"{}_L2L3Residual_AK4PFchs.txt\".format(jec_prefix_MC)) \n",
    "    request.urlretrieve(\"https://vbs-pg-support.web.cern.ch/jme_data/{}_UncertaintySources_AK4PFchs.txt\".format(jec_prefix_MC), \"{}_UncertaintySources_AK4PFchs.txt\".format(jec_prefix_MC)) \n",
    "    request.urlretrieve(\"https://vbs-pg-support.web.cern.ch/jme_data/{}_PtResolution_AK4PFchs.txt\".format(jer_prefix_MC), \"{}_PtResolution_AK4PFchs.txt\".format(jer_prefix_MC))\n",
    "    request.urlretrieve(\"https://vbs-pg-support.web.cern.ch/jme_data/{}_SF_AK4PFchs.txt\".format(jer_prefix_MC), \"{}_SF_AK4PFchs.txt\".format(jer_prefix_MC))\n",
    "    \n",
    "    \n",
    "    ROOT.gInterpreter.Declare('{}'.format(data))\n",
    "\n",
    "    ROOT.gInterpreter.ProcessLine('#ifndef LOADING')\n",
    "    ROOT.gInterpreter.ProcessLine('#define LOADING')\n",
    "    \n",
    "    ROOT.gInterpreter.ProcessLine('reader_1_UL2017.load(calibration_UL2017, BTagEntry::FLAV_B, \"comb\");')\n",
    "    ROOT.gInterpreter.ProcessLine('reader_1_UL2017.load(calibration_UL2017, BTagEntry::FLAV_C, \"comb\");')\n",
    "    ROOT.gInterpreter.ProcessLine('reader_1_UL2017.load(calibration_UL2017, BTagEntry::FLAV_UDSG, \"incl\");')\n",
    "\n",
    "    ROOT.gInterpreter.ProcessLine('#endif')\n",
    "    \n",
    "    ROOT.gInterpreter.Declare('{}'.format(data_2))\n",
    "    \n",
    "\n",
    "    from CMSJMECalculators import loadJMESystematicsCalculators\n",
    "    loadJMESystematicsCalculators()\n",
    "\n",
    "    \n",
    "    ################################# JET #####################################################################################################\n",
    "    \n",
    "    ROOT.gROOT.ProcessLine(\"JetVariationsCalculator myJetVarCalc{};\")\n",
    "    calc = getattr(ROOT, \"myJetVarCalc\")\n",
    "\n",
    "    # redo JEC, push_back corrector parameters for different levels\n",
    "    jecParams = getattr(ROOT, \"std::vector<JetCorrectorParameters>\")()\n",
    "    jecParams.push_back(ROOT.JetCorrectorParameters(\"./{}_L1FastJet_AK4PFchs.txt\".format(jec_prefix_MC)))\n",
    "    jecParams.push_back(ROOT.JetCorrectorParameters(\"./{}_L2Relative_AK4PFchs.txt\".format(jec_prefix_MC)))\n",
    "    #jecParams.push_back(ROOT.JetCorrectorParameters(\"./{}_L3Absolute_AK4PFchs.txt\".format(jec_prefix_MC)))\n",
    "    jecParams.push_back(ROOT.JetCorrectorParameters(\"./{}_L2L3Residual_AK4PFchs.txt\".format(jec_prefix_MC)))\n",
    "    calc.setJEC(jecParams)\n",
    "\n",
    "    uncert_sources = [\"Total\"]\n",
    "    #uncert_sources = []\n",
    "    # calculate JES uncertainties\n",
    "    with open(\"./{}_UncertaintySources_AK4PFchs.txt\".format(jec_prefix_MC)) as f:\n",
    "        lines = f.read().split(\"\\n\")\n",
    "        sources = [\n",
    "            x for x in lines if x.startswith(\"[\") and x.endswith(\"]\")\n",
    "        ]\n",
    "        sources = [x[1:-1] for x in sources]\n",
    "        sources = list(filter(lambda source: source in uncert_sources, sources))\n",
    "\n",
    "    for s in sources:\n",
    "        jcp_unc = ROOT.JetCorrectorParameters(\"./{}_UncertaintySources_AK4PFchs.txt\".format(jec_prefix_MC), s)\n",
    "        calc.addJESUncertainty(s, jcp_unc)\n",
    "\n",
    "    # Smear jets, with JER uncertainty\n",
    "    calc.setSmearing(\"./{}_PtResolution_AK4PFchs.txt\".format(jer_prefix_MC), \"./{}_SF_AK4PFchs.txt\".format(jer_prefix_MC), False, True, 0.2, 3.)  # use hybrid recipe, matching parameters\n",
    "\n",
    "    jetvariations = [calc.available()[i] for i in range(0,calc.available().size())]\n",
    "    \n",
    "    ############################################ MET T1SMEAR ###########################################################################################\n",
    "    \n",
    "    ROOT.gROOT.ProcessLine(\"Type1METVariationsCalculator myType1SmearMETVarCalc{};\")\n",
    "    calc = getattr(ROOT, \"myType1SmearMETVarCalc\")\n",
    "\n",
    "    # redo JEC, push_back corrector parameters for different levels\n",
    "    jecParams = getattr(ROOT, \"std::vector<JetCorrectorParameters>\")()\n",
    "    jecParams.push_back(ROOT.JetCorrectorParameters(\"./{}_L1FastJet_AK4PFchs.txt\".format(jec_prefix_MC)))\n",
    "    jecParams.push_back(ROOT.JetCorrectorParameters(\"./{}_L2Relative_AK4PFchs.txt\".format(jec_prefix_MC)))\n",
    "    jecParams.push_back(ROOT.JetCorrectorParameters(\"./{}_L2L3Residual_AK4PFchs.txt\".format(jec_prefix_MC)))\n",
    "\n",
    "    calc.setJEC(jecParams)\n",
    "\n",
    "    jecParams = getattr(ROOT, \"std::vector<JetCorrectorParameters>\")()\n",
    "    jecParams.push_back(ROOT.JetCorrectorParameters(\"./{}_L1FastJet_AK4PFchs.txt\".format(jec_prefix_MC)))\n",
    "    calc.setL1JEC(jecParams)\n",
    "\n",
    "    calc.setIsT1SmearedMET(True)\n",
    "\n",
    "    # calculate JES uncertainties\n",
    "    with open(\"./{}_UncertaintySources_AK4PFchs.txt\".format(jec_prefix_MC)) as f:\n",
    "        lines = f.read().split(\"\\n\")\n",
    "        sources = [\n",
    "            x for x in lines if x.startswith(\"[\") and x.endswith(\"]\")\n",
    "        ]\n",
    "        sources = [x[1:-1] for x in sources]\n",
    "        sources = list(filter(lambda source: source in uncert_sources, sources))\n",
    "\n",
    "\n",
    "    for s in sources:\n",
    "        jcp_unc = ROOT.JetCorrectorParameters(\"./{}_UncertaintySources_AK4PFchs.txt\".format(jec_prefix_MC), s)\n",
    "        calc.addJESUncertainty(s, jcp_unc)\n",
    "\n",
    "    # Smear jets, with JER uncertainty\n",
    "    calc.setSmearing(\"./{}_PtResolution_AK4PFchs.txt\".format(jer_prefix_MC), \"./{}_SF_AK4PFchs.txt\".format(jer_prefix_MC), False, True, 0.2, 3.)  # use hybrid recipe, matching parameters\n",
    "\n",
    "    mett1smearvariations = [calc.available()[i] for i in range(0, calc.available().size())]\n",
    "    \n",
    "    print(mett1smearvariations)\n",
    "\n",
    "    #ROOT.gInterpreter.ProcessLine('#endif')\n",
    "\n",
    "    \n",
    "jetvariations = [\n",
    " 'nominal',\n",
    " 'jerUp',\n",
    " 'jerDown',\n",
    " 'jesTotalUp',\n",
    " 'jesTotalDown',\n",
    "]\n",
    "\n",
    "mett1smearvariations = jetvariations\n",
    "\n",
    "if distributed == True:\n",
    "    RDataFrame = ROOT.RDF.Experimental.Distributed.Dask.RDataFrame\n",
    "    client = Client(address=\"tcp://127.0.0.1:\"+str(sched_port))\n",
    "    client.restart()\n",
    "    client.register_worker_plugin(UploadFile(\"/opt/workspace/persistent-storage/proxy\"))\n",
    "    client.run(set_proxy)\n",
    "    ROOT.RDF.Experimental.Distributed.initialize(my_initialization_function)\n",
    "    \n",
    "else:\n",
    "    RDataFrame = ROOT.RDataFrame\n",
    "    my_initialization_function()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e4a19623-55fb-49be-a13a-3dab50323667",
   "metadata": {},
   "source": [
    "Define the samples to be analyzed. This makes use of auxiliary sample class defined in samplesUL.py."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "274c428d-fe3b-4515-a3c3-490852344ed9",
   "metadata": {},
   "outputs": [],
   "source": [
    "sampleDict = {'ZZtoLep_UL2016APV': 0,'ZZTo2L2Nu_UL2016APV': 1,'ZZTo4L_UL2016APV': 2,'GluGluToContinToZZTo4e_UL2016APV': 3,'GluGluToContinToZZTo2e2mu_UL2016APV': 4,'GluGluToContinToZZTo2e2tau_UL2016APV': 5,'GluGluToContinToZZTo2mu2nu_UL2016APV': 6,'GluGluToContinToZZTo4mu_UL2016APV': 7,'GluGluToContinToZZTo2mu2tau_UL2016APV': 8,'GluGluToContinToZZTo2tau2nu_UL2016APV': 9,'GluGluToContinToZZTo4tau_UL2016APV': 10,'GluGluToContinToZZTo2e2nu_UL2016APV': 11,'TT_UL2016APV': 12,'TT_SemiLep_UL2016APV': 13,'TT_Had_UL2016APV': 14,'TTTo2L2Nu_UL2016APV': 15,'TT_beff_UL2016APV': 16,'TVX_UL2016APV': 17,'TTGJets_UL2016APV': 18,'TTZToQQ_UL2016APV': 19,'TTZToLLNuNu_UL2016APV': 20,'TTWJetsToQQ_UL2016APV': 21,'TTWJetsToLNu_UL2016APV': 22,'tZq_ll_4f_UL2016APV': 23,'VG_UL2016APV': 24,'ZG_UL2016APV': 25,'WG_UL2016APV': 26,'WrongSign_UL2016APV': 27,'WWto2L2Nu_UL2016APV': 28,'GluGluToWWToENEN_UL2016APV': 29,'GluGluToWWToENMN_UL2016APV': 30,'GluGluToWWToENTN_UL2016APV': 31,'GluGluToWWToMNEN_UL2016APV': 32,'GluGluToWWToMNMN_UL2016APV': 33,'GluGluToWWToMNTN_UL2016APV': 34,'GluGluToWWToTNEN_UL2016APV': 35,'GluGluToWWToTNMN_UL2016APV': 36,'GluGluToWWToTNTN_UL2016APV': 37,'ST_tW_top_UL2016APV': 38,'ST_tW_antitop_UL2016APV': 39,'GluGluHToWWTo2L2Nu_UL2016APV': 40,'GluGluHToWWToLNuQQ_UL2016APV': 41,'GluGluHToZZTo4L_UL2016APV': 42,'GluGluHToTauTau_UL2016APV': 43,'VBFHToWWTo2L2Nu_UL2016APV': 44,'VBFHToTauTau_UL2016APV': 45,'ttHToNonbb_UL2016APV': 46,'VHToNonbb_UL2016APV': 47,'Triboson_UL2016APV': 48,'WWTo2L2Nu_DoubleScattering_UL2016': 49,'WWW_4F_UL2016APV': 50,'WWZ_4F_UL2016APV': 51,'WZZ_UL2016APV': 52,'ZZZ_UL2016APV': 53,'WWG_UL2016APV': 54,'WJets_UL2016APV': 55,'WJetsHT70to100_UL2016APV': 56,'WJetsHT100to200_UL2016APV': 57,'WJetsHT200to400_UL2016APV': 58,'WJetsHT400to600_UL2016APV': 59,'WJetsHT600to800_UL2016APV': 60,'WJetsHT800to1200_UL2016APV': 61,'WJetsHT1200to2500_UL2016APV': 62,'WJetsHT2500toInf_UL2016APV': 63,'WZ_UL2016APV': 64,'DYJetsToLL_UL2016APV': 65,'DYJetsToLL_FxFx_UL2016APV': 65,'DYJetsToLL_M10to50_UL2016APV': 66,'DYJetsToLL_M50_UL2016APV': 67,'DYJetsToLL_M50_FxFx_UL2016APV': 67,'DYJetsToLL_M50_UL2016APV_ext': 67,'WpWpJJ_EWK_UL2016APV': 68,'WpWpJJ_QCD_UL2016APV': 69,'VBS_SSWW_SM_UL2016APV': 70,'VBS_SSWW_LL_SM_UL2016APV': 71,'VBS_SSWW_TL_SM_UL2016APV': 72,'VBS_SSWW_TT_SM_UL2016APV': 73,'VBS_SSWW_cW_UL2016APV': 74,'VBS_SSWW_cW_SM_UL2016APV': 75,'VBS_SSWW_cW_BSM_UL2016APV': 76,'VBS_SSWW_cW_INT_UL2016APV': 77,'VBS_SSWW_cHW_UL2016APV': 78,'VBS_SSWW_cHW_SM_UL2016APV': 79,'VBS_SSWW_cHW_BSM_UL2016APV': 80,'VBS_SSWW_cHW_INT_UL2016APV': 81,'VBS_SSWW_cW_cHW_UL2016APV': 82,'VBS_SSWW_DIM6_UL2016APV': 83,'VBS_SSWW_DIM6_SM_UL2016APV': 84,'ZZtoLep_UL2016': 85,'ZZTo2L2Nu_UL2016': 86,'ZZTo4L_UL2016': 87,'GluGluToContinToZZTo4e_UL2016': 88,'GluGluToContinToZZTo2e2mu_UL2016': 89,'GluGluToContinToZZTo2e2tau_UL2016': 90,'GluGluToContinToZZTo2mu2nu_UL2016': 91,'GluGluToContinToZZTo4mu_UL2016': 92,'GluGluToContinToZZTo2mu2tau_UL2016': 93,'GluGluToContinToZZTo2tau2nu_UL2016': 94,'GluGluToContinToZZTo4tau_UL2016': 95,'GluGluToContinToZZTo2e2nu_UL2016': 96,'TT_UL2016': 97,'TT_SemiLep_UL2016': 98,'TT_Had_UL2016': 99,'TTTo2L2Nu_UL2016': 100,'TT_beff_UL2016': 101,'TVX_UL2016': 102,'TTGJets_UL2016': 103,'TTZToQQ_UL2016': 104,'TTZToLLNuNu_UL2016': 105,'TTWJetsToQQ_UL2016': 106,'TTWJetsToLNu_UL2016': 107,'tZq_ll_4f_UL2016': 108,'VG_UL2016': 109,'ZG_UL2016': 110,'WG_UL2016': 111,'WrongSign_UL2016': 112,'WWto2L2Nu_UL2016': 113,'GluGluToWWToENEN_UL2016': 114,'GluGluToWWToENMN_UL2016': 115,'GluGluToWWToENTN_UL2016': 116,'GluGluToWWToMNEN_UL2016': 117,'GluGluToWWToMNMN_UL2016': 118,'GluGluToWWToMNTN_UL2016': 119,'GluGluToWWToTNEN_UL2016': 120,'GluGluToWWToTNMN_UL2016': 121,'GluGluToWWToTNTN_UL2016': 122,'ST_tW_top_UL2016': 123,'ST_tW_antitop_UL2016': 124,'GluGluHToWWTo2L2Nu_UL2016': 125,'GluGluHToWWToLNuQQ_UL2016': 126,'GluGluHToZZTo4L_UL2016': 127,'GluGluHToTauTau_UL2016': 128,'VBFHToWWTo2L2Nu_UL2016': 129,'VBFHToTauTau_UL2016': 130,'ttHToNonbb_UL2016': 131,'VHToNonbb_UL2016': 132,'Triboson_UL2016': 133,'WWW_4F_UL2016': 134,'WWZ_4F_UL2016': 135,'WZZ_UL2016': 136,'ZZZ_UL2016': 137,'WWG_UL2016': 138,'WJets_UL2016': 139,'WJetsHT70to100_UL2016': 140,'WJetsHT100to200_UL2016': 141,'WJetsHT200to400_UL2016': 142,'WJetsHT400to600_UL2016': 143,'WJetsHT600to800_UL2016': 144,'WJetsHT800to1200_UL2016': 145,'WJetsHT1200to2500_UL2016': 146,'WJetsHT2500toInf_UL2016': 147,'WZ_UL2016': 148,'DYJetsToLL_UL2016': 149,'DYJetsToLL_FxFx_UL2016': 149,'DYJetsToLL_M10to50_UL2016': 150,'DYJetsToLL_M50_UL2016': 151,'DYJetsToLL_M50_FxFx_UL2016': 151,'DYJetsToLL_M50_UL2016_ext': 151,'WpWpJJ_EWK_UL2016': 152,'WpWpJJ_QCD_UL2016': 153,'VBS_SSWW_SM_UL2016': 154,'VBS_SSWW_LL_SM_UL2016': 155,'VBS_SSWW_TL_SM_UL2016': 156,'VBS_SSWW_TT_SM_UL2016': 157,'VBS_SSWW_cW_UL2016': 158,'VBS_SSWW_cW_SM_UL2016': 159,'VBS_SSWW_cW_BSM_UL2016': 160,'VBS_SSWW_cW_INT_UL2016': 161,'VBS_SSWW_cHW_UL2016': 162,'VBS_SSWW_cHW_SM_UL2016': 163,'VBS_SSWW_cHW_BSM_UL2016': 164,'VBS_SSWW_cHW_INT_UL2016': 165,'VBS_SSWW_cW_cHW_UL2016': 166,'VBS_SSWW_DIM6_UL2016': 167,'VBS_SSWW_DIM6_SM_UL2016': 168,'ZZtoLep_UL2017': 169,'ZZTo2L2Nu_UL2017': 170,'ZZTo4L_UL2017': 171,'GluGluToContinToZZTo4e_UL2017': 172,'GluGluToContinToZZTo2e2mu_UL2017': 173,'GluGluToContinToZZTo2e2tau_UL2017': 174,'GluGluToContinToZZTo2mu2nu_UL2017': 175,'GluGluToContinToZZTo4mu_UL2017': 176,'GluGluToContinToZZTo2mu2tau_UL2017': 177,'GluGluToContinToZZTo2tau2nu_UL2017': 178,'GluGluToContinToZZTo4tau_UL2017': 179,'GluGluToContinToZZTo2e2nu_UL2017': 180,'TT_UL2017': 181,'TT_SemiLep_UL2017': 182,'TT_Had_UL2017': 183,'TTTo2L2Nu_UL2017': 184,'TT_beff_UL2017': 185,'TVX_UL2017': 186,'TTGJets_UL2017': 187,'TTZToQQ_UL2017': 188,'TTZToLLNuNu_UL2017': 189,'TTWJetsToQQ_UL2017': 190,'TTWJetsToLNu_UL2017': 191,'tZq_ll_4f_UL2017': 192,'VG_UL2017': 193,'ZG_UL2017': 194,'WG_UL2017': 195,'WrongSign_UL2017': 196,'WWto2L2Nu_UL2017': 197,'GluGluToWWToENEN_UL2017': 198,'GluGluToWWToENMN_UL2017': 199,'GluGluToWWToENTN_UL2017': 200,'GluGluToWWToMNEN_UL2017': 201,'GluGluToWWToMNMN_UL2017': 202,'GluGluToWWToMNTN_UL2017': 203,'GluGluToWWToTNEN_UL2017': 204,'GluGluToWWToTNMN_UL2017': 205,'GluGluToWWToTNTN_UL2017': 206,'ST_tW_top_UL2017': 207,'ST_tW_antitop_UL2017': 208,'GluGluHToWWTo2L2Nu_UL2017': 209,'GluGluHToWWToLNuQQ_UL2017': 210,'GluGluHToZZTo4L_UL2017': 211,'GluGluHToTauTau_UL2017': 212,'VBFHToWWTo2L2Nu_UL2017': 213,'VBFHToTauTau_UL2017': 214,'ttHToNonbb_UL2017': 215,'VHToNonbb_UL2017': 216,'Triboson_UL2017': 217,'WWTo2L2Nu_DoubleScattering_UL2017': 218,'WWW_4F_UL2017': 219,'WWZ_4F_UL2017': 220,'WZZ_UL2017': 221,'ZZZ_UL2017': 222,'WWG_UL2017': 223,'WJets_UL2017': 224,'WJetsHT70to100_UL2017': 225,'WJetsHT100to200_UL2017': 226,'WJetsHT200to400_UL2017': 227,'WJetsHT400to600_UL2017': 228,'WJetsHT600to800_UL2017': 229,'WJetsHT800to1200_UL2017': 230,'WJetsHT1200to2500_UL2017': 231,'WJetsHT2500toInf_UL2017': 232,'WZ_UL2017': 233,'DYJetsToLL_UL2017': 234,'DYJetsToLL_FxFx_UL2017': 234,'DYJetsToLL_M10to50_UL2017': 235,'DYJetsToLL_M50_UL2017': 236,'DYJetsToLL_M50_FxFx_UL2017': 236,'DYJetsToLL_M50_UL2017_ext': 236,'WpWpJJ_EWK_UL2017': 237,'WpWpJJ_QCD_UL2017': 238,'VBS_SSWW_SM_UL2017': 239,'VBS_SSWW_LL_SM_UL2017': 240,'VBS_SSWW_TL_SM_UL2017': 241,'VBS_SSWW_TT_SM_UL2017': 242,'VBS_SSWW_cW_UL2017': 243,'VBS_SSWW_cW_SM_UL2017': 244,'VBS_SSWW_cW_BSM_UL2017': 245,'VBS_SSWW_cW_INT_UL2017': 246,'VBS_SSWW_cHW_UL2017': 247,'VBS_SSWW_cHW_SM_UL2017': 248,'VBS_SSWW_cHW_BSM_UL2017': 249,'VBS_SSWW_cHW_INT_UL2017': 250,'VBS_SSWW_cW_cHW_UL2017': 251,'VBS_SSWW_DIM6_UL2017': 252,'VBS_SSWW_DIM6_SM_UL2017': 253,'ZZtoLep_UL2018': 254,'ZZTo2L2Nu_UL2018': 255,'ZZTo4L_UL2018': 256,'GluGluToContinToZZTo4e_UL2018': 257,'GluGluToContinToZZTo2e2mu_UL2018': 258,'GluGluToContinToZZTo2e2tau_UL2018': 259,'GluGluToContinToZZTo2mu2nu_UL2018': 260,'GluGluToContinToZZTo4mu_UL2018': 261,'GluGluToContinToZZTo2mu2tau_UL2018': 262,'GluGluToContinToZZTo2tau2nu_UL2018': 263,'GluGluToContinToZZTo4tau_UL2018': 264,'GluGluToContinToZZTo2e2nu_UL2018': 265,'TT_UL2018': 266,'TT_SemiLep_UL2018': 267,'TT_Had_UL2018': 268,'TTTo2L2Nu_UL2018': 269,'TT_beff_UL2018': 270,'TVX_UL2018': 271,'TTGJets_UL2018': 272,'TTZToQQ_UL2018': 273,'TTZToLLNuNu_UL2018': 274,'TTWJetsToQQ_UL2018': 275,'TTWJetsToLNu_UL2018': 276,'tZq_ll_4f_UL2018': 277,'VG_UL2018': 278,'ZG_UL2018': 279,'WG_UL2018': 280,'WrongSign_UL2018': 281,'WWto2L2Nu_UL2018': 282,'GluGluToWWToENEN_UL2018': 283,'GluGluToWWToENMN_UL2018': 284,'GluGluToWWToENTN_UL2018': 285,'GluGluToWWToMNEN_UL2018': 286,'GluGluToWWToMNMN_UL2018': 287,'GluGluToWWToMNTN_UL2018': 288,'GluGluToWWToTNEN_UL2018': 289,'GluGluToWWToTNMN_UL2018': 290,'GluGluToWWToTNTN_UL2018': 291,'ST_tW_top_UL2018': 292,'ST_tW_antitop_UL2018': 293,'GluGluHToWWTo2L2Nu_UL2018': 294,'GluGluHToWWToLNuQQ_UL2018': 295,'GluGluHToZZTo4L_UL2018': 296,'GluGluHToTauTau_UL2018': 297,'VBFHToWWTo2L2Nu_UL2018': 298,'VBFHToTauTau_UL2018': 299,'ttHToNonbb_UL2018': 300,'VHToNonbb_UL2018': 301,'Triboson_UL2018': 302,'WWTo2L2Nu_DoubleScattering_UL2018': 303,'WWW_4F_UL2018': 304,'WWZ_4F_UL2018': 305,'WZZ_UL2018': 306,'ZZZ_UL2018': 307,'WWG_UL2018': 308,'WJets_UL2018': 309,'WJetsHT70to100_UL2018': 310,'WJetsHT100to200_UL2018': 311,'WJetsHT200to400_UL2018': 312,'WJetsHT400to600_UL2018': 313,'WJetsHT600to800_UL2018': 314,'WJetsHT800to1200_UL2018': 315,'WJetsHT1200to2500_UL2018': 316,'WJetsHT2500toInf_UL2018': 317,'WZ_UL2018': 318,'DYJetsToLL_UL2018': 319,'DYJetsToLL_FxFx_UL2018': 319,'DYJetsToLL_M10to50_UL2018': 320,'DYJetsToLL_M50_UL2018': 321,'DYJetsToLL_M50_FxFx_UL2018': 321,'DYJetsToLL_M50_UL2018_ext': 321,'WpWpJJ_EWK_UL2018': 322,'WpWpJJ_QCD_UL2018': 323,'VBS_SSWW_SM_UL2018': 324,'VBS_SSWW_LL_SM_UL2018': 325,'VBS_SSWW_TL_SM_UL2018': 326,'VBS_SSWW_TT_SM_UL2018': 327,'VBS_SSWW_cW_UL2018': 328,'VBS_SSWW_cW_BSM_UL2018': 329,'VBS_SSWW_cW_SM_UL2018': 330,'VBS_SSWW_cW_INT_UL2018': 331,'VBS_SSWW_cHW_UL2018': 332,'VBS_SSWW_cHW_SM_UL2018': 333,'VBS_SSWW_cHW_BSM_UL2018': 334,'VBS_SSWW_cHW_INT_UL2018': 335,'VBS_SSWW_cW_cHW_UL2018': 336,'VBS_SSWW_DIM6_UL2018': 337,'VBS_SSWW_DIM6_SM_UL2018': 338,'DataMu_UL2016APV': 339,'DataMuB1_UL2016APV': 340,'DataMuB2_UL2016APV': 341,'DataMuC_UL2016APV': 342,'DataMuD_UL2016APV': 343,'DataMuE_UL2016APV': 344,'DataMuF_UL2016APV': 345,'DataMu_UL2016': 346,'DataMuF_UL2016': 347,'DataMuG_UL2016': 348,'DataMuH_UL2016': 349,'DataMu_UL2017': 350,'DataMuB_UL2017': 351,'DataMuC_UL2017': 352,'DataMuD_UL2017': 353,'DataMuE_UL2017': 354,'DataMuF_UL2017': 355,'DataMu_UL2018': 356,'DataMuA_UL2018': 357,'DataMuB_UL2018': 358,'DataMuC_UL2018': 359,'DataMuD_UL2018': 360,'DataEle_UL2016APV': 361,'DataEleB1_UL2016APV': 362,'DataEleB2_UL2016APV': 363,'DataEleC_UL2016APV': 364,'DataEleD_UL2016APV': 365,'DataEleE_UL2016APV': 366,'DataEleF_UL2016APV': 367,'DataEle_UL2016': 368,'DataEleF_UL2016': 369,'DataEleG_UL2016': 370,'DataEleH_UL2016': 371,'DataEle_UL2017': 372,'DataEleB_UL2017': 373,'DataEleC_UL2017': 374,'DataEleD_UL2017': 375,'DataEleE_UL2017': 376,'DataEleF_UL2017': 377,'DataEle_UL2018': 378,'DataEleA_UL2018': 379,'DataEleB_UL2018': 380,'DataEleC_UL2018': 381,'DataEleD_UL2018': 382,'DataHT_UL2016APV': 383,'DataHTB1_UL2016APV': 384,'DataHTB2_UL2016APV': 385,'DataHTC_UL2016APV': 386,'DataHTD_UL2016APV': 387,'DataHTE_UL2016APV': 388,'DataHTF_UL2016APV': 389,'DataHT_UL2016': 390,'DataHTF_UL2016': 391,'DataHTG_UL2016': 392,'DataHTH_UL2016': 393,'DataHT_UL2017': 394,'DataHTB_UL2017': 395,'DataHTC_UL2017': 396,'DataHTD_UL2017': 397,'DataHTE_UL2017': 398,'DataHTF_UL2017': 399,'DataHT_UL2018': 400,'DataHTA_UL2018': 401,'DataHTB_UL2018': 402,'DataHTC_UL2018': 403,'DataHTD_UL2018': 404,'SampleHTFake_UL2016APV': 405,'SampleHTFake_UL2016': 406,'SampleHTFake_UL2017': 407,'SampleHTFake_UL2018': 408,}\n",
    "\n",
    "aggregated_samples_UL2017 = {\n",
    "    'VG':  [ZG_UL2017, WG_UL2017],\n",
    "    'TVX': [TTGJets_UL2017, TTZToQQ_UL2017, TTZToLLNuNu_UL2017, TTWJetsToQQ_UL2017, TTWJetsToLNu_UL2017, tZq_ll_4f_UL2017], \n",
    "    'Triboson': [WWW_4F_UL2017, WWZ_4F_UL2017, WZZ_UL2017, ZZZ_UL2017, WWG_UL2017], \n",
    "    'TTTo2L2Nu': [TTTo2L2Nu_UL2017], \n",
    "    'WZ': [WZ_UL2017], \n",
    "    'DYJetsToLL_FxFx': [DYJetsToLL_M50_FxFx_UL2017],\n",
    "    'WrongSign': [WWto2L2Nu_UL2017, GluGluToWWToENEN_UL2017, GluGluToWWToENMN_UL2017, GluGluToWWToENTN_UL2017, GluGluToWWToMNEN_UL2017, GluGluToWWToMNMN_UL2017, GluGluToWWToMNTN_UL2017, GluGluToWWToTNEN_UL2017, \n",
    "                  GluGluToWWToTNMN_UL2017, \n",
    "                  GluGluToWWToTNTN_UL2017, ST_tW_top_UL2017, \n",
    "                  ST_tW_antitop_UL2017, \n",
    "                  GluGluHToWWTo2L2Nu_UL2017, \n",
    "                  GluGluHToZZTo4L_UL2017, GluGluHToTauTau_UL2017, VBFHToWWTo2L2Nu_UL2017, VBFHToTauTau_UL2017, ttHToNonbb_UL2017, VHToNonbb_UL2017\n",
    "    ], \n",
    "    'ZZtoLep': [ZZTo2L2Nu_UL2017, ZZTo4L_UL2017, GluGluToContinToZZTo2e2nu_UL2017, GluGluToContinToZZTo2e2mu_UL2017, GluGluToContinToZZTo2e2tau_UL2017, GluGluToContinToZZTo2mu2nu_UL2017, GluGluToContinToZZTo2mu2tau_UL2017, GluGluToContinToZZTo4e_UL2017, GluGluToContinToZZTo4mu_UL2017, GluGluToContinToZZTo4tau_UL2017], \n",
    "    'VBS_SSWW_SM': [VBS_SSWW_LL_SM_UL2017, VBS_SSWW_TL_SM_UL2017, VBS_SSWW_TT_SM_UL2017], \n",
    "}\n",
    "\n",
    "aggregated_samples = aggregated_samples_UL2017"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e205f9d7-a087-42ab-89de-d838de753b21",
   "metadata": {},
   "source": [
    "The lis of files to be read is built via Rucio."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "be9029ea-cb88-47dc-9f38-51cd3e766bec",
   "metadata": {},
   "outputs": [],
   "source": [
    "from rucio.client.client import Client\n",
    "import os\n",
    "\n",
    "import logging\n",
    "logger = logging.getLogger(__name__)\n",
    "logging.disable('WARNING')\n",
    "\n",
    "os.environ[\"RUCIO_HOME\"] = \"/cvmfs/cms.cern.ch/rucio/current/\"\n",
    "os.environ['X509_CERT_DIR'] = \"/cvmfs/grid.cern.ch/etc/grid-security/certificates/\"\n",
    "os.environ['X509_USER_PROXY'] = \"/opt/workspace/persistent-storage/proxy\"\n",
    "RUCIO_SCOPE = \"cms\"\n",
    "ACCOUNT = \"ttedesch\"\n",
    "\n",
    "SITENAME = \"T2_IT_Legnaro\"\n",
    "\n",
    "c = Client(account=ACCOUNT, auth_type=\"x509_proxy\")\n",
    "\n",
    "chain = []\n",
    "chain_gluglu = []\n",
    "chain_WZ = []\n",
    "\n",
    "proxy = redirector\n",
    "max_lines = maxNfilespersample\n",
    "\n",
    "for k,v in aggregated_samples.items():\n",
    "    for s in v:\n",
    "        CMS_DATASET_NAME = s.dataset\n",
    "        files = c.list_files(scope=RUCIO_SCOPE, name=CMS_DATASET_NAME)\n",
    "        lfns = []\n",
    "        for f in files:\n",
    "            lfns.append(RUCIO_SCOPE + ':' + f['name'])\n",
    "        for i,line in enumerate(lfns):\n",
    "            if i < max_lines and not line[4:] in []:\n",
    "                if k != 'Data':\n",
    "                    if s.name in [\"GluGluToWWToTNMN_UL2017\", \"GluGluHToWWTo2L2Nu_UL2017\"]:\n",
    "                        chain_gluglu.append(redirector + line[4:])\n",
    "                    elif s.name in [\"WZ_UL2017\"]:\n",
    "                        chain_WZ.append(redirector + line[4:])\n",
    "                    else:\n",
    "                        chain.append(redirector + line[4:])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fa6438e0-b4fd-4698-bdb9-1d8aee1b2b67",
   "metadata": {},
   "source": [
    "Define the list of branches to be saved in the Snapshot."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "1664bd3c-6f56-4dc3-864d-3afcc8601c63",
   "metadata": {},
   "outputs": [],
   "source": [
    "branchlist = [\n",
    "    \n",
    "       \"nLHEPdfWeight\",\n",
    "    \"LHEPdfWeight\",\n",
    "    \"nLHEScaleWeight\",\n",
    "    \"LHEScaleWeight\",\n",
    "    \n",
    "    \"Generator_weight\",\n",
    "    \"nPSWeight\",\n",
    "    \"PSWeight\",\n",
    "    \n",
    "    \"Sample\",\n",
    "    \n",
    "    \"HLT_IsoMu27\",\n",
    "    \"HLT_Ele35_WPTight_Gsf\",\n",
    "    \"HLT_Mu50\", \n",
    "    \"HLT_Ele32_WPTight_Gsf_L1DoubleEG\", \n",
    "    \"HLT_Photon200\", \n",
    "    \"HLT_PFHT250\", \n",
    "    \"HLT_PFHT350\",\n",
    "    \n",
    "    \"nJet\",\n",
    "    \"Jet_jetId\", \n",
    "    \"Jet_eta\", \n",
    "    \"Jet_pt\", \n",
    "    \"Jet_puId\",\n",
    "    \"Jet_phi\", \n",
    "    \"Jet_mass\",\n",
    "    \"Jet_partonFlavour\", \n",
    "    \"Jet_btagDeepFlavB\", \n",
    "    \"Jet_btagSF_deepjet_M_up\", \n",
    "    \"Jet_btagSF_deepjet_M_down\", \n",
    "    \"Jet_btagSF_deepjet_M\", \n",
    "    \"Jet_btagDeepB\",\n",
    "    \n",
    "    \"nElectron\",\n",
    "    \"Electron_pt\",\n",
    "    \"Electron_eta\", \n",
    "    \"Electron_phi\",\n",
    "    \"Electron_mass\",\n",
    "    \"Electron_jetRelIso\", \n",
    "    \"Electron_mvaFall17V2Iso_WPL\", \n",
    "    \"Electron_mvaFall17V2Iso_WP90\",\n",
    "    \"Electron_pdgId\",\n",
    "    \"Electron_genPartFlav\",\n",
    "    \"Electron_charge\",\n",
    "    \n",
    "    \"nMuon\",\n",
    "    \"Muon_corrected_pt\", \n",
    "    \"Muon_eta\", \n",
    "    \"Muon_phi\",\n",
    "    \"Muon_mass\",\n",
    "    \"Muon_tightId\", \n",
    "    \"Muon_looseId\", \n",
    "    \"Muon_pfRelIso04_all\",\n",
    "    \"Muon_pdgId\",\n",
    "    \"Muon_genPartFlav\",\n",
    "    \"Muon_charge\",\n",
    "    \n",
    "    \"nTau\",\n",
    "    \"Tau_pt\",\n",
    "    \"Tau_eta\", \n",
    "    \"Tau_phi\", \n",
    "    \"Tau_idDeepTau2017v2p1VSjet\", \n",
    "    \"Tau_idDeepTau2017v2p1VSe\", \n",
    "    \"Tau_idDeepTau2017v2p1VSmu\", \n",
    "    \"Tau_mass\",\n",
    "    \"Tau_charge\",\n",
    "    \"Tau_leadTkPtOverTauPt\",\n",
    "    \"Tau_decayMode\",\n",
    "    \"Tau_neutralIso\",\n",
    "    \"Tau_genPartFlav\",\n",
    "    \"Tau_jetIdx\",\n",
    "    \n",
    "    \"Electron_effSF\", \n",
    "    \"Muon_effSF\",\n",
    "    \"Electron_effSF_errUp\", \n",
    "    \"Muon_effSF_errUp\",\n",
    "    \"Electron_effSF_errDown\", \n",
    "    \"Muon_effSF_errDown\",\n",
    "    \n",
    "    \"Jet_pt_nom\",\n",
    "    \"Jet_pt_jerDown\",\n",
    "    \"Jet_pt_jerUp\",\n",
    "    \"Jet_pt_jesTotalDown\",\n",
    "    \"Jet_pt_jesTotalUp\",\n",
    "    \"Jet_mass_nom\",\n",
    "    \"Jet_mass_jerDown\",\n",
    "    \"Jet_mass_jerUp\",\n",
    "    \"Jet_mass_jesTotalDown\",\n",
    "    \"Jet_mass_jesTotalUp\",\n",
    "    \n",
    "    \"MET_T1Smear_pt\",\n",
    "    \"MET_T1Smear_pt_jerDown\",\n",
    "    \"MET_T1Smear_pt_jerUp\",\n",
    "    \"MET_T1Smear_pt_jesTotalDown\",\n",
    "    \"MET_T1Smear_pt_jesTotalUp\",\n",
    "    \"MET_T1Smear_phi\",\n",
    "    \"MET_T1Smear_phi_jerDown\",\n",
    "    \"MET_T1Smear_phi_jerUp\",\n",
    "    \"MET_T1Smear_phi_jesTotalDown\",\n",
    "    \"MET_T1Smear_phi_jesTotalUp\",\n",
    "    \n",
    "    \"PrefireWeight\",\n",
    "    \"PrefireWeight_Down\",\n",
    "    \"PrefireWeight_Up\",\n",
    "    \"puWeight\", \n",
    "    \"puWeightDown\", \n",
    "    \"puWeightUp\",\n",
    "    \n",
    "    \"Tau_rawDeepTau2017v2p1VSjet\",\n",
    "    \"Tau_rawDeepTau2017v2p1VSe\",\n",
    "    \"Tau_rawDeepTau2017v2p1VSmu\",\n",
    "    \"Tau_leadTkDeltaPhi\",\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9daeeac7-da1d-4259-b41f-afc1843913de",
   "metadata": {},
   "source": [
    "Build the computational graph. Output files are saved locally."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "668843b9-d83d-4ae3-9e6a-f3965508d545",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 7 µs, sys: 0 ns, total: 7 µs\n",
      "Wall time: 10.7 µs\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "def execute_MC(chain, branchlist_, outFilePath = \"./preselectionUL.root\", outTreeName = \"Events\",  nPart = nPartitions, useFlag_BadPFMuonDzFilter = True, LHE = True, label = \"distrdf\" ):\n",
    "    if distributed == True:\n",
    "        #df = RDataFrame(\"Events\", chain, npartitions=nPart, daskclient=client, monitor_label = label)  #when using root version with monitoring features (/cvmfs/images.dodas.infn.it/registry.hub.docker.com/dodasts/root-in-docker:ubuntu22-kernel-v1-monitoring)\n",
    "        df = RDataFrame(\"Events\", chain, npartitions=nPart, daskclient=client)  #when using standard root versions\n",
    "\n",
    "    else:\n",
    "        #df = RDataFrame(\"Events\", chain) #to run on all\n",
    "        df = RDataFrame(\"Events\", chain[0])\n",
    "\n",
    "    df_year = df.DefinePerSample(\"Year\",\"GetYear(rdfslot_, rdfsampleinfo_)\")\n",
    "    df_sample = df_year.DefinePerSample(\"Sample\", \"GetSample(rdfslot_, rdfsampleinfo_)\")\n",
    "\n",
    "    if useFlag_BadPFMuonDzFilter == True:\n",
    "        df_METHLTFilter = df_sample.Define(\"METHLTFilter\", \"MET_HLT_Filter_UL2017(Year, Flag_goodVertices, Flag_HBHENoiseFilter, Flag_HBHENoiseIsoFilter, Flag_EcalDeadCellTriggerPrimitiveFilter, Flag_BadPFMuonFilter, Flag_globalSuperTightHalo2016Filter, HLT_IsoMu27, HLT_Mu50, HLT_OldMu100, HLT_TkMu100,  HLT_Ele35_WPTight_Gsf, HLT_Ele32_WPTight_Gsf_L1DoubleEG, HLT_Photon200, Flag_ecalBadCalibFilter, Flag_BadPFMuonDzFilter, L1_SingleIsoEG30er2p1, L1_SingleIsoEG32, L1_SingleEG40, Flag_eeBadScFilter)\")\\\n",
    "                                   .Filter(\"METHLTFilter == true\", \"MET HLT Filter\")\n",
    "        \n",
    "    else:\n",
    "        df_METHLTFilter = df_sample.Define(\"METHLTFilter\", \"MET_HLT_Filter_UL2017_nodz(Year, Flag_goodVertices, Flag_HBHENoiseFilter, Flag_HBHENoiseIsoFilter, Flag_EcalDeadCellTriggerPrimitiveFilter, Flag_BadPFMuonFilter, Flag_globalSuperTightHalo2016Filter, HLT_IsoMu27, HLT_Mu50, HLT_OldMu100, HLT_TkMu100,  HLT_Ele35_WPTight_Gsf, HLT_Ele32_WPTight_Gsf_L1DoubleEG, HLT_Photon200, Flag_ecalBadCalibFilter, L1_SingleIsoEG30er2p1, L1_SingleIsoEG32, L1_SingleEG40, Flag_eeBadScFilter)\")\\\n",
    "                                   .Filter(\"METHLTFilter == true\", \"MET HLT Filter\")\n",
    "    #### preselection #####\n",
    "    df_preselection = df_METHLTFilter.Filter(\"PV_ndof> 4 && abs(PV_z) < 24 && hypot(PV_x, PV_y)<2\", \"Good vertex\")\\\n",
    "                                     .Define(\"HT_eventHT\", \"GetEventHT(Jet_pt, Jet_eta, Jet_phi, Jet_mass)\")\\\n",
    "                                     .Define(\"PassMinReq\", \"Pass_min_req(Muon_looseId, Muon_pt, Muon_pfRelIso04_all, Muon_eta, Electron_mvaFall17V2Iso_WPL, Electron_jetRelIso, Electron_pt, Electron_eta, Tau_idDeepTau2017v2p1VSjet, Tau_idDeepTau2017v2p1VSe, Tau_idDeepTau2017v2p1VSmu, Tau_pt, Tau_eta, Jet_pt, Jet_eta, Jet_puId)\")\\\n",
    "                                     .Filter(\"PassMinReq == true\", \"Minimum requests\")\n",
    "\n",
    "    #### lepSF #####\n",
    "    df_LepSF = df_preselection.Define(\"ElectronSFs\", \"ElectronSFs(Electron_pt, Electron_eta, Electron_pdgId, Year)\")\\\n",
    "                              .Define(\"Electron_effSF\", \"getFlattenedMatrixColumn(ElectronSFs, 3, 0)\")\\\n",
    "                              .Define(\"Electron_effSF_errUp\", \"getFlattenedMatrixColumn(ElectronSFs, 3, 1)\")\\\n",
    "                              .Define(\"Electron_effSF_errDown\", \"getFlattenedMatrixColumn(ElectronSFs, 3, 2)\")\\\n",
    "                              .Define(\"MuonSFs\", \"MuonSFs(Muon_pt, Muon_eta, Muon_pdgId, Year)\")\\\n",
    "                              .Define(\"Muon_effSF\", \"getFlattenedMatrixColumn(MuonSFs, 3, 0)\")\\\n",
    "                              .Define(\"Muon_effSF_errUp\", \"getFlattenedMatrixColumn(MuonSFs, 3, 1)\")\\\n",
    "                              .Define(\"Muon_effSF_errDown\", \"getFlattenedMatrixColumn(MuonSFs, 3, 2)\")\n",
    "\n",
    "    #### puWeight #####\n",
    "    #df_puWeight = df_mht\n",
    "    df_puWeight = df_LepSF.Define(\"puWeights\", \"puWeight(Year, Pileup_nTrueInt)\")\\\n",
    "                        .Define(\"puWeight\", \"puWeights[0]\")\\\n",
    "                        .Define(\"puWeightUp\", \"puWeights[1]\")\\\n",
    "                        .Define(\"puWeightDown\", \"puWeights[2]\")\n",
    "\n",
    "    #### prefCorr ####\n",
    "    df_prefCorr = df_puWeight.Define(\"prefCorrs\",\"PrefCorr(Photon_pt, Photon_eta, Photon_jetIdx, Photon_electronIdx, Electron_pt, Electron_eta, Electron_jetIdx, Electron_photonIdx, Jet_pt, Jet_eta)\")\\\n",
    "                             .Define(\"PrefireWeight\", \"prefCorrs[0]\")\\\n",
    "                             .Define(\"PrefireWeight_Up\", \"prefCorrs[1]\")\\\n",
    "                             .Define(\"PrefireWeight_Down\", \"prefCorrs[2]\")\n",
    "\n",
    "\n",
    "\n",
    "    #### btagSF ##### preselection part2\n",
    "    df_btagSF = df_prefCorr.Define(\"btagSFs\",  'btagSF(Jet_pt, Jet_eta, Jet_hadronFlavour, Jet_btagDeepFlavB, Year, \\\"M\\\")')\\\n",
    "                     .Define(\"Jet_btagSF_deepjet_M\", \"getMatrixColumn(btagSFs, 0)\")\\\n",
    "                     .Define(\"Jet_btagSF_deepjet_M_up\", \"getMatrixColumn(btagSFs, 1)\")\\\n",
    "                     .Define(\"Jet_btagSF_deepjet_M_down\", \"getMatrixColumn(btagSFs, 2)\")\n",
    "\n",
    "\n",
    "    #### muonScaleRes ####\n",
    "    df_muonScaleRes = df_btagSF.Define(\"muonCorrectedPTs\", \"muonScaleRes(Muon_pt, Muon_eta, Muon_phi, Muon_charge, Muon_nTrackerLayers, Muon_genPartIdx, GenPart_pt, Year)\")\\\n",
    "                               .Define(\"Muon_corrected_pt\", \"getFlattenedMatrixColumn(muonCorrectedPTs, 3, 0)\")\\\n",
    "                               .Define(\"Muon_correctedUp_pt\", \"getFlattenedMatrixColumn(muonCorrectedPTs, 3, 1)\")\\\n",
    "                               .Define(\"Muon_correctedDown_pt\", \"getFlattenedMatrixColumn(muonCorrectedPTs, 3, 2)\")\n",
    "\n",
    "\n",
    "    #### metCorrector and fatJetCorrector #### ####\n",
    "    df_jme = df_muonScaleRes.Define(\"jetVars\", \"myJetVarCalc.produce(Jet_pt, Jet_eta, Jet_phi, Jet_mass, Jet_rawFactor, Jet_area, Jet_jetId, fixedGridRhoFastjetAll, (run<<20) + (luminosityBlock<<10) + event + 1 + int(Jet_eta[0]/.01), GenJet_pt, GenJet_eta, GenJet_phi, GenJet_mass)\")\\\n",
    "                            .Define(\"Jet_pt_nom\", \"jetVars.pt(0)\")\\\n",
    "                            .Define(\"Jet_mass_nom\", \"jetVars.mass(0)\")\\\n",
    "                            .Define(\"metsVars\", \"myType1SmearMETVarCalc.produce(Jet_pt, Jet_eta, Jet_phi, Jet_mass, Jet_rawFactor, Jet_area, Jet_muonSubtrFactor, Jet_neEmEF, Jet_chEmEF, Jet_jetId, fixedGridRhoFastjetAll, (run<<20) + (luminosityBlock<<10) + event + 1 + int(Jet_eta[0]/.01), GenJet_pt, GenJet_eta, GenJet_phi, GenJet_mass, RawMET_phi , RawMET_pt, MET_MetUnclustEnUpDeltaX, MET_MetUnclustEnUpDeltaY, CorrT1METJet_rawPt, CorrT1METJet_eta, CorrT1METJet_phi, CorrT1METJet_area, CorrT1METJet_muonSubtrFactor, {}, {})\")\\\n",
    "                            .Define(\"MET_T1Smear_pt\", \"metsVars.pt(0)\")\\\n",
    "                            .Define(\"MET_T1Smear_phi\", \"metsVars.phi(0)\")\n",
    "\n",
    "    #jetvariations = variations\n",
    "    variations = jetvariations\n",
    "    for n,v in enumerate(jetvariations[1:]):\n",
    "        #print(\"Jet_pt_{}\".format(v))\n",
    "        if v in variations:\n",
    "            #print(\"Jet_pt_{}\".format(v))\n",
    "            df_jme = df_jme.Define(\"Jet_pt_{}\".format(v), \"jetVars.pt({})\".format(1+n)).Define(\"Jet_mass_{}\".format(v), \"jetVars.mass({})\".format(1+n))\n",
    "    for n,v in enumerate(mett1smearvariations[1:]):\n",
    "        if v in variations:\n",
    "            df_jme = df_jme.Define(\"MET_T1Smear_pt_{}\".format(v), \"metsVars.pt({})\".format(1+n)).Define(\"MET_T1Smear_phi_{}\".format(v), \"metsVars.phi({})\".format(1+n))  \n",
    "    \n",
    "    ### book snapshot ####\n",
    "    opts = ROOT.RDF.RSnapshotOptions()\n",
    "    opts.fLazy = True\n",
    "\n",
    "    if LHE == False:\n",
    "        branches = branchlist_[4:]\n",
    "    else:\n",
    "        branches = branchlist_\n",
    "    if distributed == True:\n",
    "        df_jme_lazy = df_jme.Snapshot(outTreeName, outFilePath, branches, opts)\n",
    "    else:\n",
    "        df_jme_lazy = df_jme.Snapshot(outTreeName, outFilePath, branches, opts)\n",
    "    counting = df_jme.Count()\n",
    "\n",
    "    return df_jme_lazy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "7a9c6cee-7e51-4bc6-8664-ee604a0d87f6",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Warning in <TClass::Init>: no dictionary for class edm::Hash<1> is available\n",
      "Warning in <TClass::Init>: no dictionary for class edm::ProcessHistory is available\n",
      "Warning in <TClass::Init>: no dictionary for class edm::ProcessConfiguration is available\n",
      "Warning in <TClass::Init>: no dictionary for class edm::ParameterSetBlob is available\n",
      "Warning in <TClass::Init>: no dictionary for class pair<edm::Hash<1>,edm::ParameterSetBlob> is available\n"
     ]
    }
   ],
   "source": [
    "df_sn = execute_MC(chain, branchlist, outFilePath = \"./preselectionUL.root\", outTreeName = \"Events\", label = \"main\") "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "24d62acf-446b-48a5-9729-35aa559ae50f",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_sn_gluglu = execute_MC(chain_gluglu, branchlist, outFilePath = \"./preselectionUL_GluGlu.root\",  outTreeName = \"Events\", nPart = 5, useFlag_BadPFMuonDzFilter = False, label = \"gluglu\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "2475b60e-3a33-4778-a73d-3ed122ad6b11",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_sn_WZ = execute_MC(chain_WZ, branchlist, outFilePath = \"./preselectionUL_WZ.root\", outTreeName = \"Events\", nPart = 5,  useFlag_BadPFMuonDzFilter = True, LHE = False, label = \"WZ\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "8abb446d-6592-4f71-aabf-a01768118481",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 44.2 s, sys: 574 ms, total: 44.8 s\n",
      "Wall time: 47.5 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "if distributed == True:\n",
    "    RunGraphs = ROOT.RDF.Experimental.Distributed.RunGraphs   \n",
    "else:\n",
    "    RunGraphs = ROOT.RDF.RunGraphs\n",
    "\n",
    "proxies = [\n",
    "    df_sn, \n",
    "    df_sn_gluglu, \n",
    "    df_sn_WZ\n",
    "]\n",
    "\n",
    "RunGraphs(proxies)\n",
    "dfs = [df_.GetValue() for df_ in proxies]"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python kernel from image /cvmfs/images.dodas.infn.it/registry.hub.docker.com/dodasts/root-in-docker:ubuntu22-kernel-v1-monitoring",
   "language": "python",
   "name": "python-kernel-from-image-_cvmfs_images.dodas.infn.it_registry.hub.docker.com_dodasts_root-in-dockerubuntu22-kernel-v1-monitoring"
  },
  "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.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}