Python Tree.var方法代码示例

# 需要导入模块: from heppy.statistics.tree import Tree [as 别名]
# 或者: from heppy.statistics.tree.Tree import var [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 var [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 var [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 var [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 var [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 var [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 var [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 var [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 var [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 var [as 别名]
class Bs2TauTauAnalyzer(Analyzer):
    def beginLoop(self, setup):
        self.start_time = time.time()
        self.last_timestamp = time.time()

        self.counter = 0 # Total number of processed decays
        self.pb_counter = 0 # Number of events with B momentum > 25 GeV

        gROOT.ProcessLine('.x ' + self.cfg_ana.stylepath) # nice looking plots

        # histograms to visualize cuts
        self.pb_hist = TH1F('pb_hist', 'P_{B}', 500, 0, 50)

        super(Bs2TauTauAnalyzer, self).beginLoop(setup)
        self.rootfile = TFile('/'.join([self.dirName, 'output.root']), 'recreate')

        # tree to store MC truth values and its branches
        self.mc_truth_tree = Tree(self.cfg_ana.mc_truth_tree_name, self.cfg_ana.mc_truth_tree_title)
        self.mc_truth_tree.var('n_particles')
        self.mc_truth_tree.var('event_number')
        self.mc_truth_tree.var('pv_x')
        self.mc_truth_tree.var('pv_y')
        self.mc_truth_tree.var('pv_z')
        self.mc_truth_tree.var('tv_tauplus_x')
        self.mc_truth_tree.var('tv_tauplus_y')
        self.mc_truth_tree.var('tv_tauplus_z')
        self.mc_truth_tree.var('tv_tauminus_x')
        self.mc_truth_tree.var('tv_tauminus_y')
        self.mc_truth_tree.var('tv_tauminus_z')
        self.mc_truth_tree.var('b_px')
        self.mc_truth_tree.var('b_py')
        self.mc_truth_tree.var('b_pz')
        self.mc_truth_tree.var('opposite_b_quark_mc_truth_px')
        self.mc_truth_tree.var('opposite_b_quark_mc_truth_py')
        self.mc_truth_tree.var('opposite_b_quark_mc_truth_pz')
        self.mc_truth_tree.var('tauplus_px')
        self.mc_truth_tree.var('tauplus_py')
        self.mc_truth_tree.var('tauplus_pz')
        self.mc_truth_tree.var('pi1_tauplus_px')
        self.mc_truth_tree.var('pi1_tauplus_py')
        self.mc_truth_tree.var('pi1_tauplus_pz')
        self.mc_truth_tree.var('pi1_tauplus_q')
        self.mc_truth_tree.var('pi2_tauplus_px')
        self.mc_truth_tree.var('pi2_tauplus_py')
        self.mc_truth_tree.var('pi2_tauplus_pz')
        self.mc_truth_tree.var('pi2_tauplus_q')
        self.mc_truth_tree.var('pi3_tauplus_px')
        self.mc_truth_tree.var('pi3_tauplus_py')
        self.mc_truth_tree.var('pi3_tauplus_pz')
        self.mc_truth_tree.var('pi3_tauplus_q')
        self.mc_truth_tree.var('nu_tauplus_px')
        self.mc_truth_tree.var('nu_tauplus_py')
        self.mc_truth_tree.var('nu_tauplus_pz')
        self.mc_truth_tree.var('tauminus_px')
        self.mc_truth_tree.var('tauminus_py')
        self.mc_truth_tree.var('tauminus_pz')
        self.mc_truth_tree.var('pi1_tauminus_px')
        self.mc_truth_tree.var('pi1_tauminus_py')
        self.mc_truth_tree.var('pi1_tauminus_pz')
        self.mc_truth_tree.var('pi1_tauminus_q')
        self.mc_truth_tree.var('pi2_tauminus_px')
        self.mc_truth_tree.var('pi2_tauminus_py')
        self.mc_truth_tree.var('pi2_tauminus_pz')
        self.mc_truth_tree.var('pi2_tauminus_q')
        self.mc_truth_tree.var('pi3_tauminus_px')
        self.mc_truth_tree.var('pi3_tauminus_py')
        self.mc_truth_tree.var('pi3_tauminus_pz')
        self.mc_truth_tree.var('pi3_tauminus_q')
        self.mc_truth_tree.var('nu_tauminus_px')
        self.mc_truth_tree.var('nu_tauminus_py')
        self.mc_truth_tree.var('nu_tauminus_pz')

        # same for smeared values
        self.tree = Tree(self.cfg_ana.tree_name, self.cfg_ana.tree_title)
        self.tree.var('n_particles')
        self.tree.var('event_number')
        self.tree.var('pv_x')
        self.tree.var('pv_y')
        self.tree.var('pv_z')
        self.tree.var('tv_tauplus_x')
        self.tree.var('tv_tauplus_y')
        self.tree.var('tv_tauplus_z')
        self.tree.var('tv_tauminus_x')
        self.tree.var('tv_tauminus_y')
        self.tree.var('tv_tauminus_z')
        self.tree.var('pi1_tauplus_px')
        self.tree.var('pi1_tauplus_py')
        self.tree.var('pi1_tauplus_pz')
        self.tree.var('pi1_tauplus_q')
        self.tree.var('pi2_tauplus_px')
        self.tree.var('pi2_tauplus_py')
        self.tree.var('pi2_tauplus_pz')
        self.tree.var('pi2_tauplus_q')
        self.tree.var('pi3_tauplus_px')
        self.tree.var('pi3_tauplus_py')
        self.tree.var('pi3_tauplus_pz')
        self.tree.var('pi3_tauplus_q')
        self.tree.var('pi1_tauminus_px')
        self.tree.var('pi1_tauminus_py')
        self.tree.var('pi1_tauminus_pz')
#.........这里部分代码省略.........

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

# 需要导入模块: from heppy.statistics.tree import Tree [as 别名]
# 或者: from heppy.statistics.tree.Tree import var [as 别名]
infname = sys.argv[1]

sh = shelve.open(infname)
outevents = list()

events = dict()

outfname = infname.replace(".shv", ".root")

print "input:", infname
print "output", outfname

f = TFile(outfname, "RECREATE")
tree = Tree("events", "tree for gael")

tree.var("rec1_e")
tree.var("rec1_gen_e")
tree.var("rec1_dr")
tree.var("rec2_e")
tree.var("rec2_gen_e")
tree.var("rec2_dr")
tree.var("ak1_e")
tree.var("ak1_gen_e")
tree.var("ak1_dr")
tree.var("ak2_e")
tree.var("ak2_gen_e")
tree.var("ak2_dr")
tree.var("drgen")


def process_event(ievent):