def preprocessing(jet):
jet = copy.deepcopy(jet)
jet=jet.reshape(-1,4)
n_consti=len(jet)
# find the jet (eta, phi)
center=jet.sum(axis=0)
v_jet=LorentzVector(center[1], center[2], center[3], center[0])
# centering parameters
phi=v_jet.phi()
bv = v_jet.boost_vector()
bv.set_perp(0)
for i in range(n_consti):
v = LorentzVector(jet[i,1], jet[i,2], jet[i,3], jet[i,0])
v.rotate_z(-phi)
v.boost(-bv)
jet[i, 0]=v[3] #e
jet[i, 1]=v[0] #px
jet[i, 2]=v[1] #py
jet[i, 3]=v[2] #pz
# rotating parameters
weighted_phi=0
weighted_eta=0
for i in range(n_consti):
if jet[i,0]<1e-10: # pass zero paddings
continue
v = LorentzVector(jet[i,1], jet[i,2], jet[i,3], jet[i,0])
r=np.sqrt(v.phi()**2 + v.eta()**2)
if r == 0: # in case there is only one component. In fact these data points should generally be invalid.
continue
weighted_phi += v.phi() * v.E()/r
weighted_eta += v.eta() * v.E()/r
#alpha = np.arctan2(weighted_phi, weighted_eta) # approximately align at eta
alpha = np.arctan2(weighted_eta, weighted_phi) # approximately align at phi
for i in range(n_consti):
v = LorentzVector(jet[i,1], jet[i,2], jet[i,3], jet[i,0])
#v.rotate_x(alpha) # approximately align at eta
v.rotate_x(-alpha) # approximately align at phi
jet[i, 0]=v[3]
jet[i, 1]=v[0]
jet[i, 2]=v[1]
jet[i, 3]=v[2]
#jet=jet.reshape(1,-1)
jet=jet.ravel()
return jet
def preprocessing(
jet
): # every entry would be a sequence of 4-vecs (E, px, py, pz) of jet constituents
jet = copy.deepcopy(jet)
jet = jet.reshape(-1, 4)
n_consti = len(jet)
# find the jet (eta, phi)
center = jet.sum(axis=0)
v_jet = LorentzVector(center[1], center[2], center[3], center[0])
# centering
phi = v_jet.phi()
bv = v_jet.boost_vector()
bv.set_perp(0)
# rotating
weighted_phi = 0
weighted_eta = 0
for i in range(n_consti):
if jet[i, 0] < 1e-10:
continue
v = LorentzVector(jet[i, 1], jet[i, 2], jet[i, 3], jet[i, 0])
r = np.sqrt(v.phi()**2 + v.eta()**2)
weighted_phi += v.phi() * v.E() / r
weighted_eta += v.eta() * v.E() / r
alpha = -np.arctan2(weighted_phi, weighted_eta)
for i in range(n_consti):
v = LorentzVector(jet[i, 1], jet[i, 2], jet[i, 3], jet[i, 0])
v.rotate_z(-phi)
v.boost(-bv)
v.rotate_x(alpha)
jet[i, 0] = v[3]
jet[i, 1] = v[0]
jet[i, 2] = v[1]
jet[i, 3] = v[2]
jet = jet.reshape(1, -1)
return jet
def preprocess(jet, cluster, output="kt", regression=False, R_clustering=0.3):
"""
preprocesses the data to make it usable by the recnn
Preprocessing algorithm:
1. j = the highest pt anti-kt jet (R=1)
2. run kt (R=0.3) on the constituents c of j, resulting in subjets sj1, sj2, ..., sjN
3. phi = sj1.phi(); for all c, do c.rotate_z(-phi)
4. bv = sj1.boost_vector(); bv.set_perp(0); for all c, do c.boost(-bv)
5. deltaz = sj1.pz - sj2.pz; deltay = sj1.py - sj2.py; alpha = -atan2(deltaz, deltay); for all c, do c.rotate_x(alpha)
6. if sj3.pz < 0: for all c, do c.set_pz(-c.pz)
7. finally recluster all transformed constituents c into a single jet
"""
jet = copy.deepcopy(jet)
constituents = jet["content"][jet["tree"][:, 0] == -1]
if regression:
genpt = jet["genpt"]
### run kt (R=0.3) on the constituents c of j, resulting in subjets sj1, sj2, ..., sjN ###
subjets = cluster(constituents, R=R_clustering, jet_algorithm=0)
oldeta = jet["eta"]
oldpt = jet['pt']
### Rot phi ###
# phi = sj1.phi()
# for all c, do c.rotate_z(-phi)
v = subjets[0][1][0]
v = LorentzVector(v)
phi = v.phi()
for _, content, _, _ in subjets:
for i in range(len(content)):
v = LorentzVector(content[i][:4])
v.rotate_z(-phi)
content[i, 0] = v[0]
content[i, 1] = v[1]
content[i, 2] = v[2]
content[i, 3] = v[3]
### boost ###
# bv = sj1.boost_vector()
# bv.set_perp(0)
# for all c, do c.boost(-bv)
v = subjets[0][1][0]
v = LorentzVector(v)
bv = v.boost_vector()
bv.set_perp(0)
for _, content, _, _ in subjets:
for i in range(len(content)):
v = LorentzVector(content[i][:4])
v.boost(-bv)
content[i, 0] = v[0]
content[i, 1] = v[1]
content[i, 2] = v[2]
content[i, 3] = v[3]
### Rot alpha ###
# deltaz = sj1.pz - sj2.pz
# deltay = sj1.py - sj2.py
# alpha = -atan2(deltaz, deltay)
# for all c, do c.rotate_x(alpha)
if len(subjets) >= 2:
deltaz = subjets[0][1][0, 2] - subjets[1][1][0, 2]
deltay = subjets[0][1][0, 1] - subjets[1][1][0, 1]
alpha = -np.arctan2(deltaz, deltay)
for _, content, _, _ in subjets:
for i in range(len(content)):
v = LorentzVector(content[i][:4])
v.rotate_x(alpha)
content[i, 0] = v[0]
content[i, 1] = v[1]
content[i, 2] = v[2]
content[i, 3] = v[3]
### flip if necessary ###
# if sj3.pz < 0: for all c, do c.set_pz(-c.pz)
if len(subjets) >= 3 and subjets[2][1][0, 2] < 0:
for _, content, _, _ in subjets:
for i in range(len(content)):
content[i, 2] *= -1.0
### finally recluster all transformed constituents c into a single jet ###
constituents = []
for tree, content, _, _ in subjets:
constituents.append(content[tree[:, 0] == -1])
constituents = np.vstack(constituents)
if output == "anti-kt":
subjets = cluster(constituents, R=100., jet_algorithm=1)
elif output == "kt":
subjets = cluster(constituents, R=100., jet_algorithm=0)
elif output == "cambridge":
subjets = cluster(constituents, R=100., jet_algorithm=2)
else:
raise
jet["tree"] = subjets[0][0]
jet["content"] = subjets[0][1]
v = LorentzVector(jet["content"][0])
jet["phi"] = v.phi()
jet["eta"] = v.eta()
jet["energy"] = v.E()
jet["mass"] = v.m()
jet["pt"] = v.pt()
jet["root_id"] = 0
jet['oldeta'] = oldeta
jet['oldpt'] = oldpt
if regression:
jet["genpt"] = genpt
return (jet)
