process_ROOT_split_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": 3,
   "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": 4,
   "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": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ttree.keys()"
   ]
  },
  {
   "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": "raw",
   "metadata": {},
   "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": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Remove events with no jets\n",
    "jaggedE = jaggedE[(jaggedE.counts != 0)]\n",
    "jaggedpT = jaggedpT[(jaggedpT.counts != 0)]\n",
    "jaggedphi = jaggedphi[(jaggedphi.counts != 0)]\n",
    "jaggedeta = jaggedeta[(jaggedeta.counts != 0)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(1939900,) (1939900,) (1939900,) (1939900,)\n"
     ]
    }
   ],
   "source": [
    "print(jaggedE.shape, jaggedpT.shape, jaggedphi.shape ,jaggedeta.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get indices that would sort array according to leading jet pT\n",
    "argsort = np.array(jaggedpT[:,0]).argsort()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "jaggedpT = jaggedpT[argsort]\n",
    "jaggedE = jaggedE[argsort]\n",
    "jaggedeta = jaggedeta[argsort]\n",
    "jaggedphi = jaggedphi[argsort]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "7.4080515\n",
      "161.86812\n",
      "161.86838\n",
      "207.1685\n",
      "207.16852\n",
      "279.61777\n",
      "279.6179\n",
      "20721.586\n"
     ]
    }
   ],
   "source": [
    "path_to_save_folder = '../data/all_jets_split_data/'\n",
    "\n",
    "N_datasets = 4\n",
    "split_index = len(jaggedpT) // N_datasets\n",
    "\n",
    "for n in range(N_datasets):\n",
    "    sub_jaggedpT = jaggedpT[n * split_index:(n + 1) * split_index]\n",
    "    sub_jaggedeta = jaggedeta[n * split_index:(n + 1) * split_index]\n",
    "    sub_jaggedphi = jaggedphi[n * split_index:(n + 1) * split_index]\n",
    "    sub_jaggedE = jaggedE[n * split_index:(n + 1) * split_index]\n",
    "    \n",
    "    print(sub_jaggedpT[0][0])\n",
    "    print(sub_jaggedpT[-1][0])\n",
    "\n",
    " #   e_list = []\n",
    "#    pt_list = []\n",
    "   # eta_list = []\n",
    "  #  phi_list = []\n",
    " #   for ii, arr in enumerate(sub_jaggedE):\n",
    "#        for kk, val in enumerate(arr):\n",
    "\n",
    "        #    e_list.append(val)\n",
    "       #     pt_list.append(sub_jaggedpT[ii][kk])\n",
    "      #      eta_list.append(sub_jaggedeta[ii][kk])\n",
    "     #       phi_list.append(sub_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)\n",
    "    \n",
    "    #df = pd.DataFrame({'pT': all_pt, 'eta': all_eta, 'phi': all_phi, 'E': all_e})\n",
    "    #train, test = train_test_split(df, test_size=0.2, random_state=41)\n",
    "    #train.to_pickle(path_to_save_folder + 'sub_train_all_jets_%d' % n)\n",
    "    #test.to_pickle(path_to_save_folder + 'sub_test_all_jets_%d' % n)\n"
   ]
  },
  {
   "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
}