make_template_pd.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "source": [
    "# Make templates for fit (`pandas` workflow)\n",
    "\n",
    "This notebook aims to make the ROOT-format templates for fit. It reads the yaml config file and the backuped files from the previous notebook `preprocess.ipynb` to make the templates needed for the fit.\n",
    "\n",
    "We use `pandas` dataframe for event processing. An alternative notebook `make_template_ak.ipynb` is based on `ak-array` data structure. Using `pandas` can be generally faster to produce various similar templates we need for each shape systematics, because it shrinks the original dataset to a rather small dataframe by applying all pre-selections, then makes various templates based on the dataframe. Although, by reading the original large dataset into `pandas`, it sometimes contumes large RAM (10-30 GB)."
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "from data_utils import get_hist, plot_hist"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "import uproot3\n",
    "from uproot3_methods import TLorentzVectorArray, TLorentzVector\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import os"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "## Load the config.yml\n",
    "import yaml\n",
    "with open('cards/config_bb_std.yml') as f:\n",
    "    config = yaml.safe_load(f)"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Load files\n",
    "\n",
    "Load the ROOT files into pandas DataFrame"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "year = config['year']\n",
    "\n",
    "lumi = {2016: 35.92, 2017: 41.53, 2018: 59.74}\n",
    "\n",
    "read_sample_list_map = {\n",
    "    'qcd-mg-noht': 'mc/qcd-mg_tree.root',\n",
    "    'qcd-herwig-noht': 'mc/qcd-herwig_tree.root',\n",
    "    'top-noht': 'mc/top_tree.root',\n",
    "    'v-qq-noht': 'mc/v-qq_tree.root',\n",
    "    'jetht-noht': 'data/jetht_tree.root',\n",
    "}\n",
    "if config['samples']['use_bflav']:\n",
    "    read_sample_list_map['qcd-mg-bflav-noht'] = 'mc/qcd-mg-bflav_tree.root',\n",
    "omit_herwig = 'optional' in config['samples'] and 'omit_herwig' in config['samples']['optional'] and config['samples']['optional']['omit_herwig']\n",
    "if omit_herwig:\n",
    "    read_sample_list_map.pop('qcd-herwig-noht', None)\n",
    "if 'optional' in config['samples'] and 'exclude_mc_sample_in_making_template' in config['samples']['optional']:\n",
    "    for ex_sam in config['samples']['optional']['exclude_mc_sample_in_making_template']:\n",
    "        read_sample_list_map.pop(ex_sam, None)\n",
    "print('Read samples for making templates:', read_sample_list_map.keys())\n",
    "\n",
    "minimal_branches = set([  ## minimal set of branches read into the notebook\n",
    "    \"run\", \"luminosityBlock\", \"event\", \"genWeight\", \"jetR\", \"passmetfilters\", \n",
    "    \"fj_1_pt\", \"fj_1_eta\", \"fj_1_sdmass\", \"fj_1_tau21\", \"fj_1_btagcsvv2\", \"fj_1_btagjp\", \"fj_1_sfBDT\", \"fj_1_nbhadrons\", \"fj_1_nchadrons\", \"fj_1_sj1_nbhadrons\", \"fj_1_sj1_nchadrons\", \"fj_1_sj2_nbhadrons\", \"fj_1_sj2_nchadrons\", \n",
    "    \"fj_2_pt\", \"fj_2_eta\", \"fj_2_sdmass\", \"fj_2_tau21\", \"fj_2_btagcsvv2\", \"fj_2_btagjp\", \"fj_2_sfBDT\", \"fj_2_nbhadrons\", \"fj_2_nchadrons\", \"fj_2_sj1_nbhadrons\", \"fj_2_sj1_nchadrons\", \"fj_2_sj2_nbhadrons\", \"fj_2_sj2_nchadrons\", \n",
    "    \"passHTTrig\", \"ht\", \"fj_1_is_qualified\", \"fj_2_is_qualified\", \"puWeight\", \"puWeightUp\", \"puWeightDown\", \"xsecWeight\",\n",
    "    'fj_1_sj1_sv1_ntracks', 'fj_1_sj1_sv1_dxy', 'fj_1_sj1_sv1_dlensig', 'fj_1_sj2_sv1_ntracks', 'fj_1_sj2_sv1_dxy', 'fj_1_sj2_sv1_dlensig', 'fj_2_sj1_sv1_ntracks', 'fj_2_sj1_sv1_dxy', 'fj_2_sj1_sv1_dlensig', 'fj_2_sj2_sv1_ntracks', 'fj_2_sj2_sv1_dxy', 'fj_2_sj2_sv1_dlensig',\n",
    "])\n",
    "minimal_branches |= set([config['tagger']['var'].replace('fj_x', 'fj_1'), config['tagger']['var'].replace('fj_x', 'fj_2')])\n",
    "minimal_branches |= set([ ## for test only!\n",
    "#     \"fj_1_sj1_matchallmu\", \"fj_1_sj2_matchallmu\", \"fj_2_sj1_matchallmu\", \"fj_2_sj2_matchallmu\",\n",
    "#     'fj_1_btagHbb', 'fj_1_btagDeepB', 'fj_1_btagDDBvLV2', 'fj_1_ParticleNetMD_XbbVsQCD', 'fj_1_btagDDCvLV2', 'fj_1_ParticleNetMD_XccVsQCD', 'fj_1_btagDDCvBV2', 'fj_1_ParticleNetMD_Xcc', 'fj_1_ParticleNetMD_Xbb', 'fj_2_btagHbb', 'fj_2_btagDeepB', 'fj_2_btagDDBvLV2', 'fj_2_ParticleNetMD_XbbVsQCD', 'fj_2_btagDDCvLV2', 'fj_2_ParticleNetMD_XccVsQCD', 'fj_2_btagDDCvBV2', 'fj_2_ParticleNetMD_Xcc', 'fj_2_ParticleNetMD_Xbb',\n",
    "])\n",
    "\n",
    "ext_hlt_branches = {  ## extra branches depend on year\n",
    "    2016: ['HLT_PFHT125', 'HLT_PFHT200', 'HLT_PFHT250', 'HLT_PFHT300', 'HLT_PFHT350', 'HLT_PFHT400', 'HLT_PFHT475', 'HLT_PFHT600', 'HLT_PFHT650', 'HLT_PFHT800', 'HLT_PFHT900'],\n",
    "    2017: ['HLT_PFHT180', 'HLT_PFHT250', 'HLT_PFHT370', 'HLT_PFHT430', 'HLT_PFHT510', 'HLT_PFHT590', 'HLT_PFHT680', 'HLT_PFHT780', 'HLT_PFHT890', 'HLT_PFHT1050', 'HLT_PFHT350'],\n",
    "    2018: ['HLT_PFHT180', 'HLT_PFHT250', 'HLT_PFHT370', 'HLT_PFHT430', 'HLT_PFHT510', 'HLT_PFHT590', 'HLT_PFHT680', 'HLT_PFHT780', 'HLT_PFHT890', 'HLT_PFHT1050', 'HLT_PFHT350'],\n",
    "}\n",
    "minimal_branches |= set(ext_hlt_branches[year])\n",
    "minimal_branches |= set(['nPSWeight', 'PSWeight']) if year==2018 or ('optional' in config['samples'] and config['samples']['optional'].get('use_own_psweight', None)) else set()  ## extra PSWeight branches for 2018\n",
    "minimal_branches_for_data = set(minimal_branches) - set([\"fj_1_dr_H\", \"fj_1_dr_Z\", \"fj_2_dr_H\", \"fj_2_dr_Z\", 'genWeight', \"puWeight\", \"puWeightUp\", \"puWeightDown\", \"xsecWeight\", 'nPSWeight', 'PSWeight',\n",
    "                                'fj_1_nchadrons', 'fj_1_nbhadrons','fj_2_nbhadrons','fj_1_sj1_nbhadrons','fj_2_sj1_nbhadrons','fj_1_sj2_nbhadrons','fj_2_sj2_nbhadrons',\n",
    "                                'fj_2_nchadrons','fj_1_sj1_nchadrons','fj_2_sj1_nchadrons','fj_1_sj2_nchadrons','fj_2_sj2_nchadrons'])\n",
    "\n",
    "## Read into pandas DataFrame\n",
    "sample_prefix = f\"{config['samples']['sample_prefix']}_{year}\"\n",
    "_df0 = {}\n",
    "for sam in read_sample_list_map:\n",
    "    _df0[sam] = uproot3.open(f\"{sample_prefix}/{read_sample_list_map[sam]}\")['Events'].pandas.df(minimal_branches if 'mc/' in read_sample_list_map[sam] else minimal_branches_for_data, flatten=False)"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Load backup pickels"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "## Load extra variables stored during pre-processing\n",
    "backup_name = f\"{config['samples']['name']}_SF{config['year']}\"\n",
    "\n",
    "import pickle\n",
    "import awkward1 as ak\n",
    "for sam in os.listdir(f'prep/{backup_name}'):\n",
    "    if sam in read_sample_list_map:\n",
    "        for var in os.listdir(f'prep/{backup_name}/{sam}'):\n",
    "            if var.startswith('.'):\n",
    "                continue\n",
    "            if var == 'maskdict':\n",
    "                with open(f'prep/{backup_name}/{sam}/maskdict', 'rb') as f:\n",
    "                    maskdict = pickle.load(f)\n",
    "                for key in maskdict:\n",
    "                    _df0[sam]['mask_'+key] = ak.fill_none(maskdict[key], 0)\n",
    "                print('storing...', sam, 'maskdict', maskdict.keys())\n",
    "            else:\n",
    "                with open(f'prep/{backup_name}/{sam}/{var}', 'rb') as f:\n",
    "                    _df0[sam][var] = ak.fill_none(pickle.load(f), 0)\n",
    "                print('loading...', sam, var)\n",
    "    elif not sam.startswith('.') and os.path.isfile(f'prep/{backup_name}/{sam}'):\n",
    "        with open(f'prep/{backup_name}/{sam}', 'rb') as f:\n",
    "            _df0[sam] = pickle.load(f)\n",
    "        print('loading...', sam)"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "## Combine branches fj_1/2 to fj_x in pandas dataframe\n",
    "df1 = {}\n",
    "updated_key_list = list(_df0.keys())\n",
    "for sam in updated_key_list:\n",
    "    if sam in read_sample_list_map:\n",
    "        ## To concatenate event lists where either fj_1 is qualified OR fj_2 is qualified\n",
    "        fj_branches = [key.replace('fj_2', 'fj_x') for key in _df0[sam].keys() if key.startswith('fj_2')]  ## all fj_2_ branches expect fj_2_is_qualified\n",
    "        for i, i_inv in zip(['1','2'], ['2','1']):\n",
    "            df1[sam + i] = _df0[sam].query(f'mask_fj_{i}_base')  ## select events where fj_1/fj_2 is qualified\n",
    "            df1[sam + i].drop(columns=set([key.replace('fj_x', f'fj_{i_inv}') for key in fj_branches]) | set(list(_df0[sam].filter(regex='mask_*'))), inplace=True)  ## drop fj branches for the other index\n",
    "            df1[sam + i].rename(columns={key.replace('fj_x', f'fj_{i}'): key for key in fj_branches}, inplace=True)  ## change branches name from fj_1/fj_2 to a unified name fj_x\n",
    "            df1[sam + i].loc[:, 'fj_idx'] = int(i)  ## label the jet index\n",
    "            df1[sam + i].loc[:, 'is_qcd'] = True if 'qcd' in sam else False\n",
    "        df1[sam] = pd.concat([df1[sam + '1'], df1[sam + '2']])\n",
    "        if 'mc/' in read_sample_list_map[sam]:\n",
    "            df1['subst_'+sam] = df1[sam]\n",
    "        del df1[sam + '1'], df1[sam + '2']\n",
    "#         del _df0[sam]  # to release memory usage if necessary\n",
    "    else:\n",
    "        df1[sam] = _df0[sam]"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Make ROOT templates\n",
    "\n",
    "We produce the ROOT templates using the DataFrame in this step. The outputs are ROOT files with neat structure. After the further reorganization, they can be used as the Higgs Combine input to implement the fit.\n",
    "\n",
    "As a reference, we provide an example of the output files and their structure. \n",
    "E.g., for a **given fit variable**, **given tagger WP** and a **certain jet-pT bin** for **a single fit**, the output ROOT templates should include the pass and fail MC template in the B/C/L flavors, the data template, and the MC systematics for all specified shape uncertainties. The files are organized in the following structure:\n",
    "```\n",
    "─── 20210315_SF2018_AK15_qcd_ak_pnV02_HP_msv12_dxysig_log_var22binsv2  [use variable: msv12_dxysig_log, Tight WP]\n",
    "    └── Cards\n",
    "        └── pt250to350   [given pT bin]\n",
    "            ├── bdt719   [the sfBDT cut points]\n",
    "            │   ├── nominal                    [the nominal histograms]\n",
    "            │   │   ├── inputs_fail.root           [include four TH1D: flvC, flvB, flvL, data_obs]\n",
    "            │   │   └── inputs_pass.root           [..]\n",
    "            │   ├── fracBBDown                 [shape uncertainty plots]\n",
    "            │   │   ├── inputs_fail.root           [include three TH1D: flvC_fracBBDown, flvB_fracBBDown, flvL_fracBBDown]\n",
    "            │   │   └── inputs_pass.root           [..]\n",
    "            │   ├── fracBBUp\n",
    "            │   │   ├── inputs_fail.root\n",
    "            │   │   └── inputs_pass.root\n",
    "            │   ├── fracCCDown\n",
    "            │   │   ├── inputs_fail.root\n",
    "            │   │   └── inputs_pass.root\n",
    "            │   ├── fracCCUp\n",
    "            │   │   ├── inputs_fail.root\n",
    "            │   │   └── inputs_pass.root\n",
    "            │   ├── fracLightDown\n",
    "            │   │   ├── inputs_fail.root\n",
    "            │   │   └── inputs_pass.root\n",
    "            │   ├── fracLightUp\n",
    "            │   │   ├── inputs_fail.root\n",
    "            │   │   └── inputs_pass.root\n",
    "            │   ├── fitVarRwgtDown\n",
    "            │   │   ├── inputs_fail.root\n",
    "            │   │   └── inputs_pass.root\n",
    "            │   ├── fitVarRwgtUp\n",
    "            │   │   ├── inputs_fail.root\n",
    "            │   │   └── inputs_pass.root\n",
    "            │   ├── psWeightFsrDown\n",
    "            │   │   ├── inputs_fail.root\n",
    "            │   │   └── inputs_pass.root\n",
    "            │   ├── psWeightFsrUp\n",
    "            │   │   ├── inputs_fail.root\n",
    "            │   │   └── inputs_pass.root\n",
    "            │   ├── psWeightIsrDown\n",
    "            │   │   ├── inputs_fail.root\n",
    "            │   │   └── inputs_pass.root\n",
    "            │   ├── psWeightIsrUp\n",
    "            │   │   ├── inputs_fail.root\n",
    "            │   │   └── inputs_pass.root\n",
    "            │   ├── puDown\n",
    "            │   │   ├── inputs_fail.root\n",
    "            │   │   └── inputs_pass.root\n",
    "            │   ├── puUp\n",
    "            │   │   ├── inputs_fail.root\n",
    "            │   │   └── inputs_pass.root\n",
    "            │   ├── sfBDTRwgtDown\n",
    "            │   │   ├── inputs_fail.root\n",
    "            │   │   └── inputs_pass.root\n",
    "            │   └── sfBDTRwgtUp\n",
    "            │       ├── inputs_fail.root\n",
    "            │       └── inputs_pass.root\n",
    "            └── bdt752\n",
    "                ├── nominal\n",
    "                │   ├── ...\n",
    "```"
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "The template making is organized in three nested functions."
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "#### =========================================================================== Global parameters =========================================================================== ####\n",
    "g_make_template_mode = 'main'\n",
    "r\"\"\"Options:\n",
    "        main           : the main fit\n",
    "        val_pt         : the validation fit -- to use an optional MC subsitute-to-data strategy, i.e. on pT variable only\n",
    "        val_tosig_mass : the validation fit -- additionally reweight MC & data to h->cc signal jet on mass\n",
    "        val_tosig_pt   : the validation fit -- additionally reweight MC & data to h->cc signal jet on pt  \n",
    "        val_tosig_tau21: the validation fit -- additionally reweight MC & data to h->cc signal jet on tau21\n",
    "        val_crop_bin   : the validation fit -- cropping the marginal bins for fit\n",
    "\"\"\"\n",
    "\n",
    "g_outdir_prefix = f\"{config['routine_name']}_SF{config['year']}\"\n",
    "r\"\"\"Prefix for the output dir name \"\"\"\n",
    "\n",
    "g_make_unce_types = {'nominal':True, 'pu':True, 'fracBB':True, 'fracCC':True, 'fracLight':True, 'psWeightIsr':True, 'psWeightFsr':True, 'sfBDTRwgt':True, 'fitVarRwgt':True}\n",
    "r\"\"\"The uncertainty types used in the fit. Use False or remove the key to disable an certain unce type\n",
    "    Note: \"qcdSyst\" and \"qcdKdeSyst\" is not used in this verision. \"psWeightIsr\" and \"psWeightFsr\" works fine in 2018 while in 2016/17 one need to first garantee the 2018 histograms exist\n",
    "          so the unce can be transferred.\n",
    "\"\"\"\n",
    "\n",
    "g_do_fit_for_var = [1, 2, 3]\n",
    "r\"\"\" Do fit for which variable\"\"\"\n",
    "\n",
    "g_mode_bdt_runlist = 'all'\n",
    "r\"\"\"Mode of BDT list for the run. Set 'all' for all 11 BDT values, or 'central' for the central BDT value only\"\"\"\n",
    "\n",
    "g_pt_range = config['pt_range']['range']\n",
    "r\"\"\"pT range for define separate fit points\"\"\"\n",
    "\n",
    "g_tagger_range = config['tagger']['working_points']['range']\n",
    "g_tagger_var = config['tagger']['var']\n",
    "r\"\"\"Trigger info\"\"\"\n",
    "\n",
    "g_use_bflav = config['samples']['use_bflav']\n",
    "r\"\"\"Use additional B flavor MC samples to improve the statistics for the 'b' catogory\"\"\"\n",
    "\n",
    "g_bdt_mod_factor = None\n",
    "r\"\"\"Set the sfBDT selection expr to sfBDT + 0.5*exp(g_bdt_mod_factor*(tagger-1)) in the template extraction\"\"\"\n",
    "\n",
    "g_mode_psWeight_run_templ = None\n",
    "r\"\"\"Set None for the normal run. If set to 2016 or 2017, produce the 2018 templates for psWeightIsr/Fsr unce that can be migarated to 2016/2017 conditions. sfBDT cut value set under the 2016/2017 condition.\"\"\"\n",
    "\n",
    "g_dryrun = False\n",
    "r\"\"\"Launch a test process only without writing the ROOT template files\"\"\"\n",
    "\n",
    "#### ===================================================================================================================================================================================== ####\n",
    "\n",
    "## Fit info: in the format of [ (fit var, nbins/edges, xmin/None, xmax/None, (underflow, overflow), label), outputdir lambda func ]\n",
    "g_fitinfo = {\n",
    "    1: [ ##  main fit var\n",
    "        ('fj_x_mSV12_dxysig_log', [-0.8,-0.4,0,0.1,0.2,0.3,0.4,0.5,0.6,0.7,0.8,0.9,1,1.1,1.2,1.3,1.4,1.5,1.6,1.7,1.8,2.5,3.2], None, None, (True, True), 'mSV12_dxysig_log'), \n",
    "        lambda prefix, wp, bdt, pt_range, sys_name: f'results/{prefix}_{wp}_msv12_dxysig_log_var22binsv2/Cards/pt{pt_range[0]}to{pt_range[1]}/bdt{int(bdt*1000)}/{sys_name}/'\n",
    "    ],\n",
    "    2: [ ## the other var for validation\n",
    "        ('fj_x_mSV12_ptmax_log', [-0.4,0,0.1,0.2,0.3,0.4,0.5,0.6,0.7,0.8,0.9,1,1.1,1.2,1.3,1.4,1.5,1.6,1.7,1.8,2.5,3.2,3.9], None, None, (True, True), 'mSV12_ptmax_log'), \n",
    "        lambda prefix, wp, bdt, pt_range, sys_name: f'results/{prefix}_{wp}_msv12_ptmax_log_var22binsv2/Cards/pt{pt_range[0]}to{pt_range[1]}/bdt{int(bdt*1000)}/{sys_name}/'\n",
    "    ],\n",
    "    3: [ ## the other var for validation\n",
    "        ('fj_x_btagcsvv2', [0,0.1,0.15,0.2,0.25,0.3,0.35,0.4,0.45,0.5,0.55,0.6,0.65,0.7,0.75,0.8,0.85,0.9,0.95,0.98,0.99,0.995,1], None, None, (True, True), 'CSVv2'), \n",
    "        lambda prefix, wp, bdt, pt_range, sys_name: f'results/{prefix}_{wp}_csvv2_var22binsv2/Cards/pt{pt_range[0]}to{pt_range[1]}/bdt{int(bdt*1000)}/{sys_name}/'\n",
    "    ],\n",
    "    901: [ ## crop the marginal bins for the main var as a validation\n",
    "        ('fj_x_mSV12_dxysig_log', [0,0.1,0.2,0.3,0.4,0.5,0.6,0.7,0.8,0.9,1,1.1,1.2,1.3,1.4,1.5,1.6,1.7,1.8], None, None, (False, False), 'mSV12_dxysig_log'), \n",
    "        lambda prefix, wp, bdt, pt_range, sys_name: f'results/{prefix}_{wp}_msv12_dxysig_log_var22binscrop/Cards/pt{pt_range[0]}to{pt_range[1]}/bdt{int(bdt*1000)}/{sys_name}/'\n",
    "    ],\n",
    "}\n",
    "\n",
    "## Necessary KDE parameters used in qcdKdeSyst unce\n",
    "g_custom_kde_bw = {'fj_x_btagcsvv2':15, 'mSV12_ptmax_log':4, 'mSV12_dxysig_log':4}\n",
    "g_custom_kde_binmask = {'fj_x_btagcsvv2':[0], 'mSV12_ptmax_log':[-0.4,1.8,2.5,3.2], 'mSV12_dxysig_log':[-0.8,-0.4,1.8,2.5]}\n",
    "\n",
    "## Some other global vars\n",
    "g_do_sfBDT_points = None\n",
    "g_outdir_prefix_used = None\n",
    "g_hist_qcdsyst = {}\n",
    "g_wgtstr_dm_sys_fac = {}\n",
    "g_hist_fitvar_rwgt = {}\n",
    "\n",
    "def check_consistency(): ## Consistency check for gloal params\n",
    "    assert g_make_template_mode in ['main', 'val_pt', 'val_tosig_mass', 'val_tosig_pt', 'val_tosig_tau21', 'val_vary_sfbdt', 'val_crop_bin'], \\\n",
    "        'Specified mode cannot be recognized.'\n",
    "    \n",
    "    global g_do_fit_for_var\n",
    "    if g_make_template_mode in ['val_pt', 'val_tosig_mass', 'val_tosig_pt', 'val_tosig_tau21'] and g_do_fit_for_var != [1]:\n",
    "        print('Warning: for validation fit, set the fit information to the main variable (1) only')\n",
    "        g_do_fit_for_var = [1]\n",
    "    if g_make_template_mode == 'val_crop_bin' and g_do_fit_for_var != [901]:\n",
    "        print('Warning: for validation fit on cropping the marginal bins, set the fit information to the cropped main variable (901) only')\n",
    "        g_do_fit_for_var = [901]\n",
    "    \n",
    "    global g_mode_bdt_runlist\n",
    "    if g_make_template_mode.startswith('val_') and g_mode_bdt_runlist != 'central':\n",
    "        print('Warning: for validation fit, set the BDT run list to central')\n",
    "        g_mode_bdt_runlist = 'central'\n",
    "    \n",
    "    global g_do_sfBDT_points\n",
    "    if g_mode_bdt_runlist == 'all':\n",
    "        g_do_sfBDT_points = df1[f\"bdt_seq_{config['pt_range']['name']}__{config['main_analysis_tree']['name']}\"]\n",
    "    elif g_mode_bdt_runlist == 'central':\n",
    "        _points = df1[f\"bdt_seq_{config['pt_range']['name']}__{config['main_analysis_tree']['name']}\"]\n",
    "        g_do_sfBDT_points = {k:[_points[k][int((len(_points[k])-1)/2)]] for k in _points}\n",
    "    elif g_mode_bdt_runlist != 'manual':\n",
    "        raise RuntimeError('Specified mode for BDT runlist cannot be recognized.')\n",
    "    \n",
    "    global g_outdir_prefix_used\n",
    "    g_outdir_prefix_used = g_outdir_prefix + '_' + config['tagger']['working_points']['name']\n",
    "    if g_make_template_mode.startswith('val_'):\n",
    "        g_outdir_prefix_used += '_-' + g_make_template_mode + '-'\n",
    "    if g_bdt_mod_factor is not None:\n",
    "        g_outdir_prefix_used = 'bdtmod/' + g_outdir_prefix_used\n",
    "    \n",
    "    if g_mode_psWeight_run_templ is not None:\n",
    "        assert year==2018, 'g_mode_psWeight_run_templ only set for year 2016/2017'\n",
    "        assert int(g_mode_psWeight_run_templ) in [2016, 2017], 'g_mode_psWeight_run_templ can only be 2016 or 2017'\n",
    "        import pickle\n",
    "        if g_mode_bdt_runlist != 'manual':\n",
    "            with open(f\"prep/{config['samples']['name']}_SF{g_mode_psWeight_run_templ}/bdt_seq_{config['pt_range']['name']}__{config['main_analysis_tree']['name']}\", 'rb') as f:\n",
    "                g_do_sfBDT_points = pickle.load(f)\n",
    "        g_outdir_prefix_used += f\"_psWeight{g_mode_psWeight_run_templ}\"\n",
    "        g_make_unce_types = {'nominal':True, 'psWeightIsr':True, 'psWeightFsr':True}\n",
    "\n",
    "def launch_maker():\n",
    "    r\"\"\"Depth 0: Main function to launch the fit given the global parameters\n",
    "    \"\"\"\n",
    "    check_consistency()\n",
    "    \n",
    "    print('Launch variablel list:', g_do_fit_for_var)\n",
    "    for _ifit in g_do_fit_for_var:\n",
    "        for _wp in g_tagger_range:\n",
    "            \n",
    "            ## Get fit info and output lambda func\n",
    "            fitinfo, outdir_func = g_fitinfo[_ifit]\n",
    "\n",
    "            ## The default args in the main fit\n",
    "            args = {\n",
    "                'wgtstr_dm': f'genWeight*xsecWeight*puWeight*{lumi[year]}*fj_x_htwgt', 'wgtstr_dm_data': None,\n",
    "                'sl_dm': ['subst_'+s if s!='jetht-noht' else s for s in read_sample_list_map if s not in ['qcd-herwig-noht', 'qcd-mg-bflav-noht']], # default is ['subst_qcd-mg-noht', 'subst_top-noht', 'subst_v-qq-noht', 'jetht-noht']\n",
    "                'sl_dm_herwig': ['subst_'+s if s!='jetht-noht' else s for s in read_sample_list_map if s not in ['qcd-mg-noht', 'qcd-mg-bflav-noht']], # default is ['subst_qcd-herwig-noht', 'subst_top-noht', 'subst_v-qq-noht', 'jetht-noht']\n",
    "                'config_dm': {\n",
    "                    'data':  '',\n",
    "                    'flvB':  'fj_x_nbhadrons>=1',\n",
    "                    'flvC':  'fj_x_nbhadrons==0 & fj_x_nchadrons>=1',\n",
    "                    'flvL':  'fj_x_nbhadrons==0 & fj_x_nchadrons==0',\n",
    "                },\n",
    "                'categories_dm': ['flvL', 'flvB', 'flvC', 'data'],\n",
    "                'catMap': {\n",
    "                    'pass': f'{g_tagger_var}>{g_tagger_range[_wp][0]:.3f} & {g_tagger_var}<={g_tagger_range[_wp][1]:.3f}',\n",
    "                    'fail': f'{g_tagger_var}<={g_tagger_range[_wp][0]:.3f} | {g_tagger_var}>{g_tagger_range[_wp][1]:.3f}',\n",
    "                },\n",
    "                'use_bflav': g_use_bflav, 'args_bflav': {\n",
    "                    'sl_dm_bflav': ['subst_qcd-mg-bflav-noht'], 'sl_dm_bflav_orig': ['subst_qcd-mg-noht'],\n",
    "                    'wgtstropt_bflav': lambda s: s.replace('fj_x_htwgt', '(fj_x_htwgt*fj_x_bflav_htwgt)'),\n",
    "                },\n",
    "            }\n",
    "            ## Modify args according to specified global param\n",
    "            if g_make_template_mode == 'val_pt':\n",
    "                args['wgtstr_dm'], args['wgtstr_dm_data'] = f'genWeight*xsecWeight*puWeight*{lumi[year]}*fj_x_ad_ptwgt', None\n",
    "            elif g_make_template_mode == 'val_tosig_mass':\n",
    "                args['wgtstr_dm'], args['wgtstr_dm_data'] = f'genWeight*xsecWeight*puWeight*{lumi[year]}*fj_x_htwgt*fj_x_massdatamcwgt', 'fj_x_massdatamcwgt'\n",
    "            elif g_make_template_mode == 'val_tosig_pt':\n",
    "                args['wgtstr_dm'], args['wgtstr_dm_data'] = f'genWeight*xsecWeight*puWeight*{lumi[year]}*fj_x_htwgt*fj_x_ptdatamcwgt', 'fj_x_ptdatamcwgt'\n",
    "            elif g_make_template_mode == 'val_tosig_tau21':\n",
    "                args['wgtstr_dm'], args['wgtstr_dm_data'] = f'genWeight*xsecWeight*puWeight*{lumi[year]}*fj_x_htwgt*fj_x_tau21datamcwgt', 'fj_x_tau21datamcwgt'\n",
    "\n",
    "            wrapperPt(df1, fitinfo, lambda bdt, pt_range, sys_name: outdir_func(g_outdir_prefix_used, _wp, bdt, pt_range, sys_name), args)"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "def wrapperPt(df2, fitinfo, outdir_func, args):\n",
    "    r\"\"\"Depth 1: Process the pT cut and wrap all other following steps\n",
    "    \"\"\"\n",
    "    print('Launch pT range:', g_pt_range)\n",
    "    for pt_range in g_pt_range:\n",
    "        pt_range = tuple(pt_range)\n",
    "        print ('pt range:', pt_range)\n",
    "        \n",
    "        ## df2->df2a: apply the pT cut (to speed up) (plus additional selection, if applied)\n",
    "        df2a = {}\n",
    "        for sam in ['subst_'+s if s!='jetht-noht' else s for s in read_sample_list_map]:\n",
    "            df2a[sam] = df2[sam].query(f'fj_x_pt>={pt_range[0]} & fj_x_pt<{pt_range[1]}')\n",
    "            if 'optional' in config['samples'] and 'additional_selection_in_making_template' in config['samples']['optional']:\n",
    "                df2a[sam] = df2a[sam].query(config['samples']['optional']['additional_selection_in_making_template'].replace('np.','').replace('ak.',''))\n",
    "        \n",
    "        sfBDT_list = g_do_sfBDT_points[pt_range]\n",
    "        if isinstance(sfBDT_list, dict):\n",
    "            sfBDT_list = sfBDT_list.values()\n",
    "        bdt_expr = 'fj_x_sfBDT'\n",
    "        if g_bdt_mod_factor is not None:\n",
    "            bdt_expr = f'fj_x_sfBDT + 0.5*exp({g_bdt_mod_factor}*({g_tagger_var}-1))'\n",
    "        for sfBDT_val in sfBDT_list:\n",
    "            print(' sfBDT cut at:', sfBDT_val)\n",
    "            \n",
    "            ## df2a->df3: apply the corresponding bdt cut\n",
    "            df3 = {}\n",
    "            for sam in ['subst_'+s if s!='jetht-noht' else s for s in read_sample_list_map]:\n",
    "                df3[sam] = df2a[sam].query(f'{bdt_expr}>{sfBDT_val}')\n",
    "\n",
    "            makeTemplatesWrapper(df3, fitinfo, lambda sys_name: outdir_func(sfBDT_val, pt_range, sys_name), sfBDT_val, args)"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "def makeTemplatesWrapper(df3, fitinfo, outdir_func, sfBDT_val, args):\n",
    "    r\"\"\"Depth 2: Specify which template (nominal or any shape uncertainty) to make in this step\n",
    "    \"\"\"\n",
    "    global g_wgtstr_dm_sys_fac, g_hist_qcdsyst, g_hist_fitvar_rwgt\n",
    "    g_wgtstr_dm_sys_fac, g_hist_qcdsyst = {}, {} ## clear\n",
    "    g_hist_fitvar_rwgt = {}\n",
    "    \n",
    "    wgtstr_dm = args['wgtstr_dm']\n",
    "    if 'nominal' in g_make_unce_types.keys() and g_make_unce_types['nominal']:\n",
    "        sys_name = 'nominal'; wgtstr_dm_sys = wgtstr_dm; makeTemplates(df3, fitinfo, outdir_func(sys_name), sys_name, wgtstr_dm_sys, args)\n",
    "    \n",
    "    ## Below we extract hists for all unce type\n",
    "    if 'pu' in g_make_unce_types.keys() and g_make_unce_types['pu']: \n",
    "        sys_name = 'puUp'; wgtstr_dm_sys = wgtstr_dm.replace('puWeight','puWeightUp'); makeTemplates(df3, fitinfo, outdir_func(sys_name), sys_name, wgtstr_dm_sys, args)\n",
    "        sys_name = 'puDown'; wgtstr_dm_sys = wgtstr_dm.replace('puWeight','puWeightDown'); makeTemplates(df3, fitinfo, outdir_func(sys_name), sys_name, wgtstr_dm_sys, args)\n",
    "\n",
    "    if 'fracBB' in g_make_unce_types.keys() and g_make_unce_types['fracBB']: \n",
    "        sys_name = \"fracBBUp\"; wgtstr_dm_sys = wgtstr_dm+'*(1.2*(fj_x_nbhadrons>1) + 1.0*(fj_x_nbhadrons<=1))'; makeTemplates(df3, fitinfo, outdir_func(sys_name), sys_name, wgtstr_dm_sys, args)\n",
    "        sys_name = \"fracBBDown\"; wgtstr_dm_sys = wgtstr_dm+'*(0.8*(fj_x_nbhadrons>1) + 1.0*(fj_x_nbhadrons<=1))'; makeTemplates(df3, fitinfo, outdir_func(sys_name), sys_name, wgtstr_dm_sys, args)\n",
    "    if 'fracCC' in g_make_unce_types.keys() and g_make_unce_types['fracCC']: \n",
    "        sys_name = \"fracCCUp\"; wgtstr_dm_sys = wgtstr_dm+'*(1.2*(fj_x_nbhadrons==0 & fj_x_nchadrons>1) + 1.0*(not(fj_x_nbhadrons==0 & fj_x_nchadrons>1)))'; makeTemplates(df3, fitinfo, outdir_func(sys_name), sys_name, wgtstr_dm_sys, args)\n",
    "        sys_name = \"fracCCDown\"; wgtstr_dm_sys = wgtstr_dm+'*(0.8*(fj_x_nbhadrons==0 & fj_x_nchadrons>1) + 1.0*(not(fj_x_nbhadrons==0 & fj_x_nchadrons>1)))'; makeTemplates(df3, fitinfo, outdir_func(sys_name), sys_name, wgtstr_dm_sys, args)\n",
    "    if 'fracLight' in g_make_unce_types.keys() and g_make_unce_types['fracLight']: \n",
    "        sys_name = \"fracLightUp\"; wgtstr_dm_sys = wgtstr_dm+'*(1.2*(fj_x_nbhadrons==0 & fj_x_nchadrons==0) + 1.0*(not(fj_x_nbhadrons==0 & fj_x_nchadrons==0)))'; makeTemplates(df3, fitinfo, outdir_func(sys_name), sys_name, wgtstr_dm_sys, args)\n",
    "        sys_name = \"fracLightDown\"; wgtstr_dm_sys = wgtstr_dm+'*(0.8*(fj_x_nbhadrons==0 & fj_x_nchadrons==0) + 1.0*(not(fj_x_nbhadrons==0 & fj_x_nchadrons==0)))'; makeTemplates(df3, fitinfo, outdir_func(sys_name), sys_name, wgtstr_dm_sys, args)\n",
    "\n",
    "    ## Below unce is not as easily extracted as above by specifying a different weight string. They may need *special treatment* implemented in the depth-3 function\n",
    "    if 'qcdSyst' in g_make_unce_types.keys() and g_make_unce_types['qcdSyst']: \n",
    "        sys_name = \"qcdSystUp\"; wgtstr_dm_sys = wgtstr_dm; makeTemplates(df3, fitinfo, outdir_func(sys_name), sys_name, wgtstr_dm_sys, args)\n",
    "        sys_name = \"qcdSystDown\"; wgtstr_dm_sys = wgtstr_dm; makeTemplates(df3, fitinfo, outdir_func(sys_name), sys_name, wgtstr_dm_sys, args)\n",
    "    if 'qcdKdeSyst' in g_make_unce_types.keys() and g_make_unce_types['qcdKdeSyst']: \n",
    "        sys_name = \"qcdKdeSystUp\"; wgtstr_dm_sys = wgtstr_dm; makeTemplates(df3, fitinfo, outdir_func(sys_name), sys_name, wgtstr_dm_sys, args)\n",
    "        sys_name = \"qcdKdeSystDown\"; wgtstr_dm_sys = wgtstr_dm; makeTemplates(df3, fitinfo, outdir_func(sys_name), sys_name, wgtstr_dm_sys, args)\n",
    "    if 'psWeightIsr' in g_make_unce_types.keys() and g_make_unce_types['psWeightIsr']: \n",
    "        sys_name = \"psWeightIsrUp\"; wgtstr_dm_sys = wgtstr_dm; makeTemplates(df3, fitinfo, outdir_func(sys_name), sys_name, wgtstr_dm_sys, args)\n",
    "        sys_name = \"psWeightIsrDown\"; wgtstr_dm_sys = wgtstr_dm; makeTemplates(df3, fitinfo, outdir_func(sys_name), sys_name, wgtstr_dm_sys, args)\n",
    "    if 'psWeightFsr' in g_make_unce_types.keys() and g_make_unce_types['psWeightFsr']: \n",
    "        sys_name = \"psWeightFsrUp\"; wgtstr_dm_sys = wgtstr_dm; makeTemplates(df3, fitinfo, outdir_func(sys_name), sys_name, wgtstr_dm_sys, args)\n",
    "        sys_name = \"psWeightFsrDown\"; wgtstr_dm_sys = wgtstr_dm; makeTemplates(df3, fitinfo, outdir_func(sys_name), sys_name, wgtstr_dm_sys, args)\n",
    "\n",
    "    if 'sfBDTRwgt' in g_make_unce_types.keys() and g_make_unce_types['sfBDTRwgt']: \n",
    "        sys_name = 'sfBDTRwgtUp'; wgtstr_dm_sys = wgtstr_dm;'''factors decided by special_wgtstr argument'''; makeTemplates(df3, fitinfo, outdir_func(sys_name), sys_name, wgtstr_dm_sys, args, special_wgtstr='fj_x_sfbdtwgt_g50')\n",
    "        sys_name = 'sfBDTRwgtDown'; wgtstr_dm_sys = wgtstr_dm;'''factors decided by special_wgtstr argument'''; makeTemplates(df3, fitinfo, outdir_func(sys_name), sys_name, wgtstr_dm_sys, args, special_wgtstr='fj_x_sfbdtwgt_g50')\n",
    "    \n",
    "    if 'fitVarRwgt' in g_make_unce_types.keys() and g_make_unce_types['fitVarRwgt']: \n",
    "        sys_name = 'fitVarRwgtUp'; wgtstr_dm_sys = wgtstr_dm; makeTemplates(df3, fitinfo, outdir_func(sys_name), sys_name, wgtstr_dm_sys, args)\n",
    "        sys_name = 'fitVarRwgtDown'; wgtstr_dm_sys = wgtstr_dm; makeTemplates(df3, fitinfo, outdir_func(sys_name), sys_name, wgtstr_dm_sys, args)\n",
    "    "
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "def makeTemplates(df3, fitinfo, outputdir, sys_name, wgtstr_dm_sys, args, special_wgtstr=None):\n",
    "    r\"\"\"Depth 3: The very base implementation that apply the final pass/fail cut and make the template\n",
    "    \"\"\"\n",
    "    \n",
    "    wgtstr_dm, wgtstr_dm_data, sl_dm, sl_dm_herwig, config_dm, categories_dm, catMap = args['wgtstr_dm'], args['wgtstr_dm_data'], args['sl_dm'], args['sl_dm_herwig'], args['config_dm'], args['categories_dm'], args['catMap']\n",
    "    \n",
    "    ## Create the output root file\n",
    "    if not os.path.exists(outputdir) and not g_dryrun:\n",
    "        os.makedirs(outputdir)\n",
    "\n",
    "    import ROOT, array  ## use ROOT to write file...\n",
    "    vname, nbin, xmin, xmax, (underflow, overflow), vlabel = fitinfo\n",
    "    ## Tranfer the {nbin, xmin, xmax} set to the real bin edge if necessary\n",
    "    if not isinstance(nbin, int):\n",
    "        edges = nbin\n",
    "        nbin = len(edges)-1 # reset nbin to \"real\" nbin\n",
    "        edges_inroot = (len(edges)-1, array.array('f', edges))\n",
    "    else:\n",
    "        edges = np.linspace(xmin, xmax, nbin+1)\n",
    "        edges_inroot = (nbin, xmin, xmax)\n",
    "\n",
    "    ## Impose the overall factor between MC and data\n",
    "    def extract_factor_overal(_sl, _wgtstr):\n",
    "        return np.round(df3[_sl[-1]].shape[0] * 1. / sum([df3[sam].eval(_wgtstr).sum() for sam in _sl[:-1]]), 4)\n",
    "    \n",
    "    if special_wgtstr is None: ## no special weight string provided -> use the nominal one\n",
    "        if any([_sys in sys_name for _sys in ['qcdSyst','qcdKdeSyst']]): # note that qcd syst uses the setting of the herwig sample\n",
    "            fac_overal = g_wgtstr_dm_sys_fac['qcdSystUp'] if 'qcdSystUp' in g_wgtstr_dm_sys_fac else \\\n",
    "                         g_wgtstr_dm_sys_fac['qcdKdeSystUp'] if 'qcdKdeSystUp' in g_wgtstr_dm_sys_fac else None\n",
    "            if fac_overal is None:\n",
    "                fac_overal = extract_factor_overal(sl_dm_herwig, wgtstr_dm.replace('htwgt','htwgt_herwig'))\n",
    "        else:  # nominal case\n",
    "            fac_overal = g_wgtstr_dm_sys_fac['nominal'] if 'nominal' in g_wgtstr_dm_sys_fac else None\n",
    "            if fac_overal is None:\n",
    "                fac_overal = extract_factor_overal(sl_dm, wgtstr_dm)\n",
    "        # equip the weight factor\n",
    "        g_wgtstr_dm_sys_fac[sys_name] = fac_overal\n",
    "        wgtstr_dm_sys = wgtstr_dm_sys+f'*{fac_overal}'\n",
    "\n",
    "    else: ## special weight string specified\n",
    "        if sys_name.endswith('Up'):\n",
    "            fac_overal = extract_factor_overal(sl_dm, wgtstr_dm+f'*{special_wgtstr}')\n",
    "            # equip the weight factor\n",
    "            g_wgtstr_dm_sys_fac[sys_name] = fac_overal\n",
    "            wgtstr_dm_sys = wgtstr_dm+f'*{special_wgtstr}*{fac_overal}'\n",
    "        else:\n",
    "            wgtstr_dm_sys = wgtstr_dm+f\"*(2*{g_wgtstr_dm_sys_fac['nominal']}-{special_wgtstr}*{g_wgtstr_dm_sys_fac[sys_name.replace('Down','Up')]})\"\n",
    "\n",
    "    print (fitinfo, outputdir, sys_name, wgtstr_dm_sys)\n",
    "    \n",
    "    ## Preprocess for fitVarRwgt\n",
    "    if sys_name == 'fitVarRwgtUp':\n",
    "        _df_mc = pd.concat([df3[sam] for sam in sl_dm[:-1]])\n",
    "        _df_data = df3[sl_dm[-1]]\n",
    "        _h_data = get_hist(_df_data[vname].values, bins=edges, weights=np.ones(_df_data.shape[0]) if wgtstr_dm_data==None else _df_data.eval(wgtstr_dm_data).values, underflow=underflow, overflow=overflow).view(flow=True)\n",
    "        _h_mc = get_hist(_df_mc[vname].values, bins=edges, weights=_df_mc.eval(wgtstr_dm_sys).values, underflow=underflow, overflow=overflow).view(flow=True)\n",
    "        g_hist_fitvar_rwgt[sys_name] = _h_data.value / _h_mc.value\n",
    "    \n",
    "    ## Loop over pass and fail region\n",
    "    for b in ['pass', 'fail']:\n",
    "        try:\n",
    "            if not g_dryrun:\n",
    "                fw = ROOT.TFile(outputdir+f'inputs_{b}.root', 'recreate')\n",
    "            \n",
    "            hv, hist = {}, {}\n",
    "            hname_suf = '_'+sys_name if sys_name!='nominal' else ''  ## suffix to the hist name (the Higgs Combine syntax)\n",
    "            print (' -- ', catMap[b])\n",
    "            \n",
    "            ## MC and data dataframe after applying the final selection\n",
    "            df_mc = pd.concat([df3[sam].query(catMap[b]) for sam in sl_dm[:-1]])\n",
    "            df_data = df3[sl_dm[-1]].query(catMap[b])\n",
    "            \n",
    "            ## Preprocessing for herwig related dataframe if we mean to calculate qcdSyst / qcdKdeSyst unce in this iteration\n",
    "            if 'qcdSyst' in sys_name or 'qcdKdeSyst' in sys_name:\n",
    "                df_mc_herwig = pd.concat([df3[sam].query(catMap[b]) for sam in sl_dm_herwig[:-1]])\n",
    "\n",
    "            # Loop over categories: flvC/flvB/flvL/data\n",
    "            for cat in config_dm:\n",
    "                ## hv[] holds the boosted-histogram type derived from the dataframe, hist[] holds the TH1D type to be stored in ROOT\n",
    "                if cat=='data' and sys_name == 'nominal':\n",
    "                    ## Get the data hist\n",
    "                    hv['data'] = get_hist(df_data[vname].values, bins=edges, weights=np.ones(df_data.shape[0]) if wgtstr_dm_data==None else df_data.eval(wgtstr_dm_data).values, underflow=underflow, overflow=overflow).view(flow=True)\n",
    "                    # Initialize the TH1D hist\n",
    "                    hist['data'] = ROOT.TH1D('data_obs', 'data_obs;'+vname, *edges_inroot) \n",
    "                if cat!='data':\n",
    "                    df_mc_tmp = df_mc.query(config_dm[cat]) ## category selection based on flavor\n",
    "                    ## Get the MC hist for certain flavor\n",
    "                    hv[cat] = get_hist(df_mc_tmp[vname].values, bins=edges, weights=df_mc_tmp.eval(wgtstr_dm_sys).values, underflow=underflow, overflow=overflow).view(flow=True)\n",
    "                    # Initialize the TH1D hist\n",
    "                    hist[cat] = ROOT.TH1D(cat+hname_suf, cat+hname_suf+';'+vname, *edges_inroot) # init TH1 hist\n",
    "                    hist[cat].Sumw2()\n",
    "            \n",
    "                    ## For qcdSyst / qcdKdeSyst unce that is actually related to Herwig, hv[cat] is dummy here, \n",
    "                    ## and we mean to obtain hv[cat+'_herwig.value'] that will be later filled into hist[cat]\n",
    "                    if sys_name=='qcdSystUp':\n",
    "                        ## Get the Herwig fit for certain flavor\n",
    "                        df_mc_herwig_tmp = df_mc_herwig.query(config_dm[cat]) ## cat selection\n",
    "                        wgtstr_dm_sys_herwig = wgtstr_dm_sys.replace('htwgt','htwgt_herwig').replace('sfbdtwgt_g50','sfbdtwgt_g50_herwig').replace('ad_ptwgt','ad_ptwgt_herwig').replace('datamcwgt','datamcwgt_herwig')\n",
    "                        hv[cat+'_herwig.value'] = get_hist(df_mc_herwig_tmp[vname].values, bins=edges, \n",
    "                                                     weights=df_mc_herwig_tmp.eval(wgtstr_dm_sys_herwig).values, \n",
    "                                                     underflow=underflow, overflow=overflow).view(flow=True).value\n",
    "                        ## Store the histogram into global var so we can recycle the same hist in the \"Down\" routine\n",
    "                        g_hist_qcdsyst[(sys_name, b, cat)] = hv[cat+'_herwig.value']\n",
    "                    \n",
    "                    ## Extract the KDE shape directly from herwig shape\n",
    "                    if sys_name=='qcdKdeSystUp':\n",
    "                        df_mc_herwig_tmp = df_mc_herwig.query(config_dm[cat])\n",
    "                        wgtstr_dm_sys_herwig = wgtstr_dm_sys.replace('htwgt','htwgt_herwig').replace('sfbdtwgt_g50','sfbdtwgt_g50_herwig').replace('ad_ptwgt','ad_ptwgt_herwig').replace('datamcwgt','datamcwgt_herwig')\n",
    "                        hv_herwig_orig_value = get_hist(df_mc_herwig_tmp[vname].values, bins=edges, \n",
    "                                                     weights=df_mc_herwig_tmp.eval(wgtstr_dm_sys_herwig).values, \n",
    "                                                     underflow=underflow, overflow=overflow).view(flow=True).value\n",
    "                        \n",
    "                        ## Calculate KDE shape, apply two times so that we specify a finer KDE bindwidth based on the first result\n",
    "                        from scipy.stats import gaussian_kde\n",
    "                        kde = gaussian_kde(df_mc_herwig_tmp[vname].values, weights=np.clip(df_mc_herwig_tmp.eval(wgtstr_dm_sys_herwig).values, 0, +np.inf))\n",
    "                        kde = gaussian_kde(df_mc_herwig_tmp[vname].values, weights=np.clip(df_mc_herwig_tmp.eval(wgtstr_dm_sys_herwig).values, 0, +np.inf), bw_method=kde.factor/g_custom_kde_bw[vname])\n",
    "                        kde_int = np.zeros([nbin, 2])\n",
    "                        \n",
    "                        ## Integrate the KDE function to obtain KDE histogram\n",
    "                        for i, (low, high) in enumerate(zip(edges[:-1], edges[1:])):\n",
    "                            if low in g_custom_kde_binmask[vname]:\n",
    "                                continue\n",
    "                            kde_int[i] = [kde.integrate_box_1d(low, high), hv_herwig_orig_value[i]]\n",
    "                        # print('rescale kde sum to original herwig sum: ', kde_int[:,1].sum() / kde_int[:,0].sum())\n",
    "                        kde_int[:,0] *= kde_int[:,1].sum() / kde_int[:,0].sum()\n",
    "                        \n",
    "                        ## Fill with original madgraph hist if we plan to mask the bin for KDE. \n",
    "                        ## This is based on the fact that KDE cannot model the hist well in the marginal bins\n",
    "                        hv[cat+'_herwig.value'] = np.array([kde_int[i][0] if kde_int[i][0]!=0 else hv[cat].value[i] for i in range(nbin)])\n",
    "                        \n",
    "                        ## Store the histogram into global var so we can recycle the same hist in the \"Down\" routine\n",
    "                        g_hist_qcdsyst[(sys_name, b, cat)] = hv[cat+'_herwig.value']\n",
    "            \n",
    "                    ## Extract the PSWeight histogram\n",
    "                    if 'psWeight' in sys_name:\n",
    "                        if year==2018:  ## for 2018, calculate the hist by PSWeight vars \n",
    "                            ps_idx = {'psWeightIsrUp':2, 'psWeightIsrDown':0, 'psWeightFsrUp':3, 'psWeightFsrDown':1}\n",
    "                            hv[cat] = get_hist(df_mc_tmp[vname].values, bins=edges, weights=df_mc_tmp.eval(wgtstr_dm_sys).values*df_mc_tmp['PSWeight'].map(lambda v: v[ps_idx[sys_name]]).values, underflow=underflow, overflow=overflow).view(flow=True)\n",
    "                        else:  ## for 2016/17 extract the PSWeight hist based on 2018 result (transfer the ratio for PSWeight/nominal)\n",
    "                            import re\n",
    "                            outputdir_ps_18 = re.sub('^(.+)_SF201[6-8]_%s_(.*)$' % config['tagger']['working_points']['name'], f'\\g<1>_SF2018_%s_psWeight{year}_\\g<2>' % config['tagger']['working_points']['name'], outputdir)\n",
    "                            hv_nom_18 = uproot3.open(outputdir_ps_18.replace(sys_name, 'nominal')+f'inputs_{b}.root')[cat]\n",
    "                            hv_ps_18 = uproot3.open(outputdir_ps_18+f'inputs_{b}.root')[cat+'_'+sys_name]\n",
    "                            hv[cat].value *= hv_ps_18.values / hv_nom_18.values\n",
    "                        # print (hv[cat].value)\n",
    "                    \n",
    "                    ## Extract the sfBDTFloAround histogram.\n",
    "                    ## Method: to utilize the nominal hist for sfbdt>0.95 or 0.85 and migrate the MC-to-data confidence level in the 0.90 case\n",
    "                    if 'sfBDTFloAround' in sys_name:\n",
    "                        from scipy.stats import chi2\n",
    "                        hv_data = uproot3.open(outputdir.replace(sys_name, 'nominal')+f'inputs_{b}.root')['data_obs'].values  ## nominal data hist for 0.90\n",
    "                        _bdtname = '95' if 'Up' in sys_name else '85'\n",
    "                        fr = uproot3.open(outputdir.replace(sys_name, 'nominal').replace(f'/bdt{int(g_sfBDT_val_list[-1]*1000)}/',f'/bdt{_bdtname}0/')+f'inputs_{b}.root')\n",
    "                        fr_data, fr_mc = fr['data_obs'].values, fr['flvC'].values+fr['flvB'].values+fr['flvL'].values  ## nominal data & MC hist for 0.95 or 0.85 (depends on Up or Down)\n",
    "                        \n",
    "                        ## For each bins, migrate the confidence level of MC yield F0 given data yield D0 to the target data yield D => F\n",
    "                        hv_mc = []\n",
    "                        for D, D0, F0 in zip(hv_data, fr_data, fr_mc):\n",
    "                            ## The precise calculation\n",
    "                            F = 0.5*chi2.ppf(chi2.cdf(2*F0, 2*D0+2), 2*D+2) if F0>D0 else 0.5*chi2.ppf(chi2.cdf(2*F0, 2*D0), 2*D)\n",
    "                            if F == np.inf: ## in case the formula results in inf (may occur if F0 >> D0)\n",
    "                                assert F0 > D0\n",
    "                                sigD0 = 0.5 * chi2.ppf(1-(1-0.682689492)/2, 2*D0+2) - D0\n",
    "                                sigD = 0.5 * chi2.ppf(1-(1-0.682689492)/2, 2*D+2) - D\n",
    "                                F = D + sigD/sigD0*(F0-D0)\n",
    "                            hv_mc.append(F)\n",
    "                        \n",
    "                        ## Obtain flavor template based on the flavor proportion in 0.95 or 0.85 region\n",
    "                        hv[cat].value = np.nan_to_num(hv_mc * fr[cat].values / fr_mc, nan=0)\n",
    "                    \n",
    "                    ## Modify hv[cat] based on extracted pass+fail histogram\n",
    "                    if 'fitVarRwgt' in sys_name:\n",
    "                        if sys_name == 'fitVarRwgtUp':\n",
    "                            hv[cat].value = hv[cat].value * g_hist_fitvar_rwgt['fitVarRwgtUp']\n",
    "                        else:\n",
    "                            hv[cat].value = 2 * hv[cat].value - hv[cat].value * g_hist_fitvar_rwgt['fitVarRwgtUp']\n",
    "                    \n",
    "                    ## Use bflav qcd samples to stitch the final bflav template\n",
    "                    if 'use_bflav' in args and args['use_bflav'] and cat == 'flvB' and not all([s in sys_name for s in ['qcd','Syst']]):\n",
    "                        # print('---', hv[cat])\n",
    "                        ## Get the MC hist from the new b-enriched sample\n",
    "                        df_mc_bflav = pd.concat([df3[sam].query(f'({catMap[b]}) & ({config_dm[cat]})') for sam in args['args_bflav']['sl_dm_bflav']])\n",
    "                        hv_bflav = get_hist(df_mc_bflav[vname].values, bins=edges, weights=df_mc_bflav.eval(args['args_bflav']['wgtstropt_bflav'](wgtstr_dm_sys)).values, underflow=underflow, overflow=overflow).view(flow=True)\n",
    "                        df_mc_bflav_og = pd.concat([df3[sam].query(f'({catMap[b]}) & ({config_dm[cat]})') for sam in args['args_bflav']['sl_dm_bflav_orig']])\n",
    "                        hv_bflav_og = get_hist(df_mc_bflav_og[vname].values, bins=edges, weights=df_mc_bflav_og.eval(wgtstr_dm_sys).values, underflow=underflow, overflow=overflow).view(flow=True)\n",
    "                        ## Combine histogram\n",
    "                        hv_bflav_og.variance[hv_bflav_og.variance==0] = 1e20\n",
    "                        hv_bflav.variance[hv_bflav.variance==0] = 1e20\n",
    "                        hv_bflav_comb = hv[cat].copy()\n",
    "                        hv_bflav_comb.value = (hv_bflav_og.value*(1/hv_bflav_og.variance) + hv_bflav.value*(1/hv_bflav.variance)) / (1/hv_bflav_og.variance + 1/hv_bflav.variance)\n",
    "                        hv_bflav_comb.variance = 1 / (1/hv_bflav_og.variance + 1/hv_bflav.variance)\n",
    "                        ## Further combine with the non no-QCD contribution\n",
    "                        hv_bflav_nonsubst = hv[cat].copy() # histogram constitution not to be combined (i.e. no-QCD contribution)\n",
    "                        hv_bflav_nonsubst.value -= hv_bflav_og.value\n",
    "                        hv_bflav_nonsubst.variance -= hv_bflav_og.variance\n",
    "                        hv[cat] = hv_bflav_comb + hv_bflav_nonsubst\n",
    "                        # print('+++', hv_bflav_og, hv_bflav, hv_bflav_comb, hv_bflav_nonsubst, hv[cat])\n",
    "                    \n",
    "            ## Fill the hv[cat] (for qcd*, fill hv[cat+'_herwig.value']) into TH1D and save into ROOT\n",
    "            for cat in hist.keys():\n",
    "                ## Special handling for qcdSyst / qcdKdeSyst\n",
    "                if 'qcd' in sys_name and 'SystUp' in sys_name:\n",
    "                    for i in range(nbin):\n",
    "                        hist[cat].SetBinContent(i+1, hv[cat+'_herwig.value'][i])\n",
    "                elif 'qcd' in sys_name and 'SystDown' in sys_name:\n",
    "                    hv[cat+'_herwig.value'] = g_hist_qcdsyst[(sys_name.replace('Down','Up'), b, cat)]\n",
    "                    for i in range(nbin):\n",
    "                        hist[cat].SetBinContent(i+1, 2 * hv[cat].value[i] - hv[cat+'_herwig.value'][i])\n",
    "                    g_hist_qcdsyst[(sys_name.replace('Down','Up'), b, cat)] = None\n",
    "\n",
    "                ## Normal routine\n",
    "                else:\n",
    "                    for i in range(nbin):\n",
    "                        hist[cat].SetBinContent(i+1, hv[cat].value[i])\n",
    "                        hist[cat].SetBinError(i+1, np.sqrt(hv[cat].variance[i]))\n",
    "                \n",
    "                ## Fix some buggy points\n",
    "                if cat!='data':\n",
    "                    for i in range(nbin):\n",
    "                        if hist[cat].GetBinContent(i+1) <= 1e-3:\n",
    "                            hist[cat].SetBinContent(i+1, 1e-3)\n",
    "                            hist[cat].SetBinError(i+1, 1e-3)\n",
    "                        elif hist[cat].GetBinError(i+1) > hist[cat].GetBinContent(i+1):\n",
    "                            hist[cat].SetBinError(i+1, hist[cat].GetBinContent(i+1))\n",
    "\n",
    "                if not g_dryrun:\n",
    "                    hist[cat].Write()\n",
    "        ## Close the ROOT file if error occurs (otherwise the notebook is easily corrupted)\n",
    "        finally:\n",
    "            if not g_dryrun:\n",
    "                fw.Close()"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "Now we launch the template maker"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "def launch_std_routine():\n",
    "    global g_bdt_mod_factor\n",
    "    g_bdt_mod_factor = None; launch_maker()\n",
    "    load_factor = df1[f\"bdt_mod_factor_{config['pt_range']['name']}__{config['main_analysis_tree']['name']}\"]\n",
    "    if load_factor is not None:\n",
    "        ## launch again for g_bdt_mod_factor set\n",
    "        g_bdt_mod_factor = load_factor; launch_maker()"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "## ====================================================================================================\n",
    "## Main fit routine: launch all sfBDT values, only run on 1st variable\n",
    "g_dryrun = False\n",
    "g_make_template_mode = 'main'; g_mode_bdt_runlist = 'all'\n",
    "g_make_unce_types = {'nominal':True, 'pu':True, 'fracBB':True, 'fracCC':True, 'fracLight':True, 'psWeightIsr':True, 'psWeightFsr':True, 'sfBDTRwgt':True, 'fitVarRwgt':True}\n",
    "g_mode_psWeight_run_templ = None\n",
    "g_do_fit_for_var = [1] # only run the first fit variable (2, 3 are for validation fit)\n",
    "launch_std_routine()\n",
    "# g_mode_bdt_runlist = 'manual'; g_do_sfBDT_points = {tuple(k):[0.75, 0.80, 0.85, 0.88, 0.90, 0.92, 0.94, 0.96, 0.98] for k in g_pt_range}; launch_maker() # if chooses to use fixed sfBDT points"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "--------\n",
    "**For year 2018**: you need to run the following block to provide psWeight templates for year 2016 and 2017 (otherwise for year condition 2016 and 2017 the above block will report errors)\n",
    "\n",
    "However, to acomplish the following block, you need to first run the same `preprocess.ipynb` for the corresponding 2016 and 2017 to extract the sfBDT sequence in that year condition."
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "## ====================================================================================================\n",
    "## For year 2018, extract necessary psWeight templates for 2016/2017\n",
    "if year == 2018:\n",
    "    for ext_year in [2016, 2017]:\n",
    "        g_make_template_mode = 'main'; g_mode_bdt_runlist = 'all'\n",
    "        g_make_unce_types = {'nominal':True, 'psWeightIsr':True, 'psWeightFsr':True}\n",
    "        g_mode_psWeight_run_templ = ext_year\n",
    "        g_do_fit_for_var = [1] # only run the first fit variable (2, 3 are for validation fit)\n",
    "        launch_std_routine()\n",
    "        # g_mode_bdt_runlist = 'manual'; g_do_sfBDT_points = {tuple(k):[0.75, 0.80, 0.85, 0.88, 0.90, 0.92, 0.94, 0.96, 0.98] for k in g_pt_range}; launch_maker() # if chooses to use fixed sfBDT points"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "--------\n",
    "Below are optional routines for the validation fit. No need to launch during the first run."
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "# ## ====================================================================================================\n",
    "# ## Validation on other variables\n",
    "# g_make_template_mode = 'main'; g_mode_bdt_runlist = 'all'\n",
    "# g_make_unce_types = {'nominal':True, 'pu':True, 'fracBB':True, 'fracCC':True, 'fracLight':True, 'psWeightIsr':True, 'psWeightFsr':True, 'sfBDTRwgt':True, 'fitVarRwgt':True}\n",
    "# g_mode_psWeight_run_templ = None\n",
    "# g_do_fit_for_var = [2, 3]\n",
    "# launch_maker()\n",
    "\n",
    "# ## ====================================================================================================\n",
    "# ## Multiple validations modes: only run the central sfBDT cut point is fine\n",
    "# for mode in ['val_pt', 'val_tosig_mass', 'val_tosig_pt', 'val_tosig_tau21', 'val_crop_bin']:\n",
    "#     g_make_template_mode = mode; g_mode_bdt_runlist = 'central'\n",
    "#     g_mode_psWeight_run_templ = None\n",
    "#     g_do_fit_for_var = [1]\n",
    "#     launch_maker()"
   ],
   "outputs": [],
   "metadata": {}
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}