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
