Belle II Software light-2406-ragdoll
per_event_statistics.py
1
8
9import basf2
10import numpy as np
11
12
14 """
15 A basf2 python module to export *all* module time statistics (of every event, not just averaged)
16 into a ROOT TTree written to a file.
17 """
18
19 def __init__(self, output_file_name):
20 """
21 Create a new PerEventStatisticsGetterModule. You have to give the name of the
22 file, where the TTree with the full event statistics will be saved.
23 """
24 super().__init__()
25
26
27 self.output_file_name = output_file_name
28
29 self.tfile = None
30
31 self.statistics = None
32
33
34 self.event_number = np.zeros(3, dtype=float)
35
36 self.ttree_inputs = None
37
38 self.last_time_sum = None
39
40
41 self.branches_added = False
42
43 # Always avoid the top-level 'import ROOT'.
44 import ROOT # noqa
45
46
47 self.event_meta_data = ROOT.Belle2.PyStoreObj("EventMetaData") # noqa
48
49 def initialize(self):
50 """
51 Create the needed store object pointer in the DataStore and the TFile with the TTree.
52 """
53 # Always avoid the top-level 'import ROOT'.
54 import ROOT # noqa
55
56 self.event_meta_data.isRequired()
57
58 # try to avoid side effects ...
59 old = ROOT.gDirectory
60 # Before creating the TTree, open the output file.
61 self.tfile = ROOT.TFile.Open(self.output_file_name, "RECREATE")
62 # Create the TTree.
63 self.statistics = ROOT.TTree("statistics", "Per event execution Statistics for all modules")
64 # and reset stupid global directory pointer
65 ROOT.gDirectory = old
66
67 # Add branches identifying the event. The reason these are [i:] instead
68 # of just [i] is that the latter just returns a single float object which
69 # doesn't work as a buffer to TTree::Branch. We need a buffer with an
70 # address where it can write the value. So we take a reference to the
71 # array starting at element i with [i:]
72 self.statistics.Branch("expNo", self.event_number[0:], "expNo/D")
73 self.statistics.Branch("runNo", self.event_number[1:], "runNo/D")
74 self.statistics.Branch("evtNo", self.event_number[2:], "evtNo/D")
75
76 def event(self):
77 """
78 The event loop: Store the statistics as a new row in the TTree.
79 """
80 # Always avoid the top-level 'import ROOT'.
81 import ROOT # noqa
82
83 # TTree reference (for shorter code)
84 ttree = self.statistics
85
86 # Get the statistics
87 module_stats = basf2.statistics.modules
88
89 # Fill in the event information
90 self.event_number[0] = self.event_meta_data.getExperiment()
91 self.event_number[1] = self.event_meta_data.getRun()
92 self.event_number[2] = self.event_meta_data.getEvent()
93
94 # Again, change the directory
95 old = ROOT.gDirectory
96 self.tfile.cd()
97
98 if not self.branches_added:
99
100 # make sure all this is only done in the output process
101 if ROOT.Belle2.ProcHandler.parallelProcessingUsed() and not ROOT.Belle2.ProcHandler.isOutputProcess():
102 basf2.B2FATAL("PerEventStatisticsGetterModule can only be used in single processing mode or in the output process")
103
104 self.ttree_inputs = np.zeros(len(module_stats), dtype=float)
105 self.last_time_sum = np.zeros(len(module_stats), dtype=float)
106
107 for i, stat in enumerate(module_stats):
108 # escape the module names in ROOT-safe manner. Otherwise weird stuff happens like
109 # sub-branches get created or the branch cannot be opened in the TBrowser
110 module_name = f"{ROOT.Belle2.MakeROOTCompatible.makeROOTCompatible(stat.name)}_{i}"
111 # Create the branch. See initialize() for why it's [i:] and not [i]
112 ttree.Branch(module_name, self.ttree_inputs[i:], f"{module_name}/D")
113
114 self.branches_added = True
115
116 # Fill in the module statistics into the branches
117 time_sum = np.array([m.time_sum(basf2.statistics.EVENT) for m in module_stats], dtype=float)
118 # We need to change the contents of the ttree_inputs array **without**
119 # changing its address so we cannot just assign. But we can assign to
120 # to the full array itself using the empty slice syntax.
121 self.ttree_inputs[:] = time_sum - self.last_time_sum
122 self.last_time_sum[:] = time_sum
123
124 # Send the branches to the TTree.
125 ttree.Fill()
126
127 # and back again
128 ROOT.gDirectory = old
129
130 def terminate(self):
131 """
132 Write out the merged statistics to the ROOT file.
133 This should only be called once, as we would end up with different versions otherwise.
134 """
135 # Always avoid the top-level 'import ROOT'.
136 import ROOT # noqa
137
138 # Change the path a last time
139 old = ROOT.gDirectory
140 self.tfile.cd()
141
142 self.statistics.Write()
143 self.tfile.Close()
144
145 # and back again
146 ROOT.gDirectory = old
last_time_sum
Last recored sum of event calls for all modules.
ttree_inputs
The columns for the statistics TTree (they will be filled in the event function).
tfile
Will host the pointer to the opened TFile later.
event_number
The columns to store the event number.
branches_added
A flag to indicate that we have already added the Branches to the TTree (which we will do in the firs...