Python tree.Tree类代码示例

class JetTreeProducer(Analyzer):

    def beginLoop(self, setup):
        super(JetTreeProducer, self).beginLoop(setup)
        self.rootfile = TFile('/'.join([self.dirName,
                                        'jet_tree.root']),
                              'recreate')
        self.tree = Tree( self.cfg_ana.tree_name,
                          self.cfg_ana.tree_title )
        bookJet(self.tree, 'jet1')
        bookJet(self.tree, 'jet1_gen')
        bookJet(self.tree, 'jet2')
        bookJet(self.tree, 'jet2_gen')
        
    def process(self, event):
        self.tree.reset()
        if( len(event.rec_jets)>0 ):
            jet = event.rec_jets[0]
            fillJet(self.tree, 'jet1', jet)
            if jet.gen:
                fillJet(self.tree, 'jet1_gen', jet.gen)
        if( len(event.rec_jets)>1 ):
            jet = event.rec_jets[1]
            fillJet(self.tree, 'jet2', jet)
            if jet.gen:
                fillJet(self.tree, 'jet2_gen', jet.gen)
        self.tree.tree.Fill()
        
        
    def write(self, setup):
        self.rootfile.Write()
        self.rootfile.Close()

class LeptonTreeProducer(Analyzer):

    def beginLoop(self, setup):
        super(LeptonTreeProducer, self).beginLoop(setup)
        self.rootfile = TFile('/'.join([self.dirName,
                                        'tree.root']),
                              'recreate')
        self.tree = Tree( self.cfg_ana.tree_name,
                          self.cfg_ana.tree_title )
        bookLepton(self.tree, 'lep1')
        bookLepton(self.tree, 'lep2')
        

    def process(self, event):
        self.tree.reset()
        leptons = getattr(event, self.cfg_ana.leptons)
        if len(leptons) > 0:
            fillLepton(self.tree, 'lep1', leptons[0])
        if len(leptons) > 1:
            fillLepton(self.tree, 'lep2', leptons[1])
        self.tree.tree.Fill()
        
        
    def write(self, setup):
        self.rootfile.Write()
        self.rootfile.Close()

class SimpleTreeProducer(Analyzer):

    def beginLoop(self, setup):
        super(SimpleTreeProducer, self).beginLoop(setup)
        self.rootfile = TFile('/'.join([self.dirName,
                                        'tree.root']),
                              'recreate')
        self.tree = Tree( 'events', '')
        bookParticle(self.tree, 'recoil_gen')
        bookParticle(self.tree, 'recoil_visible_gen')
        bookParticle(self.tree, 'recoil_papas')
        bookParticle(self.tree, 'recoil_visible_papas')

        
    def process(self, event):
        self.tree.reset()
        fillParticle(self.tree, 'recoil_gen', event.recoil_gen)
        fillParticle(self.tree, 'recoil_visible_gen', event.recoil_visible_gen)
        if hasattr(event, 'recoil_papas'):
            fillParticle(self.tree, 'recoil_papas', event.recoil_papas)
            fillParticle(self.tree, 'recoil_visible_papas', event.recoil_visible_papas)
        self.tree.tree.Fill()

        
    def write(self, setup):
        self.rootfile.Write()
        self.rootfile.Close()

class IsoParticleTreeProducer(Analyzer):

    def beginLoop(self, setup):
        super(IsoParticleTreeProducer, self).beginLoop(setup)
        self.rootfile = TFile('/'.join([self.dirName,
                                        'tree.root']),
                              'recreate')
        self.tree = Tree( self.cfg_ana.tree_name,
                          self.cfg_ana.tree_title )
        bookIsoParticle(self.tree, 'ptc')

    def process(self, event):
        self.tree.reset()
        leptons = getattr(event, self.cfg_ana.leptons)
        pdgids = [211, 22, 130]
        for lepton in leptons:
            for pdgid in pdgids:
                iso = getattr(lepton, 'iso_{pdgid:d}'.format(pdgid=pdgid))
                for ptc in iso.on_ptcs:
                    self.tree.reset()
                    fillIsoParticle(self.tree, 'ptc', ptc, lepton)
                    self.tree.tree.Fill()
        
    def write(self, setup):
        self.rootfile.Write()
        self.rootfile.Close()

class JetPtcTreeProducer(Analyzer):
    '''Some class doc'''

    def beginLoop(self, setup):
        super(JetPtcTreeProducer, self).beginLoop(setup)
        self.rootfile = TFile('/'.join([self.dirName,
                                        'jet_tree.root']),
                              'recreate')
        self.tree = Tree( self.cfg_ana.tree_name,
                          self.cfg_ana.tree_title )
        bookJet(self.tree, 'jet1')
        bookJet(self.tree, 'jet1_rec')
        bookJet(self.tree, 'jet2')
        bookJet(self.tree, 'jet2_rec')
        var(self.tree, 'event')
        var(self.tree, 'lumi')
        var(self.tree, 'run')
        bookParticle(self.tree, 'ptc_gen')
        bookParticle(self.tree, 'ptc_rec')


    def process(self, event):
        self.tree.reset()
        if hasattr(event, 'eventId'): 
            fill(self.tree, 'event', event.eventId)
            fill(self.tree, 'lumi', event.lumi)
            fill(self.tree, 'run', event.run)
        elif hasattr(event, 'iEv'):
            fill(self.tree, 'event', event.iEv)
        jets = getattr(event, self.cfg_ana.jets)
        if( len(jets)>0 ):
            jet = jets[0]
            comp211 = jet.constituents.get(211, None)
            if comp211: 
                if comp211.num==2:
                    import pdb; pdb.set_trace()
            fillJet(self.tree, 'jet1', jet)
            if jet.match:
                fillJet(self.tree, 'jet1_rec', jet.match)
                # if jet.e()/jet.match.e() > 2.:
                #     import pdb; pdb.set_trace()
        if( len(jets)>1 ):
            jet = jets[1]
            fillJet(self.tree, 'jet2', jet)
            if jet.match:
                fillJet(self.tree, 'jet2_rec', jet.match)
        ptc = getattr(event, self.cfg_ana.particle)
        fillParticle(self.tree, 'ptc_gen', ptc[0])
        if ptc[0].match:
            fillParticle(self.tree, 'ptc_rec', ptc[0].match)
        self.tree.tree.Fill()
        
        
    def write(self, setup):
        self.rootfile.Write()
        self.rootfile.Close()

def create_tree(filename="test_tree.root"): 
    outfile = TFile(filename, 'recreate')
    tree = Tree('test_tree', 'A test tree')
    tree.var('var1')
    for i in range(100):
        tree.fill('var1', i)
        tree.tree.Fill()
    print 'creating a tree', tree.tree.GetName(),\
        tree.tree.GetEntries(), 'entries in',\
        outfile.GetName()
    outfile.Write()

class ParticleTreeProducer(Analyzer):
    '''Fills a TTree for particle-flow studies (experts only)'''
    
    def beginLoop(self, setup):
        super(ParticleTreeProducer, self).beginLoop(setup)
        self.rootfile = TFile('/'.join([self.dirName,
                                        'tree.root']),
                              'recreate')
        self.tree = Tree('particles', '')
        bookParticle(self.tree, 'ptc')
        bookCluster(self.tree, 'ptc_ecal')
        bookParticle(self.tree, 'ptc_match')
        var(self.tree, 'dr')
        bookParticle(self.tree, 'ptc_match_211')
        var(self.tree, 'dr_211')
        bookParticle(self.tree, 'ptc_match_130')
        var(self.tree, 'dr_130')
        bookParticle(self.tree, 'ptc_match_22')
        var(self.tree, 'dr_22')

    def process(self, event):
        particles = getattr(event, self.cfg_ana.particles)
        for ptc in particles: 
            self.tree.reset()
            fillParticle(self.tree, 'ptc', ptc)
            m211 = m22 = False
            # if hasattr(ptc, 'clusters'):
            #     # sim particle
            #     ecal = ptc.clusters.get('ecal_in', None)
            #     if ecal:
            #         fillCluster(self.tree, 'ptc_ecal', ecal)
            if hasattr(ptc, 'match') and ptc.match:
                fillParticle(self.tree, 'ptc_match', ptc.match)
                fill(self.tree, 'dr', ptc.dr)
            if hasattr(ptc, 'match_211') and ptc.match_211:
                m211 = True
                fillParticle(self.tree, 'ptc_match_211', ptc.match_211)
                fill(self.tree, 'dr_211', ptc.dr_211)
            if hasattr(ptc, 'match_130') and ptc.match_130:
                fillParticle(self.tree, 'ptc_match_130', ptc.match_130)
                fill(self.tree, 'dr_130', ptc.dr_130)
            if hasattr(ptc, 'match_22') and ptc.match_22:
                m22 = True
                fillParticle(self.tree, 'ptc_match_22', ptc.match_22)
                fill(self.tree, 'dr_22', ptc.dr_22)
            # if m22 and not m211:
            #     print event
            #     import pdb; pdb.set_trace()
            self.tree.tree.Fill()
        
        
    def write(self, setup):
        self.rootfile.Write()
        self.rootfile.Close()

class JetConeTreeProducer(Analyzer):

    def beginLoop(self, setup):
        super(JetConeTreeProducer, self).beginLoop(setup)
        self.rootfile = TFile('/'.join([self.dirName,
                                        'jet_tree.root']),
                              'recreate')
        self.tree = Tree( self.cfg_ana.tree_name,
                          self.cfg_ana.tree_title )
        bookJet(self.tree, 'papasjet')
        bookJet(self.tree, 'cmsjet')
        bookJet(self.tree, 'papas_control_jet')
        bookJet(self.tree, 'cms_control_jet')
        bookJet(self.tree, 'gen_jet')
        bookJet(self.tree, 'simtrack')
        var(self.tree, 'simtrack_len')
        for i in range(20):
            bookParticle(self.tree, 'simtrack_ptc'+str(i))
        bookParticle(self.tree, 'PFCandidate')

    def process(self, event):
        self.tree.reset()
        papasjet = getattr(event, self.cfg_ana.rec_jet, None)
        cmsjet = getattr(event, self.cfg_ana.pf_jet, None)
        papas_control_jet = getattr(event, self.cfg_ana.rec_control_jet, None)
        cms_control_jet = getattr(event, self.cfg_ana.pf_control_jet, None)
        gen_jet = getattr(event, self.cfg_ana.gen_jet, None)
        if papasjet:
            fillJet(self.tree, 'papasjet', papasjet)
        if cmsjet:
            fillJet(self.tree, 'cmsjet', cmsjet)
        if papas_control_jet:
            fillJet(self.tree, 'papas_control_jet', papas_control_jet)
        if cms_control_jet:
            fillJet(self.tree, 'cms_control_jet', cms_control_jet)
        if gen_jet:
            fillJet(self.tree, 'gen_jet', gen_jet)
        sim_track_jet = getattr(event, self.cfg_ana.sim_track_jet, None)
        if sim_track_jet:
            fillJet(self.tree, 'simtrack', sim_track_jet)
        sim_track_ptcs = getattr(event, self.cfg_ana.sim_track, None)
        if sim_track_ptcs:
            for i in range(min(len(sim_track_ptcs), 20)):
                fillParticle(self.tree, 'simtrack_ptc'+str(i), sim_track_ptcs[i])
        fill(self.tree, 'simtrack_len', len(sim_track_ptcs))
        pfcandidates = getattr(event, self.cfg_ana.pfcandidates, None)
        if pfcandidates:
            fillParticle(self.tree, 'PFCandidate', pfcandidates[0])
        self.tree.tree.Fill()

    def write(self, setup):
        self.rootfile.Write()
        self.rootfile.Close()

class TTbarTreeProducer(Analyzer):

    def beginLoop(self, setup):
        super(TTbarTreeProducer, self).beginLoop(setup)
        self.rootfile = TFile('/'.join([self.dirName,
                                        'tree.root']),
                              'recreate')
        self.tree = Tree( 'events', '')
        bookParticle(self.tree, 'jet1')
        bookParticle(self.tree, 'jet2')
        bookParticle(self.tree, 'jet3')
        bookParticle(self.tree, 'jet4')
        bookParticle(self.tree, 'm3')
        bookMet(self.tree, 'met')
        bookLepton(self.tree, 'muon')
        bookLepton(self.tree, 'electron')

    def process(self, event):
        self.tree.reset()
        muons = getattr(event, self.cfg_ana.muons)        
        electrons = getattr(event, self.cfg_ana.electrons)

        if len(muons)==0 and len(electrons)==0:
            return # NOT FILLING THE TREE IF NO

        if len(muons)==1 and len(electrons)==0:
            fillLepton(self.tree, 'muon', muons[0])

        elif len(electrons)==1 and len(muons)==0:
            fillLepton(self.tree, 'electron', electrons[0])

        else:
            return # NOT FILLING THE TREE IF MORE THAN 1 LEPTON

        jets = getattr(event, self.cfg_ana.jets_30)
        if len(jets)<3:
            return # NOT FILLING THE TREE IF LESS THAN 4 JETS
        for ijet, jet in enumerate(jets):
            if ijet==4:
                break
            fillParticle(self.tree, 'jet{ijet}'.format(ijet=ijet+1), jet)
        m3 = getattr(event, self.cfg_ana.m3)
        if m3: 
            fillParticle(self.tree, 'm3', m3)

        met = getattr(event, self.cfg_ana.met)
        fillMet(self.tree, 'met', met)
        self.tree.tree.Fill()
        
    def write(self, setup):
        self.rootfile.Write()
        self.rootfile.Close()

class ZHTreeProducer(Analyzer):

    def beginLoop(self, setup):
        super(ZHTreeProducer, self).beginLoop(setup)
        self.rootfile = TFile('/'.join([self.dirName,
                                        'tree.root']),
                              'recreate')
        self.tree = Tree( 'events', '')
        if hasattr(self.cfg_ana, 'recoil'):
            bookParticle(self.tree, 'recoil')
        if hasattr(self.cfg_ana, 'zeds'):  
            bookZed(self.tree, 'zed')
        self.taggers = ['b', 'bmatch', 'bfrac']
        bookJet(self.tree, 'jet1', self.taggers)
        bookJet(self.tree, 'jet2', self.taggers)
        bookHbb(self.tree, 'higgs')
        bookParticle(self.tree, 'misenergy')
        var(self.tree, 'n_nu')
       
    def process(self, event):
        self.tree.reset()
        if hasattr(self.cfg_ana, 'recoil'):
            recoil = getattr(event, self.cfg_ana.recoil)    
            fillParticle(self.tree, 'recoil', recoil)
        if hasattr(self.cfg_ana, 'zeds'):  
            zeds = getattr(event, self.cfg_ana.zeds)
            if len(zeds)>0:
                zed = zeds[0]
                fillZed(self.tree, 'zed', zed)
        misenergy = getattr(event, self.cfg_ana.misenergy)
        fillParticle(self.tree, 'misenergy', misenergy )        
        jets = getattr(event, self.cfg_ana.jets)
        for ijet, jet in enumerate(jets):
            if ijet == 2:
                break
            fillJet(self.tree, 'jet{ijet}'.format(ijet=ijet+1), jet, self.taggers)
        higgses = getattr(event, self.cfg_ana.higgses)
        if len(higgses)>0:
            higgs = higgses[0]
            # if higgs.m() < 30:
            #    import pdb; pdb.set_trace()
            fillHbb(self.tree, 'higgs', higgs)
        neutrinos = getattr(event, 'neutrinos', None)
        if neutrinos:
            fill(self.tree, 'n_nu', len(neutrinos))
        self.tree.tree.Fill()
        
    def write(self, setup):
        self.rootfile.Write()
        self.rootfile.Close()

class ParticleTreeProducer(Analyzer):
    def beginLoop(self, setup):
        super(ParticleTreeProducer, self).beginLoop(setup)
        self.rootfile = TFile("/".join([self.dirName, "tree.root"]), "recreate")
        self.tree = Tree("particles", "")
        bookParticle(self.tree, "ptc")
        bookCluster(self.tree, "ptc_ecal")
        bookParticle(self.tree, "ptc_match")
        var(self.tree, "dr")
        bookParticle(self.tree, "ptc_match_211")
        var(self.tree, "dr_211")
        bookParticle(self.tree, "ptc_match_130")
        var(self.tree, "dr_130")
        bookParticle(self.tree, "ptc_match_22")
        var(self.tree, "dr_22")

    def process(self, event):
        particles = getattr(event, self.cfg_ana.particles)
        for ptc in particles:
            self.tree.reset()
            fillParticle(self.tree, "ptc", ptc)
            m211 = m22 = False
            # if hasattr(ptc, 'clusters'):
            #     # sim particle
            #     ecal = ptc.clusters.get('ecal_in', None)
            #     if ecal:
            #         fillCluster(self.tree, 'ptc_ecal', ecal)
            if hasattr(ptc, "match") and ptc.match:
                fillParticle(self.tree, "ptc_match", ptc.match)
                fill(self.tree, "dr", ptc.dr)
            if hasattr(ptc, "match_211") and ptc.match_211:
                m211 = True
                fillParticle(self.tree, "ptc_match_211", ptc.match_211)
                fill(self.tree, "dr_211", ptc.dr_211)
            if hasattr(ptc, "match_130") and ptc.match_130:
                fillParticle(self.tree, "ptc_match_130", ptc.match_130)
                fill(self.tree, "dr_130", ptc.dr_130)
            if hasattr(ptc, "match_22") and ptc.match_22:
                m22 = True
                fillParticle(self.tree, "ptc_match_22", ptc.match_22)
                fill(self.tree, "dr_22", ptc.dr_22)
            # if m22 and not m211:
            #     print event
            #     import pdb; pdb.set_trace()
            self.tree.tree.Fill()

    def write(self, setup):
        self.rootfile.Write()
        self.rootfile.Close()

class ZHTreeProducer(Analyzer):

    def beginLoop(self, setup):
        super(ZHTreeProducer, self).beginLoop(setup)
        self.rootfile = TFile('/'.join([self.dirName,
                                        'tree.root']),
                              'recreate')
        self.tree = Tree( 'events', '')
        bookParticle(self.tree, 'zed')
        bookParticle(self.tree, 'recoil')
        bookJet(self.tree, 'jet1')
        bookJet(self.tree, 'jet2')
        bookJet(self.tree, 'jet3')
        bookJet(self.tree, 'jet4')
        bookLepton(self.tree, 'zed_1')
        bookLepton(self.tree, 'zed_2')
        bookParticle(self.tree, 'higgs')
        bookParticle(self.tree, 'higgs_1')
        bookParticle(self.tree, 'higgs_2')
        bookParticle(self.tree, 'misenergy')
       
    def process(self, event):
        self.tree.reset()
        recoil = getattr(event, self.cfg_ana.recoil)
        fillParticle(self.tree, 'recoil', recoil)        
        misenergy = getattr(event, self.cfg_ana.misenergy)
        fillParticle(self.tree, 'misenergy', misenergy )        
        zeds = getattr(event, self.cfg_ana.zeds)
        if len(zeds)>0:
            zed = zeds[0]
            fillParticle(self.tree, 'zed', zed)
            fillLepton(self.tree, 'zed_1', zed.legs[0])
            fillLepton(self.tree, 'zed_2', zed.legs[1])
        jets = getattr(event, self.cfg_ana.jets)
        for ijet, jet in enumerate(jets):
            if ijet==4:
                break
            fillJet(self.tree, 'jet{ijet}'.format(ijet=ijet+1), jet)
        higgses = getattr(event, self.cfg_ana.higgses)
        if len(higgses)>0:
            higgs = higgses[0]
            fillParticle(self.tree, 'higgs', higgs)
            fillLepton(self.tree, 'higgs_1', higgs.legs[0])
            fillLepton(self.tree, 'higgs_2', higgs.legs[1])
        self.tree.tree.Fill()
        
    def write(self, setup):
        self.rootfile.Write()
        self.rootfile.Close()

def create_tree(filename=FNAME):
    if os.path.isfile(filename):
        return filename
    outfile = TFile(filename, 'recreate')
    tree = Tree('test_tree', 'A test tree')
    tree.var('var1')
    for i in range(200):
        tree.fill('var1', i)
        tree.tree.Fill()
    # print 'creating a tree', tree.tree.GetName(),\
    #    tree.tree.GetEntries(), 'entries in',\
    #    outfile.GetName()
    outfile.Write()
    outfile.Close()
    return outfile.GetName()

class TreeProducer(Analyzer):

    def beginLoop(self, setup):
        super(TreeProducer, self).beginLoop(setup)
        self.rootfile = TFile('/'.join([self.dirName,
                                        'tree.root']),
                              'recreate')
        self.tree = Tree( 'events', '')
        self.taggers = 'b'
        bookJet(self.tree, 'jet1', self.taggers)
        bookJet(self.tree, 'jet2', self.taggers)
        bookJet(self.tree, 'jet3', self.taggers)
        bookJet(self.tree, 'jet4', self.taggers)
        bookParticle(self.tree, 'misenergy')
        bookParticle(self.tree, 'higgs')
        bookParticle(self.tree, 'zed')
        bookLepton(self.tree, 'lepton1')
        bookLepton(self.tree, 'lepton2')
        
       
    def process(self, event):
        self.tree.reset()
        misenergy = getattr(event, self.cfg_ana.misenergy)
        fillParticle(self.tree, 'misenergy', misenergy )        
        jets = getattr(event, self.cfg_ana.jets)
        for ijet, jet in enumerate(jets):
            if ijet==4:
                break
            fillJet(self.tree, 'jet{ijet}'.format(ijet=ijet+1),
                    jet, self.taggers)
        higgs = getattr(event, self.cfg_ana.higgs)
        if higgs:
            fillParticle(self.tree, 'higgs', higgs)
        zed = getattr(event, self.cfg_ana.zed)
        if zed:
            fillParticle(self.tree, 'zed', zed)
        leptons = getattr(event, self.cfg_ana.leptons)
        for ilep, lepton in enumerate(reversed(leptons)):
            if ilep == 2:
                break
            fillLepton(self.tree,
                       'lepton{ilep}'.format(ilep=ilep+1), 
                       lepton)
        self.tree.tree.Fill()
        
    def write(self, setup):
        self.rootfile.Write()
        self.rootfile.Close()

 def beginLoop(self, setup):
     super(JetConeTreeProducer, self).beginLoop(setup)
     self.rootfile = TFile('/'.join([self.dirName,
                                     'jet_tree.root']),
                           'recreate')
     self.tree = Tree( self.cfg_ana.tree_name,
                       self.cfg_ana.tree_title )
     bookJet(self.tree, 'papasjet')
     bookJet(self.tree, 'cmsjet')
     # bookJet(self.tree, 'papas_control_jet')
     # bookJet(self.tree, 'cms_control_jet')
     bookJet(self.tree, 'gen_jet')
     # bookJet(self.tree, 'simtrack')
     var(self.tree, 'simtrack_len')
     var(self.tree, 'event_id')
     bookParticles(self.tree, 'tagged', 150)
     bookParticles(self.tree, 'nottagged', 150)
     var(self.tree, '130tagged')
     var(self.tree, '22tagged')
     var(self.tree, '130nottagged')
     var(self.tree, '22nottagged')
     var(self.tree, '11tagged')
     var(self.tree, 'othertagged')
     var(self.tree, 'threshold_pass')
     bookParticles(self.tree, 'cms_ptcs', 150)