make_template_ak.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "source": [
    "# Make templates for fit (`coffea`+`awkward` 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 `awkward-array` with `coffea` non-processor workflow for data processing. An alternative notebook `make_template_pd.ipynb` is based on `pandas` dataframe. The method in this notebook saves RAM, making it possible to run on lxplus, although it may take longer time for processing."
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "from coffea.nanoevents import NanoEventsFactory, TreeMakerSchema, BaseSchema\n",
    "import awkward1 as ak\n",
    "import uproot4 as uproot\n",
    "import numpy as np\n",
    "import math\n",
    "import os"
   ],
   "outputs": [],
   "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": [
    "## 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 lazy awkward arrays"
   ],
   "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",
    "## Read the root file into lazy awkward arrays\n",
    "arr = {}\n",
    "sample_prefix = f\"{config['samples']['sample_prefix']}_{year}\"\n",
    "for sam in read_sample_list_map:\n",
    "    arr[sam] = NanoEventsFactory.from_root(f'{sample_prefix}/{read_sample_list_map[sam]}', schemaclass=BaseSchema).events()\n",
    "\n",
    "## Store the branch names\n",
    "stored_branches = {}\n",
    "for sam in read_sample_list_map:\n",
    "    stored_branches[sam] = ak.fields(arr[sam])\n",
    "store_name = f\"{config['samples']['name']}_SF{config['year']}\""
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Load backup pickels"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "from data_utils import ExtendedNanoEventsArray\n",
    "def use_extended_nanoarray(arr, store_name):\n",
    "    for k in arr:\n",
    "        arr[k] = ExtendedNanoEventsArray(arr[k])\n",
    "        arr[k].record_awkward_items()\n",
    "        arr[k].set_backup_path(f'prep/{store_name}/{k}/') # backup directly to backup_array destination\n",
    "\n",
    "use_extended_nanoarray(arr, store_name)"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "## Fetch variables from the backup file\n",
    "## If you run the optional block above, i.e. to extend the custom coffea NanoEventsArray to ExtendedNanoEventsArray,\n",
    "## only 'maskdict' will be loaded into memory (arr[sam].maskdict) - other new variables will be loaded from disk automaticaly when you access arr[sam]['someVar']\n",
    "def load_backup_array(backup_name, read_sample_list_map):\n",
    "    r\"\"\"Load newly stored variables to the awkwary array list.\n",
    "    \n",
    "    Arguments:\n",
    "        backup_name: name of backup folder\n",
    "        read_sample_list: sample list to read.\n",
    "    \"\"\"\n",
    "\n",
    "    import pickle\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",
    "                    arr[sam].maskdict = {}\n",
    "                    with open(f'prep/{backup_name}/{sam}/maskdict', 'rb') as f:\n",
    "                        arr[sam].maskdict = pickle.load(f)\n",
    "                    print('loading...', sam, 'maskdict', arr[sam].maskdict.keys())\n",
    "                elif 'ExtendedNanoEventsArray' not in str(type(arr[sam])): # not using the extended nanoarray functionality\n",
    "                    with open(f'prep/{backup_name}/{sam}/{var}', 'rb') as f:\n",
    "                        arr[sam][var] = pickle.load(f)\n",
    "                    print('loading...', sam, var)\n",
    "            if sam != 'jetht-noht':\n",
    "                arr['subst_'+sam] = arr[sam] # make a reference\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",
    "                arr[sam] = pickle.load(f)\n",
    "            print('loading...', sam)\n",
    "\n",
    "load_backup_array(store_name, read_sample_list_map)"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "def eval_expr(ak_array, expr, mask=None):\n",
    "    \"\"\"A function that can do `eval` to the awkward array, immitating the behavior of `eval` in pandas.\"\"\"\n",
    "    \n",
    "    def get_variable_names(expr, exclude=['awkward', 'ak', 'np', 'numpy', 'math']):\n",
    "        \"\"\"Extract variables in the expr\"\"\"\n",
    "        import ast\n",
    "        root = ast.parse(expr)\n",
    "        return sorted({node.id for node in ast.walk(root) if isinstance(node, ast.Name) and not node.id.startswith('_')} - set(exclude))\n",
    "\n",
    "    tmp = {k:ak_array[k] if mask is None else ak_array[k].mask[mask] for k in get_variable_names(expr)}\n",
    "    tmp.update({'math': math, 'numpy': np, 'np': np, 'awkward': ak, 'ak': ak})\n",
    "#     print('eval expr: ', expr, '\\nvars', get_variable_names(expr))\n",
    "    return eval(expr, tmp)"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "def mask_and(arr, mask_list):\n",
    "    \"\"\"Calculate AND of given mask list\"\"\"\n",
    "    return np.logical_and.reduce([arr.maskdict[mask] for mask in mask_list])\n",
    "\n",
    "def concat_array(arrdict, expr, sam_list, filter_list):\n",
    "    \"\"\"Concatenate the awkward arrays passing the given filter list\"\"\"\n",
    "    if not isinstance(sam_list, list):\n",
    "        sam_list = [sam_list]\n",
    "    return np.concatenate([\n",
    "        np.array(eval_expr(arrdict[sam], expr)[mask_and(arrdict[sam], filter_list)]) for sam in sam_list\n",
    "    ])\n",
    "\n",
    "def mask_and_fj12(arr, mask_list):\n",
    "    \"\"\"Comibne `mask_and` result for fj_1 and fj_2\"\"\"\n",
    "    mask_list_fj1 = [ele.replace('fj_x', 'fj_1') for ele in mask_list]\n",
    "    mask_list_fj2 = [ele.replace('fj_x', 'fj_2') for ele in mask_list]\n",
    "    return np.concatenate([mask_and(arr, mask_list_fj1), mask_and(arr, mask_list_fj2)])\n",
    "\n",
    "def concat_array_fj12(arrdict, expr, sam_list, filter_list):\n",
    "    \"\"\"Comibne `concat_array` result for fj_1 and fj_2\"\"\"\n",
    "    filter_list_fj1 = [ele.replace('fj_x', 'fj_1') for ele in filter_list]\n",
    "    filter_list_fj2 = [ele.replace('fj_x', 'fj_2') for ele in filter_list]\n",
    "    return np.concatenate([concat_array(arrdict, expr.replace('fj_x', 'fj_1'), sam_list, filter_list_fj1), \n",
    "                           concat_array(arrdict, expr.replace('fj_x', 'fj_2'), sam_list, filter_list_fj2)])\n",
    "\n",
    "def calc_rwgt_akarray(arr, rwgt_edge, rwgt):\n",
    "    \"\"\"Calculate the weight ak-array based on the value ak-array of the reweight variable\"\"\"\n",
    "    arr_out = (arr<rwgt_edge[0])*rwgt[0]\n",
    "    for i in range(len(rwgt_edge)-1):\n",
    "        arr_out = arr_out + ((arr>=rwgt_edge[i]) & (arr<rwgt_edge[i+1]))*rwgt[i+1]\n",
    "    arr_out = arr_out + (arr>=rwgt_edge[-1])*rwgt[-1]\n",
    "    return arr_out"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Make ROOT templates\n",
    "\n",
    "We produce the ROOT templates using the ak-arrays 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 signal jet on mass\n",
    "        val_tosig_pt   : the validation fit -- additionally reweight MC & data to signal jet on pt  \n",
    "        val_tosig_tau21: the validation fit -- additionally reweight MC & data to 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.keys() != [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 = arr[f\"bdt_seq_{config['pt_range']['name']}__{config['main_analysis_tree']['name']}\"]\n",
    "    elif g_mode_bdt_runlist == 'central':\n",
    "        _points = arr[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",
    "    ## flavor masks\n",
    "    for sam in ['subst_'+s for s in read_sample_list_map if s not in ['jetht-noht']]:\n",
    "        for i in '12':\n",
    "            arr[sam].maskdict[f'fj_{i}_flvB'] = eval_expr(arr[sam], f'fj_{i}_nbhadrons>=1')\n",
    "            arr[sam].maskdict[f'fj_{i}_flvC'] = eval_expr(arr[sam], f'(fj_{i}_nbhadrons==0) & (fj_{i}_nchadrons>=1)')\n",
    "            arr[sam].maskdict[f'fj_{i}_flvL'] = eval_expr(arr[sam], f'(fj_{i}_nbhadrons==0) & (fj_{i}_nchadrons==0)')\n",
    "\n",
    "    for _ifit in g_do_fit_for_var:\n",
    "        for _wp in g_tagger_range:\n",
    "            \n",
    "            ## masks for applying the tagger\n",
    "            for sam in ['subst_'+s if s!='jetht-noht' else s for s in read_sample_list_map]:\n",
    "                for i in '12':\n",
    "                    arr[sam].maskdict[f'fj_{i}_tagger_pass'] = eval_expr(arr[sam], f\"({g_tagger_var.replace('fj_x', f'fj_{i}')}>{g_tagger_range[_wp][0]:.3f}) & ({g_tagger_var.replace('fj_x', f'fj_{i}')}<={g_tagger_range[_wp][1]:.3f})\")\n",
    "                    arr[sam].maskdict[f'fj_{i}_tagger_fail'] = eval_expr(arr[sam], f\"({g_tagger_var.replace('fj_x', f'fj_{i}')}<={g_tagger_range[_wp][0]:.3f}) | ({g_tagger_var.replace('fj_x', f'fj_{i}')}>{g_tagger_range[_wp][1]:.3f})\")\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",
    "                'categories_dm': ['flvL', 'flvB', 'flvC', 'data'],\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",
    "                'base_masks': {\n",
    "                    'mc': ['fj_x_base'],\n",
    "                    'data': ['fj_x_base'],\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(arr, fitinfo, lambda bdt, pt_range, sys_name: outdir_func(g_outdir_prefix_used, _wp, bdt, pt_range, sys_name), args, ext_masks=[])"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "def wrapperPt(arr, fitinfo, outdir_func, args, ext_masks):\n",
    "    r\"\"\"Depth 1: Process the pT cut and wrap all other following steps\n",
    "    \"\"\"\n",
    "    if 'optional' in config['samples'] and 'additional_selection_in_making_template' in config['samples']['optional']:\n",
    "        for sam in ['subst_'+s if s!='jetht-noht' else s for s in read_sample_list_map]:\n",
    "            for i in '12':\n",
    "                arr[sam].maskdict[f'fj_{i}_ext_ymal_sel'] = eval_expr(arr[sam], config['samples']['optional']['additional_selection_in_making_template'].replace('fj_x', f'fj_{i}'))\n",
    "        ext_ymal_sel = ['fj_x_ext_ymal_sel']\n",
    "    else:\n",
    "        ext_ymal_sel = []\n",
    "    \n",
    "    print('Launch pT range:', g_pt_range)\n",
    "    for pt_range in g_pt_range:\n",
    "        pt_range = tuple(pt_range)\n",
    "        ptlab = f'pt{pt_range[0]}to{pt_range[1]}'\n",
    "        for sam in ['subst_'+s if s!='jetht-noht' else s for s in read_sample_list_map]:\n",
    "            for i in '12':\n",
    "                arr[sam].maskdict[f'fj_{i}_{ptlab}'] = eval_expr(arr[sam], f'(fj_x_pt>={pt_range[0]}) & (fj_x_pt<{pt_range[1]})'.replace('fj_x', f'fj_{i}'))\n",
    "        print ('pt range:', pt_range)\n",
    "        \n",
    "        ## Loop over BDT varing list \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",
    "            ## masks for applying sfBDT cut\n",
    "            for sam in ['subst_'+s if s!='jetht-noht' else s for s in read_sample_list_map]:\n",
    "                for i in '12':\n",
    "                    arr[sam].maskdict[f'fj_{i}_mod_sfBDT>{sfBDT_val:.3f}'] = eval_expr(arr[sam], f'{bdt_expr}>{sfBDT_val}'.replace('fj_x', f'fj_{i}'))\n",
    "            \n",
    "            makeTemplatesWrapper(arr, fitinfo, lambda sys_name: outdir_func(sfBDT_val, pt_range, sys_name), sfBDT_val, args, ext_masks=ext_masks+ext_ymal_sel+[f'fj_x_{ptlab}', f'fj_x_mod_sfBDT>{sfBDT_val:.3f}'])"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "def makeTemplatesWrapper(arr, fitinfo, outdir_func, sfBDT_val, args, ext_masks):\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(arr, fitinfo, outdir_func(sys_name), sys_name, wgtstr_dm_sys, args, ext_masks)\n",
    "    \n",
    "    ## Below we extract hists for all unce type. Note: we only need such procedure in sfBDT>0.9 case (except for the validaiton when varying the sfBDT)\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(arr, fitinfo, outdir_func(sys_name), sys_name, wgtstr_dm_sys, args, ext_masks)\n",
    "        sys_name = 'puDown'; wgtstr_dm_sys = wgtstr_dm.replace('puWeight','puWeightDown'); makeTemplates(arr, fitinfo, outdir_func(sys_name), sys_name, wgtstr_dm_sys, args, ext_masks)\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(arr, fitinfo, outdir_func(sys_name), sys_name, wgtstr_dm_sys, args, ext_masks)\n",
    "        sys_name = \"fracBBDown\"; wgtstr_dm_sys = wgtstr_dm+'*(0.8*(fj_x_nbhadrons>1) + 1.0*(fj_x_nbhadrons<=1))'; makeTemplates(arr, fitinfo, outdir_func(sys_name), sys_name, wgtstr_dm_sys, args, ext_masks)\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*(np.logical_not((fj_x_nbhadrons==0) & (fj_x_nchadrons>1))))'; makeTemplates(arr, fitinfo, outdir_func(sys_name), sys_name, wgtstr_dm_sys, args, ext_masks)\n",
    "        sys_name = \"fracCCDown\"; wgtstr_dm_sys = wgtstr_dm+'*(0.8*((fj_x_nbhadrons==0) & (fj_x_nchadrons>1)) + 1.0*(np.logical_not((fj_x_nbhadrons==0) & (fj_x_nchadrons>1))))'; makeTemplates(arr, fitinfo, outdir_func(sys_name), sys_name, wgtstr_dm_sys, args, ext_masks)\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*(np.logical_not((fj_x_nbhadrons==0) & (fj_x_nchadrons==0))))'; makeTemplates(arr, fitinfo, outdir_func(sys_name), sys_name, wgtstr_dm_sys, args, ext_masks)\n",
    "        sys_name = \"fracLightDown\"; wgtstr_dm_sys = wgtstr_dm+'*(0.8*((fj_x_nbhadrons==0) & (fj_x_nchadrons==0)) + 1.0*(np.logical_not((fj_x_nbhadrons==0) & (fj_x_nchadrons==0))))'; makeTemplates(arr, fitinfo, outdir_func(sys_name), sys_name, wgtstr_dm_sys, args, ext_masks)\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(arr, fitinfo, outdir_func(sys_name), sys_name, wgtstr_dm_sys, args, ext_masks)\n",
    "        sys_name = \"qcdSystDown\"; wgtstr_dm_sys = wgtstr_dm; makeTemplates(arr, fitinfo, outdir_func(sys_name), sys_name, wgtstr_dm_sys, args, ext_masks)\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(arr, fitinfo, outdir_func(sys_name), sys_name, wgtstr_dm_sys, args, ext_masks)\n",
    "        sys_name = \"qcdKdeSystDown\"; wgtstr_dm_sys = wgtstr_dm; makeTemplates(arr, fitinfo, outdir_func(sys_name), sys_name, wgtstr_dm_sys, args, ext_masks)\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(arr, fitinfo, outdir_func(sys_name), sys_name, wgtstr_dm_sys, args, ext_masks)\n",
    "        sys_name = \"psWeightIsrDown\"; wgtstr_dm_sys = wgtstr_dm; makeTemplates(arr, fitinfo, outdir_func(sys_name), sys_name, wgtstr_dm_sys, args, ext_masks)\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(arr, fitinfo, outdir_func(sys_name), sys_name, wgtstr_dm_sys, args, ext_masks)\n",
    "        sys_name = \"psWeightFsrDown\"; wgtstr_dm_sys = wgtstr_dm; makeTemplates(arr, fitinfo, outdir_func(sys_name), sys_name, wgtstr_dm_sys, args, ext_masks)\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(arr, fitinfo, outdir_func(sys_name), sys_name, wgtstr_dm_sys, args, ext_masks, special_wgtstr='fj_x_sfbdtwgt_g50')\n",
    "        sys_name = 'sfBDTRwgtDown'; wgtstr_dm_sys = wgtstr_dm;'''factors decided by special_wgtstr argument'''; makeTemplates(arr, fitinfo, outdir_func(sys_name), sys_name, wgtstr_dm_sys, args, ext_masks, 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(arr, fitinfo, outdir_func(sys_name), sys_name, wgtstr_dm_sys, args, ext_masks)\n",
    "        sys_name = 'fitVarRwgtDown'; wgtstr_dm_sys = wgtstr_dm; makeTemplates(arr, fitinfo, outdir_func(sys_name), sys_name, wgtstr_dm_sys, args, ext_masks)\n",
    "    "
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "def makeTemplates(arr, fitinfo, outputdir, sys_name, wgtstr_dm_sys, args, ext_masks, special_wgtstr=None):\n",
    "    r\"\"\"Depth 3: The very base implementation that apply the final pass/fail cut and make the template\n",
    "    \"\"\"\n",
    "    print(ext_masks)\n",
    "    wgtstr_dm, wgtstr_dm_data, sl_dm, sl_dm_herwig, categories_dm, base_masks = args['wgtstr_dm'], args['wgtstr_dm_data'], args['sl_dm'], args['sl_dm_herwig'], args['categories_dm'], args['base_masks']\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",
    "        _mc_wgt = concat_array_fj12(arr, expr=_wgtstr, sam_list=_sl[:-1], filter_list=base_masks['mc']+ext_masks)\n",
    "        _data_wgt = np.ones_like(concat_array_fj12(arr, expr='ht', sam_list=[_sl[-1]], filter_list=base_masks['mc']+ext_masks))\n",
    "        return np.round(_data_wgt.sum() * 1. / _mc_wgt.sum(), 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",
    "        _content = concat_array_fj12(arr, expr=vname, sam_list=[sl_dm[-1]], filter_list=base_masks['data']+ext_masks)\n",
    "        _weights = np.ones_like(_content) if wgtstr_dm_data is None else concat_array_fj12(arr, expr=wgtstr_dm_data, sam_list=[sl_dm[-1]], filter_list=base_masks['data']+ext_masks)\n",
    "        _h_data = get_hist(_content, bins=edges, weights=_weights, underflow=underflow, overflow=overflow).view(flow=True)     \n",
    "        _content = concat_array_fj12(arr, expr=vname, sam_list=sl_dm[:-1], filter_list=base_masks['mc']+ext_masks)\n",
    "        _weights = concat_array_fj12(arr, expr=wgtstr_dm_sys, sam_list=sl_dm[:-1], filter_list=base_masks['mc']+ext_masks)\n",
    "        _h_mc = get_hist(_content, bins=edges, weights=_weights, 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 (' -- ', b)\n",
    "            \n",
    "            # Loop over categories: flvC/flvB/flvL/data\n",
    "            for cat in categories_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",
    "                    _content = concat_array_fj12(arr, expr=vname, sam_list=[sl_dm[-1]], filter_list=base_masks['data']+ext_masks+[f'fj_x_tagger_{b}'])\n",
    "                    _weights = np.ones_like(_content) if wgtstr_dm_data is None else concat_array_fj12(arr, expr=wgtstr_dm_data, sam_list=[sl_dm[-1]], filter_list=base_masks['data']+ext_masks+[f'fj_x_tagger_{b}'])\n",
    "                    hv['data'] = get_hist(_content, bins=edges, weights=_weights, 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",
    "                    ## Get the MC hist for certain flavor\n",
    "                    _content = concat_array_fj12(arr, expr=vname, sam_list=sl_dm[:-1], filter_list=base_masks['mc']+ext_masks+[f'fj_x_tagger_{b}', f'fj_x_{cat}'])\n",
    "                    _weights = concat_array_fj12(arr, expr=wgtstr_dm_sys, sam_list=sl_dm[:-1], filter_list=base_masks['mc']+ext_masks+[f'fj_x_tagger_{b}', f'fj_x_{cat}'])\n",
    "                    hv[cat] = get_hist(_content, bins=edges, weights=_weights, 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",
    "                        wgtstr_dm_sys_herwig = wgtstr_dm_sys.replace('htwgt','htwgt_herwig').replace('sfbdtwgt_g90','sfbdtwgt_g90_herwig').replace('ad_ptwgt','ad_ptwgt_herwig').replace('datamcwgt','datamcwgt_herwig')\n",
    "                        _content = concat_array_fj12(arr, expr=vname, sam_list=sl_dm_herwig[:-1], filter_list=base_masks['mc']+ext_masks+[f'fj_x_tagger_{b}', f'fj_x_{cat}'])\n",
    "                        _weights = concat_array_fj12(arr, expr=wgtstr_dm_sys_herwig, sam_list=sl_dm_herwig[:-1], filter_list=base_masks['mc']+ext_masks+[f'fj_x_tagger_{b}', f'fj_x_{cat}'])                        \n",
    "                        hv[cat+'_herwig.value'] = get_hist(_content, bins=edges, weights=_weights, 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",
    "                        wgtstr_dm_sys_herwig = wgtstr_dm_sys.replace('htwgt','htwgt_herwig').replace('sfbdtwgt_g90','sfbdtwgt_g90_herwig').replace('ad_ptwgt','ad_ptwgt_herwig').replace('datamcwgt','datamcwgt_herwig')\n",
    "                        _content = concat_array_fj12(arr, expr=vname, sam_list=sl_dm_herwig[:-1], filter_list=base_masks['mc']+ext_masks+[f'fj_x_tagger_{b}', f'fj_x_{cat}'])\n",
    "                        _weights = concat_array_fj12(arr, expr=wgtstr_dm_sys_herwig, sam_list=sl_dm_herwig[:-1], filter_list=base_masks['mc']+ext_masks+[f'fj_x_tagger_{b}', f'fj_x_{cat}'])                        \n",
    "                        hv_herwig_orig_value = get_hist(_content, bins=edges, weights=_weights, 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(_content, weights=np.clip(_weights, 0, +np.inf))\n",
    "                        kde = gaussian_kde(_content, weights=np.clip(_weights, 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",
    "                            wgtstr_dm_sys_ps = wgtstr_dm_sys + f\"*(PSWeight[:,{ps_idx[sys_name]}])\"\n",
    "                            _weights = concat_array_fj12(arr, expr=wgtstr_dm_sys_ps, sam_list=sl_dm[:-1], filter_list=base_masks['mc']+ext_masks+[f'fj_x_tagger_{b}', f'fj_x_{cat}'])\n",
    "                            hv[cat] = get_hist(_content, bins=edges, weights=_weights, 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 = uproot.open(outputdir_ps_18.replace(sys_name, 'nominal')+f'inputs_{b}.root')[cat]\n",
    "                            hv_ps_18 = uproot.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 = uproot.open(outputdir.replace(sys_name, 'nominal')+f'inputs_{b}.root')['data_obs'].values()[1:-1]  ## nominal data hist for 0.90\n",
    "                        _bdtname = '95' if 'Up' in sys_name else '85'\n",
    "                        fr = uproot.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()[1:-1], fr['flvC'].values()[1:-1]+fr['flvB'].values()[1:-1]+fr['flvL'].values()[1:-1]  ## 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()[1:-1] / 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",
    "                        _content = concat_array_fj12(arr, expr=vname, sam_list=args['args_bflav']['sl_dm_bflav'], filter_list=base_masks['mc']+ext_masks+[f'fj_x_tagger_{b}', f'fj_x_{cat}'])\n",
    "                        _weights = concat_array_fj12(arr, expr=args['args_bflav']['wgtstropt_bflav'](wgtstr_dm_sys), sam_list=args['args_bflav']['sl_dm_bflav'], filter_list=base_masks['mc']+ext_masks+[f'fj_x_tagger_{b}', f'fj_x_{cat}'])\n",
    "                        hv_bflav = get_hist(_content, bins=edges, weights=_weights, underflow=underflow, overflow=overflow).view(flow=True)\n",
    "                        _content = concat_array_fj12(arr, expr=vname, sam_list=args['args_bflav']['sl_dm_bflav_orig'], filter_list=base_masks['mc']+ext_masks+[f'fj_x_tagger_{b}', f'fj_x_{cat}'])\n",
    "                        _weights = concat_array_fj12(arr, expr=wgtstr_dm_sys, sam_list=args['args_bflav']['sl_dm_bflav_orig'], filter_list=base_masks['mc']+ext_masks+[f'fj_x_tagger_{b}', f'fj_x_{cat}'])\n",
    "                        hv_bflav_og = get_hist(_content, bins=edges, weights=_weights, 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 = arr[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
}