Belle II Software  light-2212-foldex
VariablesToNtupleModule.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/VariablesToNtuple/VariablesToNtupleModule.h>
10 
11 // analysis
12 #include <analysis/dataobjects/ParticleList.h>
13 #include <analysis/VariableManager/Manager.h>
14 #include <analysis/VariableManager/Utility.h>
15 #include <analysis/dataobjects/StringWrapper.h>
16 
17 // framework
18 #include <framework/logging/Logger.h>
19 #include <framework/pcore/ProcHandler.h>
20 #include <framework/core/ModuleParam.templateDetails.h>
21 
22 // framework - root utilities
23 #include <framework/utilities/MakeROOTCompatible.h>
24 #include <framework/utilities/RootFileCreationManager.h>
25 
26 #include <cmath>
27 
28 using namespace std;
29 using namespace Belle2;
30 
31 // Register module in the framework
32 REG_MODULE(VariablesToNtuple);
33 
34 
35 VariablesToNtupleModule::VariablesToNtupleModule() :
36  Module(), m_tree("", DataStore::c_Persistent)
37 {
38  //Set module properties
39  setDescription("Calculate variables specified by the user for a given ParticleList and save them into a TNtuple. The TNtuple is candidate-based, meaning that the variables of each candidate are saved into separate rows.");
41 
42  vector<string> emptylist;
43  addParam("particleList", m_particleList,
44  "Name of particle list with reconstructed particles. If no list is provided the variables are saved once per event (only possible for event-type variables)",
45  std::string(""));
46  addParam("variables", m_variables,
47  "List of variables (or collections) to save. Variables are taken from Variable::Manager, and are identical to those available to e.g. ParticleSelector.",
48  emptylist);
49 
50  addParam("fileName", m_fileName, "Name of ROOT file for output.", string("VariablesToNtuple.root"));
51  addParam("treeName", m_treeName, "Name of the NTuple in the saved file.", string("ntuple"));
52  addParam("basketSize", m_basketsize, "Size of baskets in Output NTuple in bytes.", 1600);
53 
54  std::tuple<std::string, std::map<int, unsigned int>> default_sampling{"", {}};
55  addParam("sampling", m_sampling,
56  "Tuple of variable name and a map of integer values and inverse sampling rate. E.g. (signal, {1: 0, 0:10}) selects all signal candidates and every 10th background candidate.",
57  default_sampling);
58 
59  addParam("signalSideParticleList", m_signalSideParticleList,
60  "Name of signal-side particle list to store the index of the signal-side particle when one calls the module in a for_each loop over the RestOfEvent",
61  std::string(""));
62 
63 }
64 
66 {
67  m_eventMetaData.isRequired();
68  if (not m_particleList.empty())
70 
71  // Initializing the output root file
72  if (m_fileName.empty()) {
73  B2FATAL("Output root file name is not set. Please set a valid root output file name (\"fileName\" module parameter).");
74  }
75  // See if there is already a file in which case add a new tree to it ...
76  // otherwise create a new file (all handled by framework)
78  if (!m_file) {
79  B2ERROR("Could not create file \"" << m_fileName <<
80  "\". Please set a valid root output file name (\"fileName\" module parameter).");
81  return;
82  }
83 
84  TDirectory::TContext directoryGuard(m_file.get());
85 
86  // check if TTree with that name already exists
87  if (m_file->Get(m_treeName.c_str())) {
88  B2FATAL("Tree with the name \"" << m_treeName
89  << "\" already exists in the file \"" << m_fileName << "\"\n"
90  << "\nYou probably want to either set the output fileName or the treeName to something else:\n\n"
91  << " from modularAnalysis import variablesToNtuple\n"
92  << " variablesToNtuple('pi+:all', ['p'], treename='pions', filename='variablesToNtuple.root')\n"
93  << " variablesToNtuple('gamma:all', ['p'], treename='photons', filename='variablesToNtuple.root') # two trees, same file\n"
94  << "\n == Or ==\n"
95  << " from modularAnalysis import variablesToNtuple\n"
96  << " variablesToNtuple('pi+:all', ['p'], filename='pions.root')\n"
97  << " variablesToNtuple('gamma:all', ['p'], filename='photons.root') # two files\n"
98  );
99  return;
100  }
101 
102  // set up tree and register it in the datastore
103  m_tree.registerInDataStore(m_fileName + m_treeName, DataStore::c_DontWriteOut);
104  m_tree.construct(m_treeName.c_str(), "");
105  m_tree->get().SetCacheSize(100000);
106 
107  // declare counter branches - pass through variable list, remove counters added by user
108  m_tree->get().Branch("__experiment__", &m_experiment, "__experiment__/I");
109  m_tree->get().Branch("__run__", &m_run, "__run__/I");
110  m_tree->get().Branch("__event__", &m_event, "__event__/i");
111  m_tree->get().Branch("__production__", &m_production, "__production__/I");
112  if (not m_particleList.empty()) {
113  m_tree->get().Branch("__candidate__", &m_candidate, "__candidate__/I");
114  m_tree->get().Branch("__ncandidates__", &m_ncandidates, "__ncandidates__/I");
115  }
116 
117  if (not m_signalSideParticleList.empty()) {
119  m_tree->get().Branch("__signalSideCandidate__", &m_signalSideCandidate, "__signalSideCandidate__/I");
120  m_tree->get().Branch("__nSignalSideCandidates__", &m_nSignalSideCandidates, "__nSignalSideCandidates__/I");
121  if (not m_roe.isOptional("RestOfEvent")) {
122  B2WARNING("The signalSideParticleList is set outside of a for_each loop over the RestOfEvent. "
123  << "__signalSideCandidates__ and __nSignalSideCandidate__ will be always -1 and 0, respectively.");
124  }
125  }
126 
127  if (m_stringWrapper.isOptional("MCDecayString"))
128  m_tree->get().Branch("__MCDecayString__", &m_MCDecayString);
129 
130  for (const auto& variable : m_variables)
131  if (Variable::isCounterVariable(variable)) {
132  B2WARNING("The counter '" << variable
133  << "' is handled automatically by VariablesToNtuple, you don't need to add it.");
134  }
135 
136  // declare branches and get the variable strings
138  // remove duplicates from list of variables but keep the previous order
139  unordered_set<string> seen;
140  auto newEnd = remove_if(m_variables.begin(), m_variables.end(), [&seen](const string & varStr) {
141  if (seen.find(varStr) != std::end(seen)) return true;
142  seen.insert(varStr);
143  return false;
144  });
145  m_variables.erase(newEnd, m_variables.end());
146 
147  m_branchAddressesDouble.resize(m_variables.size() + 1);
148  m_branchAddressesInt.resize(m_variables.size() + 1);
149  m_tree->get().Branch("__weight__", &m_branchAddressesDouble[0], "__weight__/D");
150  size_t enumerate = 1;
151  for (const string& varStr : m_variables) {
152  string branchName = MakeROOTCompatible::makeROOTCompatible(varStr);
153 
154  // Check for deprecated variables
156 
157  // also collection function pointers
159  if (!var) {
160  B2ERROR("Variable '" << varStr << "' is not available in Variable::Manager!");
161  } else {
162  if (m_particleList.empty() && var->description.find("[Eventbased]") == string::npos) {
163  B2ERROR("Variable '" << varStr << "' is not an event-based variable, "
164  "but you are using VariablesToNtuple without a decay string, i.e. in the event-wise mode.\n"
165  "If you have created an event-based alias you can wrap your alias with `eventCached` to "
166  "declare it as event based, which avoids this error.\n\n"
167  "vm.addAlias('myAliasName', 'eventCached(myAlias)')");
168  continue;
169  }
170  if (var->variabletype == Variable::Manager::VariableDataType::c_double) {
171  m_tree->get().Branch(branchName.c_str(), &m_branchAddressesDouble[enumerate], (branchName + "/D").c_str());
172  } else if (var->variabletype == Variable::Manager::VariableDataType::c_int) {
173  m_tree->get().Branch(branchName.c_str(), &m_branchAddressesInt[enumerate], (branchName + "/I").c_str());
174  } else if (var->variabletype == Variable::Manager::VariableDataType::c_bool) {
175  m_tree->get().Branch(branchName.c_str(), &m_branchAddressesInt[enumerate], (branchName + "/O").c_str());
176  }
177  m_functions.push_back(std::make_pair(var->function, var->variabletype));
178  }
179  enumerate++;
180  }
181  m_tree->get().SetBasketSize("*", m_basketsize);
182 
183  m_sampling_name = std::get<0>(m_sampling);
184  m_sampling_rates = std::get<1>(m_sampling);
185 
186  if (m_sampling_name != "") {
188  if (m_sampling_variable == nullptr) {
189  B2FATAL("Couldn't find sample variable " << m_sampling_name << " via the Variable::Manager. Check the name!");
190  }
191  for (const auto& pair : m_sampling_rates)
192  m_sampling_counts[pair.first] = 0;
193  } else {
194  m_sampling_variable = nullptr;
195  }
196 }
197 
198 
200 {
201  if (m_sampling_variable == nullptr)
202  return 1.0;
203 
204  long target = 0;
205  if (m_sampling_variable->variabletype == Variable::Manager::VariableDataType::c_double) {
206  target = std::lround(std::get<double>(m_sampling_variable->function(particle)));
207  } else if (m_sampling_variable->variabletype == Variable::Manager::VariableDataType::c_int) {
208  target = std::lround(std::get<int>(m_sampling_variable->function(particle)));
209  } else if (m_sampling_variable->variabletype == Variable::Manager::VariableDataType::c_bool) {
210  target = std::lround(std::get<bool>(m_sampling_variable->function(particle)));
211  }
212  if (m_sampling_rates.find(target) != m_sampling_rates.end() and m_sampling_rates[target] > 0) {
213  m_sampling_counts[target]++;
214  if (m_sampling_counts[target] % m_sampling_rates[target] != 0)
215  return 0;
216  else {
217  m_sampling_counts[target] = 0;
218  return m_sampling_rates[target];
219  }
220  }
221  return 1.0;
222 }
223 
225 {
226  m_event = m_eventMetaData->getEvent();
227  m_run = m_eventMetaData->getRun();
228  m_experiment = m_eventMetaData->getExperiment();
229  m_production = m_eventMetaData->getProduction();
230 
231  if (m_stringWrapper.isValid())
232  m_MCDecayString = m_stringWrapper->getString();
233  else
234  m_MCDecayString = "";
235 
236  if (not m_signalSideParticleList.empty()) {
237  if (m_roe.isValid()) {
239  auto signal = m_roe->getRelatedFrom<Particle>();
240  m_signalSideCandidate = signaSideParticleList->getIndex(signal);
241  m_nSignalSideCandidates = signaSideParticleList->getListSize();
242  } else {
245  }
246  }
247 
248  if (m_particleList.empty()) {
250  if (m_branchAddressesDouble[0] > 0) {
251  for (unsigned int iVar = 0; iVar < m_variables.size(); iVar++) {
252  auto var_result = std::get<0>(m_functions[iVar])(nullptr);
253  auto var_type = std::get<1>(m_functions[iVar]);
254  if (std::holds_alternative<double>(var_result)) {
255  if (var_type != Variable::Manager::VariableDataType::c_double)
256  B2WARNING("Wrong registered data type for variable '" + m_variables[iVar] +
257  "'. Expected Variable::Manager::VariableDataType::c_double. Exported data for this variable might be incorrect.");
258  m_branchAddressesDouble[iVar + 1] = std::get<double>(var_result);
259  } else if (std::holds_alternative<int>(var_result)) {
260  if (var_type != Variable::Manager::VariableDataType::c_int)
261  B2WARNING("Wrong registered data type for variable '" + m_variables[iVar] +
262  "'. Expected Variable::Manager::VariableDataType::c_int. Exported data for this variable might be incorrect.");
263  m_branchAddressesInt[iVar + 1] = std::get<int>(var_result);
264  } else if (std::holds_alternative<bool>(var_result)) {
265  if (var_type != Variable::Manager::VariableDataType::c_bool)
266  B2WARNING("Wrong registered data type for variable '" + m_variables[iVar] +
267  "'. Expected Variable::Manager::VariableDataType::c_bool. Exported data for this variable might be incorrect.");
268  m_branchAddressesInt[iVar + 1] = std::get<bool>(var_result);
269  }
270  }
271  m_tree->get().Fill();
272  }
273 
274  } else {
276  m_ncandidates = particlelist->getListSize();
277  for (unsigned int iPart = 0; iPart < m_ncandidates; iPart++) {
278  m_candidate = iPart;
279  const Particle* particle = particlelist->getParticle(iPart);
281  if (m_branchAddressesDouble[0] > 0) {
282  for (unsigned int iVar = 0; iVar < m_variables.size(); iVar++) {
283  auto var_result = std::get<0>(m_functions[iVar])(particle);
284  auto var_type = std::get<1>(m_functions[iVar]);
285  if (std::holds_alternative<double>(var_result)) {
286  if (var_type != Variable::Manager::VariableDataType::c_double)
287  B2WARNING("Wrong registered data type for variable '" + m_variables[iVar] +
288  "'. Expected Variable::Manager::VariableDataType::c_double. Exported data for this variable might be incorrect.");
289  m_branchAddressesDouble[iVar + 1] = std::get<double>(var_result);
290  } else if (std::holds_alternative<int>(var_result)) {
291  if (var_type != Variable::Manager::VariableDataType::c_int)
292  B2WARNING("Wrong registered data type for variable '" + m_variables[iVar] +
293  "'. Expected Variable::Manager::VariableDataType::c_int. Exported data for this variable might be incorrect.");
294  m_branchAddressesInt[iVar + 1] = std::get<int>(var_result);
295  } else if (std::holds_alternative<bool>(var_result)) {
296  if (var_type != Variable::Manager::VariableDataType::c_bool)
297  B2WARNING("Wrong registered data type for variable '" + m_variables[iVar] +
298  "'. Expected Variable::Manager::VariableDataType::c_bool. Exported data for this variable might be incorrect.");
299  m_branchAddressesInt[iVar + 1] = std::get<bool>(var_result);
300  }
301  }
302  m_tree->get().Fill();
303  }
304  }
305  }
306 }
307 
309 {
311  B2INFO("Writing NTuple " << m_treeName);
312  TDirectory::TContext directoryGuard(m_file.get());
313  m_tree->write(m_file.get());
314 
315  const bool writeError = m_file->TestBit(TFile::kWriteError);
316  m_file.reset();
317  if (writeError) {
318  B2FATAL("A write error occurred while saving '" << m_fileName << "', please check if enough disk space is available.");
319  }
320  }
321 }
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
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:74
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:95
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
void checkDeprecatedVariable(const std::string &name)
Check if a variable is deprecated.
Definition: Manager.cc:443
StoreObjPtr< RestOfEvent > m_roe
ROE object.
std::vector< std::string > m_variables
List of variables to save.
virtual void initialize() override
Initialises the module.
std::map< int, unsigned int > m_sampling_rates
Inverse sampling rates.
virtual void event() override
Method called for each event.
unsigned int m_ncandidates
total n candidates
virtual void terminate() override
Write TTree to file, and close file if necessary.
StoreObjPtr< EventMetaData > m_eventMetaData
the event information
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::vector< std::pair< Variable::Manager::FunctionPtr, Variable::Manager::VariableDataType > > m_functions
List of pairs of function pointers and respective data type corresponding to given variables.
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< int > m_branchAddressesInt
Branch addresses of variables of type int (or bool)
int m_basketsize
Size of TBaskets in the output ROOT file in bytes.
int m_production
production ID (to distinguish MC samples)
unsigned int m_nSignalSideCandidates
total n signal-side candidates
StoreObjPtr< StringWrapper > m_stringWrapper
string wrapper storing the MCDecayString
std::string m_particleList
Name of particle list with reconstructed particles.
StoreObjPtr< RootMergeable< TTree > > m_tree
The ROOT TNtuple for output.
std::shared_ptr< TFile > m_file
ROOT file for output.
std::vector< double > m_branchAddressesDouble
Branch addresses of variables of type double.
std::string m_sampling_name
Variable name of sampling variable.
float getInverseSamplingRateWeight(const Particle *particle)
Calculate inverse sampling rate weight.
std::string m_treeName
Name of the TTree.
std::string m_signalSideParticleList
Name of signal-side particle list
const Variable::Manager::Var * m_sampling_variable
Variable Pointer to target variable.
int m_signalSideCandidate
signal-side candidate counter
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:23
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