preprocess.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Pre-processing\n",
    "\n",
    "This notebook aims to do pre-processing before making the fit templates. The pre-processing includes calculating new variables and extracting reweight factors. \n",
    "\n",
    "The data analyzing is backed by `awkward-array` with `coffea` non-processor workflow.\n",
    "\n",
    "Please first setup `config.yaml`, then run the notebook in the order of blocks. The results are stored in the backup folder `prep/` with your specified \"routine name\" and \"year condition\" in the yaml file. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from data_utils import get_hist, plot_hist\n",
    "from cycler import cycler\n",
    "import boost_histogram as bh\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib as mpl\n",
    "# mpl.use('AGG') # no rendering plots in the window\n",
    "\n",
    "import mplhep as hep\n",
    "use_helvet = True  ## true: use helvetica for plots, make sure the system have the font installed\n",
    "if use_helvet:\n",
    "    CMShelvet = hep.style.CMS\n",
    "    CMShelvet['font.sans-serif'] = ['Helvetica', 'Arial']\n",
    "    plt.style.use(CMShelvet)\n",
    "else:\n",
    "    plt.style.use(hep.style.CMS)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Load the config.yml\n",
    "import yaml\n",
    "with open('cards/config_bb_std.yml') as f:\n",
    "    config = yaml.safe_load(f)\n",
    "config"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Load files\n",
    "\n",
    "Load the ROOT files into lazy awkward arrays"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "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_preprocessing' in config['samples']['optional']:\n",
    "    for ex_sam in config['samples']['optional']['exclude_mc_sample_in_preprocessing']:\n",
    "        read_sample_list_map.pop(ex_sam, None)\n",
    "print('Read samples for preprocessing:', 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",
    "def get_stored_branches(arr, read_sample_list_map):\n",
    "    stored_branches = {}\n",
    "    for sam in read_sample_list_map:\n",
    "        stored_branches[sam] = ak.fields(arr[sam])\n",
    "    return stored_branches\n",
    "stored_branches = get_stored_branches(arr, read_sample_list_map)\n",
    "stored_branches_interm = {}\n",
    "store_name = f\"{config['samples']['name']}_SF{config['year']}\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## You can optionally run this block to extend the custom coffea NanoEventsArray to ExtendedNanoEventsArray.\n",
    "## This functionality will automatically store and read NEW variables from disk instead of having them all in memory.\n",
    "\n",
    "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)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Apply baseline selections\n",
    "\n",
    "For data: apply OR of all HT trigger to enhance statistics.\n",
    "\n",
    "For MC: apply no HT trigger, based on the strategy we name it \"MC substitute\".\n",
    "\n",
    "We define an attribute `maskdict` in each sample that stores masks corresponding to different selections."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def backup_array(backup_name, stored_branches, read_sample_list_map, global_key_only=False, ext_branches=[]):\n",
    "    r\"\"\"Backup newly produced variables in the awkward array list to pickle.\n",
    "    \n",
    "    Arguments:\n",
    "        backup_name: name of backup folder\n",
    "        stored_branches: branches already stored in previous routines. New branches that does not appear in the list will be stored, and be updated in the list.\n",
    "        read_sample_list: sample list to read.\n",
    "    \"\"\"\n",
    "\n",
    "    import pickle\n",
    "    if not global_key_only:\n",
    "        for sam in read_sample_list_map:\n",
    "            if 'ExtendedNanoEventsArray' in str(type(arr[sam])):\n",
    "                ext_branches = []\n",
    "            if not os.path.exists(f'prep/{backup_name}/{sam}'):\n",
    "                os.makedirs(f'prep/{backup_name}/{sam}')\n",
    "            for var in set(ak.fields(arr[sam])) - set(stored_branches[sam]) | set([br for br in ext_branches if br in ak.fields(arr[sam])]):\n",
    "                with open(f'prep/{backup_name}/{sam}/{var}', 'wb') as fw:\n",
    "                    pickle.dump(arr[sam][var], fw)\n",
    "                print('storing...', sam, var)\n",
    "            if hasattr(arr[sam], 'maskdict'):\n",
    "                with open(f'prep/{backup_name}/{sam}/maskdict', 'wb') as fw:\n",
    "                    pickle.dump(arr[sam].maskdict, fw)\n",
    "    for key in arr.keys():\n",
    "        if key not in read_sample_list_map and not key.startswith('subst_') and key != 'real-signal':\n",
    "            with open(f'prep/{backup_name}/{key}', 'wb') as fw:\n",
    "                pickle.dump(arr[key], fw)\n",
    "            print('storing additonal keys...', key)\n",
    "\n",
    "## Fetch variables from the backup file\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)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "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)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ## If you have run some blocks in the following - use this to restore the backup variables directly\n",
    "# load_backup_array(store_name, read_sample_list_map)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### ================ Pre-processing for data  ===================\n",
    "\n",
    "## Baseline selection applied to data. \n",
    "## Note that we use the OR or all HT triggers (some are pre-scaled triggers)\n",
    "\n",
    "if 'prep-data' not in stored_branches_interm:\n",
    "    stored_branches_interm['prep-data'] = get_stored_branches(arr, read_sample_list_map)\n",
    "hlt_branches = {  ## used HLT_PFHT* 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",
    "htcut_incl = '('+' | '.join(hlt_branches[year])+')'\n",
    "basesel_ext_noht_prep = f\"passmetfilters & (fj_x_pt>200) & fj_x_is_qualified\"\n",
    "sl_prep = ['jetht-noht']\n",
    "\n",
    "for sam in sl_prep:\n",
    "    assert 'noht' in sam\n",
    "    arr[sam].maskdict = {}\n",
    "    arr[sam].maskdict['hlt'] = eval_expr(arr[sam], htcut_incl)\n",
    "    for i in '12':\n",
    "        ## The baseline selection for data\n",
    "        print('baseline selection for data: ', sam, f'jet{i}')\n",
    "        arr[sam].maskdict[f'fj_{i}_base'] = arr[sam].maskdict['hlt'] & eval_expr(arr[sam], basesel_ext_noht_prep.replace('fj_x', f'fj_{i}'))\n",
    "\n",
    "## Store new variables\n",
    "stored_branches = backup_array(store_name, stored_branches_interm['prep-data'], read_sample_list_map)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ## FOR TEST: check the xsecWeight for MG samples & genWeight for Herwig sample (to avoid extremely large values) \n",
    "# from collections import Counter\n",
    "# print(Counter(np.array(arr['qcd-mg-noht'].xsecWeight)),'\\n')\n",
    "# for i in [0.96, 0.98, 0.99]:\n",
    "#     print(np.quantile(np.array(arr['qcd-herwig-noht'].genWeight), q=i))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### ================ Pre-processing for MC substitute  ===================\n",
    "\n",
    "## Baseline selection applied to MC.\n",
    "## No HT trigger is applied, based on the \"MC substitute\" strategy\n",
    "if 'prep-mc' not in stored_branches_interm:\n",
    "    stored_branches_interm['prep-mc'] = get_stored_branches(arr, read_sample_list_map)\n",
    "\n",
    "basesel_noht_prep_subst = \"passmetfilters & (fj_x_pt>200) & fj_x_is_qualified\"\n",
    "## Mark sample name with \"subst_\" as a reminder of MC substitute. Default is ['subst_qcd-mg-noht', 'subst_qcd-herwig-noht', 'subst_top-noht', 'subst_v-qq-noht']  \n",
    "sl_prep_subst = ['subst_'+sam for sam in read_sample_list_map.keys() if 'jetht' not in sam]\n",
    "if config['samples']['use_bflav']:\n",
    "    sl_prep_subst += ['subst_qcd-mg-bflav-noht']\n",
    "for sam in sl_prep_subst:\n",
    "    assert 'noht' in sam\n",
    "    arr[sam] = arr[sam.replace('subst_','')]  ## use the name subst_ as a ref\n",
    "    arr[sam].maskdict = {}\n",
    "    for i in '12':\n",
    "        print('baseline selection for: ', sam, f'jet{i}')\n",
    "        arr[sam].maskdict[f'fj_{i}_base'] = eval_expr(arr[sam], basesel_noht_prep_subst.replace('fj_x', f'fj_{i}'))\n",
    "        ## Drop MG events with extremely large xsecWeight (coming from low HT sample in the HT-binned MG list)\n",
    "        if sam == 'subst_qcd-mg-noht':\n",
    "            arr[sam].maskdict[f'fj_{i}_base'] = arr[sam].maskdict[f'fj_{i}_base'] & eval_expr(arr[sam], 'xsecWeight<5.')\n",
    "        ## Drop Herwig events with extremely large genWeight\n",
    "        if sam == 'subst_qcd-herwig-noht':\n",
    "            arr[sam].maskdict[f'fj_{i}_base'] = arr[sam].maskdict[f'fj_{i}_base'] & eval_expr(arr[sam], 'genWeight<{}'.format(np.quantile(np.array(arr[sam].genWeight), q=0.96)))\n",
    "    ## Fix a 2016 bug: Herwig sample xsec is mistaken\n",
    "    if year == 2016 and sam == 'subst_qcd-herwig-noht' and not hasattr(arr[sam], 'xsecWeight_is_normed'):\n",
    "        arr[sam]['xsecWeight'] = arr[sam]['xsecWeight'] * 2400.\n",
    "        arr[sam]['xsecWeight_is_normed'] = True\n",
    "\n",
    "## Produce new variables used for fit\n",
    "for sam in sl_prep + sl_prep_subst:\n",
    "    for i in '12':\n",
    "        _mask = arr[sam].maskdict[f'fj_{i}_base']\n",
    "        print('calculating new vars for: ', sam, f'jet{i}')\n",
    "        arr[sam][f'fj_{i}_mSV12_ptmax'] = eval_expr(arr[sam], f'(fj_{i}_sj1_sv1_pt>fj_{i}_sj2_sv1_pt)*fj_{i}_sj1_sv1_masscor + (fj_{i}_sj1_sv1_pt<=fj_{i}_sj2_sv1_pt)*fj_{i}_sj2_sv1_masscor', mask=_mask)\n",
    "        arr[sam][f'fj_{i}_mSV12_ptmax_log'] = eval_expr(arr[sam], f'np.log(fj_{i}_mSV12_ptmax)', mask=_mask)\n",
    "        arr[sam][f'fj_{i}_mSV12_dxysig'] = eval_expr(arr[sam], f'(fj_{i}_sj1_sv1_dxysig>fj_{i}_sj2_sv1_dxysig)*fj_{i}_sj1_sv1_masscor + (fj_{i}_sj1_sv1_dxysig<=fj_{i}_sj2_sv1_dxysig)*fj_{i}_sj2_sv1_masscor', mask=_mask)\n",
    "        arr[sam][f'fj_{i}_mSV12_dxysig_log'] = eval_expr(arr[sam], f'np.log(fj_{i}_mSV12_dxysig)', mask=_mask)\n",
    "\n",
    "## Store new variables\n",
    "stored_branches = backup_array(store_name, stored_branches_interm['prep-mc'], read_sample_list_map, ext_branches=['xsecWeight'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Obtain reweight factors\n",
    "\n",
    "We need to extract some reweight factors as well as the BDT variation points specific to pT ranges. Step 1-3 are necessary for the nominal fit routine. Factors obtained from step 4-5 are for validation fits.\n",
    "\n",
    " 3-1. **MC substitute-to-data reweight factor**: reweight based on the 3D (HT, pT, jet index) grid. The goal is to bring the shape of MC (without pre-selection on the jet-HT triggers) back to the data shape (passing the logical OR of prescaled jet-HT triggers). Remember that the raw MC always yields much larger than data. New variables take the name `htwgt`, `htwgt_herwig`. (`htwgt_herwig` is derived using the Herwig QCD sample and is only used in the validation fit.)\n",
    "\n",
    " 3-2. **sfBDT central point and variation range**: a set of sfBDT cut values which are specific for different pT range. The values are derived by judging the similarity of the tagger shape between the signal and proxy jet samples.\n",
    "\n",
    " 3-3. **sfBDT reweight factor**: reweight on the sfBDT variable based on (pT, jet index) bins. The reweight factors `sfbdtwgt_g50` are obtained, which is only used to derive the systematics shape templates in the nominal fit. `sfbdtwgt_g50_herwig` is derived as well using the Herwig sample, used in the validation fit.\n",
    "\n",
    " 3-4. **Additional MC substitute-to-data reweight factor on $p_{T}$ only**: A possible replacement of the factors in step 1. This factor is only used in the validation fit to check if different reweighting schemes may affect the SF fit results. New variables take the name `ad_ptwgt` and `ad_ptwgt_herwig`.\n",
    " \n",
    " 3-5. **Proxy-to-signal reweight factor on $m_{SD}$ / $p_{T}$ / $\\tau_{21}$**: based on the shape of MC after applying the MC-to-data factors in step 1 and the H->cc signal jet shape. The factor is only used in the validation fit, in which we apply such reweight factor to both MC and data to check if the SF results are affected. New variables take the name `(mass|pt|tau21)datamcwgt` and `(mass|pt|tau21)datamcwgt_herwig`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### ================ 3-1. Reweight MC subsitute to data: stored as variable \"fj_x_htwgt\", \"fj_x_htwgt_herwig\") ===================\n",
    "\n",
    "## True: if the block has run before, we can obtain the reweight factor from the previously stored pickle output\n",
    "is_read_from_pickel = False\n",
    "\n",
    "if 'prep-3-1' not in stored_branches_interm:\n",
    "    stored_branches_interm['prep-3-1'] = get_stored_branches(arr, read_sample_list_map)\n",
    "\n",
    "def extract_source_to_target_ht_weight(arr, sl_rwgt_source, wgtstr_rwgt_source, sl_rwgt_target, wgtstr_rwgt_target, wgtname, ext_sl_rwgt_source=[], presel='', do_plot=True):\n",
    "    r\"\"\"Extract the \"MC subsisute to data\" reweight factor on HT based on (pT, jet index) bins\n",
    "    \n",
    "    Arguments:\n",
    "        arr: awkward array dict as input\n",
    "        sl_rwgt_(source|target): sample list for the source/target in this reweighting routine\n",
    "        wgtstr_rwgt_(source|target): the weight string applied to the source/target to produce the histogram in this reweighting routine\n",
    "        wgtname: the reweight name stored as a new column\n",
    "        ext_sl_rwgt_source: extra source sample list for which we also calculate the reweight factors after extracting them\n",
    "        presel: additonal pre-selection applied before reweigting\n",
    "        do_plot: if store plots of reweighting\n",
    "    \"\"\"\n",
    "\n",
    "    rwgt_var = 'ht'\n",
    "    ## The binning info for (pT, HT) grid. The adopted HT grid is based on MC shape in each pT bin\n",
    "    rwgt_edge_dic = {}\n",
    "    def linear_edge(start, end):\n",
    "        return list(np.arange(start, end-50, 50)) + [end]\n",
    "    rwgt_edge_dic[2016] = rwgt_edge_dic[2017] = rwgt_edge_dic[2018] = {\n",
    "        'jet1': {\n",
    "            'pt200to250': linear_edge(250, 1250),\n",
    "            'pt250to300': linear_edge(350, 1400),\n",
    "            'pt300to350': linear_edge(400, 1600),\n",
    "            'pt350to400': linear_edge(450, 1700),\n",
    "            'pt400to450': linear_edge(500, 1800),\n",
    "            'pt450to500': linear_edge(550, 1900),\n",
    "            'pt500to550': linear_edge(600, 1900),\n",
    "            'pt550to600': linear_edge(650, 2000),\n",
    "            'pt600to700': linear_edge(700, 2100),\n",
    "            'pt700to800': linear_edge(800, 2200),\n",
    "            'pt800to100000': linear_edge(1000, 2400),\n",
    "        },\n",
    "        'jet2': {\n",
    "            'pt200to250': linear_edge(250, 1500),\n",
    "            'pt250to300': linear_edge(350, 1600),\n",
    "            'pt300to350': linear_edge(400, 1800),\n",
    "            'pt350to400': linear_edge(450, 2000),\n",
    "            'pt400to450': linear_edge(500, 2200),\n",
    "            'pt450to500': linear_edge(550, 2400),\n",
    "            'pt500to550': linear_edge(650, 2400),\n",
    "            'pt550to600': linear_edge(750, 2400),\n",
    "            'pt600to700': linear_edge(850, 2400),\n",
    "            'pt700to800': linear_edge(1000, 2400),\n",
    "            'pt800to100000': linear_edge(1200, 2400),\n",
    "        },\n",
    "\n",
    "    }\n",
    "    \n",
    "    ## Initially fill the output column with 0, since we will fill the column iteratively for each pT bin\n",
    "    for sam in sl_rwgt_source + ext_sl_rwgt_source:\n",
    "        for i in '12':\n",
    "            arr[sam][wgtname.replace('fj_x', f'fj_{i}')] = ak.zeros_like(arr[sam][rwgt_var])\n",
    "\n",
    "    if is_read_from_pickel: ## restore info from a previously stored pickle\n",
    "        import pickle\n",
    "        with open(f'prep/{store_name}/plots/{wgtname}_{year}.pickle', 'rb') as f:\n",
    "            res = pickle.load(f)\n",
    "            ent_target, ent_source, rwgt = res['ent_target'], res['ent_source'], res['rwgt']\n",
    "    else:\n",
    "        ent_target, ent_source, rwgt = {}, {}, {}\n",
    "\n",
    "    ## Rewight separately on jet pT bins\n",
    "    ptlab_list = list(rwgt_edge_dic[year]['jet1'].keys())\n",
    "    import re\n",
    "    ptcutval_list = [re.findall('pt(\\d+)to(\\d+)', s)[0] for s in ptlab_list]\n",
    "    for ptsel, ptlab in zip([f'(fj_x_pt>={ptmin}) & (fj_x_pt<{ptmax})' for ptmin, ptmax in ptcutval_list], ptlab_list):\n",
    "        ## Reweight separately for 1st or 2nd jet\n",
    "        for i, lab in zip(['1','2'], ['jet1','jet2']):\n",
    "            print (' -- ', ptsel, lab)\n",
    "            rwgt_edge = rwgt_edge_dic[year][lab][ptlab]\n",
    "            ## Calculate the rwgt for the first time\n",
    "            if not is_read_from_pickel:\n",
    "                for sam in set(sl_rwgt_target) | set(sl_rwgt_source) | set(ext_sl_rwgt_source):\n",
    "                    arr[sam].maskdict[f'fj_{i}_{ptlab}'] = eval_expr(arr[sam], ptsel.replace('fj_x', f'fj_{i}'))\n",
    "                    if presel != '':  ## has additional preselection\n",
    "                        arr[sam].maskdict[f'_tmp_fj_{i}_presel'] = eval_expr(arr[sam], presel.replace('fj_x', f'fj_{i}'))\n",
    "\n",
    "                ## Get data and MC histogram. Note: consider underflow & overflow bins, hence len = nbins+2\n",
    "                ent_target[ptlab+lab] = get_hist(\n",
    "                    concat_array(arr, expr=rwgt_var, sam_list=sl_rwgt_target, filter_list=[f'fj_{i}_base', f'fj_{i}_{ptlab}']+([f'_tmp_fj_{i}_presel'] if presel != '' else [])),\n",
    "                    bins=rwgt_edge,\n",
    "                    weights=concat_array(arr, expr=wgtstr_rwgt_target if wgtstr_rwgt_target!='1' else '0*ht+1', sam_list=sl_rwgt_target, filter_list=[f'fj_{i}_base', f'fj_{i}_{ptlab}']+([f'_tmp_fj_{i}_presel'] if presel != '' else [])),\n",
    "                    underflow=True, overflow=True, mergeflowbin=False\n",
    "                ).view(flow=True).value\n",
    "                ent_source[ptlab+lab] = get_hist(\n",
    "                    concat_array(arr, expr=rwgt_var, sam_list=sl_rwgt_source, filter_list=[f'fj_{i}_base', f'fj_{i}_{ptlab}']+([f'_tmp_fj_{i}_presel'] if presel != '' else [])),\n",
    "                    bins=rwgt_edge,\n",
    "                    weights=concat_array(arr, expr=wgtstr_rwgt_source if wgtstr_rwgt_source!='1' else '0*ht+1', sam_list=sl_rwgt_source, filter_list=[f'fj_{i}_base', f'fj_{i}_{ptlab}']+([f'_tmp_fj_{i}_presel'] if presel != '' else [])),\n",
    "                    underflow=True, overflow=True, mergeflowbin=False\n",
    "                ).view(flow=True).value\n",
    "                ## Calculate the reweight factor\n",
    "                rwgt[ptlab+lab] = np.nan_to_num(ent_target[ptlab+lab] / ent_source[ptlab+lab], nan=0) # len=nbin+2\n",
    "                rwgt[ptlab+lab] = np.clip(rwgt[ptlab+lab], 0, 2) # clip between [0, 2]\n",
    "            print(ent_target[ptlab+lab], '\\n', rwgt[ptlab+lab])\n",
    "\n",
    "            ## Assign the reweight factor to the new column\n",
    "            for sam in sl_rwgt_source + ext_sl_rwgt_source:\n",
    "                _var = rwgt_var\n",
    "                _wgtname = wgtname.replace('fj_x', f'fj_{i}')\n",
    "                _mask = mask_and(arr[sam], mask_list=[f'fj_{i}_base', f'fj_{i}_{ptlab}'])\n",
    "                arr[sam][_wgtname] = arr[sam][_wgtname] + ak.fill_none(calc_rwgt_akarray(arr[sam][_var].mask[_mask], rwgt_edge, rwgt[ptlab+lab]), 0)\n",
    "#                 print('midpoint: ', sam, _wgtname, arr[sam][_wgtname])\n",
    "\n",
    "    if not is_read_from_pickel: ## store the info for the first run\n",
    "        import pickle\n",
    "        if not os.path.exists(f'prep/{store_name}/plots'):\n",
    "            os.makedirs(f'prep/{store_name}/plots')\n",
    "        with open(f'prep/{store_name}/plots/{wgtname}_{year}.pickle', 'wb') as fw:\n",
    "            pickle.dump({'ent_target':ent_target, 'ent_source':ent_source, 'rwgt':rwgt}, fw)\n",
    "\n",
    "    # =========== plot ===========\n",
    "    if do_plot:\n",
    "        mpl.rcParams['axes.prop_cycle'] = cycler(color=['blue', 'red', 'green', 'violet', 'darkorange', 'black', 'cyan', 'yellow'])\n",
    "        for ptlab in ptlab_list:\n",
    "            f = plt.figure(figsize=(12,12))\n",
    "            gs = mpl.gridspec.GridSpec(2, 1, height_ratios=[2, 1], hspace=0.04) \n",
    "            for lab, cm, cd in zip(['jet1', 'jet2'], ['blue', 'red'], ['royalblue', 'lightcoral']):\n",
    "                ax = f.add_subplot(gs[0])\n",
    "                hep.cms.label(data=True, paper=False, year=year, ax=ax, rlabel=r'%s $fb^{-1}$ (13 TeV)'%lumi[year], fontname='sans-serif')\n",
    "                hep.histplot(ent_source[ptlab+lab], bins=[0]+list(rwgt_edge_dic[year][lab][ptlab])+[2500], label='Jet '+lab[-1]+' (MC)', color=cm)\n",
    "                hep.histplot(ent_target[ptlab+lab], bins=[0]+list(rwgt_edge_dic[year][lab][ptlab])+[2500], label='Jet '+lab[-1]+' (Data)', color=cd, linestyle='--')\n",
    "\n",
    "                ax.set_xlim(0, 2500); ax.set_xticklabels([]); \n",
    "                ax.set_yscale('log'); ax.set_ylabel('Events', ha='right', y=1.0)\n",
    "                ax.legend()\n",
    "                ax1 = f.add_subplot(gs[1]); \n",
    "                hep.histplot(rwgt[ptlab+lab], bins=[0]+list(rwgt_edge_dic[year][lab][ptlab])+[2500], label='Jet '+lab[-1], color=cm)\n",
    "                ax1.set_xlim(0, 2500); ax1.set_xlabel('$H_{T}$ [GeV]', ha='right', x=1.0);\n",
    "                ax1.legend()\n",
    "                ax1.set_yscale('log')\n",
    "                ax1.set_ylim(5e-3, 2e0); ax1.set_ylabel('Rwgt factor', ha='right', y=1.0);  ax1.set_yticks([1e-2,1e-1,1e0,1e1]);\n",
    "                ax1.plot([0, 2500], [1, 1], 'k:')\n",
    "\n",
    "            if not os.path.exists(f'prep/{store_name}/plots'):\n",
    "                os.makedirs(f'prep/{store_name}/plots')\n",
    "            plt.savefig(f'prep/{store_name}/plots/rwgtfac_{wgtname}_{year}_{ptlab}_{lab}.pdf')\n",
    "            plt.savefig(f'prep/{store_name}/plots/rwgtfac_{wgtname}_{year}_{ptlab}_{lab}.png')\n",
    "    # ============================\n",
    "    \n",
    "    return {'ent_target':ent_target, 'ent_source':ent_source, 'rwgt':rwgt}\n",
    "\n",
    "## Calculate two sets of reweight factor: one for the MG sample list and another for Herwig sample list\n",
    "htwgt = extract_source_to_target_ht_weight(\n",
    "    arr, sl_rwgt_source=['subst_'+s for s in read_sample_list_map if s not in ['qcd-herwig-noht', 'qcd-mg-bflav-noht', 'jetht-noht']], ext_sl_rwgt_source=['subst_qcd-mg-bflav-noht'] if config['samples']['use_bflav'] else [], \n",
    "    wgtstr_rwgt_source=f\"{lumi[year]}*genWeight*xsecWeight*puWeight\",\n",
    "    sl_rwgt_target=['jetht-noht'], wgtstr_rwgt_target='1', wgtname='fj_x_htwgt',\n",
    ")\n",
    "if not omit_herwig:\n",
    "    htwgt_herwig = extract_source_to_target_ht_weight(\n",
    "        arr, sl_rwgt_source=['subst_'+s for s in read_sample_list_map if s not in ['qcd-mg-noht', 'qcd-mg-bflav-noht', 'jetht-noht']], ext_sl_rwgt_source=['subst_qcd-mg-bflav-noht'] if config['samples']['use_bflav'] else [],\n",
    "        wgtstr_rwgt_source=f\"{lumi[year]}*genWeight*xsecWeight*puWeight\", \n",
    "        sl_rwgt_target=['jetht-noht'], wgtstr_rwgt_target='1', wgtname='fj_x_htwgt_herwig',\n",
    "    )\n",
    "\n",
    "## Calculate bflav factors: reweight bflav sample to inclusive QCD (after b selection cut)\n",
    "if config['samples']['use_bflav']:\n",
    "    bflav_htwgt = extract_source_to_target_ht_weight(\n",
    "        arr, sl_rwgt_source=['subst_qcd-mg-bflav-noht'], wgtstr_rwgt_source=f\"{lumi[year]}*genWeight*xsecWeight*puWeight\",\n",
    "        sl_rwgt_target=['subst_qcd-mg-noht'], wgtstr_rwgt_target=f\"{lumi[year]}*genWeight*xsecWeight*puWeight\", wgtname='fj_x_bflav_htwgt',\n",
    "        presel='fj_x_nbhadrons>=1', do_plot=False,\n",
    "    )\n",
    "    for sam in ['subst_'+s for s in read_sample_list_map if s not in ['qcd-mg-bflav-noht', 'jetht-noht']]:\n",
    "        arr[sam]['bflav_htwgt'] = ak.ones_like(arr[sam]['ht'])\n",
    "\n",
    "## Store new variables\n",
    "stored_branches = backup_array(store_name, stored_branches_interm['prep-3-1'], read_sample_list_map)\n",
    "\n",
    "## Test output\n",
    "ak.to_pandas(arr['subst_qcd-mg-noht'][['ht', 'fj_1_pt', 'fj_1_htwgt']][arr['subst_qcd-mg-noht'].maskdict['fj_1_base']])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### ================ 3-2. Determine the optimal sfBDT cut value for each pT range  ===================\n",
    "# First load the h->cc signal ntuple. Adopt the selction used in the analysis\n",
    "import re\n",
    "arr['real-signal'] = NanoEventsFactory.from_root(config['main_analysis_tree']['path'].replace('$YEAR', str(year)), treepath='/'+config['main_analysis_tree']['treename'], schemaclass=BaseSchema).events()\n",
    "\n",
    "basecut_signal = config['main_analysis_tree']['selection']\n",
    "arr['real-signal'].maskdict = {}\n",
    "arr['real-signal'].maskdict['base'] = eval_expr(arr['real-signal'], basecut_signal)\n",
    "\n",
    "if 'prep-3-2' not in stored_branches_interm:\n",
    "    stored_branches_interm['prep-3-2'] = get_stored_branches(arr, read_sample_list_map)\n",
    "\n",
    "def extract_bdt_sequence(arr, sl_pxy, wgtstr, bdt_start, bdt_mod_factor=None, do_plot=True):\n",
    "    r\"\"\"Extract the sfBDT sequence for specified pT range, based on the signal/proxy similarity\n",
    "    \n",
    "    Arguments:\n",
    "        arr: awkward array dict as input\n",
    "        sl_pxy: proxy sample list\n",
    "        wgtstr: the weight string applied to proxy samples\n",
    "        bdt_start: starting point of sfBDT for scanning\n",
    "        bdt_mod_factor: modify the sfBDT cut by introducing a exponentially decay term from the tagger\n",
    "    \"\"\"\n",
    "    \n",
    "    ## Edges based on the tagger WPs\n",
    "    edges = [0.] + sorted([rg[0] for rg in config['tagger']['working_points']['range'].values()]) + [1.]\n",
    "    print('edges:', edges)\n",
    "    rat_pxy = {}\n",
    "    bdt_seq = {}\n",
    "\n",
    "    ## Extract the optimal sfBDT and variation cut values for each pT range\n",
    "    for ptmin, ptmax in config['pt_range']['range']:\n",
    "        ptlab = f'pt{ptmin}to{ptmax}'\n",
    "        print('pt range: ', ptmin, ptmax)\n",
    "\n",
    "        ## Calculate the proportion of LP+MP+HP over inclusive tagger score for \"signal jets\"\n",
    "        for sam in ['real-signal']:\n",
    "            if ptlab not in arr[sam].maskdict.keys():\n",
    "                arr[sam].maskdict[ptlab] = eval_expr(arr[sam], f\"({config['main_analysis_tree']['pt_var']}>={ptmin}) & ({config['main_analysis_tree']['pt_var']}<{ptmax})\")\n",
    "        h = get_hist(\n",
    "            concat_array(arr, expr=config['main_analysis_tree']['tagger'], sam_list=['real-signal'], filter_list=['base', ptlab]),\n",
    "            bins=edges,\n",
    "            weights=concat_array(arr, expr=config['main_analysis_tree']['weight'], sam_list=['real-signal'], filter_list=['base', ptlab]),\n",
    "        )\n",
    "        rat_hcc = np.array([h.view().value[-1], sum(h.view().value[1:])]) / sum(h.view().value) ## <LP, LP, MP, TP, LP+MP+TP\n",
    "\n",
    "        ## Calculate the proportion for \"proxy jets\" as sfBDT floats\n",
    "\n",
    "        if config['type'] == 'cc':\n",
    "            pxy_base_sel = f'(fj_x_nbhadrons==0) & (fj_x_nchadrons>=1) & (fj_x_sfBDT>{bdt_start})'\n",
    "        elif config['type'] == 'bb':\n",
    "            pxy_base_sel = f'(fj_x_nbhadrons>=1) & (fj_x_sfBDT>{bdt_start})'\n",
    "        for sam in sl_pxy:\n",
    "            for i in '12':\n",
    "                arr[sam].maskdict[f'fj_{i}_bdt_seq_pxy_base'] = eval_expr(arr[sam], pxy_base_sel.replace('fj_x', f'fj_{i}'))\n",
    "                if f'fj_{i}_{ptlab}' not in arr[sam].maskdict.keys():\n",
    "                    arr[sam].maskdict[f'fj_{i}_{ptlab}'] = eval_expr(arr[sam], f'(fj_{i}_pt>={ptmin}) & (fj_{i}_pt<{ptmax})')\n",
    "        ratios = [[], []]\n",
    "        bdt_scanlist = []\n",
    "        _df = ak.to_pandas(\n",
    "            concat_array_fj12(arr, expr=config['tagger']['var'], sam_list=sl_pxy, filter_list=['fj_x_base', 'fj_x_bdt_seq_pxy_base', f'fj_x_{ptlab}']),\n",
    "            anonymous='tagger',\n",
    "        ) # use pandas dataframe to speed up iterative processing\n",
    "        _df['wgt'] = concat_array_fj12(arr, expr=wgtstr, sam_list=sl_pxy, filter_list=['fj_x_base', 'fj_x_bdt_seq_pxy_base', f'fj_x_{ptlab}'])\n",
    "        _df['sfBDT'] = concat_array_fj12(arr, expr='fj_x_sfBDT', sam_list=sl_pxy, filter_list=['fj_x_base', 'fj_x_bdt_seq_pxy_base', f'fj_x_{ptlab}'])\n",
    "        # =========== plot ===========\n",
    "        if do_plot:\n",
    "            mpl.rcParams['axes.prop_cycle'] = cycler(color=['blue', 'red', 'green', 'violet', 'darkorange', 'black', 'cyan', 'yellow'])\n",
    "            f, ax = plt.subplots(figsize=(12,12))\n",
    "            hep.cms.label(data=False, paper=False, year=year, ax=ax, rlabel=r'%s $fb^{-1}$ (13 TeV)'%lumi[year], fontname='sans-serif')\n",
    "            edge_plot = np.linspace(0, 1, 51)\n",
    "            ## Signal jet hist\n",
    "            h_sig = get_hist(\n",
    "                concat_array(arr, expr=config['main_analysis_tree']['tagger'], sam_list=['real-signal'], filter_list=['base', ptlab]),\n",
    "                bins=edge_plot,\n",
    "                weights=concat_array(arr, expr=config['main_analysis_tree']['weight'], sam_list=['real-signal'], filter_list=['base', ptlab]),\n",
    "            )\n",
    "            plot_hist(h_sig, bins=edge_plot, label=config['main_analysis_tree']['label'], normed=True)\n",
    "            ## Proxy jet hist\n",
    "            for bdt in [0.5, 0.85, 0.9, 0.95]:\n",
    "                h_pxy = get_hist(_df.query(f'sfBDT>{bdt}')['tagger'], bins=edge_plot, weights=_df.query(f'sfBDT>{bdt}')['wgt'])\n",
    "                plot_hist(h_pxy, bins=edge_plot, label=f\"g({config['type']}) (sfBDT>{bdt:.2f})\", normed=True)\n",
    "            ax.set_xlabel(config['tagger']['var'].replace('fj_x_',''), ha='right', x=1.0); ax.set_ylabel('A.U.', ha='right', y=1.0);\n",
    "            ax.set_ylim(bottom=0)\n",
    "            plt.savefig(f\"prep/{store_name}/plots/tagger_shape_comp_{config['pt_range']['name']}__{config['main_analysis_tree']['name']}.pdf\")\n",
    "            plt.savefig(f\"prep/{store_name}/plots/tagger_shape_comp_{config['pt_range']['name']}__{config['main_analysis_tree']['name']}.png\")\n",
    "        # ============================\n",
    "        \n",
    "        bdt_expr = 'sfBDT'\n",
    "        if bdt_mod_factor is not None:\n",
    "            bdt_expr = f'sfBDT + 0.5*exp({bdt_mod_factor}*(tagger-1))'\n",
    "        for bdt in np.arange(bdt_start, 0.999, 0.001):  # loop oversf BDT grid\n",
    "            _df = _df.query(f'{bdt_expr}>{bdt}')\n",
    "            h = get_hist(_df['tagger'].values, bins=edges, weights=_df['wgt'].values)\n",
    "            rat = np.array([h.view().value[-1], sum(h.view().value[1:])]) / sum(h.view().value) ## <LP, LP, MP, TP, LP+MP+TP\n",
    "            if len(ratios[1]) > 0 and rat[1] < ratios[1][-1] and bdt > 0.996:\n",
    "                break\n",
    "            rat_pxy[((ptmin,ptmax), np.round(bdt,3))] = rat\n",
    "            for j in range(2):\n",
    "                ratios[j].append(rat[j])\n",
    "            bdt_scanlist.append(bdt)\n",
    "\n",
    "        ## Get sfBDT cut WP\n",
    "        from scipy.interpolate import interp1d\n",
    "        bdt_wp = interp1d(ratios[1], bdt_scanlist, fill_value=\"extrapolate\")(rat_hcc[1]) # chosen BDT WP: proxy proportion under LP+MP+TP reaches signal\n",
    "        bdt_wp_hi = interp1d(ratios[0], bdt_scanlist, fill_value=\"extrapolate\")(rat_hcc[0]) # chosen BDT WP (for 4/5's upper bound): proxy proportion under TP reaches signal\n",
    "        print('all WP:\\nsfBDT scan list:', bdt_scanlist[0], '->', bdt_scanlist[-1], 'proxy prop:', ratios[1][0], '->', ratios[1][-1], 'signal prop:', rat_hcc[1], 'interp bdt:', bdt_wp)\n",
    "        print('HP:    \\nsfBDT scan list:', bdt_scanlist[0], '->', bdt_scanlist[-1], 'proxy prop:', ratios[0][0], '->', ratios[0][-1], 'signal prop:', rat_hcc[0], 'interp bdt:', bdt_wp_hi)\n",
    "        rat_wp, rat_wp_hi = rat_hcc[1], interp1d(bdt_scanlist, ratios[1], fill_value=\"extrapolate\")(bdt_wp_hi) # corresponding LP+MP+TP proportion\n",
    "        step = (rat_wp_hi - rat_wp) / 4\n",
    "        rat_seq = np.linspace(rat_wp-step*5, rat_wp+step*5, 11) # derive an arithmetic sequence\n",
    "        bdt_seq[(ptmin,ptmax)] = interp1d(ratios[1], bdt_scanlist, fill_value=\"extrapolate\")(rat_seq)\n",
    "        if bdt_seq[(ptmin,ptmax)][-1] >= 1.0:\n",
    "            raise RuntimeError('The derived sfBDT sequence has values exceeded 1.0')\n",
    "        print('BDT seq: ', bdt_seq[(ptmin,ptmax)])\n",
    "\n",
    "    arr[f\"bdt_seq_{config['pt_range']['name']}__{config['main_analysis_tree']['name']}\"] = bdt_seq\n",
    "    arr[f\"rat_pxy_{config['pt_range']['name']}__{config['main_analysis_tree']['name']}\"] = rat_pxy\n",
    "    arr[f\"bdt_mod_factor_{config['pt_range']['name']}__{config['main_analysis_tree']['name']}\"] = bdt_mod_factor\n",
    "\n",
    "## ================ Tune this parameter if necessary ================\n",
    "## Alternatively set the sfBDT expression to sfBDT + 0.5*exp(bdt_mod_factor*(tagger-1)) in order to \n",
    "## extract a possible set of BDT sequence without *going out of interpolation range*\n",
    "bdt_mod_factor = 70\n",
    "# bdt_mod_factor = None\n",
    "# ===================================================================\n",
    "extract_bdt_sequence(arr, sl_pxy=['subst_'+s for s in read_sample_list_map if s not in ['qcd-herwig-noht', 'qcd-mg-bflav-noht', 'jetht-noht']],\n",
    "                     wgtstr='genWeight*xsecWeight*puWeight*fj_x_htwgt', bdt_start=0.5, bdt_mod_factor=bdt_mod_factor)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "stored_branches = backup_array(store_name, stored_branches_interm['prep-3-2'], read_sample_list_map, global_key_only=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### ================ 3-3. Extract the sfBDT>0.5 binned fractor: stored as variable \"sfbdtwgt_g50\"; similar for herwig ===================\n",
    "\n",
    "if 'prep-3-3' not in stored_branches_interm:\n",
    "    stored_branches_interm['prep-3-3'] = get_stored_branches(arr, read_sample_list_map)\n",
    "\n",
    "def extract_further_sfbdt_weight(arr, sl_rwgt, wgtstr_rwgt, wgtname, rwgt_info, sl_ext_rwgt=[], presel=''):\n",
    "    r\"\"\"Extract the \"MC substitute to data\" reweight factor (both overall and binned factor) further on sfBDT variable, after a sfBDT>0.9 selection\n",
    "    \n",
    "    Arguments:\n",
    "        arr: awkward array dict as input\n",
    "        sl_rwgt: sample list for MC substitute in this reweighting routine\n",
    "        wgtstr_rwgt: the weight string applied to MC to produce the histogram in this reweighting routine\n",
    "        wgtname: the reweight name (the binned factors) stored as a new column\n",
    "        rwgt_info: info of the reweight variable, in the format of (var, nbin, xmin, xmax) or (var, edges list, None, None)\n",
    "        sl_ext_rwgt: extra sample list for which we also calculate the reweight factors after extracting them\n",
    "        presel: pre-selection before reweighting\n",
    "    \"\"\"\n",
    "    \n",
    "    ## Initially fill the output column with 0, since we will fill the column iteratively for each pT bin\n",
    "    for sam in sl_rwgt + sl_ext_rwgt:\n",
    "        for i in '12':\n",
    "            arr[sam][wgtname.replace('fj_x', f'fj_{i}')] = ak.zeros_like(arr[sam]['ht'])\n",
    "    \n",
    "    ## Reweight based on given variable\n",
    "    rwgt_var, nbin, xmin, xmax = rwgt_info\n",
    "    if not isinstance(nbin, int):\n",
    "        rwgt_edge, xmin, xmax, nbin = nbin, min(nbin), max(nbin), len(nbin)\n",
    "    else:\n",
    "        rwgt_edge = np.linspace(xmin, xmax, nbin+1)\n",
    "    print('rwgt info: ', rwgt_var, rwgt_edge)\n",
    "    \n",
    "    ## Rewight separately on jet pT bins\n",
    "    ent_data, ent_mc, rwgt = {}, {}, {}\n",
    "    for pt_range in config['pt_range']['range']:\n",
    "        pt_range = tuple(pt_range)\n",
    "        ptlab = f'pt{pt_range[0]}to{pt_range[1]}'\n",
    "        rwgt_presel = f'(fj_x_pt>={pt_range[0]}) & (fj_x_pt<{pt_range[1]})'\n",
    "        for sam in sl_rwgt + sl_ext_rwgt + ['jetht-noht']:\n",
    "            for i in '12':\n",
    "                if f'fj_{i}_{ptlab}' not in arr[sam].maskdict.keys():\n",
    "                    arr[sam].maskdict[f'fj_{i}_{ptlab}'] = eval_expr(arr[sam], rwgt_presel.replace('fj_x', f'fj_{i}'))\n",
    "        \n",
    "        ## Requires the selection sfBDT>0.9 which is used in the fit region\n",
    "        if presel != '':\n",
    "            for sam in sl_rwgt + sl_ext_rwgt + ['jetht-noht']:\n",
    "                for i in '12':\n",
    "                    arr[sam].maskdict[f'_tmp_fj_{i}_presel'] = eval_expr(arr[sam], presel.replace('fj_x', f'fj_{i}'))\n",
    "\n",
    "        ## Get data and MC histogram. Note: consider underflow & overflow bins, hence len = nbins+2\n",
    "        ## does not distinguish jet1 or jet2 on this reweighting\n",
    "        filter_list = ['fj_x_base', f'fj_x_{ptlab}']+([f'_tmp_fj_{i}_presel'] if presel != '' else [])\n",
    "        ent_data[pt_range] = get_hist(concat_array_fj12(arr, expr=rwgt_var, sam_list=['jetht-noht'], filter_list=filter_list),\n",
    "                            bins=rwgt_edge, \n",
    "                            weights=np.ones(np.sum(mask_and_fj12(arr['jetht-noht'], mask_list=filter_list))), \n",
    "                            underflow=True, overflow=True, mergeflowbin=False).view(flow=True).value\n",
    "        ent_mc[pt_range]   = get_hist(concat_array_fj12(arr, expr=rwgt_var, sam_list=sl_rwgt, filter_list=filter_list),\n",
    "                            bins=rwgt_edge,\n",
    "                            weights=concat_array_fj12(arr, expr=wgtstr_rwgt, sam_list=sl_rwgt, filter_list=filter_list),\n",
    "                            underflow=True, overflow=True, mergeflowbin=False).view(flow=True).value\n",
    "        ## Calculate the reweight factor\n",
    "        rwgt[pt_range] = np.nan_to_num(ent_data[pt_range] / ent_mc[pt_range], nan=0) # len=nbin+2\n",
    "        print (ent_data[pt_range], rwgt[pt_range])\n",
    "\n",
    "        ## Assign the reweight factor to the new column\n",
    "        for sam in sl_rwgt + sl_ext_rwgt:\n",
    "            for i in '12':\n",
    "                _var = rwgt_var.replace('fj_x', f'fj_{i}')\n",
    "                _wgtname = wgtname.replace('fj_x', f'fj_{i}')\n",
    "                _mask = mask_and(arr[sam], mask_list=[f.replace('fj_x', f'fj_{i}') for f in filter_list])\n",
    "                arr[sam][_wgtname] = arr[sam][_wgtname] + ak.fill_none(calc_rwgt_akarray(arr[sam][_var].mask[_mask], rwgt_edge, rwgt[pt_range]), 0)\n",
    "        \n",
    "    ## Store reweight factors\n",
    "    import pickle\n",
    "    if not os.path.exists(f'prep/{store_name}/plots'):\n",
    "        os.makedirs(f'prep/{store_name}/plots')\n",
    "    with open(f'prep/{store_name}/plots/{wgtname}_{year}.pickle', 'wb') as fw:\n",
    "        pickle.dump({'ent_data':ent_data, 'ent_mc':ent_mc, 'rwgt':rwgt}, fw)\n",
    "    \n",
    "    return {'ent_data':ent_data, 'ent_mc':ent_mc, 'rwgt':rwgt}\n",
    "\n",
    "## Calculate two sets of reweight factor: one for the MG sample list and another for Herwig sample list\n",
    "extract_further_sfbdt_weight(arr, sl_rwgt=['subst_'+s for s in read_sample_list_map if s not in ['qcd-herwig-noht', 'qcd-mg-bflav-noht', 'jetht-noht']], sl_ext_rwgt=['subst_qcd-mg-bflav-noht'] if config['samples']['use_bflav'] else [], \n",
    "                             wgtstr_rwgt = f\"{lumi[year]}*genWeight*xsecWeight*puWeight*fj_x_htwgt\",\n",
    "                             wgtname='fj_x_sfbdtwgt_g50', rwgt_info=('fj_x_sfBDT', 25, 0.5, 1.))\n",
    "if not omit_herwig:\n",
    "    extract_further_sfbdt_weight(arr, sl_rwgt=['subst_'+s for s in read_sample_list_map if s not in ['qcd-mg-noht', 'qcd-mg-bflav-noht', 'jetht-noht']], sl_ext_rwgt=['subst_qcd-mg-bflav-noht'] if config['samples']['use_bflav'] else [], \n",
    "                                 wgtstr_rwgt = f\"{lumi[year]}*genWeight*xsecWeight*puWeight*fj_x_htwgt_herwig\",\n",
    "                                 wgtname='fj_x_sfbdtwgt_g50_herwig', rwgt_info=('fj_x_sfBDT', 25, 0.5, 1.))\n",
    "\n",
    "## Store new variables\n",
    "stored_branches = backup_array(store_name, stored_branches_interm['prep-3-3'], read_sample_list_map)\n",
    "\n",
    "## Test output\n",
    "ak.to_pandas(arr['subst_qcd-mg-noht'][['fj_1_pt', 'fj_1_sfBDT', 'fj_1_sfbdtwgt_g50']][arr['subst_qcd-mg-noht'].maskdict['fj_1_base']])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### ================ 3-4. [additional] Reweight MC subsitute to data on pT: stored as variable \"ad_ptwgt\", \"ad_ptwgt_herwig\" ===================\n",
    "\n",
    "if 'prep-3-4' not in stored_branches_interm:\n",
    "    stored_branches_interm['prep-3-4'] = get_stored_branches(arr, read_sample_list_map)\n",
    "\n",
    "def extract_mc_to_data_pt_weight(arr, sl_rwgt, wgtstr_rwgt, wgtname, sl_ext_rwgt=[]):\n",
    "    r\"\"\"Extract the \"MC subsisute to data\" reweight factor on pT as a optional choice\n",
    "    \n",
    "    Arguments:\n",
    "        arr: awkward array dict as input\n",
    "        sl_rwgt: sample list for MC substitue in this reweighting routine\n",
    "        wgtstr_rwgt: the weight string applied to MC to produce the histogram in this reweighting routine\n",
    "        wgtname: the reweight name stored as a new column\n",
    "        sl_ext_rwgt: extra sample list for which we also calculate the reweight factors after extracting them\n",
    "    \"\"\"\n",
    "    \n",
    "    # Apply simple 1D reweight to pT\n",
    "    rwgt_var, nbin, xmin, xmax  = 'fj_x_pt', 20, 200., 1200.\n",
    "    rwgt_edge = np.linspace(xmin, xmax, nbin+1)\n",
    "    \n",
    "    ## Rewight separately on 1st/2nd jet\n",
    "    for i, lab in zip(['1','2'], ['jet1','jet2']):\n",
    "        ## Get data and MC histogram. Note: consider underflow & overflow bins, hence len = nbins+2\n",
    "        ## Previously this extra factor is extracted with a presel of sfBDT>0.9. Now given that the sfBDT is optimized by pT range thus not fixed, we relax this cut\n",
    "        ent_data = get_hist(concat_array(arr, expr=rwgt_var.replace('fj_x', f'fj_{i}'), sam_list=['jetht-noht'], filter_list=[f'fj_{i}_base']),\n",
    "                            bins=rwgt_edge, \n",
    "                            weights=np.ones(np.sum(mask_and(arr['jetht-noht'], mask_list=[f'fj_{i}_base']))), \n",
    "                            underflow=True, overflow=True, mergeflowbin=False).view(flow=True).value\n",
    "        ent_mc   = get_hist(concat_array(arr, expr=rwgt_var.replace('fj_x', f'fj_{i}'), sam_list=sl_rwgt, filter_list=[f'fj_{i}_base']),\n",
    "                            bins=rwgt_edge,\n",
    "                            weights=concat_array(arr, expr=wgtstr_rwgt, sam_list=sl_rwgt, filter_list=[f'fj_{i}_base']),\n",
    "                            underflow=True, overflow=True, mergeflowbin=False).view(flow=True).value\n",
    "        ## Calculate the reweight factor\n",
    "        rwgt = np.nan_to_num(ent_data / ent_mc, nan=0) # len=nbin+2\n",
    "        print (ent_data, rwgt)\n",
    "        \n",
    "        ## assign the reweight factor to the new column\n",
    "        for sam in sl_rwgt + sl_ext_rwgt:\n",
    "            _var = rwgt_var.replace('fj_x', f'fj_{i}')\n",
    "            _wgtname = wgtname.replace('fj_x', f'fj_{i}')\n",
    "            _mask = mask_and(arr[sam], mask_list=[f'fj_{i}_base'])\n",
    "            arr[sam][_wgtname] = calc_rwgt_akarray(arr[sam][_var].mask[_mask], rwgt_edge, rwgt)  ## fill the new column directly as a masked array\n",
    "#             print('midpoint: ', sam, _wgtname, arr[sam][_wgtname])\n",
    "        \n",
    "## Calculate two sets of reweight factor: one for the MG sample list and another for Herwig sample list\n",
    "extract_mc_to_data_pt_weight(arr, sl_rwgt=['subst_'+s for s in read_sample_list_map if s not in ['qcd-herwig-noht', 'qcd-mg-bflav-noht', 'jetht-noht']], sl_ext_rwgt=['subst_qcd-mg-bflav-noht'] if config['samples']['use_bflav'] else [],\n",
    "                             wgtstr_rwgt = f\"{lumi[year]}*genWeight*xsecWeight*puWeight\", wgtname='fj_x_ad_ptwgt')\n",
    "if not omit_herwig:\n",
    "    extract_mc_to_data_pt_weight(arr, sl_rwgt=['subst_'+s for s in read_sample_list_map if s not in ['qcd-mg-noht', 'qcd-mg-bflav-noht', 'jetht-noht']], sl_ext_rwgt=['subst_qcd-mg-bflav-noht'] if config['samples']['use_bflav'] else [],\n",
    "                                 wgtstr_rwgt = f\"{lumi[year]}*genWeight*xsecWeight*puWeight\", wgtname='fj_x_ad_ptwgt_herwig')\n",
    "\n",
    "## Store new variables\n",
    "stored_branches = backup_array(store_name, stored_branches_interm['prep-3-4'], read_sample_list_map)\n",
    "\n",
    "## Test output\n",
    "ak.to_pandas(arr['subst_qcd-mg-noht'][['ht', 'fj_1_pt', 'fj_1_htwgt', 'fj_1_sfbdtwgt_g50', 'fj_1_ad_ptwgt']][arr['subst_qcd-mg-noht'].maskdict['fj_1_base']])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### ================ 3-5. [additional] Reweight MC (proxy jet) to H->cc signal jet on either mass/pT/tau21: stored as variable \"(mass|pt|tau21)datamcwgt\"; similar for herwig  ===================\n",
    "\n",
    "# First load the h->cc signal ntuple. Adopt the selction used in the analysis\n",
    "import re\n",
    "arr['real-signal'] = NanoEventsFactory.from_root(config['main_analysis_tree']['path'].replace('$YEAR', str(year)), treepath='/'+config['main_analysis_tree']['treename'], schemaclass=BaseSchema).events()\n",
    "\n",
    "basecut_signal = config['main_analysis_tree']['selection']\n",
    "arr['real-signal'].maskdict = {}\n",
    "arr['real-signal'].maskdict['base'] = eval_expr(arr['real-signal'], basecut_signal)\n",
    "\n",
    "if 'prep-3-5' not in stored_branches_interm:\n",
    "    stored_branches_interm['prep-3-5'] = get_stored_branches(arr, read_sample_list_map)\n",
    "\n",
    "def extract_mc_to_signal_weight(arr, sl_rwgt, wgtstr_rwgt, wgtname, rwgt_info, sl_ext_rwgt=[]):\n",
    "    r\"\"\"Extract the \"MC subsisute (proxy) to H->cc signal jet\" reweight factor on possible variable\n",
    "    \n",
    "    Arguments:\n",
    "        arr: awkward array dict as input\n",
    "        sl_rwgt: sample list for MC substitue in this reweighting routine\n",
    "        wgtstr_rwgt: the weight string applied to MC to produce the histogram in this reweighting routine\n",
    "        wgtname: the reweight name stored as a new column\n",
    "        rwgt_info: variable and binning info for this reweighting routine\n",
    "        sl_ext_rwgt: extra sample list for which we also calculate the reweight factors after extracting them\n",
    "    \"\"\"\n",
    "    \n",
    "    # Reweight info extracted from the function argument\n",
    "    rwgt_var, nbin, xmin, xmax, rwgt_var_nom  = rwgt_info\n",
    "    print('rwgt info: ', rwgt_var, nbin, xmin, xmax)\n",
    "    rwgt_edge = np.linspace(xmin, xmax, nbin+1)\n",
    "    \n",
    "    ## Requires the selection sfBDT>0.9 which is (averagely) used in the fit region\n",
    "    rwgt_sel = 'fj_x_sfBDT>0.9'\n",
    "    for sam in sl_rwgt + sl_ext_rwgt:\n",
    "        for i in '12':\n",
    "            arr[sam].maskdict[f'_tmp_fj_{i}_presel'] = eval_expr(arr[sam], rwgt_sel.replace('fj_x', f'fj_{i}'))\n",
    "        \n",
    "    ## Get MC and h->cc signal histogram. Note: consider underflow & overflow bins, hence len = nbins+2\n",
    "    wgt_mc = concat_array_fj12(arr, expr=wgtstr_rwgt, sam_list=sl_rwgt, filter_list=['fj_x_base', '_tmp_fj_x_presel'])\n",
    "    yield_mc = wgt_mc.sum()\n",
    "    ent_mc  = get_hist(concat_array_fj12(arr, expr=rwgt_var, sam_list=sl_rwgt, filter_list=['fj_x_base', '_tmp_fj_x_presel']),\n",
    "                       bins=rwgt_edge,\n",
    "                       weights=wgt_mc,\n",
    "                       underflow=True, overflow=True, mergeflowbin=False).view(flow=True).value      \n",
    "    \n",
    "    wgt_hcc = concat_array(arr, expr=config['main_analysis_tree']['weight'], sam_list=['real-signal'], filter_list=['base'])\n",
    "    yield_hcc = wgt_hcc.sum()\n",
    "    ent_hcc = get_hist(concat_array(arr, expr=rwgt_var_nom, sam_list=['real-signal'], filter_list=['base']),\n",
    "                       bins=rwgt_edge,\n",
    "                       weights=wgt_hcc,\n",
    "                       underflow=True, overflow=True, mergeflowbin=False).view(flow=True).value\n",
    "    \n",
    "    ## Calculate the reweight factors for the two normalized histograms, and clip to (0, 50)\n",
    "    rwgt = np.nan_to_num((ent_hcc/yield_hcc) / (ent_mc/yield_mc), nan=0) # len=nbin+2\n",
    "    rwgt = np.clip(rwgt, 0, 50)\n",
    "    print (ent_hcc, rwgt)\n",
    "\n",
    "    ## assign the reweight factor to the new column (to both MC and data)\n",
    "    for sam in sl_rwgt + ['jetht-noht']:\n",
    "        for i in '12':\n",
    "            _var = rwgt_var.replace('fj_x', f'fj_{i}')\n",
    "            _wgtname = wgtname.replace('fj_x', f'fj_{i}')\n",
    "            _mask = mask_and(arr[sam], mask_list=[f'fj_{i}_base'])\n",
    "            arr[sam][_wgtname] = calc_rwgt_akarray(arr[sam][_var].mask[_mask], rwgt_edge, rwgt)  ## fill the new column directly as a masked array\n",
    "#             print('midpoint: ', sam, _wgtname, arr[sam][_wgtname])\n",
    "    \n",
    "## For each reweight variable, calculate two sets of reweight factor: one for the MG sample list and another for Herwig sample list\n",
    "extract_mc_to_signal_weight(arr, sl_rwgt=['subst_'+s for s in read_sample_list_map if s not in ['qcd-herwig-noht', 'qcd-mg-bflav-noht', 'jetht-noht']], sl_ext_rwgt=['subst_qcd-mg-bflav-noht'] if config['samples']['use_bflav'] else [], \n",
    "                            wgtstr_rwgt = f\"{lumi[year]}*genWeight*xsecWeight*puWeight*fj_x_htwgt\",\n",
    "                            wgtname='fj_x_massdatamcwgt', rwgt_info=('fj_x_sdmass', 15, 50, 200, config['main_analysis_tree']['addition_var']['mass']))\n",
    "extract_mc_to_signal_weight(arr, sl_rwgt=['subst_'+s for s in read_sample_list_map if s not in ['qcd-herwig-noht', 'qcd-mg-bflav-noht', 'jetht-noht']], sl_ext_rwgt=['subst_qcd-mg-bflav-noht'] if config['samples']['use_bflav'] else [], \n",
    "                            wgtstr_rwgt = f\"{lumi[year]}*genWeight*xsecWeight*puWeight*fj_x_htwgt\",\n",
    "                            wgtname='fj_x_ptdatamcwgt', rwgt_info=('fj_x_pt', 20, 200, 1200, config['main_analysis_tree']['pt_var']))\n",
    "extract_mc_to_signal_weight(arr, sl_rwgt=['subst_'+s for s in read_sample_list_map if s not in ['qcd-herwig-noht', 'qcd-mg-bflav-noht', 'jetht-noht']], sl_ext_rwgt=['subst_qcd-mg-bflav-noht'] if config['samples']['use_bflav'] else [], \n",
    "                            wgtstr_rwgt = f\"{lumi[year]}*genWeight*xsecWeight*puWeight*fj_x_htwgt\",\n",
    "                            wgtname='fj_x_tau21datamcwgt', rwgt_info=('fj_x_tau21', 20, 0, 1, config['main_analysis_tree']['addition_var']['tau21']))\n",
    "if not omit_herwig:\n",
    "    extract_mc_to_signal_weight(arr, sl_rwgt=['subst_'+s for s in read_sample_list_map if s not in ['qcd-mg-noht', 'qcd-mg-bflav-noht', 'jetht-noht']], sl_ext_rwgt=['subst_qcd-mg-bflav-noht'] if config['samples']['use_bflav'] else [], \n",
    "                                wgtstr_rwgt = f\"{lumi[year]}*genWeight*xsecWeight*puWeight*fj_x_htwgt_herwig\",\n",
    "                                wgtname='fj_x_massdatamcwgt_herwig', rwgt_info=('fj_x_sdmass', 15, 50, 200, config['main_analysis_tree']['addition_var']['mass']))\n",
    "    extract_mc_to_signal_weight(arr, sl_rwgt=['subst_'+s for s in read_sample_list_map if s not in ['qcd-mg-noht', 'qcd-mg-bflav-noht', 'jetht-noht']], sl_ext_rwgt=['subst_qcd-mg-bflav-noht'] if config['samples']['use_bflav'] else [], \n",
    "                                wgtstr_rwgt = f\"{lumi[year]}*genWeight*xsecWeight*puWeight*fj_x_htwgt_herwig\",\n",
    "                                wgtname='fj_x_ptdatamcwgt_herwig', rwgt_info=('fj_x_pt', 20, 200, 1200, config['main_analysis_tree']['pt_var']))\n",
    "    extract_mc_to_signal_weight(arr, sl_rwgt=['subst_'+s for s in read_sample_list_map if s not in ['qcd-mg-noht', 'qcd-mg-bflav-noht', 'jetht-noht']], sl_ext_rwgt=['subst_qcd-mg-bflav-noht'] if config['samples']['use_bflav'] else [], \n",
    "                                wgtstr_rwgt = f\"{lumi[year]}*genWeight*xsecWeight*puWeight*fj_x_htwgt_herwig\",\n",
    "                                wgtname='fj_x_tau21datamcwgt_herwig', rwgt_info=('fj_x_tau21', 20, 0, 1, config['main_analysis_tree']['addition_var']['tau21']))\n",
    "\n",
    "## Store new variables\n",
    "stored_branches = backup_array(store_name, stored_branches_interm['prep-3-5'], read_sample_list_map)\n",
    "\n",
    "## Test output\n",
    "ak.to_pandas(arr['jetht-noht'][['fj_1_sdmass', 'fj_1_massdatamcwgt', 'fj_1_pt', 'fj_1_ptdatamcwgt', 'fj_1_tau21', 'fj_1_tau21datamcwgt']][arr['jetht-noht'].maskdict['fj_1_base']])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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
}