Skip to content
Snippets Groups Projects
Commit d4136bda authored by Nikolai.Hartmann's avatar Nikolai.Hartmann
Browse files

Make add_friend function importable

parent 47ef6c18
No related branches found
No related tags found
No related merge requests found
#!/usr/bin/env python
import argparse
import ROOT
parser = argparse.ArgumentParser(description='add a friend tree to a tree in another file')
parser.add_argument("infile", help="input file that contains the friend tree")
parser.add_argument("intree", help="name of the friend tree")
parser.add_argument("outfile", help="output file where the friend tree should be added")
parser.add_argument("outtree", help="name of the tree (in output file) to which the friend should be added")
args = parser.parse_args()
outfile = ROOT.TFile.Open(args.outfile, "UPDATE")
infile = ROOT.TFile.Open(args.infile)
friend_name = args.outtree+"_friend_"+args.intree
for k in outfile.GetListOfKeys():
if k.GetName() == friend_name:
raise ValueError("Tree with name {} already exists in outputfile".format(args.intree))
outfile.cd()
outtree = outfile.Get(args.outtree)
if not outtree:
raise KeyError("Tree {} not found in file {}".format(args.outtree, args.outfile))
if outtree.GetListOfFriends():
for k in outtree.GetListOfFriends():
if k.GetName() == friend_name:
raise ValueError("Tree with name {} is already friend of {}".format(args.intree, args.outtree))
infile.cd()
intree = infile.Get(args.intree)
if not intree:
raise KeyError("Tree {} not found in file {}".format(args.intree, args.infile))
# Add friend and write friend tree and original tree to outfile
outfile.cd()
clonetree = intree.CloneTree(-1, "fast")
clonetree.SetName(friend_name)
clonetree.SetTitle(friend_name)
clonetree.Write(friend_name)
outtree.AddFriend(clonetree)
outtree.Write(outtree.GetName())
infile.Close()
outfile.Close()
#!/usr/bin/env python
import ROOT
def add_friend(infile, intree, outfile, outtree):
root_outfile = ROOT.TFile.Open(outfile, "UPDATE")
root_infile = ROOT.TFile.Open(infile)
friend_name = outtree+"_friend_"+intree
for k in root_outfile.GetListOfKeys():
if k.GetName() == friend_name:
raise ValueError("Tree with name {} already exists in outputfile".format(intree))
root_outfile.cd()
root_outtree = root_outfile.Get(outtree)
if not root_outtree:
raise KeyError("Tree {} not found in file {}".format(outtree, outfile))
if root_outtree.GetListOfFriends():
for k in root_outtree.GetListOfFriends():
if k.GetName() == friend_name:
raise ValueError("Tree with name {} is already friend of {}".format(intree, outtree))
root_infile.cd()
root_intree = root_infile.Get(intree)
if not root_intree:
raise KeyError("Tree {} not found in file {}".format(intree, infile))
# Add friend and write friend tree and original tree to outfile
root_outfile.cd()
clonetree = root_intree.CloneTree(-1, "fast")
clonetree.SetName(friend_name)
clonetree.SetTitle(friend_name)
clonetree.Write(friend_name)
root_outtree.AddFriend(clonetree)
root_outtree.Write(root_outtree.GetName())
root_infile.Close()
root_outfile.Close()
if __name__ == "__main__":
import argparse
parser = argparse.ArgumentParser(description='add a friend tree to a tree in another file')
parser.add_argument("infile", help="input file that contains the friend tree")
parser.add_argument("intree", help="name of the friend tree")
parser.add_argument("outfile", help="output file where the friend tree should be added")
parser.add_argument("outtree", help="name of the tree (in output file) to which the friend should be added")
args = parser.parse_args()
add_friend(args.infile, args.intree, args.outfile, args.outtree)
......@@ -608,6 +608,7 @@ class ClassificationProject(object):
f = ROOT.TFile.Open(source_filename)
tree = f.Get(source_treename)
entries = tree.GetEntries()
logger.info("Write friend tree for {} in {}".format(source_treename, source_filename))
if os.path.exists(target_filename):
raise IOError("{} already exists, if you want to recreate it, delete it first".format(target_filename))
for start in range(0, entries, batch_size):
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment