def RotateEvent(lep, jets, met_phi, phi):
"""Takes in LorentzVector lep, jets and rotates each along the Z axis by
an angle phi
@==========================================================
@ Parameters
lep: LorentzVector containing lepton information
jets: Array of LorentzVectors containing jet information
phi: Angle between 0 and 2 pi
@==========================================================
@ Return
A rotated LorentzVector
"""
# Missing Azimuthal Energy
met_phi = TVector2.Phi_mpi_pi(met_phi + phi)
# Lepton
lep_new = TLorentzVector(lep)
lep_new.RotateZ(phi)
# Jets
jets_new = []
for j in jets:
jets_new += [TLorentzVector(j)]
j_new = jets_new[-1]
j_new.btag = j.btag
j_new.RotateZ(phi)
return lep_new, jets_new, met_phi
def jet_processing(jet):
# Find the jet (eta, phi)
center=jet.sum(axis=0)
v_jet=TLorentzVector(center[1], center[2], center[3], center[0])
# Centering parameters
phi=v_jet.Phi()
bv = v_jet.BoostVector()
bv.SetPerp(0)
for n in np.arange(len(jet)):
if np.sum(jet[n,:]) != 0:
v = TLorentzVector(jet[n,1], jet[n,2], jet[n,3], jet[n,0])
v.RotateZ(-phi)
v.Boost(-bv)
jet[n,:] = v[3], v[0], v[1], v[2] #(E,Px,Py,Pz)
# Rotating parameters
weighted_phi=0
weighted_eta=0
for n in np.arange(len(jet)):
if np.sum(jet[n,:]) != 0:
v = TLorentzVector(jet[n,1], jet[n,2], jet[n,3], jet[n,0])
r = np.sqrt(v.Phi()**2 + v.Eta()**2)
if r != 0: #in case there is only one component
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 n in np.arange(len(jet)):
if np.sum(jet[n,:]) != 0:
v = TLorentzVector(jet[n,1], jet[n,2], jet[n,3], jet[n,0])
#v.rotate_x(alpha) #approximately align at eta
v.RotateX(-alpha) #approximately align at phi
jet[n,:] = v[3], v[0], v[1], v[2] #(E,Px,Py,Pz)
return jet