Skip to content

Instantly share code, notes, and snippets.

@benkrikler
Created August 9, 2018 14:49
Show Gist options
  • Save benkrikler/e300722832a9bb1e5a01bb5f87edc822 to your computer and use it in GitHub Desktop.
Save benkrikler/e300722832a9bb1e5a01bb5f87edc822 to your computer and use it in GitHub Desktop.
from math import pi, cos, sin, sinh, sqrt, asinh, atan2, fabs
class LorentzVector(object):
def __init__(self, x, y, z, T, r2=None):
self.x = x
self.y = y
self.z = z
self.T = T
if r2:
self.__r2 = r2
else:
self.__r2 = x**2 + y**2 + z**2
@classmethod
def fromPtEtaPhiM(self, pt, eta, phi, mass):
sinh_eta = sinh(eta)
x = pt * cos(phi)
y = pt * sin(phi)
z = pt * sinh_eta
r2 = pt**2*(1. + sinh_eta**2)
T = sqrt(r2 + mass**2)
return LorentzVector(x, y, z, T, r2)
@property
def r2(self):
return self.__r2
@property
def Pt2(self):
return self.__r2 - self.z**2
@property
def mass(self):
return sqrt(self.T**2 - self.__r2)
@property
def phi(self):
return atan2(self.y, self.x)
@property
def eta(self):
return asinh(self.z/sqrt(self.Pt2))
def __add__(self, rhs):
return LorentzVector(self.x + rhs.x, self.y + rhs.y, self.z + rhs.z, self.T + rhs.T)
def __sub__(self, rhs):
return LorentzVector(self.x - rhs.x, self.y - rhs.y, self.z - rhs.z, self.T - rhs.T)
@benkrikler
Copy link
Author

To calculate the dijet invariant mass:

from inv_mass_calc import LorentzVector
jet1 = LorentzVector.fromPtEtaPhiM(pt1, eta1, phi1, mass1)
jet2 = LorentzVector.fromPtEtaPhiM(pt2, eta2, phi2, mass2)
inv_mass = (jet_1 + jet2).mass

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment