def test_fill_genevent_from_hepevt(evt):
h = hep.HEPEVT()
h.from_genevent(evt)
evt2 = hep.GenEvent()
hep.fill_genevent_from_hepevt(evt2,
event_number=h.event_number,
status=h.status(),
pid=h.pid(),
p=h.pm()[:, :4] * 2,
m=h.pm()[:, 4] * 2,
v=h.v() * 5,
parents=h.parents(),
momentum_scaling=2,
vertex_scaling=5)
# hepevt has no run_info, so we add it articially
evt2.run_info = evt.run_info
assert evt == evt2
with pytest.raises(ValueError, match="exceeds HepMC3 buffer size"):
n = h.max_size + 1
x = np.zeros(n, dtype=int)
p = np.zeros((n, 2))
v = np.zeros((n, 4), dtype=float)
hep.fill_genevent_from_hepevt(evt2,
event_number=h.event_number,
status=x,
pid=x,
p=v,
m=v[:, 0],
v=v,
parents=p)
with pytest.raises(KeyError):
hep.fill_genevent_from_hepevt(evt2)
def test_fill_genevent_from_hepevt_ptr():
evt1 = prepare_event()
h = prepare_hepevt(evt1)
evt2 = hep.GenEvent()
hep.fill_genevent_from_hepevent_ptr(evt2, h.ptr, h.max_size)
assert evt1.particles == evt2.particles
assert evt1.vertices == evt2.vertices
assert evt1 == evt2
def test_hepevt(evt):
particles = evt.particles
h = hep.HEPEVT()
h.from_genevent(evt)
evt2 = hep.GenEvent()
h.to_genevent(evt2)
# hepevt has no run_info, so we add it articially
evt2.run_info = evt.run_info
assert evt == evt2
def next_event(self):
event = pyhepmc_ng.GenEvent()
self.reader.read_event(event)
if self.reader.failed():
return False
else:
ev = event_hepmc.EventHepMC(event)
return ev
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 test_sequence_access():
evt = hep.GenEvent()
evt.add_particle(hep.GenParticle())
evt.particles[0].momentum = (1, 2, 3, 4)
evt.particles[0].pid = 5
evt.add_vertex(hep.GenVertex())
evt.vertices[0].position = (1, 2, 3, 4)
assert len(evt.particles) == 1
assert evt.particles[0].id == 1
assert evt.particles[0].pid == 5
assert evt.particles[0].momentum == (1, 2, 3, 4)
assert len(evt.vertices) == 1
assert evt.vertices[0].id == -1
assert evt.vertices[0].position == (1, 2, 3, 4)
def test_weights():
evt = hep.GenEvent()
assert evt.weights == []
with pytest.raises(RuntimeError,
match=".*requires the event to have a GenRunInfo"):
evt.weight_names
evt.run_info = hep.GenRunInfo()
evt.run_info.weight_names = ["a", "b"]
evt.weights = [2, 3]
assert evt.weight(1) == 3
assert evt.weight("a") == 2
with pytest.raises(IndexError):
evt.weight(2)
with pytest.raises(RuntimeError, match="no weight with given name"):
evt.set_weight("c", 4)
def test_read_write(evt):
oss = hep.stringstream()
with hep.WriterAscii(oss) as f:
f.write_event(evt)
evt2 = hep.GenEvent()
assert evt != evt2
with hep.ReaderAscii(oss) as f:
f.read_event(evt2)
assert evt.event_number == evt2.event_number
assert evt.momentum_unit == evt2.momentum_unit
assert evt.length_unit == evt2.length_unit
assert evt.particles == evt2.particles
assert evt.vertices == evt2.vertices
assert evt == evt2
def test_sequence_access():
evt = hep.GenEvent()
evt.add_particle(hep.GenParticle())
evt.particles[0].momentum = (1, 2, 3, 4)
evt.particles[0].pid = 5
evt.add_vertex(hep.GenVertex())
evt.vertices[0].position = (1, 2, 3, 4)
assert len(evt.particles) == 1
assert evt.particles[0].id == 1
assert evt.particles[0].pid == 5
assert evt.particles[0].momentum == (1, 2, 3, 4)
assert len(evt.vertices) == 1
assert evt.vertices[0].id == -1
assert evt.vertices[0].position == (1, 2, 3, 4)
assert repr(
evt
) == "GenEvent(momentum_unit=1, length_unit=0, event_number=0, particles=[GenParticle(FourVector(1, 2, 3, 4), status=0, id=1, production_vertex=0, end_vertex=-1)], vertices=[GenVertex(FourVector(1, 2, 3, 4), status=0, id=-1, particles_in=[], particles_out=[])], run_info=None)"
def main():
input_file = "$HOME/data/jetscape/test_out.hepmc"
if len(sys.argv) > 1:
input_file = sys.argv[1]
input_file = os.path.expandvars(input_file)
print('[i] reading from:', input_file)
# input = pyhepmc_ng.ReaderAsciiHepMC2(input_file)
input = pyhepmc_ng.ReaderAscii(input_file)
if input.failed():
print("[error] unable to read from {}".format(input_file))
return
nevents = 1000
# 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.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 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 >= nevents:
break
pbar.close()
def test_read_write_file():
evt1 = prepare_event()
with hep.WriterAscii("test_read_write_file.dat") as f:
f.write_event(evt1)
evt2 = hep.GenEvent()
assert evt1 != evt2
with hep.ReaderAscii("test_read_write_file.dat") as f:
ok = f.read_event(evt2)
assert ok
assert evt1.particles == evt2.particles
assert evt1.vertices == evt2.vertices
assert evt1 == evt2
import os
os.unlink("test_read_write_file.dat")
def main():
input_file = "$HOME/data/jetscape/test_out.hepmc"
if len(sys.argv) > 1:
input_file = sys.argv[1]
input_file = os.path.expandvars(input_file)
print('[i] reading from:', input_file)
input = pyhepmc_ng.ReaderAscii(input_file)
if input.failed():
print("[error] unable to read from {}".format(input_file))
return
event = pyhepmc_ng.GenEvent()
while not input.failed():
e = input.read_event(event)
if input.failed():
break
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 test_fill_genevent_from_hepevt():
evt1 = prepare_event()
h = prepare_hepevt(evt1)
evt2 = hep.GenEvent()
hep.fill_genevent_from_hepevt(
evt2,
h.event_number,
h.pm()[:, :4] * 2,
h.pm()[:, 4] * 2,
h.v() * 5,
h.pid(),
h.parents(),
h.children(),
h.status(), # particle status
np.zeros_like(h.pid()), # vertex status
2,
5)
assert evt1.particles == evt2.particles
assert evt1.vertices == evt2.vertices
assert evt1 == evt2
def test_read_write_stream():
evt1 = prepare_event()
oss = hep.stringstream()
with hep.WriterAscii(oss) as f:
f.write_event(evt1)
assert str(oss) == """HepMC::Version 3.01.00
HepMC::Asciiv3-START_EVENT_LISTING
E 1 4 8
U GEV MM
W 1.0000000000000000000000e+00
P 1 0 2212 0.0000000000000000e+00 0.0000000000000000e+00 7.0000000000000000e+03 7.0000000000000000e+03 9.3799999999999994e-01 1
P 2 0 2212 0.0000000000000000e+00 0.0000000000000000e+00 -7.0000000000000000e+03 7.0000000000000000e+03 9.3799999999999994e-01 2
V -1 0 [1] @ 1.0000000000000000e+00 1.0000000000000000e+00 1.0000000000000000e+00 1.0000000000000000e+00
P 3 -1 1 7.5000000000000000e-01 -1.5690000000000000e+00 3.2191000000000003e+01 3.2238000000000000e+01 0.0000000000000000e+00 3
V -2 0 [2] @ 2.0000000000000000e+00 2.0000000000000000e+00 2.0000000000000000e+00 2.0000000000000000e+00
P 4 -2 -2 -3.0470000000000002e+00 -1.9000000000000000e+01 -5.4628999999999998e+01 5.7920000000000002e+01 0.0000000000000000e+00 4
V -3 0 [3,4] @ 3.0000000000000000e+00 3.0000000000000000e+00 3.0000000000000000e+00 3.0000000000000000e+00
P 5 -3 -24 1.5169999999999999e+00 -2.0680000000000000e+01 -2.0605000000000000e+01 8.5924999999999997e+01 8.0799000000000007e+01 5
P 6 -3 22 -3.8130000000000002e+00 1.1300000000000000e-01 -1.8330000000000000e+00 4.2329999999999997e+00 0.0000000000000000e+00 6
V -4 0 [5] @ 4.0000000000000000e+00 4.0000000000000000e+00 4.0000000000000000e+00 4.0000000000000000e+00
P 7 -4 1 -2.4449999999999998e+00 2.8815999999999999e+01 6.0819999999999999e+00 2.9552000000000000e+01 1.0000000000000000e-02 7
P 8 -4 -2 3.9620000000000002e+00 -4.9497999999999998e+01 -2.6687000000000001e+01 5.6372999999999998e+01 6.0000000000000001e-03 8
HepMC::Asciiv3-END_EVENT_LISTING
"""
evt2 = hep.GenEvent()
assert evt1 != evt2
with hep.ReaderAscii(oss) as f:
f.read_event(evt2)
assert evt1.event_number == evt2.event_number
assert evt1.momentum_unit == evt2.momentum_unit
assert evt1.length_unit == evt2.length_unit
assert evt1.particles == evt2.particles
assert evt1.vertices == evt2.vertices
assert evt1 == evt2
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 prepare_event():
#
# In this example we will place the following event into HepMC "by hand"
#
# name status pdg_id parent Px Py Pz Energy Mass
# 1 !p+! 3 2212 0,0 0.000 0.000 7000.000 7000.000 0.938
# 2 !p+! 3 2212 0,0 0.000 0.000-7000.000 7000.000 0.938
#=========================================================================
# 3 !d! 3 1 1,1 0.750 -1.569 32.191 32.238 0.000
# 4 !u~! 3 -2 2,2 -3.047 -19.000 -54.629 57.920 0.000
# 5 !W-! 3 -24 3,4 1.517 -20.68 -20.605 85.925 80.799
# 6 !gamma! 1 22 3,4 -3.813 0.113 -1.833 4.233 0.000
# 7 !d! 1 1 5,5 -2.445 28.816 6.082 29.552 0.010
# 8 !u~! 1 -2 5,5 3.962 -49.498 -26.687 56.373 0.006
# now we build the graph, which will looks like
# p7 #
# p1 / #
# \v1__p3 p5---v4 #
# \_v3_/ \ #
# / \ p8 #
# v2__p4 \ #
# / p6 #
# p2 #
# #
evt = hep.GenEvent(hep.Units.GEV, hep.Units.MM)
evt.event_number = 1
# px py pz e pdgid status
p1 = hep.GenParticle((0.0, 0.0, 7000.0, 7000.0), 2212, 1)
p1.generated_mass = 0.938
p2 = hep.GenParticle((0.0, 0.0, -7000.0, 7000.0), 2212, 2)
p2.generated_mass = 0.938
p3 = hep.GenParticle((0.750, -1.569, 32.191, 32.238), 1, 3)
p3.generated_mass = 0
p4 = hep.GenParticle((-3.047, -19.0, -54.629, 57.920), -2, 4)
p4.generated_mass = 0
p5 = hep.GenParticle((1.517, -20.68, -20.605, 85.925), -24, 5)
p5.generated_mass = 80.799
p6 = hep.GenParticle((-3.813, 0.113, -1.833, 4.233), 22, 6)
p6.generated_mass = 0
p7 = hep.GenParticle((-2.445, 28.816, 6.082, 29.552), 1, 7)
p7.generated_mass = 0.01
p8 = hep.GenParticle((3.962, -49.498, -26.687, 56.373), -2, 8)
p8.generated_mass = 0.006
evt.add_particle(p1)
evt.add_particle(p2)
evt.add_particle(p3)
evt.add_particle(p4)
evt.add_particle(p5)
evt.add_particle(p6)
evt.add_particle(p7)
evt.add_particle(p8)
# make sure vertex is not optimized away by WriterAscii
v1 = hep.GenVertex((1.0, 1.0, 1.0, 1.0))
v1.add_particle_in(p1)
v1.add_particle_out(p3)
evt.add_vertex(v1)
# make sure vertex is not optimized away by WriterAscii
v2 = hep.GenVertex((2.0, 2.0, 2.0, 2.0))
v2.add_particle_in(p2)
v2.add_particle_out(p4)
evt.add_vertex(v2)
# make sure vertex is not optimized away by WriterAscii
v3 = hep.GenVertex((3.0, 3.0, 3.0, 3.0))
v3.add_particle_in(p3)
v3.add_particle_in(p4)
v3.add_particle_out(p5)
v3.add_particle_out(p6)
evt.add_vertex(v3)
# make sure vertex is not optimized away by WriterAscii
v4 = hep.GenVertex((4.0, 4.0, 4.0, 4.0))
v4.add_particle_in(p5)
v4.add_particle_out(p7)
v4.add_particle_out(p8)
evt.add_vertex(v4)
return evt
def main():
parser = argparse.ArgumentParser(description='jetscape in python', \
prog=os.path.basename(__file__))
parser.add_argument('-i', '--input', help='input file', \
default='low', type=str, required=True)
parser.add_argument('--nev', help='number of events', \
default=1000, type=int)
args = parser.parse_args()
# Use pyhepmc_ng to parse the HepMC file
input_hepmc = pyhepmc_ng.ReaderAscii(args.input)
if input_hepmc.failed():
print("[error] unable to read from {}".format(args.input))
sys.exit(1)
# Create a histogram with ROOT
lbins = logbins(1., 500, 50)
hJetPt04 = ROOT.TH1D("hJetPt04", "hJetPt04", 50, lbins)
# jet finder
fj.ClusterSequence.print_banner()
print()
jet_R0 = 0.4
jet_def = fj.JetDefinition(fj.antikt_algorithm, jet_R0)
jet_selector = fj.SelectorPtMin(50.0) & fj.SelectorPtMax(
200.0) & fj.SelectorAbsEtaMax(3)
# Loop through events
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():
nstop = pbar.n
pbar.close()
print('End of HepMC file at event {} '.format(nstop))
break
jets_hepmc = find_jets_hepmc(jet_def, jet_selector, event_hepmc)
all_jets.extend(jets_hepmc)
pbar.update()
# Fill histogram
[fill_jet_histogram(hJetPt04, jet) for jet in all_jets]
if pbar.n >= args.nev:
pbar.close()
print('{} event limit reached'.format(args.nev))
break
# Plot and save histogram
print('Creating ROOT file...')
c = ROOT.TCanvas('c', 'c', 600, 450)
c.cd()
c.SetLogy()
hJetPt04.SetMarkerStyle(21)
hJetPt04.Sumw2()
hJetPt04.Draw('E P')
output_filename = './AnalysisResult.root'
c.SaveAs(output_filename)
fout = ROOT.TFile(output_filename, "update")
fout.cd()
hJetPt04.Write()
fout.Close()
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 test_read_empty_stream():
oss = hep.stringstream()
with hep.ReaderAscii(oss) as f:
evt = hep.GenEvent()
ok = f.read_event(evt)
assert ok == True # reading empty stream is ok in HepMC
def main():
parser = argparse.ArgumentParser(description='jetscape in python', \
prog=os.path.basename(__file__))
parser.add_argument('-i', '--input', help='input file', \
default='low', type=str, required=True)
parser.add_argument('--nev', help='number of events', \
default=1000, type=int)
args = parser.parse_args()
plot_hadrons = False
plot_final_state_partons = True
plot_nth_partons = False
include_thermal_background = False
# Initialize histogram dictionary
hDict = initializeHistograms()
# Use pyhepmc_ng to parse the HepMC file
input_hepmc = pyhepmc_ng.ReaderAscii(args.input)
if input_hepmc.failed():
print("[error] unable to read from {}".format(args.input))
sys.exit(1)
# jet finder
fj.ClusterSequence.print_banner()
print()
jetR = 0.4
jet_def = fj.JetDefinition(fj.antikt_algorithm, jetR)
jet_selector = fj.SelectorPtMin(50.0) & fj.SelectorAbsEtaMax(2)
# Jet re-clustering definition, for primary Lund plane
jet_def_lund = fj.JetDefinition(fj.cambridge_algorithm, jetR)
lund_gen = fjcontrib.LundGenerator(jet_def_lund)
# Load thermal background: array of dataframes, one per event
if include_thermal_background:
thermal_background_array = get_thermal_background()
# Loop through events
all_jets_hadron = []
all_jets_parton = []
kt_shower_list = []
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():
nstop = pbar.n
pbar.close()
print('End of HepMC file at event {} '.format(nstop))
break
if plot_hadrons:
hadrons = get_hadrons(event_hepmc)
jets_hadron = find_jets(jet_def, jet_selector, hadrons)
all_jets_hadron.extend(jets_hadron)
if plot_final_state_partons:
parton_list = get_final_partons(event_hepmc, hDict)
elif plot_nth_partons:
# Get the first parton in the shower (it is apparently the child of the initiating parton)
first_parton = []
first_parton.append(
get_first_parton(event_hepmc, hDict).children[0])
#print('-----------------------------')
#print('First parton E: {}'.format(first_parton[0].momentum.e))
n = 3
parton_list = get_nth_partons(first_parton, n)
#print('Total number of partons: {}'.format(len(parton_list)))
fj_particles = []
for parton in parton_list:
fj_particles.append(
fj.PseudoJet(parton.momentum.px, parton.momentum.py,
parton.momentum.pz, parton.momentum.e))
# Append thermal background
if include_thermal_background:
thermal_background_df = thermal_background_array[0]
thermal_background_fjparticles = get_fjparticles(
thermal_background_df)
fj_particles.extend(thermal_background_fjparticles)
jets_parton = find_jets(jet_def, jet_selector, fj_particles)
all_jets_parton.extend(jets_parton)
# Analyze shower history
max_kt = get_max_kt_shower(event_hepmc, hDict)
kt_shower_list.append(max_kt)
pbar.update()
if pbar.n >= args.nev:
pbar.close()
print('{} event limit reached'.format(args.nev))
break
print('Constructing histograms...')
if plot_hadrons:
n_jets_hadron = len(all_jets_hadron)
print('n_jets_hadron: {}'.format(n_jets_hadron))
n_jets_parton = len(all_jets_parton)
print('n_jets_parton: {}'.format(n_jets_parton))
# Fill histogram
if plot_hadrons:
[fill_jet_histogram(hDict, jet) for jet in all_jets_hadron]
# Fill Lund diagram
# Note: l in lunds, l is the list of splittings in a given jet (following hardest splitting)
lunds_hadron = [lund_gen.result(jet) for jet in all_jets_hadron]
[
fill_lund_histogram(hDict, "hLundHadron", splitting_list)
for splitting_list in lunds_hadron
]
hDict['hLundHadron'].Scale(1. / n_jets_hadron, "width")
lunds_parton = [lund_gen.result(jet) for jet in all_jets_parton]
[
fill_lund_histogram(hDict, "hLundParton", splitting_list)
for splitting_list in lunds_parton
]
hDict['hLundParton'].Scale(1. / n_jets_parton, "width")
for splitting_list in lunds_parton:
max_kt_recluster = 0
for split in splitting_list:
kt = split.kt()
if kt > max_kt_recluster:
max_kt_recluster = kt
hDict['hMaxKT_2D'].Fill(np.log(max_kt_recluster), np.log(kt))
plot_histograms(hDict)
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()
