From d4136bda4f56a88aae68bc15a435de5fa927fa5f Mon Sep 17 00:00:00 2001
From: Nikolai Hartmann <Nikolai.Hartmann@physik.uni-muenchen.de>
Date: Wed, 9 May 2018 12:02:58 +0200
Subject: [PATCH] Make add_friend function importable
---
addFriend.py | 51 --------------------------------------------
add_friend.py | 58 +++++++++++++++++++++++++++++++++++++++++++++++++++
toolkit.py | 1 +
3 files changed, 59 insertions(+), 51 deletions(-)
delete mode 100755 addFriend.py
create mode 100755 add_friend.py
diff --git a/addFriend.py b/addFriend.py
deleted file mode 100755
index a04f6d7..0000000
--- a/addFriend.py
+++ /dev/null
@@ -1,51 +0,0 @@
-#!/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()
diff --git a/add_friend.py b/add_friend.py
new file mode 100755
index 0000000..0f1807f
--- /dev/null
+++ b/add_friend.py
@@ -0,0 +1,58 @@
+#!/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)
diff --git a/toolkit.py b/toolkit.py
index 5edaad7..01b17e7 100755
--- a/toolkit.py
+++ b/toolkit.py
@@ -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):
--
GitLab