def __create_branches__(self):
if self._created_branches:
return
self._model = TreeModelMeta('MyTreeModel', (TreeModel, ),
self._branches)
self._tree = Tree(self._tree_name, model=self._model)
self._created_branches = True
def create_tree():
f = TemporaryFile()
tree = Tree("tree", model=create_model())
# fill the tree
for i in xrange(1000):
assert_equal(tree.a_vect, LorentzVector(0, 0, 0, 0))
random_vect = LorentzVector(gauss(0.5, 1.0), gauss(0.5, 1.0), gauss(0.5, 1.0), gauss(0.5, 1.0))
tree.a_vect.copy_from(random_vect)
assert_equal(tree.a_vect, random_vect)
tree.a_x = gauss(0.5, 1.0)
tree.a_y = gauss(0.3, 2.0)
tree.a_z = gauss(13.0, 42.0)
tree.b_n = randint(1, 5)
for j in xrange(tree.b_n):
vect = LorentzVector(gauss(0.5, 1.0), gauss(0.5, 1.0), gauss(0.5, 1.0), gauss(0.5, 1.0))
tree.b_vect.push_back(vect)
tree.b_x.push_back(randint(1, 10))
tree.b_y.push_back(gauss(0.3, 2.0))
tree.i = i
assert_equal(tree.b_n, tree.b_vect.size())
assert_equal(tree.b_n, tree.b_x.size())
assert_equal(tree.b_n, tree.b_y.size())
tree.fill(reset=True)
tree.write()
# TFile.Close the file but keep the underlying
# tempfile file descriptor open
ROOT.TFile.Close(f)
FILES.append(f)
FILE_PATHS.append(f.GetName())
def create_tree():
f = TemporaryFile()
tree = Tree("tree", model=create_model())
# fill the tree
for i in xrange(1000):
tree.a_x = gauss(.5, 1.)
tree.a_y = gauss(.3, 2.)
tree.a_z = gauss(13., 42.)
tree.b_vect.clear()
tree.b_x.clear()
tree.b_y.clear()
tree.b_n = randint(1, 5)
for j in xrange(tree.b_n):
vect = LorentzVector(
gauss(.5, 1.),
gauss(.5, 1.),
gauss(.5, 1.),
gauss(.5, 1.))
tree.b_vect.push_back(vect)
tree.b_x.push_back(randint(1, 10))
tree.b_y.push_back(gauss(.3, 2.))
tree.i = i
tree.fill()
tree.write()
# TFile.Close the file but keep the underlying
# tempfile file descriptor open
ROOT.TFile.Close(f)
FILES.append(f)
FILE_PATHS.append(f.GetName())
def np2root(data, column_names, outname="output.root",tname="tree",dtype=float):
"""
converts numpy array to ROOT TTree and file.
:param data: the 2D array containing M variables for N events
:param column_names: M variables
:param outname: name of the output root file
:param dtype: float or int or list or dictionary. will map columns to data types in ROOT tree.
:return:
"""
# adding support for different types.
branches = {}
if not (isinstance(dtype,dict) or isinstance(dtype,list)):
assert dtype in [float, int], "dtype not understood"
mtype = FloatCol
if dtype == int: mtype = IntCol
branches = {col: mtype() for col in column_names}
elif isinstance(dtype,dict):
my_map = { col : FloatCol if val == float else IntCol for col,val in dtype.iteritems()}
branches = {col: my_map[col]() for col in column_names}
else:
my_map = [ FloatCol if val == float else IntCol for val in dtype]
branches = {col: my_map[i]() for i,col in enumerate(column_names)}
fOut = root_open(outname,"RECREATE")
tree = Tree(tname)
tree.create_branches(branches)
rows, cols = shape(data)
for i in range(0, rows):
for j in range(0, cols):
exec("tree.{col} = {val}".format(col=column_names[j], val=data[i,j])) in locals()
tree.Fill()
fOut.Write()
fOut.Close()
print 'wrote ROOT file {name}'.format(name=outname)
def test_file_assoc():
with TemporaryFile() as f1:
t = Tree()
with TemporaryFile() as f2:
pass
#f1.cd() <== this should not be needed!
# the tree should "remember" what file it was created in
t.Write()
def create_tree(num_branches, num_entries):
branches = dict([(random_name(10), 'F') for x in xrange(num_branches)])
tree = Tree()
tree.create_branches(branches)
# fill the tree with zeros
for i in xrange(num_entries):
tree.fill()
return tree
def create_test_tree():
tree = Tree("test")
tree.create_branches({'x': 'F', 'y': 'F', 'z': 'F', 'i': 'I'})
for i in xrange(10000):
tree.x = gauss(.5, 1.)
tree.y = gauss(.3, 2.)
tree.z = gauss(13., 42.)
tree.i = i
tree.fill()
return tree
class rootifier(object):
'''
This class uses the digital pulse processing tools found in the KData software package
Please email [email protected] for access to KData tools if you wish to use these.
However, you could also look into using the scipy pulse processing tools.
'''
def __init__(self, outputFile, pause=False):
self.data = {}
self.file = open(outputFile, 'recreate')
self.tree = Tree("t", model=Event)
self.epoch = datetime.datetime(1995, 1, 1) #jan 1 1995 is ROOT epoch
def write(self):
self.file.cd()
self.tree.write()
def close(self):
self.file.Close()
def handleEvent(self, event):
#
# here is where you get each event, which is all four digitized pulses
# do analysis here, or convert to ROOT file.
#
#
if int(event.event) % 100 == 0:
print 'event', event.event, event.eventTime
for chan in event.chanData.iterkeys():
if event.chanData[chan]['misread'] == True:
return #skip this event and go to the next...
try:
#eventTime = datetime.datetime(event.eventTime)
tempTime = event.eventTime.split(".")[0]
eventTime = datetime.datetime.strptime(tempTime,
'%Y/%m/%d %H:%M:%S')
td = eventTime - self.epoch
self.tree.date = (
td.microseconds +
(td.seconds + td.days * 24 * 3600) * 10**6) / 10**6 + int(
event.eventTime.split(".")[1]) * 10**-3
except:
print event.eventTime, "is bad"
pass
for chan in event.chanData.iterkeys():
if chan == "CHN1":
for value in event.chanData[chan]["y"]:
self.tree.c1y_val.push_back(value)
for value in event.chanData[chan]["x"]:
self.tree.c1x_val.push_back(value)
self.tree.fill(reset=True)
def rootCombine(file1, file2, fileOut):
fIn = TFile(file1, 'READ')
tree1 = fIn.Get('tree') #ONLY FOR TREES NAMED TREE
fIn2 = TFile(file2, 'READ')
tree2 = fIn.Get('tree') #ONLY FOR TREES NAMED TREE
# fOut = TFile(fileOut, "recreate")
fOut = root_open(fileOut, 'recreate')
tree = Tree('tree', model = Event)
# tree.num_vals = 120
# print tree.num_vals
ctr = 0
for aEvent, bEvent in izip(tree1, tree2):
ctr += 1
if (ctr % 100 == 0):
print('Event: %d' %ctr)
tree.iEvents = aEvent.iEvents
tree.nFinalParticles = (aEvent.nFinalParticles + bEvent.nFinalParticles)
if (tree.nFinalParticles > ARRAY_LENGTH):
print ('event: %d contains too many particles: %d' %(aEvent.iEvents, tree.nFinalParticles))
break
for i in range(aEvent.nFinalParticles):
# print aEvent.px[i]
tree.px[i] = aEvent.px[i]
tree.py[i] = aEvent.py[i]
tree.pz[i] = aEvent.pz[i]
tree.energy[i] = aEvent.energy[i]
# tree.num_vals+=1
for i in range(bEvent.nFinalParticles):
tree.px[aEvent.nFinalParticles + i] = bEvent.px[i]
tree.py[aEvent.nFinalParticles + i] = bEvent.py[i]
tree.pz[aEvent.nFinalParticles + i] = bEvent.pz[i]
tree.energy[aEvent.nFinalParticles + i] = bEvent.energy[i]
# tree.num_vals+=1
tree.fill()
fOut.write()
tree.px.reset()
tree.py.reset()
tree.pz.reset()
tree.energy.reset()
# tree.csv()
fOut.close()
print 'Done'
def __init__(self, outdir, outfile, inputfiles, issignal):
# Setup
self.outdir = outdir
self.outfile = outfile
self.issignal = issignal
# Create TChain reading from ntuples
self.intree = r.TChain('LDMX_Events')
for f in inputfiles:
self.intree.Add(f)
# Setup to read the collections we care about
self.evHeader = r.ldmx.EventHeader()
self.ecalVetoRes = r.TClonesArray('ldmx::EcalVetoResult')
#self.ecalVetoResFernand = r.TClonesArray('ldmx::EcalVetoResult')
self.hcalVetoRes = r.TClonesArray('ldmx::HcalVetoResult')
self.trackerVetoRes = r.TClonesArray('ldmx::TrackerVetoResult')
self.hcalhits = r.TClonesArray('ldmx::HcalHit') #added by Jack
self.trigRes = r.TClonesArray('ldmx::TriggerResult')
self.hcalhits = r.TClonesArray('ldmx::HcalHit') #added by Jack
self.ecalhits = r.TClonesArray('ldmx::EcalHit')
self.simParticles = r.TClonesArray('ldmx::SimParticle')
self.spHits = r.TClonesArray('ldmx::SimTrackerHit')
self.targetSPHits = r.TClonesArray('ldmx::SimTrackerHit')
self.intree.SetBranchAddress('EventHeader', r.AddressOf(self.evHeader))
self.intree.SetBranchAddress('EcalVeto_recon',
r.AddressOf(self.ecalVetoRes))
#self.intree.SetBranchAddress('EcalVetoFernand_recon',r.AddressOf(self.ecalVetoResFernand))
self.intree.SetBranchAddress('HcalVeto_recon',
r.AddressOf(self.hcalVetoRes))
self.intree.SetBranchAddress('TrackerVeto_recon',
r.AddressOf(self.trackerVetoRes))
self.intree.SetBranchAddress('ecalDigis_recon',
r.AddressOf(self.ecalhits))
self.intree.SetBranchAddress('hcalDigis_recon',
r.AddressOf(self.hcalhits))
self.intree.SetBranchAddress('SimParticles_sim',
r.AddressOf(self.simParticles))
self.intree.SetBranchAddress('HcalScoringPlaneHits_sim',
r.AddressOf(self.spHits))
self.intree.SetBranchAddress('TargetScoringPlaneHits_sim',
r.AddressOf(self.targetSPHits))
#if not self.issignal:
#self.intree.SetBranchAddress('Trigger_recon',r.AddressOf(self.trigRes))
# Create output file and tree
self.tfile = root_open(self.outfile, 'recreate')
self.tree = Tree('EcalVeto', model=EcalVetoEvent)
# Initialize event count
self.event_count = 0
class rootifier(object):
'''
This class uses the digital pulse processing tools found in the KData software package
Please email [email protected] for access to KData tools if you wish to use these.
However, you could also look into using the scipy pulse processing tools.
'''
def __init__(self, outputFile, pause=False):
self.data = {}
self.file = open(outputFile,'recreate')
self.tree = Tree("t", model=Event)
self.epoch = datetime.datetime(1995, 1, 1) #jan 1 1995 is ROOT epoch
def write(self):
self.file.cd()
self.tree.write()
def close(self):
self.file.Close()
def handleEvent(self, event):
#
# here is where you get each event, which is all four digitized pulses
# do analysis here, or convert to ROOT file.
#
#
if int(event.event) % 100 == 0:
print 'event', event.event, event.eventTime
for chan in event.chanData.iterkeys():
if event.chanData[chan]['misread'] == True:
return #skip this event and go to the next...
try:
#eventTime = datetime.datetime(event.eventTime)
tempTime = event.eventTime.split(".")[0]
eventTime = datetime.datetime.strptime(tempTime,'%Y/%m/%d %H:%M:%S')
td = eventTime - self.epoch
self.tree.date = (td.microseconds + (td.seconds + td.days * 24 * 3600) * 10**6) / 10**6 + int(event.eventTime.split(".")[1])*10**-3
except:
print event.eventTime, "is bad"
pass
for chan in event.chanData.iterkeys():
if chan == "CHN1":
for value in event.chanData[chan]["y"]:
self.tree.c1y_val.push_back(value)
for value in event.chanData[chan]["x"]:
self.tree.c1x_val.push_back(value)
self.tree.fill(reset=True)
def __init__(self,
files=None,
type='MC',
maxevents=-1,
channel=['2mu2e', '4mu'],
**kwargs):
super(MyEvents, self).__init__(files=files,
type=type,
maxevents=maxevents,
channel=channel,
**kwargs)
self._outf = root_open(self.OutName, 'recreate')
self._outt = Tree("egmlj")
self._outt.create_branches({
's_MXX': 'F',
's_MA': 'F',
's_LXY': 'F',
'dp_lxy': 'F',
'dp_lz': 'F',
'dp_pt': 'F',
'dp_eta': 'F',
'dp_daudr': 'F',
'rho': 'F',
'lj_pt': 'F',
'lj_gendr': 'F',
'lj_nele': 'I',
'lj_npho': 'I',
'lj_ndau': 'I',
'lj_area': 'F',
'j_pt': 'F',
'j_rawpt': 'F',
'j_gendr': 'F',
'j_hadfrac': 'F',
'j_emfrac': 'F',
'j_neuhadfrac': 'F',
'j_neuemfrac': 'F',
'j_id': 'I',
'j_ncands': 'I',
'ele_gendr': 'F',
'ele_pt': 'F',
'ele_idbit': 'I',
'ele_grade': 'I',
'pho_gendr': 'F',
'pho_pt': 'F',
'pho_isconv': 'I',
'pho_haspix': 'I',
'pho_idbit': 'I',
'pho_grade': 'I',
'ljsrc_pt': 'F',
'ljsrc_gendr': 'F',
'ljsrc_type': 'I', # 2: electron, 4: photon
})
def add_tree_to_file(self,
file_name,
table_name,
table_title,
headers,
data,
mode="update"):
"""
:param table_name
:param table_title
:param headers: should be a dictionary defining each column name and it's data type.
Valid data types are Floats (F) and Integers (I).
So to define a simple X, Y plot where X is an Integer and Y is a Float,
we'd define {'X': 'I', 'Y': 'F'}
:param data: defined as an array of dictionary objects defining the values for each record.
For example, X and Y in the above case would be [{'X':0,'Y':23.12}, {'X':1,'Y':20.12}, ...]
"""
f = root_open(file_name, mode)
new_tree = Tree(name=table_name, title=table_title)
new_tree.create_branches(headers)
for record in data:
for key, value in record.iteritems():
new_tree[key] = value
new_tree.fill()
new_tree.write()
f.close()
def setup_class(cls):
class ObjectA(TreeModel):
# A simple tree object
x = FloatCol()
y = FloatCol()
z = FloatCol()
class ObjectB(TreeModel):
# A tree object collection
x = stl.vector('int')
y = stl.vector('float')
vect = stl.vector('TLorentzVector')
# collection size
n = IntCol()
class Event(ObjectA.prefix('a_') + ObjectB.prefix('b_')):
i = IntCol()
for i in range(5):
f = TemporaryFile()
tree = Tree("tree", model=Event)
# fill the tree
for i in xrange(10000):
tree.a_x = gauss(.5, 1.)
tree.a_y = gauss(.3, 2.)
tree.a_z = gauss(13., 42.)
tree.b_vect.clear()
tree.b_x.clear()
tree.b_y.clear()
tree.b_n = randint(1, 5)
for j in xrange(tree.b_n):
vect = LorentzVector(
gauss(.5, 1.),
gauss(.5, 1.),
gauss(.5, 1.),
gauss(.5, 1.))
tree.b_vect.push_back(vect)
tree.b_x.push_back(randint(1, 10))
tree.b_y.push_back(gauss(.3, 2.))
tree.i = i
tree.fill()
tree.write()
# TFile.Close the file but keep the underlying
# tempfile file descriptor open
ROOT.TFile.Close(f)
cls.files.append(f)
cls.file_paths.append(f.GetName())
def __init__(self,outputFileName,treeName) :
# Open/recreate output file
self.theTreeFile = root_open(outputFileName,"RECREATE")
# Create tree with given name and branches structure
self.theTree = Tree(treeName)
def setup_class(cls):
cls.temp_file = TemporaryFile()
cls.temp_file_path = cls.temp_file.GetName()
class ObjectA(TreeModel):
# A simple tree object
x = FloatCol()
y = FloatCol()
z = FloatCol()
class ObjectB(TreeModel):
# A tree object collection
x = ROOT.vector('int')
y = ROOT.vector('float')
vect = ROOT.vector('TLorentzVector')
# collection size
n = IntCol()
class Event(ObjectA.prefix('a_') + ObjectB.prefix('b_')):
i = IntCol()
tree = Tree("tree", model=Event)
# fill the tree
for i in xrange(10000):
tree.a_x = gauss(.5, 1.)
tree.a_y = gauss(.3, 2.)
tree.a_z = gauss(13., 42.)
tree.b_vect.clear()
tree.b_x.clear()
tree.b_y.clear()
tree.b_n = randint(1, 5)
for j in xrange(tree.b_n):
vect = LorentzVector(
gauss(.5, 1.),
gauss(.5, 1.),
gauss(.5, 1.),
gauss(.5, 1.))
tree.b_vect.push_back(vect)
tree.b_x.push_back(randint(1, 10))
tree.b_y.push_back(gauss(.3, 2.))
tree.i = i
tree.fill()
tree.write()
cls.temp_file.write()
def create_test_tree(filename='test.root'):
with File.open (filename, 'recreate') as f:
tree = Tree("test")
tree.create_branches(
{'x': 'F',
'y': 'F',
'z': 'F',
'i': 'I',
'EventWeight': "F"})
for i in xrange(10000):
tree.x = gauss(.5, 1.)
tree.y = gauss(.3, 2.)
tree.z = gauss(13., 42.)
tree.i = i
tree.EventWeight = 1.
tree.fill()
f.write()
def create_test_tree():
tree = Tree("test")
tree.create_branches(
{'x': 'F',
'y': 'F',
'z': 'F',
'i': 'I'})
for i in xrange(10000):
tree.x = gauss(.5, 1.)
tree.y = gauss(.3, 2.)
tree.z = gauss(13., 42.)
tree.i = i
tree.fill()
return tree
def __init__(self,
files=None,
type='MC',
maxevents=-1,
channel=['2mu2e', '4mu'],
**kwargs):
super(MyEvents, self).__init__(files=files,
type=type,
maxevents=maxevents,
channel=channel,
**kwargs)
self._outf = root_open(self.OutName, 'recreate')
self._outt = Tree("XX")
self._outt.create_branches({
's_MXX': 'F',
's_MA': 'F',
's_LXY': 'F',
'xx_mass': 'F',
})
def save_hists_to_root_file(root_filename, hists_to_save):
f = TFile(root_filename, "RECREATE")
main_tree = Tree(name="main_tree")
for var in hists_to_save:
a = hists_to_save[var][0].hist()
print a
a = pickle.dump(a, open("plot.p", "wb"))
def setup_class(cls):
class ObjectA(TreeModel):
# A simple tree object
x = FloatCol()
y = FloatCol()
z = FloatCol()
class ObjectB(TreeModel):
# A tree object collection
x = stl.vector("int")
y = stl.vector("float")
vect = stl.vector("TLorentzVector")
# collection size
n = IntCol()
class Event(ObjectA.prefix("a_") + ObjectB.prefix("b_")):
i = IntCol()
for i in range(5):
f = TemporaryFile()
tree = Tree("tree", model=Event)
# fill the tree
for i in xrange(10000):
tree.a_x = gauss(0.5, 1.0)
tree.a_y = gauss(0.3, 2.0)
tree.a_z = gauss(13.0, 42.0)
tree.b_vect.clear()
tree.b_x.clear()
tree.b_y.clear()
tree.b_n = randint(1, 5)
for j in xrange(tree.b_n):
vect = LorentzVector(gauss(0.5, 1.0), gauss(0.5, 1.0), gauss(0.5, 1.0), gauss(0.5, 1.0))
tree.b_vect.push_back(vect)
tree.b_x.push_back(randint(1, 10))
tree.b_y.push_back(gauss(0.3, 2.0))
tree.i = i
tree.fill()
tree.write()
cls.files.append(f)
cls.file_paths.append(f.GetName())
def replaceBranch(str_infile, str_intree, float_factor, str_outfile):
"""
Replaces a branch in a tree by multiplying it
It only works for .EventWeight . hardcoded
Inputs : infilename , intreename, multiplicative factor, oufilename
"""
from rootpy.tree import Tree, TreeModel, TreeChain
from rootpy.io import root_open
from rootpy import asrootpy
chain = TreeChain(name=str_intree, files=str_infile)
f_copy = root_open(str_outfile,'recreate')
tree_copy = Tree(str_intree)
tree_copy.set_buffer(chain._buffer, create_branches=True)
for entry in chain:
entry.EventWeight = entry.EventWeight *float_factor
tree_copy.Fill()
tree_copy.Write()
f_copy.Close()
def open_file(self, filename):
self.__t_file_out = File(filename, "recreate")
self.__t_tree_ppd = Tree("t_ppd", "platform parameters data")
self.__t_tree_ppd.create_branches({
"pitch_angle" : "D" ,
"yaw_angle" : "D" ,
"roll_angle" : "D" ,
"pitch_angle_v" : "D" ,
"yaw_angle_v" : "D" ,
"roll_angle_v" : "D" ,
"orbit_agl_v" : "D" ,
"longitude" : "D" ,
"latitude" : "D" ,
"geocentric_d" : "D" ,
"ship_time_sec" : "D" ,
"utc_time_sec" : "D" ,
"utc_time_str" : "C[32]" ,
"flag_of_pos" : "I" ,
"wgs84_x" : "D" ,
"wgs84_y" : "D" ,
"wgs84_z" : "D" ,
"wgs84_x_v" : "D" ,
"wgs84_y_v" : "D" ,
"wgs84_z_v" : "D" ,
"det_z_lat" : "D" ,
"det_z_lon" : "D" ,
"det_z_ra" : "D" ,
"det_z_dec" : "D" ,
"det_x_lat" : "D" ,
"det_x_lon" : "D" ,
"det_x_ra" : "D" ,
"det_x_dec" : "D" ,
"earth_lat" : "D" ,
"earth_lon" : "D" ,
"earth_ra" : "D" ,
"earth_dec" : "D" ,
"sun_ra" : "D" ,
"sun_dec" : "D"
})
def mergeChannel(channel, fileList):
'''
Merge one channel for all files, return merged tree.
Should be run with a file open.
'''
if isinstance(fileList, str):
fileList = [fileList]
chain = TreeChain('{}/ntuple'.format(channel), fileList)
out = Tree('ntuple'.format(channel))
out.set_buffer(chain._buffer, create_branches=True)
found = set()
for ev in chain:
evID = (ev.run, ev.lumi, ev.evt)
if evID in found:
continue
out.fill()
found.add(evID)
return out
def create_tree(num_branches, num_entries):
branches = dict([(random_name(10), 'F') for x in xrange(num_branches)])
tree = Tree()
tree.create_branches(branches)
# fill the tree with zeros
for i in xrange(num_entries):
tree.fill()
return tree
def __init__( self, root_branch = None):
if root_branch is None:
self.name = 'unknown'
self.size = 0
self.zipped_bytes = 0
self.type = 'unknown'
else:
self.name = root_branch.GetName()
self.size = root_branch.GetTotalSize()
self.zipped_bytes = root_branch.GetZipBytes()
self.type = Tree.branch_type(root_branch)
self.subitems = {}
if self.has_group():
token = self.name.split( BranchInfo.group_delimiter )
self.group = '.'.join(token[:-1])
else:
self.group = self.name
class MyEvents(SignalEvents):
def __init__(self,
files=None,
type='MC',
maxevents=-1,
channel=['2mu2e', '4mu'],
**kwargs):
super(MyEvents, self).__init__(files=files,
type=type,
maxevents=maxevents,
channel=channel,
**kwargs)
self._outf = root_open(self.OutName, 'recreate')
self._outt = Tree("XX")
self._outt.create_branches({
's_MXX': 'F',
's_MA': 'F',
's_LXY': 'F',
'xx_mass': 'F',
})
def processEvent(self, event, aux):
if aux['channel'] not in self.Channel: return
chan = aux['channel']
_leptonjets = []
for lj in event.leptonjets:
if not lj.passSelection(event): continue
_leptonjets.append(lj)
if len(_leptonjets) < 2: return
self._outt.s_MXX = self.SignalParam['MXX']
self._outt.s_MA = self.SignalParam['MA']
self._outt.s_LXY = self.SignalParam['LXY']
self._outt.xx_mass = (_leptonjets[0].p4 + _leptonjets[1].p4).M()
self._outt.fill()
def postProcess(self):
super(MyEvents, self).postProcess()
self._outt.write()
self._outf.close()
def mergeChannel(channel, fileList):
'''
Merge one channel for all files, return merged tree.
Should be run with a file open.
'''
if isinstance(fileList, str):
fileList = [fileList]
chain = TreeChain('{}/ntuple'.format(channel), fileList)
out = Tree('ntuple'.format(channel))
out.set_buffer(chain._buffer, create_branches=True)
found = set()
for ev in chain:
evID = (ev.run, ev.lumi, ev.evt)
if evID in found:
continue
out.fill()
found.add(evID)
return out
class Writer:
def __init__(self,outputFileName,treeName) :
# Open/recreate output file
self.theTreeFile = root_open(outputFileName,"RECREATE")
# Create tree with given name and branches structure
self.theTree = Tree(treeName)
def addBranches(self,branches) :
self.theTree.create_branches(branches)
def fill(self) :
self.theTree.fill()
def writeAndClose(self) :
# Write the tree to the file
self.theTree.write()
self.theTreeFile.close()
def getTree(self) :
return self.theTree
from rootpy.vector import LorentzVector
from rootpy.tree import Tree, TreeModel, IntCol
from rootpy.io import root_open
from rootpy import stl
from random import gauss
f = root_open("test.root", "recreate")
# define the model
class Event(TreeModel):
x = stl.vector('TLorentzVector')
i = IntCol()
tree = Tree("test", model=Event)
# fill the tree
for i in xrange(100):
tree.x.clear()
for j in xrange(5):
vect = LorentzVector(
gauss(.5, 1.),
gauss(.5, 1.),
gauss(.5, 1.),
gauss(.5, 1.))
tree.x.push_back(vect)
tree.i = i
tree.fill()
tree.write()
a_y = FloatCol()
a_z = FloatCol()
# properties of particle "b"
b_x = FloatCol()
b_y = FloatCol()
b_z = FloatCol()
# a collection of particles
col_x = stl.vector("float")
col_y = stl.vector("float")
col_z = stl.vector("float")
col_n = IntCol()
# a TLorentzVector
p = LorentzVector
i = IntCol()
tree = Tree("test", model=Event)
# fill the tree
for i in range(10):
tree.a_x = gauss(.5, 1.)
tree.a_y = gauss(.3, 2.)
tree.a_z = gauss(13., 42.)
tree.b_x = gauss(.5, 1.)
tree.b_y = gauss(.3, 2.)
tree.b_z = gauss(13., 42.)
n = randint(1, 10)
for j in range(n):
tree.col_x.push_back(gauss(.5, 1.))
tree.col_y.push_back(gauss(.3, 2.))
from rootpy.io import open
from rootpy.types import *
from random import gauss
f = open("test.root", "recreate")
# define the model
class Event(TreeModel):
x = FloatCol()
y = FloatCol()
z = FloatCol()
i = IntCol()
tree = Tree("test", model=Event)
# fill the tree
for i in xrange(10000):
tree.x = gauss(.5, 1.)
tree.y = gauss(.3, 2.)
tree.z = gauss(13., 42.)
tree.i = i
tree.fill()
tree.write()
# write tree in CSV format
tree.csv()
f.close()
"filename", help="ROOT file that stores decoded platform parameters data")
parser.add_argument(
"-o",
dest="outfile",
help="ROOT file to store platform parameters data of Level 1",
default="TG2_PPD_file_L1.root")
args = parser.parse_args()
t_file_in = File(args.filename, 'read')
t_tree_ppd_in = t_file_in.get('t_ppd')
t_tree_ppd_in.create_buffer()
m_shipspan = t_file_in.get('m_shipspan')
m_utc_span = t_file_in.get('m_utc_span')
t_file_out = File(args.outfile, 'recreate')
t_tree_ppd_out = Tree("t_ppd", "platform parameters data")
t_tree_ppd_out.create_branches({
"longitude": "D",
"latitude": "D",
"geocentric_d": "D",
"ship_time_sec": "D",
"utc_time_sec": "D",
"utc_time_str": "C[32]",
"flag_of_pos": "I",
"det_z_ra": "D",
"det_z_dec": "D",
"det_x_ra": "D",
"det_x_dec": "D",
"earth_ra": "D",
"earth_dec": "D"
})
import glob
# -- hack
files = args.files # -- glob.glob(args.files[0] + '/*.root')
ROOTfile = None
# ROOTfile = None
for i, fname in enumerate(files):
logger.info('ON: CHUNK #{}'.format(CURRENT_CHUNK))
try:
if args.dump and ROOTfile is None:
logger.info('Making ROOT file: {}'.format(
args.dump + '-chnk{}.root'.format(CURRENT_CHUNK)))
ROOTfile = root_open(
args.dump + '-chnk{}.root'.format(CURRENT_CHUNK), "recreate")
tree = Tree('images', model=JetImage)
except Exception:
continue
logger.info('({} of {}) working on file: {}'.format(
i, len(files), fname))
try:
with root_open(fname) as f:
df = f.EventTree.to_array()
n_entries = df.shape[0]
pix = df[0]['Intensity'].shape[0]
if not perfectsquare(pix):
raise ValueError('shape of image array must be square.')
def tokenize(line):
tokens = line.strip().split()
tokens = [i.strip() for i in tokens if i.strip()]
ret = []
for token in tokens:
if token.startswith('#'):
break
else:
ret.append(token)
return ret
with root_open('test.root', 'w') as outfile:
with open('Summer15_25nsV6_MC_L1FastJet_AK8PFchs.txt') as infile:
print "opening file"
corr_tree = Tree('L1')
corr_tree.create_branches({
'min': 'F',
'max': 'F',
})
fcn = None
bin_vars = None
fcn_vars = None
n_bin_vars = None
n_fcn_vars = None
for line in infile:
tokens = tokenize(line)
print "parsing line", line
if not fcn:
tokens = [i.strip('{}') for i in tokens]
n_bin_vars = int(tokens[0])
# define the model
class Event(TreeModel):
s = CharCol()
string = CharArrayCol(5)
x = FloatCol()
y = FloatCol()
z = FloatCol()
f = FloatArrayCol(5)
num_vals = IntCol()
# variable-length array
vals = FloatArrayCol(5, length_name='num_vals')
tree = Tree("test", model=Event)
# fill the tree
for i in range(5):
tree.s = ord(choice(ascii_letters))
tree.string = (u''.join(sample(ascii_letters, 4))).encode('ascii')
tree.x = gauss(.5, 1.)
tree.y = gauss(.3, 2.)
tree.z = gauss(13., 42.)
for j in range(5):
tree.f[j] = gauss(-2, 5)
tree.num_vals = i
for j in range(i):
tree.vals[j] = j
tree.fill()
tree.write()
def __init__(self, fns, outn, outd, model):
print "Initializing Container!"
#self.tin = r.TChain('EcalVeto')
#for fn in fns:
# self.tin.Add(fn)
#self.tfile = root_open(fn,'r+')
#with root_open(fn,'r+') as f:
#self.tin = self.tfile.EcalVeto
#self.tin.Print()
self.tin = TreeChain('EcalVeto', fns)
self.tin.create_branches({'discValue_gabrielle': 'F'})
self.model = model
self.outdir = outd
self.outname = outn
self.outfile = root_open(outn, 'RECREATE')
self.tout = Tree('EcalVeto')
self.tout.set_buffer(self.tin._buffer, create_branches=True)
self.events = []
#print self.tin.GetEntries()
for event in self.tin:
#if len(self.events)>10:
# continue
evt = []
################################### Features #######################################
evt.append(event.nReadoutHits)
evt.append(event.summedDet)
evt.append(event.summedTightIso)
evt.append(event.maxCellDep)
evt.append(event.showerRMS)
evt.append(event.xStd)
evt.append(event.yStd)
evt.append(event.avgLayerHit)
evt.append(event.deepestLayerHit)
evt.append(event.stdLayerHit)
#new features
evt.append(event.ele68ContEnergy)
evt.append(event.ele68x2ContEnergy)
evt.append(event.ele68x3ContEnergy)
evt.append(event.ele68x4ContEnergy)
evt.append(event.ele68x5ContEnergy)
evt.append(event.photon68ContEnergy)
evt.append(event.photon68x2ContEnergy)
evt.append(event.photon68x3ContEnergy)
evt.append(event.photon68x4ContEnergy)
evt.append(event.photon68x5ContEnergy)
evt.append(event.outside68ContEnergy)
evt.append(event.outside68x2ContEnergy)
evt.append(event.outside68x3ContEnergy)
evt.append(event.outside68x4ContEnergy)
evt.append(event.outside68x5ContEnergy)
evt.append(event.outside68ContNHits)
evt.append(event.outside68x2ContNHits)
evt.append(event.outside68x3ContNHits)
evt.append(event.outside68x4ContNHits)
evt.append(event.outside68x5ContNHits)
evt.append(event.outside68ContXstd)
evt.append(event.outside68x2ContXstd)
evt.append(event.outside68x3ContXstd)
evt.append(event.outside68x4ContXstd)
evt.append(event.outside68x5ContXstd)
evt.append(event.outside68ContYstd)
evt.append(event.outside68x2ContYstd)
evt.append(event.outside68x3ContYstd)
evt.append(event.outside68x4ContYstd)
evt.append(event.outside68x5ContYstd)
evt.append(event.ecalBackEnergy)
evtarray = np.array([evt])
pred = float(model.predict(xgb.DMatrix(evtarray))[0])
#print pred
event.discValue_gabrielle = pred
self.tout.Fill()
######################################################################################
self.events.append(evt)
if (len(self.events) % 10000 == 0 and len(self.events) > 0):
print 'The shape of events = ', np.shape(self.events)
self.outfile.cd()
self.tout.Write()
self.outfile.Close()
#self.tfile.Close()
print 'cp %s %s' % (self.outname, self.outdir)
os.system('cp %s %s' % (self.outname, self.outdir))
def testPhotonBomb(self):
# Run only one photon at a time
#nphotons = 7200000
#nphotons = 256*10000
nphotons = 256 * 1000
dphi = np.random.uniform(0, 2.0 * np.pi, nphotons)
dcos = np.random.uniform(-1.0, 1.0, nphotons)
dir = np.array(zip(
np.sqrt(1 - dcos[:] * dcos[:]) * np.cos(dphi[:]),
np.sqrt(1 - dcos[:] * dcos[:]) * np.sin(dphi[:]), dcos[:]),
dtype=np.float32)
pos = np.tile([-200.0, -400.0, -500.0],
(nphotons, 1)).astype(np.float32)
pol = np.zeros_like(pos)
phi = np.random.uniform(0, 2 * np.pi, nphotons).astype(np.float32)
pol[:, 0] = np.cos(phi)
pol[:, 1] = np.sin(phi)
pol = np.cross(pol, dir)
for n, p in enumerate(pol):
norm = np.sqrt(p[0] * p[0] + p[1] * p[1] + p[2] * p[2])
p /= norm
t = np.zeros(nphotons,
dtype=np.float32) + 100.0 # Avoid negative photon times
wavelengths = np.empty(nphotons, np.float32)
wavelengths.fill(128.0)
photons = Photons(pos=pos,
dir=dir,
pol=pol,
t=t,
wavelengths=wavelengths)
hit_charges = []
if has_root:
root_file = root_open("output_test_lbne35t_nowires.root",
"recreate")
root_tree = Tree("PhotonData", model=PhotonData)
root_tree.reset()
for ev in self.sim.simulate((photons for i in xrange(1)),
keep_photons_end=True,
keep_photons_beg=False):
ev.photons_end.dump_history()
# lht = ev.photons_end[0].last_hit_triangles
# nhits = ev.channels.hit[ np.arange(0,30)[:] ]
#
# print "nchannels: ",len(ev.channels.hit)
# print nhits
# print ev.channels.q
# print ev.channels.t
if True:
# Fill Tree
#print "save info for ",len(ev.photons_end)
for photon in ev.photons_end:
root_tree.end_x = photon.pos[0]
root_tree.end_y = photon.pos[1]
root_tree.end_z = photon.pos[2]
root_tree.reflect_diffuse = int(event.REFLECT_DIFFUSE
& photon.flags)
root_tree.reflect_specular = int(event.REFLECT_SPECULAR
& photon.flags)
root_tree.bulk_scatter = int(event.RAYLEIGH_SCATTER
& photon.flags)
root_tree.bulk_absorb = int(event.BULK_ABSORB
& photon.flags)
root_tree.surface_detect = int(event.SURFACE_DETECT
& photon.flags)
root_tree.surface_absorb = int(event.SURFACE_ABSORB
& photon.flags)
root_tree.surface_reemit = int(event.SURFACE_REEMIT
& photon.flags)
root_tree.fill()
if has_root:
root_tree.write()
from rootpy.math.physics.vector import LorentzVector
from rootpy.tree import Tree, TreeModel, IntCol
from rootpy.io import root_open
from rootpy import stl
from random import gauss
f = root_open("test.root", "recreate")
# define the model
class Event(TreeModel):
x = stl.vector('TLorentzVector')
i = IntCol()
tree = Tree("test", model=Event)
# fill the tree
for i in xrange(100):
tree.x.clear()
for j in xrange(5):
vect = LorentzVector(gauss(.5, 1.), gauss(.5, 1.), gauss(.5, 1.),
gauss(.5, 1.))
tree.x.push_back(vect)
tree.i = i
tree.fill()
tree.write()
f.close()
class Event(TreeModel):
"""Event model definition"""
x = FloatCol()
y = FloatCol()
z = FloatCol()
i = IntCol()
# first create several example trees in separate files
fnames = ["test_{0:d}.root".format(i) for i in range(5)]
for fname in fnames:
with root_open(fname, "recreate") as f:
tree = Tree("test", model=Event)
# fill the tree
for i in range(100):
tree.x = gauss(.5, 1.)
tree.y = gauss(.3, 2.)
tree.z = gauss(13., 42.)
tree.i = i
tree.fill()
tree.write()
"""
This section below takes the example trees and copies it while overwriting a
branch with new values.
"""
# first define the chain of trees
This example demonstrates how to create a tree with a variable-length array.
"""
print(__doc__)
from rootpy.tree import Tree, TreeModel, IntCol, FloatArrayCol
from rootpy.io import root_open
class Event(TreeModel):
num_vals = IntCol()
vals = FloatArrayCol(10, length_name='num_vals')
rfile = root_open('test.root', 'w')
tree = Tree('events', model=Event)
for i in range(10):
tree.num_vals = i + 1
for j in range(i + 1):
tree.vals[j] = j
tree.fill()
tree.write()
tree.vals.reset()
tree.csv()
rfile.close()
print("===")
# CSV output from tree read from file should match above output
root_open('test.root', 'r').events.csv()
class bdtTreeMaker:
def __init__(self, outdir, outfile, inputfiles, issignal):
# Setup
self.outdir = outdir
self.outfile = outfile
self.issignal = issignal
# Create TChain reading from ntuples
self.intree = r.TChain('LDMX_Events')
for f in inputfiles:
self.intree.Add(f)
# Setup to read the collections we care about
self.evHeader = r.ldmx.EventHeader()
self.ecalVetoRes = r.TClonesArray('ldmx::EcalVetoResult')
#self.ecalVetoResFernand = r.TClonesArray('ldmx::EcalVetoResult')
self.hcalVetoRes = r.TClonesArray('ldmx::HcalVetoResult')
self.trackerVetoRes = r.TClonesArray('ldmx::TrackerVetoResult')
self.hcalhits = r.TClonesArray('ldmx::HcalHit') #added by Jack
self.trigRes = r.TClonesArray('ldmx::TriggerResult')
self.hcalhits = r.TClonesArray('ldmx::HcalHit') #added by Jack
self.ecalhits = r.TClonesArray('ldmx::EcalHit')
self.simParticles = r.TClonesArray('ldmx::SimParticle')
self.spHits = r.TClonesArray('ldmx::SimTrackerHit')
self.targetSPHits = r.TClonesArray('ldmx::SimTrackerHit')
self.intree.SetBranchAddress('EventHeader', r.AddressOf(self.evHeader))
self.intree.SetBranchAddress('EcalVeto_recon',
r.AddressOf(self.ecalVetoRes))
#self.intree.SetBranchAddress('EcalVetoFernand_recon',r.AddressOf(self.ecalVetoResFernand))
self.intree.SetBranchAddress('HcalVeto_recon',
r.AddressOf(self.hcalVetoRes))
self.intree.SetBranchAddress('TrackerVeto_recon',
r.AddressOf(self.trackerVetoRes))
self.intree.SetBranchAddress('ecalDigis_recon',
r.AddressOf(self.ecalhits))
self.intree.SetBranchAddress('hcalDigis_recon',
r.AddressOf(self.hcalhits))
self.intree.SetBranchAddress('SimParticles_sim',
r.AddressOf(self.simParticles))
self.intree.SetBranchAddress('HcalScoringPlaneHits_sim',
r.AddressOf(self.spHits))
self.intree.SetBranchAddress('TargetScoringPlaneHits_sim',
r.AddressOf(self.targetSPHits))
#if not self.issignal:
#self.intree.SetBranchAddress('Trigger_recon',r.AddressOf(self.trigRes))
# Create output file and tree
self.tfile = root_open(self.outfile, 'recreate')
self.tree = Tree('EcalVeto', model=EcalVetoEvent)
# Initialize event count
self.event_count = 0
# Loop over all input events
def run(self):
while self.event_count < self.intree.GetEntries():
self.intree.GetEntry(self.event_count)
if self.event_count % 1000 == 0:
print 'Processing event ', self.event_count
if self.event_count > 1000:
return
self.process()
# Process an event
def process(self):
self.event_count += 1
self.tree.evtNum = self.evHeader.getEventNumber()
#self.tree.trigPass = 1 if self.issignal else self.trigRes[0].passed()
self.tree.trigPass = 1 #self.trigRes[0].passed()
# BDT input variables
self.tree.nReadoutHits = self.ecalVetoRes[0].getNReadoutHits()
self.tree.summedDet = self.ecalVetoRes[0].getSummedDet()
self.tree.summedTightIso = self.ecalVetoRes[0].getSummedTightIso()
self.tree.maxCellDep = self.ecalVetoRes[0].getMaxCellDep()
self.tree.showerRMS = self.ecalVetoRes[0].getShowerRMS()
self.tree.xStd = self.ecalVetoRes[0].getXStd()
self.tree.yStd = self.ecalVetoRes[0].getYStd()
self.tree.avgLayerHit = self.ecalVetoRes[0].getAvgLayerHit()
self.tree.stdLayerHit = self.ecalVetoRes[0].getStdLayerHit()
self.tree.deepestLayerHit = self.ecalVetoRes[0].getDeepestLayerHit()
kemax = 0
thetamax = 0
pidmax = -1
for particle in self.simParticles:
if particle.getParentCount() > 0 and particle.getParent(
0).getPdgID() == 22:
p = particle.getMomentum()
ke = math.sqrt(p[0]**2 + p[1]**2 + p[2]**2)
if ke > kemax:
kemax = ke
thetamax = math.acos(p[2] / ke)
pidmax = particle.getPdgID()
self.tree.leadhadpid = pidmax
self.tree.leadhadke = kemax
self.tree.leadhadthetaz = thetamax * 180 / math.pi
# Stored value of bdt discriminator
self.tree.discValue = self.ecalVetoRes[0].getDisc()
#self.tree.discValueFernand = self.ecalVetoResFernand[0].getDisc()
#print self.event_count,self.evHeader.getEventNumber(),self.ecalVetoRes[0].getDisc(),self.trigRes[0].passed()
# HCal MaxPE value, needed for HCAL veto
self.tree.hcalMaxPE = self.hcalVetoRes[0].getMaxPEHit().getPE()
self.tree.passHcalVeto = self.hcalVetoRes[0].passesVeto()
self.tree.passTrackerVeto = self.trackerVetoRes[0].passesVeto()
# Need to update to get this from first layer of tracker
recoilPx = self.ecalVetoRes[0].getRecoilMomentum()[0]
recoilPy = self.ecalVetoRes[0].getRecoilMomentum()[1]
recoilPz = self.ecalVetoRes[0].getRecoilMomentum()[2]
recoilX = self.ecalVetoRes[0].getRecoilX()
recoilY = self.ecalVetoRes[0].getRecoilY()
recoilfX = CallX(ecalFaceZ, recoilX, recoilY, scoringPlaneZ, recoilPx,
recoilPy, recoilPz)
recoilfY = CallY(ecalFaceZ, recoilX, recoilY, scoringPlaneZ, recoilPx,
recoilPy, recoilPz)
# Fiducial or not
inside = False
if not recoilX == -9999 and not recoilY == -9999 and not recoilPx == -9999 and not recoilPy == -9999 and not recoilPz == -9999:
for x in cellMap:
xdis = recoilfY - x[2]
ydis = recoilfX - x[1]
celldis = np.sqrt(xdis**2 + ydis**2)
if celldis <= cell_radius:
inside = True
break
self.tree.recoilPx = recoilPx
self.tree.recoilPy = recoilPy
self.tree.recoilPz = recoilPz
self.tree.recoilX = recoilX
self.tree.recoilY = recoilY
self.tree.fiducial = inside
# find electrons at scoring plane
electronsAtSP = getElectronSPHits(self.simParticles, self.spHits,
self.targetSPHits)
#print electronsAtSP
self.tree.nelectrons = len(electronsAtSP)
electronLayerIntercepts = []
photonLayerIntercepts = []
recoilangle = -1
recoil_p = -1
if len(electronsAtSP) > 0:
eleinfo = electronsAtSP[0]
ecalSPHit = eleinfo[0]
targetSPHit = eleinfo[1]
simParticle = eleinfo[2]
pvec = ecalSPHit.getMomentum()
pos = ecalSPHit.getPosition()
recoilangle = getTheta(pvec) * 180. / math.pi
recoil_p = getMagnitude(pvec)
pvec0 = [0, 0, 0]
pos0 = [0, 0, 0]
self.tree.eleP = [pvec[0], pvec[1], pvec[2]]
self.tree.elePosSP = [pos[0], pos[1], pos[2]]
pvec0 = targetSPHit.getMomentum() if targetSPHit != None else [
0, 0, 0
]
pos0 = targetSPHit.getPosition() if targetSPHit != None else [
0, 0, 0
]
photonP = [-pvec0[0], -pvec0[1], 4000.0 - pvec0[2]]
self.tree.elePTarget = [pvec0[0], pvec0[1], pvec0[2]]
self.tree.photonP = [photonP[0], photonP[1], photonP[2]]
self.tree.photonPosSP = [pos0[0], pos0[1], pos0[2]]
#print 'Photon momentum: ',photonP
electronLayerIntercepts = getElectronTrajectory(pvec, pos)
if pvec0 != [0, 0, 0]:
photonLayerIntercepts = getElectronTrajectory(photonP, pos0)
#print electronLayerIntercepts
#print photonLayerIntercepts
# compute new variables based on containment regions
ele68Totals = np.zeros(34)
ele68ContEnergy = 0
ele68x2ContEnergy = 0
ele68x3ContEnergy = 0
ele68x4ContEnergy = 0
ele68x5ContEnergy = 0
photon68Totals = np.zeros(34)
photon68ContEnergy = 0
photon68x2ContEnergy = 0
photon68x3ContEnergy = 0
photon68x4ContEnergy = 0
photon68x5ContEnergy = 0
overlap68Totals = np.zeros(34)
overlap68ContEnergy = 0
overlap68x2ContEnergy = 0
overlap68x3ContEnergy = 0
overlap68x4ContEnergy = 0
overlap68x5ContEnergy = 0
outside68Totals = np.zeros(34)
outside68ContEnergy = 0
outside68x2ContEnergy = 0
outside68x3ContEnergy = 0
outside68x4ContEnergy = 0
outside68x5ContEnergy = 0
outside68NHits = np.zeros(34)
outside68ContNHits = 0
outside68x2ContNHits = 0
outside68x3ContNHits = 0
outside68x4ContNHits = 0
outside68x5ContNHits = 0
outside68Xmean = np.zeros(34)
outside68ContXmean = 0
outside68x2ContXmean = 0
outside68x3ContXmean = 0
outside68x4ContXmean = 0
outside68x5ContXmean = 0
outside68Ymean = np.zeros(34)
outside68ContYmean = 0
outside68x2ContYmean = 0
outside68x3ContYmean = 0
outside68x4ContYmean = 0
outside68x5ContYmean = 0
outside68Xstd = np.zeros(34)
outside68ContXstd = 0
outside68x2ContXstd = 0
outside68x3ContXstd = 0
outside68x4ContXstd = 0
outside68x5ContXstd = 0
outside68Ystd = np.zeros(34)
outside68ContYstd = 0
outside68x2ContYstd = 0
outside68x3ContYstd = 0
outside68x4ContYstd = 0
outside68x5ContYstd = 0
outside68HitPositions = []
outside68x2HitPositions = []
outside68x3HitPositions = []
outside68x4HitPositions = []
outside68x5HitPositions = []
outside68WgtCentroidCoords = (0, 0)
outside68x2WgtCentroidCoords = (0, 0)
outside68x3WgtCentroidCoords = (0, 0)
outside68x4WgtCentroidCoords = (0, 0)
outside68x5WgtCentroidCoords = (0, 0)
trigEnergy = 0
ecalBackEnergy = 0
ir = -1
if recoilangle == -1 or recoil_p == -1:
ir = 1
elif recoilangle < 10 and recoil_p < 500:
ir = 1
elif recoilangle < 10 and recoil_p >= 500:
ir = 2
elif recoilangle <= 20:
ir = 3
else:
ir = 4
ip = 1
for hit in self.ecalhits:
hitE = hit.getEnergy()
if not hitE > 0:
continue
cell, module, layer = mask(hit.getID())
if layer < 20:
trigEnergy += hitE
else:
ecalBackEnergy += hitE
#print 'layer:',layer,hit.getLayer()
mcid_val = 10 * cell + module
hitX = cellMap[mcid.index(mcid_val)][1]
hitY = cellMap[mcid.index(mcid_val)][2]
# get distances of hit from projected positions of electrons and photon
distanceEle = math.sqrt(
(hitX - electronLayerIntercepts[layer][0])**2 +
(hitY - electronLayerIntercepts[layer][1])**2
) if len(electronLayerIntercepts) > 0 else -1
distancePhoton = math.sqrt(
(hitX - photonLayerIntercepts[layer][0])**2 +
(hitY - photonLayerIntercepts[layer][1])**2
) if len(photonLayerIntercepts) > 0 else -1
# add to containment totals depending on distance relative to containment radii
# add to containment totals depending on distance relative to containment radii
# i selects which radius of containment to use
if distanceEle < radius_68[ir][layer] and distanceEle > 0:
ele68Totals[layer] += hitE
ele68ContEnergy += hitE
elif distanceEle >= radius_68[ir][
layer] and distanceEle < 2 * radius_68[ir][layer]:
ele68x2ContEnergy += hitE
elif distanceEle >= 2 * radius_68[ir][
layer] and distanceEle < 3 * radius_68[ir][layer]:
ele68x3ContEnergy += hitE
elif distanceEle >= 3 * radius_68[ir][
layer] and distanceEle < 4 * radius_68[ir][layer]:
ele68x4ContEnergy += hitE
elif distanceEle >= 4 * radius_68[ir][
layer] and distanceEle < 5 * radius_68[ir][layer]:
ele68x5ContEnergy += hitE
if distancePhoton < radius_68[ip][layer] and distancePhoton > 0:
photon68Totals[layer] += hitE
photon68ContEnergy += hitE
elif distancePhoton >= radius_68[ip][
layer] and distancePhoton < 2 * radius_68[ip][layer]:
photon68x2ContEnergy += hitE
elif distancePhoton >= 2 * radius_68[ip][
layer] and distancePhoton < 3 * radius_68[ip][layer]:
photon68x3ContEnergy += hitE
elif distancePhoton >= 3 * radius_68[ip][
layer] and distancePhoton < 4 * radius_68[ip][layer]:
photon68x4ContEnergy += hitE
elif distancePhoton >= 4 * radius_68[ip][
layer] and distancePhoton < 5 * radius_68[ip][layer]:
photon68x5ContEnergy += hitE
if distanceEle < radius_68[ir][
layer] and distancePhoton < radius_68[ip][
layer] and distanceEle > 0 and distancePhoton > 0:
overlap68Totals[layer] += hitE
overlap68ContEnergy += hitE
if distanceEle < 2 * radius_68[ir][
layer] and distancePhoton < 2 * radius_68[ip][
layer] and distanceEle > 0 and distancePhoton > 0:
overlap68x2ContEnergy += hitE
if distanceEle < 3 * radius_68[ir][
layer] and distancePhoton < 3 * radius_68[ip][
layer] and distanceEle > 0 and distancePhoton > 0:
overlap68x3ContEnergy += hitE
if distanceEle < 4 * radius_68[ir][
layer] and distancePhoton < 4 * radius_68[ip][
layer] and distanceEle > 0 and distancePhoton > 0:
overlap68x4ContEnergy += hitE
if distanceEle < 5 * radius_68[ir][
layer] and distancePhoton < 5 * radius_68[ip][
layer] and distanceEle > 0 and distancePhoton > 0:
overlap68x5ContEnergy += hitE
if distanceEle > radius_68[ir][
layer] and distancePhoton > radius_68[ip][layer]:
outside68Totals[layer] += hitE
outside68NHits[layer] += 1
outside68Xmean[layer] += hitX * hitE
outside68Ymean[layer] += hitY * hitE
outside68ContEnergy += hitE
outside68ContNHits += 1
outside68ContXmean += hitX * hitE
outside68ContYmean += hitY * hitE
outside68WgtCentroidCoords = (outside68WgtCentroidCoords[0] +
hitX * hitE,
outside68WgtCentroidCoords[1] +
hitY * hitE)
outside68HitPositions.append((hitX, hitY, layer, hitE))
if distanceEle > 2 * radius_68[ir][
layer] and distancePhoton > 2 * radius_68[ip][layer]:
outside68x2ContEnergy += hitE
outside68x2ContNHits += 1
outside68x2ContXmean += hitX * hitE
outside68x2ContYmean += hitY * hitE
outside68x2WgtCentroidCoords = (
outside68x2WgtCentroidCoords[0] + hitX * hitE,
outside68x2WgtCentroidCoords[1] + hitY * hitE)
outside68x2HitPositions.append((hitX, hitY, layer, hitE))
if distanceEle > 3 * radius_68[ir][
layer] and distancePhoton > 3 * radius_68[ip][layer]:
outside68x3ContEnergy += hitE
outside68x3ContNHits += 1
outside68x3ContXmean += hitX * hitE
outside68x3ContYmean += hitY * hitE
outside68x3WgtCentroidCoords = (
outside68x3WgtCentroidCoords[0] + hitX * hitE,
outside68x3WgtCentroidCoords[1] + hitY * hitE)
outside68x3HitPositions.append((hitX, hitY, layer, hitE))
if distanceEle > 4 * radius_68[ir][
layer] and distancePhoton > 4 * radius_68[ip][layer]:
outside68x4ContEnergy += hitE
outside68x4ContNHits += 1
outside68x4ContXmean += hitX * hitE
outside68x4ContYmean += hitY * hitE
outside68x4WgtCentroidCoords = (
outside68x4WgtCentroidCoords[0] + hitX * hitE,
outside68x4WgtCentroidCoords[1] + hitY * hitE)
outside68x4HitPositions.append((hitX, hitY, layer, hitE))
if distanceEle > 5 * radius_68[ir][
layer] and distancePhoton > 5 * radius_68[ip][layer]:
outside68x5ContEnergy += hitE
outside68x5ContNHits += 1
outside68x5ContXmean += hitX * hitE
outside68x5ContYmean += hitY * hitE
outside68x5WgtCentroidCoords = (
outside68x5WgtCentroidCoords[0] + hitX * hitE,
outside68x5WgtCentroidCoords[1] + hitY * hitE)
outside68x5HitPositions.append((hitX, hitY, layer, hitE))
outside68ContShowerRMS = 0
# mean and std deviations of hits outside containment regions
for hit in outside68HitPositions:
distanceCentroid = math.sqrt(
(hit[0] - outside68WgtCentroidCoords[0])**2 +
(hit[1] - outside68WgtCentroidCoords[1])**2)
outside68ContShowerRMS += distanceCentroid * hit[3]
layer = hit[2]
if outside68Totals[layer] > 0:
outside68Xstd[layer] += (
(hit[0] - (outside68Xmean[layer] / outside68Totals[layer]))
**2) * hit[3]
outside68Ystd[layer] += (
(hit[1] - (outside68Ymean[layer] / outside68Totals[layer]))
**2) * hit[3]
if outside68ContEnergy > 0:
outside68ContXstd += (
(hit[0] -
(outside68ContXmean / outside68ContEnergy))**2) * hit[3]
outside68ContYstd += (
(hit[1] -
(outside68ContYmean / outside68ContEnergy))**2) * hit[3]
outside68x2ContShowerRMS = 0
# mean and std deviations of hits outside containment regions
for hit in outside68x2HitPositions:
distanceCentroid = math.sqrt(
(hit[0] - outside68x2WgtCentroidCoords[0])**2 +
(hit[1] - outside68x2WgtCentroidCoords[1])**2)
outside68x2ContShowerRMS += distanceCentroid * hit[3]
if outside68x2ContEnergy > 0:
outside68x2ContXstd += (
(hit[0] - (outside68x2ContXmean / outside68x2ContEnergy))**
2) * hit[3]
outside68x2ContYstd += (
(hit[1] - (outside68x2ContYmean / outside68x2ContEnergy))**
2) * hit[3]
outside68x3ContShowerRMS = 0
# mean and std deviations of hits outside containment regions
for hit in outside68x3HitPositions:
distanceCentroid = math.sqrt(
(hit[0] - outside68x3WgtCentroidCoords[0])**2 +
(hit[1] - outside68x3WgtCentroidCoords[1])**2)
outside68x3ContShowerRMS += distanceCentroid * hit[3]
if outside68x3ContEnergy > 0:
outside68x3ContXstd += (
(hit[0] - (outside68x3ContXmean / outside68x3ContEnergy))**
2) * hit[3]
outside68x3ContYstd += (
(hit[1] - (outside68x3ContYmean / outside68x3ContEnergy))**
2) * hit[3]
outside68x4ContShowerRMS = 0
# mean and std deviations of hits outside containment regions
for hit in outside68x4HitPositions:
distanceCentroid = math.sqrt(
(hit[0] - outside68x4WgtCentroidCoords[0])**2 +
(hit[1] - outside68x4WgtCentroidCoords[1])**2)
outside68x4ContShowerRMS += distanceCentroid * hit[3]
if outside68x4ContEnergy > 0:
outside68x4ContXstd += (
(hit[0] - (outside68x4ContXmean / outside68x4ContEnergy))**
2) * hit[3]
outside68x4ContYstd += (
(hit[1] - (outside68x4ContYmean / outside68x4ContEnergy))**
2) * hit[3]
outside68x5ContShowerRMS = 0
# mean and std deviations of hits outside containment regions
for hit in outside68x5HitPositions:
distanceCentroid = math.sqrt(
(hit[0] - outside68x5WgtCentroidCoords[0])**2 +
(hit[1] - outside68x5WgtCentroidCoords[1])**2)
outside68x5ContShowerRMS += distanceCentroid * hit[3]
if outside68x5ContEnergy > 0:
outside68x5ContXstd += (
(hit[0] - (outside68x5ContXmean / outside68x5ContEnergy))**
2) * hit[3]
outside68x5ContYstd += (
(hit[1] - (outside68x5ContYmean / outside68x5ContEnergy))**
2) * hit[3]
if outside68ContEnergy > 0:
outside68ContXmean /= outside68ContEnergy
outside68ContYmean /= outside68ContEnergy
outside68ContXstd = math.sqrt(outside68ContXstd /
outside68ContEnergy)
outside68ContYstd = math.sqrt(outside68ContYstd /
outside68ContEnergy)
if outside68x2ContEnergy > 0:
outside68x2ContXmean /= outside68x2ContEnergy
outside68x2ContYmean /= outside68x2ContEnergy
outside68x2ContXstd = math.sqrt(outside68x2ContXstd /
outside68x2ContEnergy)
outside68x2ContYstd = math.sqrt(outside68x2ContYstd /
outside68x2ContEnergy)
if outside68x3ContEnergy > 0:
outside68x3ContXmean /= outside68x3ContEnergy
outside68x3ContYmean /= outside68x3ContEnergy
outside68x3ContXstd = math.sqrt(outside68x3ContXstd /
outside68x3ContEnergy)
outside68x3ContYstd = math.sqrt(outside68x3ContYstd /
outside68x3ContEnergy)
if outside68x4ContEnergy > 0:
outside68x4ContXmean /= outside68x4ContEnergy
outside68x4ContYmean /= outside68x4ContEnergy
outside68x4ContXstd = math.sqrt(outside68x4ContXstd /
outside68x4ContEnergy)
outside68x4ContYstd = math.sqrt(outside68x4ContYstd /
outside68x4ContEnergy)
if outside68x5ContEnergy > 0:
outside68x5ContXmean /= outside68x5ContEnergy
outside68x5ContYmean /= outside68x5ContEnergy
outside68x5ContXstd = math.sqrt(outside68x5ContXstd /
outside68x5ContEnergy)
outside68x5ContYstd = math.sqrt(outside68x5ContYstd /
outside68x5ContEnergy)
for i in range(34):
if outside68Totals[i] > 0:
outside68Xmean[i] /= outside68Totals[i]
outside68Ymean[i] /= outside68Totals[i]
outside68Xstd[i] = math.sqrt(outside68Xstd[i] /
outside68Totals[i])
outside68Ystd[i] = math.sqrt(outside68Ystd[i] /
outside68Totals[i])
self.tree.trigEnergy = trigEnergy
self.tree.ecalBackEnergy = ecalBackEnergy
self.tree.ele68TotalEnergies = ele68Totals
self.tree.photon68TotalEnergies = photon68Totals
self.tree.overlap68TotalEnergies = overlap68Totals
self.tree.outside68TotalEnergies = outside68Totals
self.tree.outside68TotalNHits = outside68NHits
self.tree.outside68Xmeans = outside68Xmean
self.tree.outside68Ymeans = outside68Ymean
self.tree.outside68Xstds = outside68Xstd
self.tree.outside68Ystds = outside68Ystd
self.tree.ele68ContEnergy = ele68ContEnergy
self.tree.ele68x2ContEnergy = ele68x2ContEnergy
self.tree.ele68x3ContEnergy = ele68x3ContEnergy
self.tree.ele68x4ContEnergy = ele68x4ContEnergy
self.tree.ele68x5ContEnergy = ele68x5ContEnergy
self.tree.photon68ContEnergy = photon68ContEnergy
self.tree.photon68x2ContEnergy = photon68x2ContEnergy
self.tree.photon68x3ContEnergy = photon68x3ContEnergy
self.tree.photon68x4ContEnergy = photon68x4ContEnergy
self.tree.photon68x5ContEnergy = photon68x5ContEnergy
self.tree.overlap68ContEnergy = overlap68ContEnergy
self.tree.overlap68x2ContEnergy = overlap68x2ContEnergy
self.tree.overlap68x3ContEnergy = overlap68x3ContEnergy
self.tree.overlap68x4ContEnergy = overlap68x4ContEnergy
self.tree.overlap68x5ContEnergy = overlap68x5ContEnergy
self.tree.outside68ContEnergy = outside68ContEnergy
self.tree.outside68x2ContEnergy = outside68x2ContEnergy
self.tree.outside68x3ContEnergy = outside68x3ContEnergy
self.tree.outside68x4ContEnergy = outside68x4ContEnergy
self.tree.outside68x5ContEnergy = outside68x5ContEnergy
self.tree.outside68ContNHits = outside68ContNHits
self.tree.outside68x2ContNHits = outside68x2ContNHits
self.tree.outside68x3ContNHits = outside68x3ContNHits
self.tree.outside68x4ContNHits = outside68x4ContNHits
self.tree.outside68x5ContNHits = outside68x5ContNHits
self.tree.outside68ContXmean = outside68ContXmean
self.tree.outside68x2ContXmean = outside68x2ContXmean
self.tree.outside68x3ContXmean = outside68x3ContXmean
self.tree.outside68x4ContXmean = outside68x4ContXmean
self.tree.outside68x5ContXmean = outside68x5ContXmean
self.tree.outside68ContYmean = outside68ContYmean
self.tree.outside68x2ContYmean = outside68x2ContYmean
self.tree.outside68x3ContYmean = outside68x3ContYmean
self.tree.outside68x4ContYmean = outside68x4ContYmean
self.tree.outside68x5ContYmean = outside68x5ContYmean
self.tree.outside68ContXstd = outside68ContXstd
self.tree.outside68x2ContXstd = outside68x2ContXstd
self.tree.outside68x3ContXstd = outside68x3ContXstd
self.tree.outside68x4ContXstd = outside68x4ContXstd
self.tree.outside68x5ContXstd = outside68x5ContXstd
self.tree.outside68ContYstd = outside68ContYstd
self.tree.outside68x2ContYstd = outside68x2ContYstd
self.tree.outside68x3ContYstd = outside68x3ContYstd
self.tree.outside68x4ContYstd = outside68x4ContYstd
self.tree.outside68x5ContYstd = outside68x5ContYstd
self.tree.outside68ContShowerRMS = outside68ContShowerRMS
self.tree.outside68x2ContShowerRMS = outside68x2ContShowerRMS
self.tree.outside68x3ContShowerRMS = outside68x3ContShowerRMS
self.tree.outside68x4ContShowerRMS = outside68x4ContShowerRMS
self.tree.outside68x5ContShowerRMS = outside68x5ContShowerRMS
# Fill the tree with values for this event
self.tree.fill(reset=True)
# Write and copy the output
def end(self):
self.tree.write()
self.tfile.close()
print 'cp %s %s' % (self.outfile, self.outdir)
os.system('cp %s %s' % (self.outfile, self.outdir))
class Event(TreeModel):
"""Event model definition"""
x = FloatCol()
y = FloatCol()
z = FloatCol()
i = IntCol()
# first create several example trees in separate files
fnames = ["test_%d.root" % i for i in xrange(10)]
for fname in fnames:
with root_open(fname, "recreate") as f:
tree = Tree("test", model=Event)
# fill the tree
for i in xrange(10000):
tree.x = gauss(0.5, 1.0)
tree.y = gauss(0.3, 2.0)
tree.z = gauss(13.0, 42.0)
tree.i = i
tree.fill()
tree.write()
"""
This section below takes the example trees and copies it while overwriting a
branch with new values.
"""
"""
print __doc__
import rootpy
rootpy.log.basic_config_colorized()
from random import gauss
from rootpy.io import root_open
from rootpy.tree import Tree, TreeChain
from rootpy.plotting import Hist
# Make two files, each with a Tree called "test"
print "Creating test tree in chaintest1.root"
f = root_open("chaintest1.root", "recreate")
tree = Tree("test")
branches = {"x": "F", "y": "F", "z": "F", "i": "I"}
tree.create_branches(branches)
for i in xrange(10000):
tree.x = gauss(0.5, 1.0)
tree.y = gauss(0.3, 2.0)
tree.z = gauss(13.0, 42.0)
tree.i = i
tree.fill()
# Make a histogram of x when y > 1
hist1 = Hist(100, -10, 10, name="hist1")
tree.Draw("x", "y > 1", hist=hist1)
hist1.SetDirectory(0) # memory resident
print "The first tree has %f entries where y > 1" % hist1.Integral()
if len(sys.argv) < 4:
print "USAGE: " + basename(sys.argv[0]) + " <time_win_mat.root> <decoded_file.root> <merged_file.root>"
exit(1)
time_win_filename = sys.argv[1]
decoded_data_filename = sys.argv[2]
merged_filename = sys.argv[3]
t_file_time_win = File(time_win_filename, "read")
begin_time_mat = t_file_time_win.get("begin_time_mat")
end_time_mat = t_file_time_win.get("end_time_mat")
max_count_mat = t_file_time_win.get("max_count_mat")
t_file_time_win.close()
t_file_merged_out = File(merged_filename, "recreate")
t_beam_event_tree = Tree("t_beam_event", "Beam Event Data")
t_beam_event_tree.create_branches({
'type': 'I',
'trig_accepted': 'B[25]',
'time_aligned': 'B[25]',
'pkt_count': 'I',
'lost_count': 'I',
'trigger_bit': 'B[1600]',
'trigger_n': 'I',
'multiplicity': 'I[25]',
'energy_adc': 'F[1600]',
'compress': 'I[25]',
'common_noise': 'F[25]',
'bar_beam': 'B[1600]' })
t_file_merged_out.cd()
def test_file(inname, outname, train_file, variables, testEvery):
log.info('START processing: %s', os.path.basename(inname))
bdt = get_training(train_file, variables)
with io.root_open(outname, 'w') as tout:
out_tree = Tree('tree')
out_tree.create_branches({
'flavour': 'F',
'vertexCategory': 'I',
'jetPt': 'F',
'jetEta': 'F',
'BDTG': 'F',
})
with io.root_open(inname) as tin:
in_tree = tin.tree
in_tree.SetBranchStatus('*', 0)
in_tree.SetBranchStatus('flavour', 1)
in_tree.SetBranchStatus('vertexCategory', 1)
in_tree.SetBranchStatus('jetPt', 1)
in_tree.SetBranchStatus('jetEta', 1)
for var in variables:
in_tree.SetBranchStatus(var, 1)
for idx, entry in enumerate(in_tree):
if testEvery != 1 and (idx % testEvery) == 0: continue
var_vals = [getattr(entry, i) for i in variables]
btd_out = bdt.predict_proba([var_vals])[0][1]
out_tree.flavour = entry.flavour
out_tree.vertexCategory = entry.vertexCategory
out_tree.jetPt = entry.jetPt
out_tree.jetEta = entry.jetEta
out_tree.BDTG = btd_out
out_tree.fill()
out_tree.write()
log.info('DONE processing: %s', os.path.basename(inname))
import signals
import optfilter
import rootpy
rootpy.log.basic_config_colorized()
from random import gauss
from rootpy.io import open as ropen
from rootpy.tree import Tree, TreeChain
from rootpy.plotting import Hist
import numpy as np
import random
outfilename = 'ampInAmpOutHeatArbitraryNumpyFilter.root'
rootFileOut = ropen(outfilename, 'recreate')
tree = Tree("t")
branches = {
'amp_in':'F', 'amp_out': 'F', 'amp_out_maxtime': 'F'
}
tree.create_branches(branches)
numEvents = 1000
_pulselength = 512
_pulseStartTime = _pulselength/2
_tukey_alpha = 0.3
amps = [-10, -20, -50, -100, -200, -400, -500, -800, -1000, -1400, -1800, -2000]
#amps = [-100]
#signal
simpulse = signals.HeatSignal(length = _pulselength)
scaleFactor = -1.0*simpulse().min()
};
""", ["Thingy"])
# alternatively you can ROOT.gSystem.Load() your library
# define the model
class Event(TreeModel):
event_number = IntCol()
thingy = ObjectCol(C.Thingy)
f = root_open("test.root", "recreate")
tree = Tree("test", model=Event)
# fill the tree
for i in xrange(20):
tree.event_number = i
tree.thingy.i = i
tree.thingy.x = gauss(.3, 2.)
tree.thingy.y = gauss(13., 42.)
tree.fill()
tree.write()
f.close()
# now to read the same tree
with root_open("test.root") as f:
=============================================
This example demonstrates how to use the TreeChain class, a more Python-friendly
TChain replacement.
"""
print __doc__
from random import gauss
from rootpy.io import open
from rootpy.tree import Tree, TreeChain
# Make two files, each with a Tree called "test"
print "Creating test tree in chaintest1.root"
f = open("chaintest1.root", "recreate")
tree = Tree("test")
tree.create_branches([('x', 'F'),
('y', 'F'),
('z', 'F'),
('i', 'I')])
for i in xrange(10000):
tree.x = gauss(.5, 1.)
tree.y = gauss(.3, 2.)
tree.z = gauss(13., 42.)
tree.i = i
tree.fill()
# Make a histogram of x when y > 1
hist1 = tree.Draw('x', 'y > 1', min=-10, max=10, bins=100).Clone()
hist1.SetName('hist1')
inputFile = root_open('data/bltuples/output_{0}_{1}.root'.format(
channel, period))
outputFile = root_open(
'data/step1_sfs/output_{0}_{1}.root'.format(channel, period),
'recreate')
if period == 2017 or period == 2018:
pt_weights = ur.open(
'data/from_ming-yan/pt_weights/pt_weights_{0}_{1}.root'.format(
channel, period))['hd']
for dataset in tqdm(datasets):
tree = inputFile['signal/tree_{0}'.format(dataset)]
hist = inputFile['TotalEvents_{0}'.format(dataset)]
tree.create_buffer()
outTree = Tree(dataset)
outTree.set_buffer(tree._buffer, create_branches=True)
outTree.create_branches({'mc_sf': 'F'})
outTree.create_branches({'pt_weight': 'F'})
outTree.create_branches({'useTMVA': 'F'})
outTree.create_branches({'corrPhotonMVA': 'F'})
outTree.create_branches({'llgMKinMY': 'F'})
outTree.create_branches({'isMingYanData': 'I'})
# compute the MC scale factor
n_gen = inputFile['TotalEvents_{0}'.format(dataset)].GetBinContent(
1)
n_neg = inputFile['TotalEvents_{0}'.format(dataset)].GetBinContent(
30)
sf = luminosity[period] * xsec[dataset] * br[dataset] / (
float(n_gen) - 2. * float(n_neg))
# test
from rootpy.io import root_open
from rootpy.tree import Tree, TreeModel
from rootpy.tree import FloatCol
from random import gauss
class Event(TreeModel):
variable = FloatCol(default=-1111)
with root_open('output.root', 'recreate'):
tree = Tree(name='mytree', model=Event)
for i in xrange(100):
tree.variable = gauss(0, 1)
tree.Fill()
tree.Write()
b_y = FloatCol()
b_z = FloatCol()
# a collection of particles
col_x = stl.vector("float")
col_y = stl.vector("float")
col_z = stl.vector("float")
col_n = IntCol()
# a TLorentzVector
p = LorentzVector
i = IntCol()
tree = Tree("test", model=Event)
# fill the tree
for i in xrange(10):
tree.a_x = gauss(0.5, 1.0)
tree.a_y = gauss(0.3, 2.0)
tree.a_z = gauss(13.0, 42.0)
tree.b_x = gauss(0.5, 1.0)
tree.b_y = gauss(0.3, 2.0)
tree.b_z = gauss(13.0, 42.0)
n = randint(1, 10)
for j in xrange(n):
tree.col_x.push_back(gauss(0.5, 1.0))
tree.col_y.push_back(gauss(0.3, 2.0))
def predict (**kwargs):
filenames = kwargs.setdefault ("filenames", ["training","test","check_correlation","check_agreement"])
variableOrder = kwargs.setdefault ("variable_order", ["id", "signal", "mass", "min_ANNmuon", "prediction"])
prediction_name = kwargs.setdefault ("prediction_name", "prediction")
regression_targets = kwargs.setdefault ("regression_targets", None)
execute_tests = kwargs.setdefault ("execute_tests",False)
# default values
variablesForFiles = {
"training" : ["prediction","id","signal","mass","min_ANNmuon"],
"test" : ["prediction","id"],
"check_agreement" : ["signal","weight","prediction"],
"check_correlation" : ["mass","prediction"]
}
variablesForFiles = kwargs.setdefault ("variablesForFiles", variablesForFiles)
createForFiles = {
"training" : ["csv"],
"test" : ["csv"],
"check_correlation" : ["csv"],
"check_agreement" : ["csv"]
}
input_variables = kwargs.setdefault ("variables", None)
input_spectators = kwargs.setdefault ("spectators", [])
method_name = kwargs.setdefault ("method_name", None)
weightfile_name = kwargs.setdefault ("weightfile_name", None)
print "------------ prediction ---------------"
for key,value in kwargs.iteritems ():
print key," = ",value
ROOT.TMVA.Tools.Instance ()
reader = ROOT.TMVA.Reader( "!Color:!Silent" )
variables = []
varIndex = {}
for idx, var_name in enumerate (input_variables):
tmp = array.array('f',[0])
variables.append (tmp)
if type(var_name) == str:
reader.AddVariable (var_name, tmp)
varIndex[var_name] = idx
else:
varcomposed = var_name[0] + ":=" + var_name[1]
reader.AddVariable (varcomposed, tmp)
spectators = []
specIndex = {}
for idx, spec_name in enumerate (input_spectators):
tmp = array.array('f',[0])
spectators.append (tmp)
if type(spec_name) == str:
reader.AddVariable (spec_name, tmp)
specIndex[spec_name] = idx
else:
speccomposed = spec_name[0] + ":=" + spec_name[1]
reader.AddVariable (speccomposed, tmp)
reader.BookMVA (method_name, weightfile_name)
regTag = "_r_"
denom = "_regression_"
doRegression = True
if regression_targets == None:
doRegression = False
regTag = "_p_"
denom = "_prediction_"
returnValues = {}
for currentFileName in filenames:
print "predict for file : ",currentFileName
doCSV = "csv" in createForFiles[currentFileName]
doROOT = ("root" in createForFiles[currentFileName]) or doCSV
if not doCSV and not doROOT:
continue
fileName = default_path + currentFileName + ".root"
# define variables
ID = array.array ('i',[0])
outputVariables = variablesForFiles[currentFileName]
prediction = array.array ('f',[0])
weight = array.array ('f',[0])
min_ANNmuon = array.array ('f',[0])
mass = array.array ('f',[0])
signal = array.array ('f',[0])
# --- open input file
f = rootpy.io.File.Open (fileName)
tree = f.Get("data");
for v in outputVariables:
if v != "prediction":
cmd = setbranch (v)
#print cmd
exec (cmd)
# create tree formulas
formulas = []
for idx,var in enumerate (input_variables):
if type(var) == str:
fml = None
tree.SetBranchAddress (var, variables[varIndex[var]])
else:
fml = ROOT.TTreeFormula (var[0], var[1], tree)
formulas.append (fml)
# ---- make ROOT file if requested
outTree = None
if doROOT:
rootFileName = currentFileName + denom + method_name + ".root"
outRootFile = rootpy.io.File (rootFileName, "RECREATE")
outTree = Tree ("data","data")
for var in variableOrder:
if var in variablesForFiles[currentFileName]:
altName = ""
if var in regression_targets:
altName = "t_"+var
if var == "prediction":
altName = prediction_name
if doRegression:
altName = regression_targets[0]
cmd = branch (var, altName)
exec (cmd)
curr = currentFileName + denom + "root"
returnValues[curr] = rootFileName
#
tmstmp = time.time ()
for ievt in xrange (tree.GetEntries()):
tree.GetEntry (ievt)
# predict
for idx,fml in enumerate (formulas):
if fml != None:
variables[idx][0] = fml.EvalInstance ()
# this will create a harmless warning
# https://root.cern.ch/phpBB3/viewtopic.php?f=14&t=14213
if doRegression:
targets = reader.EvaluateRegression (method_name)
prediction[0] = targets[0]
else:
prediction[0] = reader.EvaluateMVA (method_name)
prediction[0] = 1.0/(1.0+exp (-prediction[0]))
#print prediction
outTree.fill ()
if ievt%10000 == 0:
tmp = tmstmp
tmstmp = time.time ()
print ievt," t = %f"%(tmstmp-tmp)
if doROOT:
outRootFile.Write ()
# ---- prepare csv file
#csvfile = None
writer = None
if doCSV:
print "prepare csv"
csvFileName = currentFileName + denom + method_name + ".csv"
csvFile = open (csvFileName, 'w')
outTree.csv (",", stream=csvFile);
csvFile.close ()
curr = currentFileName + denom + "csv"
returnValues[curr] = csvFileName
f.Close ()
#if doCSV:
# csvfile.close ()
if doROOT:
outRootFile.Close ()
if execute_tests:
cmd = "os.system ('python tests.py %s %s %s')"%(returnValues["check_agreement_prediction_csv"],returnValues["check_correlation_prediction_csv"],returnValues["training_prediction_csv"])
exec (cmd)
return returnValues
class ppd_file_w:
def __init__(self):
self.__t_file_out = None
self.__t_tree_ppd = None
def open_file(self, filename):
self.__t_file_out = File(filename, "recreate")
self.__t_tree_ppd = Tree("t_ppd", "platform parameters data")
self.__t_tree_ppd.create_branches({
"pitch_angle" : "D" ,
"yaw_angle" : "D" ,
"roll_angle" : "D" ,
"pitch_angle_v" : "D" ,
"yaw_angle_v" : "D" ,
"roll_angle_v" : "D" ,
"orbit_agl_v" : "D" ,
"longitude" : "D" ,
"latitude" : "D" ,
"geocentric_d" : "D" ,
"ship_time_sec" : "D" ,
"utc_time_sec" : "D" ,
"utc_time_str" : "C[32]" ,
"flag_of_pos" : "I" ,
"wgs84_x" : "D" ,
"wgs84_y" : "D" ,
"wgs84_z" : "D" ,
"wgs84_x_v" : "D" ,
"wgs84_y_v" : "D" ,
"wgs84_z_v" : "D" ,
"det_z_lat" : "D" ,
"det_z_lon" : "D" ,
"det_z_ra" : "D" ,
"det_z_dec" : "D" ,
"det_x_lat" : "D" ,
"det_x_lon" : "D" ,
"det_x_ra" : "D" ,
"det_x_dec" : "D" ,
"earth_lat" : "D" ,
"earth_lon" : "D" ,
"earth_ra" : "D" ,
"earth_dec" : "D" ,
"sun_ra" : "D" ,
"sun_dec" : "D"
})
def fill_data(self, ppd_obj):
self.__t_tree_ppd.pitch_angle = ppd_obj.pitch_angle
self.__t_tree_ppd.yaw_angle = ppd_obj.yaw_angle
self.__t_tree_ppd.roll_angle = ppd_obj.roll_angle
self.__t_tree_ppd.pitch_angle_v = ppd_obj.pitch_angle_v
self.__t_tree_ppd.yaw_angle_v = ppd_obj.yaw_angle_v
self.__t_tree_ppd.roll_angle_v = ppd_obj.roll_angle_v
self.__t_tree_ppd.orbit_agl_v = ppd_obj.orbit_agl_v
self.__t_tree_ppd.longitude = ppd_obj.longitude
self.__t_tree_ppd.latitude = ppd_obj.latitude
self.__t_tree_ppd.geocentric_d = ppd_obj.geocentric_d
self.__t_tree_ppd.ship_time_sec = ppd_obj.ship_time_sec
self.__t_tree_ppd.utc_time_sec = ppd_obj.utc_time_sec
self.__t_tree_ppd.utc_time_str = str(ppd_obj.utc_time_str)
self.__t_tree_ppd.flag_of_pos = ppd_obj.flag_of_pos
self.__t_tree_ppd.wgs84_x = ppd_obj.wgs84_x
self.__t_tree_ppd.wgs84_y = ppd_obj.wgs84_y
self.__t_tree_ppd.wgs84_z = ppd_obj.wgs84_z
self.__t_tree_ppd.wgs84_x_v = ppd_obj.wgs84_x_v
self.__t_tree_ppd.wgs84_y_v = ppd_obj.wgs84_y_v
self.__t_tree_ppd.wgs84_z_v = ppd_obj.wgs84_z_v
self.__t_tree_ppd.det_z_lat = ppd_obj.det_z_lat
self.__t_tree_ppd.det_z_lon = ppd_obj.det_z_lon
self.__t_tree_ppd.det_z_ra = ppd_obj.det_z_ra
self.__t_tree_ppd.det_z_dec = ppd_obj.det_z_dec
self.__t_tree_ppd.det_x_lat = ppd_obj.det_x_lat
self.__t_tree_ppd.det_x_lon = ppd_obj.det_x_lon
self.__t_tree_ppd.det_x_ra = ppd_obj.det_x_ra
self.__t_tree_ppd.det_x_dec = ppd_obj.det_x_dec
self.__t_tree_ppd.earth_lat = ppd_obj.earth_lat
self.__t_tree_ppd.earth_lon = ppd_obj.earth_lon
self.__t_tree_ppd.earth_ra = ppd_obj.earth_ra
self.__t_tree_ppd.earth_dec = ppd_obj.earth_dec
self.__t_tree_ppd.sun_ra = ppd_obj.sun_ra
self.__t_tree_ppd.sun_dec = ppd_obj.sun_dec
self.__t_tree_ppd.fill()
def write_tree(self):
self.__t_file_out.cd()
self.__t_tree_ppd.write()
def write_meta(self, key, value):
self.__t_file_out.cd()
ROOT.TNamed(key, value).Write()
def close_file(self):
self.__t_file_out.close()
self.__t_file_out = None
self.__t_tree_ppd = None
def manufacturePredictor (**kwargs):
filenames = kwargs.setdefault ("filenames", ["training","test","check_correlation","check_agreement"])
variableOrder = kwargs.setdefault ("variable_order", ["id", "signal", "mass", "min_ANNmuon", "prediction"])
prediction_name = kwargs.setdefault ("prediction_name", "prediction")
prediction_formula = kwargs.setdefault ("prediction_formula", "prediction_formula")
prediction_formula = ROOT.TFormula ('pFml',prediction_formula)
#print prediction_formula
#return
execute_tests = kwargs.setdefault ("execute_tests",False)
# default values
variablesForFiles = {
"training" : ["prediction","id","signal","mass","min_ANNmuon"],
"test" : ["prediction","id"],
"check_agreement" : ["signal","weight","prediction"],
"check_correlation" : ["mass","prediction"]
}
variablesForFiles = kwargs.setdefault ("variablesForFiles", variablesForFiles)
createForFiles = {
"training" : ["csv"],
"test" : ["csv"],
"check_correlation" : ["csv"],
"check_agreement" : ["csv"]
}
input_variables = kwargs.setdefault ("variables", None)
input_spectators = kwargs.setdefault ("spectators", [])
method_names = kwargs.setdefault ("method_names", None)
weightfile_names = kwargs.setdefault ("weightfile_names", None)
print "------------ prediction ---------------"
for key,value in kwargs.iteritems ():
print key," = ",value
ROOT.TMVA.Tools.Instance ()
readers = []
for weightfile in weightfile_names:
readers.append (ROOT.TMVA.Reader( "!Color:!Silent" ))
variables = []
varIndex = {}
spectators = []
specIndex = {}
for idxReader, reader in enumerate (readers):
for idx, var_name in enumerate (input_variables):
tmpVarName = ""
if type(var_name) == str:
tmpVarName = var_name
else:
varcomposed = var_name[0] + ":=" + var_name[1]
tmpVarName = varcomposed
tmp = None
if tmpVarName in varIndex:
tmp = variables[varIndex[tmpVarName]]
else:
tmp = array.array('f',[0])
variables.append (tmp)
varIndex[tmpVarName] = len(variables)-1
reader.AddVariable (tmpVarName, tmp)
for idx, var_name in enumerate (input_spectators):
tmpVarName = ""
if type(var_name) == str:
tmpVarName = var_name
else:
varcomposed = var_name[0] + ":=" + var_name[1]
tmpVarName = varcomposed
tmp = None
if tmpVarName in varIndex:
tmp = variables[varIndex[tmpVarName]]
else:
tmp = array.array('f',[0])
variables.append (tmp)
varIndex[tmpVarName] = len(variables)-1
reader.AddSpectator (tmpVarName, tmp)
print "reader: book mva: ",method_names[idxReader]," from weightfile ",weightfile_names[idxReader]
reader.BookMVA (method_names[idxReader], weightfile_names[idxReader])
returnValues = {}
for currentFileName in filenames:
print "predict for file : ",currentFileName
doCSV = "csv" in createForFiles[currentFileName]
doROOT = ("root" in createForFiles[currentFileName]) or doCSV
if not doCSV and not doROOT:
continue
fileName = default_path + currentFileName + ".root"
# define variables
ID = array.array ('i',[0])
outputVariables = variablesForFiles[currentFileName]
prediction = array.array ('f',[0])
weight = array.array ('f',[0])
min_ANNmuon = array.array ('f',[0])
mass = array.array ('f',[0])
signal = array.array ('f',[0])
# --- open input file
f = rootpy.io.File.Open (fileName)
tree = f.Get("data");
for v in outputVariables:
if v != "prediction":
cmd = setbranch (v)
#print cmd
exec (cmd)
# create tree formulas
formulas = []
for idx,var in enumerate (input_variables):
if type(var) == str:
fml = None
tree.SetBranchAddress (var, variables[varIndex[var]])
else:
fml = ROOT.TTreeFormula (var[0], var[1], tree)
formulas.append (fml)
# ---- make ROOT file if requested
outTree = None
manu_name = ""
for m in method_names:
manu_name += m
if doROOT:
rootFileName = currentFileName + "_p_" + manu_name + ".root"
print "prepare root file : ",rootFileName
outRootFile = rootpy.io.File (rootFileName, "RECREATE")
outTree = Tree ("data","data")
for var in variableOrder:
if var in variablesForFiles[currentFileName]:
altName = ""
if var == "prediction":
altName = prediction_name
if doRegression:
altName = regression_targets[0]
cmd = branch (var, altName)
exec (cmd)
curr = currentFileName + denom + "root"
returnValues[curr] = rootFileName
#
tmstmp = time.time ()
for ievt in xrange (tree.GetEntries()):
tree.GetEntry (ievt)
# predict
for idx,fml in enumerate (formulas):
if fml != None:
variables[idx][0] = fml.EvalInstance ()
# this will create a harmless warning
# https://root.cern.ch/phpBB3/viewtopic.php?f=14&t=14213
prediction_bases = [] #array.array ('f',[0])
for idx, meth in enumerate (method_names):
p = readers[idx].EvaluateMVA (meth)
prediction_bases.append (p)
#print "idx ",idx," meth ",meth," p ",p," pred_bases ",prediction_bases
#for idxP, p in enumerate (prediction_bases):
# prediction_formula.SetParameter (idxP, p)
prediction[0] = prediction_formula.EvalPar (numpy.array(prediction_bases))
prediction[0] = 1.0/(1.0+exp (-prediction[0]))
#print prediction_bases," ",prediction
#print prediction
outTree.fill ()
if ievt%10000 == 0:
tmp = tmstmp
tmstmp = time.time ()
print ievt," t = %f"%(tmstmp-tmp)
if doROOT:
outRootFile.Write ()
# ---- prepare csv file
#csvfile = None
writer = None
if doCSV:
csvFileName = currentFileName + "_p_" + manu_name + ".csv"
print "prepare csv : ",csvFileName
csvFile = open (csvFileName, 'w')
outTree.csv (",", stream=csvFile);
csvFile.close ()
curr = currentFileName + "_prediction_csv"
returnValues[curr] = csvFileName
f.Close ()
#if doCSV:
# csvfile.close ()
if doROOT:
outRootFile.Close ()
if execute_tests:
cmd = "os.system ('python tests.py %s %s %s')"%(returnValues["check_agreement_prediction_csv"],returnValues["check_correlation_prediction_csv"],returnValues["training_prediction_csv"])
print cmd
exec (cmd)
return returnValues
"""
This first section of code only creates an example tree.
"""
# define the model
class Event(TreeModel):
x = FloatCol()
y = FloatCol()
z = FloatCol()
i = IntCol()
# first create a tree "test" in a file "test.root"
f = ropen("test.root", "recreate")
tree = Tree("test", model=Event)
# fill the tree
for i in xrange(10000):
tree.x = gauss(.5, 1.)
tree.y = gauss(.3, 2.)
tree.z = gauss(13., 42.)
tree.i = i
tree.fill()
tree.write()
"""
This section below takes the example tree and copies it while overwriting a
branch with new values.
"""
#!/usr/bin/env python
"""
=====================
A simple Tree example
=====================
This example demonstrates how to create a simple tree.
"""
print __doc__
from rootpy.tree import Tree
from rootpy.io import root_open
from random import gauss
f = root_open("test.root", "recreate")
tree = Tree("test")
tree.create_branches(
{'x': 'F',
'y': 'F',
'z': 'F',
'i': 'I'})
for i in xrange(10000):
tree.x = gauss(.5, 1.)
tree.y = gauss(.3, 2.)
tree.z = gauss(13., 42.)
tree.i = i
tree.fill()
tree.write()
f.close()
def work(self):
# get argument values
local = self.args.local
syst_terms = self.args.syst_terms
datatype = self.metadata.datatype
year = self.metadata.year
verbose = self.args.student_verbose
very_verbose = self.args.student_very_verbose
redo_selection = self.args.redo_selection
nominal_values = self.args.nominal_values
# get the dataset name
dsname = os.getenv('INPUT_DATASET_NAME', None)
if dsname is None:
# attempt to guess dsname from dirname
if self.files:
dsname = os.path.basename(os.path.dirname(self.files[0]))
# is this a signal sample?
# if so we will also keep some truth information in the output below
is_signal = datatype == datasets.MC and (
'_VBFH' in dsname or '_ggH' in dsname or '_ZH' in dsname
or '_WH' in dsname or '_ttH' in dsname)
log.info("DATASET: {0}".format(dsname))
log.info("IS SIGNAL: {0}".format(is_signal))
# is this an inclusive signal sample for overlap studies?
is_inclusive_signal = is_signal and '_inclusive' in dsname
# is this a BCH-fixed sample? (temporary)
is_bch_sample = 'r5470_r4540_p1344' in dsname
if is_bch_sample:
log.warning("this is a BCH-fixed r5470 sample")
# onfilechange will contain a list of functions to be called as the
# chain rolls over to each new file
onfilechange = []
count_funcs = {}
if datatype != datasets.DATA:
# count the weighted number of events
if local:
def mc_weight_count(event):
return event.hh_mc_weight
else:
def mc_weight_count(event):
return event.mc_event_weight
count_funcs = {
'mc_weight': mc_weight_count,
}
# three instances of the pileup reweighting tool are created to write
# out the nominal, high and low pileup weights
pileup_tool = None
pileup_tool_high = None
pileup_tool_low = None
if local:
# local means running on the skims, the output of this script
# running on the grid
if datatype == datasets.DATA:
# merge the GRL fragments
merged_grl = goodruns.GRL()
def update_grl(student, grl, name, file, tree):
grl |= str(
file.Get('Lumi/%s' %
student.metadata.treename).GetString())
onfilechange.append((update_grl, (
self,
merged_grl,
)))
if datatype == datasets.DATA:
merged_cutflow = Hist(1, 0, 1, name='cutflow', type='D')
else:
merged_cutflow = Hist(2, 0, 2, name='cutflow', type='D')
def update_cutflow(student, cutflow, name, file, tree):
# record a cut-flow
year = student.metadata.year
datatype = student.metadata.datatype
cutflow[1].value += file.cutflow_event[1].value
if datatype != datasets.DATA:
cutflow[2].value += file.cutflow_event_mc_weight[1].value
onfilechange.append((update_cutflow, (
self,
merged_cutflow,
)))
else:
# get pileup reweighting tool
pileup_tool = get_pileup_reweighting_tool(year=year,
use_defaults=True)
pileup_tool_high = get_pileup_reweighting_tool(year=year,
use_defaults=True,
systematic='high')
pileup_tool_low = get_pileup_reweighting_tool(year=year,
use_defaults=True,
systematic='low')
if datatype not in (datasets.EMBED, datasets.MCEMBED):
# merge TrigConfTrees
metadirname = '%sMeta' % self.metadata.treename
trigconfchain = ROOT.TChain('%s/TrigConfTree' % metadirname)
map(trigconfchain.Add, self.files)
metadir = self.output.mkdir(metadirname)
metadir.cd()
trigconfchain.Merge(self.output, -1, 'fast keep')
self.output.cd()
if datatype == datasets.DATA:
# merge GRL XML strings
merged_grl = goodruns.GRL()
for fname in self.files:
with root_open(fname) as f:
for key in f.Lumi.keys():
merged_grl |= goodruns.GRL(str(
key.ReadObj().GetString()),
from_string=True)
lumi_dir = self.output.mkdir('Lumi')
lumi_dir.cd()
xml_string = ROOT.TObjString(merged_grl.str())
xml_string.Write(self.metadata.treename)
self.output.cd()
self.output.cd()
# create the output tree
model = get_model(datatype,
dsname,
prefix=None if local else 'hh_',
is_inclusive_signal=is_inclusive_signal)
log.info("Output Model:\n\n{0}\n\n".format(model))
outtree = Tree(name=self.metadata.treename, model=model)
if local:
tree = outtree
else:
tree = outtree.define_object(name='tree', prefix='hh_')
tree.define_object(name='tau', prefix='tau_')
tree.define_object(name='tau1', prefix='tau1_')
tree.define_object(name='tau2', prefix='tau2_')
tree.define_object(name='truetau1', prefix='truetau1_')
tree.define_object(name='truetau2', prefix='truetau2_')
tree.define_object(name='jet1', prefix='jet1_')
tree.define_object(name='jet2', prefix='jet2_')
tree.define_object(name='jet3', prefix='jet3_')
mmc_objects = [
tree.define_object(name='mmc0', prefix='mmc0_'),
tree.define_object(name='mmc1', prefix='mmc1_'),
tree.define_object(name='mmc2', prefix='mmc2_'),
]
for mmc_obj in mmc_objects:
mmc_obj.define_object(name='resonance', prefix='resonance_')
trigger_emulation = TauTriggerEmulation(year=year,
passthrough=local
or datatype != datasets.MC
or year > 2011,
count_funcs=count_funcs)
if not trigger_emulation.passthrough:
onfilechange.append((update_trigger_trees, (
self,
trigger_emulation,
)))
trigger_config = None
if datatype not in (datasets.EMBED, datasets.MCEMBED):
# trigger config tool to read trigger info in the ntuples
trigger_config = get_trigger_config()
# update the trigger config maps on every file change
onfilechange.append((update_trigger_config, (trigger_config, )))
# define the list of event filters
if local and syst_terms is None and not redo_selection:
event_filters = None
else:
tau_ntrack_recounted_use_ntup = False
if year > 2011:
# peek at first tree to determine if the extended number of
# tracks is already stored
with root_open(self.files[0]) as test_file:
test_tree = test_file.Get(self.metadata.treename)
tau_ntrack_recounted_use_ntup = ('tau_out_track_n_extended'
in test_tree)
event_filters = EventFilterList([
averageIntPerXingPatch(
passthrough=(local or year < 2012
or datatype != datasets.MC),
count_funcs=count_funcs),
PileupTemplates(
year=year,
passthrough=(local or is_bch_sample or datatype
not in (datasets.MC, datasets.MCEMBED)),
count_funcs=count_funcs),
RandomSeed(datatype=datatype, count_funcs=count_funcs),
RandomRunNumber(tree=tree,
datatype=datatype,
pileup_tool=pileup_tool,
passthrough=local,
count_funcs=count_funcs),
PileupReweight(
year=year,
tool=pileup_tool,
tool_high=pileup_tool_high,
tool_low=pileup_tool_low,
tree=tree,
passthrough=(local
or (datatype
not in (datasets.MC, datasets.MCEMBED))),
count_funcs=count_funcs),
TruthMatching(passthrough=datatype == datasets.DATA,
count_funcs=count_funcs),
JetIsPileup(
passthrough=(local or year < 2012 or datatype
not in (datasets.MC, datasets.MCEMBED)),
count_funcs=count_funcs),
HiggsPT(year=year,
tree=tree,
passthrough=not is_signal or local,
count_funcs=count_funcs),
MCWeight(datatype=datatype,
tree=tree,
passthrough=local or datatype == datasets.DATA,
count_funcs=count_funcs),
ClassifyInclusiveHiggsSample(
tree=tree,
passthrough=not is_inclusive_signal,
count_funcs=count_funcs),
])
# set the event filters
self.filters['event'] = event_filters
# peek at first tree to determine which branches to exclude
with root_open(self.files[0]) as test_file:
test_tree = test_file.Get(self.metadata.treename)
ignore_branches = test_tree.glob(hhbranches.REMOVE,
exclude=hhbranches.KEEP)
ignore_branches_output = test_tree.glob(
hhbranches.REMOVE_OUTPUT, exclude=hhbranches.KEEP_OUTPUT)
# initialize the TreeChain of all input files
chain = TreeChain(self.metadata.treename,
files=self.files,
ignore_branches=ignore_branches,
events=self.events,
onfilechange=onfilechange,
filters=event_filters,
cache=True,
cache_size=50000000,
learn_entries=100)
if local:
copied = [
'EventNumber',
]
hh_buffer = TreeBuffer()
buffer = TreeBuffer()
for name, value in chain._buffer.items():
if name.startswith('hh_'):
hh_buffer[name[3:]] = value
elif name in copied:
buffer[name] = value
outtree.set_buffer(hh_buffer, create_branches=False, visible=True)
outtree.set_buffer(buffer, create_branches=True, visible=False)
else:
# additional decorations on existing objects
if year > 2011 and datatype in (datasets.MC, datasets.MCEMBED):
class Decorations(TreeModel):
jet_ispileup = stl.vector('bool')
chain.set_buffer(Decorations(), create_branches=True)
# include the branches in the input chain in the output tree
# set branches to be removed in ignore_branches
outtree.set_buffer(chain._buffer,
ignore_branches=ignore_branches +
ignore_branches_output,
create_branches=True,
ignore_duplicates=True,
transfer_objects=True,
visible=False)
# define tree objects
define_objects(chain, year)
# create the MMC
mmc = mass.MMC(year=year)
# report which packages have been loaded
externaltools.report()
self.output.cd()
# The main event loop
# the event filters above are automatically run for each event and only
# the surviving events are looped on
for event in chain:
if local and syst_terms is None and not redo_selection:
outtree.Fill()
continue
# sort taus and jets in decreasing order by pT
event.taus.sort(key=lambda tau: tau.pt, reverse=True)
event.jets.sort(key=lambda jet: jet.pt, reverse=True)
# tau1 is the leading tau
# tau2 is the subleading tau
taus = list(event.taus)
if len(taus) >= 2:
tau1, tau2 = taus[0], taus[1]
jets = list(event.jets)
jet1, jet2, jet3 = None, None, None
beta = None
if len(jets) >= 2:
jet1, jet2 = jets[:2]
# determine boost of system
# determine jet CoM frame
beta = (jet1.fourvect + jet2.fourvect).BoostVector()
tree.jet_beta.copy_from(beta)
jet1.fourvect_boosted.copy_from(jet1.fourvect)
jet2.fourvect_boosted.copy_from(jet2.fourvect)
jet1.fourvect_boosted.Boost(beta * -1)
jet2.fourvect_boosted.Boost(beta * -1)
tau1.fourvect_boosted.copy_from(tau1.fourvect)
tau2.fourvect_boosted.copy_from(tau2.fourvect)
tau1.fourvect_boosted.Boost(beta * -1)
tau2.fourvect_boosted.Boost(beta * -1)
tau1.min_dr_jet = min(tau1.fourvect.DeltaR(jet1.fourvect),
tau1.fourvect.DeltaR(jet2.fourvect))
tau2.min_dr_jet = min(tau2.fourvect.DeltaR(jet1.fourvect),
tau2.fourvect.DeltaR(jet2.fourvect))
# sphericity, aplanarity = eventshapes.sphericity_aplanarity(
# [tau1.fourvect,
# tau2.fourvect,
# jet1.fourvect,
# jet2.fourvect])
# sphericity
# tree.sphericity = sphericity
# aplanarity
# tree.aplanarity = aplanarity
# sphericity_boosted, aplanarity_boosted = eventshapes.sphericity_aplanarity(
# [tau1.fourvect_boosted,
# tau2.fourvect_boosted,
# jet1.fourvect_boosted,
# jet2.fourvect_boosted])
# sphericity
# tree.sphericity_boosted = sphericity_boosted
# aplanarity
# tree.aplanarity_boosted = aplanarity_boosted
# tau centrality (degree to which they are between the two jets)
tau1.centrality = eventshapes.eta_centrality(
tau1.fourvect.Eta(), jet1.fourvect.Eta(),
jet2.fourvect.Eta())
tau2.centrality = eventshapes.eta_centrality(
tau2.fourvect.Eta(), jet1.fourvect.Eta(),
jet2.fourvect.Eta())
# boosted tau centrality
tau1.centrality_boosted = eventshapes.eta_centrality(
tau1.fourvect_boosted.Eta(),
jet1.fourvect_boosted.Eta(),
jet2.fourvect_boosted.Eta())
tau2.centrality_boosted = eventshapes.eta_centrality(
tau2.fourvect_boosted.Eta(),
jet1.fourvect_boosted.Eta(),
jet2.fourvect_boosted.Eta())
# 3rd leading jet
if len(jets) >= 3:
jet3 = jets[2]
jet3.fourvect_boosted.copy_from(jet3.fourvect)
jet3.fourvect_boosted.Boost(beta * -1)
elif len(jets) == 1:
jet1 = jets[0]
tau1.min_dr_jet = tau1.fourvect.DeltaR(jet1.fourvect)
tau2.min_dr_jet = tau2.fourvect.DeltaR(jet1.fourvect)
# sphericity, aplanarity = eventshapes.sphericity_aplanarity(
# [tau1.fourvect,
# tau2.fourvect,
# jet1.fourvect])
# sphericity
# tree.sphericity = sphericity
# aplanarity
#tree.aplanarity = aplanarity
RecoJetBlock.set(tree, jet1, jet2, jet3, local=local)
# mass of ditau + leading jet system
if jet1 is not None:
tree.mass_tau1_tau2_jet1 = (tau1.fourvect + tau2.fourvect +
jet1.fourvect).M()
# full sphericity and aplanarity
# sphericity_full, aplanarity_full = eventshapes.sphericity_aplanarity(
# [tau1.fourvect, tau2.fourvect] + [jet.fourvect for jet in jets])
# tree.sphericity_full = sphericity_full
# tree.aplanarity_full = aplanarity_full
# ####################################
# number of tracks from PV minus taus
# ####################################
ntrack_pv = 0
ntrack_nontau_pv = 0
for vxp in event.vertices:
# primary vertex
if vxp.type == 1:
ntrack_pv = vxp.nTracks
ntrack_nontau_pv = ntrack_pv - tau1.numTrack - tau2.numTrack
break
tree.ntrack_pv = ntrack_pv
tree.ntrack_nontau_pv = ntrack_nontau_pv
# ########################
# MET variables
# ########################
METx = event.MET.etx
METy = event.MET.ety
MET = event.MET.et
MET_vect = Vector2(METx, METy)
MET_4vect = LorentzVector()
MET_4vect.SetPxPyPzE(METx, METy, 0., MET)
MET_4vect_boosted = LorentzVector()
MET_4vect_boosted.copy_from(MET_4vect)
if beta is not None:
MET_4vect_boosted.Boost(beta * -1)
tree.MET_et = MET
tree.MET_etx = METx
tree.MET_ety = METy
tree.MET_phi = event.MET.phi
dPhi_tau1_tau2 = abs(tau1.fourvect.DeltaPhi(tau2.fourvect))
dPhi_tau1_MET = abs(tau1.fourvect.DeltaPhi(MET_4vect))
dPhi_tau2_MET = abs(tau2.fourvect.DeltaPhi(MET_4vect))
tree.dPhi_tau1_tau2 = dPhi_tau1_tau2
tree.dPhi_tau1_MET = dPhi_tau1_MET
tree.dPhi_tau2_MET = dPhi_tau2_MET
tree.dPhi_min_tau_MET = min(dPhi_tau1_MET, dPhi_tau2_MET)
tree.MET_bisecting = is_MET_bisecting(dPhi_tau1_tau2,
dPhi_tau1_MET,
dPhi_tau2_MET)
sumET = event.MET.sumet
tree.MET_sumet = sumET
if sumET != 0:
tree.MET_sig = (
(2. * MET / GeV) /
(utils.sign(sumET) * sqrt(abs(sumET / GeV))))
else:
tree.MET_sig = -1.
tree.MET_centrality = eventshapes.phi_centrality(
tau1.fourvect, tau2.fourvect, MET_vect)
tree.MET_centrality_boosted = eventshapes.phi_centrality(
tau1.fourvect_boosted, tau2.fourvect_boosted,
MET_4vect_boosted)
tree.number_of_good_vertices = len(event.vertices)
# #########################
# Jet and sum pt variables
# #########################
tree.numJets = len(event.jets)
# sum pT with only the two leading jets
tree.sum_pt = sum([tau1.pt, tau2.pt] +
[jet.pt for jet in jets[:2]])
# sum pT with all selected jets
tree.sum_pt_full = sum([tau1.pt, tau2.pt] +
[jet.pt for jet in jets])
# vector sum pT with two leading jets and MET
tree.vector_sum_pt = sum([tau1.fourvect, tau2.fourvect] +
[jet.fourvect for jet in jets[:2]] +
[MET_4vect]).Pt()
# vector sum pT with all selected jets and MET
tree.vector_sum_pt_full = sum([tau1.fourvect, tau2.fourvect] +
[jet.fourvect for jet in jets] +
[MET_4vect]).Pt()
# resonance pT
tree.resonance_pt = sum(
[tau1.fourvect, tau2.fourvect, MET_4vect]).Pt()
# ############################
# tau <-> vertex association
# ############################
tree.tau_same_vertex = (tau1.privtx_x == tau2.privtx_x
and tau1.privtx_y == tau2.privtx_y
and tau1.privtx_z == tau2.privtx_z)
tau1.vertex_prob = ROOT.TMath.Prob(tau1.privtx_chiSquared,
int(tau1.privtx_numberDoF))
tau2.vertex_prob = ROOT.TMath.Prob(tau2.privtx_chiSquared,
int(tau2.privtx_numberDoF))
# #########################
# MMC Mass
# #########################
mmc_result = mmc.mass(tau1,
tau2,
METx,
METy,
sumET,
njets=len(event.jets))
for mmc_method, mmc_object in enumerate(mmc_objects):
mmc_mass, mmc_resonance, mmc_met = mmc_result[mmc_method]
if verbose:
log.info("MMC (method %d): %f" %
(mmc_method, mmc_mass))
mmc_object.mass = mmc_mass
mmc_object.MET_et = mmc_met.Mod()
mmc_object.MET_etx = mmc_met.X()
mmc_object.MET_ety = mmc_met.Y()
mmc_object.MET_phi = math.pi - mmc_met.Phi()
if mmc_mass > 0:
FourMomentum.set(mmc_object.resonance, mmc_resonance)
# ###########################
# collinear and visible mass
# ###########################
vis_mass, collin_mass, tau1_x, tau2_x = mass.collinearmass(
tau1, tau2, METx, METy)
tree.mass_vis_tau1_tau2 = vis_mass
tree.mass_collinear_tau1_tau2 = collin_mass
tau1.collinear_momentum_fraction = tau1_x
tau2.collinear_momentum_fraction = tau2_x
###########################
# Match jets to VBF partons
###########################
#if datatype == datasets.MC and 'VBF' in dsname and year == 2011:
# # get partons (already sorted by eta in hepmc) FIXME!!!
# parton1, parton2 = hepmc.get_VBF_partons(event)
# tree.mass_true_quark1_quark2 = (parton1.fourvect + parton2.fourvect).M()
# # order here needs to be revised since jets are no longer
# # sorted by eta but instead by pT
# PartonBlock.set(tree, parton1, parton2)
# if len(jets) >= 2:
# jet1, jet2 = jets[:2]
# for i, jet in zip((1, 2), (jet1, jet2)):
# for parton in (parton1, parton2):
# if utils.dR(jet.eta, jet.phi, parton.eta, parton.phi) < .8:
# setattr(tree, 'jet%i_matched' % i, True)
# Fill the tau block
# This must come after the RecoJetBlock is filled since
# that sets the jet_beta for boosting the taus
RecoTauBlock.set(event,
tree,
datatype,
tau1,
tau2,
local=local)
if datatype != datasets.DATA:
TrueTauBlock.set(tree, tau1, tau2)
# fill the output tree
outtree.Fill(reset=True)
externaltools.report()
# flush any baskets remaining in memory to disk
self.output.cd()
outtree.FlushBaskets()
outtree.Write()
if local:
if datatype == datasets.DATA:
xml_string = ROOT.TObjString(merged_grl.str())
xml_string.Write('lumi')
merged_cutflow.Write()
from rootpy.types import FloatCol, IntCol
from random import gauss
f = open("test.root", "recreate")
# define the model
class Event(TreeModel):
x = FloatCol()
y = FloatCol()
z = FloatCol()
i = IntCol()
tree = Tree("test", model=Event)
# fill the tree
for i in xrange(100000):
tree.x = gauss(.5, 1.)
tree.y = gauss(.3, 2.)
tree.z = gauss(13., 42.)
tree.i = i
tree.fill()
tree.write()
# convert tree into a numpy record array
from root_numpy import tree2rec
array = tree2rec(tree)
print array
print array.x
