mickypaganini · February 1, 2016 13:19
diff --git a/angles.py b/angles.py
 from root_numpy.tmva import add_classification_events, evaluate_reader
 from rootpy.io import root_open
 from rootpy import stl, log
 from rootpy.tree import Tree
 from rootpy.vector import Vector3, LorentzVector
 from ROOT import TMVA, TFile
 import xml.etree.ElementTree as ET
 from array import array
 import pandas as pd
 import pandautils as pup
 import numpy as np
 import math
 import os
 import sys
 from rootpy.plotting.style import get_style, set_style
 import matplotlib.pyplot as plt
 import fit
 import logging 

 class Jet(object):
    """
    Better Jet Class!

    Example:

    >>> j = Jet(100.1, 1.2, 1.0, 12, 1) # instantiation with assignment
    >>> j.lv.Pt() # print pt
    100.1
    >>> j.pt = 205 # assign pt
    >>> j.lv.Pt()
    205

    """
    def __init__(self, pt = 0 , eta = 0, phi = 0, m = 0, btag = 0):
        super(Jet, self).__init__()

        self.btag = btag
        
        # -- hidden to most usage
        self._pt = pt
        self._eta = eta
        self._phi = phi
        self._m = m

        # -- Lorentz 4-vector
        self.lv = LorentzVector()
        self.lv.SetPtEtaPhiM(self._pt, self._eta, self._phi, self._m)

    def _internal_update(self):
        self.lv.SetPtEtaPhiM(self._pt, self._eta, self._phi, self._m)

    @property
    def pt(self):
        return self._pt

    @pt.setter
    def pt(self, value):
        self._pt = value
        self._internal_update()
        
    @property
    def eta(self):
        return self._eta

    @eta.setter
    def eta(self, value):
        self._eta = value
        self._internal_update()        

    @property
    def phi(self):
        return self._phi

    @phi.setter
    def phi(self, value):
        self._phi = value
        self._internal_update()
        
    @property
    def m(self):
        return self._m

    @m.setter
    def m(self, value):
        self._m = value
        self._internal_update()



 class Photon(object):
    """
    Photon Class
    """
    def __init__(self, pt = 0 , eta = 0, phi = 0, m = 0):
        super(Photon, self).__init__()
        
        # -- hidden to most usage
        self._pt = pt
        self._eta = eta
        self._phi = phi
        self._m = m

        # -- Lorentz 4-vector
        self.lv = LorentzVector()
        self.lv.SetPtEtaPhiM(self._pt, self._eta, self._phi, self._m)

    def _internal_update(self):
        self.lv.SetPtEtaPhiM(self._pt, self._eta, self._phi, self._m)

    @property
    def pt(self):
        return self._pt

    @pt.setter
    def pt(self, value):
        self._pt = value
        self._internal_update()
        
    @property
    def eta(self):
        return self._eta

    @eta.setter
    def eta(self, value):
        self._eta = value
        self._internal_update()        

    @property
    def phi(self):
        return self._phi

    @phi.setter
    def phi(self, value):
        self._phi = value
        self._internal_update()
        
    @property
    def m(self):
        return self._m

    @m.setter
    def m(self, value):
        self._m = value
        self._internal_update()



 class Muon(object):
    """
    Muon Class
    """
    def __init__(self, pt = 0 , eta = 0, phi = 0, m = 0):
        super(Muon, self).__init__()
        
        # -- hidden to most usage
        self._pt = pt
        self._eta = eta
        self._phi = phi
        self._m = m

        # -- Lorentz 4-vector
        self.lv = LorentzVector()
        self.lv.SetPtEtaPhiM(self._pt, self._eta, self._phi, self._m)

    def _internal_update(self):
        self.lv.SetPtEtaPhiM(self._pt, self._eta, self._phi, self._m)

    @property
    def pt(self):
        return self._pt

    @pt.setter
    def pt(self, value):
        self._pt = value
        self._internal_update()
        
    @property
    def eta(self):
        return self._eta

    @eta.setter
    def eta(self, value):
        self._eta = value
        self._internal_update()        

    @property
    def phi(self):
        return self._phi

    @phi.setter
    def phi(self, value):
        self._phi = value
        self._internal_update()
        
    @property
    def m(self):
        return self._m

    @m.setter
    def m(self, value):
        self._m = value
        self._internal_update()





 def main(signalfile, bkgfile):

    # -- remove root warnings saying that containers cannot be opened (because we don't care, anyways)
    logging.basicConfig(level=logging.ERROR)
    log["/ROOT.TClass.TClass"].setLevel(log.ERROR)

    set_style('ATLAS', mpl=True)

    # -- only extract the variables we care about from the .root files
    features = get_features(signalfile)

    # -- initialize dictionary for signal and background variables
    val_dict = {}


    # -- perform the same study on both signal and background
    for filename in [signalfile, bkgfile]:

        print 'Working on {} ...'.format(filename.rpartition('23LO_')[2].rpartition('.merge.')[0])
        sys.stdout.flush()

        # -- put .root files into pandas dataframes
        signal_df = load_df(filename, features)

        # -- hack to only consider events that pass the cutflow (excluding the m_bb cut, which gets opened up in a larger window)
        signal_df = signal_df[signal_df['HH2yybbEventInfoAuxDyn.m_yyjj'] != -99]

        # -- jet
        jets = [Jet(pt, eta, phi, m, btag) for (pt, eta, phi, m, btag) in zip(
            pup.flatten(signal_df['HGamAntiKt4EMTopoJetsAuxDyn.pt']),
            pup.flatten(signal_df['HGamAntiKt4EMTopoJetsAuxDyn.eta']),
            pup.flatten(signal_df['HGamAntiKt4EMTopoJetsAuxDyn.phi']),
            pup.flatten(signal_df['HGamAntiKt4EMTopoJetsAuxDyn.m']),
            pup.flatten(signal_df['HGamAntiKt4EMTopoJetsAuxDyn.MV2c20_85'])
            )]
        # -- reshape jets to match the right quantity per event
        jets = pup.match_shape(np.array(jets), signal_df['HGamAntiKt4EMTopoJetsAuxDyn.pt'])

        # -- photons
        photons = [Photon(pt, eta, phi, m) for (pt, eta, phi, m) in zip(
            pup.flatten(signal_df['HGamPhotonsAuxDyn.pt']),
            pup.flatten(signal_df['HGamPhotonsAuxDyn.eta']),
            pup.flatten(signal_df['HGamPhotonsAuxDyn.phi']),
            pup.flatten(signal_df['HGamPhotonsAuxDyn.m'])
            )]
        # -- reshape photons to match the right quantity per event
        photons = pup.match_shape(np.array(photons), signal_df['HGamPhotonsAuxDyn.pt'])

        # -- muons
        muons = [Muon(pt, eta, phi, 105.65837) for (pt, eta, phi) in zip(
            pup.flatten(signal_df['HGamMuonsAuxDyn.pt']),
            pup.flatten(signal_df['HGamMuonsAuxDyn.eta']),
            pup.flatten(signal_df['HGamMuonsAuxDyn.phi'])
            )]
        # -- reshape muon to match the right quantity per event
        muons = pup.match_shape(np.array(muons), signal_df['HGamMuonsAuxDyn.pt'])

        # -- cut out b-jets only
        bjets = [((np.array(jets[evn]))[signal_df['HGamAntiKt4EMTopoJetsAuxDyn.MV2c20_85'][ev] == 1]).tolist() for (evn,ev) 
                 in enumerate(signal_df.index.values) ]

        # -- select events with == 2 b-tagged jets (2-Tag Category)
        a = [(len(bjets[ev]) == 2) for ev in xrange(len(bjets))] 
        bjets_2 = np.array(bjets)[np.array(a)]
        photons_2 = np.array(photons)[np.array(a)]
        muons_2 = np.array(muons)[np.array(a)] 

        # -- apply corrections
        bjets_2_corr04 = apply_correction(0.4, bjets_2, muons_2)
        bjets_2_corr02 = apply_correction(0.2, bjets_2, muons_2)
        bjets_2_corr01 = apply_correction(0.1, bjets_2, muons_2)

        # -- uncorrected four vectors for dijet, diphoton and dihiggs systems
        bb, yy, yybb, bb04, yybb04, bb02, yybb02, bb01, yybb01 = get_composite_sys(
            bjets_2, photons_2, bjets_2_corr04, bjets_2_corr02, bjets_2_corr01)

        # -- plot muon in jet correction results (m_bb, m_yybb) for different radii of association
        #dr_plots(yybb04, yybb02, yybb01, yybb, xmin = 250, xmax = 500, sys = '{\gamma \gamma b b}')
        #dr_plots(bb04, bb02, bb01, bb, xmin = 80, xmax = 150, bins = 15)

        # -- define a bunch of useful variables:
        hh = [(h1,h2) for (h1,h2) in zip(yy,bb04)] # with muon corrections
        m_bb = np.array([bb04[ev].M() for ev in xrange(len(bb04))])
        mass_window = np.logical_and( (m_bb > 80000), (m_bb < 135000) )
        #h_pts = np.array([np.sort([hh[ev][h].pt() for h in xrange(2)]) for ev in xrange(len(hh))]) # sorted!!

        # -- get DW's variables as computed in their code:
        m12, m34, mTX, pTX, yX, cos_theta_star, phi, phi1, cos_theta1, cos_theta2 = DW_variables(
            mass_window, bb04, yy, yybb04, bjets_2_corr04, photons_2)

        val_dict[filename] = {'m12':np.array(m12)/1000, 'm34':np.array(m34)/1000, 'mTX':np.array(mTX)/1000, 
           'pTX':np.array(pTX)/1000, 'yX':yX, 'cos_theta_star':cos_theta_star, 'phi':phi, 'phi1':np.abs(phi1)-(math.pi/2), 'cos_theta1':cos_theta1, 'cos_theta2':cos_theta2, 'mass_window':mass_window}
        
    # -- hardcoding the binning strategy
    bin_dict = {'m12':np.linspace(70,140,30), 
    'm34':np.linspace(100,160,50), 'mTX':np.linspace(200,650,50), 'pTX':np.linspace(0,400,50), 'yX':np.linspace(-6,6,20), 
        'cos_theta_star':np.linspace(-1,1,20), 'phi':np.linspace(-4, 4, 30), 'phi1':np.linspace(-1.8,0.2,30), 'cos_theta1':np.linspace(-1.5,1.5,20), 
        'cos_theta2':np.linspace(-1.5,1.5,20)}

    # -- plotting
    print 'Plotting ...'
    for key in val_dict[signalfile].keys():
        if key == 'mass_window': continue
        else:    
            plot_DW_vars(val_dict, signalfile, bkgfile, key, bin_dict)

    return 0


 def get_features(signalfile):
    f = root_open(signalfile, "read")
    t = f["CollectionTree"]
    features = [key.GetName() for key in t.GetListOfBranches() if (
            ('HGamMuonsAuxDyn' in key.GetName()) 
            or ('HGamAntiKt4EMTopoJetsAuxDyn' in key.GetName()) 
            or ('HGamPhotonsAuxDyn' in key.GetName()) 
            or ('HH2yybbEventInfoAuxDyn' in key.GetName())
        )]
    f.Close()
    return features


 def load_df(filename, features):
    df = pup.root2panda(filename, "CollectionTree", branches = features)
    return df


 def deltaR(eta1, eta2, phi1, phi2):
    import math
    '''
    Definition:
    -----------
        Function that calculates DR between two objects given their etas and phis
    Args:
    -----
        eta1 = float, eta of first object
        eta2 = float, eta of second object
        phi1 = float, phi of first object
        phi2 = float, phi of second object
    Output:
    -------
        deltaR = float, distance between the two objects 
    '''
    DEta = abs(eta1-eta2)
    DPhi = math.acos(math.cos( abs( phi1-phi2 ) ) ) # hack to avoid |phi1-phi2| larger than 180 degrees
    return math.sqrt( pow(DEta,2) + pow(DPhi,2) ) 


 # -- now we need to corrcet the jets by the muons, if the muons are inside of them
 # -- start by calculating the distances between each muon and each jet. If the distance is less than a threshold, correct the jet by the muon
 def apply_correction(DRMAX, bjets_2, muons_2):

    corr_bjets_2 = []
    for ev in xrange(len(bjets_2)): # loop thru events
        corr_bjets_2_perevent = []
        for jet in bjets_2[ev]: # loop thru jets
            muon_add = Muon() # save all muons before adding them to the jet
            if (len(muons_2[ev]) != 0): # if there is at least one muon, do correction
                for muon in muons_2[ev]: # loop thru muons
                    dr = deltaR(jet.eta, muon.eta, jet.phi, muon.phi) # calculate dr between muon and original jet
                if ((dr < DRMAX) and (muon.lv.Pt() > 4000)): # if it's closer that the fixed DR and pT>4GeV, add in the muon
                    muon_add.lv += muon.lv
            corr_bjets_2_perevent.append(Jet(
                    (jet.lv + muon_add.lv).pt(),
                    (jet.lv + muon_add.lv).eta(),
                    (jet.lv + muon_add.lv).phi(),
                    (jet.lv + muon_add.lv).m(),
                    jet.btag
                ))
        corr_bjets_2.append(corr_bjets_2_perevent)
    
    return corr_bjets_2

 def get_composite_sys(bjets_2, photons_2, bjets_2_corr04, bjets_2_corr02, bjets_2_corr01):

    bb = []
    yy = []
    yybb = []
    bb04 = []
    yybb04 = []
    bb02 = []
    yybb02 = []
    bb01 = []
    yybb01 = []

    for ev in xrange(len(bjets_2)):
        bb.append(bjets_2[ev][0].lv + bjets_2[ev][1].lv)
        yy.append(photons_2[ev][0].lv + photons_2[ev][1].lv)
        yybb.append(bjets_2[ev][0].lv + bjets_2[ev][1].lv + photons_2[ev][0].lv + photons_2[ev][1].lv)
        
        bb04.append(bjets_2_corr04[ev][0].lv + bjets_2_corr04[ev][1].lv)
        yybb04.append(bjets_2_corr04[ev][0].lv + bjets_2_corr04[ev][1].lv + photons_2[ev][0].lv + photons_2[ev][1].lv)
        bb02.append(bjets_2_corr02[ev][0].lv + bjets_2_corr02[ev][1].lv)
        yybb02.append(bjets_2_corr02[ev][0].lv + bjets_2_corr02[ev][1].lv + photons_2[ev][0].lv + photons_2[ev][1].lv)
        bb01.append(bjets_2_corr01[ev][0].lv + bjets_2_corr01[ev][1].lv)
        yybb01.append(bjets_2_corr01[ev][0].lv + bjets_2_corr01[ev][1].lv + photons_2[ev][0].lv + photons_2[ev][1].lv)

    return bb, yy, yybb, bb04, yybb04, bb02, yybb02, bb01, yybb01


 def dr_plots(corrected_recomass_04, corrected_recomass_02, corrected_recomass_01, recomass, xmin = 80, xmax = 160, 
             bins = 20, sys = '{bb}'):

    hista = [corrected_recomass_04[ev].M()/1000 for ev in xrange(len(corrected_recomass_04))]
    histb = [corrected_recomass_02[ev].M()/1000 for ev in xrange(len(corrected_recomass_02))]
    histd = [corrected_recomass_01[ev].M()/1000 for ev in xrange(len(corrected_recomass_01))] 
    histc = [recomass[ev].M()/1000 for ev in xrange(len(recomass))]   

    a, bina = np.histogram(hista, bins = 1000, range = (xmin, xmax), normed=True)
    b, binb = np.histogram(histb, bins = 1000, range = (xmin, xmax), normed=True)
    d, bind = np.histogram(histd, bins = 1000, range = (xmin, xmax), normed=True)
    c, binc = np.histogram(histc, bins = 1000, range = (xmin, xmax), normed=True)
    
    plt.hist(hista, label = 'Corrected All Muons, DR = 0.4', linewidth = 1,
               bins = np.linspace(xmin, xmax, bins), histtype = 'step', alpha = 0.5, color = 'blue', normed = True)
    
    plt.hist(histb, label = 'Corrected All Muons, DR = 0.2', linewidth = 1,
               bins = np.linspace(xmin, xmax, bins), histtype = 'step', alpha = 0.5, color = 'green', normed = True)
    
    plt.hist(histd, label = 'Corrected All Muons, DR = 0.1', linewidth = 1,
               bins = np.linspace(xmin, xmax, bins), histtype = 'step', alpha = 0.5, color = 'red', normed = True)
    
    plt.hist(histc, label = 'No Correction', linewidth = 1,
               bins = np.linspace(xmin, xmax, bins), histtype = 'step', alpha = 0.5, color = 'black', normed = True)

    
    bincentersa = 0.5*(bina[1:]+bina[:-1])
    bincentersb = 0.5*(binb[1:]+binb[:-1])
    bincentersd = 0.5*(bind[1:]+bind[:-1])
    bincentersc = 0.5*(binc[1:]+binc[:-1])
    (xfa,yfa),coeffa,erra,_ = fit.fit(fit.crystal_ball, bincentersa, a)
    (xfb,yfb),coeffb,errb,_ = fit.fit(fit.crystal_ball, bincentersb, b)
    (xfd,yfd),coeffd,errd,_ = fit.fit(fit.crystal_ball, bincentersd, d)
    (xfc,yfc),coeffc,errc,_ = fit.fit(fit.crystal_ball, bincentersc, c)
    
    plt.plot(xfa,yfa, color = 'blue')
    plt.plot(xfb,yfb, color = 'green')
    plt.plot(xfd,yfd, color = 'red')
    plt.plot(xfc,yfc, color = 'black')
    
    #plt.xlim(xmin=80, xmax=160)
    #plt.ylim( ymax = 0.015)
    plt.legend(loc = 'best', fontsize = 10)
    plt.ylabel('Normalized')
    plt.xlabel(r'$m_{}\ (GeV)$'.format(sys))
    plt.title(r'Invariant Mass of ${}$ system, 2-Tag Category'.format(sys))
    plt.show()
    
    print r'DR = 0.4, Peak = {0:.2f} +- {1:.2f} GeV, Sigma = {2:.2f} +- {3:.2f} GeV'.format(coeffa[3], erra[3], coeffa[4], erra[4] )
    print r'DR = 0.2, Peak = {0:.2f} +- {1:.2f} GeV, Sigma = {2:.2f} +- {3:.2f} GeV'.format(coeffb[3], errb[3], coeffb[4], errb[4])
    print r'DR = 0.1, Peak = {0:.2f} +- {1:.2f} GeV, Sigma = {2:.2f} +- {3:.2f} GeV'.format(coeffd[3], errd[3], coeffd[4], errd[4])
    print r'No Correction, Peak = {0:.2f} +- {1:.2f} GeV, Sigma = {2:.2f} +- {3:.2f} GeV'.format(coeffc[3], errc[3], coeffc[4], errc[4])


 def DW_variables(mass_window, bb04, yy, yybb04, bjets_2_corr04, photons_2):
    import math
    # -- start calculating values in David Wardrope's paper
    m12 = [bb04[ev].m() for ev in xrange(len(bb04))] 
    m34 = [yy[ev].m() for ev in xrange(len(yy))]
    mTX = [yybb04[ev].m() for ev in xrange(len(yybb04))] 
    pTX = [yybb04[ev].pt() for ev in xrange(len(yybb04))] 
    yX = [yybb04[ev].eta() for ev in xrange(len(yybb04))] # units problem?

    # -- this crap below should follow exactly what they do in their code...
    q1_lab = bb04 # bb system
    q2_lab = yy   # yy system
    X_lab = [(q1_lab[ev] + q2_lab[ev]) for ev in xrange(len(q1_lab))] # yybb system

    # -- BOOST: bb, yy, yybb, individual jets and photons --> to the rest frame of the yybb system
    [q1_lab[ev].Boost(-X_lab[ev].BoostVector().x, -X_lab[ev].BoostVector().y, -X_lab[ev].BoostVector().z) 
        for ev in xrange(len(q1_lab))]
    q1 = q1_lab 


    [bjets_2_corr04[ev][0].lv.Boost(-X_lab[ev].BoostVector().x, -X_lab[ev].BoostVector().y, -X_lab[ev].BoostVector().z) 
        for ev in xrange(len(bjets_2_corr04))]
    q11 = [bjets_2_corr04[ev][0].lv for ev in xrange(len(bjets_2_corr04))]


    [bjets_2_corr04[ev][1].lv.Boost(-X_lab[ev].BoostVector().x, -X_lab[ev].BoostVector().y, -X_lab[ev].BoostVector().z) 
        for ev in xrange(len(bjets_2_corr04))]
    q12 = [bjets_2_corr04[ev][1].lv for ev in xrange(len(bjets_2_corr04))]


    [q2_lab[ev].Boost(-X_lab[ev].BoostVector().x, -X_lab[ev].BoostVector().y, -X_lab[ev].BoostVector().z) 
        for ev in xrange(len(q2_lab))]
    q2 = q1_lab


    [photons_2[ev][0].lv.Boost(-X_lab[ev].BoostVector().x, -X_lab[ev].BoostVector().y, -X_lab[ev].BoostVector().z) 
        for ev in xrange(len(photons_2))]
    q21 = [photons_2[ev][0].lv for ev in xrange(len(photons_2))]


    [photons_2[ev][1].lv.Boost(-X_lab[ev].BoostVector().x, -X_lab[ev].BoostVector().y, -X_lab[ev].BoostVector().z) 
        for ev in xrange(len(photons_2))]
    q22 = [photons_2[ev][1].lv for ev in xrange(len(photons_2))]


    # -- cos(\Theta^*)
    nZ = Vector3(0., 0., 1.) # useless, because that's already where the angles are calculated from, but ok
    cos_theta_star = [math.cos( (Vector3(q1[ev].px, q1[ev].py, q1[ev].pz)).theta()-nZ.theta() ) for ev in xrange(len(q1))]

    # -- normal vector to bb plane
    n1 = [ Vector3(
        (Vector3(q11[ev].px, q11[ev].py, q11[ev].pz)).cross( Vector3(q12[ev].px, q12[ev].py, q12[ev].pz) ).Unit() ) for ev in xrange(len(q11))] 

    # -- normal vector to yy plane
    n2 = [ Vector3(
            (Vector3(q21[ev].px, q21[ev].py, q21[ev].pz)).cross( Vector3(q22[ev].px, q22[ev].py, q22[ev].pz) ).Unit() ) 
          for ev in xrange(len(q21))] 

    # -- normal vector to the plane that contains the z axis and the first Higgs boson
    nSC = [ Vector3(
            nZ.cross( Vector3(q1[ev].px, q1[ev].py, q1[ev].pz) ).Unit() ) 
          for ev in xrange(len(q1))] 

    # -- acos actually calculates the angle (why a - sign??)
    # -- the first factor just calculates the sign with respect to the first Higgs boson (the one that -->bb)
    phi  = [np.sign( Vector3(q1[ev].px, q1[ev].py, q1[ev].pz).dot(n1[ev].cross(n2[ev])) )  *  math.acos(-n1[ev].dot(n2[ev])) for ev in xrange(len(q1))]

    # -- I don't even know
    phi1 = [np.sign( Vector3(q1[ev].px, q1[ev].py, q1[ev].pz).dot(n1[ev].cross(nSC[ev])) ) *  math.acos(n1[ev].dot(nSC[ev])) for ev in xrange(len(q1))]
    phi1_plot = abs(np.array(phi1))-(math.pi/2)

    # -- BOOST: individual jets and photons --> to the rest frame of the corresponding parent Higgs
    [bjets_2_corr04[ev][0].lv.Boost(-q1_lab[ev].BoostVector().x, -q1_lab[ev].BoostVector().y, -q1_lab[ev].BoostVector().y ) for ev in xrange(len(bjets_2_corr04))]
    q11_H1 = [bjets_2_corr04[ev][0].lv for ev in xrange(len(bjets_2_corr04))]

    [photons_2[ev][0].lv.Boost(-q2_lab[ev].BoostVector().x, -q2_lab[ev].BoostVector().y, -q2_lab[ev].BoostVector().y ) for ev in xrange(len(photons_2))]
    q21_H2 = [photons_2[ev][0].lv for ev in xrange(len(photons_2))]

    # -- cos(angle) between the first Higgs in the yybb rest frame and the first b jet in the rest fram of its parent Higgs (???)
    cos_theta1 = [Vector3(q1[ev].px, q1[ev].py, q1[ev].pz).dot(Vector3(q11_H1[ev].px, q11_H1[ev].py, q11_H1[ev].pz)) / math.sqrt(Vector3(q1[ev].px, q1[ev].py, q1[ev].pz).mag2() * Vector3(q11_H1[ev].px, q11_H1[ev].py, q11_H1[ev].pz).mag2()) for ev in xrange(len(q1))]

    # -- cos(angle) between the second Higgs in the yybb rest frame and the first y in the rest fram of its parent Higgs (???)
    cos_theta2 = [Vector3(q2[ev].px, q2[ev].py, q2[ev].pz).dot(Vector3(q21_H2[ev].px, q21_H2[ev].py, q21_H2[ev].pz)) / math.sqrt(Vector3(q2[ev].px, q2[ev].py, q2[ev].pz).mag2() * Vector3(q21_H2[ev].px, q21_H2[ev].py, q21_H2[ev].pz).mag2()) for ev in xrange(len(q2))]

    return m12, m34, mTX, pTX, yX, cos_theta_star, phi, phi1_plot, cos_theta1, cos_theta2




 def plot_DW_vars(val_dict, signalfile, bkgfile, key, bin_dict):

    bins = bin_dict[key]
    _ = plt.hist(np.array(val_dict[signalfile][key])[val_dict[signalfile]['mass_window']], bins = bins, alpha = 0.5, histtype = 'stepfilled', normed = True, label = 'X350') 
    _ = plt.hist(np.array(val_dict[bkgfile][key])[val_dict[bkgfile]['mass_window']], bins = bins, alpha = 1, histtype = 'step', normed = True, label = 'yybb',
                color = 'red', hatch = '/////') 
    plt.xlabel(key)
    plt.ylabel('Normalized')
    plt.legend()
    plt.show()




 if __name__ == '__main__':

    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument("signalfile", help="input .root file name for signal")
    parser.add_argument("bkgfile", help="input .root file name for background")
    args = parser.parse_args()

    sys.exit( main(args.signalfile, args.bkgfile) )
	from root_numpy.tmva import add_classification_events, evaluate_reader
	from rootpy.io import root_open
	from rootpy import stl, log
	from rootpy.tree import Tree
	from rootpy.vector import Vector3, LorentzVector
	from ROOT import TMVA, TFile
	import xml.etree.ElementTree as ET
	from array import array
	import pandas as pd
	import pandautils as pup
	import numpy as np
	import math
	import os
	import sys
	from rootpy.plotting.style import get_style, set_style
	import matplotlib.pyplot as plt
	import fit
	import logging

	class Jet(object):
	"""
	Better Jet Class!

	Example:

	>>> j = Jet(100.1, 1.2, 1.0, 12, 1) # instantiation with assignment
	>>> j.lv.Pt() # print pt
	100.1
	>>> j.pt = 205 # assign pt
	>>> j.lv.Pt()
	205

	"""
	def __init__(self, pt = 0 , eta = 0, phi = 0, m = 0, btag = 0):
	super(Jet, self).__init__()

	self.btag = btag

	# -- hidden to most usage
	self._pt = pt
	self._eta = eta
	self._phi = phi
	self._m = m

	# -- Lorentz 4-vector
	self.lv = LorentzVector()
	self.lv.SetPtEtaPhiM(self._pt, self._eta, self._phi, self._m)

	def _internal_update(self):
	self.lv.SetPtEtaPhiM(self._pt, self._eta, self._phi, self._m)

	@property
	def pt(self):
	return self._pt

	@pt.setter
	def pt(self, value):
	self._pt = value
	self._internal_update()

	@property
	def eta(self):
	return self._eta

	@eta.setter
	def eta(self, value):
	self._eta = value
	self._internal_update()

	@property
	def phi(self):
	return self._phi

	@phi.setter
	def phi(self, value):
	self._phi = value
	self._internal_update()

	@property
	def m(self):
	return self._m

	@m.setter
	def m(self, value):
	self._m = value
	self._internal_update()



	class Photon(object):
	"""
	Photon Class
	"""
	def __init__(self, pt = 0 , eta = 0, phi = 0, m = 0):
	super(Photon, self).__init__()

	# -- hidden to most usage
	self._pt = pt
	self._eta = eta
	self._phi = phi
	self._m = m

	# -- Lorentz 4-vector
	self.lv = LorentzVector()
	self.lv.SetPtEtaPhiM(self._pt, self._eta, self._phi, self._m)

	def _internal_update(self):
	self.lv.SetPtEtaPhiM(self._pt, self._eta, self._phi, self._m)

	@property
	def pt(self):
	return self._pt

	@pt.setter
	def pt(self, value):
	self._pt = value
	self._internal_update()

	@property
	def eta(self):
	return self._eta

	@eta.setter
	def eta(self, value):
	self._eta = value
	self._internal_update()

	@property
	def phi(self):
	return self._phi

	@phi.setter
	def phi(self, value):
	self._phi = value
	self._internal_update()

	@property
	def m(self):
	return self._m

	@m.setter
	def m(self, value):
	self._m = value
	self._internal_update()



	class Muon(object):
	"""
	Muon Class
	"""
	def __init__(self, pt = 0 , eta = 0, phi = 0, m = 0):
	super(Muon, self).__init__()

	# -- hidden to most usage
	self._pt = pt
	self._eta = eta
	self._phi = phi
	self._m = m

	# -- Lorentz 4-vector
	self.lv = LorentzVector()
	self.lv.SetPtEtaPhiM(self._pt, self._eta, self._phi, self._m)

	def _internal_update(self):
	self.lv.SetPtEtaPhiM(self._pt, self._eta, self._phi, self._m)

	@property
	def pt(self):
	return self._pt

	@pt.setter
	def pt(self, value):
	self._pt = value
	self._internal_update()

	@property
	def eta(self):
	return self._eta

	@eta.setter
	def eta(self, value):
	self._eta = value
	self._internal_update()

	@property
	def phi(self):
	return self._phi

	@phi.setter
	def phi(self, value):
	self._phi = value
	self._internal_update()

	@property
	def m(self):
	return self._m

	@m.setter
	def m(self, value):
	self._m = value
	self._internal_update()





	def main(signalfile, bkgfile):

	# -- remove root warnings saying that containers cannot be opened (because we don't care, anyways)
	logging.basicConfig(level=logging.ERROR)
	log["/ROOT.TClass.TClass"].setLevel(log.ERROR)

	set_style('ATLAS', mpl=True)

	# -- only extract the variables we care about from the .root files
	features = get_features(signalfile)

	# -- initialize dictionary for signal and background variables
	val_dict = {}


	# -- perform the same study on both signal and background
	for filename in [signalfile, bkgfile]:

	print 'Working on {} ...'.format(filename.rpartition('23LO_')[2].rpartition('.merge.')[0])
	sys.stdout.flush()

	# -- put .root files into pandas dataframes
	signal_df = load_df(filename, features)

	# -- hack to only consider events that pass the cutflow (excluding the m_bb cut, which gets opened up in a larger window)
	signal_df = signal_df[signal_df['HH2yybbEventInfoAuxDyn.m_yyjj'] != -99]

	# -- jet
	jets = [Jet(pt, eta, phi, m, btag) for (pt, eta, phi, m, btag) in zip(
	pup.flatten(signal_df['HGamAntiKt4EMTopoJetsAuxDyn.pt']),
	pup.flatten(signal_df['HGamAntiKt4EMTopoJetsAuxDyn.eta']),
	pup.flatten(signal_df['HGamAntiKt4EMTopoJetsAuxDyn.phi']),
	pup.flatten(signal_df['HGamAntiKt4EMTopoJetsAuxDyn.m']),
	pup.flatten(signal_df['HGamAntiKt4EMTopoJetsAuxDyn.MV2c20_85'])
	)]
	# -- reshape jets to match the right quantity per event
	jets = pup.match_shape(np.array(jets), signal_df['HGamAntiKt4EMTopoJetsAuxDyn.pt'])

	# -- photons
	photons = [Photon(pt, eta, phi, m) for (pt, eta, phi, m) in zip(
	pup.flatten(signal_df['HGamPhotonsAuxDyn.pt']),
	pup.flatten(signal_df['HGamPhotonsAuxDyn.eta']),
	pup.flatten(signal_df['HGamPhotonsAuxDyn.phi']),
	pup.flatten(signal_df['HGamPhotonsAuxDyn.m'])
	)]
	# -- reshape photons to match the right quantity per event
	photons = pup.match_shape(np.array(photons), signal_df['HGamPhotonsAuxDyn.pt'])

	# -- muons
	muons = [Muon(pt, eta, phi, 105.65837) for (pt, eta, phi) in zip(
	pup.flatten(signal_df['HGamMuonsAuxDyn.pt']),
	pup.flatten(signal_df['HGamMuonsAuxDyn.eta']),
	pup.flatten(signal_df['HGamMuonsAuxDyn.phi'])
	)]
	# -- reshape muon to match the right quantity per event
	muons = pup.match_shape(np.array(muons), signal_df['HGamMuonsAuxDyn.pt'])

	# -- cut out b-jets only
	bjets = [((np.array(jets[evn]))[signal_df['HGamAntiKt4EMTopoJetsAuxDyn.MV2c20_85'][ev] == 1]).tolist() for (evn,ev)
	in enumerate(signal_df.index.values) ]

	# -- select events with == 2 b-tagged jets (2-Tag Category)
	a = [(len(bjets[ev]) == 2) for ev in xrange(len(bjets))]
	bjets_2 = np.array(bjets)[np.array(a)]
	photons_2 = np.array(photons)[np.array(a)]
	muons_2 = np.array(muons)[np.array(a)]

	# -- apply corrections
	bjets_2_corr04 = apply_correction(0.4, bjets_2, muons_2)
	bjets_2_corr02 = apply_correction(0.2, bjets_2, muons_2)
	bjets_2_corr01 = apply_correction(0.1, bjets_2, muons_2)

	# -- uncorrected four vectors for dijet, diphoton and dihiggs systems
	bb, yy, yybb, bb04, yybb04, bb02, yybb02, bb01, yybb01 = get_composite_sys(
	bjets_2, photons_2, bjets_2_corr04, bjets_2_corr02, bjets_2_corr01)

	# -- plot muon in jet correction results (m_bb, m_yybb) for different radii of association
	#dr_plots(yybb04, yybb02, yybb01, yybb, xmin = 250, xmax = 500, sys = '{\gamma \gamma b b}')
	#dr_plots(bb04, bb02, bb01, bb, xmin = 80, xmax = 150, bins = 15)

	# -- define a bunch of useful variables:
	hh = [(h1,h2) for (h1,h2) in zip(yy,bb04)] # with muon corrections
	m_bb = np.array([bb04[ev].M() for ev in xrange(len(bb04))])
	mass_window = np.logical_and( (m_bb > 80000), (m_bb < 135000) )
	#h_pts = np.array([np.sort([hh[ev][h].pt() for h in xrange(2)]) for ev in xrange(len(hh))]) # sorted!!

	# -- get DW's variables as computed in their code:
	m12, m34, mTX, pTX, yX, cos_theta_star, phi, phi1, cos_theta1, cos_theta2 = DW_variables(
	mass_window, bb04, yy, yybb04, bjets_2_corr04, photons_2)

	val_dict[filename] = {'m12':np.array(m12)/1000, 'm34':np.array(m34)/1000, 'mTX':np.array(mTX)/1000,
	'pTX':np.array(pTX)/1000, 'yX':yX, 'cos_theta_star':cos_theta_star, 'phi':phi, 'phi1':np.abs(phi1)-(math.pi/2), 'cos_theta1':cos_theta1, 'cos_theta2':cos_theta2, 'mass_window':mass_window}

	# -- hardcoding the binning strategy
	bin_dict = {'m12':np.linspace(70,140,30),
	'm34':np.linspace(100,160,50), 'mTX':np.linspace(200,650,50), 'pTX':np.linspace(0,400,50), 'yX':np.linspace(-6,6,20),
	'cos_theta_star':np.linspace(-1,1,20), 'phi':np.linspace(-4, 4, 30), 'phi1':np.linspace(-1.8,0.2,30), 'cos_theta1':np.linspace(-1.5,1.5,20),
	'cos_theta2':np.linspace(-1.5,1.5,20)}

	# -- plotting
	print 'Plotting ...'
	for key in val_dict[signalfile].keys():
	if key == 'mass_window': continue
	else:
	plot_DW_vars(val_dict, signalfile, bkgfile, key, bin_dict)

	return 0


	def get_features(signalfile):
	f = root_open(signalfile, "read")
	t = f["CollectionTree"]
	features = [key.GetName() for key in t.GetListOfBranches() if (
	('HGamMuonsAuxDyn' in key.GetName())
	or ('HGamAntiKt4EMTopoJetsAuxDyn' in key.GetName())
	or ('HGamPhotonsAuxDyn' in key.GetName())
	or ('HH2yybbEventInfoAuxDyn' in key.GetName())
	)]
	f.Close()
	return features


	def load_df(filename, features):
	df = pup.root2panda(filename, "CollectionTree", branches = features)
	return df


	def deltaR(eta1, eta2, phi1, phi2):
	import math
	'''
	Definition:
	-----------
	Function that calculates DR between two objects given their etas and phis
	Args:
	-----
	eta1 = float, eta of first object
	eta2 = float, eta of second object
	phi1 = float, phi of first object
	phi2 = float, phi of second object
	Output:
	-------
	deltaR = float, distance between the two objects
	'''
	DEta = abs(eta1-eta2)
	DPhi = math.acos(math.cos( abs( phi1-phi2 ) ) ) # hack to avoid \|phi1-phi2\| larger than 180 degrees
	return math.sqrt( pow(DEta,2) + pow(DPhi,2) )


	# -- now we need to corrcet the jets by the muons, if the muons are inside of them
	# -- start by calculating the distances between each muon and each jet. If the distance is less than a threshold, correct the jet by the muon
	def apply_correction(DRMAX, bjets_2, muons_2):

	corr_bjets_2 = []
	for ev in xrange(len(bjets_2)): # loop thru events
	corr_bjets_2_perevent = []
	for jet in bjets_2[ev]: # loop thru jets
	muon_add = Muon() # save all muons before adding them to the jet
	if (len(muons_2[ev]) != 0): # if there is at least one muon, do correction
	for muon in muons_2[ev]: # loop thru muons
	dr = deltaR(jet.eta, muon.eta, jet.phi, muon.phi) # calculate dr between muon and original jet
	if ((dr < DRMAX) and (muon.lv.Pt() > 4000)): # if it's closer that the fixed DR and pT>4GeV, add in the muon
	muon_add.lv += muon.lv
	corr_bjets_2_perevent.append(Jet(
	(jet.lv + muon_add.lv).pt(),
	(jet.lv + muon_add.lv).eta(),
	(jet.lv + muon_add.lv).phi(),
	(jet.lv + muon_add.lv).m(),
	jet.btag
	))
	corr_bjets_2.append(corr_bjets_2_perevent)

	return corr_bjets_2

	def get_composite_sys(bjets_2, photons_2, bjets_2_corr04, bjets_2_corr02, bjets_2_corr01):

	bb = []
	yy = []
	yybb = []
	bb04 = []
	yybb04 = []
	bb02 = []
	yybb02 = []
	bb01 = []
	yybb01 = []

	for ev in xrange(len(bjets_2)):
	bb.append(bjets_2[ev][0].lv + bjets_2[ev][1].lv)
	yy.append(photons_2[ev][0].lv + photons_2[ev][1].lv)
	yybb.append(bjets_2[ev][0].lv + bjets_2[ev][1].lv + photons_2[ev][0].lv + photons_2[ev][1].lv)

	bb04.append(bjets_2_corr04[ev][0].lv + bjets_2_corr04[ev][1].lv)
	yybb04.append(bjets_2_corr04[ev][0].lv + bjets_2_corr04[ev][1].lv + photons_2[ev][0].lv + photons_2[ev][1].lv)
	bb02.append(bjets_2_corr02[ev][0].lv + bjets_2_corr02[ev][1].lv)
	yybb02.append(bjets_2_corr02[ev][0].lv + bjets_2_corr02[ev][1].lv + photons_2[ev][0].lv + photons_2[ev][1].lv)
	bb01.append(bjets_2_corr01[ev][0].lv + bjets_2_corr01[ev][1].lv)
	yybb01.append(bjets_2_corr01[ev][0].lv + bjets_2_corr01[ev][1].lv + photons_2[ev][0].lv + photons_2[ev][1].lv)

	return bb, yy, yybb, bb04, yybb04, bb02, yybb02, bb01, yybb01


	def dr_plots(corrected_recomass_04, corrected_recomass_02, corrected_recomass_01, recomass, xmin = 80, xmax = 160,
	bins = 20, sys = '{bb}'):

	hista = [corrected_recomass_04[ev].M()/1000 for ev in xrange(len(corrected_recomass_04))]
	histb = [corrected_recomass_02[ev].M()/1000 for ev in xrange(len(corrected_recomass_02))]
	histd = [corrected_recomass_01[ev].M()/1000 for ev in xrange(len(corrected_recomass_01))]
	histc = [recomass[ev].M()/1000 for ev in xrange(len(recomass))]

	a, bina = np.histogram(hista, bins = 1000, range = (xmin, xmax), normed=True)
	b, binb = np.histogram(histb, bins = 1000, range = (xmin, xmax), normed=True)
	d, bind = np.histogram(histd, bins = 1000, range = (xmin, xmax), normed=True)
	c, binc = np.histogram(histc, bins = 1000, range = (xmin, xmax), normed=True)

	plt.hist(hista, label = 'Corrected All Muons, DR = 0.4', linewidth = 1,
	bins = np.linspace(xmin, xmax, bins), histtype = 'step', alpha = 0.5, color = 'blue', normed = True)

	plt.hist(histb, label = 'Corrected All Muons, DR = 0.2', linewidth = 1,
	bins = np.linspace(xmin, xmax, bins), histtype = 'step', alpha = 0.5, color = 'green', normed = True)

	plt.hist(histd, label = 'Corrected All Muons, DR = 0.1', linewidth = 1,
	bins = np.linspace(xmin, xmax, bins), histtype = 'step', alpha = 0.5, color = 'red', normed = True)

	plt.hist(histc, label = 'No Correction', linewidth = 1,
	bins = np.linspace(xmin, xmax, bins), histtype = 'step', alpha = 0.5, color = 'black', normed = True)


	bincentersa = 0.5*(bina[1:]+bina[:-1])
	bincentersb = 0.5*(binb[1:]+binb[:-1])
	bincentersd = 0.5*(bind[1:]+bind[:-1])
	bincentersc = 0.5*(binc[1:]+binc[:-1])
	(xfa,yfa),coeffa,erra,_ = fit.fit(fit.crystal_ball, bincentersa, a)
	(xfb,yfb),coeffb,errb,_ = fit.fit(fit.crystal_ball, bincentersb, b)
	(xfd,yfd),coeffd,errd,_ = fit.fit(fit.crystal_ball, bincentersd, d)
	(xfc,yfc),coeffc,errc,_ = fit.fit(fit.crystal_ball, bincentersc, c)

	plt.plot(xfa,yfa, color = 'blue')
	plt.plot(xfb,yfb, color = 'green')
	plt.plot(xfd,yfd, color = 'red')
	plt.plot(xfc,yfc, color = 'black')

	#plt.xlim(xmin=80, xmax=160)
	#plt.ylim( ymax = 0.015)
	plt.legend(loc = 'best', fontsize = 10)
	plt.ylabel('Normalized')
	plt.xlabel(r'$m_{}\ (GeV)$'.format(sys))
	plt.title(r'Invariant Mass of ${}$ system, 2-Tag Category'.format(sys))
	plt.show()

	print r'DR = 0.4, Peak = {0:.2f} +- {1:.2f} GeV, Sigma = {2:.2f} +- {3:.2f} GeV'.format(coeffa[3], erra[3], coeffa[4], erra[4] )
	print r'DR = 0.2, Peak = {0:.2f} +- {1:.2f} GeV, Sigma = {2:.2f} +- {3:.2f} GeV'.format(coeffb[3], errb[3], coeffb[4], errb[4])
	print r'DR = 0.1, Peak = {0:.2f} +- {1:.2f} GeV, Sigma = {2:.2f} +- {3:.2f} GeV'.format(coeffd[3], errd[3], coeffd[4], errd[4])
	print r'No Correction, Peak = {0:.2f} +- {1:.2f} GeV, Sigma = {2:.2f} +- {3:.2f} GeV'.format(coeffc[3], errc[3], coeffc[4], errc[4])


	def DW_variables(mass_window, bb04, yy, yybb04, bjets_2_corr04, photons_2):
	import math
	# -- start calculating values in David Wardrope's paper
	m12 = [bb04[ev].m() for ev in xrange(len(bb04))]
	m34 = [yy[ev].m() for ev in xrange(len(yy))]
	mTX = [yybb04[ev].m() for ev in xrange(len(yybb04))]
	pTX = [yybb04[ev].pt() for ev in xrange(len(yybb04))]
	yX = [yybb04[ev].eta() for ev in xrange(len(yybb04))] # units problem?

	# -- this crap below should follow exactly what they do in their code...
	q1_lab = bb04 # bb system
	q2_lab = yy # yy system
	X_lab = [(q1_lab[ev] + q2_lab[ev]) for ev in xrange(len(q1_lab))] # yybb system

	# -- BOOST: bb, yy, yybb, individual jets and photons --> to the rest frame of the yybb system
	[q1_lab[ev].Boost(-X_lab[ev].BoostVector().x, -X_lab[ev].BoostVector().y, -X_lab[ev].BoostVector().z)
	for ev in xrange(len(q1_lab))]
	q1 = q1_lab


	[bjets_2_corr04[ev][0].lv.Boost(-X_lab[ev].BoostVector().x, -X_lab[ev].BoostVector().y, -X_lab[ev].BoostVector().z)
	for ev in xrange(len(bjets_2_corr04))]
	q11 = [bjets_2_corr04[ev][0].lv for ev in xrange(len(bjets_2_corr04))]


	[bjets_2_corr04[ev][1].lv.Boost(-X_lab[ev].BoostVector().x, -X_lab[ev].BoostVector().y, -X_lab[ev].BoostVector().z)
	for ev in xrange(len(bjets_2_corr04))]
	q12 = [bjets_2_corr04[ev][1].lv for ev in xrange(len(bjets_2_corr04))]


	[q2_lab[ev].Boost(-X_lab[ev].BoostVector().x, -X_lab[ev].BoostVector().y, -X_lab[ev].BoostVector().z)
	for ev in xrange(len(q2_lab))]
	q2 = q1_lab


	[photons_2[ev][0].lv.Boost(-X_lab[ev].BoostVector().x, -X_lab[ev].BoostVector().y, -X_lab[ev].BoostVector().z)
	for ev in xrange(len(photons_2))]
	q21 = [photons_2[ev][0].lv for ev in xrange(len(photons_2))]


	[photons_2[ev][1].lv.Boost(-X_lab[ev].BoostVector().x, -X_lab[ev].BoostVector().y, -X_lab[ev].BoostVector().z)
	for ev in xrange(len(photons_2))]
	q22 = [photons_2[ev][1].lv for ev in xrange(len(photons_2))]


	# -- cos(\Theta^*)
	nZ = Vector3(0., 0., 1.) # useless, because that's already where the angles are calculated from, but ok
	cos_theta_star = [math.cos( (Vector3(q1[ev].px, q1[ev].py, q1[ev].pz)).theta()-nZ.theta() ) for ev in xrange(len(q1))]

	# -- normal vector to bb plane
	n1 = [ Vector3(
	(Vector3(q11[ev].px, q11[ev].py, q11[ev].pz)).cross( Vector3(q12[ev].px, q12[ev].py, q12[ev].pz) ).Unit() ) for ev in xrange(len(q11))]

	# -- normal vector to yy plane
	n2 = [ Vector3(
	(Vector3(q21[ev].px, q21[ev].py, q21[ev].pz)).cross( Vector3(q22[ev].px, q22[ev].py, q22[ev].pz) ).Unit() )
	for ev in xrange(len(q21))]

	# -- normal vector to the plane that contains the z axis and the first Higgs boson
	nSC = [ Vector3(
	nZ.cross( Vector3(q1[ev].px, q1[ev].py, q1[ev].pz) ).Unit() )
	for ev in xrange(len(q1))]

	# -- acos actually calculates the angle (why a - sign??)
	# -- the first factor just calculates the sign with respect to the first Higgs boson (the one that -->bb)
	phi = [np.sign( Vector3(q1[ev].px, q1[ev].py, q1[ev].pz).dot(n1[ev].cross(n2[ev])) ) * math.acos(-n1[ev].dot(n2[ev])) for ev in xrange(len(q1))]

	# -- I don't even know
	phi1 = [np.sign( Vector3(q1[ev].px, q1[ev].py, q1[ev].pz).dot(n1[ev].cross(nSC[ev])) ) * math.acos(n1[ev].dot(nSC[ev])) for ev in xrange(len(q1))]
	phi1_plot = abs(np.array(phi1))-(math.pi/2)

	# -- BOOST: individual jets and photons --> to the rest frame of the corresponding parent Higgs
	[bjets_2_corr04[ev][0].lv.Boost(-q1_lab[ev].BoostVector().x, -q1_lab[ev].BoostVector().y, -q1_lab[ev].BoostVector().y ) for ev in xrange(len(bjets_2_corr04))]
	q11_H1 = [bjets_2_corr04[ev][0].lv for ev in xrange(len(bjets_2_corr04))]

	[photons_2[ev][0].lv.Boost(-q2_lab[ev].BoostVector().x, -q2_lab[ev].BoostVector().y, -q2_lab[ev].BoostVector().y ) for ev in xrange(len(photons_2))]
	q21_H2 = [photons_2[ev][0].lv for ev in xrange(len(photons_2))]

	# -- cos(angle) between the first Higgs in the yybb rest frame and the first b jet in the rest fram of its parent Higgs (???)
	cos_theta1 = [Vector3(q1[ev].px, q1[ev].py, q1[ev].pz).dot(Vector3(q11_H1[ev].px, q11_H1[ev].py, q11_H1[ev].pz)) / math.sqrt(Vector3(q1[ev].px, q1[ev].py, q1[ev].pz).mag2() * Vector3(q11_H1[ev].px, q11_H1[ev].py, q11_H1[ev].pz).mag2()) for ev in xrange(len(q1))]

	# -- cos(angle) between the second Higgs in the yybb rest frame and the first y in the rest fram of its parent Higgs (???)
	cos_theta2 = [Vector3(q2[ev].px, q2[ev].py, q2[ev].pz).dot(Vector3(q21_H2[ev].px, q21_H2[ev].py, q21_H2[ev].pz)) / math.sqrt(Vector3(q2[ev].px, q2[ev].py, q2[ev].pz).mag2() * Vector3(q21_H2[ev].px, q21_H2[ev].py, q21_H2[ev].pz).mag2()) for ev in xrange(len(q2))]

	return m12, m34, mTX, pTX, yX, cos_theta_star, phi, phi1_plot, cos_theta1, cos_theta2




	def plot_DW_vars(val_dict, signalfile, bkgfile, key, bin_dict):

	bins = bin_dict[key]
	_ = plt.hist(np.array(val_dict[signalfile][key])[val_dict[signalfile]['mass_window']], bins = bins, alpha = 0.5, histtype = 'stepfilled', normed = True, label = 'X350')
	_ = plt.hist(np.array(val_dict[bkgfile][key])[val_dict[bkgfile]['mass_window']], bins = bins, alpha = 1, histtype = 'step', normed = True, label = 'yybb',
	color = 'red', hatch = '/////')
	plt.xlabel(key)
	plt.ylabel('Normalized')
	plt.legend()
	plt.show()




	if __name__ == '__main__':

	import argparse
	parser = argparse.ArgumentParser()
	parser.add_argument("signalfile", help="input .root file name for signal")
	parser.add_argument("bkgfile", help="input .root file name for background")
	args = parser.parse_args()

	sys.exit( main(args.signalfile, args.bkgfile) )