forked from rootpy/rootpy
-
Notifications
You must be signed in to change notification settings - Fork 1
Expand file tree
/
Copy pathroot-skim
More file actions
executable file
·191 lines (175 loc) · 7.77 KB
/
Copy pathroot-skim
File metadata and controls
executable file
·191 lines (175 loc) · 7.77 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
#!/usr/bin/env python
from optparse import OptionParser
"""
Need to keep this script compatible with Python 2.4
and only dependent on the default PyROOT, not rootpy.
This script must be able to run alone in a very limited environment.
"""
parser = OptionParser()
parser.add_option('-r','--branches-regex',action="store",type="str",dest="branch_regex",default=None,
help="file containing one regex per line for branch selection")
parser.add_option('-g','--branches-glob',action="store",type="str",dest="branch_glob",default=None,
help="file containing one glob per line for branch selection")
parser.add_option('-b','--branches',action="store",type="str",dest="branch_literal",default=None,
help="file containing one branch name per line for branch selection")
parser.add_option('-t','--tree',action="store",type="str",dest="tree",default=None,
help="name of tree (including path) to optimize common to all files")
parser.add_option('--prefix',action="store",type="str",dest="prefix",default="skimmed",
help="prefix of output file/directory name")
parser.add_option('-s','--selection',action="store",type="str",dest="selection",default='',
help="only entries satisfying this cut will be kept")
parser.add_option('-o','--or-if-exists',action="store",type="str",dest="or_if_exists",default=None,
help="OR conditions on all branches existing in current file. " \
"List one branch name and condition on branch separated by whitespace per line")
parser.add_option('--grid',action="store_true",dest="grid",default=False,
help="Running on the grid mode. Input files are separated by commas.")
parser.add_option('--overwrite',action="store_true",dest="overwrite",default=False,
help="Overwrite existing output.")
options, args = parser.parse_args()
import sys
if options.tree is None:
sys.exit("You must specify a tree name")
if not options.prefix:
sys.exit("Invalid prefix")
if not args:
sys.exit("No input ROOT files specified")
if options.grid:
if len(args) != 1:
sys.exit("Expected exactly one argument containing list of files separated by commas")
args = args[0].split(',')
import ROOT
ROOT.PyConfig.IgnoreCommandLineOptions = True
import re
import os
import fnmatch
import uuid
branches_keep_regex = []
if options.branch_regex is not None:
try:
branches_file = open(options.branch_regex,'r')
branches_keep_regex += [re.compile(regex.strip()) for regex in branches_file.readlines()]
branches_file.close()
except Exception, ex:
print ex
sys.exit("Could not parse branch regex file %s"% options.branch_regex)
if options.branch_glob is not None:
try:
branches_file = open(options.branch_glob,'r')
branches_keep_regex += [re.compile(fnmatch.translate(regex.strip())) for regex in branches_file.readlines()]
branches_file.close()
except Exception, ex:
print ex
sys.exit("Could not parse branch glob file %s"% options.branch_regex)
branches_keep_literal = []
if options.branch_literal is not None:
try:
branches_file = open(options.branch_literal,'r')
branches_keep_literal = [line.strip() for line in branches_file.readlines()]
branches_file.close()
except Exception, ex:
print ex
sys.exit("Could not parse branch file %s"% options.branch_literal)
if options.or_if_exists:
or_file = open(options.or_if_exists, 'r')
options.or_if_exists = [line.strip() for line in or_file.readlines()]
or_file.close()
infiles = args
for infilename in infiles:
sub_infiles = [(infilename, "%s."% options.prefix)]
isdir = False
if os.path.isdir(infilename):
isdir = True
infilename = os.path.normpath(infilename)
sub_infiles = []
for root, dirnames, filenames in os.walk(infilename):
base = ("%s."% options.prefix) + os.path.basename(root)
for filename in fnmatch.filter(filenames, '*.root*'):
sub_infiles.append((os.path.join(root, filename), os.path.join(os.path.dirname(root), base) + os.sep))
for filename, out_path in sub_infiles:
if isdir:
if not os.path.exists(out_path):
os.makedirs(out_path)
outfilename = out_path+os.path.basename(filename)
if os.path.exists(outfilename):
if options.overwrite:
print "overwriting %s" % outfilename
outFile = ROOT.TFile(outfilename,'RECREATE')
else:
print "skipping %s which already exists"% outfilename
continue
else:
outFile = ROOT.TFile(outfilename,'NEW')
inFile = ROOT.TFile.Open(filename)
if not inFile:
print "skipping %s which could not be opened"% filename
continue
print "skimming %s ..." % filename
try:
tree = inFile.Get(options.tree)
if not tree:
print "skipping %s where tree %s could not be found" % (
filename, options.tree)
inFile.Close()
outFile.Close()
os.remove(outfilename)
continue
if tree.GetNbranches() == 0:
print "skipping %s where tree %s does not contain any branches"% (
filename, options.tree)
inFile.Close()
outFile.Close()
os.remove(outfilename)
continue
if branches_keep_regex:
tree.SetBranchStatus('*',0)
branches = [b.GetName() for b in tree.GetListOfBranches()]
for branch in branches:
for branch_regex in branches_keep_regex:
if re.search(branch_regex, branch):
print "keeping branch %s" % branch
tree.SetBranchStatus(branch, 1)
break
if branches_keep_literal:
tree.SetBranchStatus('*',0)
for branch in branches_keep_literal:
if tree.GetBranch(branch):
print "keeping branch %s" % branch
tree.SetBranchStatus(branch, 1)
selection = options.selection
if options.or_if_exists:
for branch in options.or_if_exists:
if tree.GetBranch(branch):
tree.SetBranchStatus(branch, 1)
if selection:
selection = '(%s)||(%s)' % (selection, branch)
else:
selection = branch
if selection:
tempfile = ROOT.TFile(uuid.uuid4().hex+".root",'NEW')
tempfile.cd()
print "applying selection %s on tree %s in %s ..." % \
(selection, tree.GetName(), filename)
tree = tree.CopyTree(selection)
tempfile.Write()
tempfile.Flush()
outFile.cd()
print "optimizing tree %s in %s ..." % (tree.GetName(), filename)
newTree = tree.CloneTree(-1, "fast SortBasketsByEntry")
newTree.OptimizeBaskets()
newTree.Write()
outFile.Close()
inFile.Close()
if options.selection:
tempfile.Close()
os.remove(tempfile.GetName())
except KeyboardInterrupt, SystemExit:
print "cleaning up ..."
try:
inFile.Close()
outFile.Close()
os.remove(outfilename)
tempfile.Close()
os.remove(tempfile.GetName())
except:
pass
sys.exit()