Belle II Software light-2406-ragdoll
VariablesToEventBasedTreeModule.cc
1/**************************************************************************
2 * basf2 (Belle II Analysis Software Framework) *
3 * Author: The Belle II Collaboration *
4 * *
5 * See git log for contributors and copyright holders. *
6 * This file is licensed under LGPL-3.0, see LICENSE.md. *
7 **************************************************************************/
8
9#include <analysis/modules/VariablesToEventBasedTree/VariablesToEventBasedTreeModule.h>
10
11// analysis
12#include <analysis/dataobjects/ParticleList.h>
13#include <analysis/dataobjects/StringWrapper.h>
14#include <analysis/VariableManager/Manager.h>
15#include <analysis/VariableManager/Utility.h>
16
17// framework
18#include <framework/logging/Logger.h>
19#include <framework/pcore/ProcHandler.h>
20#include <framework/utilities/MakeROOTCompatible.h>
21#include <framework/utilities/RootFileCreationManager.h>
22#include <framework/core/ModuleParam.templateDetails.h>
23#include <framework/core/Environment.h>
24
25#include <cmath>
26
27using namespace std;
28using namespace Belle2;
29
30// Register module in the framework
31REG_MODULE(VariablesToEventBasedTree);
32
33
35 Module(), m_tree("", DataStore::c_Persistent)
36{
37 //Set module properties
38 setDescription("Calculate variables specified by the user for a given ParticleList and save them into a TTree. The Tree is event-based, meaning that the variables of each candidate for each event are saved in an array of a branch of the Tree.");
40
41 vector<string> emptylist;
42 addParam("particleList", m_particleList,
43 "Name of particle list with reconstructed particles. An empty ParticleList is not supported. Use the VariablesToNtupleModule for this use-case",
44 std::string(""));
45 addParam("variables", m_variables,
46 "List of variables (or collections) to save for each candidate. Variables are taken from Variable::Manager, and are identical to those available to e.g. ParticleSelector.",
47 emptylist);
48
49 addParam("event_variables", m_event_variables,
50 "List of variables (or collections) to save for each event. Variables are taken from Variable::Manager, and are identical to those available to e.g. ParticleSelector. Only event-based variables are allowed here.",
51 emptylist);
52
53 addParam("fileName", m_fileName, "Name of ROOT file for output. Can be overridden using the -o argument of basf2.",
54 string("VariablesToEventBasedTree.root"));
55 addParam("treeName", m_treeName, "Name of the NTuple in the saved file.", string("tree"));
56 addParam("maxCandidates", m_maxCandidates, "The maximum number of candidates in the ParticleList per entry of the Tree.", 100u);
57
58 std::tuple<std::string, std::map<int, unsigned int>> default_sampling{"", {}};
59 addParam("sampling", m_sampling,
60 "Tuple of variable name and a map of integer values and inverse sampling rate. E.g. (signal, {1: 0, 0:10}) selects all signal events and every 10th background event. Variable must be event-based.",
61 default_sampling);
62
63 addParam("fileNameSuffix", m_fileNameSuffix, "The suffix of the output ROOT file to be appended before ``.root``.",
64 string(""));
65 addParam("ignoreCommandLineOverride", m_ignoreCommandLineOverride,
66 "Ignore override of file name via command line argument -o. Useful if you have multiple output modules in one path.", false);
67
68 addParam("storeEventType", m_storeEventType,
69 "If true, the branch __eventType__ is added. The eventType information is available from MC16 on.", true);
70}
71
73{
74 m_eventMetaData.isRequired();
76
77 // override the output file name with what's been provided with the -o option
79 const std::string& outputFileArgument = Environment::Instance().getOutputFileOverride();
80 if (!outputFileArgument.empty())
81 m_fileName = outputFileArgument;
82 }
83
84 if (!m_fileNameSuffix.empty())
85 m_fileName = m_fileName.insert(m_fileName.rfind(".root"), m_fileNameSuffix);
86
87 // See if there is already a file in which case add a new tree to it ...
88 // otherwise create a new file (all handled by framework)
90 if (!m_file) {
91 B2ERROR("Could not create file \"" << m_fileName <<
92 "\". Please set a valid root output file name (\"fileName\" module parameter).");
93 return;
94 }
95
96 m_file->cd();
97
98 // check if TTree with that name already exists
99 if (m_file->Get(m_treeName.c_str())) {
100 B2FATAL("Tree with the name " << m_treeName << " already exists in the file " << m_fileName);
101 return;
102 }
103
105 // remove duplicates from list of variables but keep the previous order
106 unordered_set<string> seen;
107 auto newEnd = remove_if(m_variables.begin(), m_variables.end(), [&seen](const string & varStr) {
108 if (seen.find(varStr) != std::end(seen)) return true;
109 seen.insert(varStr);
110 return false;
111 });
112 m_variables.erase(newEnd, m_variables.end());
113
115 // remove duplicates from list of variables but keep the previous order
116 unordered_set<string> seenEventVariables;
117 auto eventVariablesEnd = remove_if(m_event_variables.begin(),
118 m_event_variables.end(), [&seenEventVariables](const string & varStr) {
119 if (seenEventVariables.find(varStr) != std::end(seenEventVariables)) return true;
120 seenEventVariables.insert(varStr);
121 return false;
122 });
123 m_event_variables.erase(eventVariablesEnd, m_event_variables.end());
124
126 m_tree.construct(m_treeName.c_str(), "");
127
128 m_valuesDouble.resize(m_variables.size());
129 m_valuesInt.resize(m_variables.size());
132
133 m_tree->get().Branch("__event__", &m_event, "__event__/i");
134 m_tree->get().Branch("__run__", &m_run, "__run__/I");
135 m_tree->get().Branch("__experiment__", &m_experiment, "__experiment__/I");
136 m_tree->get().Branch("__production__", &m_production, "__production__/I");
137 m_tree->get().Branch("__ncandidates__", &m_ncandidates, "__ncandidates__/I");
138 m_tree->get().Branch("__weight__", &m_weight, "__weight__/F");
139
140 if (m_stringWrapper.isOptional("MCDecayString"))
141 m_tree->get().Branch("__MCDecayString__", &m_MCDecayString);
142
143 if (m_storeEventType) {
144 m_tree->get().Branch("__eventType__", &m_eventType);
145 if (not m_eventExtraInfo.isOptional())
146 B2INFO("EventExtraInfo is not registered. __eventType__ will be empty. The eventType is available from MC16 on.");
147 }
148
149
150 for (unsigned int i = 0; i < m_event_variables.size(); ++i) {
151 auto varStr = m_event_variables[i];
152
153 if (Variable::isCounterVariable(varStr)) {
154 B2WARNING("The counter '" << varStr
155 << "' is handled automatically by VariablesToEventBasedTree, you don't need to add it.");
156 continue;
157 }
158
159 //also collection function pointers
161 if (!var) {
162 B2ERROR("Variable '" << varStr << "' is not available in Variable::Manager!");
163 } else {
164 if (var->variabletype == Variable::Manager::VariableDataType::c_double) {
165 m_tree->get().Branch(MakeROOTCompatible::makeROOTCompatible(varStr).c_str(), &m_event_valuesDouble[i],
166 (MakeROOTCompatible::makeROOTCompatible(varStr) + "/D").c_str());
167 } else if (var->variabletype == Variable::Manager::VariableDataType::c_int) {
168 m_tree->get().Branch(MakeROOTCompatible::makeROOTCompatible(varStr).c_str(), &m_event_valuesInt[i],
169 (MakeROOTCompatible::makeROOTCompatible(varStr) + "/I").c_str());
170 } else if (var->variabletype == Variable::Manager::VariableDataType::c_bool) {
171 m_tree->get().Branch(MakeROOTCompatible::makeROOTCompatible(varStr).c_str(), &m_event_valuesInt[i],
172 (MakeROOTCompatible::makeROOTCompatible(varStr) + "/O").c_str());
173 }
174 m_event_functions.push_back(var->function);
175 }
176 }
177
178 for (unsigned int i = 0; i < m_variables.size(); ++i) {
179 auto varStr = m_variables[i];
181 m_valuesInt[i].resize(m_maxCandidates);
182
183 //also collection function pointers
185 if (!var) {
186 B2ERROR("Variable '" << varStr << "' is not available in Variable::Manager!");
187 } else {
188 if (var->variabletype == Variable::Manager::VariableDataType::c_double) {
189 m_tree->get().Branch(MakeROOTCompatible::makeROOTCompatible(varStr).c_str(), &m_valuesDouble[i][0],
190 (MakeROOTCompatible::makeROOTCompatible(varStr) + "[__ncandidates__]/D").c_str());
191 } else if (var->variabletype == Variable::Manager::VariableDataType::c_int) {
192 m_tree->get().Branch(MakeROOTCompatible::makeROOTCompatible(varStr).c_str(), &m_valuesInt[i][0],
193 (MakeROOTCompatible::makeROOTCompatible(varStr) + "[__ncandidates__]/I").c_str());
194 } else if (var->variabletype == Variable::Manager::VariableDataType::c_bool) {
195 m_tree->get().Branch(MakeROOTCompatible::makeROOTCompatible(varStr).c_str(), &m_valuesInt[i][0],
196 (MakeROOTCompatible::makeROOTCompatible(varStr) + "[__ncandidates__]/O").c_str());
197 }
198 m_functions.push_back(var->function);
199 }
200 }
201
202 m_sampling_name = std::get<0>(m_sampling);
203 m_sampling_rates = std::get<1>(m_sampling);
204
205 if (m_sampling_name != "") {
207 if (m_sampling_variable == nullptr) {
208 B2FATAL("Couldn't find sample variable " << m_sampling_name << " via the Variable::Manager. Check the name!");
209 }
210 for (const auto& pair : m_sampling_rates)
211 m_sampling_counts[pair.first] = 0;
212 } else {
213 m_sampling_variable = nullptr;
214 }
215
216}
217
218
220{
221
222 if (m_sampling_variable == nullptr)
223 return 1.0;
224
225 long target = 0;
226 if (m_sampling_variable->variabletype == Variable::Manager::VariableDataType::c_double) {
227 target = std::lround(std::get<double>(m_sampling_variable->function(nullptr)));
228 } else if (m_sampling_variable->variabletype == Variable::Manager::VariableDataType::c_int) {
229 target = std::lround(std::get<int>(m_sampling_variable->function(nullptr)));
230 } else if (m_sampling_variable->variabletype == Variable::Manager::VariableDataType::c_bool) {
231 target = std::lround(std::get<bool>(m_sampling_variable->function(nullptr)));
232 }
233
234 if (m_sampling_rates.find(target) != m_sampling_rates.end() and m_sampling_rates[target] > 0) {
235 m_sampling_counts[target]++;
236 if (m_sampling_counts[target] % m_sampling_rates[target] != 0)
237 return 0;
238 else {
239 m_sampling_counts[target] = 0;
240 return m_sampling_rates[target];
241 }
242 }
243
244 return 1.0;
245}
246
248{
249 // get counter numbers
250 m_event = m_eventMetaData->getEvent();
251 m_run = m_eventMetaData->getRun();
252 m_experiment = m_eventMetaData->getExperiment();
253 m_production = m_eventMetaData->getProduction();
254
255 if (m_stringWrapper.isValid())
256 m_MCDecayString = m_stringWrapper->getString();
257 else
258 m_MCDecayString = "";
259
260 if (m_storeEventType and m_eventExtraInfo.isValid())
261 m_eventType = m_eventExtraInfo->getEventType();
262 else
263 m_eventType = "";
264
266 m_ncandidates = particlelist->getListSize();
268 if (m_weight > 0) {
269 for (unsigned int iVar = 0; iVar < m_event_functions.size(); iVar++) {
270 if (std::holds_alternative<double>(m_event_functions[iVar](nullptr))) {
271 m_event_valuesDouble[iVar] = std::get<double>(m_event_functions[iVar](nullptr));
272 } else if (std::holds_alternative<int>(m_event_functions[iVar](nullptr))) {
273 m_event_valuesInt[iVar] = std::get<int>(m_event_functions[iVar](nullptr));
274 } else if (std::holds_alternative<bool>(m_event_functions[iVar](nullptr))) {
275 m_event_valuesInt[iVar] = std::get<bool>(m_event_functions[iVar](nullptr));
276 }
277 }
278 for (unsigned int iPart = 0; iPart < m_ncandidates; iPart++) {
279
280 if (iPart >= m_maxCandidates) {
281 B2WARNING("Maximum number of candidates exceeded in VariablesToEventBasedTree module. I will skip additional candidates");
283 break;
284 }
285
286 const Particle* particle = particlelist->getParticle(iPart);
287 for (unsigned int iVar = 0; iVar < m_functions.size(); iVar++) {
288 if (std::holds_alternative<double>(m_functions[iVar](particle))) {
289 m_valuesDouble[iVar][iPart] = std::get<double>(m_functions[iVar](particle));
290 } else if (std::holds_alternative<int>(m_functions[iVar](particle))) {
291 m_valuesInt[iVar][iPart] = std::get<int>(m_functions[iVar](particle));
292 } else if (std::holds_alternative<bool>(m_functions[iVar](particle))) {
293 m_valuesInt[iVar][iPart] = std::get<bool>(m_functions[iVar](particle));
294 }
295 }
296 }
297 m_tree->get().Fill();
298 }
299}
300
302{
304 B2INFO("Writing TTree " << m_treeName);
305 TDirectory::TContext directoryGuard(m_file.get());
306 m_tree->write(m_file.get());
307
308 const bool writeError = m_file->TestBit(TFile::kWriteError);
309 m_file.reset();
310 if (writeError) {
311 B2FATAL("A write error occurred while saving '" << m_fileName << "', please check if enough disk space is available.");
312 }
313 }
314}
In the store you can park objects that have to be accessed by various modules.
Definition: DataStore.h:51
@ c_DontWriteOut
Object/array should be NOT saved by output modules.
Definition: DataStore.h:71
const std::string & getOutputFileOverride() const
Return overriden output file name, or "" if none was set.
Definition: Environment.h:127
static Environment & Instance()
Static method to get a reference to the Environment instance.
Definition: Environment.cc:28
static std::string makeROOTCompatible(std::string str)
Remove special characters that ROOT dislikes in branch names, e.g.
Base class for Modules.
Definition: Module.h:72
void setDescription(const std::string &description)
Sets the description of the module.
Definition: Module.cc:214
void setPropertyFlags(unsigned int propertyFlags)
Sets the flags for the module properties.
Definition: Module.cc:208
@ c_ParallelProcessingCertified
This module can be run in parallel processing mode safely (All I/O must be done through the data stor...
Definition: Module.h:80
@ c_TerminateInAllProcesses
When using parallel processing, call this module's terminate() function in all processes().
Definition: Module.h:83
Class to store reconstructed particles.
Definition: Particle.h:75
static bool isOutputProcess()
Return true if the process is an output process.
Definition: ProcHandler.cc:232
static bool parallelProcessingUsed()
Returns true if multiple processes have been spawned, false in single-core mode.
Definition: ProcHandler.cc:226
bool isRequired(const std::string &name="")
Ensure this array/object has been registered previously.
Type-safe access to single objects in the data store.
Definition: StoreObjPtr.h:96
std::vector< std::string > resolveCollections(const std::vector< std::string > &variables)
Resolve Collection Returns variable names corresponding to the given collection or if it is not a col...
Definition: Manager.cc:179
const Var * getVariable(std::string name)
Get the variable belonging to the given key.
Definition: Manager.cc:57
static Manager & Instance()
get singleton instance.
Definition: Manager.cc:25
std::vector< std::string > m_variables
List of variables to save.
unsigned int m_maxCandidates
maximum number of candidates which is written out
std::vector< std::string > m_event_variables
List of event variables to save.
virtual void initialize() override
Initialises the module.
std::map< int, unsigned int > m_sampling_rates
Inverse sampling rates.
std::vector< int > m_event_valuesInt
Values of type int corresponding to given event variables.
virtual void event() override
Method called for each event.
unsigned int m_ncandidates
number of candidates in this event
virtual void terminate() override
Write TTree to file, and close file if necessary.
StoreObjPtr< EventMetaData > m_eventMetaData
event metadata (get event number etc)
std::map< int, unsigned long int > m_sampling_counts
Current number of samples with this value.
std::string m_fileName
Name of ROOT file for output.
std::tuple< std::string, std::map< int, unsigned int > > m_sampling
Tuple of variable name and a map of integer values and inverse sampling rate.
std::vector< std::vector< double > > m_valuesDouble
Values of type double corresponding to given variables.
float getInverseSamplingRateWeight()
Calculate inverse sampling rate weight.
int m_production
production ID (to distinguish MC samples)
StoreObjPtr< StringWrapper > m_stringWrapper
string wrapper storing the MCDecayString
bool m_storeEventType
If true, the branch eventType is added.
std::vector< Variable::Manager::FunctionPtr > m_functions
List of function pointers corresponding to given variables.
std::string m_particleList
Name of particle list with reconstructed particles.
bool m_ignoreCommandLineOverride
if true, ignore override of filename
StoreObjPtr< RootMergeable< TTree > > m_tree
The ROOT TNtuple for output.
std::shared_ptr< TFile > m_file
ROOT file for output.
std::string m_sampling_name
Variable name of sampling variable.
StoreObjPtr< EventExtraInfo > m_eventExtraInfo
pointer to EventExtraInfo
std::vector< std::vector< int > > m_valuesInt
Values of type int corresponding to given variables.
std::vector< double > m_event_valuesDouble
Values of type double corresponding to given event variables.
const Variable::Manager::Var * m_sampling_variable
Variable Pointer to target variable.
std::vector< Variable::Manager::FunctionPtr > m_event_functions
List of function pointers corresponding to given event variables.
std::string m_fileNameSuffix
Suffix to be appended to the output file name.
std::string m_MCDecayString
MC decay string to be filled.
void addParam(const std::string &name, T &paramVariable, const std::string &description, const T &defaultValue)
Adds a new parameter to the module.
Definition: Module.h:560
std::shared_ptr< TFile > getFile(std::string, bool ignoreErrors=false)
Get a file with a specific name, if is does not exist it will be created.
static RootFileCreationManager & getInstance()
Interface for the FileManager.
#define REG_MODULE(moduleName)
Register the given module (without 'Module' suffix) with the framework.
Definition: Module.h:650
Abstract base class for different kinds of events.
Definition: ClusterUtils.h:24
STL namespace.
VariableDataType variabletype
data type of variable
Definition: Manager.h:133
A variable returning a floating-point value for a given Particle.
Definition: Manager.h:146
FunctionPtr function
Pointer to function.
Definition: Manager.h:147