def dr23(pf):
mu2_p4 = TLorentzVectorArray.from_ptetaphim(pf.mu2pt, pf.mu2eta, pf.mu2phi,
pf.mu2mass)
mu_p4 = TLorentzVectorArray.from_ptetaphim(pf.kpt, pf.keta, pf.kphi,
pf.kmass)
pf.copy()
pf['dr13'] = mu_p4.delta_r(mu2_p4)
return pf
def inv_mass(pf, k= True):
print("Adding the new invariant mass column...")
if (k==True):
k_p4= TLorentzVectorArray.from_ptetaphim(pf.kpt,pf.keta,pf.kphi,0.493677)
else:
k_p4= TLorentzVectorArray.from_ptetaphim(pf.kpt,pf.keta,pf.kphi,pf.kmass)
mu1_p4= TLorentzVectorArray.from_ptetaphim(pf.mu1pt,pf.mu1eta,pf.mu1phi,pf.mu1mass)
mu2_p4= TLorentzVectorArray.from_ptetaphim(pf.mu2pt,pf.mu2eta,pf.mu2phi,pf.mu2mass)
sum2 = mu1_p4 + mu2_p4
pf['inv_mass']= (k_p4+sum2).mass
def DR_jpsimu(pf):
print("Adding DR between jpsi and mu branch...")
mu1_p4 = TLorentzVectorArray.from_ptetaphim(pf.mu1pt,pf.mu1eta,pf.mu1phi,pf.mu1mass)
mu2_p4 = TLorentzVectorArray.from_ptetaphim(pf.mu2pt,pf.mu2eta,pf.mu2phi,pf.mu2mass)
jpsi_p4= mu1_p4 + mu2_p4
# jpsi_p4 = TLorentzVectorArray.from_ptetaphim((pf.mu1pt+pf.mu2pt),(pf.mu1eta,+pf.mu2eta),(pf.mu1phi+pf.mu2phi),(pf.mu1mass+pf.mu2mass))
mu_p4 = TLorentzVectorArray.from_ptetaphim(pf.kpt,pf.keta,pf.kphi,pf.kmass)
# print(jpsi_p4.delta_r(mu_p4))
pf.copy()
pf['DR_jpsimu'] = jpsi_p4.delta_r(mu_p4)
return pf
def jpsi_branches(pf):
print("Adding jpsi four momentum branches...")
mu1_p4 = TLorentzVectorArray.from_ptetaphim(pf.mu1pt,pf.mu1eta,pf.mu1phi,pf.mu1mass)
mu2_p4 = TLorentzVectorArray.from_ptetaphim(pf.mu2pt,pf.mu2eta,pf.mu2phi,pf.mu2mass)
jpsi_p4= mu1_p4 + mu2_p4
pf = pf.copy()
pf['jpsi_pt'] = jpsi_p4.pt
pf['jpsi_eta'] = jpsi_p4.eta
pf['jpsi_phi'] = jpsi_p4.phi
pf['jpsi_mass'] = jpsi_p4.mass
return pf
def _make_sv(self, table):
data = {}
all_svp4 = TLorentzVectorArray.from_ptetaphim(
table[self.sv_branch + '_pt'],
table[self.sv_branch + '_eta'],
table[self.sv_branch + '_phi'],
table[self.sv_branch + '_mass'],
)
jet_cross_sv = self.jetp4.cross(all_svp4, nested=True)
match = jet_cross_sv.i0.delta_r2(jet_cross_sv.i1) < self.jet_r2
def sv(var_name):
sv_arr = table[self.sv_branch + '_' + var_name]
return self.jetp4.eta.cross(sv_arr, nested=True).unzip()[1][match]
svp4 = TLorentzVectorArray.from_ptetaphim(sv('pt'), sv('eta'),
sv('phi'), sv('mass'))
data['sv_phirel'] = svp4.delta_phi(self.jetp4)
data['sv_etarel'] = self.eta_sign * (svp4.eta - self.jetp4.eta)
data['sv_abseta'] = np.abs(svp4.eta)
data['sv_mass'] = svp4.mass
data['sv_pt_log'] = np.log(svp4.pt)
data['sv_mask'] = data['sv_pt_log'].ones_like()
data['sv_ntracks'] = sv('ntracks')
data['sv_normchi2'] = sv('chi2')
data['sv_dxy'] = sv('dxy')
data['sv_dxysig'] = sv('dxySig')
data['sv_d3d'] = sv('dlen')
data['sv_d3dsig'] = sv('dlenSig')
data['sv_costhetasvpv'] = -np.cos(sv('pAngle'))
dxysig = sv('dxySig')
dxysig.content[~np.isfinite(dxysig.content)] = 0
pos = dxysig.argsort()
for k in data:
data[k] = data[k][pos]
self._finalize_data(data)
self.data.update(data)
def mcor(pf):
#https://cds.cern.ch/record/2697350/files/1910.13404.pdf
#only for bto3mu and bto2mutrk
print("Adding mcor variable...")
b_dir_vec = TVector3Array(pf.jpsivtx_vtx_x - pf.pv_x,pf.jpsivtx_vtx_y - pf.pv_y,pf.jpsivtx_vtx_z - pf.pv_z )
b_dir = b_dir_vec/np.sqrt(b_dir_vec.mag2)
jpsimu_p4 = TLorentzVectorArray.from_ptetaphim(pf.Bpt,pf.Beta,pf.Bphi,pf.Bmass)
jpsimu_p3 = jpsimu_p4.p3
p_parallel = jpsimu_p3.dot(b_dir)
p_perp = np.sqrt(jpsimu_p3.mag2 - p_parallel * p_parallel)
mcor = np.sqrt(pf.Bmass * pf.Bmass + p_perp* p_perp) + p_perp
pf['mcor'] = mcor
return pf
def convert(self, table):
self.data = {}
self.jetp4 = TLorentzVectorArray.from_ptetaphim(
table[self.fatjet_branch + '_pt'],
table[self.fatjet_branch + '_eta'],
table[self.fatjet_branch + '_phi'],
table[self.fatjet_branch + '_mass'],
)
self.eta_sign = self.jetp4.eta.ones_like()
self.eta_sign[self.jetp4.eta <= 0] = -1
self._make_pfcands(table)
self._make_sv(table)
self.data['_jetp4'] = self.jetp4
return self.data
def ana(files,returnplots=False):
#%%################
# Plots and Setup #
###################
## Define what pdgId we expect the A to have
Aid = 9000006
## How many resolved jets we want to target with our analysis
resjets = 4
#Aid = 36
## Make a dictionary of histogram objects
bjplots = {}
for i in range(1,5):
bjplots.update({
"s_beta"+str(i): Hist(33 ,(-3.3,3.3) ,'GEN b '+str(i)+' Eta (ranked by pT)','Events','upplots/s_beta'+str(i)),
"s_bpT"+str(i): Hist(60 ,(0,120) ,'GEN pT of b '+str(i)+' (ranked by pT)','Events','upplots/s_bpT'+str(i)),
"s_bjetpT"+str(i): Hist(60 ,(0,120) ,'Matched RECO jet '+str(i)+' pT (ranked by b pT)','Events','upplots/s_RjetpT'+str(i)),
"s_bjeteta"+str(i): Hist(33 ,(-3.3,3.3) ,'Matched RECO jet '+str(i)+' Eta (ranked by b pT)','Events','upplots/s_Rjeteta'+str(i)),
"s_bjdR"+str(i): Hist(90 ,(0,3) ,'GEN b '+str(i)+' (ranked by pT) to matched jet dR','Events','upplots/s_bjdR'+str(i))
})
plots = {
"HpT": Hist(60 ,(0,320) ,'GEN Higgs pT','Events','upplots/HpT'),
#"HAdR": Hist(100,(0,2) ,'GEN Higgs to A dR','Events','upplots/HAdR'),
#'HAdeta': Hist(66 ,(-3.3,3.3) ,'GEN Higgs to A deta','Events','upplots/HAdeta'),
#'HAdphi': Hist(66 ,(-3.3,3.3) ,'GEN Higgs to A dphi','Events','upplots/HAdphi'),
"A1pT": Hist(80 ,(0,160) ,'Highest GEN A pT','Events','upplots/A1pT'),
"A2pT": Hist(80 ,(0,160) ,'Lowest GEN A pT','Events','upplots/A2pT'),
"AdR": Hist(50 ,(0,5) ,'GEN A1 to A2 dR','Events','upplots/AdR'),
"bdRA1": Hist(50 ,(0,5) ,'GEN dR between highest pT A child bs','Events','upplots/bdRA1'),
"bdRA2": Hist(50 ,(0,5) ,'GEN dR between lowest pT A child bs','Events','upplots/bdRA2'),
"bdetaA1": Hist(34 ,(0,3.4) ,'GEN |deta| between highest-A child bs','Events','upplots/bdetaA1'),
"bdetaA2": Hist(34 ,(0,3.4) ,'GEN |deta| between lowest-A child bs','Events','upplots/bdetaA2'),
"bdphiA1": Hist(34 ,(0,3.4) ,'GEN |dphi| between highest-A child bs','Events','upplots/bdphiA1'),
"bdphiA2": Hist(34 ,(0,3.4) ,'GEN |dphi| between lowest-A child bs','Events','upplots/bdphiA2'),
"bphi": Hist(66 ,(-3.3,3.3) ,'GEN b Phi','Events','upplots/bphi'),
"bjdR": Hist(100,(0,2) ,'All GEN bs to matched jet dR','Events','upplots/bjdR'),
"RjetpT": Hist(100,(0,100) ,'RECO matched jet pT','Events','upplots/RjetpT'),
"Rjeteta": Hist(66 ,(-3.3,3.3) ,'RECO matched jet eta','Events','upplots/Rjeteta'),
"RjetCSVV2":Hist(140 ,(-12,2) ,'RECO matched jet btagCSVV2 score','events','upplots/RjetCSVV2'),
"RjetDeepB":Hist(40 ,(-2.5,1.5) ,'RECO matched jet btagDeepB score','events','upplots/RjetDeepB'),
"RjetDeepFB" :Hist(24 ,(0,1.2) ,'RECO matched jet btagDeepFlavB score','events','upplots/RjetDeepFB'),
"RA1pT": Hist(80 ,(0,160) ,'pT of RECO A1 objects constructed from matched jets','Events','upplots/RA1pT'),
"RA2pT": Hist(80 ,(0,160) ,'pT of RECO A2 objects constructed from matched jets','Events','upplots/RA2pT'),
"RA1mass": Hist(40 ,(0,80) ,'reconstructed mass of A1 objects from matched jets','Events','upplots/RA1mass'),
"RA2mass": Hist(40 ,(0,80) ,'reconstructed mass of A2 objects from matched jets','Events','upplots/RA2mass'),
"RA1dR": Hist(50 ,(0,5) ,'dR between jet children of reconstructed A1 object','Events','upplots/RA1dR'),
"RA2dR": Hist(50 ,(0,5) ,'dR between jet children of reconstructed A2 object','Events','upplots/RA2dR'),
"RA1deta": Hist(33 ,(0,3.3) ,'|deta| between jet children of reconstructed A1 object','Events','upplots/RA1deta'),
"RA2deta": Hist(33 ,(0,3.3) ,'|deta| between jet children of reconstructed A2 object','Events','upplots/RA2deta'),
"RA1dphi": Hist(33 ,(0,3.3) ,'|dphi| between jet children of reconstructed A1 object','Events','upplots/RA1dphi'),
"RA2dphi": Hist(33 ,(0,3.3) ,'|dphi| between jet children of reconstructed A2 object','Events','upplots/RA2dphi'),
"RHmass": Hist(80 ,(0,160) ,'reconstructed mass of Higgs object from reconstructed As','Events','upplots/RHmass'),
"RHpT": Hist(100,(0,200) ,'pT of reconstructed higgs object from reconstructed As','Events','upplots/RHpT'),
"RHdR": Hist(50 ,(0,5) ,'dR between A children of reconstructed higgs object','Events','upplots/RHdR'),
"RHdeta": Hist(33 ,(0,3.3) ,'|deta| between A children of reconstructed higgs object','Events','upplots/RHdeta'),
"RHdphi": Hist(33 ,(0,3.3) ,'|dphi| between A children of reconstructed higgs object','Events','upplots/RHdphi'),
##
"RalljetpT": Hist(100,(0,100),'All RECO jet pT','Events','upplots/RalljetpT'),
"bjdRvlogbpT1": Hist2d([80,200],[[0,8],[0,2]],'log2(GEN b pT)','dR from 1st pT GEN b to matched RECO jet','upplots/bjdRvlogbpT1'),
"bjdRvlogbpT2": Hist2d([80,200],[[0,8],[0,2]],'log2(GEN b pT)','dR from 2nd pT GEN b to matched RECO jet','upplots/bjdRvlogbpT2'),
"bjdRvlogbpT3": Hist2d([80,200],[[0,8],[0,2]],'log2(GEN b pT)','dR from 3rd pT GEN b to matched RECO jet','upplots/bjdRvlogbpT3'),
"bjdRvlogbpT4": Hist2d([80,200],[[0,8],[0,2]],'log2(GEN b pT)','dR from 4th pT GEN b to matched RECO jet','upplots/bjdRvlogbpT4'),
"jetoverbpTvlogbpT1": Hist2d([60,40],[[2,8],[0,4]],'log2(GEN b pT)','RECO jet pT / 1st GEN b pT for matched jets','upplots/jetoverbpTvlogbpT1'),
"jetoverbpTvlogbpT2": Hist2d([60,40],[[2,8],[0,4]],'log2(GEN b pT)','RECO jet pT / 2nd GEN b pT for matched jets','upplots/jetoverbpTvlogbpT2'),
"jetoverbpTvlogbpT3": Hist2d([60,40],[[2,8],[0,4]],'log2(GEN b pT)','RECO jet pT / 3rd GEN b pT for matched jets','upplots/jetoverbpTvlogbpT3'),
"jetoverbpTvlogbpT4": Hist2d([60,40],[[2,8],[0,4]],'log2(GEN b pT)','RECO jet pT / 4th GEN b pT for matched jets','upplots/jetoverbpTvlogbpT4'),
"npassed": Hist(1 ,(0.5,1.5) ,'','Number of events that passed cuts', 'upplots/npassed'),
"genAmass": Hist(40 ,(0,80) ,'GEN mass of A objects','Events','upplots/Amass_g'),
"cutAmass": Hist(40 ,(0,80) ,'GEN mass of A objects that pass cuts','Events','upplots/Amass_c')
}
for plot in bjplots:
bjplots[plot].title = files[0]
for plot in plots:
plots[plot].title = files[0]
## Create an internal figure for pyplot to write to
plt.figure(1)
## Loop over input files
for fnum in range(len(files)):
#####################
# Loading Variables #
#####################
print('Opening '+files[fnum])
## Open our file and grab the events tree
f = uproot.open(files[fnum])#'nobias.root')
events = f.get('Events')
pdgida = events.array('GenPart_pdgId')
paridxa = events.array('GenPart_genPartIdxMother')
parida = pdgida[paridxa]
bs = PhysObj('bs')
## Removes all particles that do not have A parents
## from the GenPart arrays, then removes all particles
## that are not bs after resizing the pdgid array to be a valid mask
bs.oeta = pd.DataFrame(events.array('GenPart_eta')[abs(parida)==Aid][abs(pdgida)[abs(parida)==Aid]==5]).rename(columns=inc)
bs.ophi = pd.DataFrame(events.array('GenPart_phi')[abs(parida)==Aid][abs(pdgida)[abs(parida)==Aid]==5]).rename(columns=inc)
bs.opt = pd.DataFrame(events.array('GenPart_pt' )[abs(parida)==Aid][abs(pdgida)[abs(parida)==Aid]==5]).rename(columns=inc)
## Test b order corresponds to As
testbs = pd.DataFrame(events.array('GenPart_genPartIdxMother')[abs(parida)==Aid][abs(pdgida)[abs(parida)==Aid]==5]).rename(columns=inc)
## The first term checks b4 has greater idx than b1, the last two check that the bs are paired
if ((testbs[4]-testbs[1]).min() <= 0) or ((abs(testbs[2]-testbs[1]) + abs(testbs[4])-testbs[3]).min() != 0):
print('b to A ordering violated - time to do it the hard way')
sys.exit()
As = PhysObj('As')
As.oeta = pd.DataFrame(events.array('GenPart_eta')[abs(parida)==25][abs(pdgida)[abs(parida)==25]==Aid]).rename(columns=inc)
As.ophi = pd.DataFrame(events.array('GenPart_phi')[abs(parida)==25][abs(pdgida)[abs(parida)==25]==Aid]).rename(columns=inc)
As.opt = pd.DataFrame(events.array('GenPart_pt' )[abs(parida)==25][abs(pdgida)[abs(parida)==25]==Aid]).rename(columns=inc)
As.omass =pd.DataFrame(events.array('GenPart_mass')[abs(parida)==25][abs(pdgida)[abs(parida)==25]==Aid]).rename(columns=inc)
higgs = PhysObj('higgs')
higgs.eta = pd.DataFrame(events.array('GenPart_eta')[abs(parida)!=25][abs(pdgida)[abs(parida)!=25]==25]).rename(columns=inc)
higgs.phi = pd.DataFrame(events.array('GenPart_phi')[abs(parida)!=25][abs(pdgida)[abs(parida)!=25]==25]).rename(columns=inc)
higgs.pt = pd.DataFrame(events.array('GenPart_pt' )[abs(parida)!=25][abs(pdgida)[abs(parida)!=25]==25]).rename(columns=inc)
jets = PhysObj('jets')
jets.eta= pd.DataFrame(events.array('Jet_eta')).rename(columns=inc)
jets.phi= pd.DataFrame(events.array('Jet_phi')).rename(columns=inc)
jets.pt = pd.DataFrame(events.array('Jet_pt')).rename(columns=inc)
jets.mass=pd.DataFrame(events.array('Jet_mass')).rename(columns=inc)
jets.CSVV2 = pd.DataFrame(events.array('Jet_btagCSVV2')).rename(columns=inc)
jets.DeepB = pd.DataFrame(events.array('Jet_btagDeepB')).rename(columns=inc)
jets.DeepFB= pd.DataFrame(events.array('Jet_btagDeepFlavB')).rename(columns=inc)
print('Processing ' + str(len(bs.oeta)) + ' events')
## Figure out how many bs and jets there are
nb = bs.oeta.shape[1]
njet= jets.eta.shape[1]
na = As.oeta.shape[1]
if na != 2:
print("More than two As per event, found "+str(na)+", halting")
sys.exit()
## Create sorted versions of A values by pt
for prop in ['eta','phi','pt','mass']:
As[prop] = pd.DataFrame()
for i in range(1,3):
As[prop][i] = As['o'+prop][As.opt.rank(axis=1,ascending=False,method='first')==i].max(axis=1)
## Clean up original ordered dataframes; we don't really need them
#del As['o'+prop]
## Reorder out b dataframes to match sorted A parents
tframe = pd.DataFrame()
tframe[1] = (As.opt.rank(axis=1,ascending=False,method='first')==1)[1]
tframe[2] = (As.opt.rank(axis=1,ascending=False,method='first')==1)[1]
tframe[3] = (As.opt.rank(axis=1,ascending=False,method='first')==1)[2]
tframe[4] = (As.opt.rank(axis=1,ascending=False,method='first')==1)[2]
for prop in ['eta','phi','pt']:
bs[prop] = pd.DataFrame()
bs[prop][1] = bs['o'+prop][tframe][1].dropna().append(bs['o'+prop][tframe][3].dropna()).sort_index()
bs[prop][2] = bs['o'+prop][tframe][2].dropna().append(bs['o'+prop][tframe][4].dropna()).sort_index()
bs[prop][3] = bs['o'+prop][~tframe][1].dropna().append(bs['o'+prop][~tframe][3].dropna()).sort_index()
bs[prop][4] = bs['o'+prop][~tframe][2].dropna().append(bs['o'+prop][~tframe][4].dropna()).sort_index()
## Clean up original ordered dataframes; we don't really need them.
#del bs['o'+prop]
## Sort our b dataframes in descending order of pt
for prop in ['spt','seta','sphi']:
bs[prop] = pd.DataFrame()
#bs.spt, bs.seta, bs.sphi = pd.DataFrame(), pd.DataFrame(), pd.DataFrame()
for i in range(1,nb+1):
bs[prop][i] = bs[prop[1:]][bs.pt.rank(axis=1,ascending=False,method='first')==i].max(axis=1)
#bs.seta[i] = bs.eta[bs.pt.rank(axis=1,ascending=False,method='first')==i].max(axis=1)
#bs.sphi[i] = bs.phi[bs.pt.rank(axis=1,ascending=False,method='first')==i].max(axis=1)
plots['genAmass'].dfill(As.mass)
ev = Event(bs,jets,As,higgs)
jets.cut(jets.pt>0)
bs.cut(bs.pt>0)
ev.sync()
##############################
# Processing and Calculation #
##############################
## Create our dR dataframe by populating its first column and naming it accordingly
jbdr2 = pd.DataFrame(np.power(jets.eta[1]-bs.eta[1],2) + np.power(jets.phi[1]-bs.phi[1],2)).rename(columns={1:'Jet 1 b 1'})
sjbdr2= pd.DataFrame(np.power(jets.eta[1]-bs.seta[1],2) + np.power(jets.phi[1]-bs.sphi[1],2)).rename(columns={1:'Jet 1 b 1'})
## Loop over jet x b combinations
jbstr = []
for j in range(1,njet+1):
for b in range(1,nb+1):
## Make our column name
jbstr.append("Jet "+str(j)+" b "+str(b))
if (j+b==2):
continue
## Compute and store the dr of the given b and jet for every event at once
jbdr2[jbstr[-1]] = pd.DataFrame(np.power(jets.eta[j]-bs.eta[b],2) + np.power(jets.phi[j]-bs.phi[b],2))
sjbdr2[jbstr[-1]]= pd.DataFrame(np.power(jets.eta[j]-bs.seta[b],2) + np.power(jets.phi[j]-bs.sphi[b],2))
## Create a copy array to collapse in jets instead of bs
blist = []
sblist = []
for b in range(nb):
blist.append(np.sqrt(jbdr2.filter(like='b '+str(b+1))))
blist[b] = blist[b][blist[b].rank(axis=1,method='first') == 1]
blist[b] = blist[b].rename(columns=lambda x:int(x[4:6]))
sblist.append(np.sqrt(sjbdr2.filter(like='b '+str(b+1))))
sblist[b] = sblist[b][sblist[b].rank(axis=1,method='first') == 1]
sblist[b] = sblist[b].rename(columns=lambda x:int(x[4:6]))
## Trim resolved jet objects
if resjets==3:
for i in range(nb):
for j in range(nb):
if i != j:
blist[i] = blist[i][np.logical_not(blist[i] > blist[j])]
blist[i] = blist[i][blist[i]<0.4]
## Cut our events to only resolved 4jet events with dR<0.4
rjets = blist[0][blist[0]<0.4].fillna(0)
for i in range(1,4):
rjets = np.logical_or(rjets,blist[i][blist[i]<0.4].fillna(0))
rjets = rjets.sum(axis=1)
rjets = rjets[rjets==resjets].dropna()
jets.trimTo(rjets)
ev.sync()
#############################
# Constructing RECO objects #
#############################
for prop in ['bpt','beta','bphi','bmass']:
jets[prop] = pd.DataFrame()
for i in range(nb):
jets[prop][i+1] = jets[prop[1:]][blist[i]>0].max(axis=1)
jets.bdr = pd.DataFrame()
for i in range(nb):
jets.bdr[i+1] = blist[i][blist[i]>0].max(axis=1)
ev.sync()
if resjets==3:
pidx = [2,1,4,3]
for prop in ['bpt','beta','bphi','bmass']:
jets[prop]['merged'], jets[prop]['missing'] = (jets[prop][1]==jets[prop][3]),(jets[prop][1]==jets[prop][3])
for i in range(1,nb+1):
jets[prop]['merged']=jets[prop]['merged']+jets[prop].fillna(0)[i][(jets.bmass[i]>=15) & (jets.bmass[i]+jets.bmass[pidx[i-1]]==jets.bmass[i])]
jets[prop]['missing']=jets[prop]['missing']+jets[prop].fillna(0)[i][(jets.bmass[i]<15) & (jets.bmass[i]+jets.bmass[pidx[i-1]]==jets.bmass[i])]
#jets[prop][i] = jets[prop][i]+(0*jets[prop][pidx])
bvec = []
for i in range(1,nb+1):
bvec.append(TLorentzVectorArray.from_ptetaphim(jets.bpt[i],jets.beta[i],jets.bphi[i],jets.bmass[i]))
avec = []
for i in range(0,nb,2):
avec.append(bvec[i]+bvec[i+1])
for prop in ['apt','aeta','aphi','amass']:
jets[prop] = pd.DataFrame()
for i in range(na):
jets.apt[i+1] = avec[i].pt
jets.aeta[i+1] = avec[i].eta
jets.aphi[i+1] = avec[i].phi
jets.amass[i+1]= avec[i].mass
for prop in ['apt','aeta','aphi','amass']:
jets[prop].index = jets.pt.index
hvec = [avec[0]+avec[1]]
for prop in ['hpt','heta','hphi','hmass']:
jets[prop] = pd.DataFrame()
jets.hpt[1] = hvec[0].pt
jets.heta[1] = hvec[0].eta
jets.hphi[1] = hvec[0].phi
jets.hmass[1]= hvec[0].mass
for prop in ['hpt','heta','hphi','hmass']:
jets[prop].index = jets.eta.index
################
# Filling Data #
################
for i in range(4):
plots['bjdRvlogbpT'+str(i+1)].dfill(np.log2(bs.spt[[i+1]]),bs.trim(sblist[i]))
plots['RjetCSVV2'].dfill(jets.CSVV2[blist[i]>0])
plots['RjetDeepB'].dfill(jets.DeepB[blist[i]>0])
plots['RjetDeepFB'].dfill(jets.DeepFB[blist[i]>0])
yval = np.divide(jets.pt[sblist[i]>0].melt(value_name=0).drop('variable',axis=1).dropna().reset_index(drop=True)[0],bs.spt[[i+1]].dropna().reset_index(drop=True)[i+1])
xval = np.log2(bs.spt[[i+1]]).melt(value_name=0).drop('variable',axis=1).dropna().reset_index(drop=True)[0]
plots['jetoverbpTvlogbpT'+str(i+1)].fill(xval,yval)
bjplots['s_bpT'+str(i+1)].dfill(bs.spt[[i+1]])
bjplots['s_beta'+str(i+1)].dfill(bs.seta[[i+1]])
bjplots['s_bjetpT'+str(i+1)].dfill(jets.pt[sblist[i]>0])
bjplots['s_bjeteta'+str(i+1)].dfill(jets.eta[sblist[i]>0])
bjplots['s_bjdR'+str(i+1)].dfill(sblist[i][sblist[i]!=0])
plots['HpT'].dfill(higgs.pt)
plots['A1pT'].fill(As.pt[1])
plots['A2pT'].fill(As.pt[2])
plots['AdR'].fill(np.sqrt(np.power(As.eta[2]-As.eta[1],2) + np.power(As.phi[2]-As.phi[1],2)))
plots['bdRA1'].fill(np.sqrt(np.power(bs.eta[2]-bs.eta[1],2) + np.power(bs.phi[2]-bs.phi[1],2)))
plots['bdRA2'].fill(np.sqrt(np.power(bs.eta[4]-bs.eta[3],2) + np.power(bs.phi[4]-bs.phi[3],2)))
plots['bdetaA1'].fill(abs(bs.eta[2]-bs.eta[1]))
plots['bdetaA2'].fill(abs(bs.eta[4]-bs.eta[3]))
plots['bdphiA1'].fill(abs(bs.phi[2]-bs.phi[1]))
plots['bdphiA2'].fill(abs(bs.phi[4]-bs.phi[3]))
plots['bphi'].dfill(bs.phi)
plots['RalljetpT'].dfill(jets.pt)
plots['bjdR'].dfill(jets.bdr)
plots['RjetpT'].dfill(jets.bpt)
plots['Rjeteta'].dfill(jets.beta)
for i in range(1,3):
plots['RA'+str(i)+'pT' ].fill(jets.apt[i])
plots['RA'+str(i)+'mass'].fill(jets.amass[i])
#lots['RA'+str(i)+'deta'].fill(abs(jets.beta[2*i]-jets.beta[(2*i)-1]))
plots['RA'+str(i)+'dR' ].fill(np.sqrt(np.power(jets.beta[2*i]-jets.beta[(2*i)-1],2)+np.power(jets.bphi[2*i]-jets.bphi[(2*i)-1],2)))
plots['RA'+str(i)+'deta'].fill(abs(jets.beta[2*i]-jets.beta[(2*i)-1]))
plots['RA'+str(i)+'dphi'].fill(abs(jets.bphi[2*i]-jets.bphi[(2*i)-1]))
plots['RHpT' ].fill(jets.hpt[1])
plots['RHmass'].fill(jets.hmass[1])
plots['RHdR' ].fill(np.sqrt(np.power(jets.aeta[2]-jets.aeta[1],2)+np.power(jets.aphi[2]-jets.aphi[1],2)))
plots['RHdeta'].fill(abs(jets.aeta[2]-jets.aeta[1]))
plots['RHdphi'].fill(abs(jets.aphi[2]-jets.aphi[1]))
plots['npassed'].fill(jets.hpt[1]/jets.hpt[1])
plots['cutAmass'].dfill(As.mass)
############
# Plotting #
############
plt.clf()
#plots.pop('bjdR').plot(logv=True)
for i in range(1,5):
bjplots.pop('s_bjdR'+str(i)).plot(logv=True)
for p in plots:
plots[p].plot()
for p in bjplots:
bjplots[p].plot()
#%%
if returnplots==True:
return plots
else:
sys.exit()
def transform(self, jet):
"""
precondition - jet is a JaggedCandidateArray with additional attributes:
- 'ptRaw'
- 'massRaw'
xformer = JetTransformer(name1=corrL1,...)
xformer.transform(jet)
postcondition - jet.pt, jet.mass, jet.p4 are updated to represent the corrected jet
based on the input correction set
"""
if not isinstance(jet, JaggedCandidateArray):
raise Exception('Input data must be a JaggedCandidateArray!')
if ('ptRaw' not in jet.columns or 'massRaw' not in jet.columns):
raise Exception(
'Input JaggedCandidateArray must have "ptRaw" & "massRaw"!')
#initialize the jet momenta to raw values
_update_jet_ptm(1.0, jet, fromRaw=True)
#below we work in numpy arrays, JaggedCandidateArray knows how to convert them
args = {
key: getattr(jet, _signature_map[key]).content
for key in self._jec.signature
}
jec = self._jec.getCorrection(**args)
_update_jet_ptm(jec, jet, fromRaw=True)
junc_up = np.ones_like(jec)
junc_down = np.ones_like(jec)
if self._junc is not None:
args = {
key: getattr(jet, _signature_map[key]).content
for key in self._junc.signature
}
junc = self._junc.getUncertainty(**args)
junc_up = junc[:, 0]
junc_down = junc[:, 1]
#if there's a jer and sf to apply we have to update the momentum too
#right now only use stochastic smearing
if self._jer is not None and self._jersf is not None:
args = {
key: getattr(jet, _signature_map[key]).content
for key in self._jer.signature
}
jer = self._jer.getResolution(**args)
args = {
key: getattr(jet, _signature_map[key]).content
for key in self._jersf.signature
}
jersf = self._jersf.getScaleFactor(**args)
jsmear_cen = 1. + np.sqrt(
jersf[:, 0]**2 - 1.0) * jer * np.random.normal(size=jer.size)
jsmear_up = 1. + np.sqrt(
jersf[:, 1]**2 - 1.0) * jer * np.random.normal(size=jer.size)
jsmear_down = 1. + np.sqrt(
jersf[:, -1]**2 - 1.0) * jer * np.random.normal(size=jer.size)
# need to apply up and down jer-smear before applying central correction
jet.add_attributes(pt_jer_up=jsmear_up * jet.pt.content,
mass_jer_up=jsmear_up * jet.mass.content,
pt_jer_down=jsmear_down * jet.pt.content,
mass_jer_down=jsmear_down * jet.mass.content)
#finally, update the central value
_update_jet_ptm(jsmear_cen, jet)
# have to apply central jersf before calculating junc
if self._junc is not None:
jet.add_attributes(pt_jes_up=junc_up * jet.pt.content,
mass_jes_up=junc_up * jet.mass.content,
pt_jes_down=junc_down * jet.pt.content,
mass_jes_down=junc_down * jet.mass.content)
#hack to update the jet p4, we have the fully updated pt and mass here
jet._content._contents['p4'] = TLorentzVectorArray.from_ptetaphim(
jet.pt.content, jet.eta.content, jet.phi.content, jet.mass.content)
def transform(self, jet, met=None):
"""
precondition - jet is a JaggedCandidateArray with additional attributes:
- 'ptRaw'
- 'massRaw'
xformer = JetTransformer(name1=corrL1,...)
xformer.transform(jet)
postcondition - jet.pt, jet.mass, jet.p4 are updated to represent the corrected jet
based on the input correction set
"""
if not isinstance(jet, JaggedCandidateArray):
raise Exception('Input data must be a JaggedCandidateArray!')
if ('ptRaw' not in jet.columns or 'massRaw' not in jet.columns):
raise Exception('Input JaggedCandidateArray must have "ptRaw" & "massRaw"!')
if met is not None:
if 'p4' not in met.columns:
raise Exception('Input met must have a p4 column!')
if not isinstance(met['p4'], TLorentzVectorArray):
raise Exception('Met p4 must be a TLorentzVectorArray!')
initial_p4 = jet['p4'].copy() # keep a copy for fixing met
# initialize the jet momenta to raw values
_update_jet_ptm(1.0, jet, fromRaw=True)
# below we work in numpy arrays, JaggedCandidateArray knows how to convert them
args = {key: getattr(jet, _signature_map[key]).content for key in self._jec.signature}
jec = self._jec.getCorrection(**args)
_update_jet_ptm(jec, jet, fromRaw=True)
juncs = None
if self._junc is not None:
args = {key: getattr(jet, _signature_map[key]).content for key in self._junc.signature}
juncs = self._junc.getUncertainty(**args)
# if there's a jer and sf to apply we have to update the momentum too
# right now only use stochastic smearing
if self._jer is not None and self._jersf is not None:
args = {key: getattr(jet, _signature_map[key]).content for key in self._jer.signature}
jer = self._jer.getResolution(**args)
args = {key: getattr(jet, _signature_map[key]).content for key in self._jersf.signature}
jersf = self._jersf.getScaleFactor(**args)
jersmear = jer * np.random.normal(size=jer.size)
jsmear_cen = 1. + np.sqrt(jersf[:, 0]**2 - 1.0) * jersmear
jsmear_up = 1. + np.sqrt(jersf[:, 1]**2 - 1.0) * jersmear
jsmear_down = 1. + np.sqrt(jersf[:, -1]**2 - 1.0) * jersmear
# need to apply up and down jer-smear before applying central correction
jet.add_attributes(pt_jer_up=jsmear_up * jet.pt.content,
mass_jer_up=jsmear_up * jet.mass.content,
pt_jer_down=jsmear_down * jet.pt.content,
mass_jer_down=jsmear_down * jet.mass.content)
# finally, update the central value
_update_jet_ptm(jsmear_cen, jet)
# have to apply central jersf before calculating junc
if self._junc is not None:
for name, values in juncs:
jet.add_attributes(**{
'pt_{0}_up'.format(name): values[:, 0] * jet.pt.content,
'mass_{0}_up'.format(name): values[:, 0] * jet.mass.content,
'pt_{0}_down'.format(name): values[:, 1] * jet.pt.content,
'mass_{0}_down'.format(name): values[:, 1] * jet.mass.content
})
# hack to update the jet p4, we have the fully updated pt and mass here
jet._content._contents['p4'] = TLorentzVectorArray.from_ptetaphim(jet.pt.content,
jet.eta.content,
jet.phi.content,
jet.mass.content)
if met is None:
return
# set MET values
new_x = met['p4'].x - (initial_p4.x - jet['p4'].x).sum()
new_y = met['p4'].y - (initial_p4.y - jet['p4'].y).sum()
met.base['p4'] = TLorentzVectorArray.from_ptetaphim(
np.sqrt(new_x**2 + new_y**2), 0,
np.arctan2(new_y, new_x), 0
)
if 'MetUnclustEnUpDeltaX' in met.columns:
px_up = met['p4'].x + met['MetUnclustEnUpDeltaX']
py_up = met['p4'].y + met['MetUnclustEnUpDeltaY']
met.base['pt_UnclustEn_up'] = np.sqrt(px_up**2 + py_up**2)
met.base['phi_UnclustEn_up'] = np.arctan2(py_up, px_up)
px_down = met['p4'].x - met['MetUnclustEnUpDeltaX']
py_down = met['p4'].y - met['MetUnclustEnUpDeltaY']
met.base['pt_UnclustEn_down'] = np.sqrt(px_down**2 + py_down**2)
met.base['phi_UnclustEn_down'] = np.arctan2(py_down, px_down)
if self._junc is not None:
jets_sin = np.sin(jet['p4'].phi)
jets_cos = np.cos(jet['p4'].phi)
for name, _ in juncs:
for shift in ['up', 'down']:
px = met['p4'].x - (initial_p4.x - jet['pt_{0}_{1}'.format(name, shift)] * jets_cos).sum()
py = met['p4'].y - (initial_p4.y - jet['pt_{0}_{1}'.format(name, shift)] * jets_sin).sum()
met.base['pt_{0}_{1}'.format(name, shift)] = np.sqrt(px**2 + py**2)
met.base['phi_{0}_{1}'.format(name, shift)] = np.arctan2(py, px)
def ana(files,returnplots=False):
#%%################
# Plots and Setup #
###################
## Make a dictionary of histogram objects
plots = {
"RjetpT": Hist(100,(0,100) ,'RECO matched jet pT','Events','recplots/RjetpT'),
"Rjeteta": Hist(66 ,(-3.3,3.3) ,'RECO matched jet eta','Events','recplots/Rjeteta'),
"RjetCSVV2":Hist([0,0.1241,0.4184,0.7527,1],None,'RECO matched jet btagCSVV2 score','events','recplots/RjetCSVV2'),
#"RjetDeepB":Hist([0,0.0494,0.2770,0.7264,1],None,'RECO matched jet btagDeepB score','events','recplots/RjetDeepB'),
"RjetDeepFB":Hist([0,0.0494,0.2770,0.7264,1],None,'RECO matched jet btagDeepFlavB score','events','recplots/RjetDeepFB'),
"RA1pT": Hist(80 ,(0,160) ,'pT of RECO A1 objects constructed from matched jets','Events','recplots/RA1pT'),
"RA2pT": Hist(80 ,(0,160) ,'pT of RECO A2 objects constructed from matched jets','Events','recplots/RA2pT'),
"RA1mass": Hist(40 ,(0,80) ,'reconstructed mass of A1 objects from matched jets','Events','recplots/RA1mass'),
"RA2mass": Hist(40 ,(0,80) ,'reconstructed mass of A2 objects from matched jets','Events','recplots/RA2mass'),
"RA1dR": Hist(50 ,(0,5) ,'dR between jet children of reconstructed A1 object','Events','recplots/RA1dR'),
"RA2dR": Hist(50 ,(0,5) ,'dR between jet children of reconstructed A2 object','Events','recplots/RA2dR'),
"RA1deta": Hist(33 ,(0,3.3) ,'|deta| between jet children of reconstructed A1 object','Events','recplots/RA1deta'),
"RA2deta": Hist(33 ,(0,3.3) ,'|deta| between jet children of reconstructed A2 object','Events','recplots/RA2deta'),
"RA1dphi": Hist(33 ,(0,3.3) ,'|dphi| between jet children of reconstructed A1 object','Events','recplots/RA1dphi'),
"RA2dphi": Hist(33 ,(0,3.3) ,'|dphi| between jet children of reconstructed A2 object','Events','recplots/RA2dphi'),
"RHmass": Hist(80 ,(0,160) ,'reconstructed mass of Higgs object from reconstructed As','Events','recplots/RHmass'),
"RHpT": Hist(100,(0,200) ,'pT of reconstructed higgs object from reconstructed As','Events','recplots/RHpT'),
"RHdR": Hist(50 ,(0,5) ,'dR between A children of reconstructed higgs object','Events','recplots/RHdR'),
"RHdeta": Hist(33 ,(0,3.3) ,'|deta| between A children of reconstructed higgs object','Events','recplots/RHdeta'),
"RHdphi": Hist(33 ,(0,3.3) ,'|dphi| between A children of reconstructed higgs object','Events','recplots/RHdphi'),
##
"RalljetpT": Hist(100,(0,100),'All RECO jet pT','Events','recplots/RalljetpT'),
"npassed": Hist(1 ,(0.5,1.5) ,'','Number of events that passed cuts', 'recplots/npassed')
}
for plot in plots:
plots[plot].title = files[0]
## Create an internal figure for pyplot to write to
plt.figure(1)
## Loop over input files
for fnum in range(len(files)):
#####################
# Loading Variables #
#####################
print('Opening '+files[fnum])
## Open our file and grab the events tree
f = uproot.open(files[fnum])#'nobias.root')
events = f.get('Events')
jets = PhysObj('jets')
jets.eta= pd.DataFrame(events.array('Jet_eta')).rename(columns=inc)
jets.phi= pd.DataFrame(events.array('Jet_phi')).rename(columns=inc)
jets.pt = pd.DataFrame(events.array('Jet_pt')).rename(columns=inc)
jets.mass=pd.DataFrame(events.array('Jet_mass')).rename(columns=inc)
jets.CSVV2 = pd.DataFrame(events.array('Jet_btagCSVV2')).rename(columns=inc)
jets.DeepB = pd.DataFrame(events.array('Jet_btagDeepB')).rename(columns=inc)
jets.DeepFB= pd.DataFrame(events.array('Jet_btagDeepFlavB')).rename(columns=inc)
print('Processing ' + str(len(jets.eta)) + ' events')
## Figure out how many bs and jets there are
njet= jets.eta.shape[1]
if njet > 10:
njet = 10
ev = Event(jets)
jets.cut(jets.pt>15)
jets.cut(abs(jets.eta)<2.4)
jets.cut(jets.DeepB > 0.4184 )
ev.sync()
##############################
# Processing and Calculation #
##############################
## Create our dR dataframe by populating its first column and naming it accordingly
jjdr2 = pd.DataFrame(np.power(jets.eta[1]-jets.eta[2],2) + np.power(jets.phi[1]-jets.phi[2],2)).rename(columns={0:'Jet 1 x Jet 2'})
jjmass = pd.DataFrame(jets.mass[1] + jets.mass[2]).rename(columns={0:'Jet 1 x Jet 2'})
## Loop over jet x b combinations
jjstr = []
for j in range(1,njet+1):
for i in range(j+1,njet+1):
## Make our column name
jjstr.append("Jet "+str(j)+" x Jet "+str(i))
#if (i+j == 3):
# continue
## Compute and store the dr of the given b and jet for every event at once
jjdr2[jjstr[-1]] = pd.DataFrame(np.power(jets.eta[j]-jets.eta[i],2) + np.power(jets.phi[j]-jets.phi[i],2))
jjmass[jjstr[-1]] = pd.DataFrame(jets.mass[j] + jets.mass[i])
#if (j==i):
# jjdr2[jjstr[-1]] = jjdr2[jjstr[-1]] * np.nan
# jjmass[jjstr[-1]] = jjmass[jjstr[-1]] * np.nan
print('jjs done')
j4mass = pd.DataFrame(jets.mass[1]+jets.mass[2]+jets.mass[3]+jets.mass[4]).rename(columns={0:"J1 x J2 J3 J4"})
j4str = []
for a in range(1,njet+1):
for b in range(a+1,njet+1):
for c in range(b+1,njet+1):
for d in range(c+1,njet+1):
j4str.append("J"+str(a)+" J"+str(b)+" J"+str(c)+" J"+str(d))
#if (a+b+c+d == 10):
# continue
j4mass[j4str[-1]] = pd.DataFrame(jets.mass[a]+jets.mass[b]+jets.mass[c]+jets.mass[d])
#if (a==b or a==c or a==d or b==c or b==d or c==d):
# j4mass[j4str[-1]] = j4mass[j4str[-1]] * np.nan
print('j4s done')
## Create a copy array to collapse in jets into
drlist = []
mmlist = []
m4list = []
for j in range(njet):
drlist.append(np.sqrt(jjdr2.filter(like='Jet '+str(j+1))))
mmlist.append(jjmass.filter(like='Jet '+str(j+1)))
m4list.append(j4mass.filter(like='J'+str(j+1)))
#jlist[j] = jlist[j][jlist[j].rank(axis=1,method='first') == 1]
#drlist[j] = jlist[j].rename(columns=lambda x:int(x[4:6]))
#mmlist[j] = mmlist[j].rename(columns=lambda x:int(x[4:6]))
print('jlist done')
## Cut our events to only resolved 4jet events with dR<0.4
djets = mmlist[0][(mmlist[0]>25) & (mmlist[0]<65)]
qjets = m4list[0][(m4list[0]>90) & (m4list[0]<150)]
for i in range(1,njet):
djets = djets.combine_first(mmlist[i][(mmlist[i]>25) & (mmlist[i]<65)])
qjets = qjets.combine_first(m4list[i][(m4list[i]>90) & (m4list[i]<150)])
djets = djets / djets
qjets = qjets / qjets
djets = djets.sum(axis=1)
qjets = qjets.sum(axis=1)
djets = djets[djets>=2].dropna()
qjets = qjets[qjets>=1].dropna()
jets.trimTo(djets)
jets.trimTo(qjets)
ev.sync()
print('trimming done')
#############################
# Constructing RECO objects #
#############################
for prop in ['bpt','beta','bphi','bmass']:
jets[prop] = pd.DataFrame()
for i in range(1,5):
jets[prop][i] = jets[prop[1:]][jets.mass.rank(axis=1,method='first',ascending=False) == i].max(axis=1)
#jets.bdr = pd.DataFrame()
#for i in range(nb):
# jets.bdr[i+1] = blist[i][blist[i]>0].max(axis=1)
#ev.sync()
bvec = []
for i in range(1,5):
bvec.append(TLorentzVectorArray.from_ptetaphim(jets.bpt[i],jets.beta[i],jets.bphi[i],jets.bmass[i]))
avec = []
for i in range(0,4,2):
avec.append(bvec[i]+bvec[i+1])
for prop in ['apt','aeta','aphi','amass']:
jets[prop] = pd.DataFrame()
for i in range(2):
jets.apt[i+1] = avec[i].pt
jets.aeta[i+1] = avec[i].eta
jets.aphi[i+1] = avec[i].phi
jets.amass[i+1]= avec[i].mass
for prop in ['apt','aeta','aphi','amass']:
jets[prop].index = jets.pt.index
hvec = [avec[0]+avec[1]]
for prop in ['hpt','heta','hphi','hmass']:
jets[prop] = pd.DataFrame()
jets.hpt[1] = hvec[0].pt
jets.heta[1] = hvec[0].eta
jets.hphi[1] = hvec[0].phi
jets.hmass[1]= hvec[0].mass
for prop in ['hpt','heta','hphi','hmass']:
jets[prop].index = jets.eta.index
################
# Filling Data #
################
plots['RalljetpT'].dfill(jets.pt)
#plots['bjdR'].dfill(jets.bdr)
plots['RjetpT'].dfill(jets.bpt)
plots['Rjeteta'].dfill(jets.beta)
for i in range(1,3):
plots['RA'+str(i)+'pT' ].fill(jets.apt[i])
plots['RA'+str(i)+'mass'].fill(jets.amass[i])
#lots['RA'+str(i)+'deta'].fill(abs(jets.beta[2*i]-jets.beta[(2*i)-1]))
plots['RA'+str(i)+'dR' ].fill(np.sqrt(np.power(jets.beta[2*i]-jets.beta[(2*i)-1],2)+np.power(jets.bphi[2*i]-jets.bphi[(2*i)-1],2)))
plots['RA'+str(i)+'deta'].fill(abs(jets.beta[2*i]-jets.beta[(2*i)-1]))
plots['RA'+str(i)+'dphi'].fill(abs(jets.bphi[2*i]-jets.bphi[(2*i)-1]))
plots['RHpT' ].fill(jets.hpt[1])
plots['RHmass'].fill(jets.hmass[1])
plots['RHdR' ].fill(np.sqrt(np.power(jets.aeta[2]-jets.aeta[1],2)+np.power(jets.aphi[2]-jets.aphi[1],2)))
plots['RHdeta'].fill(abs(jets.aeta[2]-jets.aeta[1]))
plots['RHdphi'].fill(abs(jets.aphi[2]-jets.aphi[1]))
plots['npassed'].fill(jets.hpt[1]/jets.hpt[1])
plots['RjetCSVV2'].dfill(jets.CSVV2)
plots['RjetDeepFB'].dfill(jets.DeepFB)
############
# Plotting #
############
plt.clf()
#plots.pop('bjdR').plot(logv=True)
for p in plots:
plots[p].plot()
#%%
if returnplots==True:
return plots
else:
sys.exit()
def process(self, df):
dataset = df['dataset']
isRealData = 'genWeight' not in df
isSignal = 'htautau' in dataset
output = self.accumulator.identity()
# select at least one jet and one muon ( this is Pre-Selection! )
events = buildevents(df, fatjet='CustomAK8Puppi')
good = (
(events.muons.counts >= 1)
& (events.fatjets.counts >= 1)
)
events = events[good]
selection = processor.PackedSelection()
# trigger
trigger = np.ones(df.size, dtype='bool')
for t in self._triggers[self._year+'_'+self._trigger]:
trigger &= df[t]
selection.add('trigger', trigger[good])
# muon selection
goodmuon = (
(events.muons.p4.pt > 10)
& (np.abs(events.muons.p4.eta) < 2.4)
& (events.muons.sip3d < 4)
& (np.abs(events.muons.dz) < 0.1)
& (np.abs(events.muons.dxy) < 0.05)
& (events.muons.mvaId == 2)
)
nmuons = goodmuon.sum()
leadingmuon = events.muons[goodmuon][:, 0:1]
# fatjet closest to lepton
leadingmuon = events.muons[:, 0]
mujet_dR = leadingmuon.p4.delta_r(events.fatjets.p4)
mu_in_cone = mujet_dR.min() < 0.8 # this I am not sure we have to put as a selection...
mujet_bestidx = mujet_dR.argmin()
leadingjet_mu = events.fatjets[mujet_bestidx]
selection.add('jetkin', (
(leadingjet_mu.p4.pt > 300)
& (leadingjet_mu.p4.eta < 2.4)
& (leadingjet_mu.msoftdrop > 10.)
).any())
selection.add('jetid', (leadingjet_mu.jetId & 2).any()) # tight id
# lepton inside jet?
selection.add('muinside', mu_in_cone.astype(bool))
selection.add('LSF3muinside', (leadingjet_mu.electronIdx3SJ == 0).any())
selection.add('LSF3medium', (leadingjet_mu.lsf3>0.78).any())
# veto b-tag in opposite side
jets = events.jets[
(events.jets.p4.pt > 30.)
& (events.jets.jetId & 2) # tight id
]
ak4_ak8_pair = jets.cross(leadingjet_mu, nested=True)
dphi = ak4_ak8_pair.i0.p4.delta_phi(ak4_ak8_pair.i1.p4)
ak4_opposite = jets[(np.abs(dphi) > np.pi / 2).all()]
selection.add('antiak4btagMediumOppHem', ak4_opposite.deepcsvb.max() < self._btagWPs['med'][self._year])
# b-tag in same side
#subjets = events.subjets[:, leadingjet_mu.subJetIdx1]
# final lepton selection
nelectrons = (
(events.electrons.p4.pt > 10)
& (np.abs(events.electrons.p4.eta) < 2.5)
& (events.electrons.cutBased & (1 << 2)).astype(bool) # 2017V2 loose
).sum()
selection.add('onemuon', (nmuons == 1) & (nelectrons == 0)) # should we veto taus?
selection.add('muonkin', (
(leadingmuon.p4.pt > 27.)
& (np.abs(leadingmuon.p4.eta) < 2.4)
))
# building variables
leadingjet_mu = leadingjet_mu.flatten()
mm = (leadingjet_mu.p4 - leadingmuon.p4).mass2
jmass = (mm>0)*np.sqrt(np.maximum(0, mm)) + (mm<0)*leadingjet_mu.p4.mass # (jet - lep).M
met = events.met
joffshell = jmass < 62.5
massassumption = 80.*joffshell + (125 - 80.)*~joffshell
x = massassumption**2/(2*leadingmuon.p4.pt*met.rho) + np.cos(leadingmuon.p4.phi - met.phi)
met_eta = (
(x < 1)*np.arcsinh(x*np.sinh(leadingmuon.p4.eta))
+ (x >= 1)*(
leadingmuon.p4.eta
- np.sign(leadingmuon.p4.eta)*np.arccosh(np.maximum(1., x))
)
)
met_p4 = TLorentzVectorArray.from_ptetaphim(met.rho, met_eta, met.phi, np.zeros(met.size))
# filling missing columns
df['jet_pt'] = leadingjet_mu.p4.pt
df['jet_lsf3'] = leadingjet_mu.lsf3
df['jet_mmass'] = jmass
df['jet_hmass'] = (met_p4 + leadingjet_mu.p4).mass
df['jet_oppbtag'] = ak4_opposite.deepcsvb.max()
df['muon_pt'] = leadingmuon.p4.pt
df['muon_miso'] = leadingmuon.miniPFRelIso_all
df['met_pt'] = met.rho
df['met_eta'] = met_eta
# fill cutflow
cutflow = ['trigger', 'jetkin', 'jetid', 'antiak4btagMediumOppHem', 'onemuon', 'muonkin', 'muinside', 'LSF3muinside','LSF3muinside']
allcuts = set()
output['cutflow']['none'] += len(events)
for cut in cutflow:
allcuts.add(cut)
output['cutflow'][cut] += selection.all(*allcuts).sum()
weights = processor.Weights(len(events))
if not isRealData:
weights.add('genweight', events.genWeight)
regions = {}
regions['presel'] = {'trigger', 'jetkin', 'jetid', 'antiak4btagMediumOppHem', 'onemuon', 'muonkin'}
regions['muinjet'] = {'trigger', 'jetkin', 'jetid', 'antiak4btagMediumOppHem', 'onemuon', 'muonkin', 'muinside', 'LSF3muinside','LSF3muinside'}
for histname, h in output.items():
if not isinstance(h, hist.Hist):
continue
if not all(k in df or k == 'systematic' for k in h.fields):
print("Missing fields %r from %r" % (set(h.fields) - set(df.keys()), h))
continue
fields = {k: df[k] for k in h.fields if k in df}
region = [r for r in regions.keys() if r in histname.split('_')]
if len(region) == 1:
region = region[0]
cut = selection.all(*regions[region])
h.fill(**fields, weight=cut)
elif len(region) > 1:
raise ValueError("Histogram '%s' has a name matching multiple region definitions: %r" % (histname, region))
else:
raise ValueError("Histogram '%s' does not fall into any region definitions." % (histname, ))
return output
def to_define(df):
df['ip3d_sig'] = df['ip3d'] / df['ip3d_e']
df['decay_time_ps'] = df['decay_time'] * 1e12
df['abs_mu1_dxy'] = abs(df['mu1_dxy'])
df['abs_mu2_dxy'] = abs(df['mu2_dxy'])
df['mu1_dxy_sig'] = abs(df['mu1_dxy'] // df['mu1_dxyErr'])
df['mu2_dxy_sig'] = abs(df['mu2_dxy'] / df['mu2_dxyErr'])
df['abs_k_dxy'] = abs(df['k_dxy'])
df['k_dxy_sig'] = abs(df['k_dxy'] / df['k_dxyErr'])
df['abs_mu1_dz'] = abs(df['mu1_dz'])
df['abs_mu2_dz'] = abs(df['mu2_dz'])
df['mu1_dz_sig'] = abs(df['mu1_dz'] / df['mu1_dzErr'])
df['mu2_dz_sig'] = abs(df['mu2_dz'] / df['mu2_dzErr'])
df['abs_k_dz'] = abs(df['k_dz'])
df['k_dz_sig'] = abs(df['k_dz'] / df['k_dzErr'])
df['abs_mu1mu2_dz'] = abs(df['mu1_dz'] - df['mu2_dz'])
df['abs_mu1k_dz'] = abs(df['mu1_dz'] - df['k_dz'])
df['abs_mu2k_dz'] = abs(df['mu2_dz'] - df['k_dz'])
#df['bvtx_log10_svprob' ] = TMath::Log10(df['bvtx_svprob'])
#df['jpsivtx_log10_svprob' ] = TMath::Log10(df['jpsivtx_svprob'])
'''df['bvtx_log10_lxy' ] = TMath::Log10(df['bvtx_lxy)'
df['jpsivtx_log10_lxy' ] = TMath::Log10(df['jpsivtx_lxy)'
df['bvtx_log10_lxy_sig' ] = TMath::Log10(df['bvtx_lxy_sig)'
df['jpsivtx_log10_lxy_sig' ] = TMath::Log10(df['jpsivtx_lxy_sig)'
'''
df['b_iso03_rel'] = df['b_iso03'] / df['Bpt']
df['b_iso04_rel'] = df['b_iso04'] / df['Bpt']
df['k_iso03_rel'] = df['k_iso03'] / df['kpt']
df['k_iso04_rel'] = df['k_iso04'] / df['kpt']
df['mu1_iso03_rel'] = df['mu1_iso03'] / df['mu1pt']
df['mu1_iso04_rel'] = df['mu1_iso04'] / df['mu2pt']
df['mu2_iso03_rel'] = df['mu2_iso03'] / df['mu2pt']
df['mu2_iso04_rel'] = df['mu2_iso04'] / df['mu2pt']
mu1_p4 = TLorentzVectorArray.from_ptetaphim(df['mu1pt'], df['mu1eta'],
df['mu1phi'], df['mu1mass'])
mu2_p4 = TLorentzVectorArray.from_ptetaphim(df['mu2pt'], df['mu2eta'],
df['mu2phi'], df['mu2mass'])
mu3_p4 = TLorentzVectorArray.from_ptetaphim(df['kpt'], df['keta'],
df['kphi'], df['kmass'])
k_p4 = TLorentzVectorArray.from_ptetaphim(df['kpt'], df['keta'],
df['kphi'], 0.493677)
mmm_p4 = mu1_p4 + mu2_p4 + mu3_p4
df['m12'] = (mu1_p4 + mu2_p4).mass
df['m23'] = (mu2_p4 + mu3_p4).mass
df['m13'] = (mu1_p4 + mu3_p4).mass
#df['bct' ] = df['bvtx_lxy*6.275/Bpt_reco'
#df['jpsiK_p4' ] = df['mu1_p4+mu2_p4+kaon_p4'
#df['jpsiK_mass' ] = df['jpsiK_p4.mass()'
#df['jpsiK_pt' ] = df['jpsiK_p4.pt()'
#df['jpsiK_eta' ] = df['jpsiK_p4.eta()'
#df['jpsiK_phi' ] = df['jpsiK_p4.phi()'
#df['pion_p4' ] = df['ROOT::Math::PtEtaPhiMVector(kpt, keta, kphi, 0.13957018)' # this is at the correct pion mass
#df['jpsipi_p4' ] = df['mu1_p4+mu2_p4+pion_p4'
#df['jpsipi_mass' ] = df['jpsipi_p4.mass()'
#df['jpsipi_pt' ] = df['jpsipi_p4.pt()'
#df['jpsipi_eta' ] = df['jpsipi_p4.eta()'
#df['jpsipi_phi' ] = df['jpsipi_p4.phi()'
jpsi_p4 = mu1_p4 + mu2_p4
df['jpsi_pt'] = jpsi_p4.pt
df['jpsi_eta'] = jpsi_p4.eta
df['jpsi_phi'] = jpsi_p4.phi
df['jpsi_mass'] = jpsi_p4.mass
#df['q2jpsik' ] = df['(((jpsiK_p4 * 6.275/jpsiK_p4.mass()) - jpsi_p4).M2())'
#df['q2jpsipi' ] = df['(((jpsipi_p4 * 6.275/jpsipi_p4.mass()) - jpsi_p4).M2())'
#df['m2missjpsik' ] = df['(((jpsiK_p4 * 6.275/jpsiK_p4.mass()) - jpsi_p4 - kaon_p4).M2())'
#df['m2missjpsipi' ] = df['(((jpsipi_p4 * 6.275/jpsipi_p4.mass()) - jpsi_p4 - pion_p4).M2())'
#df['dr12' ] = df['ROOT::Math::VectorUtil::DeltaR(mu1_p4.Vect( mu2_p4.Vect())'
#df['dr13' ] = df['ROOT::Math::VectorUtil::DeltaR(mu1_p4.Vect( mu3_p4.Vect())'
#df['dr23' ] = df['ROOT::Math::VectorUtil::DeltaR(mu2_p4.Vect( mu3_p4.Vect())'
#df['dr_jpsi_mu' ] = df['ROOT::Math::VectorUtil::DeltaR(jpsi_p4.Vect( mu3_p4.Vect())'
#df['norm' ] = df['0.5'
# is there a better way?
#df['maxdr' ] = df['dr12*(dr12>dr13 & dr12>dr23) + dr13*(dr13>dr12 & dr13>dr23) + dr23*(dr23>dr12 & dr23>dr13)'
#df['mindr' ] = df['dr12*(dr12<dr13 & dr12<dr23) + dr13*(dr13<dr12 & dr13<dr23) + dr23*(dr23<dr12 & dr23<dr13)'
#df['pv_to_sv' ] = df['ROOT::Math::XYZVector((bvtx_vtx_x - pv_x (bvtx_vtx_y - pv_y (bvtx_vtx_z - pv_z))'
#df['Bdirection' ] = df['pv_to_sv/sqrt(pv_to_sv.Mag2())'
#df['Bdir_eta' ] = df['Bdirection.eta()'
#df['Bdir_phi' ] = df['Bdirection.phi()'
#df['mmm_p4_par' ] = df['mmm_p4.Vect().Dot(Bdirection)'
#df['mmm_p4_perp' ] = df['sqrt(mmm_p4.Vect().Mag2()-mmm_p4_par*mmm_p4_par)'
#df['mcorr' ] = df['sqrt(mmm_p4.mass()*mmm_p4.mass() + mmm_p4_perp*mmm_p4_perp) + mmm_p4_perp' # Eq. 3 https://cds.cern.ch/record/2697350/files/1910.13404.pdf
return df
def bregcorr(jets):
return TLorentzVectorArray.from_ptetaphim(jets.pt * jets.bRegCorr,
jets.eta, jets.phi,
jets.mass * jets.bRegCorr)
def process(self, events):
dataset = events.metadata['dataset']
print('process dataset', dataset)
isRealData = 'genWeight' not in events.columns
selection = processor.PackedSelection()
weights = processor.Weights(len(events))
output = self.accumulator.identity()
if (len(events) == 0): return output
if not isRealData:
output['sumw'][dataset] += events.genWeight.sum()
# trigger paths
if isRealData:
trigger_fatjet = np.zeros(events.size, dtype='bool')
for t in self._triggers[self._year]:
try:
trigger_fatjet = trigger_fatjet | events.HLT[t]
except:
print('trigger %s not available' % t)
continue
trigger_muon = np.zeros(events.size, dtype='bool')
for t in self._muontriggers[self._year]:
trigger_muon = trigger_muon | events.HLT[t]
else:
trigger_fatjet = np.ones(events.size, dtype='bool')
trigger_muon = np.ones(events.size, dtype='bool')
selection.add('fatjet_trigger', trigger_fatjet)
selection.add('muon_trigger', trigger_muon)
#jet corrected kinematics
gru = events.GRU
IN = events.IN
fatjets = events.FatJet
fatjets['msdcorr'] = corrected_msoftdrop(fatjets)
fatjets['rhocorr'] = 2 * np.log(fatjets.msdcorr / fatjets.pt)
fatjets['gruddt'] = gru.v25 - shift(
fatjets, algo='gruddt', year=self._year)
fatjets['gru'] = gru.v25
fatjets['in_v3'] = IN.v3
fatjets['in_v3_ddt'] = IN.v3 - shift(
fatjets, algo='inddt', year=self._year)
fatjets['in_v3_ddt_90pctl'] = IN.v3 - shift(
fatjets, algo='inddt90pctl', year=self._year)
fatjets['n2ddt'] = fatjets.n2b1 - n2ddt_shift(fatjets, year=self._year)
fatjets["genMatchFull"] = genmatch(events, dataset)
#else: fatjets["genMatchFull"] = fatjets.pt.zeros_like() #np.zeros(events.size, dtype='bool')
candidatejet = fatjets[:, :1]
candidatemuon = events.Muon[:, :5]
# run model on PFCands associated to FatJet (FatJetPFCands)
#events.FatJet.array.content["PFCands"] = type(events.FatJetPFCands.array).fromcounts(events.FatJet.nPFConstituents.flatten(), events.FatJetPFCands.flatten())
#events.FatJet.array.content["twoProngGru"] = run_model(events.FatJet.flatten())
selection.add('pt', (candidatejet.pt > 525).any())
selection.add('msdcorr', (candidatejet.msdcorr > 40).any())
# basic jet selection
goodjet_sel = ((candidatejet.pt > 525)
& (abs(candidatejet.eta) < 2.5)
& (candidatejet.msoftdrop > 40.)
& (candidatejet.rhocorr > -5.5)
& (candidatejet.rhocorr < -2)
&
(candidatejet.genMatchFull if
('WJetsToQQ' in dataset or 'ZJetsToQQ' in dataset) else
(1 == 1))).any()
vselection_goodjet_sel = ((candidatejet.pt > 200)
& (abs(candidatejet.eta) < 2.5)
& (candidatejet.msoftdrop > 40.)).any()
#& (candidatejet.genMatchFull if ('TTTo' in dataset) else (1==1))).any()
#& (candidatejet.rhocorr > -5.5)
#& (candidatejet.rhocorr < -2)).any()
selection.add('vselection_jetkin', vselection_goodjet_sel)
#goodmuon sel for muon CR (lep vetos below)
goodmuon_sel = ((candidatemuon.pt > 55)
& (abs(candidatemuon.eta) < 2.1)
& (candidatemuon.looseId).astype(bool)
& (candidatemuon.pfRelIso04_all < 0.15)).any()
vselection_goodmuon_sel = ((candidatemuon.pt > 53)
& (abs(candidatemuon.eta) < 2.1)
& (candidatemuon.tightId).astype(bool))
#& (candidatemuon.pfRelIso04_all < 0.15))
vselection_goodmuon_sel_loose = ((candidatemuon.pt > 20)
& (candidatemuon.looseId).astype(bool)
& (abs(candidatemuon.eta) < 2.4))
selection.add('vselection_muonkin', vselection_goodmuon_sel.any())
selection.add('vselection_onetightmuon',
vselection_goodmuon_sel.sum() == 1)
selection.add('vselection_oneloosemuon',
vselection_goodmuon_sel_loose.sum() == 1)
candidatemuon = candidatemuon[:, 0:1]
selection.add('muonkin', goodmuon_sel)
selection.add('jetkin', goodjet_sel)
selection.add('n2ddt', (candidatejet.n2ddt < 0.).any())
selection.add('jetid', candidatejet.isTight.any())
selection.add('met', events.MET.pt > 40.)
muon_ak8_pair = candidatemuon.cross(candidatejet, nested=True)
selection.add('muonDphiAK8',
(abs(muon_ak8_pair.i0.delta_phi(muon_ak8_pair.i1)) >
2 * np.pi / 3).all().all())
selection.add('vselection_muonDphiAK8', (abs(
muon_ak8_pair.i0.delta_phi(muon_ak8_pair.i1)) > 1).all().all())
#ak4 puppi jet for CR
jets = events.Jet[((events.Jet.pt > 50.)
& (abs(events.Jet.eta) < 2.5))][:, :10]
# only consider first 4 jets to be consistent with old framework
ak4_ak8_pair = jets.cross(candidatejet, nested=True)
dr = abs(ak4_ak8_pair.i0.delta_r(ak4_ak8_pair.i1))
dphi = abs(ak4_ak8_pair.i0.delta_phi(ak4_ak8_pair.i1))
ak4_away = jets[(dr > 0.8).all()]
selection.add('ak4btagMedium08', ak4_away.btagCSVV2.max() > 0.8838)
ak4_opposite = jets[(dphi > np.pi / 2).all()]
selection.add('antiak4btagMediumOppHem',
ak4_opposite.btagCSVV2.max() < 0.8838)
mu_p4 = TLorentzVectorArray.from_ptetaphim(
candidatemuon.pt.fillna(0), candidatemuon.eta.fillna(0),
candidatemuon.phi.fillna(0), candidatemuon.mass.fillna(0))
met_p4 = TLorentzVectorArray.from_ptetaphim(
awkward.JaggedArray.fromiter([[v] for v in events.MET.pt]),
awkward.JaggedArray.fromiter([[v] for v in np.zeros(events.size)]),
awkward.JaggedArray.fromiter([[v] for v in events.MET.phi]),
awkward.JaggedArray.fromiter([[v] for v in np.zeros(events.size)]))
met_candidatemuon_pair = met_p4.cross(mu_p4)
Wleptoniccandidate = met_candidatemuon_pair.i0 + met_candidatemuon_pair.i1
selection.add('Wleptonic_candidate',
(Wleptoniccandidate.pt > 200).any())
vselection_jets = events.Jet[((events.Jet.pt > 30.)
& (abs(events.Jet.eta) < 2.4))]
vselection_ak4_ak8_pair = vselection_jets.cross(candidatejet,
nested=True)
muon_ak4_pair = vselection_jets.cross(candidatemuon, nested=True)
dr_ak8 = abs(
vselection_ak4_ak8_pair.i0.delta_r(vselection_ak4_ak8_pair.i1))
dr_muon = abs(muon_ak4_pair.i0.delta_r(muon_ak4_pair.i1))
ak4_away = vselection_jets[(dr_ak8 > 0.8).all()]
selection.add('vselection_ak4btagMedium08',
ak4_away.btagCSVV2.max() > 0.8838)
ak4_away = vselection_jets[(dr_muon > 0.3).all()]
selection.add('vselection_muonDphiAK4',
ak4_away.btagCSVV2.max() > 0.8838)
nelectrons = ((
(events.Electron.pt > 10.)
& (abs(events.Electron.eta) < 2.5)
#& (events.Electron.cutBased >= events.Electron.LOOSE))
#& (events.Electron.cutBased_Fall17_V1 >= 1))
& (events.Electron.cutBased >= 2))).sum()
nmuons = (((events.Muon.pt > 10)
& (abs(events.Muon.eta) < 2.1)
#& (events.Muon.pfRelIso04_all < 0.4)
& (events.Muon.looseId).astype(bool))).sum()
ntaus = (((events.Tau.pt > 20.)
#& (events.Tau.idMVAnewDM2017v2 >=4))
& (events.Tau.idDecayMode).astype(bool)
& (events.Tau.rawIso < 5)
& (abs(events.Tau.eta) < 2.3))).sum()
selection.add('noleptons',
(nmuons == 0) & (nelectrons == 0) & (ntaus == 0))
selection.add('noelectron_notau', (nelectrons == 0) & (ntaus == 0))
#weights.add('metfilter', events.Flag.METFilters)
if isRealData:
genflavor = candidatejet.pt.zeros_like().pad(
1, clip=True).fillna(-1).flatten()
if not isRealData:
weights.add('genweight', events.genWeight)
add_pileup_weight(weights, events.Pileup.nPU, self._year, dataset)
#add_jetTriggerWeight(weights, candidatejet.msdcorr, candidatejet.pt, self._year) #signal region only
#add_singleMuTriggerWeight(weights, abs(candidatemuon.eta), candidatemuon.pt, self._year)
bosons = getBosons(events)
genBosonPt = bosons.pt.pad(1, clip=True).fillna(0)
add_VJets_NLOkFactor(weights, genBosonPt, self._year, dataset)
genflavor = matchedBosonFlavor(candidatejet, bosons).pad(
1, clip=True).fillna(-1).flatten()
#b-tag weights
regions = {
'signal': [
'fatjet_trigger',
'jetkin',
'noleptons',
'jetid',
'antiak4btagMediumOppHem',
],
'ttbar_muoncontrol': [
'muon_trigger',
'pt',
'msdcorr',
'jetid',
'jetkin',
'muonkin',
'muonDphiAK8',
'ak4btagMedium08',
'noelectron_notau',
],
'vselection': [
'muon_trigger', 'vselection_jetkin', 'vselection_muonkin',
'vselection_onetightmuon', 'vselection_oneloosemuon',
'vselection_muonDphiAK8', 'vselection_ak4btagMedium08',
'vselection_muonDphiAK4', 'Wleptonic_candidate', 'met'
],
'noselection':
[], #'vselection_muoncontrol' : ['muon_trigger', 'v_selection_jetkin', 'genmatch', 'jetid', 'ak4btagMedium08', 'muonkin','met'],
}
allcuts_signal = set()
output['cutflow_signal'][dataset]['none'] += float(
weights.weight().sum())
allcuts_ttbar_muoncontrol = set()
output['cutflow_ttbar_muoncontrol'][dataset]['none'] += float(
weights.weight().sum())
allcuts_vselection = set()
output['cutflow_vselection'][dataset]['none'] += float(
weights.weight().sum())
for cut in regions['signal']:
allcuts_signal.add(cut)
output['cutflow_signal'][dataset][cut] += float(
weights.weight()[selection.all(*allcuts_signal)].sum())
for cut in regions['ttbar_muoncontrol']:
allcuts_ttbar_muoncontrol.add(cut)
output['cutflow_ttbar_muoncontrol'][dataset][cut] += float(
weights.weight()[selection.all(
*allcuts_ttbar_muoncontrol)].sum())
for cut in regions['vselection']:
allcuts_vselection.add(cut)
output['cutflow_vselection'][dataset][cut] += float(
weights.weight()[selection.all(*allcuts_vselection)].sum())
def normalize(val, cut):
return val[cut].pad(1, clip=True).fillna(0).flatten()
def fill(region, systematic=None, wmod=None):
print('filling %s' % region)
selections = regions[region]
cut = selection.all(*selections)
weight = weights.weight()[cut]
output['templates'].fill(
dataset=dataset,
region=region,
pt=normalize(candidatejet.pt, cut),
msd=normalize(candidatejet.msdcorr, cut),
n2ddt=normalize(candidatejet.n2ddt, cut),
gruddt=normalize(candidatejet.gruddt, cut),
in_v3_ddt=normalize(candidatejet.in_v3_ddt_90pctl, cut),
weight=weight,
),
output['event'].fill(
dataset=dataset,
region=region,
MET=events.MET.pt[cut],
nJet=fatjets.counts[cut],
nPFConstituents=normalize(candidatejet.nPFConstituents, cut),
weight=weight,
),
output['muon'].fill(
dataset=dataset,
region=region,
mu_pt=normalize(candidatemuon.pt, cut),
mu_pfRelIso04_all=normalize(candidatemuon.pfRelIso04_all, cut),
weight=weight,
),
output['deepAK8'].fill(
dataset=dataset,
region=region,
deepTagMDWqq=normalize(candidatejet.deepTagMDWqq, cut),
deepTagMDZqq=normalize(candidatejet.deepTagMDZqq, cut),
msd=normalize(candidatejet.msdcorr, cut),
genflavor=genflavor[cut],
weight=weight,
),
output['in_v3'].fill(
dataset=dataset,
region=region,
genflavor=genflavor[cut],
in_v3=normalize(candidatejet.in_v3, cut),
n2=normalize(candidatejet.n2b1, cut),
gru=normalize(candidatejet.gru, cut),
weight=weight,
)
for region in regions:
fill(region)
return output
def process(self, df):
dataset = df.metadata['dataset']
isRealData = 'genWeight' not in df.columns
output = self.accumulator.identity()
selection = processor.PackedSelection()
output = self.accumulator.identity()
good = False
goodMuon = ((df.Muon.pt > 27.) & (np.abs(df.Muon.eta) < 2.4))
nmuons = goodMuon.sum()
goodElectron = ((df.Electron.pt > 30.)
& (np.abs(df.Electron.eta) < 2.5))
nelectrons = goodElectron.sum()
df.FatJet['msdcorr'] = corrected_msoftdrop(df.FatJet)
goodFatJet = ((df.FatJet.pt > 300.)
& (np.abs(df.FatJet.eta) < 2.4)
& (df.FatJet.msdcorr > 10.)
& (df.FatJet.isTight))
nfatjets = goodFatJet.sum()
if self._channel == 'muon':
good = ((nmuons >= 1) & (nfatjets >= 1))
else:
good = ((nelectrons >= 1) & (nfatjets >= 1))
events = df[good]
if not isRealData:
output['sumw'][dataset] += events.genWeight.sum()
# trigger
trigger = np.zeros(df.size, dtype='bool')
for t in self._triggers[self._year + '_' + self._trigger]:
try:
trigger = trigger | df.HLT[t]
except:
warnings.warn("Missing trigger %s" % t, RuntimeWarning)
selection.add('trigger', trigger[good])
# Muons
candidatemuon = events.Muon[:, 0:1]
nmuons = events.Muon.counts
# Electrons
candidateelectron = events.Electron[:, 0:1]
nelectrons = events.Electron.counts
if self._channel == 'muon':
candidatelep = candidatemuon
selection.add('nootherlepton', (nelectrons == 0))
else:
candidatelep = candidateelectron
selection.add('nootherlepton', (nmuons == 0))
selection.add('iplepton', ((np.abs(candidatelep.dz) < 0.1)
& (np.abs(candidatelep.dxy) < 0.05)).any())
# FatJets
ak8_lep_pair = candidatelep.cross(events.FatJet)
ak8_lep_dR = ak8_lep_pair.i0.delta_r(ak8_lep_pair.i1)
candidatejet = events.FatJet[ak8_lep_dR.argmin()]
leadingjet = events.FatJet[:, 0:1]
ak8_lep_dR_closest = candidatelep.delta_r(candidatejet)
selection.add('jetkin', (candidatejet.pt > self._fjetptMIN).any())
selection.add('jetmsd', (candidatejet.msdcorr > 20).any())
selection.add('LSF3medium', (candidatejet.lsf3 > 0.7).any())
selection.add('LSF3tight', (candidatejet.lsf3 > 0.78).any())
selection.add('lepnearjet', (ak8_lep_dR.min() < 1.5))
selection.add('lepinjet', (ak8_lep_dR.min() < 0.8))
# FatJet substracted Lepton
# sj1_sj2_btagDeepB_pair = candidatejet.LSsubJet1btagDeepB.cross(candidatejet.LSsubJet2btagDeepB)
# fls_btagDeepB_max = max(sj1_sj2_btagDeepB_pair.i0,sj1_sj2_btagDeepB_pair.i1)
# Jets
jets = events.Jet[(events.Jet.pt > 30.)
& (abs(events.Jet.eta) < 2.5)
& (events.Jet.isTight)]
ak4_ak8_pair = jets.cross(candidatejet, nested=True)
ak4_ak8_dphi = abs(ak4_ak8_pair.i0.delta_phi(ak4_ak8_pair.i1))
ak4_opposite = jets[(ak4_ak8_dphi > np.pi / 2).all()]
ak4_away = jets[(ak4_ak8_dphi > 0.8).all()]
selection.add(
'antiak4btagMediumOppHem',
ak4_opposite.btagDeepB.max() < self._btagWPs['med'][self._year])
selection.add(
'ak4btagMedium08',
ak4_away.btagDeepB.max() < self._btagWPs['med'][self._year])
# MET
met = events.MET
# MET eta with mass assumption
mm = (candidatejet - candidatelep).mass2
jmass = (mm > 0) * np.sqrt(np.maximum(
0, mm)) + (mm < 0) * candidatejet.mass
joffshell = jmass < 62.5
massassumption = 80. * joffshell + (125 - 80.) * ~joffshell
x = massassumption**2 / (2 * candidatelep.pt *
met.pt) + np.cos(candidatelep.phi - met.phi)
met_eta = ((x < 1) * np.arcsinh(x * np.sinh(candidatelep.eta)) +
(x > 1) *
(candidatelep.eta -
np.sign(candidatelep.eta) * np.arccosh(candidatelep.eta)))
met_p4 = TLorentzVectorArray.from_ptetaphim(np.array([0.]),
np.array([0.]),
np.array([0.]),
np.array([0.]))
if met.size > 0:
met_p4 = TLorentzVectorArray.from_ptetaphim(
met.pt, met_eta.fillna(0.), met.phi, np.zeros(met.size))
hmass = (candidatejet + met_p4).mass
else:
hmass = candidatejet.pt.zeros_like()
# weights
weights = processor.Weights(len(events), storeIndividual=True)
if isRealData:
genflavor = candidatejet.pt.zeros_like()
else:
try:
weights.add('genweight', events.genWeight)
add_pileup_weight(weights, events.Pileup.nPU, self._year)
#print("Weight statistics: %r" % weights._weightStats)
except:
print('no gen weight')
if 'TTTo' in dataset:
genW, genW_idx = getParticles(
events, 24, ['fromHardProcess', 'isLastCopy'])
genb, genb_idx = getParticles(
events, 5, ['fromHardProcess', 'isLastCopy'])
genflavorW = matchedParticleFlavor(candidatelep, genW, 'child',
0.4)
genflavorb = matchedParticleFlavor(candidatelep, genb, 'mom',
0.4)
genflavor = getFlavor(genflavorW, genflavorb)
elif (('hww_2017' in dataset) or ('GluGluHToWW' in dataset)):
genH, genH_idx = getParticles(
events, 25, ['fromHardProcess', 'isLastCopy'])
genW, genW_idx = getParticles(
events, 24, ['fromHardProcess', 'isLastCopy'])
genE, genE_idx = getParticles(
events, 11, ['fromHardProcess', 'isFirstCopy'], 1)
genM, genM_idx = getParticles(
events, 13, ['fromHardProcess', 'isFirstCopy'], 1)
genT, genT_idx = getParticles(
events, 15, ['fromHardProcess', 'isFirstCopy'], 1)
genQ, genQ_idx = getParticles(
events, [0, 5], ['fromHardProcess', 'isFirstCopy'])
ishWW_qqelev = (genH.counts == 1) & (genW.counts == 2) & (
genE.counts == 1) & (genM.counts == 0) & (genT.counts == 0)
ishWW_qqmuv = (genH.counts == 1) & (genW.counts == 2) & (
genM.counts == 1) & (genE.counts == 0) & (genT.counts == 0)
ishWW_qqtauv = (genH.counts == 1) & (genW.counts == 2) & (
genT.counts == 1) & (genM.counts == 0) & (genE.counts == 0)
ishWW_qqqq = (genH.counts == 1) & (genW.counts == 2) & (
genQ.counts == 4) & (genM.counts == 0) & (genE.counts == 0)
ishWW_muvelev = (genH.counts == 1) & (genW.counts == 2) & (
genE.counts == 1) & (genM.counts == 1)
ishWW_elevelev = (genH.counts == 1) & (genW.counts == 2) & (
genE.counts == 2) & (genM.counts == 0)
ishWW_tauvtauv = (genH.counts == 1) & (genW.counts == 2) & (
genT.counts == 2) & (genM.counts == 0) & (genE.counts == 0)
ishWW_muvmuv = (genH.counts == 1) & (genW.counts == 2) & (
genE.counts == 0) & (genM.counts == 2)
genflavor = ((ishWW_qqelev) * 8 + (ishWW_qqmuv) * 9)
else:
genflavor = candidatejet.pt.zeros_like()
# fill cutflow
cutflow = [
'trigger', 'jetkin', 'jetmsd', 'lepnearjet', 'lepinjet',
'antiak4btagMediumOppHem', 'nootherlepton', 'iplepton',
'LSF3medium', 'LSF3tight'
]
allcuts = set()
output['cutflow']['none'] += len(events)
for cut in cutflow:
allcuts.add(cut)
output['cutflow'][cut] += selection.all(*allcuts).sum()
regions = {}
regions['presel'] = {'trigger', 'jetkin', 'jetmsd', 'lepinjet'}
regions['antibtag'] = {
'trigger', 'jetkin', 'jetmsd', 'antiak4btagMediumOppHem'
}
regions['noinjet'] = {
'trigger', 'jetkin', 'jetmsd', 'lepnearjet',
'antiak4btagMediumOppHem'
}
regions['nolsf'] = {
'trigger', 'jetkin', 'jetmsd', 'lepinjet',
'antiak4btagMediumOppHem'
} #,'nootherlepton'}
regions['lsf'] = {
'trigger', 'jetkin', 'jetmsd', 'lepinjet', 'LSF3tight'
}
regions['bopp'] = {
'trigger', 'jetkin', 'jetmsd', 'lepinjet', 'LSF3tight',
'antiak4btagMediumOppHem'
}
regions['lep'] = {
'trigger', 'jetkin', 'jetmsd', 'lepinjet', 'LSF3tight',
'antiak4btagMediumOppHem', 'nootherlepton', 'iplepton'
}
for region in self._regions:
selections = regions[region]
cut = selection.all(*selections)
weight = weights.weight()[cut]
def normalize(val):
try:
return val[cut].pad(1, clip=True).fillna(0).flatten()
except:
try:
return val[cut].flatten()
except:
return val[cut]
# output['%s_fjetprop'%region].fill(#fjet_pt = normalize(candidatejet.pt),
# fjet_msd = normalize(candidatejet.msdcorr),
# fjet_lsf3 = normalize(candidatejet.lsf3),
# #jet_oppbtag = normalize(ak4_opposite.btagDeepB.max()),
# genflavor = normalize(genflavor),
# dataset=dataset,
# weight=weight
# )
# output['%s_fjetextraprop'%region].fill(fjet_t41 = normalize(candidatejet.tau4/candidatejet.tau1),
# fjet_t42 = normalize(candidatejet.tau4/candidatejet.tau2),
# fjet_t31 = normalize(candidatejet.tau3/candidatejet.tau1),
# dataset=dataset,
# weight=weight
# )
# output['%s_jetprop'%region].fill(jet_oppbtag = normalize(ak4_opposite.btagDeepB.max()),
# genflavor = normalize(genflavor),
# dataset=dataset,
# weight=weight
# )
output['%s_fmmjetprop' % region].fill(
fjet_pt=normalize(candidatejet.pt),
#fjet_mmass = normalize(jmass),
#fjet_hmass = normalize(hmass),
lep_pt=normalize(candidatelep.pt),
fjet_lsf3=normalize(candidatejet.lsf3),
genflavor=normalize(genflavor),
dataset=dataset,
weight=weight)
output['%s_fmmjetprop2' % region].fill(
fjet_mmass=normalize(jmass),
fjet_lsf3=normalize(candidatejet.lsf3),
genflavor=normalize(genflavor),
dataset=dataset,
weight=weight)
# output['%s_flsjetprop'%region].fill(#flsjet_pt = normalize(candidatejet.LSpt),
# flsjet_msd = normalize(candidatejet.LSmsoftdrop),
# #flsjet_n2b1 = normalize(candidatejet.LSn2b1),
# #flsjet_n3b1 = normalize(candidatejet.LSn3b1),
# #flsjet_t21 = normalize(candidatejet.LStau2/candidatejet.LStau1),
# #flsjet_t32 = normalize(candidatejet.LStau3/candidatejet.LStau2),
# genflavor = normalize(genflavor),
# dataset=dataset,
# weight=weight)
#output['%s_metprop'%region].fill(met_pt = normalize(met.pt),
# met_phi = normalize(met.phi),
# dataset=dataset,
# weight=weight)
# output['%s_weight'%region].fill(puweight=weights.partial_weight(include=["pileup_weight"])[cut],
# genweight=weights.partial_weight(include=["genweight"])[cut],
# dataset=dataset,
# )
# if self._channel=='muon':
# output['%s_muonprop'%region].fill(muon_pt = normalize(candidatemuon.pt),
# muon_miso = normalize(candidatemuon.miniPFRelIso_all),
# muon_sip = normalize(candidatemuon.sip3d),
# dataset=dataset,
# weight=weight)
# output['%s_muonextraprop'%region].fill(nmuons = normalize(nmuons),
# nelectrons = normalize(nelectrons),
# muon_dz = normalize(candidatemuon.dz),
# muon_dxy = normalize(candidatemuon.dxy),
# dataset=dataset,
# weight=weight)
# else:
# output['%s_electronprop'%region].fill(electron_pt = normalize(candidateelectron.pt),
# electron_miso = normalize(candidateelectron.miniPFRelIso_all),
# electron_sip = normalize(candidateelectron.sip3d),
# dataset=dataset,
# weight=weight)
# output['%s_electronextraprop'%region].fill(nmuons = normalize(nmuons),
# nelectrons = normalize(nelectrons),
# electron_dz = normalize(candidateelectron.dz),
# electron_dxy = normalize(candidateelectron.dxy),
# dataset=dataset,
# weight=weight)
return output
def process(self, events):
dataset = events.metadata['dataset']
selected_regions = []
for region, samples in self._samples.items():
for sample in samples:
if sample not in dataset: continue
selected_regions.append(region)
isData = 'genWeight' not in events.columns
selection = processor.PackedSelection()
weights = {}
hout = self.accumulator.identity()
###
#Getting corrections, ids from .coffea files
### Sunil Need to check why we need corrections
#get_msd_weight = self._corrections['get_msd_weight']
get_ttbar_weight = self._corrections['get_ttbar_weight']
get_nlo_weight = self._corrections['get_nlo_weight'][self._year]
get_nnlo_weight = self._corrections['get_nnlo_weight']
get_nnlo_nlo_weight = self._corrections['get_nnlo_nlo_weight']
get_adhoc_weight = self._corrections['get_adhoc_weight']
get_pu_weight = self._corrections['get_pu_weight'][self._year]
get_met_trig_weight = self._corrections['get_met_trig_weight'][self._year]
get_met_zmm_trig_weight = self._corrections['get_met_zmm_trig_weight'][self._year]
get_ele_trig_weight = self._corrections['get_ele_trig_weight'][self._year]
get_pho_trig_weight = self._corrections['get_pho_trig_weight'][self._year]
get_ele_loose_id_sf = self._corrections['get_ele_loose_id_sf'][self._year]
get_ele_tight_id_sf = self._corrections['get_ele_tight_id_sf'][self._year]
get_ele_loose_id_eff = self._corrections['get_ele_loose_id_eff'][self._year]
get_ele_tight_id_eff = self._corrections['get_ele_tight_id_eff'][self._year]
get_pho_tight_id_sf = self._corrections['get_pho_tight_id_sf'][self._year]
get_mu_tight_id_sf = self._corrections['get_mu_tight_id_sf'][self._year]
get_mu_loose_id_sf = self._corrections['get_mu_loose_id_sf'][self._year]
get_ele_reco_sf = self._corrections['get_ele_reco_sf'][self._year]
get_mu_tight_iso_sf = self._corrections['get_mu_tight_iso_sf'][self._year]
get_mu_loose_iso_sf = self._corrections['get_mu_loose_iso_sf'][self._year]
get_ecal_bad_calib = self._corrections['get_ecal_bad_calib']
get_deepflav_weight = self._corrections['get_btag_weight']['deepflav'][self._year]
Jetevaluator = self._corrections['Jetevaluator']
isLooseElectron = self._ids['isLooseElectron']
isTightElectron = self._ids['isTightElectron']
isLooseMuon = self._ids['isLooseMuon']
isTightMuon = self._ids['isTightMuon']
isLooseTau = self._ids['isLooseTau']
isLoosePhoton = self._ids['isLoosePhoton']
isTightPhoton = self._ids['isTightPhoton']
isGoodJet = self._ids['isGoodJet']
#isGoodFatJet = self._ids['isGoodFatJet']
isHEMJet = self._ids['isHEMJet']
match = self._common['match']
deepflavWPs = self._common['btagWPs']['deepflav'][self._year]
deepcsvWPs = self._common['btagWPs']['deepcsv'][self._year]
###
# Derive jet corrector for JEC/JER
###
JECcorrector = FactorizedJetCorrector(**{name: Jetevaluator[name] for name in self._jec[self._year]})
JECuncertainties = JetCorrectionUncertainty(**{name:Jetevaluator[name] for name in self._junc[self._year]})
JER = JetResolution(**{name:Jetevaluator[name] for name in self._jr[self._year]})
JERsf = JetResolutionScaleFactor(**{name:Jetevaluator[name] for name in self._jersf[self._year]})
Jet_transformer = JetTransformer(jec=JECcorrector,junc=JECuncertainties, jer = JER, jersf = JERsf)
###
#Initialize global quantities (MET ecc.)
###
met = events.MET
met['T'] = TVector2Array.from_polar(met.pt, met.phi)
met['p4'] = TLorentzVectorArray.from_ptetaphim(met.pt, 0., met.phi, 0.)
calomet = events.CaloMET
###
#Initialize physics objects
###
e = events.Electron
e['isloose'] = isLooseElectron(e.pt,e.eta,e.dxy,e.dz,e.cutBased,self._year)
e['istight'] = isTightElectron(e.pt,e.eta,e.dxy,e.dz,e.cutBased,self._year)
e['T'] = TVector2Array.from_polar(e.pt, e.phi)
#e['p4'] = TLorentzVectorArray.from_ptetaphim(e.pt, e.eta, e.phi, e.mass)
e_loose = e[e.isloose.astype(np.bool)]
e_tight = e[e.istight.astype(np.bool)]
e_ntot = e.counts
e_nloose = e_loose.counts
e_ntight = e_tight.counts
leading_e = e[e.pt.argmax()]
leading_e = leading_e[leading_e.istight.astype(np.bool)]
mu = events.Muon
mu['isloose'] = isLooseMuon(mu.pt,mu.eta,mu.pfRelIso04_all,mu.looseId,self._year)
mu['istight'] = isTightMuon(mu.pt,mu.eta,mu.pfRelIso04_all,mu.tightId,self._year)
mu['T'] = TVector2Array.from_polar(mu.pt, mu.phi)
#mu['p4'] = TLorentzVectorArray.from_ptetaphim(mu.pt, mu.eta, mu.phi, mu.mass)
mu_loose=mu[mu.isloose.astype(np.bool)]
mu_tight=mu[mu.istight.astype(np.bool)]
mu_ntot = mu.counts
mu_nloose = mu_loose.counts
mu_ntight = mu_tight.counts
leading_mu = mu[mu.pt.argmax()]
leading_mu = leading_mu[leading_mu.istight.astype(np.bool)]
tau = events.Tau
tau['isclean']=~match(tau,mu_loose,0.5)&~match(tau,e_loose,0.5)
tau['isloose']=isLooseTau(tau.pt,tau.eta,tau.idDecayMode,tau.idMVAoldDM2017v2,self._year)
tau_clean=tau[tau.isclean.astype(np.bool)]
tau_loose=tau_clean[tau_clean.isloose.astype(np.bool)]
tau_ntot=tau.counts
tau_nloose=tau_loose.counts
pho = events.Photon
pho['isclean']=~match(pho,mu_loose,0.5)&~match(pho,e_loose,0.5)
_id = 'cutBasedBitmap'
if self._year=='2016': _id = 'cutBased'
pho['isloose']=isLoosePhoton(pho.pt,pho.eta,pho[_id],self._year)
pho['istight']=isTightPhoton(pho.pt,pho.eta,pho[_id],self._year)
pho['T'] = TVector2Array.from_polar(pho.pt, pho.phi)
#pho['p4'] = TLorentzVectorArray.from_ptetaphim(pho.pt, pho.eta, pho.phi, pho.mass)
pho_clean=pho[pho.isclean.astype(np.bool)]
pho_loose=pho_clean[pho_clean.isloose.astype(np.bool)]
pho_tight=pho_clean[pho_clean.istight.astype(np.bool)]
pho_ntot=pho.counts
pho_nloose=pho_loose.counts
pho_ntight=pho_tight.counts
leading_pho = pho[pho.pt.argmax()]
leading_pho = leading_pho[leading_pho.isclean.astype(np.bool)]
leading_pho = leading_pho[leading_pho.istight.astype(np.bool)]
j = events.Jet
j['isgood'] = isGoodJet(j.pt, j.eta, j.jetId, j.neHEF, j.neEmEF, j.chHEF, j.chEmEF)
j['isHEM'] = isHEMJet(j.pt, j.eta, j.phi)
j['isclean'] = ~match(j,e_loose,0.4)&~match(j,mu_loose,0.4)&~match(j,pho_loose,0.4)
#j['isiso'] = ~match(j,fj_clean,1.5) # What is this ?????
j['isdcsvL'] = (j.btagDeepB>deepcsvWPs['loose'])
j['isdflvL'] = (j.btagDeepFlavB>deepflavWPs['loose'])
j['T'] = TVector2Array.from_polar(j.pt, j.phi)
j['p4'] = TLorentzVectorArray.from_ptetaphim(j.pt, j.eta, j.phi, j.mass)
j['ptRaw'] =j.pt * (1-j.rawFactor)
j['massRaw'] = j.mass * (1-j.rawFactor)
j['rho'] = j.pt.ones_like()*events.fixedGridRhoFastjetAll.array
j_good = j[j.isgood.astype(np.bool)]
j_clean = j_good[j_good.isclean.astype(np.bool)] # USe this instead of j_iso Sunil
#j_iso = j_clean[j_clean.isiso.astype(np.bool)]
j_iso = j_clean[j_clean.astype(np.bool)] #Sunil changed
j_dcsvL = j_iso[j_iso.isdcsvL.astype(np.bool)]
j_dflvL = j_iso[j_iso.isdflvL.astype(np.bool)]
j_HEM = j[j.isHEM.astype(np.bool)]
j_ntot=j.counts
j_ngood=j_good.counts
j_nclean=j_clean.counts
j_niso=j_iso.counts
j_ndcsvL=j_dcsvL.counts
j_ndflvL=j_dflvL.counts
j_nHEM = j_HEM.counts
leading_j = j[j.pt.argmax()]
leading_j = leading_j[leading_j.isgood.astype(np.bool)]
leading_j = leading_j[leading_j.isclean.astype(np.bool)]
###
#Calculating derivatives
###
ele_pairs = e_loose.distincts()
diele = ele_pairs.i0+ele_pairs.i1
diele['T'] = TVector2Array.from_polar(diele.pt, diele.phi)
leading_ele_pair = ele_pairs[diele.pt.argmax()]
leading_diele = diele[diele.pt.argmax()]
mu_pairs = mu_loose.distincts()
dimu = mu_pairs.i0+mu_pairs.i1
dimu['T'] = TVector2Array.from_polar(dimu.pt, dimu.phi)
leading_mu_pair = mu_pairs[dimu.pt.argmax()]
leading_dimu = dimu[dimu.pt.argmax()]
###
# Calculate recoil
### HT, LT, dPhi, mT_{W}, MT_misET
um = met.T+leading_mu.T.sum()
ue = met.T+leading_e.T.sum()
umm = met.T+leading_dimu.T.sum()
uee = met.T+leading_diele.T.sum()
ua = met.T+leading_pho.T.sum()
#Need help from Matteo
u = {}
u['sr']=met.T
u['wecr']=ue
u['tecr']=ue
u['wmcr']=um
u['tmcr']=um
u['zecr']=uee
u['zmcr']=umm
u['gcr']=ua
###
#Calculating weights
###
if not isData:
###
# JEC/JER
###
#j['ptGenJet'] = j.matched_gen.pt
#Jet_transformer.transform(j)
gen = events.GenPart
#Need to understand this part Sunil
gen['isb'] = (abs(gen.pdgId)==5)&gen.hasFlags(['fromHardProcess', 'isLastCopy'])
gen['isc'] = (abs(gen.pdgId)==4)&gen.hasFlags(['fromHardProcess', 'isLastCopy'])
gen['isTop'] = (abs(gen.pdgId)==6)&gen.hasFlags(['fromHardProcess', 'isLastCopy'])
gen['isW'] = (abs(gen.pdgId)==24)&gen.hasFlags(['fromHardProcess', 'isLastCopy'])
gen['isZ'] = (abs(gen.pdgId)==23)&gen.hasFlags(['fromHardProcess', 'isLastCopy'])
gen['isA'] = (abs(gen.pdgId)==22)&gen.hasFlags(['fromHardProcess', 'isLastCopy'])
genTops = gen[gen.isTop]
genWs = gen[gen.isW]
genZs = gen[gen.isZ]
genAs = gen[gen.isA]
nlo = np.ones(events.size)
nnlo = np.ones(events.size)
nnlo_nlo = np.ones(events.size)
adhoc = np.ones(events.size)
if('TTJets' in dataset):
nlo = np.sqrt(get_ttbar_weight(genTops[:,0].pt.sum()) * get_ttbar_weight(genTops[:,1].pt.sum()))
#elif('GJets' in dataset):
# nlo = get_nlo_weight['a'](genAs.pt.max())
elif('WJets' in dataset):
#nlo = get_nlo_weight['w'](genWs.pt.max())
#if self._year != '2016': adhoc = get_adhoc_weight['w'](genWs.pt.max())
#nnlo = get_nnlo_weight['w'](genWs.pt.max())
nnlo_nlo = get_nnlo_nlo_weight['w'](genWs.pt.max())*(genWs.pt.max()>100).astype(np.int) + (genWs.pt.max()<=100).astype(np.int)
elif('DY' in dataset):
#nlo = get_nlo_weight['z'](genZs.pt.max())
#if self._year != '2016': adhoc = get_adhoc_weight['z'](genZs.pt.max())
#nnlo = get_nnlo_weight['dy'](genZs.pt.max())
nnlo_nlo = get_nnlo_nlo_weight['dy'](genZs.pt.max())*(genZs.pt.max()>100).astype(np.int) + (genZs.pt.max()<=100).astype(np.int)
elif('ZJets' in dataset):
#nlo = get_nlo_weight['z'](genZs.pt.max())
#if self._year != '2016': adhoc = get_adhoc_weight['z'](genZs.pt.max())
#nnlo = get_nnlo_weight['z'](genZs.pt.max())
nnlo_nlo = get_nnlo_nlo_weight['z'](genZs.pt.max())*(genZs.pt.max()>100).astype(np.int) + (genZs.pt.max()<=100).astype(np.int)
###
# Calculate PU weight and systematic variations
###
pu = get_pu_weight['cen'](events.PV.npvs)
#puUp = get_pu_weight['up'](events.PV.npvs)
#puDown = get_pu_weight['down'](events.PV.npvs)
###
# Trigger efficiency weight
###
ele1_trig_weight = get_ele_trig_weight(leading_ele_pair.i0.eta.sum(),leading_ele_pair.i0.pt.sum())
ele2_trig_weight = get_ele_trig_weight(leading_ele_pair.i1.eta.sum(),leading_ele_pair.i1.pt.sum())
# Need Help from Matteo
trig = {}
trig['sre'] = get_ele_trig_weight(leading_e.eta.sum(), leading_e.pt.sum())
trig['srm'] = #Need be fixed in Util first
trig['ttbare'] = get_ele_trig_weight(leading_e.eta.sum(), leading_e.pt.sum())
trig['ttbarm'] = #Need be fixed in Util first
trig['wjete'] = get_ele_trig_weight(leading_e.eta.sum(), leading_e.pt.sum())
trig['wjetm'] = #Need be fixed in Util first
trig['dilepe'] = 1 - (1-ele1_trig_weight)*(1-ele2_trig_weight)
#trig['dilepm'] = Need be fixed in Util first
# For muon ID weights, SFs are given as a function of abs(eta), but in 2016
##
mueta = abs(leading_mu.eta.sum())
mu1eta=abs(leading_mu_pair.i0.eta.sum())
mu2eta=abs(leading_mu_pair.i1.eta.sum())
if self._year=='2016':
mueta=leading_mu.eta.sum()
mu1eta=leading_mu_pair.i0.eta.sum()
mu2eta=leading_mu_pair.i1.eta.sum()
###
# Calculating electron and muon ID SF and efficiencies (when provided)
###
mu1Tsf = get_mu_tight_id_sf(mu1eta,leading_mu_pair.i0.pt.sum())
mu2Tsf = get_mu_tight_id_sf(mu2eta,leading_mu_pair.i1.pt.sum())
mu1Lsf = get_mu_loose_id_sf(mu1eta,leading_mu_pair.i0.pt.sum())
mu2Lsf = get_mu_loose_id_sf(mu2eta,leading_mu_pair.i1.pt.sum())
e1Tsf = get_ele_tight_id_sf(leading_ele_pair.i0.eta.sum(),leading_ele_pair.i0.pt.sum())
e2Tsf = get_ele_tight_id_sf(leading_ele_pair.i1.eta.sum(),leading_ele_pair.i1.pt.sum())
e1Lsf = get_ele_loose_id_sf(leading_ele_pair.i0.eta.sum(),leading_ele_pair.i0.pt.sum())
e2Lsf = get_ele_loose_id_sf(leading_ele_pair.i1.eta.sum(),leading_ele_pair.i1.pt.sum())
e1Teff= get_ele_tight_id_eff(leading_ele_pair.i0.eta.sum(),leading_ele_pair.i0.pt.sum())
e2Teff= get_ele_tight_id_eff(leading_ele_pair.i1.eta.sum(),leading_ele_pair.i1.pt.sum())
e1Leff= get_ele_loose_id_eff(leading_ele_pair.i0.eta.sum(),leading_ele_pair.i0.pt.sum())
e2Leff= get_ele_loose_id_eff(leading_ele_pair.i1.eta.sum(),leading_ele_pair.i1.pt.sum())
# Need Help from Matteo
ids={}
ids['sre'] = get_ele_tight_id_sf(leading_e.eta.sum(),leading_e.pt.sum())
ids['srm'] = get_mu_tight_id_sf(mueta,leading_mu.pt.sum())
ids['ttbare'] = get_ele_tight_id_sf(leading_e.eta.sum(),leading_e.pt.sum())
ids['ttbarm'] = get_mu_tight_id_sf(mueta,leading_mu.pt.sum())
ids['wjete'] = get_ele_tight_id_sf(leading_e.eta.sum(),leading_e.pt.sum())
ids['wjetm'] = get_mu_tight_id_sf(mueta,leading_mu.pt.sum())
ids['dilepe'] = e1Lsf*e2Lsf
ids['dilepm'] = mu1Lsf*mu2Lsf
###
# Reconstruction weights for electrons
###
e1sf_reco = get_ele_reco_sf(leading_ele_pair.i0.eta.sum(),leading_ele_pair.i0.pt.sum())
e2sf_reco = get_ele_reco_sf(leading_ele_pair.i1.eta.sum(),leading_ele_pair.i1.pt.sum())
# Need Help from Matteo
reco = {}
reco['sre'] = get_ele_reco_sf(leading_e.eta.sum(),leading_e.pt.sum())
reco['srm'] = np.ones(events.size)
reco['ttbare'] = get_ele_reco_sf(leading_e.eta.sum(),leading_e.pt.sum())
reco['ttbarm'] = np.ones(events.size)
reco['wjete'] = get_ele_reco_sf(leading_e.eta.sum(),leading_e.pt.sum())
reco['wjetm'] = np.ones(events.size)
reco['dilepe'] = e1sf_reco * e2sf_reco
reco['dilepm'] = np.ones(events.size)
###
# Isolation weights for muons
###
mu1Tsf_iso = get_mu_tight_iso_sf(mu1eta,leading_mu_pair.i0.pt.sum())
mu2Tsf_iso = get_mu_tight_iso_sf(mu2eta,leading_mu_pair.i1.pt.sum())
mu1Lsf_iso = get_mu_loose_iso_sf(mu1eta,leading_mu_pair.i0.pt.sum())
mu2Lsf_iso = get_mu_loose_iso_sf(mu2eta,leading_mu_pair.i1.pt.sum())
# Need Help from Matteo
isolation = {}
isolation['sre'] = np.ones(events.size)
isolation['srm'] = get_mu_tight_iso_sf(mueta,leading_mu.pt.sum())
isolation['ttbare'] = np.ones(events.size)
isolation['ttbarm'] = get_mu_tight_iso_sf(mueta,leading_mu.pt.sum())
isolation['wjete'] = np.ones(events.size)
isolation['wjetm'] = get_mu_tight_iso_sf(mueta,leading_mu.pt.sum())
isolation['dilepe'] = np.ones(events.size)
isolation['dilepm'] = mu1Lsf_iso*mu2Lsf_iso
###
# AK4 b-tagging weights
###
btag = {}
btagUp = {}
btagDown = {}
# Need Help from Matteo
btag['sr'], btagUp['sr'], btagDown['sr'] = get_deepflav_weight['loose'](j_iso.pt,j_iso.eta,j_iso.hadronFlavour,'0')
btag['wmcr'], btagUp['wmcr'], btagDown['wmcr'] = get_deepflav_weight['loose'](j_iso.pt,j_iso.eta,j_iso.hadronFlavour,'0')
btag['tmcr'], btagUp['tmcr'], btagDown['tmcr'] = get_deepflav_weight['loose'](j_iso.pt,j_iso.eta,j_iso.hadronFlavour,'-1')
btag['wecr'], btagUp['wecr'], btagDown['wecr'] = get_deepflav_weight['loose'](j_iso.pt,j_iso.eta,j_iso.hadronFlavour,'0')
btag['tecr'], btagUp['tecr'], btagDown['tecr'] = get_deepflav_weight['loose'](j_iso.pt,j_iso.eta,j_iso.hadronFlavour,'-1')
btag['zmcr'], btagUp['zmcr'], btagDown['zmcr'] = np.ones(events.size), np.ones(events.size), np.ones(events.size)#get_deepflav_weight['loose'](j_iso.pt,j_iso.eta,j_iso.hadronFlavour,'0')
btag['zecr'], btagUp['zecr'], btagDown['zecr'] = np.ones(events.size), np.ones(events.size), np.ones(events.size)#get_deepflav_weight['loose'](j_iso.pt,j_iso.eta,j_iso.hadronFlavour,'0')
btag['gcr'], btagUp['gcr'], btagDown['gcr'] = np.ones(events.size), np.ones(events.size), np.ones(events.size)#get_deepflav_weight['loose'](j_iso.pt,j_iso.eta,j_iso.hadronFlavour,'0')
for r in selected_regions:
weights[r] = processor.Weights(len(events))
weights[r].add('genw',events.genWeight)
weights[r].add('nlo',nlo)
#weights[r].add('adhoc',adhoc)
#weights[r].add('nnlo',nnlo)
weights[r].add('nnlo_nlo',nnlo_nlo)
weights[r].add('pileup',pu)#,puUp,puDown)
weights[r].add('trig', trig[r])
weights[r].add('ids', ids[r])
weights[r].add('reco', reco[r])
weights[r].add('isolation', isolation[r])
weights[r].add('btag',btag[r], btagUp[r], btagDown[r])
#leading_fj = fj[fj.pt.argmax()]
#leading_fj = leading_fj[leading_fj.isgood.astype(np.bool)]
#leading_fj = leading_fj[leading_fj.isclean.astype(np.bool)]
###
#Importing the MET filters per year from metfilters.py and constructing the filter boolean
###
met_filters = np.ones(events.size, dtype=np.bool)
for flag in AnalysisProcessor.met_filter_flags[self._year]:
met_filters = met_filters & events.Flag[flag]
selection.add('met_filters',met_filters)
triggers = np.zeros(events.size, dtype=np.bool)
for path in self._met_triggers[self._year]:
if path not in events.HLT.columns: continue
triggers = triggers | events.HLT[path]
selection.add('met_triggers', triggers)
triggers = np.zeros(events.size, dtype=np.bool)
for path in self._singleelectron_triggers[self._year]:
if path not in events.HLT.columns: continue
triggers = triggers | events.HLT[path]
selection.add('singleelectron_triggers', triggers)
triggers = np.zeros(events.size, dtype=np.bool)
for path in self._singlemuon_triggers[self._year]:
if path not in events.HLT.columns: continue
triggers = triggers | events.HLT[path]
selection.add('singlemuon_triggers', triggers)
triggers = np.zeros(events.size, dtype=np.bool)
for path in self._singlephoton_triggers[self._year]:
if path not in events.HLT.columns: continue
triggers = triggers | events.HLT[path]
selection.add('singlephoton_triggers', triggers)
noHEMj = np.ones(events.size, dtype=np.bool)
if self._year=='2018': noHEMj = (j_nHEM==0)
selection.add('iszeroL',
(e_nloose==0)&(mu_nloose==0)&(tau_nloose==0)&(pho_nloose==0)
def _make_pfcands(self, table):
data = {}
if self.idx_branch in table:
jet_cand_counts = table[self.fatjet_branch + '_nPFCand']
jet_cand_idxs = table[self.idx_branch]
else:
cand_parents = table[self.pfcand_branch + '_jetIdx']
c = Counter((cand_parents.offsets[:-1] + cand_parents).content)
jet_cand_counts = awkward.JaggedArray.fromcounts(
self.jetp4.counts, [c[k] for k in sorted(c.keys())])
ptcut = None
def pf(var_name):
branch_name = self.pfcand_branch + '_' + var_name
if var_name[:4] == 'btag' and self.idx_branch in table:
branch_name = branch_name + '_' + self.fatjet_branch
cand_arr = table[branch_name]
if self.idx_branch in table:
cand_arr = cand_arr[jet_cand_idxs]
out = jet_cand_counts.copy(content=awkward.JaggedArray.fromcounts(
jet_cand_counts.content, cand_arr.content))
if ptcut is None:
return out
else:
return out[ptcut]
if self.pfcand_ptcut > 0:
ptcut = pf('pt') > self.pfcand_ptcut
data['pfcand_VTX_ass'] = pf('pvAssocQuality')
data['pfcand_lostInnerHits'] = pf('lostInnerHits')
data['pfcand_quality'] = pf('trkQuality')
pdgId = pf('pdgId')
charge = pf('charge')
data['pfcand_isEl'] = np.abs(pdgId) == 11
data['pfcand_isMu'] = np.abs(pdgId) == 13
data['pfcand_isChargedHad'] = np.abs(pdgId) == 211
data['pfcand_isGamma'] = np.abs(pdgId) == 22
data['pfcand_isNeutralHad'] = np.abs(pdgId) == 130
data['pfcand_charge'] = charge
dz = pf('dz')
dxy = pf('d0')
# ### FIXME ###
# is_neutral = np.logical_or(charge.content == 0, np.logical_and(dz.content == -1.0, dxy.content == -1.0))
# dz.content[is_neutral] = 0
# dxy.content[is_neutral] = 0
# ### FIXME ###
data['pfcand_dz'] = dz
data['pfcand_dxy'] = dxy
data['pfcand_dzsig'] = dz / pf('dzErr')
data['pfcand_dxysig'] = dxy / pf('d0Err')
candp4 = TLorentzVectorArray.from_ptetaphim(pf('pt'), pf('eta'),
pf('phi'), pf('mass'))
data['pfcand_mask'] = charge.ones_like()
data['pfcand_phirel'] = candp4.delta_phi(self.jetp4)
data['pfcand_etarel'] = self.eta_sign * (candp4.eta - self.jetp4.eta)
data['pfcand_abseta'] = np.abs(candp4.eta)
data['pfcand_pt_log_nopuppi'] = np.log(candp4.pt)
data['pfcand_e_log_nopuppi'] = np.log(candp4.energy)
chi2 = pf('trkChi2')
chi2.content[chi2.content == -1] = 999
data['pfcand_normchi2'] = np.floor(chi2)
data['pfcand_btagEtaRel'] = pf('btagEtaRel')
data['pfcand_btagPtRatio'] = pf('btagPtRatio')
data['pfcand_btagPParRatio'] = pf('btagPParRatio')
data['pfcand_btagSip3dVal'] = pf('btagSip3dVal')
data['pfcand_btagSip3dSig'] = pf('btagSip3dSig')
data['pfcand_btagJetDistVal'] = pf('btagJetDistVal')
self._finalize_data(data)
self.data.update(data)
def get_metp4(met):
return TLorentzVectorArray.from_ptetaphim(met.pt, met.pt * 0, met.phi,
met.pt * 0)
def buildevents(df, fatjet='FatJet', subjet='SubJet', usemask=False, virtual=False):
events = ak.Table.named('event')
if virtual:
get = _getvirtual(df)
else:
get = df.__getitem__
if 'genWeight' in df:
events['genWeight'] = get('genWeight')
events['Pileup_nPU'] = get('Pileup_nPU')
events['genpart'] = ak.JaggedArray.fromcounts(
get('nGenPart'),
ak.Table.named(
'particle',
p4=TLorentzVectorArray.from_ptetaphim(
get('GenPart_pt'),
get('GenPart_eta'),
get('GenPart_phi'),
get('GenPart_mass'),
),
pdgId=get('GenPart_pdgId'),
genPartIdxMother=get('GenPart_genPartIdxMother'),
statusFlags=get('GenPart_statusFlags'),
),
)
events['fatjets'] = ak.JaggedArray.fromcounts(
get(f'n{fatjet}'),
ak.Table.named(
'fatjet',
p4=TLorentzVectorArray.from_ptetaphim(
get(f'{fatjet}_pt'),
get(f'{fatjet}_eta'),
get(f'{fatjet}_phi'),
get(f'{fatjet}_mass'),
),
msoftdrop=get(f'{fatjet}_msoftdrop'), # ak.MaskedArray(get(f'{fatjet}_msoftdrop') <= 0, get(f'{fatjet}_msoftdrop')) if usemask else np.maximum(1e-5, get(f'{fatjet}_msoftdrop')),
area=get(f'{fatjet}_area'),
n2=get(f'{fatjet}_n2b1'),
btagDDBvL=get(f'{fatjet}_btagDDBvL'),
btagDDCvL=get(f'{fatjet}_btagDDCvL'),
btagDDCvB=get(f'{fatjet}_btagDDCvB'),
jetId=get(f'{fatjet}_jetId'),
subJetIdx1=get(f'{fatjet}_subJetIdx1'),
subJetIdx2=get(f'{fatjet}_subJetIdx2'),
),
)
events['subjets'] = ak.JaggedArray.fromcounts(
get(f'n{subjet}'),
ak.Table.named(
'subjet',
p4=TLorentzVectorArray.from_ptetaphim(
get(f'{subjet}_pt'),
get(f'{subjet}_eta'),
get(f'{subjet}_phi'),
get(f'{subjet}_mass'),
),
# msoftdrop=ak.MaskedArray(df[f'{fatjet}_msoftdrop'] <= 0, df[f'{fatjet}_msoftdrop']) if usemask else np.maximum(1e-5, df[f'{fatjet}_msoftdrop']),
msoftdrop=df[f'{fatjet}_msoftdrop'],
area=df[f'{fatjet}_area'],
n2=df[f'{fatjet}_n2b1'],
jetId=df[f'{fatjet}_jetId'],
lsf3=df[f'{fatjet}_lsf3'],
muonIdx3SJ=df[f'{fatjet}_muonIdx3SJ'],
electronIdx3SJ=df[f'{fatjet}_electronIdx3SJ'],
subJetIdx1=df[f'{fatjet}_subJetIdx1'],
subJetIdx2=df[f'{fatjet}_subJetIdx2'],
),
)
events['subjets'] = ak.JaggedArray.fromcounts(
df['nCustomAK8PuppiSubJet'],
ak.Table.named(
'subjet',
p4=TLorentzVectorArray.from_ptetaphim(
df['CustomAK8PuppiSubJet_pt'],
df['CustomAK8PuppiSubJet_eta'],
df['CustomAK8PuppiSubJet_phi'],
df['CustomAK8PuppiSubJet_mass'],
),
btagDeepB=df['CustomAK8PuppiSubJet_btagDeepB'],
),
)
_embed_subjets(events)
events['jets'] = ak.JaggedArray.fromcounts(
get('nJet'),
ak.Table.named(
'jet',
p4=TLorentzVectorArray.from_ptetaphim(
get('Jet_pt'),
get('Jet_eta'),
get('Jet_phi'),
get('Jet_mass'),
),
deepcsvb=get('Jet_btagDeepB'),
hadronFlavor=get('Jet_hadronFlavour'),
jetId=get('Jet_jetId'),
),
)
events['met'] = VirtualTVector2Array(
TVector2Array.from_polar,
(get('MET_pt'), get('MET_phi')),
type=ak.type.ArrayType(df.size, RhoPhiArrayMethods),
)
events['electrons'] = ak.JaggedArray.fromcounts(
get('nElectron'),
ak.Table.named(
'electron',
p4=TLorentzVectorArray.from_ptetaphim(
get('Electron_pt'),
get('Electron_eta'),
get('Electron_phi'),
get('Electron_mass'),
),
cutBased=get('Electron_cutBased'),
),
)
events['muons'] = ak.JaggedArray.fromcounts(
get('nMuon'),
ak.Table.named(
'muon',
p4=TLorentzVectorArray.from_ptetaphim(
get('Muon_pt'),
get('Muon_eta'),
get('Muon_phi'),
get('Muon_mass'),
),
sip3d=df['Muon_sip3d'],
dxy=df['Muon_dxy'],
dz=df['Muon_dz'],
mvaId=df['Muon_mvaId'],
looseId=df['Muon_looseId'],
pfRelIso04_all=df['Muon_pfRelIso04_all'],
miniPFRelIso_all=df['Muon_miniPFRelIso_all'],
),
)
events['taus'] = ak.JaggedArray.fromcounts(
get('nTau'),
ak.Table.named(
'tau',
p4=TLorentzVectorArray.from_ptetaphim(
get('Tau_pt'),
get('Tau_eta'),
get('Tau_phi'),
get('Tau_mass'),
),
idDecayMode=get('Tau_idDecayMode'),
),
)
return events
def run_fastmtt(self, lltt, met, category, cutflow=False):
# choose the correct lepton mass
ele_mass, mu_mass = 0.511*10**-3, 0.105
l_mass = ele_mass if category[:2] == 'ee' else mu_mass
# flatten the final state leptons, assign to 4-vector arrays
l1_p4_array = TLorentzVectorArray.from_ptetaphim(lltt.i0.pt.flatten(),
lltt.i0.eta.flatten(),
lltt.i0.phi.flatten(),
l_mass*np.ones(len(lltt)))
l2_p4_array = TLorentzVectorArray.from_ptetaphim(lltt.i1.pt.flatten(),
lltt.i1.eta.flatten(),
lltt.i1.phi.flatten(),
l_mass*np.ones(len(lltt)))
# choose the correct tau decay modes
e_decay = ROOT.MeasuredTauLepton.kTauToElecDecay
m_decay = ROOT.MeasuredTauLepton.kTauToMuDecay
had_decay = ROOT.MeasuredTauLepton.kTauToHadDecay
if (category[2:]=='et'): t1_decay, t2_decay = e_decay, had_decay
if (category[2:]=='em'): t1_decay, t2_decay = e_decay, m_decay
if (category[2:]=='mt'): t1_decay, t2_decay = m_decay, had_decay
if (category[2:]=='tt'): t1_decay, t2_decay = had_decay, had_decay
# flatten the final state taus, assign to 4-vector arrays
t1_p4_array = TLorentzVectorArray.from_ptetaphim(lltt.i2.pt.flatten(),
lltt.i2.eta.flatten(),
lltt.i2.phi.flatten(),
lltt.i2.mass.flatten())
t2_p4_array = TLorentzVectorArray.from_ptetaphim(lltt.i3.pt.flatten(),
lltt.i3.eta.flatten(),
lltt.i3.phi.flatten(),
lltt.i3.mass.flatten())
# flatten MET arrays
metx = met.pt*np.cos(met.phi).flatten()
mety = met.pt*np.sin(met.phi).flatten()
metcov00, metcov11 = met.covXX.flatten(), met.covYY.flatten()
metcov01, metcov10 = met.covXY.flatten(), met.covXY.flatten()
# loop to calculate A mass
N = len(t1_p4_array)
tt_corr_masses, tt_cons_masses = np.zeros(N), np.zeros(N)
A_corr_masses, A_cons_masses = np.zeros(N), np.zeros(N)
for i in range(N):
metcov = ROOT.TMatrixD(2,2)
metcov[0][0], metcov[1][1] = metcov00[i], metcov11[i]
metcov[0][1], metcov[1][0] = metcov01[i], metcov10[i]
tau_vector = ROOT.std.vector('MeasuredTauLepton')
tau_pair = tau_vector()
t1 = ROOT.MeasuredTauLepton(t1_decay,
t1_p4_array[i].pt,
t1_p4_array[i].eta,
t1_p4_array[i].phi,
t1_p4_array[i].mass)
t2 = ROOT.MeasuredTauLepton(t2_decay,
t2_p4_array[i].pt,
t2_p4_array[i].eta,
t2_p4_array[i].phi,
t2_p4_array[i].mass)
tau_pair.push_back(t1)
tau_pair.push_back(t2)
# run SVfit algorithm
fastmtt = ROOT.FastMTT()
fastmtt.run(tau_pair, metx[i], mety[i], metcov, False) # unconstrained
tt_corr = fastmtt.getBestP4()
tt_corr_p4 = ROOT.TLorentzVector()
tt_corr_p4.SetPtEtaPhiM(tt_corr.Pt(), tt_corr.Eta(),
tt_corr.Phi(), tt_corr.M())
fastmtt.run(tau_pair, metx[i], mety[i], metcov, True) # constrained
tt_cons = fastmtt.getBestP4()
tt_cons_p4 = ROOT.TLorentzVector()
tt_cons_p4.SetPtEtaPhiM(tt_cons.Pt(), tt_cons.Eta(),
tt_cons.Phi(), tt_cons.M())
tt_corr_masses[i] = tt_corr_p4.M()
tt_cons_masses[i] = tt_cons_p4.M()
l1, l2 = ROOT.TLorentzVector(), ROOT.TLorentzVector()
l1.SetPtEtaPhiM(l1_p4_array[i].pt, l1_p4_array[i].eta,
l1_p4_array[i].phi, l1_p4_array[i].mass)
l2.SetPtEtaPhiM(l2_p4_array[i].pt, l2_p4_array[i].eta,
l2_p4_array[i].phi, l2_p4_array[i].mass)
A_corr_p4 = (l1 + l2 + tt_corr_p4)
A_cons_p4 = (l1 + l2 + tt_cons_p4)
A_corr_masses[i] = A_corr_p4.M()
A_cons_masses[i] = A_cons_p4.M()
return tt_corr_masses, tt_cons_masses, A_corr_masses, A_cons_masses
