def staff():
staff = ROOT.staff_t()
# The input file cern.dat is a copy of the CERN staff data base
# from 1988
f = TFile('staff.root', 'RECREATE')
tree = TTree('T', 'staff data from ascii file')
tree.Branch('staff', staff,
'Category/I:Flag:Age:Service:Children:Grade:Step:Hrweek:Cost')
tree.Branch('Divisions', AddressOf(staff, 'Division'), 'Division/C')
tree.Branch('Nation', AddressOf(staff, 'Nation'), 'Nation/C')
# note that the branches Division and Nation cannot be on the first branch
fname = os.path.expandvars('$ROOTSYS/tutorials/tree/cernstaff.dat')
for line in open(fname).readlines():
t = filter(lambda x: x, re.split('\s+', line))
staff.Category = int(t[0]) # assign as integers
staff.Flag = int(t[1])
staff.Age = int(t[2])
staff.Service = int(t[3])
staff.Children = int(t[4])
staff.Grade = int(t[5])
staff.Step = int(t[6])
staff.Hrweek = int(t[7])
staff.Cost = int(t[8])
staff.Division = t[9] # assign as strings
staff.Nation = t[10]
tree.Fill()
tree.Print()
tree.Write()
def makeTTree():
"""Create ROOT TTree filled with a Gaussian distribution in x
and a uniform distribution in y"""
tree = TTree("tree", "tree")
px = array('d', [0])
py = array('d', [0])
tree.Branch("px", px, "px/D")
tree.Branch("py", py, "py/D")
for i in range(500):
px[0] = gRandom.Gaus(0, 3)
py[0] = gRandom.Uniform() * 30 - 15
tree.Fill()
return tree
def open(self, update=True):
print 'NtupleHandler.open() called'
if os.path.exists(self.fileName) and update:
print 'Opening %s for updating' % self.fileName
self.updatemode = True
self.file = TFile(self.fileName, 'update')
self.tree = self.file.Get(self.treeName)
else:
print 'Creating %s for writing' % self.fileName
self.updatemode = False
self.file = TFile(self.fileName, 'create')
self.tree = TTree(self.treeName, self.treeName)
def test13WriteMisnamedLeaf( self ):
"""Test writing of an differently named leaf"""
f = TFile( self.fname, 'RECREATE' )
t = TTree( self.tname, self.ttitle )
s = SomeDataStruct()
t.Branch( 'cpu_packet_time', AddressOf(s, 'NLabel'), 'time/I' );
for i in range(self.N):
s.NLabel = i
t.Fill()
f.Write()
f.Close()
def test_minimizer(fin_name='gen_toy_10bin_highstat', fout_name='result'):
fout = TFile('deltas/' + fout_name + '.root', 'RECREATE')
tout = TTree('tree', 'tree')
bin = array('i', [0])
n = array('f', [0])
delta_n_u = array('f', [0.])
delta_n_d = array('f', [0.])
delta_m_u = array('f', [0.])
delta_m_d = array('f', [0.])
var_delta_m = array('f', [0.])
tout.Branch('bin', bin, 'bin/I')
tout.Branch('n', n, 'n/F')
tout.Branch('delta_n_u', delta_n_u, 'delta_n_u/F')
tout.Branch('delta_n_d', delta_n_d, 'delta_n_d/F')
tout.Branch('delta_m_u', delta_m_u, 'delta_m_u/F')
tout.Branch('delta_m_d', delta_m_d, 'delta_m_d/F')
tout.Branch('var_delta_m', var_delta_m, 'var_delta_m/F')
f = TFile.Open('deltas/' + fin_name + ".root")
tree = f.Get("tree")
tree.SetBranchStatus("*", True)
for iev in range(tree.GetEntries()):
tree.GetEntry(iev)
ev = tree
n_toy = ev.n_toy
n_bins = ev.n_bin
n_deltas = ev.n_deltas
deltas = ev.deltas
bin_n = ev.bin_n
bin_u = ev.bin_u
bin_d = ev.bin_d
print deltas
fix_params = []
#fix_params.extend([1, 2+3*n_bins])
#fix_params.extend(range(n_bins+2, 2*n_bins+2))
minimize_deltas = Minimze_deltas(deltas=deltas,
nbins=n_bins,
step=0.01,
fix_params=fix_params)
res = minimize_deltas.run()
for b in range(n_bins):
bin[0] = b + 1
n[0] = bin_n[b]
delta_n_u[0] = res[0][b]
delta_n_d[0] = res[1][b]
delta_m_u[0] = res[2][b]
delta_m_d[0] = res[3][b]
var_delta_m[0] = res[4][b]
#print "Bin %s: [%s, %s] ==> %s +/- %s" % ( b+1, delta_n_u[0], delta_n_d[0], abs(delta_m_u[0]), math.sqrt(var_delta_m[0]))
tout.Fill()
f.Close()
fout.cd()
fout.Write()
fout.Close()
class TreeProducer(object):
"""Base class to create and prepare a custom output file & tree for analysis modules."""
def __init__(self, filename, module, **kwargs):
self.filename = filename
self.module = module
self.outfile = TFile(filename, 'RECREATE')
ncuts = kwargs.get('ncuts', 25)
self.cutflow = Cutflow('cutflow', ncuts)
self.display = kwargs.get('display_cutflow', True)
self.pileup = TH1D('pileup', 'pileup', 100, 0, 100)
self.tree = TTree('tree', 'tree')
def addBranch(self,
name,
dtype='f',
default=None,
title=None,
arrname=None):
"""Add branch with a given name, and create an array of the same name as address."""
if hasattr(self, name):
raise IOError("Branch of name '%s' already exists!" % (name))
if not arrname:
arrname = name
if isinstance(dtype, str): # Set correct data type for numpy:
if dtype == 'F': # 'F' = 'complex64', which do not work for filling float branches
dtype = 'float32' # 'f' = 'float32' -> 'F' -> Float_t
if isinstance(dtype, str): # Set correct data type for numpy:
if dtype == 'D': # 'D' = 'complex128', which do not work for filling float branches
dtype = 'float64' # 'd' = 'float64' -> 'D' -> Double_t
setattr(self, arrname, np.zeros(1, dtype=dtype))
branch = self.tree.Branch(name, getattr(self, arrname),
'%s/%s' % (name, root_dtype[dtype]))
if default != None:
getattr(self, name)[0] = default
if title:
branch.SetTitle(title)
return branch
def fill(self):
"""Fill tree."""
return self.tree.Fill()
def endJob(self):
"""Write and close files after the job ends."""
if self.display:
self.cutflow.display()
print ">>> Write %s..." % (self.outfile.GetName())
self.outfile.Write()
self.outfile.Close()
def ROC(stree, btree, selection, score_retrieval, outname):
bhist = histfill(btree, selection, score_retrieval)
shist = histfill(stree, selection, score_retrieval)
outtree = TTree(outname, outname)
tpr = np.zeros(1)
fpr = np.zeros(1)
outtree.Branch("tpr", tpr, "tpr/D")
outtree.Branch("fpr", fpr, "fpr/D")
n_signal = shist.Integral()
n_background = bhist.Integral()
for i in range(10000):
tpr[0] = 1. - (shist.Integral(0, i + 1) / n_signal)
fpr[0] = 1. - (bhist.Integral(0, i + 1) / n_background)
outtree.Fill()
outtree.Write()
def CreateOutput(self, nam):
#create tree output associated with the hits
self.otree = TTree(nam, nam)
self.out_en = c_double(0)
self.out_x = c_double(0)
self.out_y = c_double(0)
self.out_z = c_double(0)
self.out_pdg = c_int(0)
self.otree.Branch("en", self.out_en, "en/D")
self.otree.Branch("x", self.out_x, "x/D")
self.otree.Branch("y", self.out_y, "y/D")
self.otree.Branch("z", self.out_z, "z/D")
self.otree.Branch("pdg", self.out_pdg, "pdg/I")
def test01WriteStdVector(self):
"""Test writing of a single branched TTree with an std::vector<double>"""
f = TFile(self.fname, 'RECREATE')
t = TTree(self.tname, self.ttitle)
v = std.vector('double')()
t.Branch('mydata', v.__class__.__name__, v)
for i in range(self.N):
for j in range(self.M):
v.push_back(i * self.M + j)
t.Fill()
v.clear()
f.Write()
f.Close()
def create_singles_TTree(host_file):
from ROOT import TTree
host_file.cd()
title = 'Single events by Sam Kohn (git: {})'.format(
translate.git_describe())
out = TTree('singles', title)
buf = TreeBuffer()
initialize_basic_TTree(out, buf)
buf.fQuad = float_value()
buf.fMax = float_value()
buf.fPSD_t1 = float_value()
buf.fPSD_t2 = float_value()
buf.f2inch_maxQ = float_value()
buf.x = float_value()
buf.y = float_value()
buf.z = float_value()
buf.x_AdTime = float_value()
buf.y_AdTime = float_value()
buf.z_AdTime = float_value()
buf.fID = float_value()
buf.fPSD = float_value()
buf.num_nearby_events = unsigned_int_value()
buf.nearby_dt = int_value(10)
buf.nearby_energy = float_value(10)
def branch(name, typecode):
out.Branch(name, getattr(buf, name), '{}/{}'.format(name,
typecode))
return
branch('fQuad', 'F')
branch('fMax', 'F')
branch('fPSD_t1', 'F')
branch('fPSD_t2', 'F')
branch('f2inch_maxQ', 'F')
branch('x', 'F')
branch('y', 'F')
branch('z', 'F')
branch('x_AdTime', 'F')
branch('y_AdTime', 'F')
branch('z_AdTime', 'F')
branch('fID', 'F')
branch('fPSD', 'F')
branch('num_nearby_events', 'I')
out.Branch('nearby_dt', buf.nearby_dt, 'nearby_dt[num_nearby_events]/I')
out.Branch('nearby_energy', buf.nearby_energy,
'nearby_energy[num_nearby_events]/F')
return out, buf
def splitFile(input_file_name, input_tree_name, splitting_branch_name):
input_file = uproot.open(input_file_name)
input_tree = input_file[input_tree_name]
branch_collection = BranchInfoCollection(input_tree)
#determine possible values of branch according to which the file will be split
splitting_branch = input_tree[splitting_branch_name]
split_values = possibleValues(splitting_branch)
#make new file with split trees
split_file_name = '{}_split_{}.root'.format(
removePathAndExtension(input_file_name), splitting_branch_name)
split_file = TFile(split_file_name, 'RECREATE')
split_trees = {}
for value in split_values:
name = '{}_{}_{}'.format(input_tree_name, splitting_branch_name, value)
split_trees[value] = TTree(name, name)
branch_collection.addBranches(split_trees[value])
for array_map in TreeIterator(input_tree):
split_array = array_map[splitting_branch_name]
for j in range(len(array_map)):
branch_collection.fillArrays(array_map, j)
split_trees[split_array[j]].Fill()
for tree in split_trees.values():
tree.Write()
split_file.Close()
def main():
# setting for reduced DS to be used on the output of makeSubset.C
inputfile_name = "small.root"
tree_name = "upsTree"
# settings for full dataset (long processing time for unbinned lh)
#inputfile_name = "/data1/chibdata/collision/v2/2012_AllData_v2.root"
#tree_name = "rootuple/upsTree"
print "Opening file"
inputfile = TFile.Open(inputfile_name, "READ")
print "Importing tree"
tree = TTree()
inputfile.GetObject(tree_name, tree)
mass = RooRealVar("ups_mass", "ups_mass", 7, 11)
y = RooRealVar("ups_rapidity", "ups_rapidity", -3, 3)
pt = RooRealVar("ups_pt", "ups_pt", 0, 100)
print "Assigning dataset"
dataArgSet = RooArgSet(mass, y, pt)
dataSet = RooDataSet("yds", "Y data set", tree, dataArgSet)
cuts= "abs(ups_rapidity) < 1.25"+\
"&& ups_pt>9.5"\
reduced_ds = dataSet.reduce(RooFit.Cut(cuts))
print "Performing likelihood analysis"
dofit(reduced_ds, "Y3S")
def fillTree(self, name, event, specific = None):
#determine if we need all event info or only part
keys = event.keys()
if specific is not None:
keys = specific
#if we haven't made the tree, create it and set it up
if name not in self.myTrees.keys():
print 'creating tree %s'%name
self.myTrees[name] = TTree(name,name)
gDirectory.Add(self.myTrees[name])
self.myTrees[name] = gDirectory.Get(name)
#create the struct we bind to our branches
stct = self._makeStruct(name,keys,event)
if len(stct) > 0:
gROOT.ProcessLine('.L %s_vars.C'%(name))
else:
gROOT.ProcessLine(stct)
gROOT.ProcessLine('%s_vars %s_vars_holder;'%(name,name))
self.myStructs[name] = getattr(ROOT,"%s_vars"%(name))()
#bind branches to the members of the struct we've created
self._bindBranches(name,keys,event)
self.myTrees[name].SetCacheSize()
#we are now done setting up the tree (if needed) -> set values and fill
if isinstance(event,dict):
for key in keys:
if hasattr(self.myStructs[name],key):
setattr(self.myStructs[name],key,event[key])
else:
self.myDicts[name][key] = event[key]
self.myTrees[name].SetBranchAddress(key,
self.myDicts[name][key])
self.myTrees[name].Fill()
def test_duplicate_branch_name():
from array import array
tree = TTree('tree', 'tree')
d = array('d', [0.])
tree.Branch('double', d, 'double/D')
tree.Branch('double', d, 'double/D')
tree.Fill()
# check that a warning was emitted
with warnings.catch_warnings(record=True) as w:
warnings.simplefilter("always")
a = rnp.tree2array(tree)
assert_equal(len(w), 1)
assert_true(issubclass(w[-1].category, RuntimeWarning))
assert_true("ignoring duplicate branch named" in str(w[-1].message))
assert_equal(a.dtype, [('double', '<f8')])
def __init__(self, name, isData=False):
print 'TreeProducerSimple is called', name
self.name = name
self.outputfile = TFile(name, 'RECREATE')
self.tree = TTree('tree', 'tree')
self.addBranch('event', int)
self.addBranch('njets', float)
self.addBranch('jpt_1', float)
self.addBranch('jeta_1', float)
self.addBranch('jphi_1', float)
self.addBranch('jpt_2', float)
self.addBranch('jeta_2', float)
self.addBranch('jphi_2', float)
self.addBranch('met', float)
self.addBranch('metphi', float)
if not isData:
jetuncs = ['jer', 'jes']
metuncs = ['jer', 'jes', 'unclEn']
for unc in jetuncs:
for var in ['Up', 'Down']:
self.addBranch('jpt_1_%s' % (unc + var), float)
self.addBranch('jpt_2_%s' % (unc + var), float)
for unc in metuncs:
for var in ['Up', 'Down']:
self.addBranch('met_%s' % (unc + var), float)
def calcAcceptance(chib_state):
rootupla = rootupla_gen_1 if chib_state == 1 else rootupla_gen_2
in_file = TFile.Open(rootupla,'READ')
tree = TTree()
#in_file.GetObject('rootuple/GenParticlesTree', tree)
in_file.GetObject('rootuple/chibTree', tree)
N0 = 0
N1 = 0
cont = 0
for event in tree:
if cont == 10000:
break
# if event.Upsilon_pt > cuts.upsilon_pt_lcut and event.Upsilon_pt < cuts.upsilon_pt_hcut and abs(event.photon_eta) < cuts.photon_eta_cut and abs(event.Upsilon_rapidity) < cuts.upsilon_rapidity_cut:
if event.Upsilon_p4.Pt() > cuts.upsilon_pt_lcut and event.Upsilon_p4.Pt() < cuts.upsilon_pt_hcut and abs(event.photon_p4.Eta()) < cuts.photon_eta_cut and abs(event.Upsilon_p4.Rapidity()) < cuts.upsilon_rapidity_cut:
muP_pt = event.muP_p4.Pt()
muP_eta = event.muP_p4.Eta()
ups_dir, mu_dir = upsilonMuDirections(event.chib_p4, event.Upsilon_p4, event.muP_p4,'hx')
weight = angDist(ups_dir, mu_dir, chib_state, None)
N0 += weight
cont += 1
#if abs(muP_eta) < 2.4 and abs(event.muM_p4.Eta()) < 2.4 and muP_pt > 2.5 and event.muM_p4.Pt() > 2.5:
if (abs(muP_eta) > 1.4 and abs(muP_eta) < 1.6 and abs(event.muM_p4.Eta()) > 1.4 and abs(event.muM_p4.Eta()) < 1.6 and muP_pt > 3 and event.muM_p4.Pt() > 3) \
or (abs(muP_eta) > 1.2 and abs(muP_eta) < 1.4 and abs(event.muM_p4.Eta()) > 1.2 and abs(event.muM_p4.Eta()) < 1.4 and muP_pt > 3.5 and event.muM_p4.Pt() > 3.5) \
or (abs(muP_eta) < 1.1 and abs(event.muM_p4.Eta()) < 1.1 and muP_pt > 4.5 and event.muM_p4.Pt() > 4.5):
N1 += weight
return N0, N1
def __init__(self, treeName, fileName, scale=1.0, cuts=""):
self.fileName = fileName
self.treeName = treeName
self.scale = scale
self.cuts = cuts
self.tfile = TFile()
self.ttree = TTree()
def initialize(self):
self.msg.info('initializing [%s]...', self.name())
self.sg = PyAthena.py_svc('StoreGateSvc')
self.hsvc = PyAthena.py_svc('THistSvc')
self.trigDec = PyAthena.py_tool(
'Trig::TrigDecisionTool/TrigDecisionTool')
#
# TTree
#
self.DATA = {}
self.BRANCHES = {}
self.hsvc['/turnontree/dijets'] = TTree("dijets", "dijets")
self.tree = self.hsvc['/turnontree/dijets']
#
for i in ['event_number', 'run_number', 'lumi_block']:
self.treewrap(i, 0)
for i in [
'%s%d' % (i, j)
for i in ['px', 'py', 'pz', 'e', 'rapidity', 'pt']
for j in [0, 1]
]:
self.treewrap(i, 0.0)
for i in ['roi_%s%d' % (i, j) for i in ['eta', 'phi'] for j in [0, 1]]:
self.treewrap(i, 0.0)
return StatusCode.Success
def create_event_TTree(host_file):
from ROOT import TTree
host_file.cd()
title = 'AD events by Sam Kohn (git: {})'.format(translate.git_describe())
out = TTree('events', title)
buf = TreeBuffer()
initialize_basic_TTree(out, buf)
buf.fQuad = float_value()
buf.fMax = float_value()
buf.fPSD_t1 = float_value()
buf.fPSD_t2 = float_value()
buf.f2inch_maxQ = float_value()
buf.x = float_value()
buf.y = float_value()
buf.z = float_value()
buf.fID = float_value()
buf.fPSD = float_value()
def branch(name, typecode):
out.Branch(name, getattr(buf, name), '{}/{}'.format(name, typecode))
return
branch('fQuad', 'F')
branch('fMax', 'F')
branch('fPSD_t1', 'F')
branch('fPSD_t2', 'F')
branch('f2inch_maxQ', 'F')
branch('x', 'F')
branch('y', 'F')
branch('z', 'F')
branch('fID', 'F')
branch('fPSD', 'F')
return out, buf
def copy_directory(newdir, olddir, condition=None):
"""Reads all objects from olddir and writes them to newdir.
newdir, olddir: Directories (inheriting from TDirectory).
condition: Function that takes key name and returns whether the file should
be kept or not (optional).
"""
for key in olddir.GetListOfKeys():
if condition is not None and (not condition(key) or
key.GetName().startswith('ProcessID')):
continue
cl = gROOT.GetClass(key.GetClassName())
if not cl:
continue
if cl.InheritsFrom(TDirectory.Class()):
newsub = newdir.mkdir(key.GetName())
oldsub = olddir.GetDirectory(key.GetName())
copy_directory(newsub, oldsub)
elif cl.InheritsFrom(TTree.Class()):
oldtree = olddir.Get(key.GetName())
newdir.cd()
newtree = oldtree.CloneTree(-1, 'fast')
newtree.Write()
else:
olddir.cd()
obj = key.ReadObj()
newdir.cd()
obj.Write(key.GetName())
del obj
class TreeProducerCommon(object):
"""Class to create a custom output file & tree; as well as create and contain branches."""
def __init__(self, filename, module, **kwargs):
print 'TreeProducerCommon is called for', name
ncuts = kwargs.get('ncuts',25)
self.filename = filename
self.module = module
self.outfile = TFile(filename,'RECREATE')
self.tree = TTree('tree','tree')
self.cutflow = TH1D('cutflow','cutflow',ncuts,0,ncuts)
def addBranch(self, name, dtype='f', default=None, arrname=None):
"""Add branch with a given name, and create an array of the same name as address."""
if hasattr(self,name):
raise IOError("Branch of name '%s' already exists!"%(name))
if not arrname:
arrname = name
if isinstance(dtype,str):
if dtype.lower()=='f': # 'f' is only a 'float32', and 'F' is a 'complex64', which do not work for filling float branches
dtype = float # float is a 'float64' ('f8')
elif dtype.lower()=='i': # 'i' is only a 'int32'
dtype = int # int is a 'int64' ('i8')
setattr(self,arrname,np.zeros(1,dtype=dtype))
self.tree.Branch(name, getattr(self,arrname), '%s/%s'%(name,root_dtype[dtype]))
if default!=None:
getattr(self,name)[0] = default
def endJob(self):
"""Write and close files after the job ends."""
self.outfile.Write()
self.outfile.Close()
def newRootFile(filenumber):
rootfile = {'number': filenumber}
rootfile['name'] = path + '/' + O['detector'][0] + '_' + \
O['dataset'][0] + '/' + O['detector'][0] + '_' \
+ O['dataset'][0] + '_' + str(filenumber) + \
'.root'
print '<<< Create ROOT file:', rootfile['name']
rootfile['file'] = TFile(rootfile['name'], 'RECREATE')
for inttype in ['b', 'h', 'i', 'l']:
if array(inttype, [0]).itemsize == 4:
break
rootfile['time'] = array(inttype, [0])
rootfile['data'] = array('f', [0.0])
rootfile['datas'] = nBX * [0.0]
rootfile['bx'] = array(inttype, [0])
rootfile['bxs'] = [bx for bx in O['crossings']]
rootfile['fill'] = array(inttype, [0])
rootfile['run'] = array(inttype, [0])
rootfile['ls'] = array(inttype, [0])
rootfile['tree'] = TTree(O['treename']['hd5files'], \
O['detector'][1]+' data')
rootfile['tree'].Branch(O['timename']['hd5files'], \
rootfile['time'], 'timestamp/I')
rootfile['tree'].Branch('data', rootfile['data'], 'data/F')
rootfile['tree'].Branch('bx', rootfile['bx'], 'bx/I')
rootfile['tree'].Branch('fill', rootfile['fill'], 'fill/I')
rootfile['tree'].Branch('run', rootfile['run'], 'run/I')
rootfile['tree'].Branch('ls', rootfile['ls'], 'ls/I')
return rootfile
def __init__(self, name, isMC=False, year=2016):
print 'TreeProducerTriggers is called', name
self.name = name
self.outputfile = TFile(name, 'RECREATE')
self.tree = TTree('tree', 'tree')
self.year = year
self.addBranch('jj_mass_widejet', float)
self.addBranch('jj_deltaEta_widejet', float)
self.addBranch('HLT_AK8PFJet500', int)
self.addBranch('HLT_PFJet500', int)
self.addBranch('HLT_CaloJet500_NoJetID', int)
self.addBranch('HLT_PFHT900', int)
self.addBranch('HLT_AK8PFJet550', int)
self.addBranch('HLT_PFJet550', int)
self.addBranch('HLT_CaloJet550_NoJetID', int)
self.addBranch('HLT_PFHT1050', int)
self.addBranch(
'HLT_DoublePFJets128MaxDeta1p6_DoubleCaloBTagDeepCSV_p71', int)
self.addBranch(
'HLT_DoublePFJets116MaxDeta1p6_DoubleCaloBTagDeepCSV_p71', int)
self.addBranch('HLT_DoublePFJets40_CaloBTagDeepCSV_p71', int)
self.addBranch('HLT_DoublePFJets100_CaloBTagDeepCSV_p71', int)
self.addBranch('HLT_DoublePFJets200_CaloBTagDeepCSV_p71', int)
self.addBranch('HLT_DoublePFJets350_CaloBTagDeepCSV_p71', int)
def __init__(self, name, isMC=False, year=2016):
print 'TreeProducerDijetMass is called', name
self.name = name
self.outputfile = TFile(name, 'RECREATE')
self.tree = TTree('tree','tree')
self.year = year
self.events = TH1F('Events', 'Events', 1,0,1)
self.addBranch('njets' , int)
self.addBranch('jid_1' , int)
self.addBranch('jid_2' , int)
self.addBranch('jj_mass' , float)
self.addBranch('jj_mass_lepcorr', float)
self.addBranch('jj_mass_metcorr', float)
self.addBranch('jj_mass_widejet', float)
self.addBranch('jj_deltaEta_widejet', float)
self.addBranch('jj_deltaEta' , float)
self.addBranch('nelectrons' , int)
self.addBranch('nmuons' , int)
self.addBranch('fatjetmass_1' , float)
self.addBranch('fatjetmass_2' , float)
self.addBranch('HLT_AK8PFJet500' , int)
self.addBranch('HLT_PFJet500' , int)
self.addBranch('HLT_CaloJet500_NoJetID' , int)
self.addBranch('HLT_PFHT900' , int)
self.addBranch('HLT_AK8PFJet550' , int)
self.addBranch('HLT_PFJet550' , int)
self.addBranch('HLT_CaloJet550_NoJetID' , int)
self.addBranch('HLT_PFHT1050' , int)
self.addBranch('isMC' , int)
def skim_data(weekend):
rdata = [[1, 3, 4, 5], [9, 13], [27], [33], [34]]
runList = rdata[weekend]
thresh = {}
thresh[1] = {0: 80, 1: 60, 2: 70, 3: 80}
thresh[1] = {0: 80, 1: 60, 2: 70, 3: 80}
thresh[3] = {0: 80, 1: 60, 2: 70, 3: 80}
thresh[4] = {0: 80, 1: 60, 2: 70, 3: 80}
thresh[5] = {0: 80, 1: 60, 2: 70, 3: 80}
thresh[9] = {0: 80, 1: 70, 2: 100, 3: 90}
thresh[13] = {0: 80, 1: 70, 2: 100, 3: 90}
thresh[27] = {0: 80, 1: 75, 2: 70, 3: 90}
thresh[33] = {0: 80, 1: 230, 2: 300, 3: 430}
thresh[34] = {0: 80, 1: 230, 2: 300, 3: 430}
# print(thresh)
# for key in thresh:
# print(key, thresh[key])
fileDir, skimDir = "./data", "./skim"
for run in runList:
filelist = glob.glob("{}/run_{}/FILTERED/compassF_run_{}".format(
fileDir, run, run) + "*.root")
print(filelist)
for f in sorted(filelist):
fcut = "(Channel==0 && Energy > {}) || ".format(thresh[run][0])
fcut += "(Channel==1 && Energy > {}) || ".format(thresh[run][1])
fcut += "(Channel==2 && Energy > {}) || ".format(thresh[run][2])
fcut += "(Channel==3 && Energy > {})".format(thresh[run][3])
ch = TChain("Data_F")
ch.Add(f)
print("Found %d entries" % (ch.GetEntries()))
outName = "{}/run_{}.root".format(skimDir, run)
outFile = TFile(outName, "RECREATE")
outTree = TTree()
outTree = ch.CopyTree(fcut)
outTree.Write()
outFile.Close()
f2 = TFile(outName)
tt2 = f2.Get("Data_F")
print(tt2.GetEntries())
def execute(self, log_out, log_err):
# Use local ROOT imports so ROOT env isn't necessary just to load the module.
from ROOT import gROOT, TFile, TTree
output_file = self.outputs[0]
input_files = self.inputs
logger.info("Creating output ROOT tuple '%s'" % output_file)
logger.info("Input text files: %s" % str(input_files))
treeFile = TFile(output_file, "RECREATE")
tree = TTree("ntuple", "data from text tuple " + input_files[0])
if len(input_files) > 1:
inputfile = tempfile.NamedTemporaryFile(delete=False)
print inputfile.name
firstfile = True
for filename in input_files:
if os.path.isfile(filename):
f = open(filename, 'r')
firstline = True
for i in f:
if firstline:
if firstfile:
branchdescriptor = i
inputfile.write(i)
else:
if branchdescriptor != i:
print "branch descriptor doesn't match"
sys.exit(-1)
else:
inputfile.write(i)
firstline = False
f.close()
firstfile = False
else:
logger.warn("Ignoring non-existant input file '%s'" %
filename)
inputfile.close()
print tree.ReadFile(inputfile.name)
os.remove(inputfile.name)
else:
print tree.ReadFile(input_files[0])
tree.Write()
def createdummyroot(fname, nevts=10000):
print ">>> Creating dummy ROOT file %r..." % (fname)
file = TFile(fname, 'RECREATE')
tree = TTree('tree', 'tree')
hist = TH1F('hist', 'hist', 50, -2, 2)
pt = array('d', [0])
phi = array('d', [0])
tree.Branch("pt", pt, 'normal/D')
tree.Branch("phi", phi, 'uniform/D')
for i in xrange(nevts):
pt[0] = gRandom.Landau(40, 20)
phi[0] = gRandom.Uniform(-1.57, 1.57)
hist.Fill(gRandom.Landau(0, 1))
tree.Fill()
file.Write()
file.Close()
return fname
class tree_manager:
def __init__(self, name):
self.tree = TTree(name, '')
self.arrays = {}
def Branch(self, key):
self.arrays[key] = array.array('d', [0.0])
return self.tree.Branch(key, self.arrays[key], key + '/D')
def __setitem__(self, key, value):
self.arrays[key][0] = value
def Fill(self):
return self.tree.Fill()
def Write(self):
return self.tree.Write()
def test11WriteTObjArray(self):
"""Test writing of a TObjArray"""
f = TFile(self.fname, 'RECREATE')
t = TTree(self.tname, self.ttitle)
o = TObjArray()
t.Branch('mydata', o)
nameds = [TNamed(name, name) for name in self.testnames]
for name in nameds:
o.Add(name)
self.assertEqual(len(o), len(self.testnames))
t.Fill()
f.Write()
f.Close()
def write_rootfile(self):
"""
Put the signal and bkg histos and the config as a json string
into the currently open rootfile.
rootfile - writable TFile
"""
self.histo.Write()
self.bkg_histo.Write()
config_str = json.dumps(self.config.to_dict())
config_array = array.array('B', config_str)
t = TTree("configtree", "BumpHunter Config")
t.Branch("config", config_array, "config[%s]/C" % len(config_array))
t.Fill()
t.Write()
def __init__(self,name='tree',title='tree',ttree=None):
""" empty constructor with a name and title.
"""
#if (ttree == None): self._constructor_empty(name,title)
#else: self._constructor_copy(self,ttree)
self.name = name
self.labels = []
self.varnames = []
self.vars = {}
self.nentries = 0
self.ttree = TTree.__init__(self,name,title)
