bcands = nf['BTommm']
hlt = nf['HLT']
gen = nf['GenPart']
bcands['event'] = nf['event']
bcands['run'] = nf['run']
bcands['luminosityBlock'] = nf['luminosityBlock']
bcands['l_xy_sig'] = bcands.l_xy / np.sqrt(bcands.l_xy_unc)
bcands['pv'] = nf['PV']
#number of events processed
nprocessedDataset += hlt.shape[0]
nprocessedAll += hlt.shape[0]
#add muon infos
mu1 = JaggedCandidateArray.zip({
n: muons[bcands['l1Idx']][n]
for n in muons[bcands['l1Idx']].columns
})
mu2 = JaggedCandidateArray.zip({
n: muons[bcands['l2Idx']][n]
for n in muons[bcands['l2Idx']].columns
})
k = JaggedCandidateArray.zip({
n: muons[bcands['kIdx']][n]
for n in muons[bcands['kIdx']].columns
})
# the MC ONia needs a specific treatment for gen, because there are cases in which thegen collection is zero! And in these cases we can not directly add the gen branch to the muons branch. So, consdiering that to use the MC Onia we need an additional requirement: that the third muon iha pdgId==+-13, we can mask the events in which k.genPrtIdx==-1
if (dataset == args.mc_x):
mask = (k.genPartIdx != -1)
mu1_new = mu1[mask]
mu2_new = mu2[mask]
bcands['is_psi2s_mu'] = nf['DecayFlag_is_psi2s_mu']
bcands['is_psi2s_tau'] = nf['DecayFlag_is_psi2s_tau']
bcands['is_chic0_mu'] = nf['DecayFlag_is_chic0_mu']
bcands['is_chic1_mu'] = nf['DecayFlag_is_chic1_mu']
bcands['is_chic2_mu'] = nf['DecayFlag_is_chic2_mu']
bcands['is_hc_mu'] = nf['DecayFlag_is_hc_mu']
bcands['is_jpsi_3pi'] = nf['DecayFlag_is_jpsi_3pi']
bcands['is_jpsi_hc'] = nf['DecayFlag_is_jpsi_hc']
#number of events processed
nprocessedDataset += hlt.shape[0]
nprocessedAll+=hlt.shape[0]
#add muon infos
mu1 = JaggedCandidateArray.zip({n: muons[bcands['mu1Idx']][n] for n in muons[bcands['mu1Idx']].columns})
mu2 = JaggedCandidateArray.zip({n: muons[bcands['mu2Idx']][n] for n in muons[bcands['mu2Idx']].columns})
if channel == 'BTo3Mu':
k = JaggedCandidateArray.zip({n: muons[bcands['kIdx']][n] for n in muons[bcands['kIdx']].columns})
else:
tracks = nf['ProbeTracks']
if channel == 'BTo2Mu3P':
pi1 = JaggedCandidateArray.zip({n: tracks[bcands['pi1Idx']][n] for n in tracks[bcands['pi1Idx']].columns})
pi2 = JaggedCandidateArray.zip({n: tracks[bcands['pi2Idx']][n] for n in tracks[bcands['pi2Idx']].columns})
pi3 = JaggedCandidateArray.zip({n: tracks[bcands['pi3Idx']][n] for n in tracks[bcands['pi3Idx']].columns})
if channel == 'BTo2MuP' or channel == 'BTo2MuK':
k = JaggedCandidateArray.zip({n: tracks[bcands['kIdx']][n] for n in tracks[bcands['kIdx']].columns})
# the MC ONia needs a specific treatment for gen, because there are cases in which thegen collection is zero! And in these cases we can not directly add the gen branch to the muons branch. So, consdiering that to use the MC Onia we need an additional requirement: that the third muon iha pdgId==+-13, we can mask the events in which k.genPrtIdx==-1