{ "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 }