Python Tree.fill方法代码示例

# 需要导入模块: from heppy.statistics.tree import Tree [as 别名]
# 或者: from heppy.statistics.tree.Tree import fill [as 别名]
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()

# 需要导入模块: from heppy.statistics.tree import Tree [as 别名]
# 或者: from heppy.statistics.tree.Tree import fill [as 别名]
 def test_fill(self):
     fi = TFile('tree.root','RECREATE')
     tr = Tree('test_tree', 'A test tree')
     tr.var('a')
     tr.var('b')
     tr.fill('a', 3)
     tr.fill('a', 4)
     tr.fill('b', 5)
     tr.tree.Fill()
     fi.Write()
     fi.Close()

# 需要导入模块: from heppy.statistics.tree import Tree [as 别名]
# 或者: from heppy.statistics.tree.Tree import fill [as 别名]
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()

# 需要导入模块: from heppy.statistics.tree import Tree [as 别名]
# 或者: from heppy.statistics.tree.Tree import fill [as 别名]
class HTo4lGenTreeProducer(Analyzer):

    def beginLoop(self, setup):
        super(HTo4lGenTreeProducer, self).beginLoop(setup)
        self.rootfile = TFile('/'.join([self.dirName,
                                        'tree.root']),
                              'recreate')
        self.tree = Tree( 'events', '')
        self.tree.var('weight', float)

        bookParticle(self.tree, 'lep1vsPt')
        bookParticle(self.tree, 'lep2vsPt')
        bookParticle(self.tree, 'lep3vsPt')
        bookParticle(self.tree, 'lep4vsPt')

        bookParticle(self.tree, 'lep1vsEta')
        bookParticle(self.tree, 'lep2vsEta')
        bookParticle(self.tree, 'lep3vsEta')
        bookParticle(self.tree, 'lep4vsEta')


    def process(self, event):
        self.tree.reset()
        gen_leptons = getattr(event, self.cfg_ana.leptons)

        self.tree.fill('weight' , event.weight )
        
        if len(gen_leptons) >= 4:

            gen_leptons.sort(key=lambda x: x.pt(), reverse=True)

            fillParticle(self.tree, 'lep1vsPt', gen_leptons[0])
            fillParticle(self.tree, 'lep2vsPt', gen_leptons[1])
            fillParticle(self.tree, 'lep3vsPt', gen_leptons[2])
            fillParticle(self.tree, 'lep4vsPt', gen_leptons[3])

            gen_leptons.sort(key=lambda x: abs(x.eta()))

            fillParticle(self.tree, 'lep1vsEta', gen_leptons[0])
            fillParticle(self.tree, 'lep2vsEta', gen_leptons[1])
            fillParticle(self.tree, 'lep3vsEta', gen_leptons[2])
            fillParticle(self.tree, 'lep4vsEta', gen_leptons[3])

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

# 需要导入模块: from heppy.statistics.tree import Tree [as 别名]
# 或者: from heppy.statistics.tree.Tree import fill [as 别名]
def create_tree(filename=FNAME, nentries=None):
    if not nentries: 
        if os.path.isfile(filename):
            #default number of entries, file exists
            return filename
        else: 
            nentries = 200
    nentries = int(nentries)
    outfile = TFile(filename, 'recreate')
    tree = Tree('test_tree', 'A test tree')
    tree.var('var1')
    for i in range(nentries):
        tree.fill('var1', i)
        tree.tree.Fill()
    outfile.Write()
    outfile.Close()
    return outfile.GetName()

# 需要导入模块: from heppy.statistics.tree import Tree [as 别名]
# 或者: from heppy.statistics.tree.Tree import fill [as 别名]
class SimpleTreeProducer(Analyzer):

    def beginLoop(self, setup):
        super(SimpleTreeProducer, self).beginLoop(setup)
        self.rootfile = TFile('/'.join([self.dirName,
                                        'simple_tree.root']),
                              'recreate')
        self.tree = Tree( self.cfg_ana.tree_name,
                          self.cfg_ana.tree_title )
        self.tree.var('test_variable')

    def process(self, event):
        self.tree.fill('test_variable', event.input.var1)
        self.tree.tree.Fill()

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

# 需要导入模块: from heppy.statistics.tree import Tree [as 别名]
# 或者: from heppy.statistics.tree.Tree import fill [as 别名]
class SimpleTreeProducer(Analyzer):
    '''Test analyzer creating a simple root tree.
    
    Example::
    
        tree = cfg.Analyzer(
          SimpleTreeProducer,
          tree_name = 'events',
          tree_title = 'A simple test tree'
        )
    
    The TTree is written to the file C{simple_tree.root} in the analyzer directory.
    
    @param tree_name: Name of the tree (Key in the output root file).
    @param tree_title: Title of the tree.
    '''
    def beginLoop(self, setup):
        super(SimpleTreeProducer, self).beginLoop(setup)
        self.rootfile = TFile('/'.join([self.dirName,
                                        'simple_tree.root']),
                              'recreate')
        self.tree = Tree( self.cfg_ana.tree_name,
                          self.cfg_ana.tree_title )
        self.tree.var('test_variable')
        self.tree.var('test_variable_random')

    def process(self, event):
        '''Process the event.
        
        The input data must contain a variable called "var1",
        which is the case of the L{test tree<heppy.utils.debug_tree>}. 
        
        The event must contain:
         - var_random, which is the case if the L{RandomAnalyzer<heppy.analyzers.examples.simple.RandomAnalyzer.RandomAnalyzer>}
         has processed the event. 
         
        '''
        self.tree.fill('test_variable', event.input.var1)
        self.tree.fill('test_variable_random', event.var_random)
        self.tree.tree.Fill()

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

# 需要导入模块: from heppy.statistics.tree import Tree [as 别名]
# 或者: from heppy.statistics.tree.Tree import fill [as 别名]
def create_tree(filename=FNAME, nentries=None):
    '''Create the test tree in file FNAME.'''
    if not nentries:
        file_good = False
        if os.path.isfile(filename):
            rfile = TFile(filename)
            if not rfile.IsZombie():
                file_good = True
        if file_good:
            return filename
        else:
            # file needs to be regenerated so setting default
            # number of entries
            nentries = 200
    nentries = int(nentries)
    outfile = TFile(filename, 'recreate')
    tree = Tree('test_tree', 'A test tree')
    tree.var('var1')
    for i in range(nentries):
        tree.fill('var1', i)
        tree.tree.Fill()
    outfile.Write()
    outfile.Close()
    return outfile.GetName()

# 需要导入模块: from heppy.statistics.tree import Tree [as 别名]
# 或者: from heppy.statistics.tree.Tree import fill [as 别名]
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.tree.var('tau1', float)
        self.tree.var('tau2', float)
        self.tree.var('tau3', float)
        self.tree.var('tau32', float)
        self.tree.var('tau31', float)
        self.tree.var('tau21', float)

        bookParticle(self.tree, 'Jet')
        bookParticle(self.tree, 'softDroppedJet')
        bookParticle(self.tree, 'leadingSoftDroppedSubJet')
        bookParticle(self.tree, 'trailingSoftDroppedSubJet')

    def process(self, event):
        self.tree.reset()
        jets = getattr(event, self.cfg_ana.fatjets)
        
        # store leading (fat) jet observables
        if len(jets) > 0 and len(jets[0].subjetsSoftDrop) > 2:
            self.tree.fill('tau1' , jets[0].tau1 )
            self.tree.fill('tau2' , jets[0].tau2 )
            self.tree.fill('tau3' , jets[0].tau3 )
            self.tree.fill('tau31' , jets[0].tau3/jets[0].tau1 )
            self.tree.fill('tau32' , jets[0].tau3/jets[0].tau2 )
            self.tree.fill('tau21' , jets[0].tau2/jets[0].tau1 )

            fillParticle(self.tree, 'Jet', jets[0])
            
            # first subjet entry is the cleaned jet itself
            fillParticle(self.tree, 'softDroppedJet', jets[0].subjetsSoftDrop[0])
            fillParticle(self.tree, 'leadingSoftDroppedSubJet', jets[0].subjetsSoftDrop[1])
            fillParticle(self.tree, 'trailingSoftDroppedSubJet', jets[0].subjetsSoftDrop[2])
            
            self.tree.tree.Fill()

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

# 需要导入模块: from heppy.statistics.tree import Tree [as 别名]
# 或者: from heppy.statistics.tree.Tree import fill [as 别名]

#.........这里部分代码省略.........
                        if abs(ptc_gen3.pdgid) == 211 and ptc_gen3.start_vertex == tauminus_mc_truth.end_vertex:
                            pis_tauminus_mc_truth.append(ptc_gen3)

                        if ptc_gen3.pdgid == -16:
                            nu_tauminus_mc_truth = ptc_gen3

                    if len(pis_tauminus_mc_truth) == 3:
                        pi1_tauminus_mc_truth, pi2_tauminus_mc_truth, pi3_tauminus_mc_truth = pis_tauminus_mc_truth[0], pis_tauminus_mc_truth[1], pis_tauminus_mc_truth[2]
                        pi1_tauminus, pi2_tauminus, pi3_tauminus = copy.deepcopy(pi1_tauminus_mc_truth), copy.deepcopy(pi2_tauminus_mc_truth), copy.deepcopy(pi3_tauminus_mc_truth)

                    # applying smearing
                    if self.cfg_ana.smear_momentum:
                        pi1_tauplus.p = smear_momentum(pi1_tauplus.p, self.cfg_ana.momentum_x_resolution, self.cfg_ana.momentum_y_resolution, self.cfg_ana.momentum_z_resolution)
                        pi2_tauplus.p = smear_momentum(pi2_tauplus.p, self.cfg_ana.momentum_x_resolution, self.cfg_ana.momentum_y_resolution, self.cfg_ana.momentum_z_resolution)
                        pi3_tauplus.p = smear_momentum(pi3_tauplus.p, self.cfg_ana.momentum_x_resolution, self.cfg_ana.momentum_y_resolution, self.cfg_ana.momentum_z_resolution)
                        pi1_tauminus.p = smear_momentum(pi1_tauminus.p, self.cfg_ana.momentum_x_resolution, self.cfg_ana.momentum_y_resolution, self.cfg_ana.momentum_z_resolution)
                        pi2_tauminus.p = smear_momentum(pi2_tauminus.p, self.cfg_ana.momentum_x_resolution, self.cfg_ana.momentum_y_resolution, self.cfg_ana.momentum_z_resolution)
                        pi3_tauminus.p = smear_momentum(pi3_tauminus.p, self.cfg_ana.momentum_x_resolution, self.cfg_ana.momentum_y_resolution, self.cfg_ana.momentum_z_resolution)
                    if self.cfg_ana.smear_pv:
                        pv = smear_vertex(pv, self.cfg_ana.pv_x_resolution, self.cfg_ana.pv_y_resolution, self.cfg_ana.pv_z_resolution)
                    if self.cfg_ana.smear_tv:
                        tv_tauplus = smear_vertex(tv_tauplus, self.cfg_ana.tv_x_resolution, self.cfg_ana.tv_y_resolution, self.cfg_ana.tv_z_resolution)
                        tv_tauminus = smear_vertex(tv_tauminus, self.cfg_ana.tv_x_resolution, self.cfg_ana.tv_y_resolution, self.cfg_ana.tv_z_resolution)

                        # to keep consistency
                        pi1_tauplus.start_vertex, pi2_tauplus.start_vertex, pi3_tauplus.start_vertex = tv_tauplus, tv_tauplus, tv_tauplus
                        pi1_tauminus.start_vertex, pi2_tauminus.start_vertex, pi3_tauminus.start_vertex = tv_tauminus, tv_tauminus, tv_tauminus

                    if pi1_tauplus.is_valid() and pi2_tauplus.is_valid() and pi3_tauplus.is_valid() and pi1_tauminus.is_valid() and pi2_tauminus.is_valid() and pi3_tauminus.is_valid():
                        # filling histogram
                        self.pb_hist.Fill(pb)

                        # filling MC truth information
                        self.mc_truth_tree.fill('event_number', event_number)
                        self.mc_truth_tree.fill('n_particles', n_particles)

                        self.mc_truth_tree.fill('pv_x', pv_mc_truth.x)
                        self.mc_truth_tree.fill('pv_y', pv_mc_truth.y)
                        self.mc_truth_tree.fill('pv_z', pv_mc_truth.z)
                        self.mc_truth_tree.fill('tv_tauplus_x', tv_tauplus_mc_truth.x)
                        self.mc_truth_tree.fill('tv_tauplus_y', tv_tauplus_mc_truth.y)
                        self.mc_truth_tree.fill('tv_tauplus_z', tv_tauplus_mc_truth.z)
                        self.mc_truth_tree.fill('tv_tauminus_x', tv_tauminus_mc_truth.x)
                        self.mc_truth_tree.fill('tv_tauminus_y', tv_tauminus_mc_truth.y)
                        self.mc_truth_tree.fill('tv_tauminus_z', tv_tauminus_mc_truth.z)

                        self.mc_truth_tree.fill('b_px', b_mc_truth.p.px)
                        self.mc_truth_tree.fill('b_py', b_mc_truth.p.py)
                        self.mc_truth_tree.fill('b_pz', b_mc_truth.p.pz)

                        self.mc_truth_tree.fill('opposite_b_quark_mc_truth_px', opposite_b_quark_mc_truth.p.px)
                        self.mc_truth_tree.fill('opposite_b_quark_mc_truth_py', opposite_b_quark_mc_truth.p.py)
                        self.mc_truth_tree.fill('opposite_b_quark_mc_truth_pz', opposite_b_quark_mc_truth.p.pz)

                        self.mc_truth_tree.fill('tauplus_px', tauplus_mc_truth.p.px)
                        self.mc_truth_tree.fill('tauplus_py', tauplus_mc_truth.p.py)
                        self.mc_truth_tree.fill('tauplus_pz', tauplus_mc_truth.p.pz)

                        self.mc_truth_tree.fill('pi1_tauplus_q', pi1_tauplus_mc_truth.charge)
                        self.mc_truth_tree.fill('pi1_tauplus_px', pi1_tauplus_mc_truth.p.px)
                        self.mc_truth_tree.fill('pi1_tauplus_py', pi1_tauplus_mc_truth.p.py)
                        self.mc_truth_tree.fill('pi1_tauplus_pz', pi1_tauplus_mc_truth.p.pz)

                        self.mc_truth_tree.fill('pi2_tauplus_q', pi2_tauplus_mc_truth.charge)
                        self.mc_truth_tree.fill('pi2_tauplus_px', pi2_tauplus_mc_truth.p.px)
                        self.mc_truth_tree.fill('pi2_tauplus_py', pi2_tauplus_mc_truth.p.py)

# 需要导入模块: from heppy.statistics.tree import Tree [as 别名]
# 或者: from heppy.statistics.tree.Tree import fill [as 别名]
from ROOT import TFile
from heppy.statistics.tree import Tree

outfile = TFile('test_tree.root', '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()