inv_mass_calc.py

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)

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