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 __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()
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