-
Notifications
You must be signed in to change notification settings - Fork 3
/
Copy pathdistill.py
145 lines (122 loc) · 5 KB
/
distill.py
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
import ROOT
import glob
import sys
RDF = ROOT.ROOT.RDataFrame
# Branches clasified by diagonal
diagonals = {
# return a tuple: ([left] verticals in 45, [right] verticals in 56))
"d45b_56t" : (["track_rp_5", "track_rp_21", "track_rp_25"],
["track_rp_104", "track_rp_120", "track_rp_124"]),
"ad45b_56b" : (["track_rp_5", "track_rp_21", "track_rp_25"],
["track_rp_105", "track_rp_121", "track_rp_125"]),
"d45t_56b" : (["track_rp_4", "track_rp_20", "track_rp_24"],
["track_rp_105", "track_rp_121", "track_rp_125"]),
"ad45t_56t" : (["track_rp_4", "track_rp_20", "track_rp_24"],
["track_rp_104", "track_rp_120", "track_rp_124"])
}
DS = {
'DS1' : '4495',
'DS2' : '4496',
'DS3' : '4499',
'DS4' : '4505',
'DS5' : '4509',
'DS6' : '4510',
'DS7' : '4511'
}
threads_description = "no_MT"
if len(sys.argv) < 3:
print('Usage: python distill.py <diagonal> <DS> [threads number]')
sys.exit(1) # no diagonal specified
if len(sys.argv) == 4:
if int(sys.argv[3]) < 1:
print('Threads number should be > 0')
sys.exit(1) # wrong threads number
ROOT.ROOT.EnableImplicitMT(int(sys.argv[3]))
threads_description = "threads_" + sys.argv[3]
# Select branches
selected_diagonal = sys.argv[1]
selected_DS = sys.argv[2]
if selected_diagonal not in diagonals.keys():
print('Invalid diagonal: %s' % selected_diagonal)
print('Choose between: %s' % diagonals.keys())
sys.exit(1)
if selected_DS not in DS.keys():
print('DS not available: %s' % selected_DS)
print('Choose between: %s' % DS.keys())
sys.exit(1)
rp_left, rp_right = diagonals[selected_diagonal]
# Extracted from: DS1/block1/input_files.h
source_file = "input_files_{}.txt".format(selected_DS)
input_ntuple_name = "TotemNtuple"
prefix = "root://eostotem.cern.ch//eos/totem/data/cmstotem/2015/90m/Totem/Ntuple/version2/{}/".format(DS[selected_DS])
input_files = [prefix + line.rstrip('\n') for line in open(source_file)]
# Convert to PyROOT vector
vec_input_files = ROOT.vector('string')()
[vec_input_files.push_back(f) for f in input_files]
# Columns per branch
attributes = ['valid', 'x', 'y']
full_branches = ["{}.{}".format(c,a) for a in attributes for c in rp_left+rp_right ]
# Split left and right branch on valid, x and y
valids = [ "(unsigned int) {}".format(v) for v in full_branches[0:6]]
xs = full_branches[6:12]
ys = full_branches[12:18]
print("Selected branches: \n" + "\n\t".join(full_branches))
# Filter and define output branches
filter_code = """({0}.valid + {1}.valid + {2}.valid ) >= 2 &&
({3}.valid + {4}.valid + {5}.valid ) >= 2
""".format(*(rp_left+rp_right))
print("Filter code: \n" + filter_code)
# Input tree
treename= "TotemNtuple"
rdf = RDF(treename, vec_input_files)
# Output tree, file and branches
outTreeName = "distilled"
outFileName = "distill_{}_{}_{}_new.root".format(selected_DS, threads_description, selected_diagonal)
branchList = ["v_L_1_F", "x_L_1_F", "y_L_1_F",
"v_L_2_N", "x_L_2_N", "y_L_2_N",
"v_L_2_F", "x_L_2_F", "y_L_2_F",
"v_R_1_F", "x_R_1_F", "y_R_1_F",
"v_R_2_N", "x_R_2_N", "y_R_2_N",
"v_R_2_F", "x_R_2_F", "y_R_2_F",
"timestamp",
"run_num",
"bunch_num",
"event_num",
"trigger_num",
"trigger_bits"
]
# Convert to PyROOT vector
vec_outbranchlist = ROOT.vector('string')()
[vec_outbranchlist.push_back(b) for b in branchList]
# Filter and define output branches
r = rdf.Filter(filter_code) \
.Define("v_L_1_F", valids[0]) \
.Define("x_L_1_F", xs[0]) \
.Define("y_L_1_F", ys[0]) \
.Define("v_L_2_N", valids[1]) \
.Define("x_L_2_N", xs[1]) \
.Define("y_L_2_N", ys[1]) \
.Define("v_L_2_F", valids[2]) \
.Define("x_L_2_F", xs[2]) \
.Define("y_L_2_F", ys[2]) \
.Define("v_R_1_F", valids[3]) \
.Define("x_R_1_F", xs[3]) \
.Define("y_R_1_F", ys[3]) \
.Define("v_R_2_N", valids[4]) \
.Define("x_R_2_N", xs[4]) \
.Define("y_R_2_N", ys[4]) \
.Define("v_R_2_F", valids[5]) \
.Define("x_R_2_F", xs[5]) \
.Define("y_R_2_F", ys[5]) \
.Define("timestamp", "(unsigned int) (event_info.timestamp - 1444860000)") \
.Define("run_num", "(unsigned int) event_info.run_no") \
.Define("bunch_num", "trigger_data.bunch_num") \
.Define("event_num", "trigger_data.event_num") \
.Define("trigger_num", "trigger_data.trigger_num") \
.Define("trigger_bits", "trigger_data.input_status_bits")
# All above actions are not executed at the moment they are called,
# but they are lazy, i.e. delayed until the moment one of their results
# is accessed (in this case by .GetValue() )
print("Distilled events: %s" % r.Count().GetValue())
# Save output tree
r.Snapshot(outTreeName, outFileName, vec_outbranchlist)