Belle II Software development
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/Utility.h>
15
16// framework
17#include <framework/logging/Logger.h>
18#include <framework/pcore/ProcHandler.h>
19#include <framework/utilities/MakeROOTCompatible.h>
20#include <framework/utilities/RootFileCreationManager.h>
21#include <framework/core/ModuleParam.templateDetails.h>
22#include <framework/core/Environment.h>
23
24#include <cmath>
25
26using namespace std;
27using namespace Belle2;
28
29// Register module in the framework
30REG_MODULE(VariablesToEventBasedTree);
31
32
34 Module(), m_tree("", DataStore::c_Persistent)
35{
36 //Set module properties
37 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.");
39
40 vector<string> emptylist;
41 addParam("particleList", m_particleList,
42 "Name of particle list with reconstructed particles. An empty ParticleList is not supported. Use the VariablesToNtupleModule for this use-case",
43 std::string(""));
44 addParam("variables", m_variables,
45 "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.",
46 emptylist);
47
48 addParam("event_variables", m_event_variables,
49 "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.",
50 emptylist);
51
52 addParam("fileName", m_fileName, "Name of ROOT file for output. Can be overridden using the -o argument of basf2.",
53 string("VariablesToEventBasedTree.root"));
54 addParam("treeName", m_treeName, "Name of the NTuple in the saved file.", string("tree"));
55 addParam("maxCandidates", m_maxCandidates, "The maximum number of candidates in the ParticleList per entry of the Tree.", 100u);
56
57 std::tuple<std::string, std::map<int, unsigned int>> default_sampling{"", {}};
58 addParam("sampling", m_sampling,
59 "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.",
60 default_sampling);
61
62 addParam("fileNameSuffix", m_fileNameSuffix, "The suffix of the output ROOT file to be appended before ``.root``.",
63 string(""));
64 addParam("ignoreCommandLineOverride", m_ignoreCommandLineOverride,
65 "Ignore override of file name via command line argument -o. Useful if you have multiple output modules in one path.", false);
66
67 addParam("storeEventType", m_storeEventType,
68 "If true, the branch __eventType__ is added. The eventType information is available from MC16 on.", true);
69}
70
72{
73 m_eventMetaData.isRequired();
75
76 // override the output file name with what's been provided with the -o option
78 const std::string& outputFileArgument = Environment::Instance().getOutputFileOverride();
79 if (!outputFileArgument.empty())
80 m_fileName = outputFileArgument;
81 }
82
83 if (!m_fileNameSuffix.empty())
84 m_fileName = m_fileName.insert(m_fileName.rfind(".root"), m_fileNameSuffix);
85
86 // See if there is already a file in which case add a new tree to it ...
87 // otherwise create a new file (all handled by framework)
89 if (!m_file) {
90 B2ERROR("Could not create file \"" << m_fileName <<
91 "\". Please set a valid root output file name (\"fileName\" module parameter).");
92 return;
93 }
94
95 m_file->cd();
96
97 // check if TTree with that name already exists
98 if (m_file->Get(m_treeName.c_str())) {
99 B2FATAL("Tree with the name " << m_treeName << " already exists in the file " << m_fileName);
100 return;
101 }
102
104 // remove duplicates from list of variables but keep the previous order
105 unordered_set<string> seen;
106 auto newEnd = remove_if(m_variables.begin(), m_variables.end(), [&seen](const string & varStr) {
107 if (seen.find(varStr) != std::end(seen)) return true;
108 seen.insert(varStr);
109 return false;
110 });
111 m_variables.erase(newEnd, m_variables.end());
112
114 // remove duplicates from list of variables but keep the previous order
115 unordered_set<string> seenEventVariables;
116 auto eventVariablesEnd = remove_if(m_event_variables.begin(),
117 m_event_variables.end(), [&seenEventVariables](const string & varStr) {
118 if (seenEventVariables.find(varStr) != std::end(seenEventVariables)) return true;
119 seenEventVariables.insert(varStr);
120 return false;
121 });
122 m_event_variables.erase(eventVariablesEnd, m_event_variables.end());
123
125 m_tree.construct(m_treeName.c_str(), "");
126
127 m_valuesDouble.resize(m_variables.size());
128 m_valuesInt.resize(m_variables.size());
131
132 m_tree->get().Branch("__event__", &m_event, "__event__/i");
133 m_tree->get().Branch("__run__", &m_run, "__run__/I");
134 m_tree->get().Branch("__experiment__", &m_experiment, "__experiment__/I");
135 m_tree->get().Branch("__production__", &m_production, "__production__/I");
136 m_tree->get().Branch("__ncandidates__", &m_ncandidates, "__ncandidates__/I");
137 m_tree->get().Branch("__weight__", &m_weight, "__weight__/F");
138
139 if (m_stringWrapper.isOptional("MCDecayString"))
140 m_tree->get().Branch("__MCDecayString__", &m_MCDecayString);
141
142 if (m_storeEventType) {
143 m_tree->get().Branch("__eventType__", &m_eventType);
144 if (not m_eventExtraInfo.isOptional())
145 B2INFO("EventExtraInfo is not registered. __eventType__ will be empty. The eventType is available from MC16 on.");
146 }
147
148
149 for (unsigned int i = 0; i < m_event_variables.size(); ++i) {
150 auto varStr = m_event_variables[i];
151
152 if (Variable::isCounterVariable(varStr)) {
153 B2WARNING("The counter '" << varStr
154 << "' is handled automatically by VariablesToEventBasedTree, you don't need to add it.");
155 continue;
156 }
157
158 //also collection function pointers
160 if (!var) {
161 B2ERROR("Variable '" << varStr << "' is not available in Variable::Manager!");
162 } else {
163 if (var->variabletype == Variable::Manager::VariableDataType::c_double) {
164 m_tree->get().Branch(MakeROOTCompatible::makeROOTCompatible(varStr).c_str(), &m_event_valuesDouble[i],
165 (MakeROOTCompatible::makeROOTCompatible(varStr) + "/D").c_str());
166 } else if (var->variabletype == Variable::Manager::VariableDataType::c_int) {
167 m_tree->get().Branch(MakeROOTCompatible::makeROOTCompatible(varStr).c_str(), &m_event_valuesInt[i],
168 (MakeROOTCompatible::makeROOTCompatible(varStr) + "/I").c_str());
169 } else if (var->variabletype == Variable::Manager::VariableDataType::c_bool) {
170 m_tree->get().Branch(MakeROOTCompatible::makeROOTCompatible(varStr).c_str(), &m_event_valuesInt[i],
171 (MakeROOTCompatible::makeROOTCompatible(varStr) + "/O").c_str());
172 }
173 m_event_functions.push_back(var->function);
174 }
175 }
176
177 for (unsigned int i = 0; i < m_variables.size(); ++i) {
178 auto varStr = m_variables[i];
180 m_valuesInt[i].resize(m_maxCandidates);
181
182 //also collection function pointers
184 if (!var) {
185 B2ERROR("Variable '" << varStr << "' is not available in Variable::Manager!");
186 } else {
187 if (var->variabletype == Variable::Manager::VariableDataType::c_double) {
188 m_tree->get().Branch(MakeROOTCompatible::makeROOTCompatible(varStr).c_str(), &m_valuesDouble[i][0],
189 (MakeROOTCompatible::makeROOTCompatible(varStr) + "[__ncandidates__]/D").c_str());
190 } else if (var->variabletype == Variable::Manager::VariableDataType::c_int) {
191 m_tree->get().Branch(MakeROOTCompatible::makeROOTCompatible(varStr).c_str(), &m_valuesInt[i][0],
192 (MakeROOTCompatible::makeROOTCompatible(varStr) + "[__ncandidates__]/I").c_str());
193 } else if (var->variabletype == Variable::Manager::VariableDataType::c_bool) {
194 m_tree->get().Branch(MakeROOTCompatible::makeROOTCompatible(varStr).c_str(), &m_valuesInt[i][0],
195 (MakeROOTCompatible::makeROOTCompatible(varStr) + "[__ncandidates__]/O").c_str());
196 }
197 m_functions.push_back(var->function);
198 }
199 }
200
201 m_sampling_name = std::get<0>(m_sampling);
202 m_sampling_rates = std::get<1>(m_sampling);
203
204 if (m_sampling_name != "") {
206 if (m_sampling_variable == nullptr) {
207 B2FATAL("Couldn't find sample variable " << m_sampling_name << " via the Variable::Manager. Check the name!");
208 }
209 for (const auto& pair : m_sampling_rates)
210 m_sampling_counts[pair.first] = 0;
211 } else {
212 m_sampling_variable = nullptr;
213 }
214
215}
216
217
219{
220
221 if (m_sampling_variable == nullptr)
222 return 1.0;
223
224 long target = 0;
225 if (m_sampling_variable->variabletype == Variable::Manager::VariableDataType::c_double) {
226 target = std::lround(std::get<double>(m_sampling_variable->function(nullptr)));
227 } else if (m_sampling_variable->variabletype == Variable::Manager::VariableDataType::c_int) {
228 target = std::lround(std::get<int>(m_sampling_variable->function(nullptr)));
229 } else if (m_sampling_variable->variabletype == Variable::Manager::VariableDataType::c_bool) {
230 target = std::lround(std::get<bool>(m_sampling_variable->function(nullptr)));
231 }
232
233 if (m_sampling_rates.find(target) != m_sampling_rates.end() and m_sampling_rates[target] > 0) {
234 m_sampling_counts[target]++;
235 if (m_sampling_counts[target] % m_sampling_rates[target] != 0)
236 return 0;
237 else {
238 m_sampling_counts[target] = 0;
239 return m_sampling_rates[target];
240 }
241 }
242
243 return 1.0;
244}
245
247{
248 // get counter numbers
249 m_event = m_eventMetaData->getEvent();
250 m_run = m_eventMetaData->getRun();
251 m_experiment = m_eventMetaData->getExperiment();
252 m_production = m_eventMetaData->getProduction();
253
254 if (m_stringWrapper.isValid())
255 m_MCDecayString = m_stringWrapper->getString();
256 else
257 m_MCDecayString = "";
258
259 if (m_storeEventType and m_eventExtraInfo.isValid())
260 m_eventType = m_eventExtraInfo->getEventType();
261 else
262 m_eventType = "";
263
265 m_ncandidates = particlelist->getListSize();
267 if (m_weight > 0) {
268 for (unsigned int iVar = 0; iVar < m_event_functions.size(); iVar++) {
269 if (std::holds_alternative<double>(m_event_functions[iVar](nullptr))) {
270 m_event_valuesDouble[iVar] = std::get<double>(m_event_functions[iVar](nullptr));
271 } else if (std::holds_alternative<int>(m_event_functions[iVar](nullptr))) {
272 m_event_valuesInt[iVar] = std::get<int>(m_event_functions[iVar](nullptr));
273 } else if (std::holds_alternative<bool>(m_event_functions[iVar](nullptr))) {
274 m_event_valuesInt[iVar] = std::get<bool>(m_event_functions[iVar](nullptr));
275 }
276 }
277 for (unsigned int iPart = 0; iPart < m_ncandidates; iPart++) {
278
279 if (iPart >= m_maxCandidates) {
280 B2WARNING("Maximum number of candidates exceeded in VariablesToEventBasedTree module. I will skip additional candidates");
282 break;
283 }
284
285 const Particle* particle = particlelist->getParticle(iPart);
286 for (unsigned int iVar = 0; iVar < m_functions.size(); iVar++) {
287 if (std::holds_alternative<double>(m_functions[iVar](particle))) {
288 m_valuesDouble[iVar][iPart] = std::get<double>(m_functions[iVar](particle));
289 } else if (std::holds_alternative<int>(m_functions[iVar](particle))) {
290 m_valuesInt[iVar][iPart] = std::get<int>(m_functions[iVar](particle));
291 } else if (std::holds_alternative<bool>(m_functions[iVar](particle))) {
292 m_valuesInt[iVar][iPart] = std::get<bool>(m_functions[iVar](particle));
293 }
294 }
295 }
296 m_tree->get().Fill();
297 }
298}
299
301{
303 B2INFO("Writing TTree " << m_treeName);
304 TDirectory::TContext directoryGuard(m_file.get());
305 m_tree->write(m_file.get());
306
307 const bool writeError = m_file->TestBit(TFile::kWriteError);
308 m_file.reset();
309 if (writeError) {
310 B2FATAL("A write error occurred while saving '" << m_fileName << "', please check if enough disk space is available.");
311 }
312 }
313}
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 overridden output file name, or "" if none was set.
static Environment & Instance()
Static method to get a reference to the Environment instance.
static std::string makeROOTCompatible(std::string str)
Remove special characters that ROOT dislikes in branch names, e.g.
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
Module()
Constructor.
Definition Module.cc:30
@ 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:76
static bool isOutputProcess()
Return true if the process is an output process.
static bool parallelProcessingUsed()
Returns true if multiple processes have been spawned, false in single-core mode.
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:180
const Var * getVariable(std::string name)
Get the variable belonging to the given key.
Definition Manager.cc:58
static Manager & Instance()
get singleton instance.
Definition Manager.cc:26
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:559
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:649
Abstract base class for different kinds of events.
STL namespace.
A variable returning a floating-point value for a given Particle.
Definition Manager.h:145