def process(self, event):
constits = []
if self.pseudojets:
for p in event[1:]:
constits.append(fj.PseudoJet(p['px'], p['py'], p['pz'],
p['E']))
jets = self.jet_def(constits)
if (len(jets) > 0):
return jets[0]
return fj.PseudoJet()
else:
for p in event[1:]:
constits.append([p['px'], p['py'], p['pz'], p['E']])
return constits
def pseudojet(particle):
p = [float(item) for item in particle]
e = pow(p[0], 2) + pow(p[1], 2) + pow(p[2], 2) + pow(p[3], 2)
e = sqrt(e)
temp = fj.PseudoJet(p[0], p[1], p[2], e)
temp.set_user_index(int(p[4]))
return temp
def run(data, n_events=1000):
out = []
# Loop over events
for ievt in range(n_events):
# Build a list of all particles
pjs = []
for i in range(data.shape[1]):
pj = fj.PseudoJet()
pj.reset_PtYPhiM(data[ievt, i, 0], data[ievt, i, 1], data[ievt, i,
2], 0)
pjs.append(pj)
# run jet clustering with AntiKt, R=1.0
R = 1.0
jet_def = fj.JetDefinition(fj.antikt_algorithm, R)
# Save the two leading jets
jets = jet_def(pjs)
jets = [j for j in jets if j.pt() > 30.]
out.append([jets[0], jets[1]])
return out
def read_event(file_or_filename):
"""
Routine that can take either an existing opened file object, or a
filename (which it will open itself) and then reads an event from that
file. An event is deemed to end when the file comes to an end or when
the reader encounters the string "#END".
The event is converted to a python list of PseudoJets
"""
# open the file if necessary
if (isinstance(file_or_filename, str)): f = open(file_or_filename, 'r')
else: f = file_or_filename
# create an empty list
event = []
while True:
line = f.readline()
# exit if the file has come to an end
if (not line): break
# or if we reach the string "#END"
if (len(line) >= 4 and line[0:4] == '#END'):
break
# ignore comment lines or empty lines
elif (line[0] == '#' or len(line) <= 1):
continue
# assume we have a good line and split it into px, py, pz, E
p = line.split()
# and append the PseudoJet
event.append(
fj.PseudoJet(float(p[0]), float(p[1]), float(p[2]), float(p[3])))
return event
def find_jets_hepmc(jet_def, jet_selector, hepmc_reader, final=True): fjparts = [] for i, part in enumerate(hepmc_reader.HepMCParticles(final)): psj = fj.PseudoJet(part.momentum().px(), part.momentum().py(), part.momentum().pz(), part.momentum().e()) # psj.set_user_index(i) fjparts.append(psj) jets = jet_selector(jet_def(fjparts)) return jets
def draw_parts_thermal(parts, Rabs=0.4):
jet_R0 = Rabs * 1.0
jet_def = fj.JetDefinition(fj.antikt_algorithm, jet_R0)
parts_v = [fj.PseudoJet(p.px(), p.py(), p.pz(), p.e()) for p in parts]
jets = jet_def(parts_v)
for j in jets:
draw_jet(j, jet_R0)
break
def old_loop(e):
parts = []
aleph_parts = e.get_particles()
# print(len(aleph_parts))
for p in e.get_particles():
psj = fj.PseudoJet(p.px(), p.py(), p.pz(), p.e())
parts.append(psj)
jets = jet_selector(jet_def(parts))
def make_pseudojet(event):
"""Create a fastjet pseudojet from the 4-vector entries"""
event_particles = list()
for n, t in enumerate(event):
# Data Format:(E,px,py,pz)
# FastJet pseudojet format should be: (px,py,pz,E)
temp = fj.PseudoJet(t[1], t[2], t[3], t[0])
event_particles.append(temp)
return event_particles
def read_event(filename):
f = open(filename, 'r')
event = []
for line in f:
p = line.split()
event.append(
fj.PseudoJet(float(p[0]), float(p[1]), float(p[2]), float(p[3])))
return event
def find_jets_hepmc(jet_def, jet_selector, hepmc_event):
fjparts = []
for i, p in enumerate(hepmc_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))
return jets
def parts_from_grid(self):
for ieta in range(1, self.grid_eta_phi.GetNbinsX() + 1):
for iphi in range(1, self.grid_eta_phi.GetNbinsY() + 1):
if self.grid_eta_phi.GetBinContent(ieta, iphi):
psj = fj.PseudoJet()
psj.reset_PtYPhiM(
self.grid_eta_phi.GetBinContent(ieta, iphi),
self.grid_eta_phi.GetXaxis().GetBinCenter(ieta),
self.grid_eta_phi.GetYaxis().GetBinCenter(iphi), 0.)
yield psj
def analyze_event(self, parts):
self.gridR.Reset()
_tmp = [
self.gridR.fillParticle(p.eta(), p.phi(), p.e()) for p in parts
]
self.particles = [
fj.PseudoJet(p.px(), p.py(), p.pz(), p.e())
for p in self.gridR.getGridParticles()
]
JetAnalysis.analyze_event(self, self.particles)
def _groom(self, j, constits, pseudojets):
j1 = fj.PseudoJet()
j2 = fj.PseudoJet()
if j.has_parents(j1, j2):
# order the parents in pt
if (j2.pt() > j1.pt()):
j1, j2 = j2, j1
delta = j1.delta_R(j2)
z = j2.pt() / (j1.pt() + j2.pt())
remove_soft = (z < self.zcut * pow(delta / self.R0, self.beta))
if remove_soft:
self._groom(j1, constits, pseudojets)
else:
self._groom(j1, constits, pseudojets)
self._groom(j2, constits, pseudojets)
else:
if pseudojets:
constits.append(fj.PseudoJet(j.px(), j.py(), j.pz(), j.E()))
else:
constits.append([j.px(), j.py(), j.pz(), j.E()])
def __init__(self, pseudojet, child=None):
"""Initialize a new node, and create its two parents if they exist."""
self.harder = None
self.softer = None
self.lundCoord = None
# first define the current node
self.node = np.array([pseudojet.px(),pseudojet.py(),pseudojet.pz(),pseudojet.E()])
# if it has a direct child (i.e. one level further up in the
# tree), give a link to the corresponding tree object here
self.child = child
j1 = fj.PseudoJet()
j2 = fj.PseudoJet()
if pseudojet and pseudojet.has_parents(j1,j2):
# order the parents in pt
if (j2.pt() > j1.pt()):
j1,j2=j2,j1
# then create two new tree nodes with j1 and j2
self.harder = JetTree(j1, self)
self.softer = JetTree(j2, self)
self.lundCoord = LundCoordinates(j1, j2)
def read_event(file_or_filename):
"""
Routine that can take either an existing opened file object, or a
filename (which it will open itself) and then reads an event from that
file. An event is deemed to end when the file comes to an end or when
the reader encounters the string "#END".
The event is converted to a python list of PseudoJets
"""
# open the file if necessary
if (isinstance(file_or_filename, str)): f = open(file_or_filename, 'r')
else: f = file_or_filename
# create an empty list
event = []
particle_index = 0
subevent_index = -1
while True:
line = f.readline()
# exit if the file has come to an end
if (not line): break
# or if we reach the string "#END"
if (len(line) >= 4 and line[0:4] == '#END'): break
# #SUBSTART means that we start a new subevent
if (len(line) >= 9 and line[0:9] == '#SUBSTART'):
subevent_index += 1
continue
# ignore comment lines or empty lines
elif (line[0] == '#' or len(line) <= 1):
continue
# assume we have a good line and split it into px, py, pz, E
p = line.split()
# create the PseudoJet
pj = fj.PseudoJet(float(p[0]), float(p[1]), float(p[2]), float(p[3]))
# then create its user info
info = ParticleInfo(subevent_index, particle_index)
# optionally with a pdg_id (if it was in the file)
if (len(p) > 4): info.set_pdg_id(int(p[4]))
# finally assign the user info and append the particle the event
pj.set_python_info(info)
event.append(pj)
# don't forget to increment it
particle_index += 1
global event_index
event_index += 1
return event
def __init__(self, pseudojet, child=None):
"""Initialize a new node, and create its two parents if they exist."""
self.harder = None
self.softer = None
self.delta2 = 0.0
self.lundCoord = None
# if it has a direct child (i.e. one level further up in the
# tree), give a link to the corresponding tree object here
self.child = child
while True:
j1 = fj.PseudoJet()
j2 = fj.PseudoJet()
if pseudojet and pseudojet.has_parents(j1, j2):
# order the parents in pt
if (j2.pt() > j1.pt()):
j1, j2 = j2, j1
# check if we satisfy cuts
delta = j1.delta_R(j2)
kt = j2.pt() * delta
if (delta < JetTree.deltamin):
break
# then create two new tree nodes with j1 and j2
if kt >= JetTree.ktmin:
self.harder = JetTree(j1, child=self)
self.softer = JetTree(j2, child=self)
self.delta2 = np.float32(delta * delta)
self.lundCoord = LundCoordinates(j1, j2)
break
else:
pseudojet = j1
else:
break
# finally define the current node
self.node = np.array(
[pseudojet.px(),
pseudojet.py(),
pseudojet.pz(),
pseudojet.E()],
dtype='float32')
def fill_tree_data(jet, tw, sd, rho, iev=None, weight=None, sigma=None):
tw.clear()
if iev:
tw.fill_branch('ev_id', iev)
if weight:
tw.fill_branch('weight', weight)
if sigma:
tw.fill_branch('sigma', sigma)
good_jet = 0.0
if len(jet.constituents()) > 0:
if fj.sorted_by_pt(jet.constituents())[0].pt() < 100.:
good_jet = 1.0
tw.fill_branch('good', good_jet)
sd_jet = sd.result(jet)
sd_info_jet = fjcontrib.get_SD_jet_info(sd_jet)
tw.fill_branch('rho', rho)
tw.fill_branch('j', jet)
tw.fill_branch('j_ptc', jet.pt() - jet.area() * rho)
tw.fill_branch('sd_j', sd_jet)
tw.fill_branch('sd_j_cpt', sd_jet.pt() - sd_jet.area() * rho)
tw.fill_branch('sd_j_z', sd_info_jet.z)
tw.fill_branch('sd_j_dR', sd_info_jet.dR)
pe1 = fj.PseudoJet()
pe2 = fj.PseudoJet()
has_parents = sd_jet.has_parents(pe1, pe2)
tw.fill_branch('j_p1', pe1)
tw.fill_branch('j_p2', pe2)
if has_parents:
tw.fill_branch('j_p1_ptc', pe1.pt() - pe1.area() * rho)
tw.fill_branch('j_p2_ptc', pe2.pt() - pe2.area() * rho)
else:
tw.fill_branch('j_p1_ptc', -1000)
tw.fill_branch('j_p2_ptc', -1000)
tw.fill_tree()
return jet
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 getDummies(ghosts):
res = []
for ghost in ghosts:
eta = ghost.eta()
phi = ghost.phi()
e = 1e-6
px = e * cos(phi) / cosh(eta)
py = e * sin(phi) / cosh(eta)
pz = e * tanh(eta)
res.append(fj.PseudoJet(px, py, pz, e))
return res
def fj_parts_from_tracks(tracks):
fjparts = []
_pts = tracks['ParticlePt']
_etas = tracks['ParticleEta']
_phis = tracks['ParticlePhi']
for index, row in tracks.iterrows():
lv = r.Math.PtEtaPhiMVector(row['ParticlePt'], row['ParticleEta'],
row['ParticlePhi'], 0.0)
psj = fj.PseudoJet(lv.Px(), lv.Py(), lv.Pz(), lv.E())
psj.set_user_index(index)
fjparts.append(psj)
return fjparts
def do_subtraction(self, parts): self.sparts = [] for p in parts: if p.pt() > self.max_pt_subtract: self.sparts.append(p) continue else: fraction = self.funbg.Eval(p.pt()) * self.total_pt if p.pt() - fraction > 0: newp = fj.PseudoJet() newp.reset_PtYPhiM(p.pt() - fraction, p.rap(), p.phi(), p.m()) self.sparts.append(newp)
def subtracted_particles(self, parts): sparts = [] for p in parts: self.tlv.SetPtEtaPhiM(p.pt(), p.eta(), p.phi(), p.m()) ieta = self.grid_eta_phi.GetXaxis().FindBin(p.eta()) iphi = self.grid_eta_phi.GetYaxis().FindBin(p.phi()) spt = self.grid_eta_phi.GetBinContent(ieta, iphi) if spt <= self.tlv.Pt(): self.tlvS.SetPtEtaPhiM(spt, p.eta(), p.phi(), p.m()) self.tlv = self.tlv - self.tlvS sp = fj.PseudoJet(self.tlv.Px(), self.tlv.Py(), self.tlv.Pz(), self.tlv.E()) sparts.append(sp) return sparts
def example():
tw = RTreeWriter()
print(tw)
tw.fill_branch('b', 10)
tw.fill_branch('b', 12.)
tw.fill_branch('bl', [1, 2, 3], do_enumerate=True)
tw.fill_branch('bt', (10, 20, 30.))
psj = fj.PseudoJet()
tw.fill_branch('jet', psj)
tw.fill_branch('jet', psj)
v = fj.vectorPJ()
_v = fj.PseudoJet(1, 2, 3, 4)
v.push_back(_v)
v.push_back(_v)
v.push_back(_v)
tw.fill_branch('jets', v)
tw.fill_branch('bd', {'x': 10, 'y': 20, 'z': 30.})
tw.fill_tree()
tw.write_and_close()
def find_jets_hepmc(jet_def, jet_selector, hepmc_reader, accept_status = []): fjparts = [] for i, part in enumerate(hepmc_reader.HepMCParticles()): if part.status() not in accept_status: continue mom = part.getMomentum() # print(part) # print(mom, mom.px(), mom.py(), mom.pz(), mom.e()) psj = fj.PseudoJet(mom.px(), mom.py(), mom.pz(), mom.e()) # psj.set_user_index(i) fjparts.append(psj) jets = jet_selector(jet_def(fjparts)) return jets
def fillFastJetConstituents(self, hadrons):
# Create a vector of fastjet::PseudoJets from arrays of px,py,pz,e
#fj_particles = fjext.vectorize_px_py_pz_e(px, py, pz, e)
fj_particles = [
fj.PseudoJet(hadron.px, hadron.py, hadron.pz, hadron.E)
for hadron in hadrons
]
for i in range(len(fj_particles)):
fj_particles[i].set_user_index(
int(saveParticleInfo(hadrons[i].pid, hadrons[i].status)))
return fj_particles
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 mergejet(jet1_scin, jet2_scin, jet1_cher, jet2_cher): print "merging clusters" deltar1 = jet1_scin.delta_R(jet1_cher) deltar2 = jet1_scin.delta_R(jet2_cher) c = 0.34 #chi factor if deltar1 < deltar2: jet1Px = (jet1_scin.px()-c*jet1_cher.px())/(1-c) jet1Py = (jet1_scin.py()-c*jet1_cher.py())/(1-c) jet1Pz = (jet1_scin.pz()-c*jet1_cher.pz())/(1-c) jet1E = (jet1_scin.e()-c*jet1_cher.e())/(1.-c) jet1 = fastjet.PseudoJet(jet1Px, jet1Py, jet1Pz, jet1E) jet2Px = (jet2_scin.px()-c*jet2_cher.px())/(1-c) jet2Py = (jet2_scin.py()-c*jet2_cher.py())/(1-c) jet2Pz = (jet2_scin.pz()-c*jet2_cher.pz())/(1-c) jet2E = (jet2_scin.e()-c*jet2_cher.e())/(1.-c) jet2 = fastjet.PseudoJet(jet2Px, jet2Py, jet2Pz, jet2E) else: jet1Px = (jet1_scin.px()-c*jet2_cher.px())/(1-c) jet1Py = (jet1_scin.py()-c*jet2_cher.py())/(1-c) jet1Pz = (jet1_scin.pz()-c*jet2_cher.pz())/(1-c) jet1E = (jet1_scin.e()-c*jet2_cher.e())/(1.-c) jet1 = fastjet.PseudoJet(jet1Px, jet1Py, jet1Pz, jet1E) jet2Px = (jet2_scin.px()-c*jet1_cher.px())/(1-c) jet2Py = (jet2_scin.py()-c*jet1_cher.py())/(1-c) jet2Pz = (jet2_scin.pz()-c*jet1_cher.pz())/(1-c) jet2E = (jet2_scin.e()-c*jet1_cher.e())/(1.-c) jet2 = fastjet.PseudoJet(jet2Px, jet2Py, jet2Pz, jet2E) return jet1, jet2
def get_hadrons(hepmc_event):
fjparts = []
hadrons = []
for vertex in hepmc_event.vertices:
vertex_time = vertex.position.t
if abs(vertex_time - 100) < 1e-3:
hadrons = vertex.particles_out
for hadron in hadrons:
psj = fj.PseudoJet(hadron.momentum.px, hadron.momentum.py,
hadron.momentum.pz, hadron.momentum.e)
fjparts.append(psj)
return fjparts
def find_jets_hepmc(jet_def, jet_selector, hepmc_event):
fjparts = []
hadrons = []
for vertex in hepmc_event.vertices:
vertex_time = vertex.position.t
if abs(vertex_time - 100) < 1e-3:
hadrons = vertex.particles_out
for hadron in hadrons:
psj = fj.PseudoJet(hadron.momentum.px, hadron.momentum.py,
hadron.momentum.pz, hadron.momentum.e)
fjparts.append(psj)
jets = jet_selector(jet_def(fjparts))
return jets
def generate(self, multiplicity=None, offset=0): if multiplicity: self.multiplicity = multiplicity self.particles.clear() for n in range(0, int(self.multiplicity)): _pt = self.funbg.GetRandom(0, self.max_pt) _eta = self.ROOT_random.Rndm() * self.max_eta * 2. - self.max_eta; _phi = self.ROOT_random.Rndm() * ROOT.TMath.Pi() * 2. - ROOT.TMath.Pi(); _p = fj.PseudoJet() _p.reset_PtYPhiM (_pt, _eta, _phi, 0.0) _p.set_user_index(n + offset) self.particles.push_back(_p) self.histogram_pt.Fill(_pt) self.histogram_eta.Fill(_eta) self.histogram_phi.Fill(_phi) self.nEvent = self.nEvent + 1 return self.particles
