process_ROOT_data_all_jets.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import uproot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.model_selection import train_test_split\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from torch.utils.data import TensorDataset\n",
    "from torch.utils.data import DataLoader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "path_to_data = '../data/'\n",
    "\n",
    "folder15 = 'breynold/user.breynold.data15_13TeV.00284484.physics_Main.DAOD_NTUP_JTRIG_JETM1.r9264_p3083_p3601_j042_tree.root/'\n",
    "file15 = 'user.breynold.18753218._000001.tree.root'\n",
    "folder16 = 'breynold/user.breynold.data16_13TeV.00307656.physics_Main.DAOD_NTUP_JTRIG_JETM1.r9264_p3083_p3601_j042_tree.root/'\n",
    "file16 = 'user.breynold.18797259._000001.tree.root'\n",
    "\n",
    "folder = 'doglioni/'\n",
    "file = 'data18_13TeV.00355331.calibration_DataScouting_05_Jets.merge.AOD.r10657_p3592.root'\n",
    "file = 'dijetISRphoton.root'\n",
    "file = 'dijetISRphoton36k.root'\n",
    "\n",
    "# Load a ROOT file\n",
    "filePath = path_to_data + folder16 + file16\n",
    "#filePath = path_to_data + folder + file\n",
    "ttree = uproot.open(filePath)['outTree']['nominal']\n",
    "#ttree = uproot.open(filePath)['CollectionTree']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[b'runNumber',\n",
       " b'eventNumber',\n",
       " b'lumiBlock',\n",
       " b'coreFlags',\n",
       " b'bcid',\n",
       " b'NPV',\n",
       " b'actualInteractionsPerCrossing',\n",
       " b'averageInteractionsPerCrossing',\n",
       " b'weight_pileup',\n",
       " b'correctedAverageMu',\n",
       " b'correctedAndScaledAverageMu',\n",
       " b'correctedActualMu',\n",
       " b'correctedAndScaledActualMu',\n",
       " b'passL1',\n",
       " b'passHLT',\n",
       " b'passedTriggers',\n",
       " b'disabledTriggers',\n",
       " b'triggerPrescales',\n",
       " b'isPassBits',\n",
       " b'isPassBitsNames',\n",
       " b'nAntiKt4EMTopoJets_Calib2016',\n",
       " b'AntiKt4EMTopoJets_Calib2016_E',\n",
       " b'AntiKt4EMTopoJets_Calib2016_pt',\n",
       " b'AntiKt4EMTopoJets_Calib2016_phi',\n",
       " b'AntiKt4EMTopoJets_Calib2016_eta',\n",
       " b'AntiKt4EMTopoJets_Calib2016_clean_passLooseBad',\n",
       " b'AntiKt4EMTopoJets_Calib2016_clean_passTightBad',\n",
       " b'AntiKt4EMTopoJets_Calib2016_JvtPass_Medium',\n",
       " b'nAntiKt4EMTopoJets_Calib2018',\n",
       " b'AntiKt4EMTopoJets_Calib2018_E',\n",
       " b'AntiKt4EMTopoJets_Calib2018_pt',\n",
       " b'AntiKt4EMTopoJets_Calib2018_phi',\n",
       " b'AntiKt4EMTopoJets_Calib2018_eta',\n",
       " b'AntiKt4EMTopoJets_Calib2018_clean_passLooseBad',\n",
       " b'AntiKt4EMTopoJets_Calib2018_clean_passTightBad',\n",
       " b'AntiKt4EMTopoJets_Calib2018_JvtPass_Medium',\n",
       " b'nAntiKt4EMPFlowJets_Calib2018',\n",
       " b'AntiKt4EMPFlowJets_Calib2018_E',\n",
       " b'AntiKt4EMPFlowJets_Calib2018_pt',\n",
       " b'AntiKt4EMPFlowJets_Calib2018_phi',\n",
       " b'AntiKt4EMPFlowJets_Calib2018_eta',\n",
       " b'AntiKt4EMPFlowJets_Calib2018_clean_passLooseBad',\n",
       " b'AntiKt4EMPFlowJets_Calib2018_clean_passTightBad',\n",
       " b'AntiKt4EMPFlowJets_Calib2018_JvtPass_Medium',\n",
       " b'nL1Jets',\n",
       " b'L1Jet_et8x8',\n",
       " b'L1Jet_eta',\n",
       " b'L1Jet_phi',\n",
       " b'nHLT_a4tcemsubjesFS_Jets',\n",
       " b'HLT_a4tcemsubjesFS_Jets_E',\n",
       " b'HLT_a4tcemsubjesFS_Jets_pt',\n",
       " b'HLT_a4tcemsubjesFS_Jets_phi',\n",
       " b'HLT_a4tcemsubjesFS_Jets_eta',\n",
       " b'nAntiKt10LCTopoTrimmedPtFrac5SmallR20Jets_Calib2018',\n",
       " b'AntiKt10LCTopoTrimmedPtFrac5SmallR20Jets_Calib2018_E',\n",
       " b'AntiKt10LCTopoTrimmedPtFrac5SmallR20Jets_Calib2018_pt',\n",
       " b'AntiKt10LCTopoTrimmedPtFrac5SmallR20Jets_Calib2018_phi',\n",
       " b'AntiKt10LCTopoTrimmedPtFrac5SmallR20Jets_Calib2018_eta']"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ttree.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<JaggedArray [[1 1 1 ... 1 1 1] [1 1 1 ... 1 1 1] [1 1 1 ... 0 1 1] ... [1 1 1 1 1] [1 1 1 ... 1 1 1] [1 1 1 ... 1 1 1]] at 0x00025c011978>"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ttree.array('AntiKt4EMTopoJets_Calib2016_clean_passLooseBad')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<JaggedArray [[105.0 104.0] [154.0 105.0 38.0 26.0 20.0] [279.0 257.0] ... [43.0 13.0] [301.0 112.0 86.0 35.0 17.0] [120.0]] at 0x00025bfd9160>"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ttree.array('L1Jet_et8x8')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<JaggedArray [[-0.4 -0.8] [-2.0 0.2 2.45 -1.2 -1.8] [5.551115e-17 0.6] ... [-3.9 -1.2] [1.8 -0.6 -1.4 -1.4 5.551115e-17] [-0.4]] at 0x00025bfbf9b0>"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ttree.array('L1Jet_eta')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "branchnames = ['nAntiKt4EMTopoJets_Calib2018',\n",
    "               'AntiKt4EMTopoJets_Calib2018_E',\n",
    "               'AntiKt4EMTopoJets_Calib2018_pt',\n",
    "               'AntiKt4EMTopoJets_Calib2018_phi',\n",
    "               'AntiKt4EMTopoJets_Calib2018_eta']\n",
    "\n",
    "\n",
    "#branchnames = ['nDSTrigJet',\n",
    "#               'DSTrigJet_E',\n",
    "#              'DSTrigJet_pt',\n",
    "#              'DSTrigJet_phi',\n",
    "#              'DSTrigJet_eta']\n",
    "\n",
    "#branchnames = ['n_jet',\n",
    "#               'jet_e',\n",
    "#               'jet_pt',\n",
    "#               'jet_phi',\n",
    "#               'jet_eta',\n",
    "#               'n_photons',\n",
    "#               'photons_e',\n",
    "#               'photons_pt',\n",
    "#               'photons_phi',\n",
    "#               'photons_eta']\n",
    "\n",
    "jaggedE = ttree.array(branchnames[1])\n",
    "jaggedpT = ttree.array(branchnames[2])\n",
    "jaggedphi = ttree.array(branchnames[3])\n",
    "jaggedeta = ttree.array(branchnames[4])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "ename": "IndexError",
     "evalue": "list index out of range",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mIndexError\u001b[0m                                Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-6-2b8d4d99e764>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0mjaggedE_ph\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mttree\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0marray\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mbranchnames\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m6\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      2\u001b[0m \u001b[0mjaggedpT_ph\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mttree\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0marray\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mbranchnames\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m7\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      3\u001b[0m \u001b[0mjaggedphi_ph\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mttree\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0marray\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mbranchnames\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m8\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      4\u001b[0m \u001b[0mjaggedeta_ph\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mttree\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0marray\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mbranchnames\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m9\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mIndexError\u001b[0m: list index out of range"
     ]
    }
   ],
   "source": [
    "jaggedE_ph = ttree.array(branchnames[6])\n",
    "jaggedpT_ph = ttree.array(branchnames[7])\n",
    "jaggedphi_ph = ttree.array(branchnames[8])\n",
    "jaggedeta_ph = ttree.array(branchnames[9])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "e_list = []\n",
    "pt_list = []\n",
    "eta_list = []\n",
    "phi_list = []\n",
    "for ii, arr in enumerate(jaggedE):\n",
    "    for kk, val in enumerate(arr):\n",
    "        e_list.append(val)\n",
    "        pt_list.append(jaggedpT[ii][kk])\n",
    "        eta_list.append(jaggedeta[ii][kk])\n",
    "        phi_list.append(jaggedphi[ii][kk])\n",
    "\n",
    "\n",
    "all_e = np.array(e_list)\n",
    "all_pt = np.array(pt_list)\n",
    "all_eta = np.array(eta_list)\n",
    "all_phi = np.array(phi_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(22621866,) (22621866,) (22621866,) (22621866,)\n"
     ]
    }
   ],
   "source": [
    "print(all_e.shape, all_eta.shape, all_phi.shape, all_pt.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_leading(jaggedX):\n",
    "    return jaggedX[jaggedX.counts > 0, 0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "leading_E = get_leading(jaggedE)\n",
    "leading_pT = get_leading(jaggedpT)\n",
    "leading_phi = get_leading(jaggedphi)\n",
    "leading_eta = get_leading(jaggedeta)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "leading_ph_E = get_leading(jaggedE_ph)\n",
    "leading_ph_pT = get_leading(jaggedpT_ph)\n",
    "leading_ph_phi = get_leading(jaggedphi_ph)\n",
    "leading_ph_eta = get_leading(jaggedeta_ph)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(leading_E.shape, leading_eta.shape, leading_phi.shape, leading_pT.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pd.DataFrame(data = {'jet pT': all_pt, 'jet eta': all_eta, 'jet phi': all_phi, 'jet E': all_e})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>jet pT</th>\n",
       "      <th>jet eta</th>\n",
       "      <th>jet phi</th>\n",
       "      <th>jet E</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>0</td>\n",
       "      <td>161.850494</td>\n",
       "      <td>-0.764774</td>\n",
       "      <td>2.287350</td>\n",
       "      <td>212.217300</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1</td>\n",
       "      <td>155.803253</td>\n",
       "      <td>-0.482176</td>\n",
       "      <td>-0.687477</td>\n",
       "      <td>175.156143</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2</td>\n",
       "      <td>29.825865</td>\n",
       "      <td>-1.400844</td>\n",
       "      <td>-0.953364</td>\n",
       "      <td>64.451935</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>3</td>\n",
       "      <td>22.049797</td>\n",
       "      <td>1.565698</td>\n",
       "      <td>-2.697407</td>\n",
       "      <td>55.199814</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>4</td>\n",
       "      <td>21.709501</td>\n",
       "      <td>-0.085313</td>\n",
       "      <td>0.009056</td>\n",
       "      <td>22.383587</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "       jet pT   jet eta   jet phi       jet E\n",
       "0  161.850494 -0.764774  2.287350  212.217300\n",
       "1  155.803253 -0.482176 -0.687477  175.156143\n",
       "2   29.825865 -1.400844 -0.953364   64.451935\n",
       "3   22.049797  1.565698 -2.697407   55.199814\n",
       "4   21.709501 -0.085313  0.009056   22.383587"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df.head()"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "df1 = pd.DataFrame(data = {'jet pT': leading_pT, 'jet eta': leading_eta, 'jet phi': leading_phi, 'jet E': leading_E})"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "df2 = pd.DataFrame(data = {'photon pT': leading_ph_pT, 'photon eta': leading_ph_eta, 'photon phi': leading_ph_phi, 'photon E': leading_ph_E})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df1.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df2.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "train, test = train_test_split(df, test_size=0.2, random_state=41)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(train.shape, test.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save train and test sets\n",
    "#train.to_pickle('processed_data/all_jets_train.pkl')\n",
    "#test.to_pickle('processed_data/all_jets_test.pkl')\n",
    "#train = pd.read_pickle('processed_data/all_jets_train.pkl')\n",
    "#test = pd.read_pickle('processed_data/all_jets_test.pkl')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python3.7 (fastairoot)",
   "language": "python",
   "name": "fastairoot"
  },
  "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.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}