예제 #1
0
def lorentz_trans(obt_d):
    """
    Assume that X[ievent] only contains (in order) pt, eta, phi, mass
    """
    #After chacking the functionality of this function, the matchpattern feature should also be removed.
    a0, a1 = obt_d['X'].shape
    X_new = np.zeros((a0, a1))
    for i, ix in zip(range(a0), obt_d['X']):
        for j in range(8):  #8 objects in the game
            tmp = LorentzVector()
            tmp.set_pt_eta_phi_e(ix[5 * j + 0], ix[5 * j + 1], ix[5 * j + 2],
                                 ix[5 * j + 3])
            X_new[i][5 * j + 0] = tmp.px
            X_new[i][5 * j + 1] = tmp.py
            X_new[i][5 * j + 2] = tmp.pz
            X_new[i][5 * j + 3] = tmp.e
            X_new[i][5 * j + 4] = ix[5 * j + 4]

    new_d = {}
    for key, value in obt_d.iteritems():
        new_d[key] = value

    new_d['X'] = X_new

    return new_d
예제 #2
0
def create_tree():

    f = TemporaryFile()
    tree = Tree("tree", model=create_model())
    # fill the tree
    for i in xrange(1000):
        assert_equal(tree.a_vect, LorentzVector(0, 0, 0, 0))
        random_vect = LorentzVector(gauss(.5, 1.), gauss(.5, 1.),
                                    gauss(.5, 1.), gauss(.5, 1.))
        tree.a_vect.copy_from(random_vect)
        assert_equal(tree.a_vect, random_vect)
        tree.a_x = gauss(.5, 1.)
        tree.a_y = gauss(.3, 2.)
        tree.a_z = gauss(13., 42.)
        tree.b_n = randint(1, 5)
        for j in xrange(tree.b_n):
            vect = LorentzVector(gauss(.5, 1.), gauss(.5, 1.), gauss(.5, 1.),
                                 gauss(.5, 1.))
            tree.b_vect.push_back(vect)
            tree.b_x.push_back(randint(1, 10))
            tree.b_y.push_back(gauss(.3, 2.))
        tree.i = i
        assert_equal(tree.b_n, tree.b_vect.size())
        assert_equal(tree.b_n, tree.b_x.size())
        assert_equal(tree.b_n, tree.b_y.size())
        tree.fill(reset=True)
    tree.write()
    # TFile.Close the file but keep the underlying
    # tempfile file descriptor open
    ROOT.TFile.Close(f)
    FILES.append(f)
    FILE_PATHS.append(f.GetName())
예제 #3
0
 def fourvect(self):
     vect = LorentzVector()
     vect.SetPtEtaPhiM(
             self.pt,
             self.eta,
             self.phi,
             self.m)
     #       self._particle.Mass() * GeV)
     return vect
예제 #4
0
 def fourvect_vis(self):
     vect = LorentzVector()
     try:
         vect.SetPtEtaPhiM(et2pt(self.vis_Et, self.vis_eta, self.vis_m),
                           self.eta, self.phi, self.m)
     except ValueError:
         log.warning("DOMAIN ERROR ON TRUTH 4-VECT: "
                     "Et: {0} eta: {1} m: {2}".format(
                         self.vis_Et, self.vis_eta, self.vis_m))
         vect.SetPtEtaPhiM(0, self.eta, self.phi, self.m)
     return vect
예제 #5
0
    def fourvect(self):
        if ((self.nSCTHits + self.nPixHits) < 4):
            # electron with low number of tracker hits
            eta = self.cl_eta
            phi = self.cl_phi
            et  = self.cl_E / math.cosh(self.cl_eta)
        else:
            eta = self.tracketa
            phi = self.trackphi
            et  = self.cl_E / math.cosh(self.tracketa)

        vect = LorentzVector()
        vect.SetPtEtaPhiE(et, eta, phi, self.cl_E)
        return vect
예제 #6
0
def neutrinoPz(lepton_fourVector, neutrino_pt, neutrino_phi):
    """                                                                                                                                          
    Calculate the z-component of the nu momentum by using the W-boson mass as the constraint                                                    
    General idea:                                                                                                                               
    If the discriminant is less than zero, then force it to be zero.                                                                          
    You solve with the discriminant set to zero to get a scaled value                                                                           
    for the term "mu" and "pt". The neutrino Pt will be the                                                                               
    (lepton_pz*scaled_pt)/lepton_pt                                                                                                            
    """
    m_w = 80.4e3 # mass of the W boson                                                                                                          
    delta_phi = lepton_fourVector.Phi() - neutrino_phi

    # Simplifying term you get when solve for neutrino Pz using transverse mass of W boson                                                      
    mu   = (m_w)**2/2 + np.cos(delta_phi)*lepton_fourVector.Pt()*neutrino_pt
    pz_l = lepton_fourVector.Pz() # Lepton Pz                               
    pt_l = lepton_fourVector.Pt() # lepton Pt                                                                                                  
    e_l  = lepton_fourVector.E() # Lepton energy                                                                                                
    p_l  = sqrt(pt_l**2 + pz_l**2)  # Lepton momentum                                                                                         
    
    el_px = lepton_fourVector.Px()
    el_py = lepton_fourVector.Py()
    nu_px = neutrino_pt*np.cos(neutrino_phi)
    nu_py = neutrino_pt*np.sin(neutrino_phi)
    
    if e_l == 0:
        nu = LorentzVector()
        return nu

    discriminant = ((mu**2*pz_l**2)/(e_l**2 - pz_l**2)**2) - ((e_l**2*neutrino_pt**2 - mu**2)/(e_l**2 - pz_l**2))
    if discriminant>0:
        pZ_nu_A = mu*lepton_fourVector.Pz()/(pt_l**2) + sqrt(discriminant)
        pZ_nu_B = mu*lepton_fourVector.Pz()/(pt_l**2) - sqrt(discriminant)

    elif discriminant<0:
        scaled_mu = sqrt(pt_l**2*e_l**2*neutrino_pt**2/(pz_l**2+pt_l**2))
        scaled_pt = m_w**2/(2*pt_l*(1-np.cos(delta_phi)))
        pZ_nu_A = pZ_nu_B = (pz_l*scaled_pt)/pt_l

    elif discriminant==0:
        pZ_nu_A = pZ_nu_B = mu*lepton_fourVector.Pz()/(pt_l**2)

    if abs(pZ_nu_A) < abs(pZ_nu_B):
        nu_pz = pZ_nu_A
    else:
        nu_pz = pZ_nu_B

    nu = LorentzVector()
    nu.SetPxPyPzE(nu_px, nu_py, nu_pz, neutrino_pt)

    return nu
예제 #7
0
파일: dictW.py 프로젝트: afwebb/higgs_diff
def lorentzVecsTop(nom, topIdx0, topIdx1):
    '''
    Takes the indices of two jets identified to be bjets from top decay, return their LorentzVectors 
    '''

    top0 = LorentzVector()                                                                                                   
    top0.SetPtEtaPhiE(nom.jet_pt[topIdx0], nom.jet_eta[topIdx0], nom.jet_phi[topIdx0], nom.jet_e[topIdx0])
    top1 = LorentzVector()                                                                                         
    top1.SetPtEtaPhiE(nom.jet_pt[topIdx1], nom.jet_eta[topIdx1], nom.jet_phi[topIdx1], nom.jet_e[topIdx1])

    return (top0, top1)
예제 #8
0
 def fourvect(self):
     vect = LorentzVector()
     vect.SetPtEtaPhiM(self.pt * GeV, self.eta, self.phi, self.m)
     return vect
예제 #9
0
 def __init__(self):
     self.fourvect_boosted = LorentzVector()
     super(FourMomentum, self).__init__()
예제 #10
0
 def __init__(self):
     self.fourvect_boosted = LorentzVector()
예제 #11
0
### Load data to check ###
### Importing Pyplot ###
import matplotlib.pyplot as plt
from matplotlib.colors import LogNorm
plt.rcParams["figure.figsize"] = (7,6)
fd = f+"anti-kt_test.npy"
X, _ = np.load(fd)

# In[]:
a1 = []
w1=[]
for i,j in enumerate(X):
    constituents = j["content"][j["tree"][:, 0] == -1]
#    if len(constituents)>1:
#        constituents = np.delete(constituents,0,0)
        w1.append([LorentzVector(c).pt() for c in constituents])
w1 = [item for sublist in w1 for item in sublist]

w1=100*np.array(w1)/sum(w1)
a1 = np.vstack(a1)

# In[]:
plt.close()
t=plt.hist2d(a1[:, 0], a1[:, 1], range=[(-0.5,0.5), (-0.5,0.5)], 
           bins=200,  cmap=plt.cm.jet,weights=w1,norm=LogNorm())
cbar = plt.colorbar()
plt.xlabel(r'$\eta$')
plt.ylabel(r'$\varphi$')
cbar.set_label(r'% of p$_t$')
#plt.savefig('tau_pfd_log_bis.png',dpi=600, transparent=True)
plt.show()
예제 #12
0
def mass(tau1, tau2, METpx, METpy):
    """
    Calculate and return the collinear mass and momentum fractions
    of tau1 and tau2

    TODO: set visible mass of taus. 1.2 GeV for 3p and 0.8 GeV for 1p
    """
    recTau1 = LorentzVector()
    recTau2 = LorentzVector()

    # tau 4-vector; synchronize for MMC calculation
    if tau1.nTracks() < 3:
        recTau1.SetPtEtaPhiM(tau1.pt(), tau1.eta(), tau1.phi(), 800.)  # MeV
    else:
        recTau1.SetPtEtaPhiM(tau1.pt(), tau1.eta(), tau1.phi(), 1200.)  # MeV

    if tau2.nTracks() < 3:
        recTau2.SetPtEtaPhiM(tau2.pt(), tau2.eta(), tau2.phi(), 800.)  # MeV
    else:
        recTau2.SetPtEtaPhiM(tau2.pt(), tau2.eta(), tau2.phi(), 1200.)  # MeV

    K = ROOT.TMatrixD(2, 2)
    K[0][0] = recTau1.Px()
    K[0][1] = recTau2.Px()
    K[1][0] = recTau1.Py()
    K[1][1] = recTau2.Py()

    if K.Determinant() == 0:
        return -1., -1111., -1111.

    M = ROOT.TMatrixD(2, 1)
    M[0][0] = METpx
    M[1][0] = METpy

    Kinv = K.Invert()

    X = Kinv * M

    X1 = X(0, 0)
    X2 = X(1, 0)

    x1 = 1. / (1. + X1)
    x2 = 1. / (1. + X2)

    p1 = recTau1 * (1. / x1)
    p2 = recTau2 * (1. / x2)
    m_col = (p1 + p2).M()
    m_vis = (recTau1 + recTau2).M()

    return m_vis, m_col, x1, x2
예제 #13
0
파일: mixins.py 프로젝트: sempersax/tauperf
 def fourvect_clbased(self):
     vect = LorentzVector()
     tau_numTrack = self.numTrack
     tau_nPi0s = self.pi0_n
     if tau_nPi0s == 0:
         if self.track_n > 0:
             sumTrk = LorentzVector()
             for trk_ind in xrange(0, self.track_n):
                 curTrk = LorentzVector()
                 curTrk.SetPtEtaPhiM(self.track_atTJVA_pt[trk_ind],
                                     self.track_atTJVA_eta[trk_ind],
                                     self.track_atTJVA_phi[trk_ind], 139.8)
                 sumTrk += curTrk
             vect.SetPtEtaPhiM(sumTrk.Pt(), sumTrk.Eta(), sumTrk.Phi(),
                               sumTrk.M())
         else:
             vect.SetPtEtaPhiM(self.pt, self.eta, self.phi, self.m)
     elif tau_nPi0s == 1 or tau_nPi0s == 2:
         if self.pi0_vistau_pt == 0:
             vect.SetPtEtaPhiM(self.pt, self.eta, self.phi, self.m)
         else:
             vect.SetPtEtaPhiM(self.pi0_vistau_pt, self.pi0_vistau_eta,
                               self.pi0_vistau_phi, self.pi0_vistau_m)
     else:
         vect.SetPtEtaPhiM(self.pi0_vistau_pt, self.pi0_vistau_eta,
                           self.pi0_vistau_phi, self.pi0_vistau_m)
     return vect
예제 #14
0
def run_top_mass(inputPath):

    topMassesWZ = []
    topMassestZ = []
    topMassReco = []

    f = TFile(inputPath, "READ")
    dsid = inputPath.split('/')[-1]
    dsid = dsid.replace('.root', '')
    #print(inputPath)
    nom = f.Get('nominal')

    try:
        nom.GetEntries()
    except:
        print('failed for ' + inputPath)
        return 0

    try:
        nom.Mll01
    except:
        print('failed for ' + inputPath)
        return 0

    if nom.GetEntries() == 0:
        return 0

    if hasattr(nom, "topMassReco"):
        print('already there', inputPath)
        return 0

    nEntries = nom.GetEntries()
    for idx in range(nEntries):
        if idx % 10000 == 0:
            print(str(idx) + '/' + str(nEntries))

        nom.GetEntry(idx)

        lep = LorentzVector()

        if abs(nom.Mll02 - 91.2e3) < abs(nom.Mll01 - 91.2e3):
            lep.SetPtEtaPhiE(nom.lep_Pt_1, nom.lep_Eta_1, nom.lep_Phi_1,
                             nom.lep_E_1)
        else:
            lep.SetPtEtaPhiE(nom.lep_Pt_2, nom.lep_Eta_2, nom.lep_Phi_2,
                             nom.lep_E_2)

        met = neutrinoPz(lep, nom.met_met, nom.met_phi)

        w = lep + met

        jet = LorentzVector()
        #jet.SetPtEtaPhiE( nom.jet_Pt_0, nom.jet_Eta_0, nom.jet_Phi_0, nom.jet_E_0 )
        if len(nom.jet_pt):
            jet.SetPtEtaPhiE(nom.jet_pt[0], nom.jet_eta[0], nom.jet_phi[0],
                             nom.jet_e[0])
        else:
            jet.SetPtEtaPhiE(0, 0, 0, 0)

        top = LorentzVector()

        top = w + jet

        topMassReco.append(top.M())

    f.Close()

    with root_open(inputPath, mode='a') as myfile:
        topMassReco = np.asarray(topMassReco)
        topMassReco.dtype = [('topMassReco', 'float64')]
        topMassReco.dtype.names = ['topMassReco']
        root_numpy.array2tree(topMassReco, tree=myfile.nominal)
        myfile.write()
        myfile.Close()
def transformVars(df):
    '''
    modifies the variables to create the ones that mv2 uses, inserts default values when needed, saves new variables
    in the dataframe
    Args:
    -----
        df: pandas dataframe containing all the interesting variables as extracted from the .root file
    Returns:
    --------
        modified mv2-compliant dataframe
    '''
    from rootpy.vector import LorentzVector, Vector3
    import pandautils as pup

    # -- modify features and set default values
    df['abs(jet_eta)'] = abs(df['jet_eta'])

    # -- create new IPxD features
    for (pu, pb, pc) in zip(df['jet_ip2d_pu'], df['jet_ip2d_pb'],
                            df['jet_ip2d_pc']):
        pu[np.logical_or(pu >= 10, pu < -1)] = -1
        pb[np.logical_or(pu >= 10, pu < -1)] = -1
        pc[np.logical_or(pu >= 10, pu < -1)] = -1
    for (pu, pb, pc) in zip(df['jet_ip3d_pu'], df['jet_ip3d_pb'],
                            df['jet_ip3d_pc']):
        pu[pu >= 10] = -1
        pb[pu >= 10] = -1
        pc[pu >= 10] = -1
    df['jet_ip2'] = (df['jet_ip2d_pb'] / df['jet_ip2d_pu']).apply(
        lambda x: np.log(x)).apply(lambda x: _replaceInfNaN(x, -20))
    df['jet_ip2_c'] = (df['jet_ip2d_pb'] / df['jet_ip2d_pc']).apply(
        lambda x: np.log(x)).apply(lambda x: _replaceInfNaN(x, -20))
    df['jet_ip2_cu'] = (df['jet_ip2d_pc'] / df['jet_ip2d_pu']).apply(
        lambda x: np.log(x)).apply(lambda x: _replaceInfNaN(x, -20))
    df['jet_ip3'] = (df['jet_ip3d_pb'] / df['jet_ip3d_pu']).apply(
        lambda x: np.log(x)).apply(lambda x: _replaceInfNaN(x, -20))
    df['jet_ip3_c'] = (df['jet_ip3d_pb'] / df['jet_ip3d_pc']).apply(
        lambda x: np.log(x)).apply(lambda x: _replaceInfNaN(x, -20))
    df['jet_ip3_cu'] = (df['jet_ip3d_pc'] / df['jet_ip3d_pu']).apply(
        lambda x: np.log(x)).apply(lambda x: _replaceInfNaN(x, -20))

    # -- create new IPMP features
    for (pu, pb, pc) in zip(df['jet_ipmp_pu'], df['jet_ipmp_pb'],
                            df['jet_ipmp_pc']):
        pu[pu >= 10] = -1
        pb[pu >= 10] = -1
        pc[pu >= 10] = -1
    df['jet_ip'] = (df['jet_ipmp_pb'] / df['jet_ipmp_pu']).apply(
        lambda x: np.log(x)).apply(lambda x: _replaceInfNaN(x, -20))
    df['jet_ip_c'] = (df['jet_ipmp_pb'] / df['jet_ipmp_pc']).apply(
        lambda x: np.log(x)).apply(lambda x: _replaceInfNaN(x, -20))
    df['jet_ip_cu'] = (df['jet_ipmp_pc'] / df['jet_ipmp_pu']).apply(
        lambda x: np.log(x)).apply(lambda x: _replaceInfNaN(x, -20))

    # -- SV1 features
    dx = df['jet_sv1_vtx_x'] - df['PVx']
    dy = df['jet_sv1_vtx_y'] - df['PVy']
    dz = df['jet_sv1_vtx_z'] - df['PVz']

    v_jet = LorentzVector()
    pv2sv = Vector3()
    sv1_L3d = []
    sv1_Lxy = []
    dR = []

    for index, dxi in enumerate(dx):  # loop thru events
        sv1_L3d_ev = []
        sv1L_ev = []
        dR_ev = []
        for jet in xrange(len(dxi)):  # loop thru jets
            v_jet.SetPtEtaPhiM(df['jet_pt'][index][jet],
                               df['jet_eta'][index][jet],
                               df['jet_phi'][index][jet],
                               df['jet_m'][index][jet])
            if (dxi[jet].size != 0):
                sv1_L3d_ev.append(
                    np.sqrt(
                        pow(dx[index][jet], 2) + pow(dy[index][jet], 2) +
                        pow(dz[index][jet], 2))[0])
                sv1L_ev.append(math.hypot(dx[index][jet], dy[index][jet]))

                pv2sv.SetXYZ(dx[index][jet], dy[index][jet], dz[index][jet])
                jetAxis = Vector3(v_jet.Px(), v_jet.Py(), v_jet.Pz())
                dR_ev.append(pv2sv.DeltaR(jetAxis))
            else:
                dR_ev.append(-1)
                sv1L_ev.append(-100)
                sv1_L3d_ev.append(-100)

        sv1_Lxy.append(sv1L_ev)
        dR.append(dR_ev)
        sv1_L3d.append(sv1_L3d_ev)

    df['jet_sv1_dR'] = dR
    df['jet_sv1_Lxy'] = sv1_Lxy
    df['jet_sv1_L3d'] = sv1_L3d

    # -- add more default values for sv1 variables
    sv1_vtx_ok = pup.match_shape(
        np.asarray([len(el) for event in df['jet_sv1_vtx_x'] for el in event]),
        df['jet_pt'])

    for (ok4event, sv1_ntkv4event, sv1_n2t4event, sv1_mass4event,
         sv1_efrc4event,
         sv1_sig34event) in zip(sv1_vtx_ok, df['jet_sv1_ntrkv'],
                                df['jet_sv1_n2t'], df['jet_sv1_m'],
                                df['jet_sv1_efc'], df['jet_sv1_sig3d']):
        sv1_ntkv4event[np.asarray(ok4event) == 0] = -1
        sv1_n2t4event[np.asarray(ok4event) == 0] = -1
        sv1_mass4event[np.asarray(ok4event) == 0] = -1000
        sv1_efrc4event[np.asarray(ok4event) == 0] = -1
        sv1_sig34event[np.asarray(ok4event) == 0] = -100

    # -- JF features
    jf_dR = []
    for eventN, (etas, phis, masses) in enumerate(
            zip(df['jet_jf_deta'], df['jet_jf_dphi'],
                df['jet_jf_m'])):  # loop thru events
        jf_dR_ev = []
        for m in xrange(len(masses)):  # loop thru jets
            if (masses[m] > 0):
                jf_dR_ev.append(np.sqrt(etas[m] * etas[m] + phis[m] * phis[m]))
            else:
                jf_dR_ev.append(-10)
        jf_dR.append(jf_dR_ev)
    df['jet_jf_dR'] = jf_dR

    # -- add more default values for jf variables
    for (jf_mass, jf_n2tv, jf_ntrkv, jf_nvtx, jf_nvtx1t, jf_efrc,
         jf_sig3) in zip(df['jet_jf_m'], df['jet_jf_n2t'],
                         df['jet_jf_ntrkAtVx'], df['jet_jf_nvtx'],
                         df['jet_jf_nvtx1t'], df['jet_jf_efc'],
                         df['jet_jf_sig3d']):
        jf_n2tv[jf_mass <= 0] = -1
        jf_ntrkv[jf_mass <= 0] = -1
        jf_nvtx[jf_mass <= 0] = -1
        jf_nvtx1t[jf_mass <= 0] = -1
        jf_mass[jf_mass <= 0] = -1e3
        jf_efrc[jf_mass <= 0] = -1
        jf_sig3[jf_mass <= 0] = -100

    return df
예제 #16
0
파일: hhskim.py 프로젝트: aashaqshah/hhntup
    def work(self):
        # get argument values
        local = self.args.local
        syst_terms = self.args.syst_terms
        datatype = self.metadata.datatype
        year = self.metadata.year
        verbose = self.args.student_verbose
        very_verbose = self.args.student_very_verbose
        redo_selection = self.args.redo_selection
        nominal_values = self.args.nominal_values

        # get the dataset name
        dsname = os.getenv('INPUT_DATASET_NAME', None)
        if dsname is None:
            # attempt to guess dsname from dirname
            if self.files:
                dsname = os.path.basename(os.path.dirname(self.files[0]))

        # is this a signal sample?
        # if so we will also keep some truth information in the output below
        is_signal = datatype == datasets.MC and (
            '_VBFH' in dsname or '_ggH' in dsname or '_ZH' in dsname
            or '_WH' in dsname or '_ttH' in dsname)
        log.info("DATASET: {0}".format(dsname))
        log.info("IS SIGNAL: {0}".format(is_signal))

        # is this an inclusive signal sample for overlap studies?
        is_inclusive_signal = is_signal and '_inclusive' in dsname

        # is this a BCH-fixed sample? (temporary)
        is_bch_sample = 'r5470_r4540_p1344' in dsname
        if is_bch_sample:
            log.warning("this is a BCH-fixed r5470 sample")

        # onfilechange will contain a list of functions to be called as the
        # chain rolls over to each new file
        onfilechange = []
        count_funcs = {}

        if datatype != datasets.DATA:
            # count the weighted number of events
            if local:

                def mc_weight_count(event):
                    return event.hh_mc_weight
            else:

                def mc_weight_count(event):
                    return event.TruthEvent[0].weights()[0]

            count_funcs = {
                'mc_weight': mc_weight_count,
            }

        if local:
            # local means running on the skims, the output of this script
            # running on the grid
            if datatype == datasets.DATA:
                # merge the GRL fragments
                merged_grl = goodruns.GRL()

                def update_grl(student, grl, name, file, tree):
                    grl |= str(
                        file.Get('Lumi/%s' %
                                 student.metadata.treename).GetString())

                onfilechange.append((update_grl, (
                    self,
                    merged_grl,
                )))

            if datatype == datasets.DATA:
                merged_cutflow = Hist(1, 0, 1, name='cutflow', type='D')
            else:
                merged_cutflow = Hist(2, 0, 2, name='cutflow', type='D')

            def update_cutflow(student, cutflow, name, file, tree):
                # record a cut-flow
                year = student.metadata.year
                datatype = student.metadata.datatype
                cutflow[1].value += file.cutflow_event[1].value
                if datatype != datasets.DATA:
                    cutflow[2].value += file.cutflow_event_mc_weight[1].value

            onfilechange.append((update_cutflow, (
                self,
                merged_cutflow,
            )))

        else:

            # NEED TO BE CONVERTED TO XAOD
            # if datatype not in (datasets.EMBED, datasets.MCEMBED):
            #     # merge TrigConfTrees
            #     metadirname = '%sMeta' % self.metadata.treename
            #     trigconfchain = ROOT.TChain('%s/TrigConfTree' % metadirname)
            #     map(trigconfchain.Add, self.files)
            #     metadir = self.output.mkdir(metadirname)
            #     metadir.cd()
            #     trigconfchain.Merge(self.output, -1, 'fast keep')
            #     self.output.cd()

            if datatype == datasets.DATA:
                # merge GRL XML strings
                merged_grl = goodruns.GRL()
            #     for fname in self.files:
            #         with root_open(fname) as f:
            #             for key in f.Lumi.keys():
            #                 merged_grl |= goodruns.GRL(
            #                     str(key.ReadObj().GetString()),
            #                     from_string=True)
            #     lumi_dir = self.output.mkdir('Lumi')
            #     lumi_dir.cd()
            #     xml_string= ROOT.TObjString(merged_grl.str())
            #     xml_string.Write(self.metadata.treename)
            #     self.output.cd()

        self.output.cd()

        # create the output tree
        model = get_model(datatype,
                          dsname,
                          prefix=None if local else 'hh_',
                          is_inclusive_signal=is_inclusive_signal)
        log.info("Output Model:\n\n{0}\n\n".format(model))
        outtree = Tree(name=self.metadata.treename, model=model)

        if local:
            tree = outtree
        else:
            tree = outtree.define_object(name='tree', prefix='hh_')

        #tree.define_object(name='tau', prefix='tau_')
        tree.define_object(name='tau1', prefix='tau1_')
        tree.define_object(name='tau2', prefix='tau2_')
        tree.define_object(name='truetau1', prefix='truetau1_')
        tree.define_object(name='truetau2', prefix='truetau2_')
        tree.define_object(name='jet1', prefix='jet1_')
        tree.define_object(name='jet2', prefix='jet2_')
        tree.define_object(name='jet3', prefix='jet3_')

        mmc_objects = [
            tree.define_object(name='mmc0', prefix='mmc0_'),
            tree.define_object(name='mmc1', prefix='mmc1_'),
            tree.define_object(name='mmc2', prefix='mmc2_'),
        ]

        for mmc_obj in mmc_objects:
            mmc_obj.define_object(name='resonance', prefix='resonance_')

        # NEED TO BE CONVERTED TO XAOD
        # trigger_emulation = TauTriggerEmulation(
        #     year=year,
        #     passthrough=local or datatype != datasets.MC or year > 2011,
        #     count_funcs=count_funcs)

        # if not trigger_emulation.passthrough:
        #     onfilechange.append(
        #         (update_trigger_trees, (self, trigger_emulation,)))

        # trigger_config = None

        # if datatype not in (datasets.EMBED, datasets.MCEMBED):
        #     # trigger config tool to read trigger info in the ntuples
        #     trigger_config = get_trigger_config()
        #     # update the trigger config maps on every file change
        #     onfilechange.append((update_trigger_config, (trigger_config,)))

        # define the list of event filters
        if local and syst_terms is None and not redo_selection:
            event_filters = None
        else:
            tau_ntrack_recounted_use_ntup = False
            if year > 2011:
                # peek at first tree to determine if the extended number of
                # tracks is already stored
                with root_open(self.files[0]) as test_file:
                    test_tree = test_file.Get(self.metadata.treename)
                    tau_ntrack_recounted_use_ntup = ('tau_out_track_n_extended'
                                                     in test_tree)

            log.info(self.grl)
            event_filters = EventFilterList([
                GRLFilter(self.grl,
                          passthrough=(local
                                       or (datatype not in (datasets.DATA,
                                                            datasets.EMBED))),
                          count_funcs=count_funcs),
                CoreFlags(passthrough=local, count_funcs=count_funcs),
                # NEED TO BE CONVERTED TO XAOD
                # EmbeddingPileupPatch(
                #     passthrough=(
                #         local or year > 2011 or datatype != datasets.EMBED),
                #     count_funcs=count_funcs),
                # NEED TO BE CONVERTED TO XAOD (not a priority)
                # PileupTemplates(
                #     year=year,
                #     passthrough=(
                #         local or is_bch_sample or datatype not in (
                #             datasets.MC, datasets.MCEMBED)),
                #     count_funcs=count_funcs),
                # NEED TO BE CONVERTED TO XAOD
                # RandomSeed(
                #     datatype=datatype,
                #     count_funcs=count_funcs),
                # NEED TO BE CONVERTED TO XAOD
                # BCHSampleRunNumber(
                #     passthrough=not is_bch_sample,
                #     count_funcs=count_funcs),
                # NEED TO BE CONVERTED TO XAOD
                # RandomRunNumber(
                #     tree=tree,
                #     datatype=datatype,
                #     pileup_tool=pileup_tool,
                #     passthrough=local,
                #     count_funcs=count_funcs),
                # NEED TO BE CONVERTED TO XAOD
                # trigger_emulation,
                # NEED TO BE CONVERTED TO XAOD
                # Triggers(
                #     year=year,
                #     tree=tree,
                #     datatype=datatype,
                #     passthrough=datatype in (datasets.EMBED, datasets.MCEMBED),
                #     count_funcs=count_funcs),
                # NEED TO BE CONVERTED TO XAOD
                PileupReweight_xAOD(
                    tree=tree,
                    passthrough=(local
                                 or (datatype
                                     not in (datasets.MC, datasets.MCEMBED))),
                    count_funcs=count_funcs),
                PriVertex(passthrough=local, count_funcs=count_funcs),
                LArError(passthrough=local, count_funcs=count_funcs),
                TileError(passthrough=local, count_funcs=count_funcs),
                TileTrips(passthrough=(local or datatype
                                       in (datasets.MC, datasets.MCEMBED)),
                          count_funcs=count_funcs),
                JetCalibration(datatype=datatype,
                               passthrough=local,
                               count_funcs=count_funcs),
                JetResolution(
                    passthrough=(local
                                 or (datatype
                                     not in (datasets.MC, datasets.MCEMBED))),
                    count_funcs=count_funcs),
                TauCalibration(datatype,
                               passthrough=local,
                               count_funcs=count_funcs),
                # # truth matching must come before systematics due to
                # # TES_TRUE/FAKE
                # NEED TO BE CONVERTED TO XAOD
                TrueTauSelection(passthrough=datatype == datasets.DATA,
                                 count_funcs=count_funcs),
                TruthMatching(passthrough=datatype == datasets.DATA,
                              count_funcs=count_funcs),
                # NEED TO BE CONVERTED TO XAOD
                NvtxJets(tree=tree, count_funcs=count_funcs),
                # # PUT THE SYSTEMATICS "FILTER" BEFORE
                # # ANY FILTERS THAT REFER TO OBJECTS
                # # BUT AFTER CALIBRATIONS
                # # Systematics must also come before anything that refers to
                # # thing.fourvect since fourvect is cached!
                # NEED TO BE CONVERTED TO XAOD
                # Systematics(
                #     terms=syst_terms,
                #     year=year,
                #     datatype=datatype,
                #     tree=tree,
                #     verbose=verbose,
                #     passthrough=not syst_terms,
                #     count_funcs=count_funcs),
                # NEED TO BE CONVERTED TO XAOD
                # JetIsPileup(
                #     passthrough=(
                #         local or year < 2012 or
                #         datatype not in (datasets.MC, datasets.MCEMBED)),
                #     count_funcs=count_funcs),
                JetCleaning(datatype=datatype,
                            year=year,
                            count_funcs=count_funcs),
                ElectronVeto(el_sel='Medium', count_funcs=count_funcs),
                MuonVeto(count_funcs=count_funcs),
                TauPT(2, thresh=20 * GeV, count_funcs=count_funcs),
                TauHasTrack(2, count_funcs=count_funcs),
                TauEta(2, count_funcs=count_funcs),
                TauElectronVeto(2, count_funcs=count_funcs),
                TauMuonVeto(2, count_funcs=count_funcs),
                TauCrack(2, count_funcs=count_funcs),
                # # before selecting the leading and subleading taus
                # # be sure to only consider good candidates
                TauIDMedium(2, count_funcs=count_funcs),
                # NEED TO BE CONVERTED TO XAOD
                # but not used by default
                # #TauTriggerMatchIndex(
                # #    config=trigger_config,
                # #    year=year,
                # #    datatype=datatype,
                # #    passthrough=datatype == datasets.EMBED,
                # #    count_funcs=count_funcs),
                # Select two leading taus at this point
                # 25 and 35 for data
                # 20 and 30 for MC to leave room for TES uncertainty
                TauLeadSublead(lead=(35 * GeV if datatype == datasets.DATA
                                     or local else 30 * GeV),
                               sublead=(25 * GeV if datatype == datasets.DATA
                                        or local else 20 * GeV),
                               count_funcs=count_funcs),
                # taus are sorted (in decreasing order) by pT from here on
                TauIDSelection(count_funcs=count_funcs),
                TaudR(3.2, count_funcs=count_funcs),
                # NEED TO BE CONVERTED TO XAOD
                # but not used by default
                # #TauTriggerMatchThreshold(
                # #    datatype=datatype,
                # #    tree=tree,
                # #    count_funcs=count_funcs),
                # NEED TO BE CONVERTED TO XAOD
                # TauTriggerEfficiency(
                #     year=year,
                #     datatype=datatype,
                #     tree=tree,
                #     tes_systematic=self.args.syst_terms and (
                #         Systematics.TES_TERMS & self.args.syst_terms),
                #     passthrough=datatype == datasets.DATA,
                #     count_funcs=count_funcs),
                # NEED TO BE CONVERTED TO XAOD
                PileupScale(tree=tree,
                            year=year,
                            datatype=datatype,
                            passthrough=local,
                            count_funcs=count_funcs),
                # NEED TO BE CONVERTED TO XAOD
                TauIDScaleFactors(year=year,
                                  passthrough=datatype == datasets.DATA,
                                  count_funcs=count_funcs),
                # NEED TO BE CONVERTED TO XAOD
                # TauFakeRateScaleFactors(
                #     year=year,
                #     datatype=datatype,
                #     tree=tree,
                #     tes_up=(self.args.syst_terms is not None and
                #         (Systematics.TES_FAKE_TOTAL_UP in self.args.syst_terms or
                #          Systematics.TES_FAKE_FINAL_UP in self.args.syst_terms)),
                #     tes_down=(self.args.syst_terms is not None and
                #         (Systematics.TES_FAKE_TOTAL_DOWN in self.args.syst_terms or
                #          Systematics.TES_FAKE_FINAL_DOWN in self.args.syst_terms)),
                #     passthrough=datatype in (datasets.DATA, datasets.EMBED),
                #     count_funcs=count_funcs),
                # NEED TO BE CONVERTED TO XAOD
                HiggsPT(year=year,
                        tree=tree,
                        passthrough=not is_signal or local,
                        count_funcs=count_funcs),
                # NEED TO BE CONVERTED TO XAOD
                # TauTrackRecounting(
                #     year=year,
                #     use_ntup_value=tau_ntrack_recounted_use_ntup,
                #     passthrough=local,
                #     count_funcs=count_funcs),
                # NEED TO BE CONVERTED TO XAOD
                # MCWeight(
                #     datatype=datatype,
                #     tree=tree,
                #     passthrough=local or datatype == datasets.DATA,
                #     count_funcs=count_funcs),
                # NEED TO BE CONVERTED TO XAOD
                # EmbeddingIsolation(
                #     tree=tree,
                #     passthrough=(
                #         local or year < 2012 or
                #         datatype not in (datasets.EMBED, datasets.MCEMBED)),
                #     count_funcs=count_funcs),
                # NEED TO BE CONVERTED TO XAOD
                # EmbeddingCorrections(
                #     tree=tree,
                #     year=year,
                #     passthrough=(
                #         local or
                #         datatype not in (datasets.EMBED, datasets.MCEMBED)),
                #     count_funcs=count_funcs),
                # NEED TO BE CONVERTED TO XAOD
                # EmbeddingTauSpinner(
                #     year=year,
                #     tree=tree,
                #     passthrough=(
                #         local or datatype not in (
                #             datasets.EMBED, datasets.MCEMBED)),
                #     count_funcs=count_funcs),
                # # put MET recalculation after tau selection but before tau-jet
                # # overlap removal and jet selection because of the RefAntiTau
                # # MET correction
                # NEED TO BE CONVERTED TO XAOD
                # METRecalculation(
                #     terms=syst_terms,
                #     year=year,
                #     tree=tree,
                #     refantitau=not nominal_values,
                #     verbose=verbose,
                #     very_verbose=very_verbose,
                #     count_funcs=count_funcs),
                TauJetOverlapRemoval(count_funcs=count_funcs),
                JetPreselection(count_funcs=count_funcs),
                NonIsolatedJet(tree=tree, count_funcs=count_funcs),
                JetSelection(year=year, count_funcs=count_funcs),
                RecoJetTrueTauMatching(passthrough=datatype == datasets.DATA
                                       or local,
                                       count_funcs=count_funcs),
                # NEED TO BE CONVERTED TO XAOD
                # BCHCleaning(
                #     tree=tree,
                #     passthrough=year == 2011 or local,
                #     datatype=datatype,
                #     count_funcs=count_funcs),
                # NEED TO BE CONVERTED TO XAOD
                ClassifyInclusiveHiggsSample(
                    tree=tree,
                    passthrough=not is_inclusive_signal,
                    count_funcs=count_funcs),
            ])

            # set the event filters
            self.filters['event'] = event_filters

        hh_buffer = TreeBuffer()
        if local:
            chain = TreeChain(
                self.metadata.treename,
                files=self.files,
                # ignore_branches=ignore_branches,
                events=self.events,
                onfilechange=onfilechange,
                filters=event_filters,
                cache=True,
                cache_size=50000000,
                learn_entries=100)
            buffer = TreeBuffer()
            for name, value in chain._buffer.items():
                if name.startswith('hh_'):
                    hh_buffer[name[3:]] = value
                elif name in copied:
                    buffer[name] = value
            outtree.set_buffer(hh_buffer, create_branches=False, visible=True)
            outtree.set_buffer(buffer, create_branches=True, visible=False)

        else:

            root_chain = ROOT.TChain(self.metadata.treename)
            for f in self.files:
                log.info(f)
                root_chain.Add(f)

            # if len(self.files) != 1:
            #     raise RuntimeError('lenght of files has to be 1 for now (no xAOD chaining available)')
            # self.files = self.files[0]
            # root_chain = ROOT.TFile(self.files)

            chain = xAODTree(root_chain,
                             filters=event_filters,
                             events=self.events)
            define_objects(chain, datatype=datatype)
            outtree.set_buffer(hh_buffer, create_branches=True, visible=False)

            # create the MMC
            mmc = mass.MMC(year=year)

        # report which packages have been loaded
        # externaltools.report()

        self.output.cd()

        # The main event loop
        # the event filters above are automatically run for each event and only
        # the surviving events are looped on
        for event in chain:

            if local and syst_terms is None and not redo_selection:
                outtree.Fill()
                continue

            # sort taus and jets in decreasing order by pT
            event.taus.sort(key=lambda tau: tau.pt(), reverse=True)
            event.jets.sort(key=lambda jet: jet.pt(), reverse=True)

            # tau1 is the leading tau
            # tau2 is the subleading tau
            tau1, tau2 = event.taus
            tau1.fourvect = asrootpy(tau1.p4())
            tau2.fourvect = asrootpy(tau2.p4())

            beta_taus = (tau1.fourvect + tau2.fourvect).BoostVector()
            tau1.fourvect_boosted = LorentzVector()
            tau1.fourvect_boosted.copy_from(tau1.fourvect)
            tau1.fourvect_boosted.Boost(beta_taus * -1)

            tau2.fourvect_boosted = LorentzVector()
            tau2.fourvect_boosted.copy_from(tau2.fourvect)
            tau2.fourvect_boosted.Boost(beta_taus * -1)

            jets = list(event.jets)
            for jet in jets:
                jet.fourvect = asrootpy(jet.p4())

            jet1, jet2, jet3 = None, None, None
            beta = None
            if len(jets) >= 2:
                jet1, jet2 = jets[:2]

                # determine boost of system
                # determine jet CoM frame
                beta = (jet1.fourvect + jet2.fourvect).BoostVector()
                tree.jet_beta.copy_from(beta)

                jet1.fourvect_boosted = LorentzVector()
                jet1.fourvect_boosted.copy_from(jet1.fourvect)
                jet1.fourvect_boosted.Boost(beta * -1)

                jet2.fourvect_boosted = LorentzVector()
                jet2.fourvect_boosted.copy_from(jet2.fourvect)
                jet2.fourvect_boosted.Boost(beta * -1)

                tau1.min_dr_jet = min(tau1.fourvect.DeltaR(jet1.fourvect),
                                      tau1.fourvect.DeltaR(jet2.fourvect))
                tau2.min_dr_jet = min(tau2.fourvect.DeltaR(jet1.fourvect),
                                      tau2.fourvect.DeltaR(jet2.fourvect))

                # tau centrality (degree to which they are between the two jets)
                tau1.centrality = eventshapes.eta_centrality(
                    tau1.fourvect.Eta(), jet1.fourvect.Eta(),
                    jet2.fourvect.Eta())

                tau2.centrality = eventshapes.eta_centrality(
                    tau2.fourvect.Eta(), jet1.fourvect.Eta(),
                    jet2.fourvect.Eta())

                # boosted tau centrality
                tau1.centrality_boosted = eventshapes.eta_centrality(
                    tau1.fourvect_boosted.Eta(), jet1.fourvect_boosted.Eta(),
                    jet2.fourvect_boosted.Eta())

                tau2.centrality_boosted = eventshapes.eta_centrality(
                    tau2.fourvect_boosted.Eta(), jet1.fourvect_boosted.Eta(),
                    jet2.fourvect_boosted.Eta())

                # 3rd leading jet
                if len(jets) >= 3:
                    jet3 = jets[2]
                    jet3.fourvect_boosted = LorentzVector()
                    jet3.fourvect_boosted.copy_from(jet3.fourvect)
                    jet3.fourvect_boosted.Boost(beta * -1)

            elif len(jets) == 1:
                jet1 = jets[0]

                tau1.min_dr_jet = tau1.fourvect.DeltaR(jet1.fourvect)
                tau2.min_dr_jet = tau2.fourvect.DeltaR(jet1.fourvect)

            RecoJetBlock.set(tree, jet1, jet2, jet3, local=local)

            # mass of ditau + leading jet system
            if jet1 is not None:
                tree.mass_tau1_tau2_jet1 = (tau1.fourvect + tau2.fourvect +
                                            jet1.fourvect).M()

            #####################################
            # number of tracks from PV minus taus
            #####################################
            ntrack_pv = 0
            ntrack_nontau_pv = 0
            for vxp in event.vertices:
                # primary vertex
                if vxp.vertexType() == 1:
                    ntrack_pv = vxp.nTrackParticles()
                    ntrack_nontau_pv = ntrack_pv - tau1.nTracks(
                    ) - tau2.nTracks()
                    break
            tree.ntrack_pv = ntrack_pv
            tree.ntrack_nontau_pv = ntrack_nontau_pv

            #########################
            # MET variables
            #########################
            MET = event.MET[0]
            METx = MET.mpx()
            METy = MET.mpy()
            METet = MET.met()
            MET_vect = Vector2(METx, METy)
            MET_4vect = LorentzVector()
            MET_4vect.SetPxPyPzE(METx, METy, 0., METet)
            MET_4vect_boosted = LorentzVector()
            MET_4vect_boosted.copy_from(MET_4vect)
            if beta is not None:
                MET_4vect_boosted.Boost(beta * -1)

            tree.MET_et = METet
            tree.MET_etx = METx
            tree.MET_ety = METy
            tree.MET_phi = MET.phi()
            dPhi_tau1_tau2 = abs(tau1.fourvect.DeltaPhi(tau2.fourvect))
            dPhi_tau1_MET = abs(tau1.fourvect.DeltaPhi(MET_4vect))
            dPhi_tau2_MET = abs(tau2.fourvect.DeltaPhi(MET_4vect))
            tree.dPhi_tau1_tau2 = dPhi_tau1_tau2
            tree.dPhi_tau1_MET = dPhi_tau1_MET
            tree.dPhi_tau2_MET = dPhi_tau2_MET
            tree.dPhi_min_tau_MET = min(dPhi_tau1_MET, dPhi_tau2_MET)
            tree.MET_bisecting = is_MET_bisecting(dPhi_tau1_tau2,
                                                  dPhi_tau1_MET, dPhi_tau2_MET)

            sumET = MET.sumet()
            tree.MET_sumet = sumET
            if sumET != 0:
                tree.MET_sig = ((2. * METet / GeV) /
                                (utils.sign(sumET) * sqrt(abs(sumET / GeV))))
            else:
                tree.MET_sig = -1.

            tree.MET_centrality = eventshapes.phi_centrality(
                tau1.fourvect, tau2.fourvect, MET_vect)
            tree.MET_centrality_boosted = eventshapes.phi_centrality(
                tau1.fourvect_boosted, tau2.fourvect_boosted,
                MET_4vect_boosted)

            tree.number_of_good_vertices = len(event.vertices)

            ##########################
            # Jet and sum pt variables
            ##########################
            tree.numJets = len(event.jets)

            # sum pT with only the two leading jets
            tree.sum_pt = sum([tau1.pt(), tau2.pt()] +
                              [jet.pt() for jet in jets[:2]])

            # sum pT with all selected jets
            tree.sum_pt_full = sum([tau1.pt(), tau2.pt()] +
                                   [jet.pt() for jet in jets])

            # vector sum pT with two leading jets and MET
            tree.vector_sum_pt = sum([tau1.fourvect, tau2.fourvect] +
                                     [jet.fourvect
                                      for jet in jets[:2]] + [MET_4vect]).Pt()

            # vector sum pT with all selected jets and MET
            tree.vector_sum_pt_full = sum([tau1.fourvect, tau2.fourvect] +
                                          [jet.fourvect for jet in jets] +
                                          [MET_4vect]).Pt()

            # resonance pT
            tree.resonance_pt = sum([tau1.fourvect, tau2.fourvect,
                                     MET_4vect]).Pt()

            # #############################
            # # tau <-> vertex association
            # #############################
            tree.tau_same_vertex = (tau1.vertex() == tau2.vertex())

            tau1.vertex_prob = ROOT.TMath.Prob(tau1.vertex().chiSquared(),
                                               int(tau1.vertex().numberDoF()))

            tau2.vertex_prob = ROOT.TMath.Prob(tau2.vertex().chiSquared(),
                                               int(tau2.vertex().numberDoF()))

            # ##########################
            # # MMC Mass
            # ##########################
            mmc_result = mmc.mass(tau1,
                                  tau2,
                                  METx,
                                  METy,
                                  sumET,
                                  njets=len(event.jets))

            for mmc_method, mmc_object in enumerate(mmc_objects):
                mmc_mass, mmc_resonance, mmc_met = mmc_result[mmc_method]
                if verbose:
                    log.info("MMC (method %d): %f" % (mmc_method, mmc_mass))

                mmc_object.mass = mmc_mass
                mmc_object.MET_et = mmc_met.Mod()
                mmc_object.MET_etx = mmc_met.X()
                mmc_object.MET_ety = mmc_met.Y()
                mmc_object.MET_phi = math.pi - mmc_met.Phi()
                if mmc_mass > 0:
                    FourMomentum.set(mmc_object.resonance, mmc_resonance)

            # ############################
            # # collinear and visible mass
            # ############################
            vis_mass, collin_mass, tau1_x, tau2_x = mass.collinearmass(
                tau1, tau2, METx, METy)

            tree.mass_vis_tau1_tau2 = vis_mass
            tree.mass_collinear_tau1_tau2 = collin_mass
            tau1.collinear_momentum_fraction = tau1_x
            tau2.collinear_momentum_fraction = tau2_x

            # # Fill the tau block
            # # This must come after the RecoJetBlock is filled since
            # # that sets the jet_beta for boosting the taus
            RecoTauBlock.set(event, tree, datatype, tau1, tau2, local=local)

            # NEED TO BE CONVERTED TO XAOD
            if datatype != datasets.DATA:
                TrueTauBlock.set(tree, tau1, tau2)
            # fill the output tree
            outtree.Fill(reset=True)

        # externaltools.report()

        # flush any baskets remaining in memory to disk
        self.output.cd()
        outtree.FlushBaskets()
        outtree.Write()

        if local:
            if datatype == datasets.DATA:
                xml_string = ROOT.TObjString(merged_grl.str())
                xml_string.Write('lumi')
            merged_cutflow.Write()
예제 #17
0
def create_dict(nom):
    current = 0

    events = []
    bestScores = []

    nEntries = nom.GetEntries()
    print(nEntries)
    for idx in range(nEntries):
        if idx % 10000 == 0:
            print(str(idx) + '/' + str(nEntries))

        nom.GetEntry(idx)

        higgCand = LorentzVector()

        lep4Vecs = []
        jet4Vecs = []

        btags = []

        met = LorentzVector()
        met.SetPtEtaPhiE(nom.MET_RefFinal_et, 0, nom.MET_RefFinal_phi,
                         nom.MET_RefFinal_et)

        #for i in range(2):
        lep_Pt_0 = nom.lep_Pt_0
        lep_Eta_0 = nom.lep_Eta_0
        lep_Phi_0 = nom.lep_Phi_0
        lep_E_0 = nom.lep_E_0

        lepVec_0 = LorentzVector()
        lepVec_0.SetPtEtaPhiE(lep_Pt_0, lep_Eta_0, lep_Phi_0, lep_E_0)
        lep4Vecs.append(lepVec_0)

        lep_Pt_1 = nom.lep_Pt_1
        lep_Eta_1 = nom.lep_Eta_1
        lep_Phi_1 = nom.lep_Phi_1
        lep_E_1 = nom.lep_E_1

        lepVec_1 = LorentzVector()
        lepVec_1.SetPtEtaPhiE(lep_Pt_1, lep_Eta_1, lep_Phi_1, lep_E_1)
        lep4Vecs.append(lepVec_1)

        for j in range(len(nom.m_pflow_jet_pt)):  #nom.selected_jets'][i]:
            jetVec = LorentzVector()
            jetVec.SetPtEtaPhiM(nom.m_pflow_jet_pt[j], nom.m_pflow_jet_eta[j],
                                nom.m_pflow_jet_phi[j], nom.m_pflow_jet_m[j])
            jet4Vecs.append(jetVec)

            btags.append(nom.m_pflow_jet_flavor_weight_MV2c10[j])

        combos = []
        combosTop = []

        for l in range(len(lep4Vecs)):
            for i in range(len(jet4Vecs) - 1):
                for j in range(i + 1, len(jet4Vecs)):
                    comb = [l, i, j]

                    t = topDict(jet4Vecs[i], jet4Vecs[j], lep4Vecs[0],
                                lep4Vecs[1], met, btags[i], btags[j],
                                nom.m_pflow_jet_jvt[i], nom.m_pflow_jet_jvt[j],
                                nom.m_pflow_jet_numTrk[i],
                                nom.m_pflow_jet_numTrk[j])

                    combosTop.append([t, comb])

        #loop over combinations, score them in the BDT, figure out the best result
        topDF = pd.DataFrame.from_dict([x[0] for x in combosTop])
        topMat = xgb.DMatrix(topDF, feature_names=list(topDF))

        topPred = topModel.predict(topMat)
        topBest = np.argmax(topPred)

        bestTopComb = combosTop[topBest][1]
        topMatches = bestTopComb[1:]

        for l in range(len(lep4Vecs)):
            for i in range(len(jet4Vecs) - 1):
                for j in range(i + 1, len(jet4Vecs)):
                    comb = [l, i, j]

                    if l == 0:
                        k = higgsDict(jet4Vecs[i], jet4Vecs[j], lep4Vecs[l],
                                      met, btags[i], btags[j], lep4Vecs[1],
                                      nom.m_pflow_jet_jvt[i],
                                      nom.m_pflow_jet_jvt[j],
                                      nom.m_pflow_jet_numTrk[i],
                                      nom.m_pflow_jet_numTrk[j])
                    else:
                        k = higgsDict(jet4Vecs[i], jet4Vecs[j], lep4Vecs[l],
                                      met, btags[i], btags[j], lep4Vecs[0],
                                      nom.m_pflow_jet_jvt[i],
                                      nom.m_pflow_jet_jvt[j],
                                      nom.m_pflow_jet_numTrk[i],
                                      nom.m_pflow_jet_numTrk[j])

                    combos.append([k, comb])

        ###Evaluate higgsTop BDT
        df = pd.DataFrame.from_dict([x[0] for x in combos])
        xgbMat = xgb.DMatrix(df, feature_names=list(df))

        pred = xgbModel.predict(xgbMat)
        best = np.argmax(pred)

        bestScores.append(pred[best])

        bestComb = combos[best][1]
        lepMatch = bestComb[0]
        jetMatches = bestComb[1:]

        k = {}
        #k['higgs_pt'] = nom.higgs_pt
        k['comboScore'] = pred[best]
        k['topScore'] = topPred[topBest]

        if lepMatch == 0:
            k['lep_Pt_H'] = nom.lep_Pt_0
            k['lep_Eta_H'] = nom.lep_Eta_0
            phi_0 = nom.lep_Phi_0
            k['lep_E_H'] = nom.lep_E_0

            k['lep_Pt_O'] = nom.lep_Pt_1
            k['lep_Eta_O'] = nom.lep_Eta_1
            k['lep_Phi_O'] = calc_phi(phi_0, nom.lep_Phi_1)
            k['lep_E_O'] = nom.lep_E_1

        elif lepMatch == 1:
            k['lep_Pt_H'] = nom.lep_Pt_1
            k['lep_Eta_H'] = nom.lep_Eta_1
            phi_0 = nom.lep_Phi_1
            k['lep_E_H'] = nom.lep_E_1

            k['lep_Pt_O'] = nom.lep_Pt_0
            k['lep_Eta_O'] = nom.lep_Eta_0
            k['lep_Phi_O'] = calc_phi(phi_0, nom.lep_Phi_0)
            k['lep_E_O'] = nom.lep_E_0

        n = 0
        for i in jetMatches:  #nom.nJets_OR_T):

            k['jet_Pt_h' + str(n)] = nom.m_pflow_jet_pt[i]
            k['jet_Eta_h' + str(n)] = nom.m_pflow_jet_eta[i]
            k['jet_E_h' + str(n)] = jet4Vecs[i].E()  #nom.m_pflow_jet_E[i]
            k['jet_Phi_h' + str(n)] = calc_phi(phi_0, nom.m_pflow_jet_phi[i])
            k['jet_MV2c10_h' +
              str(n)] = nom.m_pflow_jet_flavor_weight_MV2c10[i]

            n += 1

        btags = np.array(btags)

        btags[jetMatches[0]] = 0
        btags[jetMatches[1]] = 0
        bestBtags = np.argpartition(btags, -2)[-2:]

        n = 0
        for i in topMatches:  #bestBtags:#nom.nJets_OR_T):
            k['top_Pt_' + str(n)] = nom.m_pflow_jet_pt[i]
            k['top_Eta_' + str(n)] = nom.m_pflow_jet_eta[i]
            k['top_E_' + str(n)] = jet4Vecs[i].E()  #nom.m_pflow_jet_E[i]
            k['top_Phi_' + str(n)] = calc_phi(phi_0, nom.m_pflow_jet_phi[i])
            k['top_MV2c10_' + str(n)] = nom.m_pflow_jet_flavor_weight_MV2c10[i]

            n += 1

        k['MET'] = nom.MET_RefFinal_et
        k['MET_phi'] = calc_phi(phi_0, nom.MET_RefFinal_phi)

        events.append(k)

    return events
예제 #18
0
def calcTopMass(nom, topMasses):
    current = 0
    for e in nom:
        current += 1
        if current % 10000 == 0:
            print(current)
        #if current==200000:
        #    break

        if e.nJets_OR != 1:
            continue
        if e.nJets_OR_DL1_70 != 1:
            continue
        if abs(e.Mll01 - 91.2e3) > 10e3 and abs(e.Mll02 - 91.2e3) > 10e3:
            continue
        if e.trilep_type == 0: continue

        lep = LorentzVector()

        if abs(e.Mll02 - 91.2e3) < abs(e.Mll01 - 91.2e3):
            lep.SetPtEtaPhiE(e.lep_Pt_1, e.lep_Eta_1, e.lep_Phi_1, e.lep_E_1)
        else:
            lep.SetPtEtaPhiE(e.lep_Pt_2, e.lep_Eta_2, e.lep_Phi_2, e.lep_E_2)

        met = neutrinoPz(lep, e.met_met, e.met_phi)

        w = lep + met

        #w_eta = np.arccosh( abs( np.sqrt(wt.E()**2 - 80e3**2)/ wt.Pt() ) )
        #w = LorentzVector()
        #w.SetPtEtaPhiE(wt.Pt(), w_eta, wt.Phi(), wt.E())
        #print('M', w.M())

        jet = LorentzVector()
        jet.SetPtEtaPhiE(e.jets_Pt_0, e.jets_Eta_0, e.jets_Phi_0, e.jets_E_0)

        top = LorentzVector()

        top = w + jet

        topMasses.append(top.M())

        if top.M() < 0:
            print(top.M(), w.M())
예제 #19
0
    #newNom.GetEntry(current)
    current+=1
    if current%10000==0:
        print(current)
        #if current==200000:
        #    break
        
        #if e.nJets_OR_T!=1: 
        #    continue
        #if e.nJets_OR_T_MV2c10_70!=1:
        #    continue
        #if abs(e.Mll01 - 91.2e3) > 10e3 and abs(e.Mll02 - 91.2e3) > 10e3:
        #    continue
        #if e.trilep_type==0: continue

    lep = LorentzVector()
    
    if abs(e.Mll02 - 91.2e3) < abs(e.Mll01 - 91.2e3):
        lep.SetPtEtaPhiE( e.lep_Pt_1, e.lep_Eta_1, e.lep_Phi_1, e.lep_E_1 )
    else:
        lep.SetPtEtaPhiE( e.lep_Pt_2, e.lep_Eta_2, e.lep_Phi_2, e.lep_E_2 ) 
        
    met = neutrinoPz(lep, e.met_met, e.met_phi)
    
    w = lep+met
    
    jet = LorentzVector()
    jet.SetPtEtaPhiE( e.jets_Pt_0, e.jets_Eta_0, e.jets_Phi_0, e.jets_E_0 )
    
    top = LorentzVector()
    
예제 #20
0
def create_dict(nom):
    current = 0

    events1l = []
    events2l = []
    decayDicts = []
    bestScores = []

    nEntries = nom.GetEntries()
    for idx in range(nEntries):
        if idx % 10000 == 0:
            print(str(idx) + '/' + str(nEntries))
            #if current%100000==0:
            #break
        nom.GetEntry(idx)

        higgCand = LorentzVector()

        lep4Vecs = []
        jet4Vecs = []

        btags = []

        met = LorentzVector()
        met.SetPtEtaPhiE(nom.MET_RefFinal_et, 0, nom.MET_RefFinal_phi,
                         nom.MET_RefFinal_et)

        lepVec_0 = LorentzVector()
        lepVec_0.SetPtEtaPhiE(nom.lep_Pt_0, nom.lep_Eta_0, nom.lep_Phi_0,
                              nom.lep_E_0)
        lep4Vecs.append(lepVec_0)

        lepVec_1 = LorentzVector()
        lepVec_1.SetPtEtaPhiE(nom.lep_Pt_1, nom.lep_Eta_1, nom.lep_Phi_1,
                              nom.lep_E_1)
        lep4Vecs.append(lepVec_1)

        lepVec_2 = LorentzVector()
        lepVec_2.SetPtEtaPhiE(nom.lep_Pt_2, nom.lep_Eta_2, nom.lep_Phi_2,
                              nom.lep_E_2)
        lep4Vecs.append(lepVec_2)

        for j in range(len(nom.m_pflow_jet_pt)):  #nom.selected_jets'][i]:
            jetVec = LorentzVector()
            jetVec.SetPtEtaPhiM(nom.m_pflow_jet_pt[j], nom.m_pflow_jet_eta[j],
                                nom.m_pflow_jet_phi[j], nom.m_pflow_jet_m[j])
            jet4Vecs.append(jetVec)

            btags.append(nom.m_pflow_jet_flavor_weight_MV2c10[j])

        combosTop = []

        for l in range(len(lep4Vecs)):
            for i in range(len(jet4Vecs) - 1):
                for j in range(i + 1, len(jet4Vecs)):
                    comb = [l, i, j]

                    t = topDict(jet4Vecs[i], jet4Vecs[j], lep4Vecs[0],
                                lep4Vecs[1], lep4Vecs[2], met, btags[i],
                                btags[j], nom.m_pflow_jet_jvt[i],
                                nom.m_pflow_jet_jvt[j],
                                nom.m_pflow_jet_numTrk[i],
                                nom.m_pflow_jet_numTrk[j])

                    combosTop.append([t, comb])

        #loop over combinations, score them in the BDT, figure out the best result
        topDF = pd.DataFrame.from_dict([x[0] for x in combosTop])
        topMat = xgb.DMatrix(topDF, feature_names=list(topDF))

        topPred = topModel.predict(topMat)
        topBest = np.argmax(topPred)

        bestTopComb = combosTop[topBest][1]
        topMatches = bestTopComb[1:]

        combos1l = []

        for l in range(1, len(lep4Vecs)):
            for i in range(len(jet4Vecs) - 1):
                for j in range(i + 1, len(jet4Vecs)):
                    comb = [l, i, j]
                    if l == 1:
                        k = higgs1lDict(
                            jet4Vecs[i], jet4Vecs[j], lep4Vecs[l], met,
                            nom.m_pflow_jet_flavor_weight_MV2c10[i],
                            nom.m_pflow_jet_flavor_weight_MV2c10[j],
                            lep4Vecs[0], lep4Vecs[2], nom.m_pflow_jet_jvt[i],
                            nom.m_pflow_jet_jvt[j], nom.m_pflow_jet_numTrk[i],
                            nom.m_pflow_jet_numTrk[j])
                    else:
                        k = higgs1lDict(
                            jet4Vecs[i], jet4Vecs[j], lep4Vecs[l], met,
                            nom.m_pflow_jet_flavor_weight_MV2c10[i],
                            nom.m_pflow_jet_flavor_weight_MV2c10[j],
                            lep4Vecs[0], lep4Vecs[1], nom.m_pflow_jet_jvt[i],
                            nom.m_pflow_jet_jvt[j], nom.m_pflow_jet_numTrk[i],
                            nom.m_pflow_jet_numTrk[j])

                    combos1l.append([k, comb])

        combos2l = []

        possCombs = [[0, 1, 2], [0, 2, 1]]
        for comb in possCombs:
            k = higgs2lDict(lep4Vecs[comb[0]], lep4Vecs[comb[1]],
                            lep4Vecs[comb[2]], met)
            combos2l.append([k, [comb[0], comb[1]]])

        #Run 2l XGB, find best match
        df2l = pd.DataFrame.from_dict([x[0] for x in combos2l])
        xgbMat2l = xgb.DMatrix(df2l, feature_names=list(df2l))

        pred2l = higgs2lModel.predict(xgbMat2l)
        best2l = np.argmax(pred2l)

        bestComb2l = combos2l[best2l][1]
        lepMatch2l = bestComb2l[1]

        #Run 1l XGB, find best match
        df1l = pd.DataFrame.from_dict([x[0] for x in combos1l])
        xgbMat1l = xgb.DMatrix(df1l, feature_names=list(df1l))

        pred1l = higgs1lModel.predict(xgbMat1l)
        best1l = np.argmax(pred1l)

        bestComb1l = combos1l[best1l][1]
        lepMatch1l = bestComb1l[0]
        jetMatches1l = bestComb1l[1:]

        ### Add decay dict

        k = decayDict(lep4Vecs[0], lep4Vecs[1], lep4Vecs[2], met,
                      jet4Vecs[topMatches[0]], jet4Vecs[topMatches[1]])
        k['nJets'] = nom.nJets_OR_T
        k['nJets_MV2c10_70'] = nom.nJets_OR_T_MV2c10_70
        k['higgs2l_score'] = pred2l[best2l]
        k['higgs1l_score'] = pred1l[best1l]
        decayDicts.append(k)

        ### Add 2l pt prediction dict

        q = {}
        q['comboScore'] = pred2l[best2l]

        if lepMatch2l == 1:
            q['lep_Pt_0'] = nom.lep_Pt_0
            q['lep_Eta_0'] = nom.lep_Eta_0
            phi_0 = nom.lep_Phi_0
            q['lep_E_0'] = nom.lep_E_0

            q['lep_Pt_1'] = nom.lep_Pt_1
            q['lep_Eta_1'] = nom.lep_Eta_1
            q['lep_Phi_1'] = calc_phi(phi_0, nom.lep_Phi_1)
            q['lep_E_1'] = nom.lep_E_1

            q['lep_Pt_2'] = nom.lep_Pt_2
            q['lep_Eta_2'] = nom.lep_Eta_2
            q['lep_Phi_2'] = calc_phi(phi_0, nom.lep_Phi_2)
            q['lep_E_2'] = nom.lep_E_2

        elif lepMatch2l == 2:
            q['lep_Pt_0'] = nom.lep_Pt_0
            q['lep_Eta_0'] = nom.lep_Eta_0
            phi_0 = nom.lep_Phi_0
            q['lep_E_0'] = nom.lep_E_0

            q['lep_Pt_1'] = nom.lep_Pt_2
            q['lep_Eta_1'] = nom.lep_Eta_2
            q['lep_Phi_1'] = calc_phi(phi_0, nom.lep_Phi_2)
            q['lep_E_1'] = nom.lep_E_2

            q['lep_Pt_2'] = nom.lep_Pt_1
            q['lep_Eta_2'] = nom.lep_Eta_1
            q['lep_Phi_2'] = calc_phi(phi_0, nom.lep_Phi_1)
            q['lep_E_2'] = nom.lep_E_1

        n = 0
        for i in topMatches:
            q['top_Pt_' + str(n)] = nom.m_pflow_jet_pt[i]
            q['top_Eta_' + str(n)] = nom.m_pflow_jet_eta[i]
            q['top_E_' + str(n)] = jet4Vecs[i].E()  #nom.m_pflow_jet_E[i]
            q['top_Phi_' + str(n)] = calc_phi(phi_0, nom.m_pflow_jet_phi[i])
            q['top_MV2c10_' + str(n)] = nom.m_pflow_jet_flavor_weight_MV2c10[i]

            n += 1

        q['MET'] = nom.MET_RefFinal_et
        q['MET_phi'] = calc_phi(phi_0, nom.MET_RefFinal_phi)

        events2l.append(q)

        ### Add 1l Pt prediction dict

        y = {}
        #y['higgs_pt'] = nom.higgs_pt
        y['comboScore'] = pred1l[best1l]

        if lepMatch1l == 1:
            y['lep_Pt_H'] = nom.lep_Pt_1
            y['lep_Eta_H'] = nom.lep_Eta_1
            phi_0 = nom.lep_Phi_1
            y['lep_E_H'] = nom.lep_E_1

            y['lep_Pt_0'] = nom.lep_Pt_0
            y['lep_Eta_0'] = nom.lep_Eta_0
            y['lep_Phi_0'] = calc_phi(phi_0, nom.lep_Phi_0)
            y['lep_E_0'] = nom.lep_E_0

            y['lep_Pt_1'] = nom.lep_Pt_2
            y['lep_Eta_1'] = nom.lep_Eta_2
            y['lep_Phi_1'] = calc_phi(phi_0, nom.lep_Phi_2)
            y['lep_E_1'] = nom.lep_E_2

        elif lepMatch1l == 2:
            y['lep_Pt_H'] = nom.lep_Pt_2
            y['lep_Eta_H'] = nom.lep_Eta_2
            phi_0 = nom.lep_Phi_2
            y['lep_E_H'] = nom.lep_E_2

            y['lep_Pt_0'] = nom.lep_Pt_0
            y['lep_Eta_0'] = nom.lep_Eta_0
            y['lep_Phi_0'] = calc_phi(phi_0, nom.lep_Phi_0)
            y['lep_E_0'] = nom.lep_E_0

            y['lep_Pt_1'] = nom.lep_Pt_1
            y['lep_Eta_1'] = nom.lep_Eta_1
            y['lep_Phi_1'] = calc_phi(phi_0, nom.lep_Phi_1)
            y['lep_E_1'] = nom.lep_E_1

        n = 0
        for i in jetMatches1l:  #nom.nJets_OR_T):

            y['jet_Pt_h' + str(n)] = nom.m_pflow_jet_pt[i]
            y['jet_Eta_h' + str(n)] = nom.m_pflow_jet_eta[i]
            y['jet_E_h' + str(n)] = jet4Vecs[i].E()  #nom.m_pflow_jet_E[i]
            y['jet_Phi_h' + str(n)] = calc_phi(phi_0, nom.m_pflow_jet_phi[i])
            y['jet_MV2c10_h' +
              str(n)] = nom.m_pflow_jet_flavor_weight_MV2c10[i]

            n += 1

        n = 0
        for i in topMatches:  #bestBtags:#nom.nJets_OR_T):
            y['top_Pt_' + str(n)] = nom.m_pflow_jet_pt[i]
            y['top_Eta_' + str(n)] = nom.m_pflow_jet_eta[i]
            y['top_E_' + str(n)] = jet4Vecs[i].E()  #nom.m_pflow_jet_E[i]
            y['top_Phi_' + str(n)] = calc_phi(phi_0, nom.m_pflow_jet_phi[i])
            y['top_MV2c10_' + str(n)] = nom.m_pflow_jet_flavor_weight_MV2c10[i]

            n += 1

        y['MET'] = nom.MET_RefFinal_et
        y['MET_phi'] = calc_phi(phi_0, nom.MET_RefFinal_phi)

        events1l.append(y)

    return decayDicts, events1l, events2l
예제 #21
0
    def work(self):

        year = self.metadata.year
        verbose = self.args.verbose
        draw_decays = self.args.draw_decays
        args = self.args

        # initialize the TreeChain of all input files
        # only enable branches I need
        chain = TreeChain(self.metadata.treename,
                          files=self.files,
                          branches=[
                              'tau_*',
                              'mc_*',
                              'el_*',
                              'mu_staco_*',
                              'MET_RefFinal_BDTMedium_*',
                              'MET_RefFinal_STVF_*',
                              'EventNumber',
                              'RunNumber',
                              'averageIntPerXing',
                          ],
                          events=self.events,
                          read_branches_on_demand=True,
                          cache=True,
                          verbose=True)

        define_objects(chain, year)

        self.output.cd()

        # this tree will contain info pertaining to true tau decays
        # for possible use in the optimization of a missing mass calculator
        tree = Tree(name="ditaumass", model=DTMEvent)

        tree.define_object(name='resonance', prefix='resonance_')
        tree.define_object(name='radiative', prefix='radiative_')

        truetaus = [
            tree.define_object(name='truetau1', prefix='truetau1_'),
            tree.define_object(name='truetau2', prefix='truetau2_')
        ]

        taus = [
            tree.define_object(name='tau1', prefix='tau1_'),
            tree.define_object(name='tau2', prefix='tau2_')
        ]

        electrons = [
            tree.define_object(name='ele1', prefix='ele1_'),
            tree.define_object(name='ele2', prefix='ele2_')
        ]

        muons = [
            tree.define_object(name='muon1', prefix='muon1_'),
            tree.define_object(name='muon2', prefix='muon2_')
        ]

        # get the Z or Higgs
        if args.higgs:
            resonance_pdgid = 25
        else:
            resonance_pdgid = 23

        if '7TeV' in self.metadata.name:
            collision_energy = 7
        else:
            collision_energy = 8

        for event_index, event in enumerate(chain):

            try:
                tree.reset_branch_values()

                # get the Z or Higgs
                resonance = tautools.get_particles(event,
                                                   resonance_pdgid,
                                                   num_expected=1)

                if not resonance:
                    print "could not find resonance"
                    continue

                # get the resonance just before the decay
                resonance = resonance[0].last_self

                if draw_decays:
                    resonance.export_graphvis('resonance_%d.dot' %
                                              event.EventNumber)

                FourVectModel.set(tree.resonance, resonance)

                # collect decay products (taus and photons)
                tau_decays = []
                mc_photons = []
                for child in resonance.iter_children():
                    if abs(child.pdgId) == pdg.tau_minus:
                        # ignore status 3 taus in 2012 (something strange in the
                        # MC record...)
                        if year == 2012:
                            if child.status == 3:
                                continue
                        tau_decays.append(tautools.TauDecay(child))
                    elif child.pdgId == pdg.gamma:
                        mc_photons.append(child)
                    else:
                        raise TypeError(
                            'unexpected particle after resonance:\n%s' % child)

                # There should be exactly two taus
                if len(tau_decays) != 2:
                    print "found %i tau decays in MC record" % len(tau_decays)
                    for decay in tau_decays:
                        print decay
                    # skip this event
                    continue

                # check for incomplete tau decays
                invalid = False
                for decay in tau_decays:
                    if not decay.valid:
                        print "invalid tau decay:"
                        print decay
                        if draw_decays:
                            decay.init.export_graphvis('decay_invalid_%d.dot' %
                                                       event.EventNumber)
                        invalid = True
                        break
                if invalid:
                    # skip this event
                    continue

                radiative_fourvect = LorentzVector()
                for photon in mc_photons:
                    radiative_fourvect += photon.fourvect

                radiative_fourvect.fourvect = radiative_fourvect
                FourVectModel.set(tree.radiative, radiative_fourvect)
                tree.radiative_ngamma = len(mc_photons)
                tree.radiative_ngamma_5 = len(
                    [ph for ph in mc_photons if ph.pt > 5])
                tree.radiative_ngamma_10 = len(
                    [ph for ph in mc_photons if ph.pt > 10])
                tree.radiative_et_scalarsum = sum([ph.pt
                                                   for ph in mc_photons] + [0])

                all_matched = True
                matched_objects = []

                skip = False
                for i, (decay, truetau, tau, electron, muon) in enumerate(
                        zip(tau_decays, truetaus, taus, electrons, muons)):

                    if draw_decays:
                        decay.init.export_graphvis('decay%d_%d.dot' %
                                                   (i, event.EventNumber))

                    TrueTau.set(truetau, decay, verbose=verbose)

                    # match to reco taus, electrons and muons
                    if decay.hadronic:
                        recotau, dr = closest_reco_object(
                            event.taus, decay.fourvect_visible, dR=0.2)
                        if recotau is not None:
                            matched_objects.append(recotau)
                            recotau.matched = True
                            recotau.matched_dr = dr
                            RecoTau.set(tau, recotau, verbose=verbose)
                        else:
                            all_matched = False
                    elif decay.leptonic_electron:
                        recoele, dr = closest_reco_object(
                            event.electrons, decay.fourvect_visible, dR=0.2)
                        if recoele is not None:
                            matched_objects.append(recoele)
                            recoele.matched = True
                            recoele.matched_dr = dr
                            RecoElectron.set(electron, recoele)
                        else:
                            all_matched = False
                    elif decay.leptonic_muon:
                        recomuon, dr = closest_reco_object(
                            event.muons, decay.fourvect_visible, dR=0.2)
                        if recomuon is not None:
                            matched_objects.append(recomuon)
                            recomuon.matched = True
                            recomuon.matched_dr = dr
                            RecoMuon.set(muon, recomuon)
                        else:
                            all_matched = False
                    else:
                        print "unhandled invalid tau decay:"
                        print decay
                        if not draw_decays:
                            decay.init.export_graphvis('decay%d_%d.dot' %
                                                       (i, event.EventNumber))
                        # skip this event
                        skip = True
                        break
                if skip:
                    # skip this event
                    continue

                # did both decays match a reco object?
                tree.matched = all_matched

                # match collision: decays matched same reco object
                if all_matched:
                    tree.match_collision = (
                        matched_objects[0] == matched_objects[1])

                # MET
                tree.met_x = event.MET.etx
                tree.met_y = event.MET.ety
                tree.met_phi = event.MET.phi
                tree.met = event.MET.et
                tree.sum_et = event.MET.sumet

                # set extra event variables
                tree.channel = event.mc_channel_number
                tree.event = event.EventNumber
                tree.run = event.RunNumber
                tree.mu = event.averageIntPerXing
                tree.collision_energy = collision_energy

                tree.Fill()
            except:
                print "event index: %d" % event_index
                print "event number: %d" % event.EventNumber
                print "file: %s" % chain.file.GetName()
                raise

        self.output.cd()
        tree.FlushBaskets()
        tree.Write()
예제 #22
0
def lorentzVecs(nom, jetIdx0, jetIdx1, is3l):
    '''
    Initialize met, lepton, and jet lorentz vectors       
    Return jet0, jet1, met, lep0, lep1, (lep2 if is3l)
    '''

    met = LorentzVector()
    met.SetPtEtaPhiE(nom.met_met, 0, nom.met_phi, nom.met_met)

    lep0 = LorentzVector()
    lep0.SetPtEtaPhiE(nom.lep_Pt_0, nom.lep_Eta_0, nom.lep_Phi_0, nom.lep_E_0)

    lep1 = LorentzVector()
    lep1.SetPtEtaPhiE(nom.lep_Pt_1, nom.lep_Eta_1, nom.lep_Phi_1, nom.lep_E_1)

    if is3l:
        lep2 = LorentzVector()
        lep2.SetPtEtaPhiE(nom.lep_Pt_2, nom.lep_Eta_2, nom.lep_Phi_2, nom.lep_E_2)

    jet0 = LorentzVector()                                                                                                
    jet0.SetPtEtaPhiE(nom.jet_pt[jetIdx0], nom.jet_eta[jetIdx0], nom.jet_phi[jetIdx0], nom.jet_e[jetIdx0])

    jet1 = LorentzVector()                                                                                                   
    jet1.SetPtEtaPhiE(nom.jet_pt[jetIdx1], nom.jet_eta[jetIdx1], nom.jet_phi[jetIdx1], nom.jet_e[jetIdx1])

    if is3l:
        return (jet0, jet1, met, lep0, lep1, lep2)
    else:
        return (jet0, jet1, met, lep0, lep1)
예제 #23
0
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)
예제 #24
0
# In[]:
### Verification of the formating ###
### Load data to check ###
fd = f+"anti-kt_test.npy"
X, y = np.load(fd)

# In[]:
### Check for signal ###
a1 = []
w1=[]
for i,j in enumerate(X):
    constituents = j["content"][j["tree"][:, 0] == -1]
#    if len(constituents)>1:
#        constituents = np.delete(constituents,0,0)
    if y[i]==1:
        a1.append(np.array([[LorentzVector(c).eta(), 
                            LorentzVector(c).phi()] for c in constituents]))
        w1.append([LorentzVector(c).pt() for c in constituents])
w1 = [item for sublist in w1 for item in sublist]

w1=100*np.array(w1)/sum(w1)
a1 = np.vstack(a1)

# In[]:
plt.close()
t=plt.hist2d(a1[:, 0], a1[:, 1], range=[(-0.5,0.5), (-0.5,0.5)], 
           bins=200,  cmap=plt.cm.jet,weights=w1,norm=LogNorm())
cbar = plt.colorbar()
plt.xlabel(r'$\eta$')
plt.ylabel(r'$\varphi$')
cbar.set_label(r'% of p$_t$')
예제 #25
0
파일: mixins.py 프로젝트: sempersax/tauperf
 def getTruthVis4Vector(self):
     """Get the LorentzVector for the visible truth tau """
     vector = LorentzVector()
     vector.SetPtEtaPhiM(self.vis_Et, self.vis_eta, self.vis_phi,
                         self.vis_m)
     return vector
예제 #26
0
    def work(self):

        year = self.metadata.year
        verbose = self.args.verbose
        draw_decays = self.args.draw_decays
        args = self.args

        # initialize the TreeChain of all input files
        # only enable branches I need
        chain = TreeChain(
                self.metadata.treename,
                files=self.files,
                branches=[
                    'tau_*',
                    'mc_*',
                    'el_*',
                    'mu_staco_*',
                    'MET_RefFinal_BDTMedium_*',
                    'MET_RefFinal_STVF_*',
                    'EventNumber',
                    'RunNumber',
                    'averageIntPerXing',
                    ],
                events=self.events,
                read_branches_on_demand=True,
                cache=True,
                verbose=True)

        define_objects(chain, year)

        self.output.cd()

        # this tree will contain info pertaining to true tau decays
        # for possible use in the optimization of a missing mass calculator
        tree = Tree(name="ditaumass", model=DTMEvent)

        tree.define_object(name='resonance', prefix='resonance_')
        tree.define_object(name='radiative', prefix='radiative_')

        truetaus = [
            tree.define_object(name='truetau1', prefix='truetau1_'),
            tree.define_object(name='truetau2', prefix='truetau2_')]

        taus = [
            tree.define_object(name='tau1', prefix='tau1_'),
            tree.define_object(name='tau2', prefix='tau2_')]

        electrons = [
            tree.define_object(name='ele1', prefix='ele1_'),
            tree.define_object(name='ele2', prefix='ele2_')]

        muons = [
            tree.define_object(name='muon1', prefix='muon1_'),
            tree.define_object(name='muon2', prefix='muon2_')]

        # get the Z or Higgs
        if args.higgs:
            resonance_pdgid = 25
        else:
            resonance_pdgid = 23

        if '7TeV' in self.metadata.name:
            collision_energy = 7
        else:
            collision_energy = 8

        for event_index, event in enumerate(chain):

            try:
                tree.reset_branch_values()

                # get the Z or Higgs
                resonance = tautools.get_particles(event, resonance_pdgid,
                        num_expected=1)

                if not resonance:
                    print "could not find resonance"
                    continue

                # get the resonance just before the decay
                resonance = resonance[0].last_self

                if draw_decays:
                    resonance.export_graphvis('resonance_%d.dot' %
                            event.EventNumber)

                FourVectModel.set(tree.resonance, resonance)

                # collect decay products (taus and photons)
                tau_decays = []
                mc_photons = []
                for child in resonance.iter_children():
                    if abs(child.pdgId) == pdg.tau_minus:
                        # ignore status 3 taus in 2012 (something strange in the
                        # MC record...)
                        if year == 2012:
                            if child.status == 3:
                                continue
                        tau_decays.append(tautools.TauDecay(child))
                    elif child.pdgId == pdg.gamma:
                        mc_photons.append(child)
                    else:
                        raise TypeError(
                                'unexpected particle after resonance:\n%s' %
                                child)

                # There should be exactly two taus
                if len(tau_decays) != 2:
                    print "found %i tau decays in MC record" % len(tau_decays)
                    for decay in tau_decays:
                        print decay
                    # skip this event
                    continue

                # check for incomplete tau decays
                invalid = False
                for decay in tau_decays:
                    if not decay.valid:
                        print "invalid tau decay:"
                        print decay
                        if draw_decays:
                            decay.init.export_graphvis(
                                    'decay_invalid_%d.dot' %
                                    event.EventNumber)
                        invalid = True
                        break
                if invalid:
                    # skip this event
                    continue

                radiative_fourvect = LorentzVector()
                for photon in mc_photons:
                    radiative_fourvect += photon.fourvect

                radiative_fourvect.fourvect = radiative_fourvect
                FourVectModel.set(tree.radiative, radiative_fourvect)
                tree.radiative_ngamma = len(mc_photons)
                tree.radiative_ngamma_5 = len([
                    ph for ph in mc_photons if ph.pt > 5])
                tree.radiative_ngamma_10 = len([
                    ph for ph in mc_photons if ph.pt > 10])
                tree.radiative_et_scalarsum = sum([
                    ph.pt for ph in mc_photons] + [0])

                all_matched = True
                matched_objects = []

                skip = False
                for i, (decay, truetau, tau, electron, muon) in enumerate(zip(
                        tau_decays, truetaus, taus, electrons, muons)):

                    if draw_decays:
                        decay.init.export_graphvis('decay%d_%d.dot' % (
                            i, event.EventNumber))

                    TrueTau.set(truetau, decay, verbose=verbose)

                    # match to reco taus, electrons and muons
                    if decay.hadronic:
                        recotau, dr = closest_reco_object(
                                event.taus, decay.fourvect_visible, dR=0.2)
                        if recotau is not None:
                            matched_objects.append(recotau)
                            recotau.matched = True
                            recotau.matched_dr = dr
                            RecoTau.set(tau, recotau, verbose=verbose)
                        else:
                            all_matched = False
                    elif decay.leptonic_electron:
                        recoele, dr = closest_reco_object(
                                event.electrons, decay.fourvect_visible, dR=0.2)
                        if recoele is not None:
                            matched_objects.append(recoele)
                            recoele.matched = True
                            recoele.matched_dr = dr
                            RecoElectron.set(electron, recoele)
                        else:
                            all_matched = False
                    elif decay.leptonic_muon:
                        recomuon, dr = closest_reco_object(
                                event.muons, decay.fourvect_visible, dR=0.2)
                        if recomuon is not None:
                            matched_objects.append(recomuon)
                            recomuon.matched = True
                            recomuon.matched_dr = dr
                            RecoMuon.set(muon, recomuon)
                        else:
                            all_matched = False
                    else:
                        print "unhandled invalid tau decay:"
                        print decay
                        if not draw_decays:
                            decay.init.export_graphvis('decay%d_%d.dot' % (
                                i, event.EventNumber))
                        # skip this event
                        skip = True
                        break
                if skip:
                    # skip this event
                    continue

                # did both decays match a reco object?
                tree.matched = all_matched

                # match collision: decays matched same reco object
                if all_matched:
                    tree.match_collision = (
                            matched_objects[0] == matched_objects[1])

                # MET
                tree.met_x = event.MET.etx
                tree.met_y = event.MET.ety
                tree.met_phi = event.MET.phi
                tree.met = event.MET.et
                tree.sum_et = event.MET.sumet

                # set extra event variables
                tree.channel = event.mc_channel_number
                tree.event = event.EventNumber
                tree.run = event.RunNumber
                tree.mu = event.averageIntPerXing
                tree.collision_energy = collision_energy

                tree.Fill()
            except:
                print "event index: %d" % event_index
                print "event number: %d" % event.EventNumber
                print "file: %s" % chain.file.GetName()
                raise

        self.output.cd()
        tree.FlushBaskets()
        tree.Write()
예제 #27
0
파일: hhskim.py 프로젝트: qbuat/hhntup
    def work(self):
        # get argument values
        local = self.args.local
        syst_terms = self.args.syst_terms
        datatype = self.metadata.datatype
        year = self.metadata.year
        verbose = self.args.student_verbose
        very_verbose = self.args.student_very_verbose
        redo_selection = self.args.redo_selection
        nominal_values = self.args.nominal_values

        # get the dataset name
        dsname = os.getenv('INPUT_DATASET_NAME', None)
        if dsname is None:
            # attempt to guess dsname from dirname
            if self.files:
                dsname = os.path.basename(os.path.dirname(self.files[0]))

        # is this a signal sample?
        # if so we will also keep some truth information in the output below
        is_signal = datatype == datasets.MC and (
            '_VBFH' in dsname or
            '_ggH' in dsname or
            '_ZH' in dsname or
            '_WH' in dsname or
            '_ttH' in dsname)
        log.info("DATASET: {0}".format(dsname))
        log.info("IS SIGNAL: {0}".format(is_signal))

        # is this an inclusive signal sample for overlap studies?
        is_inclusive_signal = is_signal and '_inclusive' in dsname

        # is this a BCH-fixed sample? (temporary)
        is_bch_sample = 'r5470_r4540_p1344' in dsname
        if is_bch_sample:
            log.warning("this is a BCH-fixed r5470 sample")

        # onfilechange will contain a list of functions to be called as the
        # chain rolls over to each new file
        onfilechange = []
        count_funcs = {}

        if datatype != datasets.DATA:
            # count the weighted number of events
            if local:
                def mc_weight_count(event):
                    return event.hh_mc_weight
            else:
                def mc_weight_count(event):
                    return event.TruthEvent[0].weights()[0]

            count_funcs = {
                'mc_weight': mc_weight_count,
            }

        if local:
            # local means running on the skims, the output of this script
            # running on the grid
            if datatype == datasets.DATA:
                # merge the GRL fragments
                merged_grl = goodruns.GRL()

                def update_grl(student, grl, name, file, tree):
                    grl |= str(file.Get('Lumi/%s' % student.metadata.treename).GetString())

                onfilechange.append((update_grl, (self, merged_grl,)))

            if datatype == datasets.DATA:
                merged_cutflow = Hist(1, 0, 1, name='cutflow', type='D')
            else:
                merged_cutflow = Hist(2, 0, 2, name='cutflow', type='D')

            def update_cutflow(student, cutflow, name, file, tree):
                # record a cut-flow
                year = student.metadata.year
                datatype = student.metadata.datatype
                cutflow[1].value += file.cutflow_event[1].value
                if datatype != datasets.DATA:
                    cutflow[2].value += file.cutflow_event_mc_weight[1].value

            onfilechange.append((update_cutflow, (self, merged_cutflow,)))

        else:

            # NEED TO BE CONVERTED TO XAOD
            # if datatype not in (datasets.EMBED, datasets.MCEMBED):
            #     # merge TrigConfTrees
            #     metadirname = '%sMeta' % self.metadata.treename
            #     trigconfchain = ROOT.TChain('%s/TrigConfTree' % metadirname)
            #     map(trigconfchain.Add, self.files)
            #     metadir = self.output.mkdir(metadirname)
            #     metadir.cd()
            #     trigconfchain.Merge(self.output, -1, 'fast keep')
            #     self.output.cd()

            if datatype == datasets.DATA:
                # merge GRL XML strings
                merged_grl = goodruns.GRL()
            #     for fname in self.files:
            #         with root_open(fname) as f:
            #             for key in f.Lumi.keys():
            #                 merged_grl |= goodruns.GRL(
            #                     str(key.ReadObj().GetString()),
            #                     from_string=True)
            #     lumi_dir = self.output.mkdir('Lumi')
            #     lumi_dir.cd()
            #     xml_string= ROOT.TObjString(merged_grl.str())
            #     xml_string.Write(self.metadata.treename)
            #     self.output.cd()

        self.output.cd()

        # create the output tree
        model = get_model(datatype, dsname,
                          prefix=None if local else 'hh_',
                          is_inclusive_signal=is_inclusive_signal)
        log.info("Output Model:\n\n{0}\n\n".format(model))
        outtree = Tree(name=self.metadata.treename,
                       model=model)

        if local:
            tree = outtree
        else:
            tree = outtree.define_object(name='tree', prefix='hh_')

        #tree.define_object(name='tau', prefix='tau_')
        tree.define_object(name='tau1', prefix='tau1_')
        tree.define_object(name='tau2', prefix='tau2_')
        tree.define_object(name='truetau1', prefix='truetau1_')
        tree.define_object(name='truetau2', prefix='truetau2_')
        tree.define_object(name='jet1', prefix='jet1_')
        tree.define_object(name='jet2', prefix='jet2_')
        tree.define_object(name='jet3', prefix='jet3_')

        mmc_objects = [
            tree.define_object(name='mmc0', prefix='mmc0_'),
            tree.define_object(name='mmc1', prefix='mmc1_'),
            tree.define_object(name='mmc2', prefix='mmc2_'),
        ]

        for mmc_obj in mmc_objects:
            mmc_obj.define_object(name='resonance', prefix='resonance_')

        # NEED TO BE CONVERTED TO XAOD
        # trigger_emulation = TauTriggerEmulation(
        #     year=year,
        #     passthrough=local or datatype != datasets.MC or year > 2011,
        #     count_funcs=count_funcs)

        # if not trigger_emulation.passthrough:
        #     onfilechange.append(
        #         (update_trigger_trees, (self, trigger_emulation,)))

        # trigger_config = None

        # if datatype not in (datasets.EMBED, datasets.MCEMBED):
        #     # trigger config tool to read trigger info in the ntuples
        #     trigger_config = get_trigger_config()
        #     # update the trigger config maps on every file change
        #     onfilechange.append((update_trigger_config, (trigger_config,)))

        # define the list of event filters
        if local and syst_terms is None and not redo_selection:
            event_filters = None
        else:
            tau_ntrack_recounted_use_ntup = False
            if year > 2011:
                # peek at first tree to determine if the extended number of
                # tracks is already stored
                with root_open(self.files[0]) as test_file:
                    test_tree = test_file.Get(self.metadata.treename)
                    tau_ntrack_recounted_use_ntup = (
                        'tau_out_track_n_extended' in test_tree)

            log.info(self.grl)
            event_filters = EventFilterList([
                GRLFilter(
                    self.grl,
                    passthrough=(
                        local or (
                            datatype not in (datasets.DATA, datasets.EMBED))),
                    count_funcs=count_funcs),
                CoreFlags(
                    passthrough=local,
                    count_funcs=count_funcs),
                # NEED TO BE CONVERTED TO XAOD
                # EmbeddingPileupPatch(
                #     passthrough=(
                #         local or year > 2011 or datatype != datasets.EMBED),
                #     count_funcs=count_funcs),
                # NEED TO BE CONVERTED TO XAOD (not a priority)
                # PileupTemplates(
                #     year=year,
                #     passthrough=(
                #         local or is_bch_sample or datatype not in (
                #             datasets.MC, datasets.MCEMBED)),
                #     count_funcs=count_funcs),
                # NEED TO BE CONVERTED TO XAOD
                # RandomSeed(
                #     datatype=datatype,
                #     count_funcs=count_funcs),
                # NEED TO BE CONVERTED TO XAOD
                # BCHSampleRunNumber(
                #     passthrough=not is_bch_sample,
                #     count_funcs=count_funcs),
                # NEED TO BE CONVERTED TO XAOD
                # RandomRunNumber(
                #     tree=tree,
                #     datatype=datatype,
                #     pileup_tool=pileup_tool,
                #     passthrough=local,
                #     count_funcs=count_funcs),
                # NEED TO BE CONVERTED TO XAOD
                # trigger_emulation,
                # NEED TO BE CONVERTED TO XAOD
                # Triggers(
                #     year=year,
                #     tree=tree,
                #     datatype=datatype,
                #     passthrough=datatype in (datasets.EMBED, datasets.MCEMBED),
                #     count_funcs=count_funcs),
                # NEED TO BE CONVERTED TO XAOD
                PileupReweight_xAOD(
                        tree=tree,
                        passthrough=(local or (
                            datatype not in (datasets.MC, datasets.MCEMBED))),
                        count_funcs=count_funcs),
                PriVertex(
                    passthrough=local,
                    count_funcs=count_funcs),
                LArError(
                    passthrough=local,
                    count_funcs=count_funcs),
                TileError(
                    passthrough=local,
                    count_funcs=count_funcs),
                TileTrips(
                    passthrough=(
                        local or datatype in (datasets.MC, datasets.MCEMBED)),
                    count_funcs=count_funcs),
                JetCalibration(
                        datatype=datatype,
                        passthrough=local,
                        count_funcs=count_funcs),
                JetResolution(
                        passthrough=(local or (
                                datatype not in (datasets.MC, datasets.MCEMBED))),
                        count_funcs=count_funcs),
                TauCalibration(
                        datatype,
                        passthrough=local,
                        count_funcs=count_funcs),
                # # truth matching must come before systematics due to
                # # TES_TRUE/FAKE
                # NEED TO BE CONVERTED TO XAOD
                TrueTauSelection(
                        passthrough=datatype == datasets.DATA,
                        count_funcs=count_funcs),
                TruthMatching(
                    passthrough=datatype == datasets.DATA,
                    count_funcs=count_funcs),
                # NEED TO BE CONVERTED TO XAOD
                NvtxJets(
                    tree=tree,
                    count_funcs=count_funcs),
                # # PUT THE SYSTEMATICS "FILTER" BEFORE
                # # ANY FILTERS THAT REFER TO OBJECTS
                # # BUT AFTER CALIBRATIONS
                # # Systematics must also come before anything that refers to
                # # thing.fourvect since fourvect is cached!
                # NEED TO BE CONVERTED TO XAOD
                # Systematics(
                #     terms=syst_terms,
                #     year=year,
                #     datatype=datatype,
                #     tree=tree,
                #     verbose=verbose,
                #     passthrough=not syst_terms,
                #     count_funcs=count_funcs),
                # NEED TO BE CONVERTED TO XAOD
                # JetIsPileup(
                #     passthrough=(
                #         local or year < 2012 or
                #         datatype not in (datasets.MC, datasets.MCEMBED)),
                #     count_funcs=count_funcs),
                JetCleaning(
                    datatype=datatype,
                    year=year,
                    count_funcs=count_funcs),
                ElectronVeto(
                        el_sel='Medium',
                        count_funcs=count_funcs),
                MuonVeto(
                    count_funcs=count_funcs),
                TauPT(2,
                    thresh=20 * GeV,
                    count_funcs=count_funcs),
                TauHasTrack(2,
                    count_funcs=count_funcs),
                TauEta(2,
                    count_funcs=count_funcs),
                TauElectronVeto(2,
                    count_funcs=count_funcs),
                TauMuonVeto(2,
                    count_funcs=count_funcs),
                TauCrack(2,
                    count_funcs=count_funcs),
                # # before selecting the leading and subleading taus
                # # be sure to only consider good candidates
                TauIDMedium(2,
                    count_funcs=count_funcs),
                # NEED TO BE CONVERTED TO XAOD
                # but not used by default
                # #TauTriggerMatchIndex(
                # #    config=trigger_config,
                # #    year=year,
                # #    datatype=datatype,
                # #    passthrough=datatype == datasets.EMBED,
                # #    count_funcs=count_funcs),
                # Select two leading taus at this point
                # 25 and 35 for data
                # 20 and 30 for MC to leave room for TES uncertainty
                TauLeadSublead(
                    lead=(
                        35 * GeV if datatype == datasets.DATA or local
                        else 30 * GeV),
                    sublead=(
                        25 * GeV if datatype == datasets.DATA or local
                        else 20 * GeV),
                    count_funcs=count_funcs),
                # taus are sorted (in decreasing order) by pT from here on
                TauIDSelection(
                    count_funcs=count_funcs),
                TaudR(3.2,
                    count_funcs=count_funcs),
                # NEED TO BE CONVERTED TO XAOD
                # but not used by default
                # #TauTriggerMatchThreshold(
                # #    datatype=datatype,
                # #    tree=tree,
                # #    count_funcs=count_funcs),
                # NEED TO BE CONVERTED TO XAOD
                # TauTriggerEfficiency(
                #     year=year,
                #     datatype=datatype,
                #     tree=tree,
                #     tes_systematic=self.args.syst_terms and (
                #         Systematics.TES_TERMS & self.args.syst_terms),
                #     passthrough=datatype == datasets.DATA,
                #     count_funcs=count_funcs),
                # NEED TO BE CONVERTED TO XAOD
                PileupScale(
                    tree=tree,
                    year=year,
                    datatype=datatype,
                    passthrough=local,
                    count_funcs=count_funcs),
                # NEED TO BE CONVERTED TO XAOD
                TauIDScaleFactors(
                    year=year,
                    passthrough=datatype == datasets.DATA,
                    count_funcs=count_funcs),
                # NEED TO BE CONVERTED TO XAOD
                # TauFakeRateScaleFactors(
                #     year=year,
                #     datatype=datatype,
                #     tree=tree,
                #     tes_up=(self.args.syst_terms is not None and
                #         (Systematics.TES_FAKE_TOTAL_UP in self.args.syst_terms or
                #          Systematics.TES_FAKE_FINAL_UP in self.args.syst_terms)),
                #     tes_down=(self.args.syst_terms is not None and
                #         (Systematics.TES_FAKE_TOTAL_DOWN in self.args.syst_terms or
                #          Systematics.TES_FAKE_FINAL_DOWN in self.args.syst_terms)),
                #     passthrough=datatype in (datasets.DATA, datasets.EMBED),
                #     count_funcs=count_funcs),
                # NEED TO BE CONVERTED TO XAOD
                # HiggsPT(
                #     year=year,
                #     tree=tree,
                #     passthrough=not is_signal or local,
                #     count_funcs=count_funcs),
                # NEED TO BE CONVERTED TO XAOD
                # TauTrackRecounting(
                #     year=year,
                #     use_ntup_value=tau_ntrack_recounted_use_ntup,
                #     passthrough=local,
                #     count_funcs=count_funcs),
                # NEED TO BE CONVERTED TO XAOD
                # MCWeight(
                #     datatype=datatype,
                #     tree=tree,
                #     passthrough=local or datatype == datasets.DATA,
                #     count_funcs=count_funcs),
                # NEED TO BE CONVERTED TO XAOD
                # EmbeddingIsolation(
                #     tree=tree,
                #     passthrough=(
                #         local or year < 2012 or
                #         datatype not in (datasets.EMBED, datasets.MCEMBED)),
                #     count_funcs=count_funcs),
                # NEED TO BE CONVERTED TO XAOD
                # EmbeddingCorrections(
                #     tree=tree,
                #     year=year,
                #     passthrough=(
                #         local or
                #         datatype not in (datasets.EMBED, datasets.MCEMBED)),
                #     count_funcs=count_funcs),
                # NEED TO BE CONVERTED TO XAOD
                # EmbeddingTauSpinner(
                #     year=year,
                #     tree=tree,
                #     passthrough=(
                #         local or datatype not in (
                #             datasets.EMBED, datasets.MCEMBED)),
                #     count_funcs=count_funcs),
                # # put MET recalculation after tau selection but before tau-jet
                # # overlap removal and jet selection because of the RefAntiTau
                # # MET correction
                # NEED TO BE CONVERTED TO XAOD
                # METRecalculation(
                #     terms=syst_terms,
                #     year=year,
                #     tree=tree,
                #     refantitau=not nominal_values,
                #     verbose=verbose,
                #     very_verbose=very_verbose,
                #     count_funcs=count_funcs),
                TauJetOverlapRemoval(
                    count_funcs=count_funcs),
                JetPreselection(
                    count_funcs=count_funcs),
                NonIsolatedJet(
                    tree=tree,
                    count_funcs=count_funcs),
                JetSelection(
                    year=year,
                    count_funcs=count_funcs),
                RecoJetTrueTauMatching(
                    passthrough=datatype == datasets.DATA or local,
                    count_funcs=count_funcs),
                # NEED TO BE CONVERTED TO XAOD
                # BCHCleaning(
                #     tree=tree,
                #     passthrough=year == 2011 or local,
                #     datatype=datatype,
                #     count_funcs=count_funcs),
                # NEED TO BE CONVERTED TO XAOD
                ClassifyInclusiveHiggsSample(
                    tree=tree,
                    passthrough=not is_inclusive_signal,
                    count_funcs=count_funcs),
            ])

            # set the event filters
            self.filters['event'] = event_filters

        hh_buffer = TreeBuffer()
        if local:
            chain = TreeChain(
                self.metadata.treename,
                files=self.files,
                # ignore_branches=ignore_branches,
                events=self.events,
                onfilechange=onfilechange,
                filters=event_filters,
                cache=True,
                cache_size=50000000,
                learn_entries=100)
            buffer = TreeBuffer()
            for name, value in chain._buffer.items():
                if name.startswith('hh_'):
                    hh_buffer[name[3:]] = value
                elif name in copied:
                    buffer[name] = value
            outtree.set_buffer(
                hh_buffer,
                create_branches=False,
                visible=True)
            outtree.set_buffer(
                buffer,
                create_branches=True,
                visible=False)


        else:

            root_chain = ROOT.TChain(self.metadata.treename)
            for f in self.files:
                log.info(f)
                root_chain.Add(f)
            
            # if len(self.files) != 1:
            #     raise RuntimeError('lenght of files has to be 1 for now (no xAOD chaining available)')
            # self.files = self.files[0]
            # root_chain = ROOT.TFile(self.files)

            chain = xAODTree(root_chain, filters=event_filters, events=self.events)
            define_objects(chain, datatype=datatype)
            outtree.set_buffer(
                hh_buffer,
                create_branches=True,
                visible=False)

            # # create the MMC
            # mmc = mass.MMC(year=year)
            from ROOT import MissingMassTool
            mass_tool = MissingMassTool('mass_tool')
            mass_tool.initialize()


        # report which packages have been loaded
        # externaltools.report()

        self.output.cd()

        # The main event loop
        # the event filters above are automatically run for each event and only
        # the surviving events are looped on
        for event in chain:
            if local and syst_terms is None and not redo_selection:
                outtree.Fill()
                continue
            
            # sort taus and jets in decreasing order by pT
            event.taus.sort(key=lambda tau: tau.pt(), reverse=True)
            event.jets.sort(key=lambda jet: jet.pt(), reverse=True)

            # tau1 is the leading tau
            # tau2 is the subleading tau
            tau1, tau2 = event.taus
            tau1.fourvect = asrootpy(tau1.p4())
            tau2.fourvect = asrootpy(tau2.p4())

            beta_taus = (tau1.fourvect + tau2.fourvect).BoostVector()
            tau1.fourvect_boosted = LorentzVector()
            tau1.fourvect_boosted.copy_from(tau1.fourvect)
            tau1.fourvect_boosted.Boost(beta_taus * -1)
            
            tau2.fourvect_boosted = LorentzVector()
            tau2.fourvect_boosted.copy_from(tau2.fourvect)
            tau2.fourvect_boosted.Boost(beta_taus * -1)

            jets = list(event.jets)
            for jet in jets:
                jet.fourvect = asrootpy(jet.p4())

            jet1, jet2, jet3 = None, None, None
            beta = None
            if len(jets) >= 2:
                jet1, jet2 = jets[:2]

                # determine boost of system
                # determine jet CoM frame
                beta = (jet1.fourvect + jet2.fourvect).BoostVector()
                tree.jet_beta.copy_from(beta)

                jet1.fourvect_boosted = LorentzVector()
                jet1.fourvect_boosted.copy_from(jet1.fourvect)
                jet1.fourvect_boosted.Boost(beta * -1)

                jet2.fourvect_boosted = LorentzVector()
                jet2.fourvect_boosted.copy_from(jet2.fourvect)
                jet2.fourvect_boosted.Boost(beta * -1)

                tau1.min_dr_jet = min(
                    tau1.fourvect.DeltaR(jet1.fourvect),
                    tau1.fourvect.DeltaR(jet2.fourvect))
                tau2.min_dr_jet = min(
                    tau2.fourvect.DeltaR(jet1.fourvect),
                    tau2.fourvect.DeltaR(jet2.fourvect))

                # tau centrality (degree to which they are between the two jets)
                tau1.centrality = eventshapes.eta_centrality(
                    tau1.fourvect.Eta(),
                    jet1.fourvect.Eta(),
                    jet2.fourvect.Eta())

                tau2.centrality = eventshapes.eta_centrality(
                    tau2.fourvect.Eta(),
                    jet1.fourvect.Eta(),
                    jet2.fourvect.Eta())

                # boosted tau centrality
                tau1.centrality_boosted = eventshapes.eta_centrality(
                    tau1.fourvect_boosted.Eta(),
                    jet1.fourvect_boosted.Eta(),
                    jet2.fourvect_boosted.Eta())

                tau2.centrality_boosted = eventshapes.eta_centrality(
                    tau2.fourvect_boosted.Eta(),
                    jet1.fourvect_boosted.Eta(),
                    jet2.fourvect_boosted.Eta())

                # 3rd leading jet
                if len(jets) >= 3:
                    jet3 = jets[2]
                    jet3.fourvect_boosted = LorentzVector()
                    jet3.fourvect_boosted.copy_from(jet3.fourvect)
                    jet3.fourvect_boosted.Boost(beta * -1)

            elif len(jets) == 1:
                jet1 = jets[0]

                tau1.min_dr_jet = tau1.fourvect.DeltaR(jet1.fourvect)
                tau2.min_dr_jet = tau2.fourvect.DeltaR(jet1.fourvect)

            RecoJetBlock.set(tree, jet1, jet2, jet3, local=local)

            # mass of ditau + leading jet system
            if jet1 is not None:
                tree.mass_tau1_tau2_jet1 = (
                    tau1.fourvect + tau2.fourvect + jet1.fourvect).M()

            #####################################
            # number of tracks from PV minus taus
            #####################################
            ntrack_pv = 0
            ntrack_nontau_pv = 0
            for vxp in event.vertices:
                # primary vertex
                if vxp.vertexType() == 1:
                    ntrack_pv = vxp.nTrackParticles()
                    ntrack_nontau_pv = ntrack_pv - tau1.nTracks() - tau2.nTracks()
                    break
            tree.ntrack_pv = ntrack_pv
            tree.ntrack_nontau_pv = ntrack_nontau_pv

            #########################
            # MET variables
            #########################
            MET = event.MET.collection['Final']
            METx = MET.mpx()
            METy = MET.mpy()
            METet = MET.met()
            MET_vect = Vector2(METx, METy)
            MET_4vect = LorentzVector()
            MET_4vect.SetPxPyPzE(METx, METy, 0., METet)
            MET_4vect_boosted = LorentzVector()
            MET_4vect_boosted.copy_from(MET_4vect)
            if beta is not None:
                MET_4vect_boosted.Boost(beta * -1)

            tree.MET_et = METet
            tree.MET_etx = METx
            tree.MET_ety = METy
            tree.MET_phi = MET.phi()
            dPhi_tau1_tau2 = abs(tau1.fourvect.DeltaPhi(tau2.fourvect))
            dPhi_tau1_MET = abs(tau1.fourvect.DeltaPhi(MET_4vect))
            dPhi_tau2_MET = abs(tau2.fourvect.DeltaPhi(MET_4vect))
            tree.dPhi_tau1_tau2 = dPhi_tau1_tau2
            tree.dPhi_tau1_MET = dPhi_tau1_MET
            tree.dPhi_tau2_MET = dPhi_tau2_MET
            tree.dPhi_min_tau_MET = min(dPhi_tau1_MET, dPhi_tau2_MET)
            tree.MET_bisecting = is_MET_bisecting(
                dPhi_tau1_tau2,
                dPhi_tau1_MET,
                dPhi_tau2_MET)

            sumET = MET.sumet()
            tree.MET_sumet = sumET
            if sumET != 0:
                tree.MET_sig = ((2. * METet / GeV) /
                    (utils.sign(sumET) * sqrt(abs(sumET / GeV))))
            else:
                tree.MET_sig = -1.

            tree.MET_centrality = eventshapes.phi_centrality(
                tau1.fourvect,
                tau2.fourvect,
                MET_vect)
            tree.MET_centrality_boosted = eventshapes.phi_centrality(
                tau1.fourvect_boosted,
                tau2.fourvect_boosted,
                MET_4vect_boosted)

            tree.number_of_good_vertices = len(event.vertices)

            ##########################
            # Jet and sum pt variables
            ##########################
            tree.numJets = len(event.jets)

            # sum pT with only the two leading jets
            tree.sum_pt = sum(
                [tau1.pt(), tau2.pt()] +
                [jet.pt() for jet in jets[:2]])

            # sum pT with all selected jets
            tree.sum_pt_full = sum(
                [tau1.pt(), tau2.pt()] +
                [jet.pt() for jet in jets])

            # vector sum pT with two leading jets and MET
            tree.vector_sum_pt = sum(
                [tau1.fourvect, tau2.fourvect] +
                [jet.fourvect for jet in jets[:2]] +
                [MET_4vect]).Pt()

            # vector sum pT with all selected jets and MET
            tree.vector_sum_pt_full = sum(
                [tau1.fourvect, tau2.fourvect] +
                [jet.fourvect for jet in jets] +
                [MET_4vect]).Pt()

            # resonance pT
            tree.resonance_pt = sum(
                [tau1.fourvect, tau2.fourvect, MET_4vect]).Pt()

            # #############################
            # # tau <-> vertex association
            # #############################
            tree.tau_same_vertex = (
                tau1.vertex() == tau2.vertex())

            tau1.vertex_prob = ROOT.TMath.Prob(
                tau1.vertex().chiSquared(),
                int(tau1.vertex().numberDoF()))

            tau2.vertex_prob = ROOT.TMath.Prob(
                tau2.vertex().chiSquared(),
                int(tau2.vertex().numberDoF()))

            # ##########################
            # # MMC Mass
            # ##########################
            # OLD USAGE
            # mmc_result = mmc.mass(
            #     tau1, tau2,
            #     METx, METy, sumET,
            #     njets=len(event.jets))
            # for mmc_method, mmc_object in enumerate(mmc_objects):
            #     mmc_mass, mmc_resonance, mmc_met = mmc_result[mmc_method]
            #     if verbose:
            #         log.info("MMC (method %d): %f" % (mmc_method, mmc_mass))

            #     mmc_object.mass = mmc_mass
            #     mmc_object.MET_et = mmc_met.Mod()
            #     mmc_object.MET_etx = mmc_met.X()
            #     mmc_object.MET_ety = mmc_met.Y()
            #     mmc_object.MET_phi = math.pi - mmc_met.Phi()
            #     if mmc_mass > 0:
            #         FourMomentum.set(mmc_object.resonance, mmc_resonance)

            mass_tool.apply(event.EventInfo, tau1, tau2, MET, len(event.jets))
            for i, mmc_object in enumerate(mmc_objects):
                mmc_object.mass = event.EventInfo.auxdataConst('double')('mmc%s_mass' % i)
                mmc_object.MET_et = mass_tool.GetFittedMetVec(i).Mod()
                mmc_object.MET_etx = mass_tool.GetFittedMetVec(i).X()
                mmc_object.MET_ety = mass_tool.GetFittedMetVec(i).Y()
                mmc_object.MET_phi = math.pi - mass_tool.GetFittedMetVec(i).Phi()
                if mmc_object.mass > 0:
                    FourMomentum.set(mmc_object.resonance, mass_tool.GetResonanceVec(i))


            # ############################
            # # collinear and visible mass
            # ############################
            # vis_mass, collin_mass, tau1_x, tau2_x = mass.collinearmass(
            #     tau1, tau2, METx, METy)

            # tree.mass_vis_tau1_tau2 = vis_mass
            # tree.mass_collinear_tau1_tau2 = collin_mass
            # tau1.collinear_momentum_fraction = tau1_x
            # tau2.collinear_momentum_fraction = tau2_x

            # # Fill the tau block
            # # This must come after the RecoJetBlock is filled since
            # # that sets the jet_beta for boosting the taus
            RecoTauBlock.set(event, tree, datatype, tau1, tau2, local=local)

            # if datatype != datasets.DATA:
            #     TrueTauBlock.set(tree, tau1, tau2)

            # fill the output tree
            outtree.Fill(reset=True)

        # externaltools.report()

        # flush any baskets remaining in memory to disk
        self.output.cd()
        outtree.FlushBaskets()
        outtree.Write()

        if local:
            if datatype == datasets.DATA:
                xml_string = ROOT.TObjString(merged_grl.str())
                xml_string.Write('lumi')
            merged_cutflow.Write()
예제 #28
0
from rootpy.io import root_open
from rootpy import stl
from random import gauss


f = root_open("test.root", "recreate")

# define the model
class Event(TreeModel):

    x = stl.vector('TLorentzVector')
    i = IntCol()

tree = Tree("test", model=Event)

# fill the tree
for i in xrange(100):
    tree.x.clear()
    for j in xrange(5):
        vect = LorentzVector(
                gauss(.5, 1.),
                gauss(.5, 1.),
                gauss(.5, 1.),
                gauss(.5, 1.))
        tree.x.push_back(vect)
    tree.i = i
    tree.fill()

tree.write()
f.close()
예제 #29
0
파일: dictW.py 프로젝트: afwebb/higgs_diff
def lorentzVecsLeps(nom, is3l):

    '''
    Higgs decays to two jets and one lepton, or two leptons. This returns lorentzVectors for each decay product candidate
    For H -> 2j, 1l case (not isF): returns jet0, jet1, met, lep0, lep1, (lep2 if is3l)
    For H -> 2l case (isF): Return met, lep0, lep1, lep2 
    '''

    met = LorentzVector()
    met.SetPtEtaPhiE(nom.met_met, 0, nom.met_phi, nom.met_met)

    lep0 = LorentzVector()
    lep0.SetPtEtaPhiE(nom.lep_Pt_0, nom.lep_Eta_0, nom.lep_Phi_0, nom.lep_E_0)

    lep1 = LorentzVector()
    lep1.SetPtEtaPhiE(nom.lep_Pt_1, nom.lep_Eta_1, nom.lep_Phi_1, nom.lep_E_1)

    if is3l:
        lep2 = LorentzVector()
        lep2.SetPtEtaPhiE(nom.lep_Pt_2, nom.lep_Eta_2, nom.lep_Phi_2, nom.lep_E_2)


    if is3l:
        return (lep0, lep1, lep2, met)
    else:
        return (lep0, lep1, met)
예제 #30
0
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
예제 #31
0
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