proj1_TH2D.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 57,
   "id": "e3202201",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-05-25T20:34:05.018929Z",
     "start_time": "2022-05-25T20:33:52.450864Z"
    },
    "code_folding": [
     0
    ]
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Applying ATLAS style settings...\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# load style\n",
    "from ROOT import gROOT\n",
    "\n",
    "path = './../../'\n",
    "# path = '/eos/user/y/youyou/'\n",
    "gROOT.SetMacroPath(path+\"atlasstyle/\")\n",
    "gROOT.LoadMacro(\"AtlasLabels.C\")\n",
    "gROOT.LoadMacro(\"AtlasStyle.C\")\n",
    "gROOT.LoadMacro(\"AtlasUtils.C\")\n",
    "from ROOT import SetAtlasStyle, ATLASLabel, myText, myMarkerText\n",
    "SetAtlasStyle()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "id": "b23ad328",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-05-25T20:34:05.028921Z",
     "start_time": "2022-05-25T20:34:05.028890Z"
    },
    "code_folding": []
   },
   "outputs": [],
   "source": [
    "# load lib\n",
    "import ROOT\n",
    "import os, sys\n",
    "import numpy as np\n",
    "from array import array\n",
    "import pandas as pd\n",
    "\n",
    "# import pandas as pd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "id": "08bd1a85",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-05-25T20:34:05.034068Z",
     "start_time": "2022-05-25T20:34:05.034038Z"
    },
    "code_folding": []
   },
   "outputs": [],
   "source": [
    "# load files\n",
    "\n",
    "# standard model\n",
    "sm_file = path + 'EFT_MINI_TREES/SM.root'\n",
    "sm_data = ROOT.TFile(sm_file,'READ')\n",
    "sm_tree = sm_data.vbswy_tree\n",
    "\n",
    "# O_INT & O_QUAD\n",
    "eft_path = path + 'EFT_MINI_TREES/'\n",
    "files = [(eft_path+i) for i in os.listdir(eft_path)]\n",
    "files.sort()\n",
    "files.remove(sm_file)\n",
    "int_files = [i for i in files if 'INT' in i]\n",
    "quad_files = [i for i in files if 'QUAD' in i]\n",
    "\n",
    "# names of operators used\n",
    "operators = [i.split('/')[-1].split('_')[0] for i in int_files]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "id": "d098d90e",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-05-25T20:34:05.039588Z",
     "start_time": "2022-05-25T20:34:05.039560Z"
    },
    "code_folding": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'weight': 0.3747520756248783,\n",
       " 'jj_drap': 8.317871998076633,\n",
       " 'jj_m': 7583.18091078032,\n",
       " 'jj_pt': 2272.6573001713664,\n",
       " 'jj_dphi_signed': 3.1413330361247063,\n",
       " 'jj_truth_drap': 8.637828046803866,\n",
       " 'jj_truth_m': 7121.392450608517,\n",
       " 'jj_truth_pt': 2208.1405294690107,\n",
       " 'jj_truth_dphi_signed': 3.141589105129242,\n",
       " 'gam_pt': 1818.910625,\n",
       " 'gam_eta': 2.3698120117187536,\n",
       " 'gam_phi': 3.1414341926574707,\n",
       " 'gam_e': 3555.0673493398303,\n",
       " 'gam_truth_pt': 1848.5566727951737,\n",
       " 'gam_truth_eta': 2.369648674150898,\n",
       " 'gam_truth_phi': 3.141542155187217,\n",
       " 'gam_truth_e': 3457.849,\n",
       " 'lep_pt': 1340.063375,\n",
       " 'lep_eta': 2.476752758026128,\n",
       " 'lep_phi': 3.1415679454803467,\n",
       " 'lep_e': 2493.646147428611,\n",
       " 'lep_truth_pt': 1309.7969352030134,\n",
       " 'lep_truth_eta': 2.469721069156157,\n",
       " 'lep_truth_phi': 0.0,\n",
       " 'lep_truth_e': 2395.708435971912,\n",
       " 'j0_pt': 2669.738352457746,\n",
       " 'j0_eta': 4.379240989684844,\n",
       " 'j0_phi': 3.14153790473938,\n",
       " 'j0_e': 5312.517503391469,\n",
       " 'j0_truth_pt': 2709.8225,\n",
       " 'j0_truth_eta': 4.4011039733886665,\n",
       " 'j0_truth_phi': 3.141240358352661,\n",
       " 'j0_truth_e': 5265.623284840022,\n",
       " 'j1_pt': 1917.2258217950919,\n",
       " 'j1_eta': 4.408301830291587,\n",
       " 'j1_phi': 3.1415815353393555,\n",
       " 'j1_e': 4947.460295713709,\n",
       " 'j1_truth_pt': 1750.9039999999998,\n",
       " 'j1_truth_eta': 4.41820096969629,\n",
       " 'j1_truth_phi': 3.141525983810425,\n",
       " 'j1_truth_e': 4545.40688252649,\n",
       " 'w_pt': 1341.175420099982,\n",
       " 'w_eta': 2.476753139706757,\n",
       " 'w_phi': 3.1412847953898373,\n",
       " 'w_e': 2497.4805028786313,\n",
       " 'w_mt': 1302.527395824088,\n",
       " 'w_truth_pt': 1310.9090073284285,\n",
       " 'w_truth_eta': 2.469722441197965,\n",
       " 'w_truth_phi': 3.1415738609262793,\n",
       " 'w_truth_e': 2412.1763432439425,\n",
       " 'w_truth_mt': 1093.573649624575,\n",
       " 'wy_pt': 1920.5651275832138,\n",
       " 'wy_phi': 3.1412016252145953,\n",
       " 'wy_m': 2617.0071313392737,\n",
       " 'wy_truth_pt': 1949.086596822526,\n",
       " 'wy_truth_phi': 3.1414484180133275,\n",
       " 'wy_truth_m': 2581.4441913208902,\n",
       " 'dsid': 363272.0,\n",
       " 'pass_baseline_e': 1.0,\n",
       " 'pass_baseline_mu': 1.0}"
      ]
     },
     "execution_count": 67,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\n",
    "# list of branches we interested in \n",
    "names = [g.GetName() for g in sm_tree.GetListOfBranches()]\n",
    "truth_branch = [i for i in names if 'truth' in i]\n",
    "truth_branch = list(dict.fromkeys(truth_branch))\n",
    "\n",
    "is_reconst = lambda i: not(('truth' in i) or ('weight' in i) or ('pass_baseline' in i))\n",
    "reconst_branch = [i for i in names if is_reconst(i)]\n",
    "reconst_branch = list(dict.fromkeys(reconst_branch))\n",
    "\n",
    "# randomly generate an arr of bins from 0 to 16440,\n",
    "# with increasing bin size\n",
    "bins = np.linspace(1.2,1.2,40).cumprod()\n",
    "bins = np.insert(bins, 0, 0)\n",
    "\n",
    "# generate max value dictionary for sm tree\n",
    "max_sm_dict = {}\n",
    "for branch in names:\n",
    "    max_sm_dict[branch] = sm_tree.GetMaximum(branch)\n",
    "    \n",
    "max_sm_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "id": "ccbfde65",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([  0.        ,   1.05      ,   1.1025    ,   1.157625  ,\n",
       "         1.21550625,   1.27628156,   1.34009564,   1.40710042,\n",
       "         1.47745544,   1.55132822,   1.62889463,   1.71033936,\n",
       "         1.79585633,   1.88564914,   1.9799316 ,   2.07892818,\n",
       "         2.18287459,   2.29201832,   2.40661923,   2.5269502 ,\n",
       "         2.65329771,   2.78596259,   2.92526072,   3.07152376,\n",
       "         3.22509994,   3.38635494,   3.55567269,   3.73345632,\n",
       "         3.92012914,   4.1161356 ,   4.32194238,   4.53803949,\n",
       "         4.76494147,   5.00318854,   5.25334797,   5.51601537,\n",
       "         5.79181614,   6.08140694,   6.38547729,   6.70475115,\n",
       "         7.03998871,   7.39198815,   7.76158756,   8.14966693,\n",
       "         8.55715028,   8.98500779,   9.43425818,   9.90597109,\n",
       "        10.40126965,  10.92133313,  11.46739979,  12.04076978,\n",
       "        12.64280826,  13.27494868,  13.93869611,  14.63563092,\n",
       "        15.36741246,  16.13578309,  16.94257224,  17.78970085,\n",
       "        18.67918589,  19.61314519,  20.59380245,  21.62349257,\n",
       "        22.7046672 ,  23.83990056,  25.03189559,  26.28349037,\n",
       "        27.59766488,  28.97754813,  30.42642554,  31.94774681,\n",
       "        33.54513415,  35.22239086,  36.9835104 ,  38.83268592,\n",
       "        40.77432022,  42.81303623,  44.95368804,  47.20137244,\n",
       "        49.56144107,  52.03951312,  54.64148878,  57.37356322,\n",
       "        60.24224138,  63.25435344,  66.41707112,  69.73792467,\n",
       "        73.22482091,  76.88606195,  80.73036505,  84.7668833 ,\n",
       "        89.00522747,  93.45548884,  98.12826328, 103.03467645,\n",
       "       108.18641027, 113.59573078, 119.27551732, 125.23929319,\n",
       "       131.50125785])"
      ]
     },
     "execution_count": 61,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cut_bins = np.linspace(1.05,1.05,100).cumprod()\n",
    "cut_bins = np.insert(cut_bins, 0, 0)\n",
    "cut_bins"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "id": "3c998019",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-05-25T20:34:05.049427Z",
     "start_time": "2022-05-25T20:34:05.049401Z"
    },
    "code_folding": [
     7,
     11
    ]
   },
   "outputs": [],
   "source": [
    "def plot_2_branch(operator, branches, c=1, isTruth=True, cut_sm=True, isPDF=False):\n",
    "    '''\n",
    "    used default bins and assume all is log(no consider min)(x log didn't test islogy(lin50))\n",
    "    '''\n",
    "    \n",
    "    # load files\n",
    "    i = operators.index(operator)\n",
    "    int_data = ROOT.TFile(int_files[i],'READ')\n",
    "    int_tree = int_data.vbswy_tree\n",
    "    quad_data = ROOT.TFile(quad_files[i], 'READ')\n",
    "    quad_tree = quad_data.vbswy_tree\n",
    "\n",
    "    def generate_output_path():\n",
    "        folder = 'reconstructed/'\n",
    "        if isTruth:\n",
    "            folder = 'truth/'\n",
    "        format = 'png'\n",
    "        if isPDF:\n",
    "            format = 'pdf'\n",
    "        folder = './output_'+format+'/'+folder\n",
    "        if not os.path.exists(folder):\n",
    "            os.mkdir(folder)\n",
    "        if cut_sm:\n",
    "            return (folder+operator+'_cut_2D.'+format)\n",
    "        return (folder+operator+'_2D.'+format)\n",
    "\n",
    "    # ready to plot\n",
    "    nvars = len(branches)   # number of variables to plot\n",
    "    max_dict = {}\n",
    "    for branch in branches:\n",
    "        max_dict[branch] = max(int_tree.GetMaximum(branch), quad_tree.GetMaximum(branch), max_sm_dict[branch])\n",
    "\n",
    "    can = ROOT.TCanvas('can1', '',4000,4000)\n",
    "    can.Divide(nvars-1, nvars-1)\n",
    "\n",
    "    # start plot\n",
    "    for b in range(nvars):\n",
    "        for a in range(b):\n",
    "            ican = (b-1)*(nvars-1)+a+1  # use parts under diagonal of can\n",
    "            can.cd(ican).SetRightMargin(0)\n",
    "            can.cd(ican).SetLeftMargin(0)\n",
    "            can.cd(ican).SetTopMargin(0)\n",
    "            can.cd(ican).SetBottomMargin(0)\n",
    "            if a == 0:\n",
    "                can.cd(ican).SetLeftMargin(.15)\n",
    "            elif a == b-1:\n",
    "                can.cd(ican).SetRightMargin(.15)\n",
    "            if b == nvars-1:\n",
    "                can.cd(ican).SetBottomMargin(.15)      \n",
    "            \n",
    "            # variable name\n",
    "            var_x = branches[ a ]\n",
    "            var_y = branches[ b ]\n",
    "\n",
    "            can.cd(ican).SetLogz(True)\n",
    "            set_log_xy = lambda r: (r.split('_')[0] in ['lep', 'w', 'wy']) and not(isTruth)\n",
    "            can.cd(ican).SetLogx(set_log_xy(var_x))\n",
    "            can.cd(ican).SetLogy(set_log_xy(var_y))\n",
    "\n",
    "            if cut_sm:\n",
    "                divider_x = bins[-1] / (max_dict[var_x] - max_sm_dict[var_x])\n",
    "                divider_y = bins[-1] / (max_dict[var_y] - max_sm_dict[var_y])\n",
    "                bins_x = (bins/divider_x) + max_sm_dict[var_x]\n",
    "                bins_y = (bins/divider_y) + max_sm_dict[var_y]\n",
    "            else:\n",
    "                divider_x = bins[-1] / max_dict[var_x]\n",
    "                divider_y = bins[-1] / max_dict[var_y]\n",
    "                bins_x = (bins/divider_x)\n",
    "                bins_y = (bins/divider_y)\n",
    "\n",
    "            h1 = ROOT.TH2F(var_x+var_y, \";\"+var_x+\";\"+var_y, len(bins_x)-1, array('d',bins_x), len(bins_y)-1, array('d',bins_y))\n",
    "\n",
    "            if isTruth:\n",
    "                for i in range(sm_tree.GetEntries()):\n",
    "                    sm_tree.GetEntry(i)\n",
    "                    h1.Fill(getattr(sm_tree, var_x), getattr(sm_tree, var_y), sm_tree.weight)\n",
    "                for i in range(int_tree.GetEntries()):\n",
    "                    int_tree.GetEntry(i)\n",
    "                    h1.Fill(getattr(int_tree, var_x), getattr(int_tree, var_y), int_tree.weight*c)\n",
    "                for i in range(quad_tree.GetEntries()):\n",
    "                    quad_tree.GetEntry(i)\n",
    "                    h1.Fill(getattr(quad_tree, var_x), getattr(quad_tree, var_y), quad_tree.weight*c**2)\n",
    "            else:\n",
    "                for i in range(sm_tree.GetEntries()):\n",
    "                    sm_tree.GetEntry(i)\n",
    "                    tmp = int(sm_tree.pass_baseline_e) | int(sm_tree.pass_baseline_mu)\n",
    "                    if tmp:\n",
    "                        h1.Fill(getattr(sm_tree, var_x), getattr(sm_tree, var_y), sm_tree.weight)\n",
    "                for i in range(int_tree.GetEntries()):\n",
    "                    int_tree.GetEntry(i)\n",
    "                    tmp = int(int_tree.pass_baseline_e) | int(int_tree.pass_baseline_mu)\n",
    "                    if tmp:\n",
    "                        h1.Fill(getattr(int_tree, var_x), getattr(int_tree, var_y), int_tree.weight*c)\n",
    "                for i in range(quad_tree.GetEntries()):\n",
    "                    quad_tree.GetEntry(i)\n",
    "                    tmp = int(quad_tree.pass_baseline_e) | int(quad_tree.pass_baseline_mu)\n",
    "                    if tmp:\n",
    "                        h1.Fill(getattr(quad_tree, var_x), getattr(quad_tree, var_y), quad_tree.weight*c**2)\n",
    "            h1.Draw()\n",
    "            cr = h1.GetCorrelationFactor(1,2)\n",
    "            #can.CreateText('correlation coefficient is %.3f \\n'%cr)\n",
    "            '''str = 'correlation coefficient is %.3f \\n'%cr+var_x\n",
    "            h1.GetXaxis().SetTitle(str)'''\n",
    "            h1.DrawClone('Colz')\n",
    "            t = ROOT.TLatex()\n",
    "            t.SetNDC()\n",
    "            t.SetTextColor( 1 )\n",
    "            t.SetTextSize( 0.05 )\n",
    "            t.SetTextAlign( 12 )\n",
    "            t.DrawLatex( 0.2, 0.95, 'correlation coefficient is %.3f'%cr )\n",
    "    can.cd()\n",
    "    t = ROOT.TLatex()\n",
    "    t.SetNDC()\n",
    "    t.SetTextColor( 1 )\n",
    "    t.SetTextSize( 0.08 )\n",
    "    t.SetTextAlign( 12 )\n",
    "    t.DrawLatex( 0.65, 0.75, operator )\n",
    "\n",
    "    can.Update()\n",
    "    can.Draw()\n",
    "   \n",
    "    can.Print('test.png')\n",
    "    can.SaveAs(generate_output_path())\n",
    "\n",
    "    #can.Clear()\n",
    "    int_data.Close()\n",
    "    quad_data.Close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "id": "3c998019",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-05-25T20:34:05.049427Z",
     "start_time": "2022-05-25T20:34:05.049401Z"
    },
    "code_folding": [
     7,
     11
    ]
   },
   "outputs": [],
   "source": [
    "def plot_2_branch_cut_half(operator, branches, c=1, isTruth=True, isPDF=False):\n",
    "    '''\n",
    "    only y axis is cut\n",
    "    '''\n",
    "    \n",
    "    # load files\n",
    "    i = operators.index(operator)\n",
    "    int_data = ROOT.TFile(int_files[i],'READ')\n",
    "    int_tree = int_data.vbswy_tree\n",
    "    quad_data = ROOT.TFile(quad_files[i], 'READ')\n",
    "    quad_tree = quad_data.vbswy_tree\n",
    "\n",
    "    def generate_output_path():\n",
    "        folder = 'reconstructed/'\n",
    "        if isTruth:\n",
    "            folder = 'truth/'\n",
    "        format = 'png'\n",
    "        if isPDF:\n",
    "            format = 'pdf'\n",
    "        folder = './output_'+format+'/'+folder\n",
    "        if not os.path.exists(folder):\n",
    "            os.mkdir(folder)\n",
    "        return (folder+operator+'_2D.'+format)\n",
    "\n",
    "    # ready to plot\n",
    "    nvars = len(branches)   # number of variables to plot\n",
    "    max_dict = {}\n",
    "    for branch in branches:\n",
    "        max_dict[branch] = max(int_tree.GetMaximum(branch), quad_tree.GetMaximum(branch), max_sm_dict[branch])\n",
    "\n",
    "    can = ROOT.TCanvas('can1', '',4000,4000)\n",
    "    can.Divide(nvars-1, nvars-1)\n",
    "\n",
    "    # start plot\n",
    "    ican = 0\n",
    "    for b in range(nvars):\n",
    "        for a in range(nvars):\n",
    "            if a==b:\n",
    "                continue\n",
    "            ican += 1\n",
    "            can.cd(ican).SetRightMargin(0)\n",
    "            can.cd(ican).SetLeftMargin(0)\n",
    "            can.cd(ican).SetTopMargin(0)\n",
    "            can.cd(ican).SetBottomMargin(0)\n",
    "            if a == 0:\n",
    "                can.cd(ican).SetLeftMargin(.15)\n",
    "            elif a == nvars-1:\n",
    "                can.cd(ican).SetRightMargin(.15)\n",
    "            if b == nvars-1:\n",
    "                can.cd(ican).SetBottomMargin(.15)      \n",
    "            \n",
    "            # variable name\n",
    "            var_x = branches[ a ]\n",
    "            var_y = branches[ b ]\n",
    "\n",
    "            can.cd(ican).SetLogz(True)\n",
    "            set_log_xy = lambda r: (r.split('_')[0] in ['lep', 'w', 'wy']) and not(isTruth)\n",
    "            can.cd(ican).SetLogx(set_log_xy(var_x))\n",
    "            can.cd(ican).SetLogy(set_log_xy(var_y))\n",
    "        \n",
    "            divider_y = bins[-1] / (max_dict[var_y] - max_sm_dict[var_y])\n",
    "            bins_y = (bins/divider_y) + max_sm_dict[var_y]\n",
    "            divider_x = bins[-1] / max_dict[var_x]\n",
    "            bins_x = (bins/divider_x)\n",
    "\n",
    "            h1 = ROOT.TH2F(var_x+var_y, \";\"+var_x+'_cut'+\";\"+var_y, len(bins_x)-1, array('d',bins_x), len(bins_y)-1, array('d',bins_y))\n",
    "\n",
    "            if isTruth:\n",
    "                for i in range(sm_tree.GetEntries()):\n",
    "                    sm_tree.GetEntry(i)\n",
    "                    h1.Fill(getattr(sm_tree, var_x), getattr(sm_tree, var_y), sm_tree.weight)\n",
    "                for i in range(int_tree.GetEntries()):\n",
    "                    int_tree.GetEntry(i)\n",
    "                    h1.Fill(getattr(int_tree, var_x), getattr(int_tree, var_y), int_tree.weight*c)\n",
    "                for i in range(quad_tree.GetEntries()):\n",
    "                    quad_tree.GetEntry(i)\n",
    "                    h1.Fill(getattr(quad_tree, var_x), getattr(quad_tree, var_y), quad_tree.weight*c**2)\n",
    "            else:\n",
    "                for i in range(sm_tree.GetEntries()):\n",
    "                    sm_tree.GetEntry(i)\n",
    "                    tmp = int(sm_tree.pass_baseline_e) | int(sm_tree.pass_baseline_mu)\n",
    "                    if tmp:\n",
    "                        h1.Fill(getattr(sm_tree, var_x), getattr(sm_tree, var_y), sm_tree.weight)\n",
    "                for i in range(int_tree.GetEntries()):\n",
    "                    int_tree.GetEntry(i)\n",
    "                    tmp = int(int_tree.pass_baseline_e) | int(int_tree.pass_baseline_mu)\n",
    "                    if tmp:\n",
    "                        h1.Fill(getattr(int_tree, var_x), getattr(int_tree, var_y), int_tree.weight*c)\n",
    "                for i in range(quad_tree.GetEntries()):\n",
    "                    quad_tree.GetEntry(i)\n",
    "                    tmp = int(quad_tree.pass_baseline_e) | int(quad_tree.pass_baseline_mu)\n",
    "                    if tmp:\n",
    "                        h1.Fill(getattr(quad_tree, var_x), getattr(quad_tree, var_y), quad_tree.weight*c**2)\n",
    "            h1.Draw()\n",
    "            cr = h1.GetCorrelationFactor(1,2)\n",
    "            h1.DrawClone('Colz')\n",
    "            t = ROOT.TLatex()\n",
    "            t.SetNDC()\n",
    "            t.SetTextColor( 1 )\n",
    "            t.SetTextSize( 0.05 )\n",
    "            t.SetTextAlign( 12 )\n",
    "            t.DrawLatex( 0.2, 0.95, 'correlation coefficient is %.3f'%cr )\n",
    "    can.cd()\n",
    "    t = ROOT.TLatex()\n",
    "    t.SetNDC()\n",
    "    t.SetTextColor( 1 )\n",
    "    t.SetTextSize( 0.08 )\n",
    "    t.SetTextAlign( 12 )\n",
    "    t.DrawLatex( 0.65, 0.75, operator )\n",
    "\n",
    "    can.Update()\n",
    "    can.Draw()\n",
    "   \n",
    "    can.Print('test.png')\n",
    "    can.SaveAs(generate_output_path())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 73,
   "id": "3c998019",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-05-25T20:34:05.049427Z",
     "start_time": "2022-05-25T20:34:05.049401Z"
    },
    "code_folding": [
     7,
     11
    ]
   },
   "outputs": [],
   "source": [
    "def plot_2_branch_rank(operator, branches, c=1, isTruth=True, cut_sm=True, isPDF=True):\n",
    "    '''\n",
    "    ranked by correlation\n",
    "    '''\n",
    "    \n",
    "    # load files\n",
    "    i = operators.index(operator)\n",
    "    int_data = ROOT.TFile(int_files[i],'READ')\n",
    "    int_tree = int_data.vbswy_tree\n",
    "    quad_data = ROOT.TFile(quad_files[i], 'READ')\n",
    "    quad_tree = quad_data.vbswy_tree\n",
    "\n",
    "    # dictionary of maximum of each variables\n",
    "    max_dict = {}\n",
    "    for branch in branches:\n",
    "        max_dict[branch] = max(int_tree.GetMaximum(branch), quad_tree.GetMaximum(branch), max_sm_dict[branch])\n",
    "\n",
    "    hs = []\n",
    "    # add histogram\n",
    "    for b in range(len(branches)):\n",
    "        for a in range(b):    \n",
    "            \n",
    "            # variable name\n",
    "            var_x = branches[ a ]\n",
    "            var_y = branches[ b ]\n",
    "\n",
    "            if cut_sm:\n",
    "                divider_x = bins[-1] / (max_dict[var_x] - max_sm_dict[var_x])\n",
    "                divider_y = bins[-1] / (max_dict[var_y] - max_sm_dict[var_y])\n",
    "                bins_x = (bins/divider_x) + max_sm_dict[var_x]\n",
    "                bins_y = (bins/divider_y) + max_sm_dict[var_y]\n",
    "            else:\n",
    "                divider_x = bins[-1] / max_dict[var_x]\n",
    "                divider_y = bins[-1] / max_dict[var_y]\n",
    "                bins_x = (bins/divider_x)\n",
    "                bins_y = (bins/divider_y)\n",
    "\n",
    "            h1 = ROOT.TH2F(var_x+var_y, \";\"+var_x+\";\"+var_y, len(bins_x)-1, array('d',bins_x), len(bins_y)-1, array('d',bins_y))\n",
    "\n",
    "            if isTruth:\n",
    "                for i in range(sm_tree.GetEntries()):\n",
    "                    sm_tree.GetEntry(i)\n",
    "                    h1.Fill(getattr(sm_tree, var_x), getattr(sm_tree, var_y), sm_tree.weight)\n",
    "                for i in range(int_tree.GetEntries()):\n",
    "                    int_tree.GetEntry(i)\n",
    "                    h1.Fill(getattr(int_tree, var_x), getattr(int_tree, var_y), int_tree.weight*c)\n",
    "                for i in range(quad_tree.GetEntries()):\n",
    "                    quad_tree.GetEntry(i)\n",
    "                    h1.Fill(getattr(quad_tree, var_x), getattr(quad_tree, var_y), quad_tree.weight*c**2)\n",
    "            else:\n",
    "                for i in range(sm_tree.GetEntries()):\n",
    "                    sm_tree.GetEntry(i)\n",
    "                    tmp = int(sm_tree.pass_baseline_e) | int(sm_tree.pass_baseline_mu)\n",
    "                    if tmp:\n",
    "                        h1.Fill(getattr(sm_tree, var_x), getattr(sm_tree, var_y), sm_tree.weight)\n",
    "                for i in range(int_tree.GetEntries()):\n",
    "                    int_tree.GetEntry(i)\n",
    "                    tmp = int(int_tree.pass_baseline_e) | int(int_tree.pass_baseline_mu)\n",
    "                    if tmp:\n",
    "                        h1.Fill(getattr(int_tree, var_x), getattr(int_tree, var_y), int_tree.weight*c)\n",
    "                for i in range(quad_tree.GetEntries()):\n",
    "                    quad_tree.GetEntry(i)\n",
    "                    tmp = int(quad_tree.pass_baseline_e) | int(quad_tree.pass_baseline_mu)\n",
    "                    if tmp:\n",
    "                        h1.Fill(getattr(quad_tree, var_x), getattr(quad_tree, var_y), quad_tree.weight*c**2)\n",
    "            h1.Draw()\n",
    "            hs += [h1]\n",
    "\n",
    "    hs = sorted(hs, key=lambda x: x.GetCorrelationFactor(1,2), reverse=True)\n",
    "\n",
    "    #ready to plot\n",
    "    nvars = len(branches)   # number of variables to plot\n",
    "    can = ROOT.TCanvas('can1', '',4000,4000)\n",
    "    nplots = int(np.ceil(((nvars-1) * nvars/2) **0.5))\n",
    "    can.Divide(nplots, nplots)\n",
    "\n",
    "    ican = 1\n",
    "    for h in hs:\n",
    "        can.cd(ican).SetRightMargin(0)\n",
    "        can.cd(ican).SetLeftMargin(0.15)\n",
    "        can.cd(ican).SetTopMargin(0)\n",
    "        can.cd(ican).SetBottomMargin(0.15)\n",
    "        if (ican % nplots == 0) or (ican == len(hs)):\n",
    "            can.cd(ican).SetRightMargin(.15)\n",
    " \n",
    "        can.cd(ican).SetLogz(True)\n",
    "        can.cd(ican).SetLogx(1)\n",
    "        can.cd(ican).SetLogy(1)\n",
    "        h.DrawClone('Colz')\n",
    "        t = ROOT.TLatex()\n",
    "        t.SetNDC()\n",
    "        t.SetTextColor( 1 )\n",
    "        t.SetTextSize( 0.05 )\n",
    "        t.SetTextAlign( 12 )\n",
    "        t.DrawLatex( 0.2, 0.95, 'correlation coefficient is %.3f'%h.GetCorrelationFactor(1,2) )\n",
    "        ican += 1\n",
    "\n",
    "    can.Update()\n",
    "    can.Draw()\n",
    "    #can.Print('test.png')\n",
    "\n",
    "    def generate_output_path():\n",
    "        folder = 'reconstructed/'\n",
    "        if isTruth:\n",
    "            folder = 'truth/'\n",
    "        format = 'png'\n",
    "        if isPDF:\n",
    "            format = 'pdf'\n",
    "        folder = './output_'+format+'/'+folder\n",
    "        if not os.path.exists(folder):\n",
    "            os.mkdir(folder)\n",
    "        if cut_sm:\n",
    "            return (folder+operator+'_cut_2D_ranked.'+format)\n",
    "        return (folder+operator+'_2D_ranked.'+format)\n",
    "    can.SaveAs(generate_output_path())\n",
    "\n",
    "    int_data.Close()\n",
    "    quad_data.Close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "id": "0af96c8b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "M truth completed\n",
      "T truth completed\n",
      "M recon completed\n",
      "T recon completed\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Info in <TCanvas::Print>: pdf file ./output_pdf/truth/M2_cut_2D_ranked.pdf has been created\n",
      "Info in <TCanvas::Print>: pdf file ./output_pdf/truth/M5_cut_2D_ranked.pdf has been created\n",
      "Info in <TCanvas::Print>: pdf file ./output_pdf/truth/T1_cut_2D_ranked.pdf has been created\n",
      "Info in <TCanvas::Print>: pdf file ./output_pdf/truth/T6_cut_2D_ranked.pdf has been created\n",
      "Info in <TCanvas::Print>: pdf file ./output_pdf/truth/T5_cut_2D_ranked.pdf has been created\n",
      "Info in <TCanvas::Print>: pdf file ./output_pdf/reconstructed/M2_cut_2D_ranked.pdf has been created\n",
      "Info in <TCanvas::Print>: pdf file ./output_pdf/reconstructed/M5_cut_2D_ranked.pdf has been created\n",
      "Info in <TCanvas::Print>: pdf file ./output_pdf/reconstructed/T1_cut_2D_ranked.pdf has been created\n",
      "Info in <TCanvas::Print>: pdf file ./output_pdf/reconstructed/T5_cut_2D_ranked.pdf has been created\n",
      "Info in <TCanvas::Print>: pdf file ./output_pdf/reconstructed/T6_cut_2D_ranked.pdf has been created\n"
     ]
    }
   ],
   "source": [
    "M_t = ['gam_truth_e', 'gam_truth_pt', 'j0_truth_pt', 'jj_truth_pt', 'lep_truth_pt', 'w_truth_mt','w_truth_pt','wy_truth_m', 'wy_truth_pt']\n",
    "plot_2_branch_rank('M2', M_t, isTruth=True, isPDF=True)\n",
    "M_t = ['gam_truth_e', 'gam_truth_pt', 'lep_truth_pt', 'w_truth_mt','w_truth_pt', 'wy_truth_m', 'wy_truth_pt']\n",
    "plot_2_branch_rank('M5', M_t, isTruth=True, isPDF=True)\n",
    "print('M truth completed')\n",
    "\n",
    "T_t = ['gam_truth_e', 'gam_truth_pt', 'j0_truth_pt', 'jj_truth_pt', 'lep_truth_e', 'lep_truth_pt', 'w_truth_e', 'w_truth_mt', 'w_truth_pt', 'wy_truth_m', 'wy_truth_pt']\n",
    "plot_2_branch_rank('T1', T_t, isTruth=True, isPDF=True)\n",
    "plot_2_branch_rank('T6', T_t, isTruth=True, isPDF=True)\n",
    "T_t = ['gam_truth_e', 'gam_truth_pt', 'lep_truth_e', 'lep_truth_pt', 'w_truth_e', 'w_truth_mt', 'w_truth_pt', 'wy_truth_m', 'wy_truth_pt']\n",
    "plot_2_branch_rank('T5', T_t, isTruth=True, isPDF=True)\n",
    "print('T truth completed')\n",
    "\n",
    "M_r = ['gam_e', 'gam_pt', 'j0_pt', 'jj_pt', 'lep_e', 'lep_pt', 'w_e', 'w_mt', 'w_pt', 'wy_m', 'wy_pt']\n",
    "plot_2_branch_rank('M2', M_r, isTruth=False, isPDF=True)\n",
    "M_r = ['gam_e', 'gam_pt', 'j0_pt', 'jj_pt', 'lep_pt', 'w_mt', 'w_pt', 'wy_m', 'wy_pt']\n",
    "plot_2_branch_rank('M5', M_r, isTruth=False, isPDF=True)\n",
    "print('M recon completed')\n",
    "\n",
    "T_r = ['gam_e', 'gam_pt', 'lep_e', 'lep_pt', 'w_e', 'w_mt', 'w_pt', 'wy_m', 'wy_pt']\n",
    "plot_2_branch_rank('T1', T_r, isTruth=False, isPDF=True)\n",
    "plot_2_branch_rank('T5', T_r, isTruth=False, isPDF=True)\n",
    "plot_2_branch_rank('T6', T_r, isTruth=False, isPDF=True)\n",
    "print('T recon completed')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 71,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "M truth completed\n",
      "T truth completed\n",
      "M recon completed\n",
      "T recon completed\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Warning in <TCanvas::Constructor>: Deleting canvas with same name: can1\n",
      "Info in <TCanvas::Print>: png file test.png has been created\n",
      "Info in <TCanvas::Print>: pdf file ./output_pdf/truth/M2_cut_2D.pdf has been created\n",
      "Info in <TCanvas::Print>: png file test.png has been created\n",
      "Info in <TCanvas::Print>: pdf file ./output_pdf/truth/M5_cut_2D.pdf has been created\n",
      "Info in <TCanvas::Print>: png file test.png has been created\n",
      "Info in <TCanvas::Print>: pdf file ./output_pdf/truth/T1_cut_2D.pdf has been created\n",
      "Info in <TCanvas::Print>: png file test.png has been created\n",
      "Info in <TCanvas::Print>: pdf file ./output_pdf/truth/T6_cut_2D.pdf has been created\n",
      "Info in <TCanvas::Print>: png file test.png has been created\n",
      "Info in <TCanvas::Print>: pdf file ./output_pdf/truth/T5_cut_2D.pdf has been created\n",
      "Info in <TCanvas::Print>: png file test.png has been created\n",
      "Info in <TCanvas::Print>: pdf file ./output_pdf/reconstructed/M2_cut_2D.pdf has been created\n",
      "Info in <TCanvas::Print>: png file test.png has been created\n",
      "Info in <TCanvas::Print>: pdf file ./output_pdf/reconstructed/M5_cut_2D.pdf has been created\n",
      "Info in <TCanvas::Print>: png file test.png has been created\n",
      "Info in <TCanvas::Print>: pdf file ./output_pdf/reconstructed/T1_cut_2D.pdf has been created\n",
      "Info in <TCanvas::Print>: png file test.png has been created\n",
      "Info in <TCanvas::Print>: pdf file ./output_pdf/reconstructed/T5_cut_2D.pdf has been created\n",
      "Info in <TCanvas::Print>: png file test.png has been created\n",
      "Info in <TCanvas::Print>: pdf file ./output_pdf/reconstructed/T6_cut_2D.pdf has been created\n"
     ]
    }
   ],
   "source": [
    "M_t = ['gam_truth_e', 'gam_truth_pt', 'j0_truth_pt', 'jj_truth_pt', 'lep_truth_pt', 'w_truth_mt','w_truth_pt','wy_truth_m', 'wy_truth_pt']\n",
    "plot_2_branch('M2', M_t, isTruth=True, isPDF=True)\n",
    "M_t = ['gam_truth_e', 'gam_truth_pt', 'lep_truth_pt', 'w_truth_mt','w_truth_pt', 'wy_truth_m', 'wy_truth_pt']\n",
    "plot_2_branch('M5', M_t, isTruth=True, isPDF=True)\n",
    "print('M truth completed')\n",
    "\n",
    "T_t = ['gam_truth_e', 'gam_truth_pt', 'j0_truth_pt', 'jj_truth_pt', 'lep_truth_e', 'lep_truth_pt', 'w_truth_e', 'w_truth_mt', 'w_truth_pt', 'wy_truth_m', 'wy_truth_pt']\n",
    "plot_2_branch('T1', T_t, isTruth=True, isPDF=True)\n",
    "plot_2_branch('T6', T_t, isTruth=True, isPDF=True)\n",
    "T_t = ['gam_truth_e', 'gam_truth_pt', 'lep_truth_e', 'lep_truth_pt', 'w_truth_e', 'w_truth_mt', 'w_truth_pt', 'wy_truth_m', 'wy_truth_pt']\n",
    "plot_2_branch('T5', T_t, isTruth=True, isPDF=True)\n",
    "print('T truth completed')\n",
    "\n",
    "M_r = ['gam_e', 'gam_pt', 'j0_pt', 'jj_pt', 'lep_e', 'lep_pt', 'w_e', 'w_mt', 'w_pt', 'wy_m', 'wy_pt']\n",
    "plot_2_branch('M2', M_r, isTruth=False, isPDF=True)\n",
    "M_r = ['gam_e', 'gam_pt', 'j0_pt', 'jj_pt', 'lep_pt', 'w_mt', 'w_pt', 'wy_m', 'wy_pt']\n",
    "plot_2_branch('M5', M_r, isTruth=False, isPDF=True)\n",
    "print('M recon completed')\n",
    "\n",
    "T_r = ['gam_e', 'gam_pt', 'lep_e', 'lep_pt', 'w_e', 'w_mt', 'w_pt', 'wy_m', 'wy_pt']\n",
    "plot_2_branch('T1', T_r, isTruth=False, isPDF=True)\n",
    "plot_2_branch('T5', T_r, isTruth=False, isPDF=True)\n",
    "plot_2_branch('T6', T_r, isTruth=False, isPDF=True)\n",
    "print('T recon completed')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9cb2c642",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "65398ec9edf812bfa2a1cd19e30d7cdb06b1eced593dd679351ef973282302ad"
  },
  "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.9.12"
  },
  "varInspector": {
   "cols": {
    "lenName": 16,
    "lenType": 16,
    "lenVar": 40
   },
   "kernels_config": {
    "python": {
     "delete_cmd_postfix": "",
     "delete_cmd_prefix": "del ",
     "library": "var_list.py",
     "varRefreshCmd": "print(var_dic_list())"
    },
    "r": {
     "delete_cmd_postfix": ") ",
     "delete_cmd_prefix": "rm(",
     "library": "var_list.r",
     "varRefreshCmd": "cat(var_dic_list()) "
    }
   },
   "oldHeight": 399.852,
   "position": {
    "height": "40px",
    "left": "1058px",
    "right": "20px",
    "top": "120px",
    "width": "347px"
   },
   "types_to_exclude": [
    "module",
    "function",
    "builtin_function_or_method",
    "instance",
    "_Feature"
   ],
   "varInspector_section_display": "none",
   "window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}