{
"cells": [
{
"cell_type": "code",
"execution_count": 1,
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"/home/hep/davide/miniconda3/envs/root_env/lib/ROOT.py:301: RuntimeWarning: numpy.dtype size changed, may indicate binary incompatibility\n",
" return _orig_ihook( name, *args, **kwds )\n"
]
}
],
"source": [
"import ROOT as r\n",
"import ctypes\n",
"import numpy as np\n",
"from array import array\n",
"import root_numpy as rn\n",
"import matplotlib.pyplot as plt\n",
"import pickle"
]
},
{
"cell_type": "code",
"execution_count": 2,
"metadata": {},
"outputs": [],
"source": [
"l_index = 0\n",
"data_index = None \n",
"mother_index = None"
]
},
{
"cell_type": "code",
"execution_count": 3,
"metadata": {},
"outputs": [],
"source": [
"mother_ID=['Ds','Dplus']\n",
"l_flv = ['e','mu']\n",
"data_type = ['MC','data']\n",
"\n",
"def find_file_path(l_index=l_index, mother_index=mother_index, data_index=data_index): \n",
" return \"/disk/lhcb_data/davide/Rphipi/\"+data_type[data_index]+\"/\"+mother_ID[mother_index]+\"_phipi_\"+l_flv[l_index]+l_flv[l_index]+\"/\"+mother_ID[mother_index]+\"_phipi_\"+l_flv[l_index]+l_flv[l_index]+\".root\"\n",
"\n",
"mother_index=0\n",
"\n",
"data = r.TFile(find_file_path(l_index=l_index, mother_index=mother_index, data_index=1))\n",
"MC_Ds = r.TFile(find_file_path(l_index=l_index, mother_index=mother_index, data_index=0))\n",
"\n",
"tree_name_Ds = mother_ID[mother_index]+'_OfflineTree/DecayTree'\n",
"t_data = data.Get(tree_name_Ds)\n",
"t_MC_Ds = MC_Ds.Get(tree_name_Ds)\n",
"\n",
"\n",
"mother_index=1\n",
"\n",
"MC_Dplus = r.TFile(find_file_path(l_index=l_index, mother_index=mother_index, data_index=0))\n",
"\n",
"tree_name_Dplus = mother_ID[mother_index]+'_OfflineTree/DecayTree'\n",
"t_MC_Dplus = MC_Dplus.Get(tree_name_Dplus)\n"
]
},
{
"cell_type": "code",
"execution_count": 4,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"<ROOT.TTree object (\"DecayTree\") at 0x5644e0a94f60>"
]
},
"execution_count": 4,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"t_MC_Dplus"
]
},
{
"cell_type": "code",
"execution_count": 5,
"metadata": {},
"outputs": [],
"source": [
"def return_branches(mother_index=None):\n",
"\n",
" branches_needed = [\n",
" #________________________________________\n",
" #D\n",
" #________________________________________\n",
" #D Geometric variables, pT and FD\n",
" \n",
" mother_ID[mother_index]+\"_ENDVERTEX_CHI2\",\n",
" mother_ID[mother_index]+\"_ENDVERTEX_NDOF\",\n",
" mother_ID[mother_index]+\"_IPCHI2_OWNPV\",\n",
"\n",
" mother_ID[mother_index]+\"_OWNPV_CHI2\",\n",
" mother_ID[mother_index]+\"_OWNPV_NDOF\",\n",
" mother_ID[mother_index]+\"_IP_OWNPV\",\n",
" mother_ID[mother_index]+\"_DIRA_OWNPV\",\n",
" \n",
" mother_ID[mother_index]+\"_PT\",\n",
" mother_ID[mother_index]+\"_FD_OWNPV\",\n",
" mother_ID[mother_index]+\"_FDCHI2_OWNPV\",\n",
" \n",
" #D Reconstructed mass\n",
" mother_ID[mother_index]+\"_ConsD_M\",\n",
" \n",
" #D Trigger variables\n",
" mother_ID[mother_index]+\"_Hlt1TrackMVADecision_TOS\",\n",
" mother_ID[mother_index]+\"_Hlt2RareCharmD2Pi\"+l_flv[l_index].capitalize()+l_flv[l_index].capitalize()+\"OSDecision_TOS\",\n",
" mother_ID[mother_index]+\"_Hlt2Phys_TOS\",\n",
" \n",
" #________________________________________\n",
" #PHI\n",
" #________________________________________\n",
" #phi geometric variables, pT and FD\n",
" \n",
" \"phi_ENDVERTEX_CHI2\",\n",
" \"phi_ENDVERTEX_NDOF\",\n",
" \"phi_IPCHI2_OWNPV\",\n",
" \n",
" #\"phi_OWNPV_CHI2\",\n",
" #\"phi_OWNPV_NDOF\",\n",
" #\"phi_IP_OWNPV\",\n",
" #\"phi_DIRA_OWNPV\",\n",
" \n",
" \"phi_PT\",\n",
" \n",
" #phi Reconstructed mass\n",
" \n",
" \"phi_M\",\n",
" \n",
" #________________________________________\n",
" #PION\n",
" #________________________________________\n",
" #pi Geometric variables and pT\n",
" #\"pi_OWNPV_CHI2\",\n",
" #\"pi_OWNPV_NDOF\",\n",
" #'pi_IP_OWNPV',\n",
" \n",
" 'pi_PT',\n",
" \n",
" #pi PID variables\n",
" \n",
" \"pi_MC15TuneV1_ProbNNpi\",\n",
" \n",
" #________________________________________\n",
" #LEPTONS\n",
" #________________________________________\n",
" #leptons Geometric variables and pT\n",
" \n",
" #l_flv[l_index]+\"_plus_OWNPV_CHI2\",\n",
" #l_flv[l_index]+\"_plus_OWNPV_NDOF\",\n",
" #l_flv[l_index]+\"_minus_OWNPV_CHI2\",\n",
" #l_flv[l_index]+\"_minus_OWNPV_NDOF\",\n",
" #\n",
" #l_flv[l_index]+\"_plus_IP_OWNPV\",\n",
" #l_flv[l_index]+\"_minus_IP_OWNPV\",\n",
" \n",
" l_flv[l_index]+\"_plus_PT\",\n",
" l_flv[l_index]+\"_minus_PT\",\n",
" \n",
" #leptons PID variables\n",
" \n",
" l_flv[l_index]+\"_plus_MC15TuneV1_ProbNN\"+l_flv[l_index],\n",
" l_flv[l_index]+\"_minus_MC15TuneV1_ProbNN\"+l_flv[l_index],\n",
" \n",
" \n",
" ] \n",
" return branches_needed"
]
},
{
"cell_type": "code",
"execution_count": 6,
"metadata": {},
"outputs": [],
"source": [
"#Switch on only the branches that you need\n",
"t_data.SetBranchStatus(\"*\",0)\n",
"t_MC_Dplus.SetBranchStatus(\"*\",0)\n",
"t_MC_Ds.SetBranchStatus(\"*\",0)\n",
"\n",
"\n",
"for branch in return_branches(mother_index=0):\n",
" t_data.SetBranchStatus(branch, 1)\n",
" t_MC_Ds.SetBranchStatus(branch, 1)\n",
" \n",
"for branch in return_branches(mother_index=1):\n",
" t_data.SetBranchStatus(branch, 1)\n",
" t_MC_Ds.SetBranchStatus(branch, 1)"
]
},
{
"cell_type": "code",
"execution_count": 7,
"metadata": {},
"outputs": [],
"source": [
"#Create a dictionary\n",
"\n",
"#dict ={'branch_name'=[branch_value[event]]}\n",
"\n",
"MC_Ds_tuple_dict = {}\n",
"branches_needed=return_branches(mother_index=0)\n",
"for branch in branches_needed:\n",
" \n",
" MC_Ds_tuple_dict[branch] = rn.root2array(\n",
" \n",
" filenames=find_file_path(l_index, mother_index=0, data_index=0),\n",
" treename = tree_name_Ds,\n",
" branches = branch,\n",
" start=0,\n",
" stop=t_MC_Ds.GetEntries(),\n",
" )\n",
" \n",
"MC_Dplus_tuple_dict = {}\n",
"branches_needed=return_branches(mother_index=1)\n",
"for branch in branches_needed:\n",
" \n",
" MC_Dplus_tuple_dict[branch] = rn.root2array(\n",
" \n",
" filenames=find_file_path(l_index, mother_index=1, data_index=0),\n",
" treename = tree_name_Dplus,\n",
" branches = branch,\n",
" start=0,\n",
" stop=t_MC_Dplus.GetEntries(),\n",
" )\n",
" \n",
"data_tuple_dict = {}\n",
"\n",
"branches_needed=return_branches(mother_index=0)\n",
"for branch in branches_needed:\n",
" \n",
" data_tuple_dict[branch] = rn.root2array(\n",
" \n",
" filenames=find_file_path(l_index,mother_index=0, data_index=1),\n",
" treename = tree_name_Ds,\n",
" branches = branch,\n",
" start=0,\n",
" stop=t_data.GetEntries(),\n",
" )"
]
},
{
"cell_type": "code",
"execution_count": 10,
"metadata": {},
"outputs": [],
"source": [
"#HLT1 PRESELECTION\n",
"data_tuple_dict_presel_1={}\n",
"MC_Ds_tuple_dict_presel_1={}\n",
"MC_Dplus_tuple_dict_presel_1={}\n",
"\n",
"for label in return_branches(mother_index=0): \n",
"\n",
" data_tuple_dict_presel_1[label] = data_tuple_dict[label][data_tuple_dict[\"Ds_Hlt1TrackMVADecision_TOS\"]]\n",
" MC_Ds_tuple_dict_presel_1[label] = MC_Ds_tuple_dict[label][MC_Ds_tuple_dict[\"Ds_Hlt1TrackMVADecision_TOS\"]]\n",
"\n",
"for label in return_branches(mother_index=1): \n",
" MC_Dplus_tuple_dict_presel_1[label] = MC_Dplus_tuple_dict[label][MC_Dplus_tuple_dict[\"Dplus_Hlt1TrackMVADecision_TOS\"]]\n",
"\n",
"#RareCharm D2pi l l HLT2 PRESELECTION\n",
"\n",
"data_tuple_dict_presel_2={}\n",
"MC_Ds_tuple_dict_presel_2={}\n",
"MC_Dplus_tuple_dict_presel_2={}\n",
"\n",
"for label in return_branches(mother_index=0):\n",
"\n",
" data_tuple_dict_presel_2[label] = data_tuple_dict_presel_1[label][data_tuple_dict_presel_1[\"Ds_Hlt2RareCharmD2Pi\"+l_flv[l_index].capitalize()+l_flv[l_index].capitalize()+\"OSDecision_TOS\"]]\n",
" MC_Ds_tuple_dict_presel_2[label] = MC_Ds_tuple_dict_presel_1[label][MC_Ds_tuple_dict_presel_1[\"Ds_Hlt2RareCharmD2Pi\"+l_flv[l_index].capitalize()+l_flv[l_index].capitalize()+\"OSDecision_TOS\"]]\n",
"\n",
"for label in return_branches(mother_index=1): \n",
" \n",
" MC_Dplus_tuple_dict_presel_2[label] = MC_Dplus_tuple_dict_presel_1[label][MC_Dplus_tuple_dict_presel_1[\"Dplus_Hlt2RareCharmD2Pi\"+l_flv[l_index].capitalize()+l_flv[l_index].capitalize()+\"OSDecision_TOS\"]]\n",
"\n",
"#PID preselection\n",
"\n",
"#MC_PID_indices=np.where(MC_tuple_dict_presel_2[l_flv[l_index]+\"_plus_MC15TuneV1_ProbNN\"+l_flv[l_index]]>0.4)\n",
"\n",
"data_PID_indices_plus=np.where(data_tuple_dict_presel_2[l_flv[l_index]+\"_plus_MC15TuneV1_ProbNN\"+l_flv[l_index]]>0.4)\n",
"data_PID_indices_minus=np.where(data_tuple_dict_presel_2[l_flv[l_index]+\"_minus_MC15TuneV1_ProbNN\"+l_flv[l_index]]>0.4)\n",
"data_PID_indices_pi=np.where(data_tuple_dict_presel_2[\"pi_MC15TuneV1_ProbNNpi\"]>0.4)\n",
"\n",
"data_PID_indices = np.intersect1d(data_PID_indices_plus,data_PID_indices_minus)\n",
"data_PID_indices = np.intersect1d(data_PID_indices,data_PID_indices_pi)\n",
"data_tuple_dict_presel_3={}\n",
"MC_Ds_tuple_dict_presel_3={}\n",
"MC_Dplus_tuple_dict_presel_3={}\n",
"\n",
"for label in return_branches(mother_index=0):\n",
"\n",
" data_tuple_dict_presel_3[label] = data_tuple_dict_presel_2[label][data_PID_indices]\n",
" MC_Ds_tuple_dict_presel_3[label] = MC_Ds_tuple_dict_presel_2[label]#[MC_PID_indices]\n",
" \n",
"for label in return_branches(mother_index=1):\n",
" MC_Dplus_tuple_dict_presel_3[label] = MC_Dplus_tuple_dict_presel_2[label]#[MC_PID_indices]"
]
},
{
"cell_type": "code",
"execution_count": 11,
"metadata": {},
"outputs": [],
"source": [
"if l_flv[l_index]=='mu':\n",
" lower_phi_mass = 980\n",
" upper_phi_mass = 1060\n",
" \n",
"if l_flv[l_index]=='e':\n",
" lower_phi_mass = 850\n",
" upper_phi_mass = 1100\n",
" \n",
"#Retrieve mc signal and data bkg events\n",
"\n",
"MC_Dplus_indices=[]\n",
"MC_Ds_indices=[]\n",
"data_indices=[]\n",
" \n",
"for i in range(len(MC_Ds_tuple_dict_presel_3[\"Ds_ConsD_M\"])):\n",
"\n",
" phi_m = MC_Ds_tuple_dict_presel_3[\"phi_M\"][i]\n",
" #fixing a window on the phi mass\n",
" if lower_phi_mass<phi_m<upper_phi_mass:\n",
" MC_Ds_indices.append(i)\n",
" \n",
"for i in range(len(MC_Dplus_tuple_dict_presel_3[\"Dplus_ConsD_M\"])):\n",
"\n",
" phi_m = MC_Dplus_tuple_dict_presel_3[\"phi_M\"][i]\n",
" #fixing a window on the phi mass\n",
" if lower_phi_mass<phi_m<upper_phi_mass:\n",
" MC_Dplus_indices.append(i)\n",
" \n",
"for i in range(len(data_tuple_dict_presel_3[\"Ds_ConsD_M\"])):\n",
"\n",
" phi_m = data_tuple_dict_presel_3[\"phi_M\"][i]\n",
" #fixing a window on the phi mass\n",
" if lower_phi_mass<phi_m<upper_phi_mass:\n",
" data_indices.append(i)\n",
"\n",
"MC_Ds_tuple_dict_presel_4 ={}\n",
"MC_Dplus_tuple_dict_presel_4 ={}\n",
"data_tuple_dict_presel_4={}\n",
"\n",
"\n",
"for label in return_branches(mother_index=0): \n",
" \n",
" data_tuple_dict_presel_4[label] = data_tuple_dict_presel_3[label][data_indices]\n",
" MC_Ds_tuple_dict_presel_4[label] = MC_Ds_tuple_dict_presel_3[label][MC_Ds_indices]\n",
"\n",
"for label in return_branches(mother_index=1): \n",
" MC_Dplus_tuple_dict_presel_4[label] = MC_Dplus_tuple_dict_presel_3[label][MC_Dplus_indices]"
]
},
{
"cell_type": "code",
"execution_count": 13,
"metadata": {},
"outputs": [],
"source": [
"data_tuple_dict=data_tuple_dict_presel_4\n",
"MC_Ds_tuple_dict=MC_Ds_tuple_dict_presel_4\n",
"MC_Dplus_tuple_dict=MC_Dplus_tuple_dict_presel_4"
]
},
{
"cell_type": "code",
"execution_count": 16,
"metadata": {},
"outputs": [],
"source": [
"if l_flv[l_index]=='mu':\n",
" lower_Dplus_mass=1830\n",
" upper_Dplus_mass=1910\n",
"\n",
" lower_Ds_mass = 1930\n",
" upper_Ds_mass = 2010\n",
" \n",
"if l_flv[l_index]=='e':\n",
" lower_Dplus_mass=1810\n",
" upper_Dplus_mass=1900\n",
"\n",
" lower_Ds_mass = 1900\n",
" upper_Ds_mass = 2000\n",
"\n",
"#Retrieve mc signal and data bkg events\n",
"\n",
"data_bkg_indices=[]\n",
"MC_Ds_sig_indices=[]\n",
"MC_Dplus_sig_indices=[]\n",
"\n",
"\n",
"for i in range(len(data_tuple_dict[\"Ds_ConsD_M\"])):\n",
" #retrieving the Ds reconstructed mass\n",
" \n",
" #if 0<data_tuple_dict[\"phi_M\"][i]<lower_phi_mass or upper_phi_mass<data_tuple_dict[\"phi_M\"][i]:\n",
" m = data_tuple_dict[\"Ds_ConsD_M\"][i]\n",
" \n",
" #selecting the out of signal regions\n",
" if 0<m<lower_Dplus_mass or upper_Dplus_mass < m < lower_Ds_mass or upper_Ds_mass < m:\n",
" #if 0<m<lower_Ds_mass or upper_Ds_mass < m:\n",
" data_bkg_indices.append(i)\n",
" \n",
"for i in range(len(MC_Ds_tuple_dict[\"Ds_ConsD_M\"])):\n",
" \n",
" #retrieving the Ds reconstructed mass\n",
" #if lower_phi_mass<MC_tuple_dict[\"phi_M\"][i]<lower_phi_mass:\n",
" m = MC_Ds_tuple_dict[\"Ds_ConsD_M\"][i]\n",
" \n",
" #selecting the signal regions\n",
" #if lower_Dplus_mass< m <upper_Dplus_mass or lower_Ds_mass< m <upper_Ds_mass:\n",
" if lower_Ds_mass< m <upper_Ds_mass:\n",
" MC_Ds_sig_indices.append(i) \n",
" \n",
"for i in range(len(MC_Dplus_tuple_dict[\"Dplus_ConsD_M\"])):\n",
" \n",
" #retrieving the Ds reconstructed mass\n",
" #if lower_phi_mass<MC_tuple_dict[\"phi_M\"][i]<lower_phi_mass:\n",
" m = MC_Dplus_tuple_dict[\"Dplus_ConsD_M\"][i]\n",
" \n",
" #selecting the signal regions\n",
" #if lower_Dplus_mass< m <upper_Dplus_mass or lower_Ds_mass< m <upper_Ds_mass:\n",
" if lower_Dplus_mass< m <upper_Dplus_mass:\n",
" MC_Dplus_sig_indices.append(i)\n",
" \n",
"#Create the dict tuples with all MC signal and data bkg events\n",
"\n",
"data_tuple_bkg={}\n",
"MC_Ds_tuple_sig ={}\n",
"MC_Dplus_tuple_sig ={}\n",
"\n",
"MC_Ds_mass_for_fit= np.array([MC_Ds_tuple_dict[\"Ds_ConsD_M\"][i][0] for i in range(len(MC_Ds_tuple_dict[\"Ds_ConsD_M\"]))])\n",
"MC_Dplus_mass_for_fit=np.array([MC_Dplus_tuple_dict[\"Dplus_ConsD_M\"][i][0] for i in range(len(MC_Dplus_tuple_dict[\"Dplus_ConsD_M\"]))])\n",
"\n",
"for label in return_branches(mother_index=0): \n",
"\n",
" data_tuple_bkg[label] = data_tuple_dict[label][data_bkg_indices]\n",
" MC_Ds_tuple_sig[label] = MC_Ds_tuple_dict[label][MC_Ds_sig_indices]\n",
" \n",
"for label in return_branches(mother_index=1): \n",
" MC_Dplus_tuple_sig[label] = MC_Dplus_tuple_dict[label][MC_Dplus_sig_indices]"
]
},
{
"cell_type": "code",
"execution_count": 28,
"metadata": {},
"outputs": [],
"source": [
"with open('/disk/lhcb_data/davide/Rphipi/NN_for_training/'+l_flv[l_index]+l_flv[l_index]+'/MC_for_NN_Ds_'+l_flv[l_index]+l_flv[l_index]+'.pickle', 'wb') as handle:\n",
" pickle.dump(MC_Ds_tuple_sig, handle, protocol=pickle.HIGHEST_PROTOCOL)\n",
" \n",
"with open('/disk/lhcb_data/davide/Rphipi/NN_for_training/'+l_flv[l_index]+l_flv[l_index]+'/MC_for_NN_Dplus_'+l_flv[l_index]+l_flv[l_index]+'.pickle', 'wb') as handle:\n",
" pickle.dump(MC_Dplus_tuple_sig, handle, protocol=pickle.HIGHEST_PROTOCOL)\n",
"\n",
"with open('/disk/lhcb_data/davide/Rphipi/MC/for_fit/'+l_flv[l_index]+l_flv[l_index]+'/MC_Ds_'+l_flv[l_index]+l_flv[l_index]+'.pickle', 'wb') as handle:\n",
" pickle.dump(MC_Ds_mass_for_fit, handle, protocol=pickle.HIGHEST_PROTOCOL)\n",
" \n",
"with open('/disk/lhcb_data/davide/Rphipi/MC/for_fit/'+l_flv[l_index]+l_flv[l_index]+'/MC_Dplus_'+l_flv[l_index]+l_flv[l_index]+'.pickle', 'wb') as handle:\n",
" pickle.dump(MC_Dplus_mass_for_fit, handle, protocol=pickle.HIGHEST_PROTOCOL)\n",
"\n",
"with open('/disk/lhcb_data/davide/Rphipi/NN_for_training/'+l_flv[l_index]+l_flv[l_index]+'/data_for_NN_'+l_flv[l_index]+l_flv[l_index]+'.pickle', 'wb') as handle:\n",
" pickle.dump(data_tuple_bkg, handle, protocol=pickle.HIGHEST_PROTOCOL)\n",
" \n",
"with open('/disk/lhcb_data/davide/Rphipi/NN_for_selection/'+l_flv[l_index]+l_flv[l_index]+'/data_for_NN_'+l_flv[l_index]+l_flv[l_index]+'.pickle', 'wb') as handle:\n",
" pickle.dump(data_tuple_dict, handle, protocol=pickle.HIGHEST_PROTOCOL)"
]
},
{
"cell_type": "code",
"execution_count": 29,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"14066"
]
},
"execution_count": 29,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"data_tuple_bkg[\"Ds_ConsD_M\"].shape[0]"
]
},
{
"cell_type": "code",
"execution_count": 30,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"1974"
]
},
"execution_count": 30,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"MC_Dplus_tuple_sig[\"Dplus_ConsD_M\"].shape[0]"
]
},
{
"cell_type": "code",
"execution_count": 31,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"1466"
]
},
"execution_count": 31,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"MC_Ds_tuple_sig[\"Ds_ConsD_M\"].shape[0]"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 2",
"language": "python",
"name": "python2"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 2
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython2",
"version": "2.7.15"
}
},
"nbformat": 4,
"nbformat_minor": 2
}
Davide Lancierini