fail_data.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "ec6c4165-5d04-46d8-b6a3-beda838a2cff",
   "metadata": {},
   "outputs": [],
   "source": [
    "import argparse\n",
    "import glob\n",
    "import json\n",
    "import logging\n",
    "import os\n",
    "import pickle as pkl\n",
    "import warnings\n",
    "\n",
    "import hist as hist2\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import pyarrow\n",
    "import yaml\n",
    "#from systematics import get_systematic_dict, sigs\n",
    "from systematicsPass import get_systematic_dict, sigs\n",
    "\n",
    "from utils import get_common_sample_name, get_finetuned_score, get_xsecweight\n",
    "\n",
    "logging.basicConfig(level=logging.INFO)\n",
    "\n",
    "warnings.filterwarnings(\"ignore\", message=\"Found duplicate branch \")\n",
    "pd.set_option(\"mode.chained_assignment\", None)\n",
    "\n",
    "CATEGORY = 'fail'\n",
    "\n",
    "\n",
    "\n",
    "def get_templates(years, channels, samples, samples_dir, regions_sel, model_path, add_fake=False):\n",
    "\n",
    "    # add extra selections to preselection\n",
    "    presel = { \"mu\": { \"fj_mass\": \"fj_mass < 180\",}, \"ele\": {\"fj_mass\": \"fj_mass <180\",}, }\n",
    "\n",
    "        \n",
    "    mass_binning = [40,70,120,180]\n",
    "\n",
    "    \n",
    "    hists = hist2.Hist(\n",
    "        hist2.axis.StrCategory([], name=\"Sample\", growth=True), hist2.axis.StrCategory([], name=\"Systematic\", growth=True), hist2.axis.StrCategory([], name=\"Region\", growth=True),\n",
    "        hist2.axis.Variable(\n",
    "            #list(range(55, 255, mass_binning)),\n",
    "            mass_binning, name=\"mass_observable\", label=r\"V reconstructed mass [GeV]\", overflow=True,\n",
    "        ), storage=hist2.storage.Weight(),\n",
    "    )\n",
    "\n",
    "    SYST_DICT = get_systematic_dict(years)\n",
    "    for year in years:  # e.g. 2018, 2017, 2016APV, 2016    \n",
    "        for ch in channels:  # e.g. mu, ele\n",
    "            logging.info(f\"Processing year {year} and {ch} channel\")\n",
    "            with open(\"../fileset/luminosity.json\") as f:\n",
    "                luminosity = json.load(f)[ch][year]\n",
    "            for sample in os.listdir(samples_dir[year]):\n",
    "                sample_to_use = get_common_sample_name(sample)\n",
    "                if sample_to_use not in samples:\n",
    "                    continue\n",
    "                is_data = True if sample_to_use == \"Data\" else False\n",
    "\n",
    "                logging.info(f\"Finding {sample} samples and should combine them under {sample_to_use}\")\n",
    "\n",
    "                out_files = f\"{samples_dir[year]}/{sample}/outfiles/\"\n",
    "                parquet_files = glob.glob(f\"{out_files}/*_{ch}.parquet\")\n",
    "                pkl_files = glob.glob(f\"{out_files}/*.pkl\")\n",
    "                if not parquet_files:\n",
    "                    logging.info(f\"No parquet file for {sample}\")\n",
    "                    continue\n",
    "                try:\n",
    "                    data = pd.read_parquet(parquet_files)\n",
    "                except pyarrow.lib.ArrowInvalid:  # empty parquet because no event passed selection\n",
    "                    continue\n",
    "                if len(data) == 0:\n",
    "                    continue\n",
    "                data[\"THWW\"] = get_finetuned_score(data, model_path)\n",
    "                data = data[data.columns.drop(list(data.filter(regex=\"hidNeuron\")))]\n",
    "\n",
    "                for selection in presel[ch]:\n",
    "                    logging.info(f\"Applying {selection} selection on {len(data)} events\")\n",
    "                    data = data.query(presel[ch][selection])\n",
    "\n",
    "#*****             # get the xsecweight\n",
    "                xsecweight, sumgenweights, sumpdfweights, sumscaleweights = get_xsecweight( pkl_files, year, sample, sample_to_use, is_data, luminosity)\n",
    "\n",
    "                for region, region_sel in regions_sel.items():  # e.g. pass, fail, top control region, etc.\n",
    "                #need to apply the V calibration only when making a cut on V tagger; currently the TTbar CR has no cut on tagger, so don't apply it\n",
    "#**************************************************************************************\n",
    "                    df = data.copy()\n",
    "                    logging.info(f\"Applying {region} selection on {len(df)} events\")\n",
    "                    df = df.query(region_sel)\n",
    "                    logging.info(f\"Will fill the histograms with the remaining {len(df)} events\")\n",
    "\n",
    "                    # ------------------- Nominal -------------------\n",
    "                    if is_data:\n",
    "                        nominal = np.ones_like(df[\"fj_pt\"])  # for data (nominal is 1)\n",
    "                    else:\n",
    "                        nominal = df[f\"weight_{ch}\"] * xsecweight * df[\"weight_btag\"] \n",
    "                 \n",
    "                    # ------------------- QCD scale -------------------\n",
    "                hists.fill( Sample=sample_to_use, Systematic=\"fail_nominal\", Region=region, mass_observable=df[\"fj_mass\"], weight=nominal,)\n",
    "        #for variation in [\"fakes_nominal\", \"fakes_SF_Up\", \"fakes_SF_Down\", \"fakes_DR_Up\", \"fakes_DR_Down\"]:\n",
    "\n",
    "    \n",
    "    return hists\n",
    "\n",
    "\n",
    "def fix_neg_yields(h):\n",
    "    \"\"\"\n",
    "    Will set the bin yields of a process to 0 if the nominal yield is negative, and will\n",
    "    set the yield to 0 for the full Systematic axis.\n",
    "    \"\"\"\n",
    "    for region in h.axes[\"Region\"]:\n",
    "        for sample in h.axes[\"Sample\"]:\n",
    "            neg_bins = np.where(h[{\"Sample\": sample, \"Systematic\": \"fail_nominal\", \"Region\": region}].values() < 0)[0]\n",
    "\n",
    "            if len(neg_bins) > 0:\n",
    "                print('got neg bins')\n",
    "                print(f\"{region}, {sample}, has {len(neg_bins)} bins with negative yield.. will set them to 0\")\n",
    "\n",
    "                sample_index = np.argmax(np.array(h.axes[\"Sample\"]) == sample)\n",
    "                region_index = np.argmax(np.array(h.axes[\"Region\"]) == region)\n",
    "\n",
    "                for neg_bin in neg_bins:\n",
    "                    h.view(flow=True)[sample_index, :, region_index, neg_bin + 1].value = 1e-3\n",
    "                    h.view(flow=True)[sample_index, :, region_index, neg_bin + 1].variance = 1e-3\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "2676f130-3778-4ea2-b179-e9cf1f6423f1",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:root:Processing year 2017 and mu channel\n",
      "INFO:root:Finding SingleElectron_Run2017B samples and should combine them under Data\n",
      "INFO:root:Finding SingleElectron_Run2017C samples and should combine them under Data\n",
      "INFO:root:Finding SingleElectron_Run2017D samples and should combine them under Data\n",
      "INFO:root:Finding SingleElectron_Run2017E samples and should combine them under Data\n",
      "INFO:root:Finding SingleElectron_Run2017F samples and should combine them under Data\n",
      "INFO:root:Finding SingleMuon_Run2017B samples and should combine them under Data\n",
      "INFO:root:Applying fj_mass selection on 12989 events\n",
      "INFO:root:Applying SR1 selection on 12098 events\n",
      "INFO:root:Will fill the histograms with the remaining 137 events\n",
      "INFO:root:Finding SingleMuon_Run2017C samples and should combine them under Data\n",
      "INFO:root:Applying fj_mass selection on 26595 events\n",
      "INFO:root:Applying SR1 selection on 24814 events\n",
      "INFO:root:Will fill the histograms with the remaining 234 events\n",
      "INFO:root:Finding SingleMuon_Run2017D samples and should combine them under Data\n",
      "INFO:root:Applying fj_mass selection on 11933 events\n",
      "INFO:root:Applying SR1 selection on 11123 events\n",
      "INFO:root:Will fill the histograms with the remaining 127 events\n",
      "INFO:root:Finding SingleMuon_Run2017E samples and should combine them under Data\n",
      "INFO:root:Applying fj_mass selection on 26183 events\n",
      "INFO:root:Applying SR1 selection on 24452 events\n",
      "INFO:root:Will fill the histograms with the remaining 236 events\n",
      "INFO:root:Finding SingleMuon_Run2017F samples and should combine them under Data\n",
      "INFO:root:Applying fj_mass selection on 38749 events\n",
      "INFO:root:Applying SR1 selection on 36272 events\n",
      "INFO:root:Will fill the histograms with the remaining 362 events\n",
      "INFO:root:Processing year 2017 and ele channel\n",
      "INFO:root:Finding SingleElectron_Run2017B samples and should combine them under Data\n",
      "INFO:root:Applying fj_mass selection on 6925 events\n",
      "INFO:root:Applying SR1 selection on 6469 events\n",
      "INFO:root:Will fill the histograms with the remaining 56 events\n",
      "INFO:root:Finding SingleElectron_Run2017C samples and should combine them under Data\n",
      "INFO:root:Applying fj_mass selection on 19373 events\n",
      "INFO:root:Applying SR1 selection on 18197 events\n",
      "INFO:root:Will fill the histograms with the remaining 164 events\n",
      "INFO:root:Finding SingleElectron_Run2017D samples and should combine them under Data\n",
      "INFO:root:Applying fj_mass selection on 8794 events\n",
      "INFO:root:Applying SR1 selection on 8239 events\n",
      "INFO:root:Will fill the histograms with the remaining 79 events\n",
      "INFO:root:Finding SingleElectron_Run2017E samples and should combine them under Data\n",
      "INFO:root:Applying fj_mass selection on 17872 events\n",
      "INFO:root:Applying SR1 selection on 16797 events\n",
      "INFO:root:Will fill the histograms with the remaining 180 events\n",
      "INFO:root:Finding SingleElectron_Run2017F samples and should combine them under Data\n",
      "INFO:root:Applying fj_mass selection on 24916 events\n",
      "INFO:root:Applying SR1 selection on 23304 events\n",
      "INFO:root:Will fill the histograms with the remaining 222 events\n",
      "INFO:root:Finding SingleMuon_Run2017B samples and should combine them under Data\n",
      "INFO:root:Finding SingleMuon_Run2017C samples and should combine them under Data\n",
      "INFO:root:Finding SingleMuon_Run2017D samples and should combine them under Data\n",
      "INFO:root:Finding SingleMuon_Run2017E samples and should combine them under Data\n",
      "INFO:root:Finding SingleMuon_Run2017F samples and should combine them under Data\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "hists Hist(\n",
      "  StrCategory(['Data'], growth=True, name='Sample'),\n",
      "  StrCategory(['fail_nominal'], growth=True, name='Systematic'),\n",
      "  StrCategory(['SR1'], growth=True, name='Region'),\n",
      "  Variable([40, 70, 120, 180], name='mass_observable', label='V reconstructed mass [GeV]'),\n",
      "  storage=Weight()) # Sum: WeightedSum(value=1797, variance=1797)\n"
     ]
    }
   ],
   "source": [
    "#years = ['2017','2018']\n",
    "years = ['2017']\n",
    "\n",
    "#years = ['2016', '2016APV','2017', '2018']\n",
    "#years = ['2017', '2018']\n",
    "\n",
    "outdir = 'templates'\n",
    "\n",
    "channels = 'mu','ele'\n",
    "with open(\"simpleFailData.yaml\", \"r\") as stream:\n",
    "    config = yaml.safe_load(stream)\n",
    "\n",
    "if len(years) == 4:\n",
    "    save_as = \"Run2\"\n",
    "else:\n",
    "    save_as = \"_\".join(years)\n",
    "\n",
    "if len(channels) == 1:\n",
    "    save_as += f\"_{channels[0]}_\"\n",
    "\n",
    "os.system(f\"mkdir -p {outdir}\")\n",
    "\n",
    "hists = get_templates( years, channels, config[\"samples\"], config[\"samples_dir\"], config[\"regions_sel\"], config[\"model_path\"],)\n",
    "\n",
    "fix_neg_yields(hists)\n",
    "\n",
    "with open(f\"{outdir}/hists_templates_{save_as}_data_fail.pkl\", \"wb\") as fp:\n",
    "    print('hists', hists)\n",
    "    pkl.dump(hists, fp)\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "411e06fa-20da-4f40-96bd-8bce7251cc21",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}