Exemple #1
0
    def execute(self,parser,treename=''):
        """
        Main function in the miniSL class -- does all the reading/writing
        of the flat ntuples.
        Just added the systematics functionality, so we've made this
        a hybrid class that can run over the nominal tree and 
        trees for each systematics.

        @param treename The treename that we need to access for data
                        This is now an option because of the systematics
        """
        timeStamp  = strftime("%d%b%Y",localtime())

        if not treename:
            treename       = info.treename()   # 'nominal'
            systname       = ''                # for naming the output file
            cfg_name       = 'miniSL'          # which part of the config file to access
        else:
            systname       = '_'+treename
            cfg_name       = 'systematics'     # which part of the config file to access

        ## -- Configuration -- ##
        p_sel_script = parser.get(cfg_name,'selection')       # ex. pyMiniSL/SelectionBase.py
        p_nEvents    = int(parser.get(cfg_name,'nEvents'))    # ex. -1
        p_firstEvent = int(parser.get(cfg_name,'firstevent')) # ex. 0
        p_cutsfile   = parser.get(cfg_name,'cutsfile')        # ex. share/pre2loose_cuts.txt
        p_inputfile  = parser.get(cfg_name,'inputfile')       # ex. share/miniSL_ntuples.txt
        p_outputname = parser.get(cfg_name,'outputname')      # ex. loose
        ## ------------------- ##

        files              = open(p_inputfile,'r').readlines()
        sel_path,selection = p_sel_script.split('.py')[0].split('/')
        selClassName       = selection[0].upper()+selection[1:]
        LUMI               = info.LUMI()
        signalIDs          = [i for key in info.dsids()['signal'].keys() for i in info.dsids()['signal'][key]] # TTS and XX signals
        ttbarIDs           = [i for i in info.dsids()['background']['ttbar'].keys()]

        # Dynamic import the file without needing any hard-coding
        # This way, if someone makes a custom selection, this doesn't need to be changed.
        # Assumes that the class is the same as the filename just capitalized!
        pyEventFile = importlib.import_module(sel_path+"."+selection)
        pyEventSel  = getattr(pyEventFile,selClassName)

        loggingLEVEL = logging.getLogger().getEffectiveLevel()
        logging.info(" -- In file       {0}".format(os.path.basename(__file__)))
        logging.info(" -- Files: ")
        for i in files:
            logging.info("           "+i.rstrip('\n'))


        for ff,file in enumerate(files):

            ## Open the current file
            f = ROOT.TFile.Open("{0}".format(file.strip()))

            if not f or f.IsZombie():
                print
                print " FILE {0} DOES NOT EXIST".format(f.GetName())
                print " Continuing to next file in the loop. "
                print

                continue

            ## Load the TTree
            # For nominal, treename is an empty string
            # For systematics, pre2preSelection: it's the treename for a systematic
            #                  SelectionBase:    it's a dummy string
            if treename == 'systs':
                treename = f.GetListOfKeys()[0].GetName()  # only 1 tree and it's the systematic
                systname = '_'+treename
            elif treename != 'nominal':
                treename = f.GetListOfKeys()[0].GetName()  # only 1 tree and it's the systematic
                systname = '_'+treename

            t = f.Get(treename)
            t.GetEntry(0)       # to get the mcChannelNumber

            samplename   = config.getSampleName(dsid=t.mcChannelNumber,root_tree=t)
            name         = samplename['name']
            isMC         = samplename['isMC']
            sample_key   = samplename['key']

            #  check file in case we returned a different name
            tmp_r_filename = f.GetName().split('/')[-1]
            if 'noname' not in name and not tmp_r_filename.startswith(name) and not tmp_r_filename.startswith('user'):
                if tmp_r_filename.startswith(sample_key) and not name.startswith(sample_key):
                    name = sample_key
                    # ex. 'ttW_Np0' returned from config.dsid2name() but the file name
                    #     is 'ttbarV'; 'else' statement will return 'tt', not 'ttbarV'
                else:
                    candidate_name = ''
                    matched_name   = ''
                    for letter in name:
                        matched_name+=letter
                        if tmp_r_filename.startswith(matched_name):
                            candidate_name = matched_name
                        else:
                            break
                    if candidate_name.endswith('_'):
                        name = candidate_name.rstrip('_')
                    else:
                        name = candidate_name
            # for files that start with 'user.xyz' (from the grid)
            if tmp_r_filename.startswith('user'):
                tmp_r_filename = tmp_r_filename.split(".") # to get the file ID numbers
                tmp_name = [piece for piece in tmp_r_filename if piece.strip('_').isdigit()]
                name     = name+'_'+''.join(tmp_name)

            ## New File Stuff (the one to which we're writing)
            newfilename  = "data/{0}{1}_{2}.root".format(name,systname,p_outputname)
            newfile      = ROOT.TFile(newfilename,"RECREATE")
            self.newtree = ROOT.TTree(treename,"Event Information")

            ## Metadata
            pma_tag  = getoutput("git describe --tags")    # e.g., 'v2.0.0-20-g493beba'

            metadata = ROOT.TList()
            metadata.Add( ROOT.TObjString("Date Created: {0}".format(timeStamp)) )
            metadata.Add( ROOT.TObjString("PyMiniAna:    {0}".format(pma_tag)) )
            metadata.Add( ROOT.TObjString("Selection:    {0}".format(p_outputname)) )
            metadata.Add( ROOT.TObjString("Original File:{0}".format(file.strip())) )
            metadata.Add( ROOT.TObjString("Selection:    {0}".format(p_outputname)) )
            metadata.SetName("PyMetadata")
            self.newtree.GetUserInfo().Add(metadata)

            ## Log and print to the console before doing anything
            print " ++ Running over file:  {0}".format(file.strip())
            print " ++ Producing new file: {0}{1}_{2}.root".format(name,systname,p_outputname)
            logging.info(" ----")
            logging.info(" {0} Selection".format(p_outputname.title()))
            logging.info(" File: {0}".format(file.strip()))
            logging.info(" {0}".format(timeStamp))
            logging.info(" ---- \n")

           # ----------------------------------------- #
           # Initialize the branches                   #
           # ----------------------------------------- #
            self.initBranches()

           # ----------------------------------------- #
           # Event Loop                                #
           # ----------------------------------------- #
            maxEntries = t.GetEntries()
            if p_nEvents < 0 or p_nEvents > maxEntries:
                nEvents = maxEntries
                logging.info(" Selected nEvents = {0}; running {1}".format(nEvents,maxEntries))
            else:
                nEvents = p_nEvents

            ## -- Load the selection script
            evtSel = pyEventSel(t,f,parser)  # imported using 'importlib' above
            evtSel.initialize()              # setup some of the configurations

            ## -- While-loop (no list generation with for loop)
            entry = p_firstEvent    # in case the user has split the file into pieces
            while entry < nEvents:

                #### -- APPLY THE SELECTION -- ####
                results     = evtSel.event_loop(entry)
                vars        = results['objects']
                passedEvent = results['result']
                #### ------------------------- ####
                if not passedEvent:
                    logging.info(" ---- ")
                    logging.info(" miniSL::Event {0} Failed Selection ".format(entry))
                    logging.info(" ---- ")
                    entry+=1               # iterate entry to go to the next one
                    continue

                fatJets      = vars['fatjets']
                rcJets       = vars['rcjets']
                resolvedJets = vars['resjets']
                bJets        = vars['bjets']
                lepton       = vars['lepton']
                neutrino     = vars['nu']
                jets         = vars['jets']
                tjets        = [] #vars['tjets']
                MET          = vars['met']
                HT           = vars['HT']
                MC           = vars['MC']
                eventInfo    = vars['eventinfo']

                # ----------------------------------------- #
                # Fill the branches                         #
                # ----------------------------------------- #
                # Set the branches for each entry (element of a list)
                # hadtops,hadtops_subs,leptops,smallRjets,leptons,neutrinos,bTagCats
                self.m_isBoosted[0]     = int(eventInfo['isBoosted'])
                self.m_isResolved[0]    = int(eventInfo['isResolved'])

        ## boosted, hadronic W
                if len(fatJets) > 0:
                    for hadW in fatJets:
                        self.m_ljet_pt.push_back(hadW.Pt())
                        self.m_ljet_eta.push_back(hadW.Eta())
                        self.m_ljet_phi.push_back(hadW.Phi())
                        self.m_ljet_e.push_back(hadW.E())
                        self.m_ljet_m.push_back(hadW.M())
                        self.m_ljet_D2.push_back(hadW.D2)
                        self.m_ljet_tau32_wta.push_back(hadW.tau32_wta)
                        self.m_ljet_isSmoothTop50.push_back(hadW.isSmoothTop50)
                        self.m_ljet_isSmoothTop80.push_back(hadW.isSmoothTop80)
                        self.m_ljet_isWmed.push_back(hadW.isWmed)
                        self.m_ljet_isWtight.push_back(hadW.isWtight)
                        self.m_ljet_isGood.push_back(hadW.isGood)

                        # suppress fatjet output variables.
                        # These are extras we don't need in the analysis,
                        # but we might need later for control plots
                        if not self.min_fatjet_vars:
                            self.m_ljet_track_m.push_back(hadW.track_m)
                            self.m_ljet_track_pt.push_back(hadW.track_pt)
                            self.m_ljet_trackjet_btag.push_back(hadW.trackjet_btag)
                            self.m_ljet_LHtop.push_back(hadW.LHtop)
                            self.m_ljet_LHw.push_back(hadW.LHw)
                            self.m_ljet_C2.push_back(hadW.C2)
                            self.m_ljet_d12.push_back(hadW.d12)
                            self.m_ljet_d23.push_back(hadW.d23)
                            self.m_ljet_tau1.push_back(hadW.tau1)
                            self.m_ljet_tau2.push_back(hadW.tau2)
                            self.m_ljet_tau3.push_back(hadW.tau3)
                            self.m_ljet_tau21.push_back(hadW.tau21)
                            self.m_ljet_tau32.push_back(hadW.tau32)
                            self.m_ljet_tau1_wta.push_back(hadW.tau1_wta)
                            self.m_ljet_tau2_wta.push_back(hadW.tau2_wta)
                            self.m_ljet_tau3_wta.push_back(hadW.tau3_wta)
                            self.m_ljet_tau21_wta.push_back(hadW.tau21_wta)
                            self.m_ljet_isPtLowTop50.push_back(hadW.isPtLowTop50)
                            self.m_ljet_isPtHighTop50.push_back(hadW.isPtHighTop50)
                            self.m_ljet_isPtLowTop80.push_back(hadW.isPtLowTop80)
                            self.m_ljet_isPtHighTop80.push_back(hadW.isPtHighTop80)
                            self.m_ljet_isZmed.push_back(hadW.isZmed)
                            self.m_ljet_isZtight.push_back(hadW.isZtight)

                if len(rcJets) > 0:
                    rc_sub_pts     = np.zeros(self.max_Nobjects,dtype=float)
                    rc_sub_etas    = np.zeros(self.max_Nobjects,dtype=float)
                    rc_sub_phis    = np.zeros(self.max_Nobjects,dtype=float)
                    rc_sub_es      = np.zeros(self.max_Nobjects,dtype=float)
                    rc_sub_ms      = np.zeros(self.max_Nobjects,dtype=float)
                    rc_sub_mv2c20s = np.zeros(self.max_Nobjects,dtype=float)
                    for hadW in rcJets:
                        rctau = sum(s.Pt() for s in hadW.subjets)
                        rctau = rctau / max(s.Pt() for s in hadW.subjets)
                        self.m_rcjet_pt.push_back(hadW.Pt())
                        self.m_rcjet_eta.push_back(hadW.Eta())
                        self.m_rcjet_phi.push_back(hadW.Phi())
                        self.m_rcjet_e.push_back(hadW.E())
                        self.m_rcjet_m.push_back(hadW.M())
                        self.m_rcjet_D2.push_back(hadW.D2)
                        self.m_rcjet_rctau.push_back(rctau)
                        self.m_rcjet_nsub.push_back(len(hadW.subjets))

                        for s,sub in enumerate(hadW.subjets):
                            rc_sub_pts.push_back(sub.Pt())
                            rc_sub_etas.push_back(sub.Eta())
                            rc_sub_phis.push_back(sub.Phi())
                            rc_sub_es.push_back(sub.E())
                            rc_sub_ms.push_back(sub.M())
                            rc_sub_mv2c20s.push_back(sub.mv2c20)

                        self.m_rcjet_sub_pt.push_back(rc_sub_pts)
                        self.m_rcjet_sub_eta.push_back(rc_sub_etas)
                        self.m_rcjet_sub_phi.push_back(rc_sub_phis)
                        self.m_rcjet_sub_e.push_back(rc_sub_es)
                        self.m_rcjet_sub_m.push_back(rc_sub_ms)
                        self.m_rcjet_sub_mv2c20.push_back(rc_sub_mv2c20s)

                        # clear subjet vectors for next rcjet
                        rc_sub_pts.resize(0)
                        rc_sub_etas.resize(0)
                        rc_sub_phis.resize(0)
                        rc_sub_es.resize(0)
                        rc_sub_ms.resize(0)
                        rc_sub_mv2c20s.resize(0)

        ## resolved, hadronic W
                if len(resolvedJets) > 0:
                    for hadW in resolvedJets:
                        self.m_resolvedW_pt.push_back(hadW.Pt())
                        self.m_resolvedW_eta.push_back(hadW.Eta())
                        self.m_resolvedW_phi.push_back(hadW.Phi())
                        self.m_resolvedW_e.push_back(hadW.E())
                        self.m_resolvedW_m.push_back(hadW.M())

        ## b-jets
                self.m_btags_n[0] = eventInfo['nbtags']
                if bJets:  # may not have bjets if it is pre-2-pre selection (not identified yet)
                    for b in bJets:  # Keeping only 2 jets as 'b-jets' (>=1 are actually b-tagged)
                        self.m_bjet_pt.push_back(b.Pt())
                        self.m_bjet_eta.push_back(b.Eta())
                        self.m_bjet_phi.push_back(b.Phi())
                        self.m_bjet_e.push_back(b.E())
                        self.m_bjet_jvt.push_back(b.jvt)
                        self.m_bjet_mv2c20.push_back(b.mv2c20)
                        if isMC: 
                            self.m_bjet_true_flavor.push_back(b.true_flavor)
                        else:
                            self.m_bjet_true_flavor.push_back(-1)
                    self.m_mass_lb[0]    = (lepton+bJets[0]).M()

        ## Small-r jets
                TRF_jetFlavors = []
                TRF_jetPts     = [] # Needs to be in GeV!
                TRF_jetEtas    = []

                for j in jets:
                    self.m_jet_pt.push_back(j.Pt())
                    self.m_jet_eta.push_back(j.Eta())
                    self.m_jet_phi.push_back(j.Phi())
                    self.m_jet_e.push_back(j.E())
                    self.m_jet_jvt.push_back(j.jvt)
                    self.m_jet_mv2c20.push_back(j.mv2c20)
#                     self.m_jet_track_m.push_back(j.track_m)
#                     self.m_jet_track_pt.push_back(j.track_pt)
                    self.m_jet_true_flavor.push_back(j.true_flavor)
                    TRF_jetFlavors.append(j.true_flavor)
                    TRF_jetPts.append(j.Pt()/1000.)      # [GeV]
                    TRF_jetEtas.append(j.Eta())
                # TRF weight for the event
                if isMC:
                    self.m_TRFWeight[0] = TRFWeight(TRF_jetFlavors,TRF_jetPts, TRF_jetEtas)
                else:
                    self.m_TRFWeight[0] = -1.

        ## Track jets
                for j in tjets:
                    self.m_tjet_pt.push_back(j.Pt())
                    self.m_tjet_eta.push_back(j.Eta())
                    self.m_tjet_phi.push_back(j.Phi())
                    self.m_tjet_e.push_back(j.E())
                    self.m_tjet_mv2c20.push_back(j.mv2c20)
                    self.m_tjet_numConstituents.push_back(j.numConstituents)
                    self.m_tjet_true_flavor.push_back(j.true_flavor)

        ## lepton (muon or electron)
        ##   Just lepton branches: lep_e/m/pt
                self.m_ejets[0]      = eventInfo['ejets']
                self.m_mujets[0]     = eventInfo['mujets']

                self.m_lep_pt[0]     = lepton.Pt()
                self.m_lep_eta[0]    = lepton.Eta()
                self.m_lep_phi[0]    = lepton.Phi()
                self.m_lep_e[0]      = lepton.E()
                self.m_lep_charge[0] = lepton.charge

        ## neutrino (from W mass constraint)
                self.m_nu_pt[0]   = neutrino.Pt()
                self.m_nu_eta[0]  = neutrino.Eta()
                self.m_nu_phi[0]  = neutrino.Phi()
                self.m_nu_e[0]    = neutrino.E()

        ## leptonic W
                lepW = lepton+neutrino
                self.m_leptonicW_pt[0]  = lepW.Pt()
                self.m_leptonicW_eta[0] = lepW.Eta()
                self.m_leptonicW_phi[0] = lepW.Phi()
                self.m_leptonicW_e[0]   = lepW.E()
                self.m_leptonicW_m[0]   = lepW.M()


                ## Build the TTbar and Hadronic W
                if selection!='grid2preSelection':  
                    # Only define the hadronic W candidate and TTbar
                    #  if we're beyond the pre-selection.
                    # For users accessing the resulting root file,
                    #  TTbar and the Hadronic W will be saved

                    if eventInfo['isBoosted']:
                        hadW = fatJets[0]       # Should only be 1 at this stage
                    else:
                        hadW = resolvedJets[0]

                    TTbar = vlq.getTT(hadW,lepW,bJets) # returns dictionary of 'hadT' and 'lepT'
                    hadT  = TTbar['hadT'] # Note: These are defined for every sample,
                    lepT  = TTbar['lepT'] # not just the TT samples (there is a T and Tbar
                                          # candidate in each event that passes the cuts!)

                    self.m_hadronicW_pt[0]  = hadW.Pt()
                    self.m_hadronicW_eta[0] = hadW.Eta()
                    self.m_hadronicW_phi[0] = hadW.Phi()
                    self.m_hadronicW_e[0]   = hadW.E()
                    self.m_hadronicW_m[0]   = hadW.M()

                    self.m_hadronicT_pt[0]  = hadT.Pt()
                    self.m_hadronicT_eta[0] = hadT.Eta()
                    self.m_hadronicT_phi[0] = hadT.Phi()
                    self.m_hadronicT_e[0]   = hadT.E()
                    self.m_hadronicT_m[0]   = hadT.M()

                    self.m_leptonicT_pt[0]  = lepT.Pt()
                    self.m_leptonicT_eta[0] = lepT.Eta()
                    self.m_leptonicT_phi[0] = lepT.Phi()
                    self.m_leptonicT_e[0]   = lepT.E()
                    self.m_leptonicT_m[0]   = lepT.M()

        ## MET & HT
                self.m_met_met[0]       = MET.met
                self.m_met_phi[0]       = MET.phi
                self.m_mtw[0]           = MET.mtw
                self.m_HT[0]            = HT

        ## Kinematic Relationships
                self.m_DeltaR_lepnu[0]      = lepton.DeltaR(neutrino)
                if selection!='grid2preSelection':
                    self.m_minDeltaR_lepbjet[0]   = min([b.DeltaR(lepton) for b in bJets])
                    self.m_DeltaR_bjet1bjet2[0]   = bJets[0].DeltaR(bJets[1])
                    self.m_minDeltaR_hadWbjet[0]  = min([b.DeltaR(hadW) for b in bJets])
                    self.m_DeltaM_hadtopleptop[0] = fabs(hadT.M() - lepT.M())

        ## MC INFO -- only signal and ttbar
                if MC:
                    t_t    = MC['truth_t']
                    t_tbar = MC['truth_tbar']
                    t_Wp   = MC['truth_wplus']
                    t_Wm   = MC['truth_wminus']
                    t_b    = MC['truth_b']
                    t_bbar = MC['truth_bbar']
                    t_lep  = MC['truth_lep']
                    t_nu   = MC['truth_nu']
                    t_q1   = MC['truth_q1']
                    t_q2   = MC['truth_q2']

                    self.m_truth_t_pt[0]       = t_t.Pt()
                    self.m_truth_t_eta[0]      = t_t.Eta()
                    self.m_truth_t_phi[0]      = t_t.Phi()
                    self.m_truth_t_e[0]        = t_t.E()
                    self.m_truth_tbar_pt[0]    = t_tbar.Pt()
                    self.m_truth_tbar_eta[0]   = t_tbar.Eta()
                    self.m_truth_tbar_phi[0]   = t_tbar.Phi()
                    self.m_truth_tbar_e[0]     = t_tbar.E()
                    self.m_truth_Wplus_pt[0]   = t_Wp.Pt()
                    self.m_truth_Wplus_eta[0]  = t_Wp.Eta()
                    self.m_truth_Wplus_phi[0]  = t_Wp.Phi()
                    self.m_truth_Wplus_e[0]    = t_Wp.E()
                    self.m_truth_Wminus_pt[0]  = t_Wm.Pt()
                    self.m_truth_Wminus_eta[0] = t_Wm.Eta()
                    self.m_truth_Wminus_phi[0] = t_Wm.Phi()
                    self.m_truth_Wminus_e[0]   = t_Wm.E()
                    self.m_truth_b_pt[0]       = t_b.Pt()
                    self.m_truth_b_eta[0]      = t_b.Eta()
                    self.m_truth_b_phi[0]      = t_b.Phi()
                    self.m_truth_b_e[0]        = t_b.E()
                    self.m_truth_bbar_pt[0]    = t_bbar.Pt()
                    self.m_truth_bbar_eta[0]   = t_bbar.Eta()
                    self.m_truth_bbar_phi[0]   = t_bbar.Phi()
                    self.m_truth_bbar_e[0]     = t_bbar.E()
                    self.m_truth_nu_pt[0]      = t_nu.Pt()
                    self.m_truth_nu_eta[0]     = t_nu.Eta()
                    self.m_truth_nu_phi[0]     = t_nu.Phi()
                    self.m_truth_nu_e[0]       = t_nu.E()
                    self.m_truth_nu_pdgId[0]   = t_nu.pdgId
                    self.m_truth_lep_pt[0]     = t_lep.Pt()
                    self.m_truth_lep_eta[0]    = t_lep.Eta()
                    self.m_truth_lep_phi[0]    = t_lep.Phi()
                    self.m_truth_lep_e[0]      = t_lep.E()
                    self.m_truth_lep_pdgId[0]  = t_lep.pdgId
                    self.m_truth_lep_charge[0] = t_lep.charge
                    self.m_truth_q1_pt[0]      = t_q1.Pt()
                    self.m_truth_q1_eta[0]     = t_q1.Eta()
                    self.m_truth_q1_phi[0]     = t_q1.Phi()
                    self.m_truth_q1_e[0]       = t_q1.E()
                    self.m_truth_q1_pdgId[0]   = t_q1.pdgId
                    self.m_truth_q2_pt[0]      = t_q2.Pt()
                    self.m_truth_q2_eta[0]     = t_q2.Eta()
                    self.m_truth_q2_phi[0]     = t_q2.Phi()
                    self.m_truth_q2_e[0]       = t_q2.E()
                    self.m_truth_q2_pdgId[0]   = t_q2.pdgId

        ## Event Info
                self.m_weight_mc[0]       = eventInfo['mcWeight']
                self.m_weight_pileup[0]   = eventInfo['pileupWeight']
                self.m_eventNumber[0]     = eventInfo['eventNumber']
                self.m_runNumber[0]       = eventInfo['runNumber']
                self.m_mcChannelNumber[0] = eventInfo['mcChannelNumber']
                self.m_mu[0]              = eventInfo['mu']
                self.m_weight_btag_60[0]  = eventInfo['weight_btag_60']
                self.m_weight_btag_70[0]  = eventInfo['weight_btag_70']
                self.m_weight_btag_77[0]  = eventInfo['weight_btag_77']
                self.m_weight_btag_85[0]  = eventInfo['weight_btag_85']
                self.m_weight_lept_eff[0] = eventInfo['weight_lep_eff']
                self.m_vlq_evtype[0]      = eventInfo['vlq_tag']
                self.m_AMI[0]             = eventInfo['nAMI']
                self.m_XSection[0]        = eventInfo['XSection']
                self.m_KFactor[0]         = eventInfo['KFactor']
                self.m_FilterEff[0]       = eventInfo['FilterEff']
                self.m_LUMI[0]            = LUMI['lumi']  # Defined Above

                self.newtree.Fill()


           # ----------------------------------------- #
           # Clear the ROOT Vectors                    #
           # ----------------------------------------- #
                self.clearVectors()

                ## increment the while loop
                entry+=1


           # ----------------------------------------- #
           # Write the new file                        #
           # ----------------------------------------- #
            newfile.Write()
            print
            print " New output file: {0}".format(newfilename)
            print " Ended at:        {0}".format(strftime("%c",localtime()))
            print
            newfile.Close()

        return
Exemple #2
0
    def execute(self,parser,treename=''):
        """
        Main function in the miniSL class -- does all the reading/writing
        of the flat ntuples.
        Just added the systematics functionality, so we've made this
        a hybrid class that can run over the nominal tree and 
        trees for each systematics.

        @param treename The treename that we need to access for data
                        This is now an option because of the systematics
        """
        self.timeStamp = strftime("%d%b%Y",localtime())
        cfg_name  = 'miniSL' if not treename else 'systematics'

        ## -- Configuration -- ##
        p_sel_script = parser.get(cfg_name,'selection')       # ex. pyMiniSL/SelectionBase.py
        p_nEvents    = int(parser.get(cfg_name,'nEvents'))    # ex. -1
        p_firstEvent = int(parser.get(cfg_name,'firstevent')) # ex. 0
        p_cutsfile   = parser.get(cfg_name,'cutsfile')        # ex. share/pre2loose_cuts.txt
        p_inputfile  = parser.get(cfg_name,'inputfile')       # ex. share/miniSL_ntuples.txt
        p_outputname = parser.get(cfg_name,'outputname')      # ex. loose
        p_isLoose    = config.str2bool(parser.get(cfg_name,'isLoose'))  # ex. False
        ## ------------------- ##

        files              = open(p_inputfile,'r').readlines()
        sel_path,selection = p_sel_script.split('.py')[0].split('/')
        self.selection     = selection  # for access later
        selClassName       = selection[0].upper()+selection[1:]
        self.LUMI          = info.LUMI()
        signalIDs          = [i for key in info.dsids()['signal'].keys() for i in info.dsids()['signal'][key]] # TTS and XX signals
        ttbarIDs           = [i for i in info.dsids()['background']['ttbar'].keys()]

        # Dynamic import the file without needing any hard-coding
        # Assumes that the class is the same as the filename just capitalized!
        pyEventFile = importlib.import_module(sel_path+"."+selection)
        pyEventSel  = getattr(pyEventFile,selClassName)

        loggingLEVEL = logging.getLogger().getEffectiveLevel()
        logging.info(" -- In file       {0}".format(os.path.basename(__file__)))
        logging.info(" -- Files: ")
        for i in files:
            logging.info("           "+i.rstrip('\n'))


        ## -- Load the selection script
        evtSel = pyEventSel(parser)      # imported using 'importlib' above
        evtSel.initialize(cfg_name)      # setup some of the configurations
 
       # ----------------------------------------- #
       # File Loop                                 #
       # ----------------------------------------- #
        total_num_files = len(files)
        for ff,file in enumerate(files):

            ## -- Open the current file
            f = ROOT.TFile.Open("{0}".format(file.strip()))

            if not f or f.IsZombie():
                print
                print " FILE {0} DOES NOT EXIST".format(f.GetName())
                print " Continuing to next file in the loop. "
                print
                continue

            ## -- Load the TTree
            #     For nominal, treename is an empty string
            #     For systematics, pre2preSelection: it's the treename for a systematic
            #                      SelectionBase:    it's a dummy string
            if cfg_name != 'miniSL':
                if treename == 'systs':
                    treename = f.GetListOfKeys()[0].GetName()  # only 1 tree and it's the systematic
                systname     = '_'+treename
                new_treename = treename
            else:
                if p_isLoose:
                    treename     = 'nominal_Loose'
                else:
                    treename = info.treename() # nominal
                new_treename = 'nominal'
                systname = ''

            t = f.Get(treename)
            try:
                t.GetEntry(0)       # to get the mcChannelNumber
            except AttributeError:
                print " File {0} does not have the GetEntry() attribute.".format(f.GetName())
                print " Skipping and going to the next file.\n"
                continue


            ## New File Stuff (the one to which we're writing)
            name,isMC    = self.getName(t)

            newfilename  = "data/{0}{1}_{2}.root".format(name,systname,p_outputname)
            newfile      = ROOT.TFile(newfilename,"RECREATE")
            self.newtree = ROOT.TTree(new_treename,"Event Information")

            ## Metadata
            self.setMetaData(p_outputname)

            ## Log and print to the console before doing anything
            print " ++ Running over file ({0}/{1}):  {2}".format(ff+1,total_num_files,file.strip())
            print " ++ Producing new file: {0}{1}_{2}.root".format(name,systname,p_outputname)
            logging.info(" ----")
            logging.info(" {0} Selection".format(p_outputname.title()))
            logging.info(" File: {0}".format(file.strip()))
            logging.info(" {0}".format(self.timeStamp))
            logging.info(" ---- \n")

           # ----------------------------------------- #
           # Initialize the branches                   #
           # ----------------------------------------- #
            self.initBranches()

           # ----------------------------------------- #
           # Event Loop                                #
           # ----------------------------------------- #
            maxEntries = t.GetEntries()
            if p_nEvents < 0 or p_nEvents > maxEntries:
                nEvents = maxEntries
                logging.info(" Selected nEvents = {0}; running {1}".format(nEvents,maxEntries))
            else:
                nEvents = p_nEvents

            ## -- While-loop (no list generation with for loop)
            evtSel.execute(t,f)     # setup the details for this particular file in selection

            entry = p_firstEvent    # in case the user has split the file into pieces
            while entry < nEvents:

                #### -- APPLY THE SELECTION -- ####
                results = evtSel.event_loop(entry)
                #### ------------------------- ####
                if not results['result']:
                    logging.info(" ---- ")
                    logging.info(" miniSL::Event {0} Failed Selection ".format(entry))
                    logging.info(" ---- ")
                    entry+=1               # iterate entry to go to the next one
                    continue

                self.saveEvent(results['objects'])

               # ----------------------------------------- #
               # Clear the ROOT Vectors                    #
               # ----------------------------------------- #
                self.clearVectors()

                ## increment the while loop
                entry+=1


           # ----------------------------------------- #
           # Write the new file                        #
           # ----------------------------------------- #
            newfile.Write()
            print
            print " New output file: {0}".format(newfilename)
            print " Ended at:        {0}".format(strftime("%c",localtime()))
            print
            newfile.Close()

        return