def main(): INPUT_FILE_INCLUDING_PATH = "../../../root_generator/tree/root2_tree.root" FileNameList = read_file_name(INPUT_FILE_INCLUDING_PATH) BranchListAll = get_branch_list_all(FileNameList[2]) BranchListEachTree = get_branch_list_each_tree(FileNameList[2]) DicNumpyArray_branch = {} for numpyarray in BranchListAll: a = numpy.array([0], 'd') DicNumpyArray_branch[numpyarray] = a DicNumpyArray_branch = collections.OrderedDict( sorted(DicNumpyArray_branch.items()) ) # !!!the input times are ordered!!! print(DicNumpyArray_branch) DicNumpyArray_branch_w = {} for numpyarray_w in BranchListAll: a_w = numpy.array([0], 'd') DicNumpyArray_branch_w[numpyarray_w] = a_w DicNumpyArray_branch_w = collections.OrderedDict( sorted(DicNumpyArray_branch_w.items())) print(DicNumpyArray_branch_w) WCuts = WhetherAddCut(BranchListEachTree) if (WCuts == None): print("\n") print( "*********************************************************************************************" ) print("!!!! Please check input CUT !!!!! ") print(" !!!! No Output File !!!! ") print("\n") for i in range(len(BranchListEachTree)): print("compenets below :") print("The Tree name is : ", BranchListEachTree.keys()[i]) for j in range(len(BranchListEachTree.values()[i])): print(" it contains : ", BranchListEachTree.values()[i][j]) print("\n") print("!!!! Please check input CUT !!!!! ") print(" !!!! No Output File !!!! ") print( "*********************************************************************************************" ) else: print("\n") print("...Cut added and regenerating tree root file...") gBenchmark.Start("Regerating tree root") f = TFile(FileNameList[2], "READ") dirlist = f.GetListOfKeys() ITER = dirlist.MakeIterator() key = ITER.Next() outfileName = FileNameList[0] + "_cut_tree.root" outfile = TFile(outfileName, "RECREATE") ijk = 0 while key: tree = key.ReadObj() tree_f = TTree(tree.GetName() + "_f", tree.GetName() + "_f") ENTRY = tree.GetEntries() #print(ENTRY) for i in range(len(DicNumpyArray_branch)): tree.SetBranchAddress(DicNumpyArray_branch.keys()[i], DicNumpyArray_branch.values()[i]) tree_f.Branch(DicNumpyArray_branch_w.keys()[i], DicNumpyArray_branch_w.values()[i], DicNumpyArray_branch_w.keys()[i] + "/D") print("for tree", tree.GetName()) for j in range(ENTRY): tree.GetEntry(j) if (j % 5000 == 0): print("now looping", j, "th Events, total of ", ENTRY, "events") for k in range(len(DicNumpyArray_branch)): DicNumpyArray_branch_w.values( )[k][0] = DicNumpyArray_branch.values()[k][0] if ( #FIXME#FIXME#FIXME#FIXME#FIXME#FIXME#FIXME#FIXME#FIXME#FIXME ( DicNumpyArray_branch.values()[0][0] < 0.5 ) # [i][0] means "i"th branch of the tree, [0] don't change #FIXME#FIXME#FIXME#FIXME & ( DicNumpyArray_branch.values()[1][0] < 0.5 ) #FIXME#FIXME#FIXME#FIXME#FIXME#FIXME#FIXME#FIXME#FIXME#FIXME#FIXME ): ijk = ijk + 1 tree_f.Fill() else: continue print(" Event left after filtering : ", ijk, "!!!!!!!!!") print("\n") ijk = 0 if (tree_f.GetEntries() == 0): print("!!!!!!! ", tree_f.GetName(), " is Empty, would not be written !!!!!") else: tree_f.Write() key = ITER.Next() print("") print("////////////////////////////////////////////////") print("outputfile : ") print(outfileName) print("////////////////////////////////////////////////") print("") outfile.Close() f.Close() print( "*********************************************************************************************" ) gBenchmark.Show("Regerating tree root") print( "*********************************************************************************************" )
class TwoDArrayToROOT: # # for saving 2-D array to ROOT file # # ex1) array = [ [a1, b1, c1, ...], # [a2, b2, c2, ...], # ................. ] => Row = Event : use 'FillRow' # ex2) array = [ [a1, a2, a3, ...], # [b1, b2, b3, ...], # ................. ] => Column = Event : use 'FillColumn' # # ------------------------------------------------------------------------------------------------------------ # # code ex) # # from ArrayToRoot import TwoDArrayToROOT # # saveroot = TwoDArrayToROOT('WGAN_data.root') # saveroot.SetTreeName('RealData','momentums of Real Data') # namearray = ['P1_px','P1_py','P1_pz','P2_px','P2_py','P2_pz','b1_px','b1_py','b1_pz','b2_px','b2_py','b2_pz'] # saveroot.FillRow(input_data,namearray) # saveroot.SaveTree() # saveroot.SetTreeName('FakeData','momentums of Fake Data') # saveroot.FillRow(fakedata,namearray) # saveroot.SaveTree() # # output ex) # # SetFileName : WGAN_data.root # SetTreeName : RealData # Tree 'RealData' Saved # SetTreeName : FakeData # Tree 'FakeData' Saved # # [output_path]$ root WGAN_data.root # root [0] # Attaching file WGAN_data.root as _file0... # (TFile *) 0x17b92c0 # root [1] .ls # TFile** WGAN_data.root # TFile* WGAN_data.root # KEY: TTree RealData;1 momentums of Real Data # KEY: TTree FakeData;1 momentums of Fake Data # def __init__(self, filename=None): if filename != None: self.SetFileName(filename) def SetFileName(self, filename, filetype='RECREATE'): self.tfile = TFile(filename, filetype) print('SetFileName : {0}'.format(filename)) def SetTreeName(self, treename, treehelp=''): self.ttree = TTree(treename, treehelp) print('SetTreeName : {0}'.format(treename)) def Fill(self, arraytype, array, arrayname=None): if arraytype == 'R': n_evt = array.shape[0] n_val = array.shape[1] elif arraytype == 'C': n_evt = array.shape[1] n_val = array.shape[0] else: print("ARRAY TYPE ERROR") exit() val = numpy.zeros((n_val, 1), dtype=numpy.float64) for i_val in range(n_val): if arrayname == None: self.ttree.Branch('val_{0}'.format(i_val), val[i_val], 'val_{0}/D'.format(i_val)) else: self.ttree.Branch('{0}'.format(arrayname[i_val]), val[i_val], '{0}/D'.format(arrayname[i_val])) for i_evt in range(n_evt): for j_val in range(n_val): if arraytype == 'R': val[j_val] = array[i_evt, j_val] elif arraytype == 'C': val[j_val] = array[j_val, i_evt] self.ttree.Fill() def FillColumn(self, array, arrayname=None): self.Fill('C', array, arrayname) def FillRow(self, array, arrayname=None): self.Fill('R', array, arrayname) def SaveTree(self): if self.tfile == None: print('No ROOT File') exit() self.ttree.Write() print('Tree \'{0}\' Saved'.format(self.ttree.GetName()))
def REGENERATE_TREE_WITH_CUT(filename, outputpath=''): FileNameList = read_file_name(filename) # print(FileNameList) BranchListAll = get_branch_list_all(FileNameList[2]) BranchListEachTree = get_branch_list_each_tree(FileNameList[2]) # print("@#!@#!@#",BranchListEachTree) # print("@#!@#!@#",BranchListAll) DicNumpyArray_branch = {} for numpyarray in BranchListAll: a = numpy.array([0], 'd') DicNumpyArray_branch[numpyarray] = a DicNumpyArray_branch = collections.OrderedDict( sorted(DicNumpyArray_branch.items()) ) # !!!the input times are ordered!!! # print(DicNumpyArray_branch) DicNumpyArray_branch_w = {} for numpyarray_w in BranchListAll: a_w = numpy.array([0], 'd') DicNumpyArray_branch_w[numpyarray_w] = a_w DicNumpyArray_branch_w = collections.OrderedDict( sorted(DicNumpyArray_branch_w.items())) # print(DicNumpyArray_branch_w) # print(DicNumpyArray_branch.keys()) # print(DicNumpyArray_branch.values()) WCuts = WhetherAddCut(BranchListEachTree) if (WCuts == None): print("\n") print( "*********************************************************************************************" ) print("!!!! Please check input CUT !!!!! ") print(" !!!! No Output File !!!! ") print("\n") # for i in range(len(BranchListEachTree)): # print("compenets below :") # print("The Tree name is : ",BranchListEachTree.keys()[i]) # for j in range(len(BranchListEachTree.values()[i])): # print(" it contains : ", BranchListEachTree.values()[i][j]) # print("\n") print("!!!! Please check input CUT !!!!! ") print(" !!!! No Output File !!!! ") print( "*********************************************************************************************" ) else: print("\n") print("...Cut added and regenerating tree root file...") gBenchmark.Start("Regerating tree root") f = TFile(FileNameList[2], "READ") dirlist = f.GetListOfKeys() ITER = dirlist.MakeIterator() key = ITER.Next() if (outputpath == ''): outfileName = FileNameList[3] + "/" + FileNameList[0] + "_cut.root" outfileName = outfileName.replace("//", "/") # print("!@#!@!@#!@ ",outfileName) elif (outputpath[0] == "/"): outfileName = outputpath + "/" + FileNameList[0] + "_cut.root" outfileName = outfileName.replace("//", "/") # print("!@#!@!@#!@ ",outfileName) elif (outputpath[0] == "~"): outfileName = outputpath.replace( "~", os.environ['HOME']) + "/" + FileNameList[0] + "_cut.root" outfileName = outfileName.replace("//", "/") # print("!@#!@!@#!@ ",outfileName) else: outfileName = os.getcwd( ) + "/" + outputpath + "/" + FileNameList[0] + "_cut.root" outfileName = outfileName.replace("//", "/") # print("!@#!@!@#!@ ",outfileName) outfile = TFile(outfileName, "RECREATE") ijk = 0 while key: tree = key.ReadObj() tree_f = TTree(tree.GetName() + "_f", tree.GetName() + "_f") ENTRY = tree.GetEntries() #print(ENTRY) for i in range(len(DicNumpyArray_branch)): if (list(DicNumpyArray_branch.keys())[i] in BranchListEachTree[tree.GetName()]): tree.SetBranchAddress( list(DicNumpyArray_branch.keys())[i], list(DicNumpyArray_branch.values())[i]) tree_f.Branch( list(DicNumpyArray_branch_w.keys())[i], list(DicNumpyArray_branch_w.values())[i], list(DicNumpyArray_branch_w.keys())[i] + "/D") else: continue print("for tree", tree.GetName()) # tt = 0 for j in range(ENTRY): tree.GetEntry(j) if (j % 5000 == 0): print("now looping", j, "th Events, total of ", ENTRY, "events") for k in range(len(DicNumpyArray_branch)): if (list(DicNumpyArray_branch.keys())[k] in BranchListEachTree[tree.GetName()] ): ### FIXED MAYBE not correct.... pass else: continue list(DicNumpyArray_branch_w.values())[k][0] = list( DicNumpyArray_branch.values())[k][0] # if(j==0): # print(k,DicNumpyArray_branch_w.keys()[k]) # print(DicNumpyArray_branch_w.keys()[k]) # tt = raw_input("press 'Y' to proceed, 'enter' to quit!") #FIXME#FIXME#FIXME#FIXME#FIXME#FIXME#FIXME#FIXME#FIXME#FIXME if (True #FIXME#FIXME#FIXME#FIXME#FIXME#FIXME#FIXME#FIXME#FIXME#FIXME # & (DicNumpyArray_branch.values()[0][0] > 0) # [i][0] means "i+1"th branch of each tree, [0] don't change #FIXME#FIXME#FIXME#FIXME #& (DicNumpyArray_branch.values()[1][0] > 0) #FIXME#FIXME#FIXME#FIXME#FIXME#FIXME#FIXME#FIXME#FIXME#FIXME#FIXME ): ijk = ijk + 1 tree_f.Fill() else: continue print(" Event left after filtering : ", ijk, "!!!!!!!!!") print("\n") ijk = 0 if (tree_f.GetEntries() == 0): print("!!!!!!! ", tree_f.GetName(), " is Empty, would not be written !!!!!") else: tree_f.Write() key = ITER.Next() print("") print("////////////////////////////////////////////////") print("outputfile : ") print(outfileName) print("////////////////////////////////////////////////") print("") outfile.Close() f.Close() print( "*********************************************************************************************" ) gBenchmark.Show("Regerating tree root") print( "*********************************************************************************************" ) return outfileName
def SplitTree(self): FileNameList = read_file_name_root(self._infile) BranchListAll = self.get_branch_list_all() BranchListEachTree = self.get_branch_list_each_tree() DicNumpyArray_branch = {} for numpyarray in BranchListAll: a = numpy.array([0], 'd') DicNumpyArray_branch[numpyarray] = a DicNumpyArray_branch = collections.OrderedDict( sorted(DicNumpyArray_branch.items())) print(DicNumpyArray_branch) DicNumpyArray_branch_w = {} for numpyarray_w in BranchListAll: a_w = numpy.array([0], 'd') DicNumpyArray_branch_w[numpyarray_w] = a_w DicNumpyArray_branch_w = collections.OrderedDict( sorted(DicNumpyArray_branch_w.items())) print(DicNumpyArray_branch_w) gBenchmark.Start("Regerating tree root") f = TFile(FileNameList[2], "READ") dirlist = f.GetListOfKeys() ITER = dirlist.MakeIterator() key = ITER.Next() outfileName = os.getcwd( ) + "/" + FileNameList[0] + "_cut_TTTL.root" #FIXME print("CREATING... :", outfileName) outfile = TFile(outfileName, "RECREATE") ijk = 0 break_flag = 0 while key: if (break_flag == 1): break break_flag += 1 tree = key.ReadObj() tree_f = TTree(tree.GetName() + "", tree.GetName() + "") ENTRY = tree.GetEntries() #print(ENTRY) for i in range(len(DicNumpyArray_branch)): if (list(DicNumpyArray_branch.keys())[i] in BranchListEachTree[tree.GetName()]): tree.SetBranchAddress( list(DicNumpyArray_branch.keys())[i], list(DicNumpyArray_branch.values())[i]) tree_f.Branch( list(DicNumpyArray_branch_w.keys())[i], list(DicNumpyArray_branch_w.values())[i], list(DicNumpyArray_branch_w.keys())[i] + "/D") else: continue print("for tree", tree.GetName()) for j in range(ENTRY): tree.GetEntry(j) if (j % 5000 == 0): print("now looping", j, "th Events, total of ", ENTRY, "events") for k in range(len(DicNumpyArray_branch)): if (list(DicNumpyArray_branch.keys())[k] in BranchListEachTree[tree.GetName()] ): ### FIXED MAYBE not correct.... pass else: continue list(DicNumpyArray_branch_w.values())[k][0] = list( DicNumpyArray_branch.values())[k][0] if (True #& (list(DicNumpyArray_branch.values())[0][0] == 1) # LL #FIXME & (list(DicNumpyArray_branch.values())[0][0] == 0 ) # TTTL #FIXME ): ijk = ijk + 1 tree_f.Fill() else: continue print(" Event left after filtering : ", ijk, "!!!!!!!!!") print("\n") ijk = 0 if (tree_f.GetEntries() == 0): print("!!!!!!! ", tree_f.GetName(), " is Empty, would not be written !!!!!") else: tree_f.Write() key = ITER.Next() print("") print("////////////////////////////////////////////////") print("outputfile : ") print(outfileName) print("////////////////////////////////////////////////") print("") outfile.Close() f.Close() print( "*********************************************************************************************" ) gBenchmark.Show("Regerating tree root") print( "*********************************************************************************************" ) return outfileName