def main(args):
nevents = args.nevents
sconfig_pythia = get_pythia_config(args)
if args.generate:
pythia = create_and_init_pythia(sconfig_pythia)
if not pythia:
return
all_jets = []
for iEvent in tqdm(range(nevents), 'event'):
if not pythia.next(): continue
print("[i] done generating")
if args.write:
pythia = create_and_init_pythia(sconfig_pythia)
if not pythia:
return
pyhepmcwriter = mp.Pythia8HepMCWrapper(args.write)
all_jets = []
for iEvent in tqdm(range(nevents), 'event'):
if not pythia.next(): continue
pyhepmcwriter.fillEvent(pythia)
print("[i] done writing to {}".format(args.write))
if args.read:
import pyhepmc_ng
input = pyhepmc_ng.ReaderAsciiHepMC2(args.read)
if input.failed():
print("[error] unable to read from {}".format(args.read))
return
# print the banner first
fj.ClusterSequence.print_banner()
print()
jet_R0 = 0.4
jet_def = fj.JetDefinition(fj.antikt_algorithm, jet_R0)
jet_selector = fj.SelectorPtMin(100.0) & fj.SelectorPtMax(
200.0) & fj.SelectorAbsEtaMax(1)
event = pyhepmc_ng.GenEvent()
pbar = tqdm(range(nevents))
while not input.failed():
e = input.read_event(event)
if input.failed():
break
fjparts = []
for i, p in enumerate(event.particles):
if p.status == 1:
psj = fj.PseudoJet(p.momentum.px, p.momentum.py,
p.momentum.pz, p.momentum.e)
psj.set_user_index(i)
fjparts.append(psj)
jets = jet_selector(jet_def(fjparts))
pbar.update()
def main_jets(args):
outfname = '{}_output.root'.format(args.read).replace('.dat', '')
print("output file: {}".format(outfname))
foutput = r.TFile(outfname, "recreate")
foutput.cd()
lbins = logbins(1., 100, 10)
hjetpt = r.TH1F('hjetpt', 'hjetpt', 10, lbins)
input = pyhepmc_ng.ReaderAsciiHepMC2(args.read)
if input.failed():
print ("[error] unable to read from {}".format(args.read))
return
# print the banner first
fj.ClusterSequence.print_banner()
print()
jet_R0 = 0.4
jet_def = fj.JetDefinition(fj.antikt_algorithm, jet_R0)
jet_selector = fj.SelectorPtMin(0.0) & fj.SelectorPtMax(200.0) & fj.SelectorAbsEtaMax(3)
all_jets = []
event = pyhepmc_ng.GenEvent()
pbar = tqdm(range(args.nevents))
while not input.failed():
e = input.read_event(event)
if input.failed():
break
fjparts = []
for i,p in enumerate(event.particles):
if p.status == 1 and not p.end_vertex:
psj = fj.PseudoJet(p.momentum.px, p.momentum.py, p.momentum.pz, p.momentum.e)
psj.set_user_index(i)
fjparts.append(psj)
jets = jet_selector(jet_def(fjparts))
all_jets.append([ [j.pt(), j.eta()] for j in jets])
pbar.update()
for j in jets:
hjetpt.Fill(j.perp())
if pbar.n >= args.nevents:
break
foutput.Write()
foutput.Close()
joblib.dump(all_jets, outfname.replace(".root", ".joblib"))
def main():
nevents = 1000
sconfig_pythia = [
"Beams:eCM = 8000.", "HardQCD:all = on", "PhaseSpace:pTHatMin = 100.",
"PDF:pset = LHAPDF6:EPPS16nlo_CT14nlo_Pb208/0"
]
# "LHAPDF6:member = 901300"]
# "PDF:pSet = EPPS16nlo_CT14nlo_Pb208"]
if 'gen' in sys.argv:
pythia = create_and_init_pythia(sconfig_pythia)
if not pythia:
return
all_jets = []
for iEvent in tqdm(range(nevents), 'event'):
if not pythia.next(): continue
print("[i] done generating")
if 'write' in sys.argv:
pythia = create_and_init_pythia(sconfig_pythia)
if not pythia:
return
pyhepmcwriter = mp.Pythia8HepMCWrapper("test_pythia8_hepmc.dat")
all_jets = []
for iEvent in tqdm(range(nevents), 'event'):
if not pythia.next(): continue
pyhepmcwriter.fillEvent(pythia)
print("[i] done writing to test_pythia8_hepmc.dat")
if 'read' in sys.argv:
import pyhepmc_ng
input = pyhepmc_ng.ReaderAsciiHepMC2("test_pythia8_hepmc.dat")
if input.failed():
print("[error] unable to read from test_pythia8_hepmc.dat")
return
event = pyhepmc_ng.GenEvent()
pbar = tqdm(range(nevents))
while not input.failed():
p = input.read_event(event)
if input.failed():
# print ("[i] reading done.")
break
pbar.update()
def main(self):
if self.hepmc == 3:
input_hepmc = pyhepmc_ng.ReaderAscii(self.input)
if self.hepmc == 2:
input_hepmc = pyhepmc_ng.ReaderAsciiHepMC2(self.input)
if input_hepmc.failed():
print("[error] unable to read from {}".format(self.input))
sys.exit(1)
event_hepmc = pyhepmc_ng.GenEvent()
while not input_hepmc.failed():
ev = input_hepmc.read_event(event_hepmc)
if input_hepmc.failed():
break
self.fill_event(event_hepmc)
self.increment_event()
if self.nev > 0 and self.ev_id > self.nev:
break
self.finish()
def main():
parser = argparse.ArgumentParser(description='pythia8 in python',
prog=os.path.basename(__file__))
parser.add_argument('-i',
'--input',
help='input file',
default='low',
type=str,
required=True)
parser.add_argument('--hepmc',
help='what format 2 or 3',
default=2,
type=int)
parser.add_argument('--nev', help='number of events', default=10, type=int)
args = parser.parse_args()
###
# now lets read the HEPMC file and do some jet finding
if args.hepmc == 3:
input_hepmc = pyhepmc_ng.ReaderAscii(args.input)
if args.hepmc == 2:
input_hepmc = pyhepmc_ng.ReaderAsciiHepMC2(args.input)
if input_hepmc.failed():
print("[error] unable to read from {}".format(args.input))
sys.exit(1)
# jet finder
# print the banner first
fj.ClusterSequence.print_banner()
print()
jet_R0 = 0.4
jet_def = fj.JetDefinition(fj.antikt_algorithm, jet_R0)
jet_selector = fj.SelectorPtMin(10.0) & fj.SelectorPtMax(
500.0) & fj.SelectorAbsEtaMax(3)
all_jets = []
event_hepmc = pyhepmc_ng.GenEvent()
pbar = tqdm.tqdm(range(args.nev))
while not input_hepmc.failed():
ev = input_hepmc.read_event(event_hepmc)
if input_hepmc.failed():
break
jets_hepmc = find_jets_hepmc(jet_def, jet_selector, event_hepmc)
all_jets.extend(jets_hepmc)
pbar.update()
if pbar.n >= args.nev:
break
jet_def_lund = fj.JetDefinition(fj.cambridge_algorithm, 1.0)
lund_gen = fjcontrib.LundGenerator(jet_def_lund)
print('[i] making lund diagram for all jets...')
lunds = [lund_gen.result(j) for j in all_jets]
print('[i] listing lund plane points... Delta, kt - for {} selected jets'.
format(len(all_jets)))
for l in lunds:
if len(l) < 1:
continue
print('- jet pT={0:5.2f} eta={1:5.2f}'.format(l[0].pair().perp(),
l[0].pair().eta()))
print(' Deltas={}'.format([s.Delta() for s in l]))
print(' kts={}'.format([s.kt() for s in l]))
print()
print('[i] reclustering and using soft drop...')
jet_def_rc = fj.JetDefinition(fj.cambridge_algorithm, 0.1)
print('[i] Reclustering:', jet_def_rc)
all_jets_sd = []
rc = fjcontrib.Recluster(jet_def_rc, True)
sd = fjcontrib.SoftDrop(0, 0.1, 1.0)
for i, j in enumerate(all_jets):
j_rc = rc.result(j)
print()
print('- [{0:3d}] orig pT={1:10.3f} reclustered pT={2:10.3f}'.format(
i, j.perp(), j_rc.perp()))
j_sd = sd.result(j)
print(' |-> after soft drop pT={0:10.3f} delta={1:10.3f}'.format(
j_sd.perp(),
j_sd.perp() - j.perp()))
all_jets_sd.append(j_sd)
sd_info = fjcontrib.get_SD_jet_info(j_sd)
print(
" |-> SD jet params z={0:10.3f} dR={1:10.3f} mu={2:10.3f}".format(
sd_info.z, sd_info.dR, sd_info.mu))
fout = ROOT.TFile('hepmc_jetreco.root', 'recreate')
lbins = logbins(1., 500, 50)
hJetPt04 = ROOT.TH1D("hJetPt04", "hJetPt04", 50, lbins)
hJetPt04sd = ROOT.TH1D("hJetPt04sd", "hJetPt04sd", 50, lbins)
[hJetPt04.Fill(j.perp()) for j in all_jets]
[hJetPt04sd.Fill(j.perp()) for j in all_jets_sd]
hLund = ROOT.TH2D("hLund", "hLund", 60, 0, 6, 100, -4, 5)
lunds = [lund_gen.result(j) for j in all_jets if j.perp() > 100]
j100 = [j for j in all_jets if j.perp() > 100]
print('{} jets above 100 GeV/c'.format(len(j100)))
for l in lunds:
for s in l:
hLund.Fill(math.log(1. / s.Delta()), math.log(s.kt()))
fout.Write()
def main():
parser = argparse.ArgumentParser(
description='hepmc to ALICE Ntuple format',
prog=os.path.basename(__file__))
parser.add_argument('-i',
'--input',
help='input file',
default='',
type=str,
required=True)
parser.add_argument('-o',
'--output',
help='output root file',
default='',
type=str,
required=True)
parser.add_argument('--as-data',
help='write as data - tree naming convention',
action='store_true',
default=False)
parser.add_argument('--hepmc',
help='what format 2 or 3',
default=2,
type=int)
parser.add_argument('--nev', help='number of events', default=-1, type=int)
args = parser.parse_args()
if args.hepmc == 3:
input_hepmc = pyhepmc_ng.ReaderAscii(args.input)
if args.hepmc == 2:
input_hepmc = pyhepmc_ng.ReaderAsciiHepMC2(args.input)
if input_hepmc.failed():
print("[error] unable to read from {}".format(args.input))
sys.exit(1)
outf = ROOT.TFile(args.output, 'recreate')
outf.cd()
tdf = ROOT.TDirectoryFile('PWGHF_TreeCreator', 'PWGHF_TreeCreator')
tdf.cd()
if args.as_data:
t_p = ROOT.TNtuple(
'tree_Particle', 'tree_Particle',
'run_number:ev_id:ParticlePt:ParticleEta:ParticlePhi:ParticlePID')
else:
t_p = ROOT.TNtuple(
'tree_Particle_gen', 'tree_Particle_gen',
'run_number:ev_id:ParticlePt:ParticleEta:ParticlePhi:ParticlePID')
t_e = ROOT.TNtuple('tree_event_char', 'tree_event_char',
'run_number:ev_id:z_vtx_reco:is_ev_rej')
# run number will be a double - file size in MB
run_number = os.path.getsize(args.input) / 1.e6
ev_id = 0
# unfortunately pyhepmc_ng does not provide the table
# pdt = pyhepmc_ng.ParticleDataTable()
# use ROOT instead
pdg = ROOT.TDatabasePDG()
event_hepmc = pyhepmc_ng.GenEvent()
if args.nev > 0:
pbar = tqdm.tqdm(range(args.nev))
else:
pbar = tqdm.tqdm()
while not input_hepmc.failed():
ev = input_hepmc.read_event(event_hepmc)
if input_hepmc.failed():
break
fill_event(run_number, ev_id, event_hepmc, t_e, t_p, pdg)
ev_id = ev_id + 1
pbar.update()
if args.nev > 0 and ev_id > args.nev:
break
outf.Write()
outf.Close()