class HJetTree(MPBase):
def __init__(self, **kwargs):
self.configure_from_args( fout=None,
output='hjet_tree.root',
trigger_range = [6, 7],
jet_ana = None)
super(HJetTree, self).__init__(**kwargs)
if self.fout is None:
self.fout = r.TFile(self.output, 'RECREATE')
self.fout.cd()
self.tree_name='tjet_0{}_{}_{}'.format(int(10*self.jet_ana.jet_R), int(self.trigger_range[0]), int(self.trigger_range[1]))
self.tree_writer = RTreeWriter(tree_name=self.tree_name, fout=self.fout)
def analyze_event(self, jet_parts):
self.trigger_particle = None
t_selector = fj.SelectorPtRange(self.trigger_range[0], self.trigger_range[1])
t_candidates = t_selector(jet_parts)
if len(t_candidates) < 1:
return False
self.trigger_particle = random.choice(t_candidates)
return True
def fill_branches(self, ev_id, weight):
if self.jet_ana is None:
print('[e] no jet ana in hjettree...')
return False
if self.trigger_particle:
if self.trigger_particle:
self.dphis = [self.trigger_particle.delta_phi_to(j) for j in self.jet_ana.jets]
self.detas = [self.trigger_particle.eta() - j.eta() for j in self.jet_ana.jets]
else:
self.dphis = []
self.detas = []
self.tree_writer.fill_branch('ev_id', ev_id)
self.tree_writer.fill_branch('t_w', weight)
self.tree_writer.fill_branch('t_pt', self.trigger_particle.pt())
self.tree_writer.fill_branch('t_phi', self.trigger_particle.phi())
self.tree_writer.fill_branch('t_eta', self.trigger_particle.eta())
self.tree_writer.fill_branch('dphi', self.dphis)
self.tree_writer.fill_branch('deta', self.detas)
self.jet_ana.fill_branches(self.tree_writer, ev_id = None, weight = None)
def fill_tree(self):
self.tree_writer.fill_tree()
def write_and_close_file(self):
self.fout.Write()
self.fout.Close()
class EmbeddingOutput(MPBase):
def __init__(self, **kwargs):
self.configure_from_args(args=None)
super(EmbeddingOutput, self).__init__(**kwargs)
self.copy_attributes(self.args)
self.outf = None
self.sd = fjcontrib.SoftDrop(0, self.sd_zcut, self.jetR)
def initialize_output(self, output_name=None):
if output_name:
if self.output_filename != output_name:
self.output_filename = output_name
if self.outf:
if self.outf.GetName() != self.output_filename:
pinfo('closing output file', self.outf.GetName())
self.outf.Write()
self.outf.Close()
self.outf = None
else:
return True
if self.outf is None:
self.outf = ROOT.TFile(self.output_filename, 'recreate')
self.outf.cd()
self.tdet = ROOT.TTree('tdet', 'tdet')
self.twdet = RTreeWriter(
tree=self.tdet,
name='Output Tree detector level pp simulation')
self.tpp = ROOT.TTree('tpp', 'tpp')
self.twpp = RTreeWriter(tree=self.tpp,
name='Output Tree pp simulation')
self.th = ROOT.TTree('th', 'th')
self.twh = RTreeWriter(
tree=self.th,
name='Output Tree pp simulation embedded into PbPb')
pinfo('new output file', self.outf.GetName())
def close(self):
if self.outf:
pinfo('closing output file', self.outf.GetName())
self.outf.Write()
self.outf.Close()
self.outf = None
def _fill_det_level(self, jet):
self.twdet.fill_branch('pt_det', jet.pt())
self.twdet.fill_branch('pt_phi', jet.phi())
self.twdet.fill_branch('pt_eta', jet.eta())
def fill_det_level(self, iev=-1, jets=[]):
if len(jets) > 0:
self.twdet.fill_branch('iev', iev)
_tmp = [self._fill_det_level(j) for j in jets]
self.twdet.fill_tree()
def fill_pp_pairs(self, iev, jm):
self.twpp.fill_branch('iev', iev)
self.twpp.fill_branch('det', jm[0])
self.twpp.fill_branch('part', jm[1])
self.twpp.fill_branch('dpt', jm[0].pt() - jm[1].pt())
self.twpp.fill_branch('dR', jm[0].delta_R(jm[1]))
self.twpp.fill_tree()
def fill_emb_3(self, iev, jm):
# pdebug('@fill: jm[0]', jm[0], 'jm[1]', jm[1], 'jm[2]', jm[2])
self.twh.fill_branch('iev', iev)
self.twh.fill_branch('det', jm[0])
self.twh.fill_branch('part', jm[1])
self.twh.fill_branch('hybr', jm[2])
self.twh.fill_branch('dpt_pp', jm[0].pt() - jm[1].pt())
self.twh.fill_branch('dpt_emb', jm[2].pt() - jm[0].pt())
self.twh.fill_branch('dR_pp', jm[0].delta_R(jm[1]))
self.twh.fill_branch('dR_emb', jm[0].delta_R(jm[2]))
sd0 = self.sd.result(jm[0])
self.twh.fill_branch('sd_det', sd0)
sd0_pe1 = fj.PseudoJet()
sd0_pe2 = fj.PseudoJet()
sd0_has_parents = sd0.has_parents(sd0_pe1, sd0_pe2)
self.twh.fill_branch('sd_det_p1', sd0_pe1)
self.twh.fill_branch('sd_det_p2', sd0_pe2)
sdi0 = fjcontrib.get_SD_jet_info(sd0)
self.twh.fill_branch('sd_det_zg', sdi0.z)
self.twh.fill_branch('sd_det_Rg', sdi0.dR)
sd1 = self.sd.result(jm[1])
self.twh.fill_branch('sd_part', sd1)
sd1_pe1 = fj.PseudoJet()
sd1_pe2 = fj.PseudoJet()
sd1_has_parents = sd1.has_parents(sd1_pe1, sd1_pe2)
self.twh.fill_branch('sd_part_p1', sd1_pe1)
self.twh.fill_branch('sd_part_p2', sd1_pe2)
sdi1 = fjcontrib.get_SD_jet_info(sd1)
self.twh.fill_branch('sd_part_zg', sdi1.z)
self.twh.fill_branch('sd_part_Rg', sdi1.dR)
sd2 = self.sd.result(jm[2])
self.twh.fill_branch('sd_emb', sd2)
sd2_pe1 = fj.PseudoJet()
sd2_pe2 = fj.PseudoJet()
sd2_has_parents = sd2.has_parents(sd2_pe1, sd2_pe2)
self.twh.fill_branch('sd_emb_p1', sd2_pe1)
self.twh.fill_branch('sd_emb_p2', sd2_pe2)
sdi2 = fjcontrib.get_SD_jet_info(sd2)
self.twh.fill_branch('sd_emb_zg', sdi2.z)
self.twh.fill_branch('sd_emb_Rg', sdi2.dR)
m02_1 = -1
m02_2 = -1
if sd0_has_parents and sd2_has_parents:
m02_1 = fjtools.matched_pt(sd2_pe1, sd0_pe1)
m02_2 = fjtools.matched_pt(sd2_pe2, sd0_pe2)
self.twh.fill_branch('sd_det_emb_mpt1', m02_1)
self.twh.fill_branch('sd_det_emb_mpt2', m02_2)
self.twh.fill_branch('sd_det_split', sd0_has_parents)
self.twh.fill_branch('sd_part_split', sd1_has_parents)
self.twh.fill_branch('sd_emb_split', sd2_has_parents)
self.twh.fill_tree()
def main():
parser = argparse.ArgumentParser(description='pythia8 fastjet on the fly',
prog=os.path.basename(__file__))
pyconf.add_standard_pythia_args(parser)
_default_output_filename = os.path.basename(__file__).replace(
".py", "") + "_output.root"
parser.add_argument('--output', default=_default_output_filename, type=str)
parser.add_argument('--debug', default=0, type=int)
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.6
jet_def = fj.JetDefinition(fj.antikt_algorithm, jet_R0)
jet_selector = fj.SelectorPtMin(2.0) & fj.SelectorAbsEtaMax(2)
print(jet_def)
jet_def_lund = fj.JetDefinition(fj.cambridge_algorithm, jet_R0)
lund_gen = fjcontrib.LundGenerator(jet_def_lund)
print(jet_def_lund)
print(lund_gen)
outf = ROOT.TFile(args.output, 'recreate')
outf.cd()
t = ROOT.TTree('t', 't')
tw = RTreeWriter(tree=t)
# mycfg = ['PhaseSpace:pThatMin = 100']
mycfg = []
pythia = pyconf.create_and_init_pythia_from_args(args, mycfg)
if args.nev < 100:
args.nev = 100
for i in tqdm.tqdm(range(args.nev)):
if not pythia.next():
continue
if args.debug:
pwarning('-- event', i)
# parts = pythiafjext.vectorize(pythia, True, -1, 1, False)
parts = pythiafjext.vectorize_select(pythia, [pythiafjext.kFinal], 0,
True)
if args.debug > 5:
parts = pythiafjext.vectorize_select(pythia, [pythiafjext.kHadron],
0, True)
if args.debug > 10:
parts = pythiafjext.vectorize_select(pythia, [pythiafjext.kAny], 0,
True)
if args.debug > 0:
for p in parts:
pypart = pythiafjext.getPythia8Particle(p)
if pypart.name()[:2] == 'D0':
pinfo(pypart.name(), pypart.id(), pypart.status(),
'final =?', pypart.isFinal())
jets = jet_selector(jet_def(parts))
for j in jets:
isD0_lead = False
lead_part = fj.sorted_by_E(j.constituents())[0]
pypart = pythiafjext.getPythia8Particle(lead_part)
if args.debug:
pinfo('leading id is', pypart.id(), pypart.name(), 'jet', j)
if abs(pypart.id()) == 421:
# pinfo('leading D0')
isD0_lead = True
l = lund_gen.result(j)
if len(l) > 0:
tw.fill_branch('Epair', [s.pair().e() for s in l])
tw.fill_branch('z', [s.z() for s in l])
tw.fill_branch('kt', [s.kt() for s in l])
tw.fill_branch('delta', [s.Delta() for s in l])
tw.fill_branch('D0lead', isD0_lead)
tw.fill_branch('lead_id', pypart.id())
tw.fill_tree()
else:
if args.debug:
pwarning("len of a lund is less than 1?", len(l), l)
pythia.stat()
outf.Write()
outf.Close()
print('[i] written', outf.GetName())
def generate():
parser = argparse.ArgumentParser(description='pythia8 fastjet on the fly',
prog=os.path.basename(__file__))
parser.add_argument(
'-d',
'--output-dir',
help='output directory name - use with default file name',
default='.',
type=str)
parser.add_argument('-o',
'--output',
help='output file name',
default='',
type=str)
parser.add_argument('--overwrite', default=False, action='store_true')
parser.add_argument('--jet-R',
help='jet finder R',
default=0.4,
type=float)
parser.add_argument('--charged', default=False, action='store_true')
parser.add_argument('--runid', default=0, type=int)
parser.add_argument('--tranges',
default='6-7',
help='hadron trigger ranges min-max,min1-max1,...',
type=str)
group = parser.add_mutually_exclusive_group()
group.add_argument('--inel', default=False, action='store_true')
group.add_argument('--hard', default=False, action='store_true')
pyconf.add_standard_pythia_args(parser)
args = parser.parse_args()
# mycfg = ['PhaseSpace:pThatMin = 6']
mycfg = []
if '.root' not in args.output:
rs = '0{}'.format(int(args.jet_R * 10))
if args.jet_R >= 1:
rs = '{}'.format(int(args.jet_R))
args.output = '{}/h_jet_R{}'.format(args.output_dir, rs)
if args.charged:
args.output = '{}/h_jet_ch_R{}'.format(args.output_dir, rs)
if args.tranges:
args.output += '_tranges_{}'.format(args.tranges.replace(',', "_"))
if args.runid >= 0:
args.output += '_runid_{}'.format(args.runid)
args.py_seed = 1000 + args.runid
if args.py_noMPI:
args.output += '_noMPI'
if args.py_noISR:
args.output += '_noISR'
if args.py_noHadron:
args.output += '_noHadr'
if args.py_minbias > 0:
args.output += '_minbias'
if args.py_pthatmin > 0:
args.output += '_pthatmin_{}'.format(args.py_pthatmin)
if args.inel:
args.output += '_inel'
mycfg.append("SoftQCD:inelastic = on")
mycfg.append("HardQCD:all = off")
else:
if args.hard:
args.output += '_biasref_{}'.format(args.py_biasref)
args.output += '_hard'
mycfg.append("HardQCD:all = on")
else:
mycfg.append("HardQCD:all = on")
args.output += '_hard'
args.output += '.root'
print(args)
if os.path.exists(args.output):
if not args.overwrite:
print('[w] output file', args.output, 'exists - skipping.')
return
print('[i] output file', args.output)
# print the banner first
fj.ClusterSequence.print_banner()
print()
#and in the code i do not allow decay of the following particles:
mycfg.append("310:mayDecay = off") # //K0s
mycfg.append("3122:mayDecay = off") # //labda0
mycfg.append("3112:mayDecay = off") # //sigma-
mycfg.append("3212:mayDecay = off") # //sigma0
mycfg.append("3222:mayDecay = off") # //sigma+
mycfg.append("3312:mayDecay = off") # //xi-
mycfg.append("3322:mayDecay = off") # //xi+
mycfg.append("3334:mayDecay = off") # //omega-
print('[i] additional settings', mycfg)
pythia = pyconf.create_and_init_pythia_from_args(args, mycfg)
if not pythia:
return
if args.nev < 1:
args.nev = 1
tpairs = []
st = args.tranges.split(',')
try:
for _s in st:
t0 = float(_s.split('-')[0])
t1 = float(_s.split('-')[1])
tpairs.append([t0, t1])
except:
print(
'[e] something is not quite right with trigger ranges... stop here.'
)
jet_particle_eta_max = 0.9
jet_particle_selector = fj.SelectorAbsEtaMax(jet_particle_eta_max)
v0det = V0Detector()
j_ana = JetAnalysis(jet_R=args.jet_R,
particle_eta_max=jet_particle_eta_max)
fout = r.TFile(args.output, 'RECREATE')
fout.cd()
hmV0M = r.TH1F('hmV0M', 'hmV0M', 1000, 0, 1000)
hmV0A = r.TH1F('hmV0A', 'hmV0A', 1000, 0, 1000)
hmV0C = r.TH1F('hmV0C', 'hmV0C', 1000, 0, 1000)
event_output = RTreeWriter(tree_name='evT', fout=fout)
jet_output = RTreeWriter(tree_name='jetT', fout=fout)
hjet_sets = []
for t in tpairs:
hjet_output = RTreeWriter(tree_name='hjetT_{}_{}'.format(
int(t[0]), int(t[1])),
fout=fout)
hjet = HJet(trigger_range=[t[0], t[1]])
print(hjet)
hjset = HJetSet(hjet, hjet_output)
hjet_sets.append(hjset)
for ev_id in tqdm.tqdm(range(args.nev)):
if not pythia.next():
continue
if args.charged:
parts = pythiafjext.vectorize_select(
pythia, [pythiafjext.kFinal, pythiafjext.kCharged])
else:
parts = pythiafjext.vectorize_select(pythia, [pythiafjext.kFinal])
event_output.clear()
jet_output.clear()
_tmp = [s.hjet_output.clear() for s in hjet_sets]
jet_parts = jet_particle_selector(parts)
ev_s = pythia.info.sigmaGen()
ev_w = pythia.info.weight()
ev_code = pythia.info.code()
pthard = pythia.info.pTHat()
mTot = len(parts)
mCB = len(jet_parts)
v0det.analyze_event(parts)
event_output.fill_branch('ev_id', ev_id)
event_output.fill_branch('weight', ev_w)
event_output.fill_branch('sigma', ev_s)
event_output.fill_branch('code', ev_code)
event_output.fill_branch('pthard', pthard)
event_output.fill_branch('mTot', mTot)
event_output.fill_branch('mCB', mCB)
v0det.fill_branches(event_output)
event_output.fill_tree()
hmV0M.Fill(v0det.V0_mult)
hmV0A.Fill(v0det.V0A_mult)
hmV0C.Fill(v0det.V0C_mult)
if len(parts) < 1:
continue
if True in [s.analyze_event(jet_parts) for s in hjet_sets]:
j_ana.analyze_event(jet_parts)
jet_output.fill_branch('ev_id', ev_id)
jet_output.fill_branch('weight', ev_w)
jet_output.fill_branch('sigma', ev_s)
jet_output.fill_branch('code', ev_code)
jet_output.fill_branch('pthard', pthard)
jet_output.fill_branch('mTot', mTot)
jet_output.fill_branch('mCB', mCB)
j_ana.fill_branches(jet_output)
jet_output.fill_tree()
for s in hjet_sets:
if s.hjet.trigger_particle:
s.hjet_output.fill_branch('ev_id', ev_id)
s.hjet_output.fill_branch('weight', ev_w)
s.hjet_output.fill_branch('sigma', ev_s)
s.hjet_output.fill_branch('code', ev_code)
s.hjet_output.fill_branch('pthard', pthard)
s.hjet_output.fill_branch('mTot', mTot)
s.hjet_output.fill_branch('mCB', mCB)
s.hjet.fill_branches(s.hjet_output, j_ana.jets)
v0det.fill_branches(s.hjet_output)
j_ana.fill_branches(s.hjet_output)
s.hjet_output.fill_tree()
pythia.stat()
fout.Write()
fout.Close()
print('[i] written', fout.GetName())
def generate():
parser = argparse.ArgumentParser(description='pythia8 fastjet on the fly',
prog=os.path.basename(__file__))
parser.add_argument(
'-d',
'--output-dir',
help='output directory name - use with default file name',
default='.',
type=str)
parser.add_argument('-o',
'--output',
help='output file name',
default='',
type=str)
parser.add_argument('--overwrite', default=False, action='store_true')
parser.add_argument('--t-min',
help='trigger pt min',
default=6.,
type=float)
parser.add_argument('--t-max',
help='trigger pt max',
default=7.,
type=float)
parser.add_argument('--jet-R',
help='jet finder R',
default=0.4,
type=float)
parser.add_argument('--charged', default=False, action='store_true')
parser.add_argument('--runid', default=0, type=int)
pyconf.add_standard_pythia_args(parser)
args = parser.parse_args()
# mycfg = ['PhaseSpace:pThatMin = 6']
mycfg = []
if '.root' not in args.output:
rs = '0{}'.format(int(args.jet_R * 10))
if args.jet_R >= 1:
rs = '{}'.format(int(args.jet_R))
args.output = '{}/h_jet_R{}_t{}_{}'.format(args.output_dir, rs,
args.t_min, args.t_max)
if args.charged:
args.output = '{}/h_jet_ch_R{}_t{}_{}'.format(
args.output_dir, rs, int(args.t_min), int(args.t_max))
if args.py_pthatmin > 0:
args.output += '_pthatmin{}'.format(int(args.py_pthatmin))
if args.py_noMPI:
args.output += '_noMPI'
if args.py_noISR:
args.output += '_noISR'
if args.py_noHadron:
args.output += '_noHadr'
if args.py_minbias > 0:
args.output += '_minbias'
if args.runid > 0:
args.output += '_runid_{}'.format(args.runid)
args.py_seed = 1000 + args.runid
args.output += '.root'
if os.path.exists(args.output):
if not args.overwrite:
print('[w] output file', args.output, 'exists - skipping.')
return
# print the banner first
fj.ClusterSequence.print_banner()
print()
mycfg.append("SoftQCD:inelastic = on") # Andreas' recommendation
#and in the code i do not allow decay of the following particles:
mycfg.append("310:mayDecay = off") # //K0s
mycfg.append("3122:mayDecay = off") # //labda0
mycfg.append("3112:mayDecay = off") # //sigma-
mycfg.append("3212:mayDecay = off") # //sigma0
mycfg.append("3222:mayDecay = off") # //sigma+
mycfg.append("3312:mayDecay = off") # //xi-
mycfg.append("3322:mayDecay = off") # //xi+
mycfg.append("3334:mayDecay = off") # //omega-
print('[i] additional settings', mycfg)
pythia = pyconf.create_and_init_pythia_from_args(args, mycfg)
if not pythia:
return
if args.nev < 1:
args.nev = 1
jet_particle_eta_max = 0.9
jet_particle_selector = fj.SelectorAbsEtaMax(jet_particle_eta_max)
fout = r.TFile(args.output, 'RECREATE')
fout.cd()
hmV0M = r.TH1F('hmV0M', 'hmV0M', 1000, 0, 1000)
hmV0A = r.TH1F('hmV0A', 'hmV0A', 1000, 0, 1000)
hmV0C = r.TH1F('hmV0C', 'hmV0C', 1000, 0, 1000)
event_output = RTreeWriter(tree_name='evT', fout=fout)
jet_output = RTreeWriter(tree_name='jetT', fout=fout)
hjet_output = RTreeWriter(tree_name='hjetT', fout=fout)
v0det = V0Detector()
j_ana = JetAnalysis(jet_R=args.jet_R,
particle_eta_max=jet_particle_eta_max)
hjet = HJet(trigger_range=[args.t_min, args.t_max])
for ev_id in tqdm.tqdm(range(args.nev)):
if not pythia.next():
continue
if args.charged:
parts = pythiafjext.vectorize_select(
pythia, [pythiafjext.kFinal, pythiafjext.kCharged])
else:
parts = pythiafjext.vectorize_select(pythia, [pythiafjext.kFinal])
event_output.clear()
jet_output.clear()
hjet_output.clear()
jet_parts = jet_particle_selector(parts)
ev_w = pythia.info.sigmaGen()
ev_code = pythia.info.code()
pthard = pythia.info.pTHat()
mTot = len(parts)
mCB = len(jet_parts)
v0det.analyze_event(parts)
event_output.fill_branch('ev_id', ev_id)
event_output.fill_branch('weight', ev_w)
event_output.fill_branch('code', ev_code)
event_output.fill_branch('pthard', pthard)
event_output.fill_branch('mTot', mTot)
event_output.fill_branch('mCB', mCB)
v0det.fill_branches(event_output)
event_output.fill_tree()
hmV0M.Fill(v0det.V0_mult)
hmV0A.Fill(v0det.V0A_mult)
hmV0C.Fill(v0det.V0C_mult)
if len(parts) < 1:
continue
hjet.analyze_event(jet_parts)
if hjet.trigger_particle:
j_ana.analyze_event(jet_parts)
jet_output.fill_branch('ev_id', ev_id)
jet_output.fill_branch('weight', ev_w)
jet_output.fill_branch('code', ev_code)
jet_output.fill_branch('pthard', pthard)
jet_output.fill_branch('mTot', mTot)
jet_output.fill_branch('mCB', mCB)
j_ana.fill_branches(jet_output)
jet_output.fill_tree()
hjet_output.fill_branch('ev_id', ev_id)
hjet_output.fill_branch('weight', ev_w)
hjet_output.fill_branch('code', ev_code)
hjet_output.fill_branch('pthard', pthard)
hjet_output.fill_branch('mTot', mTot)
hjet_output.fill_branch('mCB', mCB)
hjet.fill_branches(hjet_output, j_ana.jets)
v0det.fill_branches(hjet_output)
j_ana.fill_branches(hjet_output)
hjet_output.fill_tree()
pythia.stat()
fout.Write()
fout.Close()
print('[i] written', fout.GetName())
def main(agrs):
fj.ClusterSequence.print_banner()
print()
ja = JetAnalysis(jet_R=args.jet_R,
jet_algorithm=fj.antikt_algorithm,
particle_eta_max=args.max_eta)
grfj = GridFastJet(grid_size=0.01,
jet_R=args.jet_R,
jet_algorithm=fj.antikt_algorithm,
particle_eta_max=args.max_eta)
print(grfj)
grfjR = GridFastJetROOT(grid_size=0.01,
jet_R=args.jet_R,
jet_algorithm=fj.antikt_algorithm,
particle_eta_max=args.max_eta)
print(grfjR)
pythia = None
if args.pythia:
mycfg = []
pythia = pyconf.create_and_init_pythia_from_args(args, mycfg)
if not pythia:
print("[e] pythia initialization failed.")
return
fout = ROOT.TFile(args.output, 'recreate')
fout.cd()
ts = RTreeWriter(tree_name='ts', fout=fout)
tg = RTreeWriter(tree_name='tg', fout=fout)
tgR = RTreeWriter(tree_name='tgR', fout=fout)
for iev in tqdm.tqdm(range(args.nev)):
if pythia:
if not pythia.next():
continue
parts_pythia = pythiafjext.vectorize_select(
pythia, [pythiafjext.kFinal, pythiafjext.kCharged])
ja.analyze_event(parts=parts_pythia)
if len(ja.jets) < 1:
continue
if ja.jets[0].pt() < 10:
continue
ts.fill_branch('j', ja.jets)
ts.fill_branch('p', ja.particles)
ts.fill_tree()
if args.grid:
grfj.analyze_event(parts=parts_pythia)
tg.fill_branch('j', grfj.jets)
tg.fill_branch('p', grfj.particles)
tg.fill_tree()
if args.gridR:
grfjR.analyze_event(parts=parts_pythia)
tgR.fill_branch('j', grfjR.jets)
tgR.fill_branch('p', grfjR.particles)
tgR.fill_tree()
ts.write_and_close()
class HJetAnalysis(MPBase):
def __init__(self, **kwargs):
self.configure_from_args(jet_particle_eta_max = 0.9,
output='hjet.root',
trigger_ranges = [ [0, 1e3], [6, 7] , [12, 22], [20, 30] ],
jet_Rs = [0.2, 0.4, 0.6])
super(HJetAnalysis, self).__init__(**kwargs)
self.fout = r.TFile(self.output, 'RECREATE')
self.fout.cd()
self.event_output = RTreeWriter(tree_name='tev', fout=self.fout)
self.hmV0M = r.TH1F('hmV0M', 'hmV0M', 1000, 0, 1000)
self.hmV0A = r.TH1F('hmV0A', 'hmV0A', 1000, 0, 1000)
self.hmV0C = r.TH1F('hmV0C', 'hmV0C', 1000, 0, 1000)
self.jet_particle_selector = fj.SelectorAbsEtaMax(self.jet_particle_eta_max)
self.jet_ans = []
self.hjet_ts = []
for jR in self.jet_Rs:
j_ana = JetAnalysis(jet_R=jR, particle_eta_max=self.jet_particle_eta_max)
self.jet_ans.append(j_ana)
for tr in self.trigger_ranges:
hjet_tree = HJetTree(fout=self.fout, trigger_range=tr, jet_ana=j_ana)
self.hjet_ts.append(hjet_tree)
def analyze_event(self, ev_id, parts, pythia=None):
# V0 multiplicity and event props
v0det = V0Detector(particles=parts)
self.hmV0M.Fill(v0det.V0_mult)
self.hmV0A.Fill(v0det.V0A_mult)
self.hmV0C.Fill(v0det.V0C_mult)
self.event_output.fill_branch('ev_id', ev_id)
self.event_output.fill_branch('mV0A', v0det.V0A_mult)
self.event_output.fill_branch('mV0C', v0det.V0C_mult)
self.event_output.fill_branch('mV0', v0det.V0_mult)
if pythia:
ev_w = pythia.info.sigmaGen()
pthard = pythia.info.pTHat()
else:
ev_w = -1
pthard = -1
self.event_output.fill_branch('weight', ev_w)
self.event_output.fill_branch('pthard', pthard)
jet_parts = self.jet_particle_selector(parts)
mTot = len(parts)
mCB = len(jet_parts)
self.event_output.fill_branch('mTot', mTot)
self.event_output.fill_branch('mCB', mCB)
self.event_output.fill_tree()
trigger_present = False
for hj in self.hjet_ts:
if hj.analyze_event(jet_parts):
trigger_present = True
if trigger_present:
for ja in self.jet_ans:
ja.analyze_event(jet_parts)
for hj in self.hjet_ts:
if hj.trigger_particle:
hj.fill_branches(ev_id, ev_w)
hj.tree_writer.fill_branch('ev_id', ev_id)
hj.tree_writer.fill_branch('mV0A', v0det.V0A_mult)
hj.tree_writer.fill_branch('mV0C', v0det.V0C_mult)
hj.tree_writer.fill_branch('mV0', v0det.V0_mult)
hj.tree_writer.fill_branch('mTot', mTot)
hj.tree_writer.fill_branch('mCB', mCB)
hj.tree_writer.fill_branch('weight', ev_w)
hj.tree_writer.fill_branch('pthard', pthard)
hj.fill_tree()
def finalize(self):
self.fout.Write()
self.fout.Close()
print('[i] written', self.fout.GetName())
def main_test(args):
if os.path.isfile(args.output):
if not args.overwrite:
print('[i] output', args.output, 'exists - use --overwrite to do just that...')
return
print(args)
# alice specific
max_eta = 0.9
jet_R0 = 0.4
mycfg = []
if args.benchmark:
mycfg = ['PhaseSpace:pThatMin = 80', 'PhaseSpace:pThatMax = -1']
jet_selector = fj.SelectorPtMin(80.0) & fj.SelectorPtMax(100.0) & fj.SelectorAbsEtaMax(max_eta - 1.05 * jet_R0)
# jet_selector_cs = fj.SelectorPtMin(50.0) & fj.SelectorAbsEtaMax(max_eta - 1.05 * jet_R0)
else:
args.py_biaspow = 4
args.py_biasref = 10
jet_selector = fj.SelectorPtMin(20) & fj.SelectorAbsEtaMax(max_eta - 1.05 * jet_R0)
# jet_selector_cs = fj.SelectorPtMin(50.0) & fj.SelectorAbsEtaMax(max_eta - 1.05 * jet_R0)
if args.ignore_mycfg:
mycfg = []
pythia = pyconf.create_and_init_pythia_from_args(args, mycfg)
if not pythia:
print("[e] pythia initialization failed.")
return
sd_zcut = 0.1
sd = fjcontrib.SoftDrop(0, sd_zcut, jet_R0)
parts_selector = fj.SelectorAbsEtaMax(max_eta)
fout = ROOT.TFile(args.output, 'recreate')
fout.cd()
twj = RTreeWriter(tree_name='jets', fout=fout)
twjm = RTreeWriter(tree_name='jetsm', fout=fout)
twp = RTreeWriter(tree_name='parts', fout=fout)
bgestim = BackgroundEstimator(input_file=args.input, name='bg_estim')
print(bgestim)
boltzmann_subtractor = BoltzmannSubtractor(max_pt_subtract=10.)
print(boltzmann_subtractor)
be = None
data = None
if len(args.data) > 0:
data = DataIO(file_list=args.data)
print(data)
else:
# be = BoltzmannEvent(mean_pt=0.6, multiplicity=2000 * max_eta * 2, max_eta=max_eta, max_pt=100)
pass
be = fjtools.BoltzmannBackground(0.7, 0.01, 1.5)
print(be.description())
parts_selector = fj.SelectorAbsEtaMax(max_eta)
if args.nev < 1:
args.nev = 1
jas = JetAnalysisWithRho(jet_R=jet_R0, jet_algorithm=fj.antikt_algorithm, particle_eta_max=max_eta)
ja = JetAnalysisWithRho(jet_R=jet_R0, jet_algorithm=fj.antikt_algorithm, particle_eta_max=max_eta)
jabg = JetAnalysisWithRho(jet_R=jet_R0, jet_algorithm=fj.antikt_algorithm, particle_eta_max=max_eta)
### EVENT LOOP STARTS HERE
pbar = tqdm.tqdm(total=args.nev)
iev = 0
while pbar.n < args.nev:
iev = iev + 1
if not pythia.next():
continue
parts_pythia = pythiafjext.vectorize_select(pythia, [pythiafjext.kFinal, pythiafjext.kCharged])
parts_gen = parts_selector(parts_pythia)
jas.analyze_event(parts_gen)
signal_jets = fj.sorted_by_pt(jet_selector(jas.jets))
if len(signal_jets) < 1:
continue
pbar.set_description('tried %i' % iev)
pbar.update(1)
if data:
bg_parts = data.load_event(offset=10000)
else:
if be:
# bg_parts = be.generate(offset=10000)
bg_parts = be.generate(int(2000 * max_eta * 2.), max_eta, 10000)
# len_bg_parts = len(bg_parts)
# subtr_parts_b = be.subtract(bg_parts, 0.7, len(bg_parts))
# print('x nparts_bg=', len(bg_parts), 'nparts_pass='******' ', be.get_formula())
# subtr_parts_b = be.subtract(bg_parts)
# print(' - nparts_bg=', len(bg_parts), 'nparts_pass='******' ', be.get_formula())
_pythia_inserts = []
_pythia_inserts_indexes = []
for sjet in signal_jets:
_pythia_inserts = [bg_parts.push_back(psj) for psj in sjet.constituents() if psj.user_index() >= 0]
_pythia_inserts_indexes.extend([psj.user_index() for psj in sjet.constituents() if psj.user_index() >= 0])
print ('pythia indexes:', len(_pythia_inserts_indexes), sorted(_pythia_inserts_indexes))
# subtr_parts_b = be.subtract(bg_parts)
# print(' - nparts_bg=', len(bg_parts), 'nparts_pass='******' ', be.get_formula())
# subtr_parts_b = be.subtract(bg_parts, 0.7, len(_pythia_inserts))
# print('. nparts_bg=', len_bg_parts, 'nparts_pass='******' ', be.get_formula())
# _mpt = mean_pt(bg_parts)
# subtr_parts_b = be.subtract(bg_parts, _mpt, len(_pythia_inserts))
# print('. nparts_bg=', len(bg_parts), 'nparts_pass='******' ', be.get_formula())
# subtr_parts = bgestim.subtracted_particles(bg_parts)
# ja.analyze_event(subtr_parts)
# for j in ja.jets:
# if matched_pt(j, sjet) > 0.5:
# tw.fill_branch('sjet_bg', j)
# tw.fill_branch('delta_pt', sjet.pt() - j.pt())
#ja.analyze_event(bg_parts)
#bg_signal_jets_b = fj.sorted_by_pt(ja.jets)
#for j in bg_signal_jets_b:
# if matched_pt(j, sjet) > 0.5:
# tw.fill_branch('sjet_hybrid', j)
# tw.fill_branch('delta_pt_hybrid', sjet.pt() - j.pt())
# subtr_parts_b = boltzmann_subtractor.subtracted_particles(bg_parts)
# _it1 = [p for p in be.subtract(bg_parts)]
# subtr_parts_b = be.subtract(_it1)
jabg.analyze_event(bg_parts)
bg_parts_noleads = remove_jets(bg_parts, jabg.jets[:2])
# print(bg_parts_noleads.size())
#for j in jabg.jets:
# pass
# continue
# if i > 2:
# break
# _tmp = [p.pt() for p in j.constituents()]
# reduce_total_pt = reduce_total_pt + sum(_tmp)
#print ('reduce total pt by', reduce_total_pt)
# subtr_parts_b = be.subtract(bg_parts)
# print(be.get_formula())
# print('nparts_bg=', len(bg_parts), 'nparts_subtr=', len(subtr_parts_b), 'nparts_pythia=', len(tmp))
_mpt = mean_pt(bg_parts_noleads)
# subtr_parts_b = be.subtract(bg_parts, _mpt, len(_pythia_inserts))
# subtr_parts_b = be.subtract(bg_parts, _mpt, len(bg_parts) - len(bg_parts_noleads))
subtr_parts_b = be.subtract(bg_parts, _mpt, len(bg_parts_noleads))
print('. nparts_bg=', len(bg_parts), 'nparts_pass='******' ', be.get_formula())
_indexes_subtr_parts_b = [p.user_index() for p in subtr_parts_b]
print ('remaining indexes:', len(_indexes_subtr_parts_b), sorted(_indexes_subtr_parts_b))
ja.analyze_event(subtr_parts_b)
bg_signal_jets_b = fj.sorted_by_pt(ja.jets)
w = pythia.info.weight()
s = pythia.info.sigmaGen()
twj.fill_branch('sjet', signal_jets)
twj.fill_tree()
print( 'py jets=', len([j for j in signal_jets if j.pt() > 0.1 ]),
'sub jets=', len([j for j in bg_signal_jets_b if j.pt() > 0.1]))
_tmp = [fill_tree_matched(sjet, j, twjm, sd, 0, iev, weight=w, sigma=s)
for sjet in signal_jets
for j in bg_signal_jets_b]
twp.fill_branch('p', bg_parts)
#twp.fill_branch('sp', subtr_parts)
twp.fill_branch('sp_b', subtr_parts_b)
twp.fill_tree()
pbar.close()
fout.cd()
fout.Write()
fout.Close()
print('[i] written', fout.GetName())