ttbar_analysis_reana.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import logging\n",
    "import time\n",
    "\n",
    "import awkward as ak\n",
    "import cabinetry\n",
    "import cloudpickle\n",
    "import correctionlib\n",
    "from coffea import processor\n",
    "from coffea.nanoevents import NanoAODSchema\n",
    "from coffea.analysis_tools import PackedSelection\n",
    "import copy\n",
    "#from func_adl import ObjectStream\n",
    "#from func_adl_servicex import ServiceXSourceUpROOT\n",
    "import hist\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import pyhf\n",
    "\n",
    "import utils  # contains code for bookkeeping and cosmetics, as well as some boilerplate\n",
    "\n",
    "logging.getLogger(\"cabinetry\").setLevel(logging.INFO)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### GLOBAL CONFIGURATION\n",
    "# input files per process, set to e.g. 10 (smaller number = faster)\n",
    "N_FILES_MAX_PER_SAMPLE = -1\n",
    "\n",
    "# enable Dask\n",
    "USE_DASK = False\n",
    "\n",
    "# enable ServiceX\n",
    "USE_SERVICEX = False\n",
    "\n",
    "### ML-INFERENCE SETTINGS\n",
    "\n",
    "# enable ML inference\n",
    "USE_INFERENCE = False \n",
    "\n",
    "# enable inference using NVIDIA Triton server\n",
    "USE_TRITON = False\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class TtbarAnalysis(processor.ProcessorABC):\n",
    "    def __init__(self, use_inference, use_triton):\n",
    "\n",
    "        # initialize dictionary of hists for signal and control region\n",
    "        self.hist_dict = {}\n",
    "        for region in [\"4j1b\", \"4j2b\"]:\n",
    "            self.hist_dict[region] = (\n",
    "                hist.Hist.new.Reg(utils.config[\"global\"][\"NUM_BINS\"], \n",
    "                                  utils.config[\"global\"][\"BIN_LOW\"], \n",
    "                                  utils.config[\"global\"][\"BIN_HIGH\"], \n",
    "                                  name=\"observable\", \n",
    "                                  label=\"observable [GeV]\")\n",
    "                .StrCat([], name=\"process\", label=\"Process\", growth=True)\n",
    "                .StrCat([], name=\"variation\", label=\"Systematic variation\", growth=True)\n",
    "                .Weight()\n",
    "            )\n",
    "        \n",
    "        self.cset = correctionlib.CorrectionSet.from_file(\"corrections.json\")\n",
    "        self.use_inference = use_inference\n",
    "        \n",
    "        # set up attributes only needed if USE_INFERENCE=True\n",
    "        if self.use_inference:\n",
    "            \n",
    "            # initialize dictionary of hists for ML observables\n",
    "            self.ml_hist_dict = {}\n",
    "            for i in range(len(utils.config[\"ml\"][\"FEATURE_NAMES\"])):\n",
    "                self.ml_hist_dict[utils.config[\"ml\"][\"FEATURE_NAMES\"][i]] = (\n",
    "                    hist.Hist.new.Reg(utils.config[\"global\"][\"NUM_BINS\"],\n",
    "                                      utils.config[\"ml\"][\"BIN_LOW\"][i],\n",
    "                                      utils.config[\"ml\"][\"BIN_HIGH\"][i],\n",
    "                                      name=\"observable\",\n",
    "                                      label=utils.config[\"ml\"][\"FEATURE_DESCRIPTIONS\"][i])\n",
    "                    .StrCat([], name=\"process\", label=\"Process\", growth=True)\n",
    "                    .StrCat([], name=\"variation\", label=\"Systematic variation\", growth=True)\n",
    "                    .Weight()\n",
    "                )\n",
    "            \n",
    "            self.use_triton = use_triton\n",
    "\n",
    "    def only_do_IO(self, events):\n",
    "        for branch in utils.config[\"benchmarking\"][\"IO_BRANCHES\"][\n",
    "            utils.config[\"benchmarking\"][\"IO_FILE_PERCENT\"]\n",
    "        ]:\n",
    "            if \"_\" in branch:\n",
    "                split = branch.split(\"_\")\n",
    "                object_type = split[0]\n",
    "                property_name = \"_\".join(split[1:])\n",
    "                ak.materialized(events[object_type][property_name])\n",
    "            else:\n",
    "                ak.materialized(events[branch])\n",
    "        return {\"hist\": {}}\n",
    "\n",
    "    def process(self, events):\n",
    "        if utils.config[\"benchmarking\"][\"DISABLE_PROCESSING\"]:\n",
    "            # IO testing with no subsequent processing\n",
    "            return self.only_do_IO(events)\n",
    "\n",
    "        # create copies of histogram objects\n",
    "        hist_dict = copy.deepcopy(self.hist_dict)\n",
    "        if self.use_inference:\n",
    "            ml_hist_dict = copy.deepcopy(self.ml_hist_dict)\n",
    "\n",
    "        process = events.metadata[\"process\"]  # \"ttbar\" etc.\n",
    "        variation = events.metadata[\"variation\"]  # \"nominal\" etc.\n",
    "\n",
    "        # normalization for MC\n",
    "        x_sec = events.metadata[\"xsec\"]\n",
    "        nevts_total = events.metadata[\"nevts\"]\n",
    "        lumi = 3378 # /pb\n",
    "        if process != \"data\":\n",
    "            xsec_weight = x_sec * lumi / nevts_total\n",
    "        else:\n",
    "            xsec_weight = 1\n",
    "\n",
    "        # setup triton gRPC client\n",
    "        if self.use_inference and self.use_triton:\n",
    "            triton_client = utils.clients.get_triton_client(utils.config[\"ml\"][\"TRITON_URL\"])\n",
    "\n",
    "\n",
    "        #### systematics\n",
    "        # jet energy scale / resolution systematics\n",
    "        # need to adjust schema to instead use coffea add_systematic feature, especially for ServiceX\n",
    "        # cannot attach pT variations to events.jet, so attach to events directly\n",
    "        # and subsequently scale pT by these scale factors\n",
    "        events[\"pt_scale_up\"] = 1.03\n",
    "        events[\"pt_res_up\"] = utils.systematics.jet_pt_resolution(events.Jet.pt)\n",
    "\n",
    "        syst_variations = [\"nominal\"]\n",
    "        jet_kinematic_systs = [\"pt_scale_up\", \"pt_res_up\"]\n",
    "        event_systs = [f\"btag_var_{i}\" for i in range(4)]\n",
    "        if process == \"wjets\":\n",
    "            event_systs.append(\"scale_var\")\n",
    "\n",
    "        # Only do systematics for nominal samples, e.g. ttbar__nominal\n",
    "        if variation == \"nominal\":\n",
    "            syst_variations.extend(jet_kinematic_systs)\n",
    "            syst_variations.extend(event_systs)\n",
    "\n",
    "        # for pt_var in pt_variations:\n",
    "        for syst_var in syst_variations:\n",
    "            ### event selection\n",
    "            # very very loosely based on https://arxiv.org/abs/2006.13076\n",
    "\n",
    "            # Note: This creates new objects, distinct from those in the 'events' object\n",
    "            elecs = events.Electron\n",
    "            muons = events.Muon\n",
    "            jets = events.Jet\n",
    "            if syst_var in jet_kinematic_systs:\n",
    "                # Replace jet.pt with the adjusted values\n",
    "                jets[\"pt\"] = jets.pt * events[syst_var]\n",
    "\n",
    "            electron_reqs = (elecs.pt > 30) & (np.abs(elecs.eta) < 2.1) & (elecs.cutBased == 4) & (elecs.sip3d < 4)\n",
    "            muon_reqs = ((muons.pt > 30) & (np.abs(muons.eta) < 2.1) & (muons.tightId) & (muons.sip3d < 4) &\n",
    "                         (muons.pfRelIso04_all < 0.15))\n",
    "            jet_reqs = (jets.pt > 30) & (np.abs(jets.eta) < 2.4) & (jets.isTightLeptonVeto)\n",
    "\n",
    "            # Only keep objects that pass our requirements\n",
    "            elecs = elecs[electron_reqs]\n",
    "            muons = muons[muon_reqs]\n",
    "            jets = jets[jet_reqs]\n",
    "\n",
    "            if self.use_inference:\n",
    "                even = (events.event%2==0)  # whether events are even/odd\n",
    "\n",
    "            B_TAG_THRESHOLD = 0.5\n",
    "\n",
    "            ######### Store boolean masks with PackedSelection ##########\n",
    "            selections = PackedSelection(dtype='uint64')\n",
    "            # Basic selection criteria\n",
    "            selections.add(\"exactly_1l\", (ak.num(elecs) + ak.num(muons)) == 1)\n",
    "            selections.add(\"atleast_4j\", ak.num(jets) >= 4)\n",
    "            selections.add(\"exactly_1b\", ak.sum(jets.btagCSVV2 > B_TAG_THRESHOLD, axis=1) == 1)\n",
    "            selections.add(\"atleast_2b\", ak.sum(jets.btagCSVV2 > B_TAG_THRESHOLD, axis=1) >= 2)\n",
    "            # Complex selection criteria\n",
    "            selections.add(\"4j1b\", selections.all(\"exactly_1l\", \"atleast_4j\", \"exactly_1b\"))\n",
    "            selections.add(\"4j2b\", selections.all(\"exactly_1l\", \"atleast_4j\", \"atleast_2b\"))\n",
    "\n",
    "            for region in [\"4j1b\", \"4j2b\"]:\n",
    "                region_selection = selections.all(region)\n",
    "                region_jets = jets[region_selection]\n",
    "                region_elecs = elecs[region_selection]\n",
    "                region_muons = muons[region_selection]\n",
    "                region_weights = np.ones(len(region_jets)) * xsec_weight\n",
    "                if self.use_inference:\n",
    "                    region_even = even[region_selection]\n",
    "\n",
    "                if region == \"4j1b\":\n",
    "                    observable = ak.sum(region_jets.pt, axis=-1)\n",
    "\n",
    "                elif region == \"4j2b\":\n",
    "\n",
    "                    # reconstruct hadronic top as bjj system with largest pT\n",
    "                    trijet = ak.combinations(region_jets, 3, fields=[\"j1\", \"j2\", \"j3\"])  # trijet candidates\n",
    "                    trijet[\"p4\"] = trijet.j1 + trijet.j2 + trijet.j3  # calculate four-momentum of tri-jet system\n",
    "                    trijet[\"max_btag\"] = np.maximum(trijet.j1.btagCSVV2, np.maximum(trijet.j2.btagCSVV2, trijet.j3.btagCSVV2))\n",
    "                    trijet = trijet[trijet.max_btag > B_TAG_THRESHOLD]  # at least one-btag in trijet candidates\n",
    "                    # pick trijet candidate with largest pT and calculate mass of system\n",
    "                    trijet_mass = trijet[\"p4\"][ak.argmax(trijet.p4.pt, axis=1, keepdims=True)].mass\n",
    "                    observable = ak.flatten(trijet_mass)\n",
    "\n",
    "                    if sum(region_selection)==0:\n",
    "                        continue\n",
    "\n",
    "                    if self.use_inference:\n",
    "                        features, perm_counts = utils.ml.get_features(\n",
    "                            region_jets,\n",
    "                            region_elecs,\n",
    "                            region_muons,\n",
    "                            max_n_jets=utils.config[\"ml\"][\"MAX_N_JETS\"],\n",
    "                        )\n",
    "                        even_perm = np.repeat(region_even, perm_counts)\n",
    "\n",
    "                        # calculate ml observable\n",
    "                        if self.use_triton:\n",
    "                            results = utils.ml.get_inference_results_triton(\n",
    "                                features,\n",
    "                                even_perm,\n",
    "                                triton_client,\n",
    "                                utils.config[\"ml\"][\"MODEL_NAME\"],\n",
    "                                utils.config[\"ml\"][\"MODEL_VERSION_EVEN\"],\n",
    "                                utils.config[\"ml\"][\"MODEL_VERSION_ODD\"],\n",
    "                            )\n",
    "\n",
    "                        else:\n",
    "                            results = utils.ml.get_inference_results_local(\n",
    "                                features,\n",
    "                                even_perm,\n",
    "                                utils.ml.model_even,\n",
    "                                utils.ml.model_odd,\n",
    "                            )\n",
    "                            \n",
    "                        results = ak.unflatten(results, perm_counts)\n",
    "                        features = ak.flatten(ak.unflatten(features, perm_counts)[\n",
    "                            ak.from_regular(ak.argmax(results,axis=1)[:, np.newaxis])\n",
    "                        ])\n",
    "                syst_var_name = f\"{syst_var}\"\n",
    "                # Break up the filling into event weight systematics and object variation systematics\n",
    "                if syst_var in event_systs:\n",
    "                    for i_dir, direction in enumerate([\"up\", \"down\"]):\n",
    "                        # Should be an event weight systematic with an up/down variation\n",
    "                        if syst_var.startswith(\"btag_var\"):\n",
    "                            i_jet = int(syst_var.rsplit(\"_\",1)[-1])   # Kind of fragile\n",
    "                            wgt_variation = self.cset[\"event_systematics\"].evaluate(\"btag_var\", direction, region_jets.pt[:,i_jet])\n",
    "                        elif syst_var == \"scale_var\":\n",
    "                            # The pt array is only used to make sure the output array has the correct shape\n",
    "                            wgt_variation = self.cset[\"event_systematics\"].evaluate(\"scale_var\", direction, region_jets.pt[:,0])\n",
    "                        syst_var_name = f\"{syst_var}_{direction}\"\n",
    "                        hist_dict[region].fill(\n",
    "                            observable=observable, process=process,\n",
    "                            variation=syst_var_name, weight=region_weights * wgt_variation\n",
    "                        )\n",
    "                        if region == \"4j2b\" and self.use_inference:\n",
    "                            for i in range(len(utils.config[\"ml\"][\"FEATURE_NAMES\"])):\n",
    "                                ml_hist_dict[utils.config[\"ml\"][\"FEATURE_NAMES\"][i]].fill(\n",
    "                                    observable=features[..., i], process=process,\n",
    "                                    variation=syst_var_name, weight=region_weights * wgt_variation\n",
    "                                )\n",
    "                else:\n",
    "                    # Should either be 'nominal' or an object variation systematic\n",
    "                    if variation != \"nominal\":\n",
    "                        # This is a 2-point systematic, e.g. ttbar__scaledown, ttbar__ME_var, etc.\n",
    "                        syst_var_name = variation\n",
    "                    hist_dict[region].fill(\n",
    "                        observable=observable, process=process,\n",
    "                        variation=syst_var_name, weight=region_weights\n",
    "                    )\n",
    "                    if region == \"4j2b\" and self.use_inference:\n",
    "                        for i in range(len(utils.config[\"ml\"][\"FEATURE_NAMES\"])):\n",
    "                            ml_hist_dict[utils.config[\"ml\"][\"FEATURE_NAMES\"][i]].fill(\n",
    "                                observable=features[..., i], process=process,\n",
    "                                variation=syst_var_name, weight=region_weights\n",
    "                            )\n",
    "\n",
    "\n",
    "        output = {\"nevents\": {events.metadata[\"dataset\"]: len(events)}, \"hist_dict\": hist_dict}\n",
    "        if self.use_inference:\n",
    "            output[\"ml_hist_dict\"] = ml_hist_dict\n",
    "\n",
    "        return output\n",
    "\n",
    "    def postprocess(self, accumulator):\n",
    "        return accumulator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "parts = sample_name.split('__')\n",
    "\n",
    "# Creating variables\n",
    "variation = parts[1]\n",
    "key_to_extract = parts[0] \n",
    "fileset = utils.file_input.construct_fileset(\n",
    "    N_FILES_MAX_PER_SAMPLE,\n",
    "    key_to_extract,\n",
    "    variation\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "original_dict = fileset\n",
    "selected_file = original_dict[sample_name]['files']\n",
    "new_dict = {sample_name: {'files': [filename], 'metadata': original_dict[sample_name]['metadata']}}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "file_parts = filename[38:] #For unl url change this number to 35\n",
    "files = file_parts[:-5]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "NanoAODSchema.warn_missing_crossrefs = False\n",
    "executor = processor.FuturesExecutor(workers=utils.config[\"benchmarking\"][\"NUM_CORES\"])\n",
    "\n",
    "NanoAODSchema.warn_missing_crossrefs = False # silences warnings about branches we will not use here\n",
    "executor = processor.FuturesExecutor(workers=utils.config[\"benchmarking\"][\"NUM_CORES\"])\n",
    "\n",
    "run = processor.Runner(\n",
    "    executor=executor, \n",
    "    schema=NanoAODSchema, \n",
    "    savemetrics=True, \n",
    "    metadata_cache={}, \n",
    "    chunksize=utils.config[\"benchmarking\"][\"CHUNKSIZE\"])\n",
    "\n",
    "treename = \"Events\"\n",
    "\n",
    "# load local models if not using Triton and models are not yet loaded\n",
    "if USE_INFERENCE and not USE_TRITON and utils.ml.model_even is None and utils.ml.model_odd is None:\n",
    "    utils.ml.load_models()\n",
    "print(new_dict)\n",
    "#print(os.environ)\n",
    "filemeta = run.preprocess(new_dict, treename=treename)  # pre-processing\n",
    "print(treename)\n",
    "print(filemeta)\n",
    "\n",
    "t0 = time.monotonic()\n",
    "all_histograms, metrics = run(\n",
    "    fileset={sample_name: new_dict[sample_name]}, \n",
    "    treename=treename, \n",
    "    processor_instance=TtbarAnalysis(USE_INFERENCE, USE_TRITON)\n",
    ")\n",
    "print(all_histograms)\n",
    "exec_time = time.monotonic() - t0\n",
    "print(f\"\\nexecution took {exec_time:.2f} seconds\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import uproot\n",
    "import os\n",
    "\n",
    "os.makedirs(os.path.split(f\"histograms/histograms_{sample_name}__{files}.root\")[0], exist_ok=True)\n",
    "# save all available histograms to disk\n",
    "with uproot.recreate(f\"histograms/histograms_{sample_name}__{files}.root\") as f:\n",
    "    for channel, histogram in all_histograms[\"hist_dict\"].items():\n",
    "        for sample in histogram.axes[1]:\n",
    "            for variation in histogram[:, sample, :].axes[1]:\n",
    "                variation_string = \"\" if variation == \"nominal\" else f\"_{variation}\"\n",
    "                if sum(histogram[:, sample, variation].values()) != 0:\n",
    "                    # only save histograms containing events\n",
    "                    # many combinations are not used (e.g. ME var for W+jets)\n",
    "                    f[f\"{channel}_{sample}{variation_string}\"] = histogram[:, sample, variation]"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}