def analysis(self, df):
djmm = fjtools.DJetMatchMaker()
djmm.set_ch_pt_eta_phi(self.df_tracks['ParticlePt'].values,
self.df_tracks['ParticleEta'].values,
self.df_tracks['ParticlePhi'].values)
djmm.set_Ds_pt_eta_phi_m(df['pt_cand'].values, df['eta_cand'].values,
df['phi_cand'].values, df['inv_mass'].values)
djmm.set_daughters0_pt_eta_phi(df['pt_prong0'].values,
df['eta_prong0'].values,
df['phi_prong0'].values)
djmm.set_daughters1_pt_eta_phi(df['pt_prong1'].values,
df['eta_prong1'].values,
df['phi_prong1'].values)
self.tw.fill_branches(dpsj=djmm.Ds)
self.tw.fill_tree()
for id0, d0 in enumerate(djmm.Ds):
_parts_and_ds = djmm.match(0.005, id0)
_parts_and_ds.push_back(d0)
ja = jet_analysis.JetAnalysis(jet_R=0.4,
particle_eta_max=0.9,
jet_pt_min=2.0)
ja.analyze_event(_parts_and_ds)
if len(ja.jets) < 1:
continue
jets = ja.jets_as_psj_vector()
djets = djmm.filter_D0_jets(jets)
if len(djets) > 0:
j = djets[0]
dcand = djmm.get_Dcand_in_jet(j)
sja = jet_analysis.JetAnalysis(jet_R=0.1,
particle_eta_max=0.9,
jet_pt_min=2.0)
sja.analyze_event(j.constituents())
lsj = fj.sorted_by_pt(sja.jets_as_psj_vector())
if len(lsj) < 1:
continue
sj_dcand = djmm.get_Dcand_in_jet(lsj[0])
is_Dsj = 0
if len(sj_dcand) > 0:
# if sj_dcand[0].m() == dcand[0].m() and sj_dcand[0].perp() == dcand[0].perp():
if sj_dcand[0].delta_R(dcand[0]) == 0.0:
is_Dsj = 1
self.twjc.fill_branches(jet=j,
dR=j.delta_R(dcand[0]),
D=dcand[0],
lsj=lsj[0],
Dsj=is_Dsj,
a10=fjext.angularity(j, 1.0, 0.4),
a15=fjext.angularity(j, 0.5, 0.4),
a20=fjext.angularity(j, 0.0, 0.4),
a30=fjext.angularity(j, -1.0, 0.4))
self.twjc.fill_tree()
if len(djets) > 1:
perror("more than one jet per D candidate?")
return True
def load_file(self, path):
if not os.path.exists(path):
pwarning('[w] file', path, 'does not exists.')
return
try:
event_tree = uproot.open(path)[self.event_tree_name]
except:
pwarning('error getting', self.event_tree_name, 'from file:', path)
return False
if not event_tree:
perror('Tree {} not found in file {}'.format(self.event_tree_name, path))
return False
event_df_orig = event_tree.pandas.df(['run_number', 'ev_id', 'z_vtx_reco','is_ev_rej'])
event_df_orig.reset_index(drop=True)
event_df = event_df_orig.query('is_ev_rej == 0')
event_df.reset_index(drop=True)
# Load track tree into dataframe
try:
track_tree = uproot.open(path)[self.tree_name]
except:
pwarning('error getting', self.tree_name, 'from file:', path)
return False
if not track_tree:
perror('Tree {} not found in file {}'.format(tree_name, path))
return False
track_df_orig = track_tree.pandas.df()
# Merge event info into track tree
track_df = pd.merge(track_df_orig, event_df, on=['run_number', 'ev_id'])
self.track_df_grouped = track_df.groupby(['run_number','ev_id'])
# (ii) Transform the DataFrameGroupBy object to a SeriesGroupBy of fastjet particles
return True
def pd_tree(self, path, tname, squery=None):
try:
tree = uproot.open(path)[tname]
except:
pwarning('error getting', tname, 'from file:', path)
return None
if not tree:
perror('Tree {} not found in file {}'.format(tname, path))
return None
df = tree.pandas.df()
if squery:
df = df.query(squery)
df.reset_index(drop=True)
return df
def pd_tree(self, path, tname, squery=None):
try:
tree = uproot.open(path)[tname]
except:
pwarning('error getting', tname, 'from file:', path)
return None
if not tree:
perror('Tree {} not found in file {}'.format(tname, path))
return None
df = uproot.concatenate(tree, library="pd")
if squery:
#df.query(squery, inplace=True)
df = df.query(squery)
df.reset_index(drop=True)
return df
def execute_analyses_on_file_list(self, file_list, nfiles=0):
print()
if os.path.exists(file_list):
with open(file_list) as f:
files = f.readlines()
if int(nfiles) > 0:
files = files[:nfiles]
for f in files:
fn = f.strip('\n')
pinfo('+file:', fn)
for f in tqdm.tqdm(files):
fn = f.strip('\n')
pinfo('file:', fn)
if self.load_file(fn):
self.execute_analyses()
else:
perror('file list does not exist', file_list)
pinfo('done.')
def analysis_gen(self, df):
m_gen_array = np.full((self.df_gen_tracks['ParticlePt'].values.size),
0.1396)
djmm_gen = fjtools.DJetMatchMaker()
djmm_gen.set_ch_pt_eta_phi_m(self.df_gen_tracks['ParticlePt'].values,
self.df_gen_tracks['ParticleEta'].values,
self.df_gen_tracks['ParticlePhi'].values,
m_gen_array)
djmm_gen.set_Ds_pt_eta_phi(df['pt_cand'].values, df['eta_cand'].values,
df['phi_cand'].values)
for id0, d0 in enumerate(djmm_gen.Ds):
_parts_and_ds = djmm_gen.ch
_parts_and_ds.push_back(d0)
ja_gen = jet_analysis.JetAnalysis(
jet_R=0.4,
jet_RecombinationScheme=fj.E_scheme,
particle_eta_max=0.9,
jet_pt_min=5.0)
ja_gen.analyze_event(_parts_and_ds)
if len(ja_gen.jets) < 1:
continue
jets = ja_gen.jets_as_psj_vector()
djets = djmm_gen.filter_D0_jets(jets)
if len(djets) > 0:
j = djets[0]
dcand = djmm_gen.get_Dcand_in_jet(j)
D_cand_type = df['cand_type'].values
self.twjc_gen.fill_branches(jet=j,
dR=j.delta_R(dcand[0]),
D=dcand[0],
Dmeson_cand_type=float(
D_cand_type[id0]))
self.twjc_gen.fill_tree()
if len(djets) > 1:
perror("more than one jet per D candidate?")
return True
def execute_analyses_on_file_list(self, file_list, nfiles=0):
self.pbar2 = tqdm.tqdm(mininterval=20, maxinterval=60)
self.pbar2_mark = None
for a in self.analyses:
a.callback = self.update_status
print()
if os.path.exists(file_list):
with open(file_list) as f:
files = f.readlines()
if int(nfiles) > 0:
files = files[:nfiles]
for f in files:
fn = f.strip('\n')
pinfo('+file:', fn)
for f in tqdm.tqdm(files):
fn = f.strip('\n')
if self.load_file(fn):
self.execute_analyses()
self.pbar2.close()
else:
perror('file list does not exist', file_list)
pinfo('done.')
def analysis(self, df):
#print(df)
m_array = np.full((self.df_tracks['ParticlePt'].values.size), 0.1396)
djmm = fjtools.DJetMatchMaker()
djmm.set_ch_pt_eta_phi_m(self.df_tracks['ParticlePt'].values,
self.df_tracks['ParticleEta'].values,
self.df_tracks['ParticlePhi'].values, m_array)
djmm.set_Ds_pt_eta_phi_m(df['pt_cand'].values, df['eta_cand'].values,
df['phi_cand'].values, df['inv_mass'].values)
djmm.set_daughters0_pt_eta_phi(df['pt_prong0'].values,
df['eta_prong0'].values,
df['phi_prong0'].values)
djmm.set_daughters1_pt_eta_phi(df['pt_prong1'].values,
df['eta_prong1'].values,
df['phi_prong1'].values)
#run for each D candidate to build jet
for id0, d0 in enumerate(djmm.Ds):
#daughter tracks matching
_parts_and_ds = djmm.match(0.005, id0)
#replacing daughter tracks with matched D0 candidate
#including D0
_parts_and_ds.push_back(d0)
#jet reconstruction with D0 and charged particle
jetR = 0.4
ja = jet_analysis.JetAnalysis(jet_R=jetR,
jet_RecombinationScheme=fj.E_scheme,
particle_eta_max=0.9,
jet_pt_min=5.0)
ja.analyze_event(_parts_and_ds)
if len(ja.jets) < 1:
continue
jets = ja.jets_as_psj_vector()
#filtering D0 jets
djets = djmm.filter_D0_jets(jets)
if len(djets) > 0:
j = djets[0]
dcand = djmm.get_Dcand_in_jet(j)
#number of constitutents > 1
#if len(j.constituents())<=1:
# continue
#jets with the winner take all axis################################
jet_def_wta = fj.JetDefinition(fj.cambridge_algorithm,
2 * jetR)
jet_def_wta.set_recombination_scheme(fj.WTA_pt_scheme)
#print('WTA jet definition is:', jet_def_wta)
reclusterer_wta = fjcontrib.Recluster(jet_def_wta)
jet_wta = reclusterer_wta.result(j)
################################
D_cosp = df['cos_p'].values
D_cosTStar = df['cos_t_star'].values
D_NormalisedDecayLength = df['norm_dl_xy'].values
D_ImpactParameterProduct = df['imp_par_prod'].values
D_cand_type = df['cand_type'].values
self.twjc.fill_branches(
jet=j,
jetWta=jet_wta,
dR=j.delta_R(dcand[0]),
dRWTA=jet_wta.delta_R(dcand[0]),
D=dcand[0],
cos_p=float(D_cosp[id0]),
D_cos_t_star=float(D_cosTStar[id0]),
D_norm_dlxy=float(D_NormalisedDecayLength[id0]),
D_imp_par_prod=float(D_ImpactParameterProduct[id0]),
Dmeson_cand_type=float(D_cand_type[id0]))
self.twjc.fill_tree()
if len(djets) > 1:
perror("more than one jet per D candidate?")
return True
def main():
parser = argparse.ArgumentParser(description='pythia8 fastjet on the fly',
prog=os.path.basename(__file__))
pyconf.add_standard_pythia_args(parser)
parser.add_argument('--ignore-mycfg',
help="ignore some settings hardcoded here",
default=False,
action='store_true')
args = parser.parse_args()
# print the banner first
fj.ClusterSequence.print_banner()
print()
# set up our jet definition and a jet selector
jet_R0 = 0.4
jet_def = fj.JetDefinition(fj.antikt_algorithm, jet_R0)
jet_selector = fj.SelectorPtMin(80.0) & fj.SelectorPtMax(
100.0) & fj.SelectorAbsEtaMax(1)
# jet_selector = fj.SelectorPtMin(40.0) & fj.SelectorPtMax(200.0) &fj.SelectorAbsEtaMax(1)
print(jet_def)
all_jets = []
mycfg = ['PhaseSpace:pThatMin = 80']
# mycfg = ['PhaseSpace:pThatMin = 40']
if args.ignore_mycfg:
mycfg = []
pythia = pyconf.create_and_init_pythia_from_args(args, mycfg)
if not pythia:
perror("pythia initialization failed.")
return
jet_def_lund = fj.JetDefinition(fj.cambridge_algorithm, 1.0)
lund_gen = fjcontrib.LundGenerator(jet_def_lund)
print(lund_gen.description())
dy_groomer = fjcontrib.DynamicalGroomer(jet_def_lund)
print(dy_groomer.description())
# sd = fjcontrib.SoftDrop(0, 0.1, 1.0)
sd = fjcontrib.SoftDrop(0, 0.2, 1.0)
print(sd)
# jet_def_rc01 = fj.JetDefinition(fj.cambridge_algorithm, 0.1)
# jet_def_rc02 = fj.JetDefinition(fj.cambridge_algorithm, 0.2)
# print (jet_def_rc01)
# print (jet_def_rc02)
# rc = fjcontrib.Recluster(jet_def_rc, True)
jet_def_rc01 = fj.JetDefinition(fj.antikt_algorithm, 0.1)
jet_def_rc02 = fj.JetDefinition(fj.antikt_algorithm, 0.2)
print(jet_def_rc01)
print(jet_def_rc02)
#rc = fjcontrib.Recluster(jet_def_rc, True)
# tw = treewriter.RTreeWriter(name = 'lsjvsx', file_name = 'leadsj_vs_x.root')
tw = treewriter.RTreeWriter(name='lsjvsx',
file_name='leadsj_vs_x_bias80.root')
if args.nev < 100:
args.nev = 100
for i in tqdm.tqdm(range(args.nev)):
if not pythia.next():
continue
# parts = pythiafjext.vectorize(pythia, True, -1, 1, False)
partons = pythiafjext.vectorize_select(pythia, [pythiafjext.kParton],
0, True)
parts = pythiafjext.vectorize_select(
pythia, [pythiafjext.kFinal, pythiafjext.kCharged], 0, False)
# parts = pythiafjext.vectorize_select(pythia, [pythiafjext.kFinal], 0, False)
jets = jet_selector(jet_def(parts))
# for j in tqdm.tqdm(jets):
for j in jets:
j_type = match_dR(j, partons, jet_R0 / 2.)
j_sd = sd.result(j)
sd_info = fjcontrib.get_SD_jet_info(j_sd)
rc_sjets01 = fj.sorted_by_pt(jet_def_rc01(j.constituents()))
rc_sjets02 = fj.sorted_by_pt(jet_def_rc02(j.constituents()))
tw.fill_branches(
j=j,
lund=[ls for ls in lund_gen.result(j)],
dyg1=dy_groomer.result(j, 1),
sd=j_sd,
sd_z=sd_info.z,
sd_mu=sd_info.mu,
sd_Delta=sd_info.dR,
lsjet01=rc_sjets01[0],
nsjet01=len(rc_sjets01),
sjet01=rc_sjets01,
lsjet02=rc_sjets02[0],
nsjet02=len(rc_sjets02),
sjet02=rc_sjets02,
ppid=j_type[0],
pquark=j_type[1],
pglue=j_type[2] # this is redundancy
)
tw.fill_tree()
pythia.stat()
tw.write_and_close()
def main():
parser = argparse.ArgumentParser(description='pythia8 fastjet on the fly', prog=os.path.basename(__file__))
pyconf.add_standard_pythia_args(parser)
parser.add_argument('--nw', help="no warn", default=False, action='store_true')
parser.add_argument('--ignore-mycfg', help="ignore some settings hardcoded here", default=False, action='store_true')
parser.add_argument('--enable-background', help="enable background calc", default=False, action='store_true')
parser.add_argument('--output', help="output file name", default='leadsj_vs_x_output.root', type=str)
# for background
parser.add_argument('--cent-bin', help="centraility bin 0 is the 0-5 percent most central bin", type=int, default=0)
parser.add_argument('--seed', help="pr gen seed", type=int, default=1111)
parser.add_argument('--harmonics', help="set harmonics flag (0 : v1 - v5) , (1 : v2 - v5) , (2: v3 - v5) , (3: v1 - v4) , (4: v1 - v3) , (5: uniform dN/dphi no harmonics) , (6 : v1 - v2 , v4 - v5) , (7 : v1 - v3 , v5) , (8 : v1 , v3 - v5) , (9 : v1 only) , (10 : v2 only) , (11 : v3 only) , (12 : v4 only) , (13 : v5 only)",
type=int, default=5)
parser.add_argument('--eta', help="set eta range must be uniform (e.g. abs(eta) < 0.9, which is ALICE TPC fiducial acceptance)",
type=float, default=0.9)
parser.add_argument('--qa', help="PrintOutQAHistos", default=False, action='store_true')
parser.add_argument('--dRmax', default=0.25, type=float)
parser.add_argument('--alpha', default=0, type=float)
args = parser.parse_args()
# print the banner first
fj.ClusterSequence.print_banner()
print()
# set up our jet definition and a jet selector
jet_R0 = 0.4
jet_def = fj.JetDefinition(fj.antikt_algorithm, jet_R0)
jet_selector = fj.SelectorPtMin(args.py_pthatmin) & fj.SelectorPtMax(1000.0) & fj.SelectorAbsEtaMax(args.eta - jet_R0)
# jet_selector = fj.SelectorPtMin(40.0) & fj.SelectorPtMax(200.0) &fj.SelectorAbsEtaMax(1)
print(jet_def)
all_jets = []
# mycfg = ['PhaseSpace:pThatMin = 80']
# mycfg = ['PhaseSpace:pThatMin = 40']
mycfg = ['']
if args.ignore_mycfg:
mycfg = []
pythia = pyconf.create_and_init_pythia_from_args(args, mycfg)
if not pythia:
perror("pythia initialization failed.")
return
jet_def_lund = fj.JetDefinition(fj.cambridge_algorithm, 1.0)
lund_gen = fjcontrib.LundGenerator(jet_def_lund)
print (lund_gen.description())
dy_groomer = fjcontrib.DynamicalGroomer(jet_def_lund)
print (dy_groomer.description())
# sd = fjcontrib.SoftDrop(0, 0.1, 1.0)
sd01 = fjcontrib.SoftDrop(0, 0.1, 1.0)
print (sd01)
sd02 = fjcontrib.SoftDrop(0, 0.2, 1.0)
print (sd02)
# jet_def_rc01 = fj.JetDefinition(fj.cambridge_algorithm, 0.1)
# jet_def_rc02 = fj.JetDefinition(fj.cambridge_algorithm, 0.2)
# print (jet_def_rc01)
# print (jet_def_rc02)
# rc = fjcontrib.Recluster(jet_def_rc, True)
jet_def_rc01 = fj.JetDefinition(fj.antikt_algorithm, 0.1)
jet_def_rc02 = fj.JetDefinition(fj.antikt_algorithm, 0.2)
print (jet_def_rc01)
print (jet_def_rc02)
#rc = fjcontrib.Recluster(jet_def_rc, True)
# tw = treewriter.RTreeWriter(name = 'lsjvsx', file_name = 'leadsj_vs_x.root')
tw = treewriter.RTreeWriter(name = 'lsjvsx', file_name = args.output)
tgbkg = None
be = None
if args.enable_background:
# ROOT.gSystem.Load("libpyjetty_TennGen.dylib")
# tgbkg = ROOT.TennGen() # //constructor
# tgbkg.SetCentralityBin(args.cent_bin) # //centraility bin 0 is the 0-5 % most central bin
# tgbkg.SetRandomSeed(args.seed) # //setting the seed
# tgbkg.SetHarmonics(args.harmonics) # // set harmonics flag (0 : v1 - v5) , (1 : v2 - v5) , (2: v3 - v5) , (3: v1 - v4) , (4: v1 - v3) , (5: uniform dN/dphi no harmonics) , (6 : v1 - v2 , v4 - v5) , (7 : v1 - v3 , v5) , (8 : v1 , v3 - v5) , (9 : v1 only) , (10 : v2 only) , (11 : v3 only) , (12 : v4 only) , (13 : v5 only)
# tgbkg.SetEtaRange(args.eta) # //set eta range must be uniform (e.g. |eta| < 0.9, which is ALICE TPC fiducial acceptance)
# tgbkg.PrintOutQAHistos(args.qa) #
# tgbkg.InitializeBackground() #
from pyjetty.mputils import BoltzmannEvent
be = BoltzmannEvent(mean_pt=0.7, multiplicity=2000 * args.eta * 2, max_eta=max_eta, max_pt=100)
print(be)
from pyjetty.mputils import CEventSubtractor, CSubtractorJetByJet
cs = CEventSubtractor(alpha=args.alpha, max_distance=args.dRmax, max_eta=args.eta, bge_rho_grid_size=0.25, max_pt_correct=100)
print(cs)
if args.nev < 100:
args.nev = 100
for i in tqdm.tqdm(range(args.nev)):
if not pythia.next():
continue
# parts = pythiafjext.vectorize(pythia, True, -1, 1, False)
partons = pythiafjext.vectorize_select(pythia, [pythiafjext.kParton], 0, True)
parts = pythiafjext.vectorize_select(pythia, [pythiafjext.kFinal, pythiafjext.kCharged], 0, False)
# parts = pythiafjext.vectorize_select(pythia, [pythiafjext.kFinal], 0, False)
jets = jet_selector(jet_def(parts))
# for j in tqdm.tqdm(jets):
for j in jets:
j_type = match_dR(j, partons, jet_R0 / 2.)
if j_type[0] is None:
if args.nw:
continue
pwarning('Jet with no parton label')
continue
j_sd02 = sd02.result(j)
sd02_info = fjcontrib.get_SD_jet_info(j_sd02)
j_sd01 = sd01.result(j)
sd01_info = fjcontrib.get_SD_jet_info(j_sd01)
rc_sjets01 = fj.sorted_by_pt(jet_def_rc01(j.constituents()))
rc_sjets02 = fj.sorted_by_pt(jet_def_rc02(j.constituents()))
tw.fill_branches( j = j,
lund = [ls for ls in lund_gen.result(j)],
dyg1 = dy_groomer.result(j, 1),
sd01 = j_sd01,
sd01_z = sd01_info.z,
sd01_mu = sd01_info.mu,
sd01_Delta = sd01_info.dR,
sd02 = j_sd02,
sd02_z = sd02_info.z,
sd02_mu = sd02_info.mu,
sd02_Delta = sd02_info.dR,
# breaking compatibility
# sd = j_sd,
# sd_z = sd_info.z,
# sd_mu = sd_info.mu,
# sd_Delta = sd_info.dR,
lsjet01 = rc_sjets01[0],
nsjet01 = len(rc_sjets01),
sjet01 = rc_sjets01,
lsjet02 = rc_sjets02[0],
nsjet02 = len(rc_sjets02),
sjet02 = rc_sjets02,
ppid = j_type[0],
pquark = j_type[1],
pglue = j_type[2], # this is redundancy
pycode = pythia.info.code(),
pysigmagen = pythia.info.sigmaGen(),
pysigmaerr = pythia.info.sigmaErr(),
pyid1 = pythia.info.id1pdf(),
pyid2 = pythia.info.id1pdf(),
pyx1 = pythia.info.x1pdf(),
pyx2 = pythia.info.x2pdf(),
pypdf1 = pythia.info.pdf1(),
pyQfac = pythia.info.QFac(),
pyalphaS = pythia.info.alphaS(),
pypthat = pythia.info.pTHat(),
pymhat = pythia.info.mHat()
)
if be:
bg_parts = be.generate(offset=10000)
full_event = bg_parts
tmp = [full_event.push_back(psj) for psj in j.constituents()]
if cs:
cs_parts = cs.process_event(full_event)
rho = cs.bge_rho.rho()
bg_jets = fj.sorted_by_pt(jet_def(cs_parts))
for bj in bg_jets:
if fjtools.matched_pt(bj, j) > 0.5:
pass
tw.fill_tree()
pythia.stat()
tw.write_and_close()
def main():
parser = argparse.ArgumentParser(description='test duplicate entries',
prog=os.path.basename(__file__))
parser.add_argument('fname', help='input file', default='', type=str)
args = parser.parse_args()
event_tree_name = 'PWGHF_TreeCreator/tree_event_char'
event_tree = uproot.open(args.fname)[event_tree_name]
if not event_tree:
perror('Tree {} not found in file {}'.format(event_tree_name,
args.fname))
return False
pinfo(args.fname)
event_df_orig = event_tree.pandas.df()
len_event_df_orig = len(event_df_orig)
df_event_accepted = event_df_orig.query('is_ev_rej == 0')
df_event_accepted.reset_index(drop=True)
len_event_df_accepted = len(df_event_accepted)
event_df_nodup = df_event_accepted.drop_duplicates()
len_event_df_nodup = len(event_df_nodup)
if len_event_df_accepted != len_event_df_nodup:
perror('original event length:', len_event_df_orig, 'accepted:',
len_event_df_accepted, 'nodup:', len_event_df_nodup)
else:
pindent('original event length:', len_event_df_orig, 'accepted:',
len_event_df_accepted, 'nodup:', len_event_df_nodup)
track_tree_name = 'PWGHF_TreeCreator/tree_Particle'
track_tree = uproot.open(args.fname)[track_tree_name]
if not track_tree:
perror('Tree {} not found in file {}'.format(tree_name, args.fname))
return False
track_df_orig = track_tree.pandas.df()
track_df = pd.merge(track_df_orig,
event_df_nodup,
on=['run_number', 'ev_id'])
len_track_df = len(track_df)
track_df_nodup = track_df.drop_duplicates()
len_track_df_nodup = len(track_df_nodup)
if len_track_df_nodup < len_track_df:
perror('track+event rows:', len_track_df, 'nodup:', len_track_df_nodup)
else:
pindent('track+event rows:', len_track_df, 'nodup:',
len_track_df_nodup)
track_df_grouped = track_df.groupby(['run_number', 'ev_id'])
len_track_df_grouped = len(track_df_grouped)
if len_track_df_grouped <= len_event_df_nodup:
pindent('track+event length grouped:', len_track_df_grouped)
else:
perror('track+event length grouped:', len_track_df_grouped)
# track_df_nodup = track_df_grouped.drop_duplicates()
# print ('track+event length no dup:', len(track_df_nodup))
# from James
# Check if there are duplicated tracks in an event.
duplicate_selection = [
'run_number', 'ev_id', 'ParticlePt', 'ParticleEta', 'ParticlePhi'
]
# if use_ev_id_ext:
# duplicate_selection.append('ev_id_ext')
duplicate_rows_df = track_df.duplicated(duplicate_selection)
for i, row in duplicate_rows_df.iteritems():
if row:
print(i, row)
# for r in duplicate_rows_df:
# print(type(r))
n_duplicates = sum(duplicate_rows_df)
pindent('2nd pass: using duplicate selection ', duplicate_selection)
if n_duplicates > 0:
perror(
'2nd pass: there appear to be {} duplicate particles in the dataframe'
.format(n_duplicates))
perror('this is: {:.2} of all tracks'.format(n_duplicates /
len_track_df))
track_df_nodup = track_df.drop_duplicates(duplicate_selection,
inplace=False)
pwarning('new count rows for particles:', len(track_df_nodup),
'old count:', len_track_df)
else:
pindent('no duplicate particles found')
def process_d0s(self, df):
self.pbar.update(1)
_n_d0s = len(df)
if _n_d0s < 1:
return
# pinfo(df)
if 'ev_id_ext' in list(self.event_df):
_ev_query = "run_number == {} & ev_id == {} & ev_id_ext == {}".format(df['run_number'].values[0], df['ev_id'].values[0], df['ev_id_ext'].values[0])
else:
_ev_query = "run_number == {} & ev_id == {}".format(df['run_number'].values[0], df['ev_id'].values[0])
_df_tracks = self.track_df.query(_ev_query)
_df_tracks.reset_index(drop=True)
djmm = fjtools.DJetMatchMaker()
djmm.set_ch_pt_eta_phi(_df_tracks['ParticlePt'].values, _df_tracks['ParticleEta'].values, _df_tracks['ParticlePhi'].values)
djmm.set_Ds_pt_eta_phi_m(df['pt_cand'].values, df['eta_cand'].values, df['phi_cand'].values, df['inv_mass'].values)
djmm.set_daughters0_pt_eta_phi(df['pt_prong0'].values, df['eta_prong0'].values, df['phi_prong0'].values)
djmm.set_daughters1_pt_eta_phi(df['pt_prong1'].values, df['eta_prong1'].values, df['phi_prong1'].values)
# if _n_d0s > 1:
# print('-- event break - 2D0 event:')
# print('pt_cand: ', df['pt_cand'].values)
# print('eta_cand: ', df['eta_cand'].values)
# print('phi_cand: ', df['phi_cand'].values)
# print('inv_mass: ', df['inv_mass'].values)
# print('pt_prong0: ', df['pt_prong0'].values)
# print('eta_prong0: ', df['eta_prong0'].values)
# print('phi_prong0: ', df['phi_prong0'].values)
# print('pt_prong1: ', df['pt_prong1'].values)
# print('eta_prong1: ', df['eta_prong1'].values)
# print('phi_prong1: ', df['phi_prong1'].values)
self.tw.fill_branches(dpsj = djmm.Ds)
self.tw.fill_tree()
for id0, d0 in enumerate(djmm.Ds):
_parts_and_ds = djmm.match(0.005, id0)
_parts_and_ds.push_back(d0)
ja = jet_analysis.JetAnalysis(jet_R = 0.4, particle_eta_max=0.9, jet_pt_min=2.0)
ja.analyze_event(_parts_and_ds)
if len(ja.jets) < 1:
continue
jets = ja.jets_as_psj_vector()
djets = djmm.filter_D0_jets(jets)
if len(djets) > 0:
j = djets[0]
dcand = djmm.get_Dcand_in_jet(j)
sja = jet_analysis.JetAnalysis(jet_R = 0.1, particle_eta_max=0.9, jet_pt_min=2.0)
sja.analyze_event(j.constituents())
lsj = fj.sorted_by_pt(sja.jets_as_psj_vector())
sj_dcand = djmm.get_Dcand_in_jet(lsj[0])
is_Dsj = 0
if len(sj_dcand) > 0:
# if sj_dcand[0].m() == dcand[0].m() and sj_dcand[0].perp() == dcand[0].perp():
if sj_dcand[0].delta_R(dcand[0]) == 0.0:
is_Dsj = 1
self.twjc.fill_branches(jet = j, dR = j.delta_R(dcand[0]), D = dcand[0], lsj = lsj[0], Dsj = is_Dsj
, a10 = fjext.angularity(j, 1.0, 0.4)
, a15 = fjext.angularity(j, 0.5, 0.4)
, a20 = fjext.angularity(j, 0.0, 0.4)
, a30 = fjext.angularity(j, -1.0, 0.4))
self.twjc.fill_tree()
if len(djets) > 1:
perror("more than one jet per D candidate?")
return True
def load_file(self, path):
if not os.path.exists(path):
pwarning('[w] file', path, 'does not exists.')
return
try:
event_tree = uproot.open(path)[self.event_tree_name]
except:
pwarning('error getting', self.event_tree_name, 'from file:', path)
return False
if not event_tree:
perror('Tree {} not found in file {}'.format(
self.event_tree_name, path))
return False
event_df_orig = event_tree.pandas.df(
['run_number', 'ev_id', 'z_vtx_reco', 'is_ev_rej'])
event_df_orig.reset_index(drop=True)
event_df = event_df_orig.query('is_ev_rej == 0')
event_df.reset_index(drop=True)
# Load gen tree into df
try:
gen_tree = uproot.open(path)[self.gen_tree_name]
except:
pwarning('error getting', self.gen_tree_name, 'from file:', path)
return False
if not gen_tree:
perror('Tree {} not found in file {}'.format(gen_tree_name, path))
return False
gen_df_orig = gen_tree.pandas.df(
['run_number', 'ev_id', 'pt_cand', 'eta_cand', 'cand_type'])
gen_df_orig.sort_values(by=['run_number', 'ev_id'], inplace=True)
df_genruns = gen_df_orig[['run_number', 'ev_id']].copy()
# Load track tree into dataframe
try:
track_tree = uproot.open(path)[self.tree_name]
except:
pwarning('error getting', self.tree_name, 'from file:', path)
return False
if not track_tree:
perror('Tree {} not found in file {}'.format(tree_name, path))
return False
track_df_orig = track_tree.pandas.df([
'run_number', 'ev_id', 'inv_mass', 'pt_cand', 'pt_prong0',
'pt_prong1', 'dca', 'cos_t_star', 'imp_par_prod', 'cos_p',
'cand_type', 'imp_par_prong0', 'imp_par_prong1', 'norm_dl_xy',
'eta_cand', 'nsigTPC_Pi_0', 'nsigTOF_Pi_0', 'nsigTPC_K_1',
'nsigTOF_K_1', 'nsigTPC_Pi_1', 'nsigTOF_Pi_1', 'nsigTPC_K_0',
'nsigTOF_K_0'
])
# Merge event info into track tree
track_df = pd.merge(track_df_orig,
event_df,
on=['run_number', 'ev_id'])
track_df.sort_values(by=['run_number', 'ev_id'], inplace=True)
df_d0runs = track_df[['run_number', 'ev_id']].copy()
df_runs = pd.merge(df_d0runs, df_genruns, on=['run_number', 'ev_id'])
df_runs.drop_duplicates(keep='first', inplace=True)
gen_df_orig = pd.merge(gen_df_orig,
df_runs,
on=['run_number', 'ev_id'])
track_df = pd.merge(track_df, df_runs, on=['run_number', 'ev_id'])
self.d0_gen = gen_df_orig.groupby(['run_number', 'ev_id'])
self.track_df_grouped = track_df.groupby(['run_number', 'ev_id'])
return True
def main():
parser = argparse.ArgumentParser(description='test the TennGen', prog=os.path.basename(__file__))
pyconf.add_standard_pythia_args(parser)
parser.add_argument('--ignore-mycfg', help="ignore some settings hardcoded here", default=False, action='store_true')
parser.add_argument('--cent-bin', help="centraility bin 0 is the 0-5 % most central bin", type=int, default=0)
parser.add_argument('--seed', help="pr gen seed", type=int, default=1111)
parser.add_argument('--harmonics', help="set harmonics flag (0 : v1 - v5) , (1 : v2 - v5) , (2: v3 - v5) , (3: v1 - v4) , (4: v1 - v3) , (5: uniform dN/dphi no harmonics) , (6 : v1 - v2 , v4 - v5) , (7 : v1 - v3 , v5) , (8 : v1 , v3 - v5) , (9 : v1 only) , (10 : v2 only) , (11 : v3 only) , (12 : v4 only) , (13 : v5 only)",
type=int, default=5)
parser.add_argument('--eta', help="set eta range must be uniform (e.g. |eta| < 0.9, which is ALICE TPC fiducial acceptance)",
type=float, default=0.9)
parser.add_argument('--qa', help="PrintOutQAHistos", default=False, action='store_true')
args = parser.parse_args()
args.py_pthatmin = 100
mycfg = ['PhaseSpace:pThatMin = {}'.format(args.py_pthatmin)]
if args.ignore_mycfg:
mycfg = []
pythia = pyconf.create_and_init_pythia_from_args(args, mycfg)
if not pythia:
perror("pythia initialization failed.")
return
tgbkg = ROOT.TennGen() # //constructor
tgbkg.SetCentralityBin(args.cent_bin) # //centraility bin 0 is the 0-5 % most central bin
tgbkg.SetRandomSeed(args.seed) # //setting the seed
tgbkg.SetHarmonics(args.harmonics) # // set harmonics flag (0 : v1 - v5) , (1 : v2 - v5) , (2: v3 - v5) , (3: v1 - v4) , (4: v1 - v3) , (5: uniform dN/dphi no harmonics) , (6 : v1 - v2 , v4 - v5) , (7 : v1 - v3 , v5) , (8 : v1 , v3 - v5) , (9 : v1 only) , (10 : v2 only) , (11 : v3 only) , (12 : v4 only) , (13 : v5 only)
tgbkg.SetEtaRange(args.eta) # //set eta range must be uniform (e.g. |eta| < 0.9, which is ALICE TPC fiducial acceptance)
tgbkg.PrintOutQAHistos(args.qa) #
tgbkg.InitializeBackground() #
jet_R0 = 0.4
jet_def = fj.JetDefinition(fj.antikt_algorithm, jet_R0)
jet_selector_pythia = fj.SelectorPtMin(args.py_pthatmin) & fj.SelectorPtMax(1000.0) &fj.SelectorAbsEtaMax(args.eta - jet_R0)
jet_selector_hybrid = fj.SelectorPtMin(10) & fj.SelectorPtMax(1000.0) &fj.SelectorAbsEtaMax(args.eta - jet_R0)
# jet_selector = fj.SelectorPtMin(40.0) & fj.SelectorPtMax(200.0) &fj.SelectorAbsEtaMax(1)
parts_selector = fj.SelectorAbsEtaMax(args.eta)
print(jet_def)
tw = treewriter.RTreeWriter(name = 'tparts', file_name = 'test_TennGen.root')
if args.nev < 100:
args.nev = 100
pbar = tqdm.tqdm(total = args.nev)
while pbar.n < args.nev:
if not pythia.next():
continue
# get pythia particles
# parts = pythiafjext.vectorize_select(pythia, [pythiafjext.kFinal, pythiafjext.kCharged], 0, False)
_py_fj_parts = parts_selector(pythiafjext.vectorize_select(pythia, [pythiafjext.kFinal], 0, False))
# get jets w/o area determination
# pythia_jets = jet_selector_pythia(jet_def(_py_fj_parts))
# with area determination
jet_area_def = fj.AreaDefinition(fj.active_area, fj.GhostedAreaSpec(args.eta))
cs = fj.ClusterSequenceArea(_py_fj_parts, jet_def, jet_area_def)
pythia_jets = jet_selector_pythia(cs.inclusive_jets())
if len(pythia_jets) < 1:
continue
pbar.update(1)
tw.fill_branches(pyj = pythia_jets)
# now generate bg
bg_tclones = tgbkg.GetBackground()
# tgbkg.GetRandomSeed()
nParticles = bg_tclones.GetEntries();
# pinfo('event', pbar.n, 'number of parts', nParticles)
# _parts = { 'pt' : [], 'eta' : [], 'phi' : [], 'kf' : []}
_parts = [[], [], [], []]
_ = [[_parts[0].append(p[0].Pt()), _parts[1].append(p[0].Eta()), _parts[2].append(p[0].Phi()), _parts[3].append(p[1])] for p in [[tlv_from_tmcparticle(_p), _p.GetKF()] for _p in bg_tclones if _p.GetEnergy()>0]]
_bg_fj_parts = fjext.vectorize_pt_eta_phi(_parts[0], _parts[1], _parts[2], 1000) #bg particles with index > 1000
# add background and pythia
_fj_parts = []
_ = [_fj_parts.append(_p) for _p in _py_fj_parts]
_ = [_fj_parts.append(_p) for _p in _bg_fj_parts]
# stream all particles
_ = [tw.fill_branches(part_pt = _pfj.perp(), part_eta = _pfj.eta(), part_phi = _pfj.phi(), part_idx=_pfj.user_index()) for _pfj in _fj_parts]
# find jets in the hybrid event
# w/o area
# jets = jet_selector_hybrid(jet_def(_fj_parts))
# w/area
cs_hybrid = fj.ClusterSequenceArea(_fj_parts, jet_def, jet_area_def)
jets = jet_selector_hybrid(cs_hybrid.inclusive_jets())
# stream jets from the hybrid event
tw.fill_branches(j = jets)
# estimate the background
bg_rho_range = fj.SelectorAbsEtaMax(args.eta * 1.1)
bg_jet_def = fj.JetDefinition(fj.kt_algorithm, jet_R0)
bg_area_def = fj.AreaDefinition(fj.active_area_explicit_ghosts, fj.GhostedAreaSpec(args.eta))
# bg_area_def = fj.AreaDefinition(fj.active_area, fj.GhostedAreaSpec(args.eta)) #active area defunct for bg estim
bg_estimator = fj.JetMedianBackgroundEstimator(bg_rho_range, bg_jet_def, bg_area_def)
bg_estimator.set_particles(_fj_parts)
if len(_fj_parts) < 0:
perror('no particles in the hybrid event?')
continue
rho = bg_estimator.rho()
sigma = bg_estimator.sigma()
corr_jet_pt = [j.pt() - j.area() * rho for j in jets]
# matches = [j.perp(), matched_jet(j, pythia_jets) for j in jets]
delta_pt = [delta_pt_matched(j, pythia_jets, rho) for j in jets]
tw.fill_branches(j_corr_pt = corr_jet_pt, dpt = delta_pt)
tw.fill_branches(rho = rho, rho_sigma = sigma)
tw.fill_tree()
bg_tclones.Clear()
pbar.close()
tgbkg.CloserFunction()
tw.write_and_close()
def main():
parser = argparse.ArgumentParser(description='pythia8 fastjet on the fly',
prog=os.path.basename(__file__))
pyconf.add_standard_pythia_args(parser)
parser.add_argument('--nw',
help="no warn",
default=True,
action='store_true')
parser.add_argument('--ignore-mycfg',
help="ignore some settings hardcoded here",
default=False,
action='store_true')
parser.add_argument('--enable-background',
help="enable background calc",
default=False,
action='store_true')
parser.add_argument('--output',
help="output file name",
default='leadsj_vs_zloss.root',
type=str)
parser.add_argument('--jetptmin',
help="jet pt minimum",
default=-1,
type=float)
parser.add_argument('--jetptmax',
help="jet pt maximum",
default=1e6,
type=float)
parser.add_argument('--eta', help="jet eta max", default=2.4, type=float)
parser.add_argument(
'--kt',
help="use kT algorithm instead of anti-kT for the subjets",
default=False,
action='store_true')
args = parser.parse_args()
# print the banner first
fj.ClusterSequence.print_banner()
print()
# set up our jet definition and a jet selector
jet_R0 = 0.4
jet_def = fj.JetDefinition(fj.antikt_algorithm, jet_R0)
jet_selector = fj.SelectorPtMin(args.py_pthatmin) & fj.SelectorPtMax(
1000.0) & fj.SelectorAbsEtaMax(args.eta - jet_R0)
mycfg = []
if args.jetptmin > 0:
mycfg = ['PhaseSpace:pThatMin = {}'.format(args.jetptmin)]
jet_selector = fj.SelectorPtMin(args.jetptmin) & fj.SelectorPtMax(
args.jetptmax) & fj.SelectorAbsEtaMax(args.eta - jet_R0)
print(jet_def)
if args.ignore_mycfg:
mycfg = []
pythia = pyconf.create_and_init_pythia_from_args(args, mycfg)
if not pythia:
perror("pythia initialization failed.")
return
nbins = 20
# sjrs = [0.001 + x * 0.04 for x in range(0, nbins)]
sjrs = logbins(0.001, jet_R0, nbins)
print(sjrs)
print('log(1/r) :', [ROOT.TMath.Log(1 / r) for r in sjrs])
sjdefs = dict()
for sjr in sjrs:
if args.kt:
_jet_def = fj.JetDefinition(fj.kt_algorithm, sjr)
else:
_jet_def = fj.JetDefinition(fj.antikt_algorithm, sjr)
sjdefs[sjr] = _jet_def
# tw = treewriter.RTreeWriter(name = 'lsjvsx', file_name = 'leadsj_vs_x.root')
tw = treewriter.RTreeWriter(name='lsjvszloss', file_name=args.output)
tw.fout.cd()
h_zloss_r_q = dict()
h_zloss_r_g = dict()
for sjr in sjrs:
sname = 'h_zloss_glue_{}'.format(sjr)
_h_zloss_r_g = ROOT.TH1F(sname, sname, len(sjrs), 0., 1.)
h_zloss_r_g[sjr] = _h_zloss_r_g
sname = 'h_zloss_quark_{}'.format(sjr)
_h_zloss_r_q = ROOT.TH1F(sname, sname, len(sjrs), 0., 1.)
h_zloss_r_q[sjr] = _h_zloss_r_q
lbins = logbins(ROOT.TMath.Log(1. / jet_R0), ROOT.TMath.Log(1. / sjrs[0]),
nbins)
print('lbins:', lbins)
sname = 'prof_zloss_vs_r_any'
prof_a = ROOT.TProfile(sname, sname, nbins, 0, jet_R0)
prof_a_log = ROOT.TProfile(sname + '_log', sname + '_log', nbins, lbins)
sname = 'prof_zloss_vs_r_glue'
prof_g = ROOT.TProfile(sname, sname, nbins, 0, jet_R0)
prof_g_log = ROOT.TProfile(sname + '_log', sname + '_log', nbins, lbins)
sname = 'prof_zloss_vs_r_quark'
prof_q = ROOT.TProfile(sname, sname, nbins, 0, jet_R0)
prof_q_log = ROOT.TProfile(sname + '_log', sname + '_log', nbins, lbins)
# prof_q_log = ROOT.TProfile(sname+'_log', sname+'_log', nbins, ROOT.TMath.Log(1./jet_R0), ROOT.TMath.Log(1./sjrs[0]))
sname = 'h2_zloss_vs_r_glue'
h2_zloss_r_g = ROOT.TH2F(sname, sname, nbins, 0., jet_R0, len(sjrs), 0.,
1.)
sname = 'h2_zloss_vs_r_quark'
h2_zloss_r_q = ROOT.TH2F(sname, sname, nbins, 0., jet_R0, len(sjrs), 0.,
1.)
# loop
if args.nev < 100:
args.nev = 100
for i in tqdm.tqdm(range(args.nev)):
if not pythia.next():
continue
# parts = pythiafjext.vectorize(pythia, True, -1, 1, False)
partons = pythiafjext.vectorize_select(pythia, [pythiafjext.kParton],
0, True)
parts = pythiafjext.vectorize_select(
pythia, [pythiafjext.kFinal, pythiafjext.kCharged], 0, False)
# parts = pythiafjext.vectorize_select(pythia, [pythiafjext.kFinal], 0, False)
jets = jet_selector(jet_def(parts))
# for j in tqdm.tqdm(jets):
for j in jets:
j_type = match_dR(j, partons, jet_R0 / 2.)
if j_type[0] is None:
if args.nw:
continue
pwarning('Jet with no parton label')
continue
tw.fill_branch("j", j)
for sjr in sjrs:
rc_jets = fj.sorted_by_pt(sjdefs[sjr](j.constituents()))
tw.fill_branch("sjr{}".format(sjr), rc_jets[0])
zloss = 1. - rc_jets[0].perp() / j.perp()
tw.fill_branch("sjr{}_zloss".format(sjr), zloss)
tw.fill_branch("ppid", j_type[0])
tw.fill_branch("pquark", j_type[1])
tw.fill_branch("pglue", j_type[2])
prof_a.Fill(sjr, zloss)
prof_a_log.Fill(ROOT.TMath.Log(1. / sjr), zloss)
if j_type[1]:
h_zloss_r_q[sjr].Fill(zloss)
h2_zloss_r_q.Fill(sjr, zloss)
prof_q.Fill(sjr, zloss)
prof_q_log.Fill(ROOT.TMath.Log(1. / sjr), zloss)
if j_type[2]:
h_zloss_r_g[sjr].Fill(zloss)
h2_zloss_r_g.Fill(sjr, zloss)
prof_g.Fill(sjr, zloss)
prof_g_log.Fill(ROOT.TMath.Log(1. / sjr), zloss)
tw.fill_tree()
pythia.stat()
tw.write_and_close()
def main():
parser = argparse.ArgumentParser(description='test groomers',
prog=os.path.basename(__file__))
parser.add_argument('-o',
'--output-filename',
default="output.root",
type=str)
parser.add_argument('datalistpp',
help='run through a file list',
default='',
type=str)
parser.add_argument('--datalistAA',
help='run through a file list - embedding mode',
default='',
type=str)
parser.add_argument('--jetR', default=0.4, type=float)
parser.add_argument('--alpha', default=0, type=float)
parser.add_argument('--dRmax', default=0.25, type=float)
parser.add_argument('--overwrite',
help="overwrite output",
default=False,
action='store_true')
parser.add_argument('--jetptcut',
help='remove jets below the cut',
default=50.,
type=float)
parser.add_argument('--nev',
help='number of events to run',
default=0,
type=int)
parser.add_argument('--max-eta',
help='max eta for particles',
default=0.9,
type=float)
parser.add_argument('--npart-cut',
help='npart cut on centrality low,high hint:' +
npart_cents,
default='325,450',
type=str)
args = parser.parse_args()
try:
npart_min = int(args.npart_cut.split(',')[0])
npart_max = int(args.npart_cut.split(',')[1])
except:
perror(
'unable to parse npart centrality selection - two integer numbers with a coma in-between needed - specified:',
args.npart_cut)
return 1
# initialize constituent subtractor
cs = None
if args.dRmax > 0:
cs = CEventSubtractor(alpha=args.alpha,
max_distance=args.dRmax,
max_eta=args.max_eta,
bge_rho_grid_size=0.25,
max_pt_correct=100)
pp_data = DataIO(name='Sim Pythia Detector level',
file_list=args.datalistpp,
random_file_order=False,
tree_name='tree_Particle_gen')
ja_pp = JetAnalysis(jet_R=args.jetR,
jet_algorithm=fj.antikt_algorithm,
jet_pt_min=50.,
particle_eta_max=args.max_eta)
if args.datalistAA:
aa_data = DataBackgroundIO(name='PbPb',
file_list=args.datalistAA,
tree_name='tree_Particle_gen')
ja_emb = JetAnalysis(jet_R=args.jetR,
jet_algorithm=fj.antikt_algorithm,
jet_pt_min=50.,
particle_eta_max=args.max_eta)
ja_aa = JetAnalysis(jet_R=args.jetR,
jet_algorithm=fj.antikt_algorithm,
jet_pt_min=50.,
particle_eta_max=args.max_eta)
dndeta_selector = fj.SelectorAbsEtaMax(1.)
# tg = thg.ThermalGenerator()
print(cs)
# print the banner first
fj.ClusterSequence.print_banner()
print()
gout = GroomerOutput(args.output_filename,
enable_aa_trees=bool(args.datalistAA))
delta_t = 0
start_t = time.time()
iev = 1
while pp_data.load_event(offset=0):
iev = iev + 1
if args.nev > 0:
if iev > args.nev:
iev = iev - 1
break
if iev % 1000 == 0:
delta_t = time.time() - start_t
pinfo('processing event', iev, ' - ev/sec =', iev / delta_t,
'elapsed =', delta_t)
# find jets on detector level
if len(pp_data.particles) < 1:
pwarning(iev, 'pp event skipped N parts', len(pp_data.particles))
continue
ja_pp.analyze_event(pp_data.particles)
if len(ja_pp.jets) < 1:
continue
# pinfo('n particles', len(pp_data.particles))
dndeta0 = dndeta_selector(pp_data.particles)
[
gout.fill_branches(j, syst=0, dndeta=len(dndeta0) / 2.)
for j in ja_pp.jets
]
# pinfo('n jets', len(ja_pp.jets))
if args.datalistAA:
while True:
aa_loaded = aa_data.load_event(offset=10000)
if aa_data.event.npart < npart_min or aa_data.event.npart >= npart_max:
continue
else:
if len(aa_data.particles) < 1:
pwarning(iev, 'AA event skipped N parts',
len(aa_data.particles))
continue
else:
break
if aa_loaded:
ja_aa.analyze_event(aa_data.particles)
dndeta1 = dndeta_selector(aa_data.particles)
if len(ja_aa.jets) > 0:
[
gout.fill_branches(j, syst=1, dndeta=len(dndeta1) / 2.)
for j in ja_aa.jets
]
else:
# pwarning('no jets in AA event?', len(ja_aa.jets), 'while dndeta=', len(dndeta1)/2.)
pass
emb_event = fj.vectorPJ()
[emb_event.push_back(p) for p in pp_data.particles]
[emb_event.push_back(p) for p in aa_data.particles]
rho = 0
if cs:
cs_parts = cs.process_event(emb_event)
rho = cs.bge_rho.rho()
ja_emb.analyze_event(cs_parts)
else:
ja_emb.analyze_event(emb_event)
# matches = [[jpp, jemb] for jpp in ja_pp.jets for jemb in ja_emb.jets if fjtools.matched_pt(jemb, jpp) > 0.5]
# for mj in matches:
# gout.fill_branches(mj[0], syst=2, dndeta=len(dndeta1)/2., rho=rho)
# gout.fill_branches(mj[1], syst=3)
[
gout.fill_branches_prong_matching(j_pp,
j_emb,
dndeta=len(dndeta1) / 2.,
rho=rho)
for j_pp in ja_pp.jets for j_emb in ja_emb.jets
]
delta_t = time.time() - start_t
pinfo('processed events', iev, ' - ev/sec =', iev / delta_t, 'elapsed =',
delta_t)
gout.write()
def main():
parser = argparse.ArgumentParser(description='pythia8 fastjet on the fly',
prog=os.path.basename(__file__))
pyconf.add_standard_pythia_args(parser)
parser.add_argument('--ignore-mycfg',
help="ignore some settings hardcoded here",
default=False,
action='store_true')
parser.add_argument('--output',
help="output file name",
default="test_hjet_parton.root",
type=str)
parser.add_argument('--no-tt',
help="do not require TT to accept the event",
default=False,
action='store_true')
parser.add_argument('--charged',
help="analyze only the charged particles of the FS",
default=False,
action='store_true')
parser.add_argument('--max-jet-pt',
help="maximum jet pT to consider",
type=float,
default=100.)
args = parser.parse_args()
# print the banner first
fj.ClusterSequence.print_banner()
print()
# set up our jet definition and a jet selector
hadron_eta_max = 2.0
jet_R0 = 0.4
jet_def = fj.JetDefinition(fj.antikt_algorithm, jet_R0)
jet_selector = fj.SelectorPtMin(10.0) & fj.SelectorPtMax(
args.max_jet_pt) & fj.SelectorAbsEtaMax(hadron_eta_max - jet_R0)
# jet_selector = fj.SelectorPtMin(40.0) & fj.SelectorPtMax(200.0) &fj.SelectorAbsEtaMax(hadron_eta_max - jet_R0)
hTT6_selector = fj.SelectorPtMin(6) & fj.SelectorPtMax(
7) & fj.SelectorAbsEtaMax(hadron_eta_max)
hTT12_selector = fj.SelectorPtMin(12) & fj.SelectorPtMax(
50) & fj.SelectorAbsEtaMax(hadron_eta_max)
hTT20_selector = fj.SelectorPtMin(20) & fj.SelectorPtMax(
50) & fj.SelectorAbsEtaMax(hadron_eta_max)
pythia_fs_part_selection = [pythiafjext.kFinal]
if args.charged is True:
pwarning('running with charged particles in the final state')
pythia_fs_part_selection.append(pythiafjext.kCharged)
print(jet_def)
all_jets = []
# mycfg = ['PhaseSpace:pThatMin = 80']
# mycfg = ['PhaseSpace:pThatMin = 6']
# mycfg = ['PhaseSpace:pThatMin = 12']
# mycfg = ['PhaseSpace:pThatMin = 40']
mycfg = []
if args.ignore_mycfg:
mycfg = []
pythia = pyconf.create_and_init_pythia_from_args(args, mycfg)
if not pythia:
perror("pythia initialization failed.")
return
jet_def_lund = fj.JetDefinition(fj.cambridge_algorithm, 1.0)
lund_gen = fjcontrib.LundGenerator(jet_def_lund)
print(lund_gen.description())
dy_groomer = fjcontrib.DynamicalGroomer(jet_def_lund)
print(dy_groomer.description())
# sd = fjcontrib.SoftDrop(0, 0.1, 1.0)
sd = fjcontrib.SoftDrop(0, 0.2, 1.0)
print(sd)
# jet_def_rc01 = fj.JetDefinition(fj.cambridge_algorithm, 0.1)
# jet_def_rc02 = fj.JetDefinition(fj.cambridge_algorithm, 0.2)
# print (jet_def_rc01)
# print (jet_def_rc02)
# rc = fjcontrib.Recluster(jet_def_rc, True)
jet_def_rc01 = fj.JetDefinition(fj.antikt_algorithm, 0.1)
jet_def_rc02 = fj.JetDefinition(fj.antikt_algorithm, 0.2)
print(jet_def_rc01)
print(jet_def_rc02)
#rc = fjcontrib.Recluster(jet_def_rc, True)
# tw = treewriter.RTreeWriter(name = 'lsjvsx', file_name = 'leadsj_vs_x.root')
tw = treewriter.RTreeWriter(name='lsjvsx', file_name=args.output)
zero_psj = fj.PseudoJet(0, 0, 10, 10)
if args.nev < 100:
args.nev = 100
t = tqdm.tqdm(total=args.nev)
while t.n < args.nev:
if not pythia.next():
continue
# information about the leading process
# print(pythia.info.code(), pythia.info.nameProc(pythia.info.code()))
# continue
# parts = pythiafjext.vectorize(pythia, True, -1, 1, False)
partons = pythiafjext.vectorize_select(pythia, [pythiafjext.kParton],
0, True)
# parts = pythiafjext.vectorize_select(pythia, [pythiafjext.kFinal, pythiafjext.kCharged], 0, False)
# parts = pythiafjext.vectorize_select(pythia, [pythiafjext.kFinal], 0, False)
parts = pythiafjext.vectorize_select(pythia, pythia_fs_part_selection,
0, False)
hTT6 = zero_psj
hTT6s = fj.sorted_by_pt(hTT6_selector(parts))
if len(hTT6s) > 0:
hTT6 = hTT6s[0]
hTT12 = zero_psj
hTT12s = fj.sorted_by_pt(hTT12_selector(parts))
if len(hTT12s) > 0:
hTT12 = hTT12s[0]
hTT20 = zero_psj
hTT20s = fj.sorted_by_pt(hTT20_selector(parts))
if len(hTT20s) > 0:
hTT20 = hTT20s[0]
if args.no_tt is False:
if hTT12.perp() < 1 and hTT6.perp() < 1 and hTT20.perp() < 1:
continue
jets = jet_selector(jet_def(parts))
# for j in tqdm.tqdm(jets):
for j in jets:
t.update(1)
j_type = match_dR(j, partons, jet_R0 / 2.)
if j_type[0] is None:
continue
j_sd = sd.result(j)
sd_info = fjcontrib.get_SD_jet_info(j_sd)
rc_sjets01 = fj.sorted_by_pt(jet_def_rc01(j.constituents()))
rc_sjets02 = fj.sorted_by_pt(jet_def_rc02(j.constituents()))
tw.fill_branches(
j=j,
mult=len(parts),
lund=[ls for ls in lund_gen.result(j)],
dyg1=dy_groomer.result(j, 1),
sd=j_sd,
sd_z=sd_info.z,
sd_mu=sd_info.mu,
sd_Delta=sd_info.dR,
lsjet01=rc_sjets01[0],
nsjet01=len(rc_sjets01),
sjet01=rc_sjets01,
lsjet02=rc_sjets02[0],
nsjet02=len(rc_sjets02),
sjet02=rc_sjets02,
hTT6=hTT6,
hTT12=hTT12,
hTT20=hTT20,
dphi6=j.delta_phi_to(hTT6),
dphi12=j.delta_phi_to(hTT12),
dphi20=j.delta_phi_to(hTT20),
ppid=j_type[0],
pquark=j_type[1],
pglue=j_type[2], # this is redundancy
pycode=pythia.info.code(),
pysigmagen=pythia.info.sigmaGen(),
pysigmaerr=pythia.info.sigmaErr(),
pyid1=pythia.info.id1pdf(),
pyid2=pythia.info.id1pdf(),
pyx1=pythia.info.x1pdf(),
pyx2=pythia.info.x2pdf(),
pypdf1=pythia.info.pdf1(),
pyQfac=pythia.info.QFac(),
pyalphaS=pythia.info.alphaS(),
pypthat=pythia.info.pTHat(),
pymhat=pythia.info.mHat())
tw.fill_tree()
t.close()
pythia.stat()
tw.write_and_close()
def run(self):
# need to change this for data to drive...
delta_t = 0
start_t = time.time()
iev = 1
# while self.det_sim.load_event() and self.part_sim.load_event():
while self.det_sim.load_event():
iev = iev + 1
if self.nev > 0:
if iev > self.nev:
iev = iev - 1
break
if iev % 1000 == 0:
delta_t = time.time() - start_t
pinfo('processing event', iev, ' - ev/sec =', iev / delta_t,
'elapsed =', delta_t)
# find jets on detector level
if len(self.det_sim.particles) < 1:
pwarning(iev, 'event skipped N detector parts',
len(self.det_sim.particles))
continue
self.ja_det.analyze_event(self.det_sim.particles)
_jets_det = self.ja_det.jets
# _x = [pdebug(' -d ', j) for j in _jets_det]
if len(_jets_det) < 1:
continue
_too_high_pt = [
p.pt() for j in _jets_det for p in j.constituents()
if p.pt() > 100.
]
if len(_too_high_pt) > 0:
pwarning(iev, 'a likely fake high pT particle(s)',
_too_high_pt, '- skipping whole event')
continue
_output_fname = os.path.expanduser(
os.path.expandvars(self.det_sim.file_io.file_input))
_output_fname = _output_fname.replace("/", "_")
self.output.initialize_output(_output_fname)
self.output.fill_det_level(iev, _jets_det)
# load the corresponding event on particle level
self.part_sim.open_afile(afile=self.det_sim.file_io.file_input)
if not self.part_sim.load_event_with_loc(
self.det_sim.event.run_number, self.det_sim.event.ev_id,
0):
perror('unable to load partL event run#:',
self.det_sim.event.run_number, 'ev_id:',
self.det_sim.event.ev_id)
continue
if self.det_sim.event.run_number != self.part_sim.event.run_number:
perror('run# missmatch detL:', self.det_sim.event.run_number,
'partL:', self.part_sim.event.run_number)
continue
if self.det_sim.event.ev_id != self.part_sim.event.ev_id:
perror('ev_id# missmatch detL:', self.det_sim.event.ev_id,
'partL:', self.part_sim.event.ev_id)
continue
# find jets on particle level
if len(self.part_sim.particles) < 1:
pwarning(iev, 'event skipped N particle parts',
len(self.part_sim.particles))
continue
self.ja_part.analyze_event(self.part_sim.particles)
_jets_part = self.ja_part.jets
# _x = [pdebug(' -p ', j) for j in _jets_part]
if len(_jets_part) < 1:
continue
# match in pp simulations
_det_part_matches = []
_n_matches = 0
_part_psjv = self.ja_part.jets_as_psj_vector()
for j_det in _jets_det:
_mactches_pp = fjtools.matched_Reta(j_det, _part_psjv,
0.6 * self.jetR)
#_mactches_pp = fjtools.matched_Ry(j_det, _part_psjv, 0.6 * self.jetR)
_n_matches = _n_matches + len(_mactches_pp)
if len(_mactches_pp) > 1:
pwarning('event:', iev, 'jet pt=', j_det.pt(),
'more than one match in pp jets',
[i for i in _mactches_pp])
if len(_mactches_pp) == 1:
j_part = _part_psjv[_mactches_pp[0]]
# pinfo('j_det', j_det, 'j_part', j_part)
_det_part_matches.append([j_det, j_part])
self.output.fill_pp_pairs(iev, [j_det, j_part])
if _n_matches < 1:
if _n_matches < 1:
pwarning('event:', iev,
'- no matched jets in simulation!?',
len(_det_part_matches))
# here embedding to PbPb data
_offset = 10000
while _offset < len(self.det_sim.particles):
_offset = _offset + 1000
pwarning('increasing bg index offset to', _offset)
_PbPb_loaded = 0
while _PbPb_loaded == 0:
if not self.dataPbPb.load_event(offset=_offset):
perror('unable to load next PbPb event')
_PbPb_loaded = -1
else:
_hybrid_event = self.dataPbPb.particles
_nparts_hybrid_no_emb = len(_hybrid_event)
if _nparts_hybrid_no_emb < 1:
pwarning(
'hybrid event with no particles! trying another one'
)
_PbPb_loaded = 0
else:
_PbPb_loaded = 1
if _PbPb_loaded < 0:
perror(
'unable to load PbPb event - permanent - bailing out here.'
)
break
_tmp = [_hybrid_event.push_back(p) for p in self.det_sim.particles]
if self.cs:
cs_parts = self.cs.process_event(_hybrid_event)
rho = self.cs.bge_rho.rho()
self.ja_hybrid.analyze_event(cs_parts)
else:
self.ja_hybrid.analyze_event(_hybrid_event)
_hybrid_matches = []
_hybrid_psjv = self.ja_hybrid.jets_as_psj_vector()
for m in _det_part_matches:
j_det = m[0]
j_part = m[1]
_mactches_hybrid = fjtools.matched_Reta(
j_det, _hybrid_psjv, 0.6 * self.jetR)
if len(_mactches_hybrid) > 1:
pwarning('event:', iev, 'jet pt=', j_det.pt(),
'more than one match in hybrid jets',
[i for i in _mactches_hybrid])
if len(_mactches_hybrid) == 1:
# m.append(_hybrid_psjv[_mactches_hybrid[0]])
j_hybr = _hybrid_psjv[_mactches_hybrid[0]]
# pdebug('L302', 'j_det', j_det, 'j_part', j_part, 'j_hybr', j_hybr)
_hybrid_matches.append([j_det, j_part, j_hybr])
self.output.fill_emb_3(iev, [j_det, j_part, j_hybr])
_n_matches_hybrid = len(_hybrid_matches)
if _n_matches_hybrid < 1:
if _n_matches_hybrid < 1:
pwarning('event:', iev, '- no matched jets in embedding!?',
_n_matches_hybrid)
delta_t = time.time() - start_t
pinfo('processed events', iev, ' - ev/sec =', iev / delta_t,
'elapsed =', delta_t)
self.output.close()
def main():
parser = argparse.ArgumentParser(description='test groomers',
prog=os.path.basename(__file__))
parser.add_argument('-o',
'--output-filename',
default="centrality_output.root",
type=str)
parser.add_argument('datalist',
help='run through a file list',
default='',
type=str)
parser.add_argument('--overwrite',
help="overwrite output",
default=False,
action='store_true')
parser.add_argument('--nev',
help='number of events to run',
default=0,
type=int)
parser.add_argument('--max-eta', help='max eta for particles', default=0.9)
parser.add_argument('--thermal',
help='enable thermal generator',
action='store_true',
default=False)
parser.add_argument('--thermal-default',
help='enable thermal generator',
action='store_true',
default=False)
parser.add_argument('--particles',
help='stream particles',
action='store_true',
default=False)
parser.add_argument('--npart-cut',
help='npart cut on centrality low,high hint:' +
npart_cents,
default='325,450',
type=str)
parser.add_argument('--nch-cut',
help='nch cut on centrality low,high hint:' +
nch_cents,
default='18467,50000',
type=str)
args = parser.parse_args()
try:
npart_min = int(args.npart_cut.split(',')[0])
npart_max = int(args.npart_cut.split(',')[1])
except:
perror(
'unable to parse npart centrality selection - two integer numbers with a coma in-between needed - specified:',
args.npart_cut)
return 1
try:
nch_min = int(args.nch_cut.split(',')[0])
nch_max = int(args.nch_cut.split(',')[1])
except:
perror(
'unable to parse nch centrality selection - two integer numbers with a coma in-between needed - specified:',
args.nch_cut)
return 1
outf = ROOT.TFile(args.output_filename, 'recreate')
outf.cd()
t = ROOT.TTree('t', 't')
tw = RTreeWriter(tree=t)
hpt_antyr = ROOT.TH1F('hpt_antyr', 'hpt_antyr', 100, 0, 100)
hpt_antyr_c = ROOT.TH1F('hpt_antyr_c', 'hpt_antyr_c', 100, 0, 100)
hpt_therm = ROOT.TH1F('hpt_therm', 'hpt_therm', 100, 0, 100)
hpt_therm_c = ROOT.TH1F('hpt_therm_c', 'hpt_therm_c', 100, 0, 100)
data = DataIO(name='Sim Pythia Detector level',
file_list=args.datalist,
random_file_order=False,
tree_name='tree_Particle_gen')
dndeta_selector = fj.SelectorAbsEtaMax(abs(
args.max_eta)) & fj.SelectorPtMin(0.15)
tg_default = None
if args.thermal_default:
tg_default = thg.ThermalGenerator()
print(tg_default)
tg_central = None
if args.thermal:
tg_central = thg.ThermalGenerator(beta=0.5, N_avg=3000, sigma_N=500)
print(tg_central)
delta_t = 0
start_t = time.time()
iev = 1
while data.load_event(offset=0):
iev = iev + 1
if args.nev > 0:
if iev > args.nev:
iev = iev - 1
break
if iev % 1000 == 0:
delta_t = time.time() - start_t
pinfo('processing event', iev, ' - ev/sec =', iev / delta_t,
'elapsed =', delta_t)
# find jets on detector level
if len(data.particles) < 1:
pwarning(iev, 'pp event skipped N parts', len(data.particles))
continue
# print(data.event)
dndeta0_parts = dndeta_selector(data.particles)
dndeta0 = len(dndeta0_parts) / (abs(args.max_eta * 2.))
[hpt_antyr.Fill(p.perp()) for p in dndeta0_parts]
if args.particles:
tw.fill_branches(dndeta=dndeta0, p=data.particles)
else:
tw.fill_branches(dndeta=dndeta0)
tw.fill_branches_attribs(
data.event, ['sigma', 'npart', 'nch', 'nchfwd', 'nchselect'],
prefix='antyr_')
if data.event.npart < npart_min or data.event.npart >= npart_max:
tw.fill_branches(cent10npart=0)
else:
tw.fill_branches(cent10npart=1)
[hpt_antyr_c.Fill(p.perp()) for p in dndeta0_parts]
if data.event.nch < nch_min or data.event.nch >= nch_max:
tw.fill_branches(cent10nch=0)
else:
tw.fill_branches(cent10nch=1)
if tg_default:
thg_particles = tg_default.load_event()
dndetathg_default = dndeta_selector(thg_particles)
if args.particles:
tw.fill_branches(dndeta_thg_0=len(dndetathg_default) /
(abs(args.max_eta * 2.)),
p_thg_0=thg_particles)
else:
tw.fill_branches(dndeta_thg_0=len(dndetathg_default) /
(abs(args.max_eta * 2.)))
if tg_central:
thg_parts_central = tg_central.load_event()
dndetathg_central = dndeta_selector(thg_parts_central)
[hpt_therm_c.Fill(p.perp()) for p in dndetathg_central]
if args.particles:
tw.fill_branches(dndeta_thg_c=len(dndetathg_central) /
(abs(args.max_eta * 2.)),
p_thg_c=thg_parts_central)
else:
tw.fill_branches(dndeta_thg_c=len(dndetathg_central) /
(abs(args.max_eta * 2.)))
tw.fill_tree()
delta_t = time.time() - start_t
pinfo('processed events', iev, ' - ev/sec =', iev / delta_t, 'elapsed =',
delta_t)
outf.Write()
outf.Close()
