Как заполнить ветки дерева? - PullRequest
Новая
       29

Как заполнить ветки дерева?

0 голосов
/ 06 марта 2019 
import sys
import ROOT
from progressbar import Bar, Percentage, ProgressBar
from time import time
from tools import duration, check_outfile_path

ECMS = 3.686
p4shw = ROOT.vector('double')()

def main ():     
    args = sys.argv[1:]

    if (len(args) < 2):
        print 'input error'

    infile = args[0]
    outfile = args[1]
    check_outfile_path(outfile)

    fin = ROOT.TFile(infile)
    t = fin.Get('ana')
    t.SetBranchAddress("p4shw", p4shw)
    entries = t.GetEntriesFast()

    fout = ROOT.TFile(outfile, "RECREATE")
    t_out = ROOT.TTree("ana","ana")
    rec_mass_gam1 = ROOT.vector('double')()
    rec_mass_gam2 = ROOT.vector('double')()
    t_out.Branch("rec_mass_gam1", rec_mass_gam1, "rec_mass_gam1/D")
    t_out.Branch("rec_mass_gam2", rec_mass_gam2, "rec_mass_gam2/D")

    pbar = ProgressBar(widgets=[Percentage(), Bar()], maxval=entries).start()
    time_start = time()
    print("checking error 2")
    cms_p4 = ROOT.TLorentzVector(0.011*ECMS, 0, 0, ECMS)
    print 'entries=', entries
    print("checking error 3")
    for k in range(entries):

        pbar.update(k+1)

        #t.GetEntry(k)
        print("indentent error checking")
        #exit()
        p4shw_gam1 = ROOT.TLorentzVector(t.p4shw[0],t.p4shw[1],t.p4shw[2],t.p4shw[3])
        p4shw_gam2 = ROOT.TLorentzVector(t.p4shw[4],t.p4shw[5],t.p4shw[6],t.p4shw[7])
        print("checking error 4")
        p4_shw_gam1 = cms_p4 - p4shw_gam1
        p4_shw_gam2 = cms_p4 - p4shw_gam2
        rec_mass_gam1 = p4_shw_gam1.M()
        rec_mass_gam2 = p4_shw_gam2.M()
        print("rec_mass_gam1", rec_mass_gam1)
        #exit()
        t_out.Fill()
        print("checking error 5")
    t_out.Write()
    fout.Close()
    pbar.finish()
    dur = duration(time()-time_start)
    sys.stdout.write(' \nDone in %s. \n' % dur)
    print("checking error 6")

if __name__ =='__main__':
    main()

1 Ответ

0 голосов
/ 21 марта 2019

Когда я сравниваю ваш код с в этом примере , вы используете ROOT.vector вместо array. Когда я делаю это изменение, ветвь заполняется как ожидалось 

#!/bin/python

import ROOT
from array import array


# doesn't work
def test1():
    t_out = ROOT.TTree("ana", "ana")
    rec_mass_gam1 = ROOT.vector('double')()
    t_out.Branch("rec_mass_gam1", rec_mass_gam1, "rec_mass_gam1/D")
    rec_mass_gam1 = 1337.
    t_out.Fill()
    t_out.Draw("rec_mass_gam1")


# works
def test2():
    t_out = ROOT.TTree("ana", "ana")
    rec_mass_gam1 = array('f', [0.])
    t_out.Branch("rec_mass_gam1", rec_mass_gam1, "rec_mass_gam1/F")
    rec_mass_gam1[0] = 1337.
    t_out.Fill()
    t_out.Draw("rec_mass_gam1")

Когда я запускаю test1, я вижу, что дерево и ветвь заполнены, но не тем значением, которое я хотел. Во втором примере желаемое значение заполняется.

Теперь, присмотревшись к тому, что происходит, в любом случае в вашем скрипте есть ошибка:

python не обрабатывает rec_mass_gam1 = p4_shw_gam1.M() как "задает значение векторной переменной rec_mass_gam1 равным числу, которое выпадает из метода M(). Он скорее создает новую переменную с плавающей точкой с именем rec_mass_gam1 и исходная векторная переменная (которая используется ветвью) остается неизменной.

Я должен признать, что не знаю, есть ли способ заполнить ветку также vector.

...