import os
import pickle
import numpy as np
import sys
import root_numpy as rn
mother_ID=['Dplus','Ds','both']
l_flv = ['e','mu']
data_type = ['MC','data']
mag_status =['Up','Down']
tree_name = 'Ds_OfflineTree/DecayTree'
def namestr(obj, namespace):
return [name for name in namespace if namespace[name] is obj]
def sel_eff(tp,threshold):
sig_eps = np.float(np.where(tp>threshold)[0].shape[0])/np.float(tp.shape[0])
return sig_eps
def load_datasets(l_index, PATH):
if sys.version_info[0]>2:
data_index=0
with open(PATH+data_type[data_index]+'/'+l_flv[l_index]+'_tuples/for_BDT_training/'+mother_ID[1]+'_phipi_'+l_flv[l_index]+l_flv[l_index]+'/'+mother_ID[1]+'_phipi_'+l_flv[l_index]+l_flv[l_index]+'_for_BDT.pickle', 'rb') as f:
MC_Ds_sig_dict=pickle.load(f, encoding='latin1')
data_index=0
with open(PATH+data_type[data_index]+'/'+l_flv[l_index]+'_tuples/for_BDT_training/'+mother_ID[0]+'_phipi_'+l_flv[l_index]+l_flv[l_index]+'/'+mother_ID[0]+'_phipi_'+l_flv[l_index]+l_flv[l_index]+'_for_BDT.pickle', 'rb') as f:
MC_Dplus_sig_dict=pickle.load(f, encoding='latin1')
data_index=1
with open(PATH+data_type[data_index]+'/'+l_flv[l_index]+'_tuples/for_BDT_training/'+mother_ID[1]+'_phipi_'+l_flv[l_index]+l_flv[l_index]+'_for_BDT.pickle', 'rb') as f:
data_bkg_dict=pickle.load(f, encoding='latin1')
else:
data_index=0
with open(PATH+data_type[data_index]+'/'+l_flv[l_index]+'_tuples/for_BDT_training/'+mother_ID[1]+'_phipi_'+l_flv[l_index]+l_flv[l_index]+'/'+mother_ID[1]+'_phipi_'+l_flv[l_index]+l_flv[l_index]+'_for_BDT.pickle', 'rb') as f:
MC_Ds_sig_dict=pickle.load(f)#, encoding='latin1')
data_index=0
with open(PATH+data_type[data_index]+'/'+l_flv[l_index]+'_tuples/for_BDT_training/'+mother_ID[0]+'_phipi_'+l_flv[l_index]+l_flv[l_index]+'/'+mother_ID[0]+'_phipi_'+l_flv[l_index]+l_flv[l_index]+'_for_BDT.pickle', 'rb') as f:
MC_Dplus_sig_dict=pickle.load(f)#, encoding='latin1')
data_index=1
with open(PATH+data_type[data_index]+'/'+l_flv[l_index]+'_tuples/for_BDT_training/'+mother_ID[1]+'_phipi_'+l_flv[l_index]+l_flv[l_index]+'_for_BDT.pickle', 'rb') as f:
data_bkg_dict=pickle.load(f)#, encoding='latin1')
return MC_Dplus_sig_dict, MC_Ds_sig_dict, data_bkg_dict
def extract_array_for_BDT(event_dict, branches_needed, examples):
features = len(branches_needed)
extracted = np.array(
[np.zeros(shape=features) for i in range(examples)]
)
for event in range(examples):
for i, key in enumerate(branches_needed):
extracted[event][i]=event_dict[key][event]
return extracted
def add_labels(ndarray, signal=True):
if signal==True:
labeled_ndarray=np.concatenate((ndarray, np.ones(shape=(ndarray.shape[0],1))), axis =1)
elif signal==False:
labeled_ndarray=np.concatenate((ndarray, np.zeros(shape=(ndarray.shape[0],1))), axis =1)
return labeled_ndarray
def to_one_hot(labels):
temp = np.zeros(shape=(labels.shape[0],2))
for i in range(labels.shape[0]):
if labels[i]==0:
temp[i][0]=1
else:
temp[i][1]=1
return temp
def k_subsets(i, k, X, Y_labels):
train_size, dim=X.shape
#divide in k subsets
k_batch_size = train_size//k
X_dict={}
Y_dict={}
for b in range(k):
X_dict[b]=X[k_batch_size*b:k_batch_size*(b+1)]
Y_dict[b]=Y_labels[k_batch_size*b:k_batch_size*(b+1)]
k_range=np.arange(k)
X_test = X_dict[i][:,0:dim]
Y_test = Y_dict[i][:,0:dim]
k_subset=np.delete(k_range, i)
X_train = np.concatenate([X_dict[j][:,0:dim] for j in k_subset],axis=0)
Y_train = np.concatenate([Y_dict[j][:,0:dim] for j in k_subset],axis=0)
return X_train, Y_train, X_test, Y_test, X_dict, Y_dict
def return_branches(data_index=None, mother_index=None, l_index=None):
#data_index =0 -> MC
if data_index==0:
if l_index==1:
branches_needed = [
#________________________________________
#D MC true info
mother_ID[mother_index]+'_MC_MOTHER_ID',
mother_ID[mother_index]+'_BKGCAT',
mother_ID[mother_index]+'_TRUEID',
#________________________________________
#D Geometric variables, pT and FD
mother_ID[mother_index]+"_ENDVERTEX_CHI2",
mother_ID[mother_index]+"_ENDVERTEX_NDOF",
mother_ID[mother_index]+"_IPCHI2_OWNPV",
mother_ID[mother_index]+"_OWNPV_CHI2",
mother_ID[mother_index]+"_OWNPV_NDOF",
mother_ID[mother_index]+"_IP_OWNPV",
mother_ID[mother_index]+"_DIRA_OWNPV",
mother_ID[mother_index]+"_PX",
mother_ID[mother_index]+"_PY",
mother_ID[mother_index]+"_PZ",
mother_ID[mother_index]+"_PT",
mother_ID[mother_index]+"_P",
mother_ID[mother_index]+"_PE",
mother_ID[mother_index]+"_FD_OWNPV",
mother_ID[mother_index]+"_FDCHI2_OWNPV",
#D Reconstructed mass
mother_ID[mother_index]+"_ConsD_M",
mother_ID[mother_index]+"_M",
#D Trigger variables
mother_ID[mother_index]+"_L0Global_TIS",
mother_ID[mother_index]+"_L0Global_TOS",
mother_ID[mother_index]+"_L0MuonDecision_TIS",
mother_ID[mother_index]+"_L0MuonDecision_TOS",
mother_ID[mother_index]+"_L0ElectronDecision_TIS",
mother_ID[mother_index]+"_L0ElectronDecision_TOS",
mother_ID[mother_index]+"_L0HadronDecision_TIS",
mother_ID[mother_index]+"_L0HadronDecision_TOS",
mother_ID[mother_index]+"_Hlt1TrackMVADecision_TOS",
mother_ID[mother_index]+"_Hlt2RareCharmD2Pi"+l_flv[l_index].capitalize()+l_flv[l_index].capitalize()+"OSDecision_TOS",
mother_ID[mother_index]+"_Hlt2Phys_TOS",
#________________________________________
#PHI MC TRUE INFO
'phi_MC_MOTHER_ID',
'phi_BKGCAT',
'phi_TRUEID',
#________________________________________
#phi geometric variables, pT and FD
"phi_ENDVERTEX_CHI2",
"phi_ENDVERTEX_NDOF",
"phi_IPCHI2_OWNPV",
#"phi_OWNPV_CHI2",
#"phi_OWNPV_NDOF",
#"phi_IP_OWNPV",
#"phi_DIRA_OWNPV",
"phi_PT",
"phi_P",
#phi Reconstructed mass
"phi_M",
#________________________________________
#PION
#Pion mother ID and bkg cat
'pi_MC_MOTHER_ID',
'pi_TRUEID',
#_____________________________________
#pi Geometric variables and pT
#"pi_OWNPV_CHI2",
#"pi_OWNPV_NDOF",
#'pi_IP_OWNPV',
"pi_PX",
"pi_PY",
"pi_PZ",
'pi_PT',
'pi_P',
'pi_PE',
#pi PID variables
"pi_MC15TuneV1_ProbNNpi",
#________________________________________
#LEPTONS
l_flv[l_index]+'_plus_MC_MOTHER_ID',
l_flv[l_index]+'_plus_TRUEID',
l_flv[l_index]+'_minus_MC_MOTHER_ID',
l_flv[l_index]+'_minus_TRUEID',
#________________________________________
#leptons Geometric variables and pT
l_flv[l_index]+"_plus_OWNPV_CHI2",
#l_flv[l_index]+"_plus_OWNPV_NDOF",
l_flv[l_index]+"_minus_OWNPV_CHI2",
#l_flv[l_index]+"_minus_OWNPV_NDOF",
#
#l_flv[l_index]+"_plus_IP_OWNPV",
#l_flv[l_index]+"_minus_IP_OWNPV",
l_flv[l_index]+"_plus_PX",
l_flv[l_index]+"_plus_PY",
l_flv[l_index]+"_plus_PZ",
l_flv[l_index]+"_plus_PE",
l_flv[l_index]+"_minus_PX",
l_flv[l_index]+"_minus_PY",
l_flv[l_index]+"_minus_PZ",
l_flv[l_index]+"_minus_PE",
l_flv[l_index]+"_plus_PT",
l_flv[l_index]+"_minus_PT",
l_flv[l_index]+"_plus_P",
l_flv[l_index]+"_minus_P",
'cos_thetal',
#leptons PID variables
l_flv[l_index]+"_plus_MC15TuneV1_ProbNN"+l_flv[l_index],
l_flv[l_index]+"_minus_MC15TuneV1_ProbNN"+l_flv[l_index],
l_flv[l_index]+"_plus_MC15TuneV1_ProbNNpi",
l_flv[l_index]+"_minus_MC15TuneV1_ProbNNpi",
]
if l_index==0:
branches_needed = [
#________________________________________
#D MC true info
mother_ID[mother_index]+'_MC_MOTHER_ID',
mother_ID[mother_index]+'_BKGCAT',
mother_ID[mother_index]+'_TRUEID',
#________________________________________
#D Geometric variables, pT and FD
mother_ID[mother_index]+"_ENDVERTEX_CHI2",
mother_ID[mother_index]+"_ENDVERTEX_NDOF",
mother_ID[mother_index]+"_IPCHI2_OWNPV",
mother_ID[mother_index]+"_OWNPV_CHI2",
mother_ID[mother_index]+"_OWNPV_NDOF",
mother_ID[mother_index]+"_IP_OWNPV",
mother_ID[mother_index]+"_DIRA_OWNPV",
mother_ID[mother_index]+"_PX",
mother_ID[mother_index]+"_PY",
mother_ID[mother_index]+"_PZ",
mother_ID[mother_index]+"_PT",
mother_ID[mother_index]+"_P",
mother_ID[mother_index]+"_PE",
mother_ID[mother_index]+"_FD_OWNPV",
mother_ID[mother_index]+"_FDCHI2_OWNPV",
#D Reconstructed mass
mother_ID[mother_index]+"_ConsD_M",
mother_ID[mother_index]+"_M",
#D Trigger variables
mother_ID[mother_index]+"_L0Global_TIS",
mother_ID[mother_index]+"_L0Global_TOS",
mother_ID[mother_index]+"_L0MuonDecision_TIS",
mother_ID[mother_index]+"_L0MuonDecision_TOS",
mother_ID[mother_index]+"_L0ElectronDecision_TIS",
mother_ID[mother_index]+"_L0ElectronDecision_TOS",
mother_ID[mother_index]+"_L0HadronDecision_TIS",
mother_ID[mother_index]+"_L0HadronDecision_TOS",
mother_ID[mother_index]+"_Hlt1TrackMVADecision_TOS",
mother_ID[mother_index]+"_Hlt2RareCharmD2Pi"+l_flv[l_index].capitalize()+l_flv[l_index].capitalize()+"OSDecision_TOS",
mother_ID[mother_index]+"_Hlt2Phys_TOS",
#________________________________________
#PHI MC TRUE INFO
'phi_MC_MOTHER_ID',
'phi_BKGCAT',
'phi_TRUEID',
#________________________________________
#phi geometric variables, pT and FD
"phi_ENDVERTEX_CHI2",
"phi_ENDVERTEX_NDOF",
"phi_IPCHI2_OWNPV",
#"phi_OWNPV_CHI2",
#"phi_OWNPV_NDOF",
#"phi_IP_OWNPV",
#"phi_DIRA_OWNPV",
"phi_PT",
"phi_P",
#phi Reconstructed mass
"phi_M",
#________________________________________
#PION
#Pion mother ID and bkg cat
'pi_MC_MOTHER_ID',
'pi_TRUEID',
#_____________________________________
#pi Geometric variables and pT
#"pi_OWNPV_CHI2",
#"pi_OWNPV_NDOF",
#'pi_IP_OWNPV',
"pi_PX",
"pi_PY",
"pi_PZ",
'pi_PT',
'pi_P',
'pi_PE',
#pi PID variables
"pi_MC15TuneV1_ProbNNpi",
#________________________________________
#LEPTONS
l_flv[l_index]+'_plus_MC_MOTHER_ID',
l_flv[l_index]+'_plus_TRUEID',
l_flv[l_index]+'_minus_MC_MOTHER_ID',
l_flv[l_index]+'_minus_TRUEID',
#________________________________________
#leptons Geometric variables and pT
l_flv[l_index]+"_plus_OWNPV_CHI2",
#l_flv[l_index]+"_plus_OWNPV_NDOF",
l_flv[l_index]+"_minus_OWNPV_CHI2",
#l_flv[l_index]+"_minus_OWNPV_NDOF",
#
#l_flv[l_index]+"_plus_IP_OWNPV",
#l_flv[l_index]+"_minus_IP_OWNPV",
l_flv[l_index]+"_plus_PX",
l_flv[l_index]+"_plus_PY",
l_flv[l_index]+"_plus_PZ",
l_flv[l_index]+"_plus_PE",
l_flv[l_index]+"_minus_PX",
l_flv[l_index]+"_minus_PY",
l_flv[l_index]+"_minus_PZ",
l_flv[l_index]+"_minus_PE",
l_flv[l_index]+"_plus_PT",
l_flv[l_index]+"_plus_P",
l_flv[l_index]+"_minus_PT",
l_flv[l_index]+"_minus_P",
'cos_thetal',
#leptons PID variables
l_flv[l_index]+"_plus_MC15TuneV1_ProbNN"+l_flv[l_index],
l_flv[l_index]+"_minus_MC15TuneV1_ProbNN"+l_flv[l_index],
l_flv[l_index]+"_plus_MC15TuneV1_ProbNNpi",
l_flv[l_index]+"_minus_MC15TuneV1_ProbNNpi",
]
#data_index =1 -> data
if data_index==1:
branches_needed = [
#________________________________
#D Geometric variables, pT and FD
mother_ID[mother_index]+"_ENDVERTEX_CHI2",
mother_ID[mother_index]+"_ENDVERTEX_NDOF",
mother_ID[mother_index]+"_IPCHI2_OWNPV",
mother_ID[mother_index]+"_OWNPV_CHI2",
mother_ID[mother_index]+"_OWNPV_NDOF",
mother_ID[mother_index]+"_IP_OWNPV",
mother_ID[mother_index]+"_DIRA_OWNPV",
mother_ID[mother_index]+"_PX",
mother_ID[mother_index]+"_PY",
mother_ID[mother_index]+"_PZ",
mother_ID[mother_index]+"_PT",
mother_ID[mother_index]+"_P",
mother_ID[mother_index]+"_PE",
mother_ID[mother_index]+"_FD_OWNPV",
mother_ID[mother_index]+"_FDCHI2_OWNPV",
#D Reconstructed mass
mother_ID[mother_index]+"_ConsD_M",
mother_ID[mother_index]+"_M",
#D Trigger variables
mother_ID[mother_index]+"_L0Global_TIS",
mother_ID[mother_index]+"_L0Global_TOS",
mother_ID[mother_index]+"_L0MuonDecision_TIS",
mother_ID[mother_index]+"_L0MuonDecision_TOS",
mother_ID[mother_index]+"_L0ElectronDecision_TIS",
mother_ID[mother_index]+"_L0ElectronDecision_TOS",
mother_ID[mother_index]+"_L0HadronDecision_TIS",
mother_ID[mother_index]+"_L0HadronDecision_TOS",
mother_ID[mother_index]+"_Hlt1TrackMVADecision_TOS",
mother_ID[mother_index]+"_Hlt2RareCharmD2Pi"+l_flv[l_index].capitalize()+l_flv[l_index].capitalize()+"OSDecision_TOS",
mother_ID[mother_index]+"_Hlt2Phys_TOS",
#________________________________________
#phi geometric variables, pT and FD
"phi_ENDVERTEX_CHI2",
"phi_ENDVERTEX_NDOF",
"phi_IPCHI2_OWNPV",
#"phi_OWNPV_CHI2",
#"phi_OWNPV_NDOF",
#"phi_IP_OWNPV",
#"phi_DIRA_OWNPV",
"phi_PT",
"phi_P",
#phi Reconstructed mass
"phi_M",
#________________________________________
#PION
#_____________________________________
#pi Geometric variables and pT
#"pi_OWNPV_CHI2",
#"pi_OWNPV_NDOF",
#'pi_IP_OWNPV',
"pi_PX",
"pi_PY",
"pi_PZ",
'pi_PT',
'pi_P',
'pi_PE',
#pi PID variables
"pi_MC15TuneV1_ProbNNpi",
#________________________________________
#LEPTONS
#________________________________________
#leptons Geometric variables and pT
l_flv[l_index]+"_plus_OWNPV_CHI2",
#l_flv[l_index]+"_plus_OWNPV_NDOF",
l_flv[l_index]+"_minus_OWNPV_CHI2",
#l_flv[l_index]+"_minus_OWNPV_NDOF",
#
#l_flv[l_index]+"_plus_IP_OWNPV",
#l_flv[l_index]+"_minus_IP_OWNPV",
l_flv[l_index]+"_plus_PX",
l_flv[l_index]+"_plus_PY",
l_flv[l_index]+"_plus_PZ",
l_flv[l_index]+"_plus_PE",
l_flv[l_index]+"_minus_PX",
l_flv[l_index]+"_minus_PY",
l_flv[l_index]+"_minus_PZ",
l_flv[l_index]+"_minus_PE",
l_flv[l_index]+"_plus_PT",
l_flv[l_index]+"_plus_P",
l_flv[l_index]+"_minus_PT",
l_flv[l_index]+"_minus_P",
'cos_thetal',
#leptons PID variables
l_flv[l_index]+"_plus_MC15TuneV1_ProbNN"+l_flv[l_index],
l_flv[l_index]+"_minus_MC15TuneV1_ProbNN"+l_flv[l_index],
l_flv[l_index]+"_plus_MC15TuneV1_ProbNNpi",
l_flv[l_index]+"_minus_MC15TuneV1_ProbNNpi",
]
return branches_needed
def return_branches_BDT(mother_index=None, l_index=None):
if l_index==1:
branches_needed = [
#________________________________
#D Geometric variables, pT and FD
# mother_ID[mother_index]+"_OWNPV_CHI2",
mother_ID[mother_index]+"_ENDVERTEX_CHI2",
mother_ID[mother_index]+"_PT",
mother_ID[mother_index]+"_FD_OWNPV",
mother_ID[mother_index]+"_IPCHI2_OWNPV",
mother_ID[mother_index]+"_DIRA_OWNPV",
#mother_ID[mother_index]+"_FDCHI2_OWNPV",
#________________________________________
#phi geometric variables, pT and FD
#"phi_ENDVERTEX_CHI2",
"phi_IPCHI2_OWNPV",
#"phi_OWNPV_CHI2",
#"phi_OWNPV_NDOF",
#"phi_IP_OWNPV",
#"phi_DIRA_OWNPV",
#"phi_PT",
#________________________________________
#PION
#_____________________________________
#pi Geometric variables and pT
#"pi_OWNPV_CHI2",
#"pi_OWNPV_NDOF",
#'pi_IP_OWNPV',
'pi_PT',
'pi_P',
#________________________________________
#LEPTONS
#________________________________________
#leptons Geometric variables and pT
#l_flv[l_index]+"_plus_OWNPV_CHI2",
#l_flv[l_index]+"_plus_OWNPV_NDOF",
#l_flv[l_index]+"_minus_OWNPV_CHI2",
#l_flv[l_index]+"_minus_OWNPV_NDOF",
#
#l_flv[l_index]+"_plus_IP_OWNPV",
#l_flv[l_index]+"_minus_IP_OWNPV",
l_flv[l_index]+"_plus_PT",
l_flv[l_index]+"_minus_PT",
l_flv[l_index]+"_plus_P",
l_flv[l_index]+"_minus_P",
"cos_thetal",
#D Reconstructed mass
# mother_ID[mother_index]+"_ConsD_M",
]
if l_index==0:
branches_needed = [
#________________________________
#D Geometric variables, pT and FD
# mother_ID[mother_index]+"_OWNPV_CHI2",
mother_ID[mother_index]+"_ENDVERTEX_CHI2",
mother_ID[mother_index]+"_PT",
mother_ID[mother_index]+"_FD_OWNPV",
mother_ID[mother_index]+"_IPCHI2_OWNPV",
mother_ID[mother_index]+"_DIRA_OWNPV",
#mother_ID[mother_index]+"_FDCHI2_OWNPV",
#________________________________________
#phi geometric variables, pT and FD
#"phi_ENDVERTEX_CHI2",
"phi_IPCHI2_OWNPV",
#"phi_OWNPV_CHI2",
#"phi_OWNPV_NDOF",
#"phi_IP_OWNPV",
#"phi_DIRA_OWNPV",
#"phi_PT",
#________________________________________
#PION
#_____________________________________
#pi Geometric variables and pT
#"pi_OWNPV_CHI2",
#"pi_OWNPV_NDOF",
#'pi_IP_OWNPV',
'pi_PT',
'pi_P',
#________________________________________
#LEPTONS
#________________________________________
#leptons Geometric variables and pT
#l_flv[l_index]+"_plus_OWNPV_CHI2",
#l_flv[l_index]+"_plus_OWNPV_NDOF",
#l_flv[l_index]+"_minus_OWNPV_CHI2",
#l_flv[l_index]+"_minus_OWNPV_NDOF",
#
#l_flv[l_index]+"_plus_IP_OWNPV",
#l_flv[l_index]+"_minus_IP_OWNPV",
l_flv[l_index]+"_plus_PT",
l_flv[l_index]+"_minus_PT",
l_flv[l_index]+"_plus_P",
l_flv[l_index]+"_minus_P",
"cos_thetal",
#D Reconstructed mass
#mother_ID[mother_index]+"_ConsD_M",
]
return branches_needed
def return_branches_TISTOS(data_index=None, l_index=None):
if l_flv[l_index]=='e':
#data_index =0 -> MC
if data_index==0:
branches_needed = [
#________________________________________
#D MC true info
'B_MC_MOTHER_ID',
'B_BKGCAT',
'B_TRUEID',
#________________________________________
#B Geometric variables, pT and FD
"B_ENDVERTEX_CHI2",
"B_ENDVERTEX_NDOF",
"B_IPCHI2_OWNPV",
"B_OWNPV_CHI2",
"B_OWNPV_NDOF",
"B_IP_OWNPV",
"B_DIRA_OWNPV",
"B_PX",
"B_PY",
"B_PZ",
"B_PT",
"B_P",
"B_PE",
"B_FD_OWNPV",
"B_FDCHI2_OWNPV",
"B_PVandJpsiDTF_B_M",
#D Reconstructed mass
#mother_ID[mother_index]+"_ConsD_M",
#D Trigger variables
"B_L0Global_TIS",
"B_L0MuonDecision_TIS",
"B_L0MuonDecision_TOS",
"B_L0ElectronDecision_TIS",
"B_L0ElectronDecision_TOS",
"B_L0HadronDecision_TIS",
"B_L0HadronDecision_TOS",
"B_Hlt1TrackMVADecision_TOS",
"B_Hlt1TrackMVADecision_TIS",
#"B_Hlt2RareCharmD2Pi"+l_flv[l_index].capitalize()+l_flv[l_index].capitalize()+"OSDecision_TOS",
#"B_Hlt2Phys_TOS",
#________________________________________
#PHI MC TRUE INFO
'Jpsi_MC_MOTHER_ID',
'Jpsi_BKGCAT',
'Jpsi_TRUEID',
#________________________________________
#phi geometric variables, pT and FD
"Jpsi_ENDVERTEX_CHI2",
"Jpsi_ENDVERTEX_NDOF",
"Jpsi_IPCHI2_OWNPV",
#"phi_OWNPV_CHI2",
#"phi_OWNPV_NDOF",
#"phi_IP_OWNPV",
#"phi_DIRA_OWNPV",
"Jpsi_PT",
#phi Reconstructed mass
"Jpsi_M",
#________________________________________
#K*
#K* mother ID and bkg cat
'Kstar_MC_MOTHER_ID',
'Kstar_TRUEID',
#_____________________________________
#Kaon Geometric variables and pT
#"pi_OWNPV_CHI2",
#"pi_OWNPV_NDOF",
#'pi_IP_OWNPV',
"Kstar_PX",
"Kstar_PY",
"Kstar_PZ",
"Kstar_PT",
"Kstar_P",
"Kstar_PE",
#________________________________________
#________________________________________
#Kaon
#Kaon mother ID and bkg cat
'K_MC_MOTHER_ID',
'K_TRUEID',
#_____________________________________
#Kaon Geometric variables and pT
#"pi_OWNPV_CHI2",
#"pi_OWNPV_NDOF",
#'pi_IP_OWNPV',
"K_PX",
"K_PY",
"K_PZ",
"K_PT",
"K_P",
"K_PE",
#Kaon PID variables
"K_MC15TuneV1_ProbNNpi",
#________________________________________
#Pion
#Pion mother ID and bkg cat
'Pi_MC_MOTHER_ID',
'Pi_TRUEID',
#_____________________________________
#Kaon Geometric variables and pT
#"pi_OWNPV_CHI2",
#"pi_OWNPV_NDOF",
#'pi_IP_OWNPV',
"Pi_PX",
"Pi_PY",
"Pi_PZ",
"Pi_PT",
"Pi_P",
"Pi_PE",
#Kaon PID variables
"Pi_MC15TuneV1_ProbNNpi",
#________________________________________
#LEPTONS
'L1_MC_MOTHER_ID',
'L1_TRUEID',
'L2_MC_MOTHER_ID',
'L2_TRUEID',
#________________________________________
#leptons Geometric variables and pT
'L1_OWNPV_CHI2',
#l_flv[l_index]+"_plus_OWNPV_NDOF",
'L2_OWNPV_CHI2',
#l_flv[l_index]+"_minus_OWNPV_NDOF",
#
#l_flv[l_index]+"_plus_IP_OWNPV",
#l_flv[l_index]+"_minus_IP_OWNPV",
"L1_PX",
"L1_PY",
"L1_PZ",
"L1_PT",
"L1_P",
"L1_PE",
'L1_CosTheta',
"L2_PX",
"L2_PY",
"L2_PZ",
"L2_PT",
"L2_P",
"L2_PE",
'L2_CosTheta',
#leptons PID variables
# "L1_MC15TuneV1_ProbNN",
# "L2_MC15TuneV1_ProbNN",
]
#data_index =1 -> data
if data_index==1:
branches_needed = [
#________________________________________
#B Geometric variables, pT and FD
"B_ENDVERTEX_CHI2",
"B_ENDVERTEX_NDOF",
"B_IPCHI2_OWNPV",
"B_OWNPV_CHI2",
"B_OWNPV_NDOF",
"B_IP_OWNPV",
"B_DIRA_OWNPV",
"B_PX",
"B_PY",
"B_PZ",
"B_PT",
"B_P",
"B_FD_OWNPV",
"B_FDCHI2_OWNPV",
"B_PVandJpsiDTF_B_M",
#B Reconstructed mass
#mother_ID[mother_index]+"_ConsD_M",
#D Trigger variables
"B_L0Global_TIS",
"B_L0MuonDecision_TIS",
"B_L0MuonDecision_TOS",
"B_L0ElectronDecision_TIS",
"B_L0ElectronDecision_TOS",
"B_L0HadronDecision_TIS",
"B_L0HadronDecision_TOS",
"B_Hlt1TrackMVADecision_TOS",
"B_Hlt1TrackMVADecision_TIS",
#"B_Hlt2RareCharmD2Pi"+l_flv[l_index].capitalize()+l_flv[l_index].capitalize()+"OSDecision_TOS",
#"B_Hlt2Phys_TOS",
#___________________________________
#Jpsi
#___________________________________
#Jpsi geometric variables, pT and FD
"Jpsi_ENDVERTEX_CHI2",
"Jpsi_ENDVERTEX_NDOF",
"Jpsi_IPCHI2_OWNPV",
#"phi_OWNPV_CHI2",
#"phi_OWNPV_NDOF",
#"phi_IP_OWNPV",
#"phi_DIRA_OWNPV",
"Jpsi_PT",
#phi Reconstructed mass
"Jpsi_M",
#________________________________________
#Kstar
#_____________________________________
#Kaon Geometric variables and pT
#"pi_OWNPV_CHI2",
#"pi_OWNPV_NDOF",
#'pi_IP_OWNPV',
"Kstar_PX",
"Kstar_PY",
"Kstar_PZ",
"Kstar_PT",
"Kstar_P",
"Kstar_PE",
#________________________________________
#Kaon
#_____________________________________
#Kaon Geometric variables and pT
#"pi_OWNPV_CHI2",
#"pi_OWNPV_NDOF",
#'pi_IP_OWNPV',
"K_PX",
"K_PY",
"K_PZ",
"K_PT",
"K_P",
"K_PE",
#Kaon PID variables
"K_MC15TuneV1_ProbNNpi",
#________________________________________
#Pion
#_____________________________________
#Kaon Geometric variables and pT
#"pi_OWNPV_CHI2",
#"pi_OWNPV_NDOF",
#'pi_IP_OWNPV',
"Pi_PX",
"Pi_PY",
"Pi_PZ",
"Pi_PT",
"Pi_P",
"Pi_PE",
#Kaon PID variables
"Pi_MC15TuneV1_ProbNNpi",
#________________________________________
#LEPTONS
#________________________________________
#leptons Geometric variables and pT
'L1_OWNPV_CHI2',
#l_flv[l_index]+"_plus_OWNPV_NDOF",
'L2_OWNPV_CHI2',
#l_flv[l_index]+"_minus_OWNPV_NDOF",
#
#l_flv[l_index]+"_plus_IP_OWNPV",
#l_flv[l_index]+"_minus_IP_OWNPV",
"L1_PX",
"L1_PY",
"L1_PZ",
"L1_PT",
"L1_P",
"L1_PE",
'L1_CosTheta',
"L2_PX",
"L2_PY",
"L2_PZ",
"L2_PT",
"L2_P",
"L2_PE",
'L2_CosTheta',
#leptons PID variables
# "L1_MC15TuneV1_ProbNN",
# "L2_MC15TuneV1_ProbNN",
]
#data_index =1 -> data
if l_flv[l_index]=='mu':
#data_index =0 -> MC
if data_index==0:
branches_needed = [
#________________________________________
#B Geometric variables, pT and FD
"B0_BKGCAT",
"B0_TRUEID",
"B0_ENDVERTEX_CHI2",
"B0_ENDVERTEX_NDOF",
"B0_IPCHI2_OWNPV",
"B0_OWNPV_CHI2",
"B0_OWNPV_NDOF",
"B0_IP_OWNPV",
"B0_DIRA_OWNPV",
"B0_PX",
"B0_PY",
"B0_PZ",
"B0_PT",
"B0_P",
"B0_PE",
"B0_FD_OWNPV",
"B0_FDCHI2_OWNPV",
"B0_M",
"B0_TRUEP_X",
"B0_TRUEP_Y",
"B0_TRUEP_Z",
"B0_TRUEP_E",
"B0_FD_OWNPV",
"B0_FDCHI2_OWNPV",
"B0_M",
#"B0_PVandJpsiDTF_B_M",
#B Reconstructed mass
#mother_ID[mother_index]+"_ConsD_M",
#D Trigger variables
"B0_L0Global_TIS",
"B0_L0MuonDecision_TIS",
"B0_L0MuonDecision_TOS",
"B0_L0ElectronDecision_TIS",
"B0_L0ElectronDecision_TOS",
"B0_L0HadronDecision_TIS",
"B0_L0HadronDecision_TOS",
"B0_Hlt1TrackMVADecision_TOS",
"B0_Hlt1TrackMVADecision_TIS",
#"B_Hlt2Phys_TOS",
#___________________________________
#Jpsi
"J_psi_BKGCAT",
"J_psi_TRUEID",
#___________________________________
#Jpsi geometric variables, pT and FD
"J_psi_ENDVERTEX_CHI2",
"J_psi_ENDVERTEX_NDOF",
"J_psi_IPCHI2_OWNPV",
#"phi_OWNPV_CHI2",
#"phi_OWNPV_NDOF",
#"phi_IP_OWNPV",
#"phi_DIRA_OWNPV",
"J_psi_PT",
#phi Reconstructed mass
"J_psi_M",
#________________________________________
#Kaon
#_____________________________________
#Kaon Geometric variables and pT
#"pi_OWNPV_CHI2",
#"pi_OWNPV_NDOF",
#'pi_IP_OWNPV',
"Kstar_TRUEID",
"Kstar_PX",
"Kstar_PY",
"Kstar_PZ",
"Kstar_TRUEP_X",
"Kstar_TRUEP_Y",
"Kstar_TRUEP_Z",
"Kstar_PT",
"Kstar_P",
"Kstar_PE",
"Kstar_TRUEP_E",
#________________________________________
#Kaon
#_____________________________________
#Kaon Geometric variables and pT
#"pi_OWNPV_CHI2",
#"pi_OWNPV_NDOF",
#'pi_IP_OWNPV',
"K_TRUEID",
"K_PX",
"K_PY",
"K_PZ",
"K_TRUEP_X",
"K_TRUEP_Y",
"K_TRUEP_Z",
"K_PT",
"K_P",
"K_PE",
"K_TRUEP_E",
#Kaon PID variables
"K_MC15TuneV1_ProbNNk",
#________________________________________
#Pion
"Pi_TRUEID",
#_____________________________________
#Kaon Geometric variables and pT
#"pi_OWNPV_CHI2",
#"pi_OWNPV_NDOF",
#'pi_IP_OWNPV',
"Pi_PX",
"Pi_PY",
"Pi_PZ",
"Pi_TRUEP_X",
"Pi_TRUEP_Y",
"Pi_TRUEP_Z",
"Pi_PT",
"Pi_P",
"Pi_PE",
"Pi_TRUEP_E",
#Kaon PID variables
"Pi_MC15TuneV1_ProbNNpi",
#________________________________________
#LEPTONS
'mu_plus_TRUEID',
'mu_minus_TRUEID',
#________________________________________
#leptons Geometric variables and pT
'mu_minus_OWNPV_CHI2',
#l_flv[l_index]+"_plus_OWNPV_NDOF",
'mu_plus_OWNPV_CHI2',
#l_flv[l_index]+"_minus_OWNPV_NDOF",
#
#l_flv[l_index]+"_plus_IP_OWNPV",
#l_flv[l_index]+"_minus_IP_OWNPV",
"mu_plus_PX",
"mu_plus_PY",
"mu_plus_PZ",
"mu_plus_TRUEP_X",
"mu_plus_TRUEP_Y",
"mu_plus_TRUEP_Z",
"mu_plus_TRUEP_E",
"mu_plus_PY",
"mu_plus_PZ",
"mu_plus_PT",
"mu_plus_P",
"mu_plus_PE",
"mu_plus_CosTheta",
"mu_minus_PX",
"mu_minus_PY",
"mu_minus_PZ",
"mu_minus_TRUEP_X",
"mu_minus_TRUEP_Y",
"mu_minus_TRUEP_Z",
"mu_minus_TRUEP_E",
"mu_minus_PT",
"mu_minus_P",
"mu_minus_PE",
"mu_minus_CosTheta",
#leptons PID variables
# "L1_MC15TuneV1_ProbNN",
# "L2_MC15TuneV1_ProbNN",
]
#data_index =1 -> data
if data_index==1:
branches_needed = [
#________________________________________
#B Geometric variables, pT and FD
"B0_ENDVERTEX_CHI2",
"B0_ENDVERTEX_NDOF",
"B0_IPCHI2_OWNPV",
"B0_OWNPV_CHI2",
"B0_OWNPV_NDOF",
"B0_IP_OWNPV",
"B0_DIRA_OWNPV",
"B0_PX",
"B0_PY",
"B0_PZ",
"B0_PT",
"B0_P",
"B0_PE",
"B0_FD_OWNPV",
"B0_FDCHI2_OWNPV",
"B0_M",
#B Reconstructed mass
#mother_ID[mother_index]+"_ConsD_M",
#D Trigger variables
"B0_L0Global_TIS",
"B0_L0MuonDecision_TIS",
"B0_L0MuonDecision_TOS",
"B0_L0ElectronDecision_TIS",
"B0_L0ElectronDecision_TOS",
"B0_L0HadronDecision_TIS",
"B0_L0HadronDecision_TOS",
"B0_Hlt1TrackMVADecision_TOS",
"B0_Hlt1TrackMVADecision_TIS",
#"B_Hlt2Phys_TOS",
#___________________________________
#Jpsi
#___________________________________
#Jpsi geometric variables, pT and FD
"J_psi_ENDVERTEX_CHI2",
"J_psi_ENDVERTEX_NDOF",
"J_psi_IPCHI2_OWNPV",
#"phi_OWNPV_CHI2",
#"phi_OWNPV_NDOF",
#"phi_IP_OWNPV",
#"phi_DIRA_OWNPV",
"J_psi_PT",
#phi Reconstructed mass
"J_psi_M",
#________________________________________
#Kstar
#_____________________________________
#Kaon Geometric variables and pT
#"pi_OWNPV_CHI2",
#"pi_OWNPV_NDOF",
#'pi_IP_OWNPV',
"Kstar_PX",
"Kstar_PY",
"Kstar_PZ",
"Kstar_PT",
"Kstar_P",
"Kstar_PE",
#________________________________________
#Kaon
#_____________________________________
#Kaon Geometric variables and pT
#"pi_OWNPV_CHI2",
#"pi_OWNPV_NDOF",
#'pi_IP_OWNPV',
"K_PX",
"K_PY",
"K_PZ",
"K_PT",
"K_P",
"K_PE",
#Kaon PID variables
"K_MC15TuneV1_ProbNNk",
#________________________________________
#Pion
#_____________________________________
#Kaon Geometric variables and pT
#"pi_OWNPV_CHI2",
#"pi_OWNPV_NDOF",
#'pi_IP_OWNPV',
"Pi_PX",
"Pi_PY",
"Pi_PZ",
"Pi_PT",
"Pi_P",
"Pi_PE",
#Kaon PID variables
"Pi_MC15TuneV1_ProbNNpi",
#________________________________________
#LEPTONS
#________________________________________
#leptons Geometric variables and pT
'mu_plus_OWNPV_CHI2',
#l_flv[l_index]+"_plus_OWNPV_NDOF",
'mu_minus_OWNPV_CHI2',
#l_flv[l_index]+"_minus_OWNPV_NDOF",
#
#l_flv[l_index]+"_plus_IP_OWNPV",
#l_flv[l_index]+"_minus_IP_OWNPV",
"mu_plus_PX",
"mu_plus_PY",
"mu_plus_PZ",
"mu_plus_PT",
"mu_plus_P",
"mu_plus_PE",
"mu_plus_CosTheta",
"mu_minus_PX",
"mu_minus_PY",
"mu_minus_PZ",
"mu_minus_PT",
"mu_minus_P",
"mu_minus_PE",
"mu_minus_CosTheta",
#leptons PID variables
# "L1_MC15TuneV1_ProbNN",
# "L2_MC15TuneV1_ProbNN",
]
#data_index =1 -> data
return branches_needed
def norm_chi2(MC_Dplus_sig_dict, MC_Ds_sig_dict, data_dict):
MC_Ds_sig_dict["Ds_ENDVERTEX_CHI2"]=MC_Ds_sig_dict["Ds_ENDVERTEX_CHI2"]/MC_Ds_sig_dict["Ds_ENDVERTEX_NDOF"]
MC_Ds_sig_dict["Ds_IPCHI2_OWNPV"]=MC_Ds_sig_dict["Ds_IPCHI2_OWNPV"]/MC_Ds_sig_dict["Ds_ENDVERTEX_NDOF"]
MC_Ds_sig_dict["Ds_FDCHI2_OWNPV"]=MC_Ds_sig_dict["Ds_FDCHI2_OWNPV"]/MC_Ds_sig_dict["Ds_OWNPV_NDOF"]
MC_Ds_sig_dict["Ds_OWNPV_CHI2"]=MC_Ds_sig_dict["Ds_OWNPV_CHI2"]/MC_Ds_sig_dict["Ds_OWNPV_NDOF"]
del MC_Ds_sig_dict["Ds_OWNPV_NDOF"]
del MC_Ds_sig_dict["Ds_ENDVERTEX_NDOF"]
MC_Dplus_sig_dict["Dplus_ENDVERTEX_CHI2"]=MC_Dplus_sig_dict["Dplus_ENDVERTEX_CHI2"]/MC_Dplus_sig_dict["Dplus_ENDVERTEX_NDOF"]
MC_Dplus_sig_dict["Dplus_IPCHI2_OWNPV"]=MC_Dplus_sig_dict["Dplus_IPCHI2_OWNPV"]/MC_Dplus_sig_dict["Dplus_ENDVERTEX_NDOF"]
MC_Dplus_sig_dict["Dplus_FDCHI2_OWNPV"]=MC_Dplus_sig_dict["Dplus_FDCHI2_OWNPV"]/MC_Dplus_sig_dict["Dplus_OWNPV_NDOF"]
MC_Dplus_sig_dict["Dplus_OWNPV_CHI2"]=MC_Dplus_sig_dict["Dplus_OWNPV_CHI2"]/MC_Dplus_sig_dict["Dplus_OWNPV_NDOF"]
del MC_Dplus_sig_dict["Dplus_OWNPV_NDOF"]
del MC_Dplus_sig_dict["Dplus_ENDVERTEX_NDOF"]
data_dict["Ds_ENDVERTEX_CHI2"]=data_dict["Ds_ENDVERTEX_CHI2"]/data_dict["Ds_ENDVERTEX_NDOF"]
data_dict["Ds_IPCHI2_OWNPV"]=data_dict["Ds_IPCHI2_OWNPV"]/data_dict["Ds_ENDVERTEX_NDOF"]
data_dict["Ds_FDCHI2_OWNPV"]=data_dict["Ds_FDCHI2_OWNPV"]/data_dict["Ds_OWNPV_NDOF"]
data_dict["Ds_OWNPV_CHI2"]=data_dict["Ds_OWNPV_CHI2"]/data_dict["Ds_OWNPV_NDOF"]
del data_dict["Ds_OWNPV_NDOF"]
del data_dict["Ds_ENDVERTEX_NDOF"]
return MC_Dplus_sig_dict, MC_Ds_sig_dict, data_dict
# def mass_cut_for_fit(lower_cut, upper_cut, mother_index_fit=None, l_index=None,
# branches_needed=None, data_dict=None,
# MC_Dplus_dict=None, MC_Ds_dict=None):
# #Applies needed mass cuts on data and
# #returns arrays of mass distr for MC and data
# data_indices=[]
# for i in range(len(data_dict["Ds_ConsD_M"])):
# D_m = data_dict["Ds_ConsD_M"][i]
# #fixing a window on the phi mass
# if lower_cut<D_m<upper_cut:
# data_indices.append(i)
# for label in branches_needed:
# data_dict[label] = data_dict[label][data_indices]
# m_plus=MC_Dplus_dict["Dplus_ConsD_M"].shape[0]
# m_s=MC_Ds_dict["Ds_ConsD_M"].shape[0]
# m=data_dict["Ds_ConsD_M"].shape[0]
# mc_Dplus_mass=np.array([MC_Dplus_dict["Dplus_ConsD_M"][i][0] for i in range(m_plus)])
# mc_Ds_mass=np.array([MC_Ds_dict["Ds_ConsD_M"][i][0] for i in range(m_s)])
# #mc_D_mass=np.concatenate((mc_Ds_mass,mc_Dplus_mass),axis=0)
# data_mass=np.array([data_dict["Ds_ConsD_M"][i][0] for i in range(m)])
# # if mother_index_fit==0:
# # mc_mass=mc_Ds_mass
# # if mother_index_fit==1:
# # mc_mass=mc_Dplus_mass
# # if mother_index_fit==2:
# # mc_mass=mc_D_mass
# return data_mass, mc_Dplus_mass, mc_Ds_mass, data_dict
def fix_mass_vector(tuple_full, branch_name):
tuple_full[branch_name+'_fixed']= [tuple_full[branch_name][i][0] for i in range(len(tuple_full[branch_name]))]
tuple_full[branch_name]= np.array(tuple_full[branch_name+'_fixed'])
tuple_full.pop(branch_name+'_fixed', None)
return tuple_full
def save_tuples(namespace, tuple_dict, l_index, SAVE_PATH, file_format=None):
if "Ds" in namestr(tuple_dict, namespace)[0]:
mother_index=1
elif "Dplus" in namestr(tuple_dict, namespace)[0]:
mother_index=0
elif "data" in namestr(tuple_dict, namespace)[0]:
mother_index=1
if "MC" in namestr(tuple_dict, namespace)[0]:
data_index=0
else:
data_index=1
append_to_name=''
stages=["presel","truthed","ghosts","for_BDT","Hlt1","Hlt2","FullHlt","BDT_selected"]#,"lPID","piPID","PID"]
for stage in stages:
if stage in namestr(tuple_dict, namespace)[0]:
append_to_name=stage
dotrigCats=False
elif 'trigCats' in namestr(tuple_dict, namespace)[0]:
dotrigCats=True
if not dotrigCats:
print(append_to_name)
assert file_format!=None, '\n Specify file format: root or picke \n'
if file_format == 'pickle':
FILE_PATH = SAVE_PATH+mother_ID[mother_index]+'_phipi_'+l_flv[l_index]+l_flv[l_index]+'_'+append_to_name+'.pickle'
if os.path.exists(FILE_PATH):
print('Overwriting pickle file at'+FILE_PATH)
else:
print('Saving pickle file at'+FILE_PATH)
with open(FILE_PATH, 'wb') as handle:
pickle.dump(tuple_dict, handle, protocol=pickle.HIGHEST_PROTOCOL)
elif file_format == 'root':
FILE_PATH = SAVE_PATH+mother_ID[mother_index]+'_phipi_'+l_flv[l_index]+l_flv[l_index]+'_'+append_to_name+'.root'
if os.path.exists(FILE_PATH):
print('Overwriting root file at'+FILE_PATH)
mode = 'recreate'
else:
print('Saving root file at'+FILE_PATH)
mode = 'create'
n_evts=len(tuple_dict[tuple_dict.keys()[0]])
branches_needed=return_branches(data_index=data_index, mother_index=mother_index, l_index=l_index)
tuple_array= np.array(
[
tuple( tuple_dict[label][k] for label in branches_needed )
for k in range(n_evts)
],
dtype=[(label, np.float32) for label in branches_needed]
)
rn.array2root(tuple_array,
filename=FILE_PATH,
treename='DecayTree',
mode=mode,
)
if dotrigCats:
assert file_format!=None, '\n Specify file format: root or picke \n'
if file_format == 'pickle':
for i in range(len(tuple_dict)):
FILE_PATH = SAVE_PATH+str(i)+'/'+mother_ID[mother_index]+'_phipi_'+l_flv[l_index]+l_flv[l_index]+'_trigCat'+str(i)+'.pickle'
if os.path.exists(FILE_PATH):
print('Overwriting pickle file at'+FILE_PATH)
else:
print('Saving pickle file at'+FILE_PATH)
with open(FILE_PATH, 'wb') as handle:
pickle.dump(tuple_dict[i], handle, protocol=pickle.HIGHEST_PROTOCOL)
elif file_format == 'root':
for i in range(len(tuple_dict)):
FILE_PATH = SAVE_PATH+str(i)+'/'+mother_ID[mother_index]+'_phipi_'+l_flv[l_index]+l_flv[l_index]+'_trigCat'+str(i)+'.root'
if os.path.exists(FILE_PATH):
print('Overwriting root file at'+FILE_PATH)
mode = 'recreate'
else:
print('Saving root file at'+FILE_PATH)
mode = 'create'
n_evts=len(tuple_dict[i][tuple_dict[i].keys()[0]])
branches_needed=return_branches(data_index=data_index, mother_index=mother_index, l_index=l_index)
tuple_array= np.array(
[
tuple( tuple_dict[i][label][k] for label in branches_needed )
for k in range(n_evts)
],
dtype=[(label, np.float32) for label in branches_needed]
)
rn.array2root(tuple_array,
filename=FILE_PATH,
treename='DecayTree',
mode=mode,
)
elif append_to_name not in stages and 'trigCats' not in namestr(tuple_dict, namespace)[0]:
print('Tuple stage not recognised')
return
def add_branch_and_save_tuples(namespace, tuple_dict, l_index, branch_name, SAVE_PATH, file_format=None):
if "Ds" in namestr(tuple_dict, namespace)[0]:
mother_index=1
elif "Dplus" in namestr(tuple_dict, namespace)[0]:
mother_index=0
elif "data" in namestr(tuple_dict, namespace)[0]:
mother_index=1
if "MC" in namestr(tuple_dict, namespace)[0]:
data_index=0
else:
data_index=1
print('Adding branch '+branch_name)
if file_format == 'pickle':
FILE_PATH = SAVE_PATH+mother_ID[mother_index]+'_phipi_'+l_flv[l_index]+l_flv[l_index]+'_'+branch_name+'.pickle'
if os.path.exists(FILE_PATH):
print('Overwriting pickle file at'+FILE_PATH)
else:
print('Saving pickle file at'+FILE_PATH)
with open(FILE_PATH, 'wb') as handle:
pickle.dump(tuple_dict, handle, protocol=pickle.HIGHEST_PROTOCOL)
elif file_format == 'root':
FILE_PATH = SAVE_PATH+mother_ID[mother_index]+'_phipi_'+l_flv[l_index]+l_flv[l_index]+'_'+branch_name+'.root'
if os.path.exists(FILE_PATH):
print('Overwriting root file at'+FILE_PATH)
mode = 'recreate'
else:
print('Saving root file at'+FILE_PATH)
mode = 'create'
n_evts=len(tuple_dict[tuple_dict.keys()[0]])
branches=return_branches(data_index=data_index, mother_index=mother_index, l_index=l_index)
branches_updated=branches+[branch_name]
tuple_array= np.array(
[
tuple( tuple_dict[label][k] for label in branches_updated )
for k in range(n_evts)
],
dtype=[(label, np.float32) for label in branches_updated]
)
rn.array2root(tuple_array,
filename=FILE_PATH,
treename='DecayTree',
mode=mode,
)
return
Davide Lancierini