def rebuild(leptons):
return JaggedCandidateArray.candidatesfromoffsets(
leptons.offsets,
pt=leptons.pt.flatten(),
eta=leptons.eta.flatten(),
phi=leptons.phi.flatten(),
mass=leptons.mass.flatten(),
charge=leptons.charge.flatten(),
pdgId=leptons.pdgId.flatten(),
# needed for electron SF
etaSC=leptons.etaSC.flatten()
if hasattr(leptons, 'etaSC') else leptons.eta.flatten(),
)
def process(self, df):
output = self.accumulator.identity()
# GenParticles
dataset = df['dataset']
if df['nGenPart'].size != 0:
GenPart = JaggedCandidateArray.candidatesfromcounts(
df['nGenPart'],
pt=df['GenPart_pt'],
eta=df['GenPart_eta'],
phi=df['GenPart_phi'],
mass=df['GenPart_mass'],
charge=df['GenPart_charge'],
pdgId=df['GenPart_pdgId'],
vx=df['GenPart_vx'],
vy=df['GenPart_vy'],
vz=df['GenPart_vz'],
mpdgId=df['GenPart_mpdgId'],
mvx=df['GenPart_mvx'],
mvy=df['GenPart_mvy'],
mvz=df['GenPart_mvz'],
)
else:
GenPart = JaggedCandidateArray.candidatesfromcounts(
np.array([]),
pt=np.array([]),
eta=np.array([]),
phi=np.array([]),
mass=np.array([]),
charge=np.array([]),
pdgId=np.array([]),
vx=np.array([]),
vy=np.array([]),
vz=np.array([]),
mpdgId=np.array([]),
mvx=np.array([]),
mvy=np.array([]),
mvz=np.array([]),
)
# Particle selection
muonid = (np.absolute(GenPart.pdgId) == 13)
Muon = GenPart[muonid]
upsilonid = (np.absolute(GenPart.pdgId) == 553)
Upsilon = GenPart[upsilonid]
Dimuon = Muon.distincts()
opposite_charge = (Dimuon.i0['charge'] * Dimuon.i1['charge'] < 0)
Dimuon = Dimuon[opposite_charge]
same_vtx = ((Dimuon.i0['vx'] == Dimuon.i1['vx']) &
(Dimuon.i0['vy'] == Dimuon.i1['vy']) &
(Dimuon.i0['vz'] == Dimuon.i1['vz']))
Dimuon = Dimuon[same_vtx]
mass_cut = ((Dimuon.mass < 12) & (Dimuon.mass > 7))
Dimuon = Dimuon[mass_cut]
leading_mu = (Dimuon.i0.pt.content > Dimuon.i1.pt.content)
Muon_lead = JaggedCandidateArray.candidatesfromoffsets(
Dimuon.offsets,
pt=np.where(leading_mu, Dimuon.i0.pt.content,
Dimuon.i1.pt.content),
eta=np.where(leading_mu, Dimuon.i0.eta.content,
Dimuon.i1.eta.content),
phi=np.where(leading_mu, Dimuon.i0.phi.content,
Dimuon.i1.phi.content),
mass=np.where(leading_mu, Dimuon.i0.mass.content,
Dimuon.i1.mass.content),
mpdgId=np.where(leading_mu, Dimuon.i0.mpdgId.content,
Dimuon.i1.mpdgId.content))
Muon_trail = JaggedCandidateArray.candidatesfromoffsets(
Dimuon.offsets,
pt=np.where(~leading_mu, Dimuon.i0.pt.content,
Dimuon.i1.pt.content),
eta=np.where(~leading_mu, Dimuon.i0.eta.content,
Dimuon.i1.eta.content),
phi=np.where(~leading_mu, Dimuon.i0.phi.content,
Dimuon.i1.phi.content),
mass=np.where(~leading_mu, Dimuon.i0.mass.content,
Dimuon.i1.mass.content),
mpdgId=np.where(leading_mu, Dimuon.i0.mpdgId.content,
Dimuon.i1.mpdgId.content))
output['Muon_lead_pt'].fill(dataset=dataset, pt=Muon_lead.pt.flatten())
output['Muon_trail_pt'].fill(dataset=dataset,
pt=Muon_trail.pt.flatten())
output['Muon_eta'].fill(dataset=dataset, eta=Muon_lead.eta.flatten())
output['Muon_eta'].fill(dataset=dataset, eta=Muon_trail.eta.flatten())
output['Muon_phi'].fill(dataset=dataset, phi=Muon_lead.phi.flatten())
output['Muon_phi'].fill(dataset=dataset, phi=Muon_trail.phi.flatten())
output['Dimuon_mass'].fill(dataset=dataset, mass=Dimuon.mass.flatten())
output['Dimuon_pt'].fill(dataset=dataset, pt=Dimuon.pt.flatten())
output['Dimuon_eta'].fill(dataset=dataset, eta=Dimuon.eta.flatten())
output['Dimuon_phi'].fill(dataset=dataset, phi=Dimuon.phi.flatten())
return output
def process(self, df):
output = self.accumulator.identity()
if df.size == 0:
return processor.dict_accumulator({
'foo':
processor.defaultdict_accumulator(int),
'cutflow':
output['cutflow']
})
# Dimu candidates
if df['nDimu'].size != 0:
Dimu = JaggedCandidateArray.candidatesfromcounts(
df['nDimu'], **get_vars_dict(df, dimu_cols))
else:
Dimu = JaggedCandidateArray.candidatesfromcounts(
np.array([]), **get_vars_dict(df, dimu_cols))
# Muon candidates
if df['nMuon'].size != 0:
Muon = JaggedCandidateArray.candidatesfromcounts(
df['nMuon'], **get_vars_dict(df, muon_cols))
else:
Muon = JaggedCandidateArray.candidatesfromcounts(
np.array([]), **get_vars_dict(df, muon_cols))
# D0 candidates
if df['nD0'].size != 0:
D0 = JaggedCandidateArray.candidatesfromcounts(
df['nD0'], mass=df['D0_mass12'], **get_vars_dict(df, d0_cols))
else:
D0 = JaggedCandidateArray.candidatesfromcounts(np.array([]),
mass=np.array([]),
**get_vars_dict(
df, d0_cols))
# Dstar candidates
if df['nDstar'].size != 0:
Dstar = JaggedCandidateArray.candidatesfromcounts(
df['nDstar'],
mass=df['Dstar_deltam'] + df['DstarD0_mass'],
**get_vars_dict(df, dstar_cols))
else:
Dstar = JaggedCandidateArray.candidatesfromcounts(
np.array([]),
mass=np.array([]),
**get_vars_dict(df, dstar_cols))
output['cutflow']['all events'] += Muon.size
output['cutflow']['all Dimu'] += Dimu.counts.sum()
output['cutflow']['all D0'] += D0.counts.sum()
output['cutflow']['all Dstar'] += Dstar.counts.sum()
############### Cuts
# Dimu cuts: charge = 0, mass cuts and chi2...
Dimu = Dimu[Dimu.charge == 0]
output['cutflow']['Dimu 0 charge'] += Dimu.counts.sum()
dimu_mass_cut = (Dimu.mass > 8.5) & (Dimu.mass < 11.5)
Dimu = Dimu[dimu_mass_cut]
output['cutflow']['Upsilon mass'] += Dimu.counts.sum()
dimu_chi2_cut = (Dimu.chi2 < 5.)
Dimu = Dimu[dimu_chi2_cut]
output['cutflow']['chi2 cut'] = Dimu.counts.sum()
############### Get the Muons from Dimu, for cuts in their params
if df['nDimu'].size != 0:
mu1Idx = (Dimu.t1_muIdx + Muon.starts).content
mu2Idx = (Dimu.t2_muIdx + Muon.starts).content
Muon1 = JaggedCandidateArray.fromoffsets(Dimu.offsets,
Muon.content[mu1Idx])
Muon2 = JaggedCandidateArray.fromoffsets(Dimu.offsets,
Muon.content[mu2Idx])
else:
Muon1 = JaggedCandidateArray.fromoffsets(
Dimu.offsets, Muon.content[np.array([], dtype='int64')])
Muon2 = JaggedCandidateArray.fromoffsets(
Dimu.offsets, Muon.content[np.array([], dtype='int64')])
# SoftId and Global Muon cuts
soft_id = (Muon1.softId > 0) & (Muon2.softId > 0)
Dimu = Dimu[soft_id]
Muon1 = Muon1[soft_id]
Muon2 = Muon2[soft_id]
output['cutflow']['Dimu muon softId'] += Dimu.counts.sum()
global_muon = (Muon1.isGlobal > 0) & (Muon2.isGlobal > 0)
Dimu = Dimu[global_muon]
Muon1 = Muon1[global_muon]
Muon2 = Muon2[global_muon]
output['cutflow']['Dimu muon global'] += Dimu.counts.sum()
# pt and eta cuts
muon_pt_cut = (Muon1.pt > 3) & (Muon2.pt > 3)
Dimu = Dimu[muon_pt_cut]
Muon1 = Muon1[muon_pt_cut]
Muon2 = Muon2[muon_pt_cut]
output['cutflow']['Dimu muon pt cut'] += Dimu.counts.sum()
muon_eta_cut = (np.absolute(Muon1.eta) <= 2.4) & (np.absolute(
Muon2.eta) <= 2.4)
Dimu = Dimu[muon_eta_cut]
Muon1 = Muon1[muon_eta_cut]
Muon2 = Muon2[muon_eta_cut]
output['cutflow']['Dimu muon eta cut'] += Dimu.counts.sum()
############### Cuts for D0
# trk cuts
D0_trk_muon_cut = ~D0.hasMuon
D0 = D0[D0_trk_muon_cut]
output['cutflow']['D0 trk muon cut'] += D0.counts.sum()
D0_trk_pt_cut = (D0.t1_pt > 0.8) & (D0.t2_pt > 0.8)
D0 = D0[D0_trk_pt_cut]
output['cutflow']['D0 trk pt cut'] += D0.counts.sum()
D0_trk_chi2_cut = (D0.t1_chindof < 2.5) & (D0.t2_chindof < 2.5)
D0 = D0[D0_trk_chi2_cut]
output['cutflow']['D0 trk chi2 cut'] += D0.counts.sum()
D0_trk_hits_cut = (D0.t1_nValid > 4) & (D0.t2_nValid > 4) & (
D0.t1_nPix > 1) & (D0.t2_nPix > 1)
D0 = D0[D0_trk_hits_cut]
output['cutflow']['D0 trk hits cut'] += D0.counts.sum()
D0_trk_dxy_cut = (D0.t1_dxy < 0.1) & (D0.t2_dxy < 0.1)
D0 = D0[D0_trk_dxy_cut]
output['cutflow']['D0 trk dxy cut'] += D0.counts.sum()
D0_trk_dz_cut = (D0.t1_dz < 1.) & (D0.t2_dz < 1.)
D0 = D0[D0_trk_dz_cut]
output['cutflow']['D0 trk dz cut'] += D0.counts.sum()
# D0 cosphi
D0_cosphi_cut = (D0.cosphi > 0.99)
D0 = D0[D0_cosphi_cut]
output['cutflow']['D0 cosphi cut'] += D0.counts.sum()
# D0 dl Significance
D0_dlSig_cut = (D0.dlSig > 5.)
D0 = D0[D0_dlSig_cut]
output['cutflow']['D0 dlSig cut'] += D0.counts.sum()
# D0 pt
D0_pt_cut = (D0.pt > 3.)
D0 = D0[D0_pt_cut]
output['cutflow']['D0 pt cut'] += D0.counts.sum()
############### Cuts for Dstar
# trks cuts
Dstar_trk_muon_cut = ~Dstar.hasMuon
Dstar = Dstar[Dstar_trk_muon_cut]
output['cutflow']['Dstar trk muon cut'] += Dstar.counts.sum()
Dstar_trk_pt_cut = (Dstar.K_pt > 0.5) & (Dstar.pi_pt > 0.5)
Dstar = Dstar[Dstar_trk_pt_cut]
output['cutflow']['Dstar trk pt cut'] += Dstar.counts.sum()
Dstar_trk_chi2_cut = (Dstar.K_chindof < 2.5) & (Dstar.pi_chindof < 2.5)
Dstar = Dstar[Dstar_trk_chi2_cut]
output['cutflow']['Dstar trk pt cut'] += Dstar.counts.sum()
Dstar_trk_hits_cut = (Dstar.K_nValid > 4) & (Dstar.pi_nValid > 4) & (
Dstar.K_nValid > 1) & (Dstar.pi_nValid > 1)
Dstar = Dstar[Dstar_trk_hits_cut]
output['cutflow']['Dstar trk hits cut'] += Dstar.counts.sum()
Dstar_trk_dxy_cut = (Dstar.K_dxy < 0.1) & (Dstar.pi_dxy < 0.1)
Dstar = Dstar[Dstar_trk_dxy_cut]
output['cutflow']['Dstar trk pt cut'] += Dstar.counts.sum()
Dstar_trk_dz_cut = (Dstar.K_dz < 1) & (Dstar.pi_dz < 1)
Dstar = Dstar[Dstar_trk_dz_cut]
output['cutflow']['Dstar trk pt cut'] += Dstar.counts.sum()
# pis cuts
Dstar_pis_pt_cut = (Dstar.pis_pt > 0.3)
Dstar = Dstar[Dstar_pis_pt_cut]
output['cutflow']['Dstar pis pt cut'] += Dstar.counts.sum()
Dstar_pis_chi2_cut = (Dstar.pis_chindof < 3)
Dstar = Dstar[Dstar_pis_chi2_cut]
output['cutflow']['Dstar pis chi2 cut'] += Dstar.counts.sum()
Dstar_pis_hits_cut = (Dstar.pis_nValid > 2)
Dstar = Dstar[Dstar_pis_hits_cut]
output['cutflow']['Dstar pis hits cut'] += Dstar.counts.sum()
# D0 of Dstar cuts
DstarD0_cosphi_cut = (Dstar.D0_cosphi > 0.99)
Dstar = Dstar[DstarD0_cosphi_cut]
output['cutflow']['Dstar D0 cosphi cut'] += Dstar.counts.sum()
DstarD0_mass_cut = (Dstar.D0_mass < D0_PDG_MASS + 0.025) & (
Dstar.D0_mass > D0_PDG_MASS - 0.025)
Dstar = Dstar[DstarD0_mass_cut]
output['cutflow']['Dstar D0 mass cut'] += Dstar.counts.sum()
DstarD0_pt_cut = (Dstar.D0_pt > 3)
Dstar = Dstar[DstarD0_pt_cut]
output['cutflow']['Dstar D0 pt cut'] += Dstar.counts.sum()
DstarD0_dlSig_cut = (Dstar.D0_dlSig > 3)
Dstar = Dstar[DstarD0_dlSig_cut]
output['cutflow']['Dstar D0 dlSig cut'] += Dstar.counts.sum()
""" Dstar_wrong_charge_cut = (Dstar.K_chg != Dstar.pi_chg)
Dstar = Dstar[Dstar_wrong_charge_cut]
output['cutflow']['Dstar wrong charge cut'] += Dstar.counts.sum() """
############### Upsilon + Dstar association
############### Final computation of number of objects
output['cutflow']['Dimu final'] += Dimu.counts.sum()
output['cutflow']['D0 final'] += D0.counts.sum()
output['cutflow']['Dstar final'] += Dstar.counts.sum()
############### Leading and Trailing muon separation
leading_mu = (Muon1.pt.content > Muon2.pt.content)
Muon_lead = JaggedCandidateArray.candidatesfromoffsets(
Dimu.offsets,
pt=np.where(leading_mu, Muon1.pt.content, Muon2.pt.content),
eta=np.where(leading_mu, Muon1.eta.content, Muon2.eta.content),
phi=np.where(leading_mu, Muon1.phi.content, Muon2.phi.content),
mass=np.where(leading_mu, Muon1.mass.content, Muon2.mass.content),
)
Muon_trail = JaggedCandidateArray.candidatesfromoffsets(
Dimu.offsets,
pt=np.where(~leading_mu, Muon1.pt.content, Muon2.pt.content),
eta=np.where(~leading_mu, Muon1.eta.content, Muon2.eta.content),
phi=np.where(~leading_mu, Muon1.phi.content, Muon2.phi.content),
mass=np.where(~leading_mu, Muon1.mass.content, Muon2.mass.content),
)
############### Create the accumulators to save output
muon_lead_acc = processor.dict_accumulator({})
for var in Muon_lead.columns:
if var == 'p4': continue
muon_lead_acc[var] = processor.column_accumulator(
np.array(Muon_lead[var].flatten()))
muon_lead_acc["nMuon"] = processor.column_accumulator(Muon_lead.counts)
output["Muon_lead"] = muon_lead_acc
muon_trail_acc = processor.dict_accumulator({})
for var in Muon_trail.columns:
if var == 'p4': continue
muon_trail_acc[var] = processor.column_accumulator(
np.array(Muon_trail[var].flatten()))
muon_trail_acc["nMuon"] = processor.column_accumulator(
Muon_trail.counts)
output["Muon_trail"] = muon_trail_acc
dimu_acc = processor.dict_accumulator({})
for var in Dimu.columns:
if (var == 'p4' or var.startswith('t')): continue
dimu_acc[var] = processor.column_accumulator(
np.array(Dimu[var].flatten()))
dimu_acc["nDimu"] = processor.column_accumulator(Dimu.counts)
output["Dimu"] = dimu_acc
D0_acc = processor.dict_accumulator({})
D0_trk_acc = processor.dict_accumulator({})
for var in D0.columns:
if (var == 'p4'): continue
elif (var.startswith('t')):
D0_trk_acc[var] = processor.column_accumulator(
np.array(D0[var].flatten()))
else:
D0_acc[var] = processor.column_accumulator(
np.array(D0[var].flatten()))
D0_acc["nD0"] = processor.column_accumulator(D0.counts)
output["D0"] = D0_acc
output["D0_trk"] = D0_trk_acc
Dstar_acc = processor.dict_accumulator({})
Dstar_D0_acc = processor.dict_accumulator({})
Dstar_trk_acc = processor.dict_accumulator({})
for var in Dstar.columns:
if (var == 'p4'): continue
elif var.startswith('D0'):
Dstar_D0_acc[var] = processor.column_accumulator(
np.array(Dstar[var].flatten()))
elif (var.startswith('K') or var.startswith('pi')):
Dstar_trk_acc[var] = processor.column_accumulator(
np.array(Dstar[var].flatten()))
else:
Dstar_acc[var] = processor.column_accumulator(
np.array(Dstar[var].flatten()))
Dstar_acc["nDstar"] = processor.column_accumulator(Dstar.counts)
output["Dstar"] = Dstar_acc
output["Dstar_D0"] = Dstar_D0_acc
output["Dstar_trk"] = Dstar_trk_acc
file_hash = str(random.getrandbits(128)) + str(df.size)
save(
output, "output/" + self.analyzer_name + "/" + self.analyzer_name +
"_" + file_hash + ".coffea")
# return dummy accumulator
return processor.dict_accumulator({
'foo':
processor.defaultdict_accumulator(int),
'cutflow':
output['cutflow']
})
