Belle II Software development
ConcreteVariablesToNtuplePersistenceManager.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 <reconstruction/persistenceManager/ConcreteVariablesToNtuplePersistenceManager.h>
10
11#include <framework/utilities/MakeROOTCompatible.h>
12#include <framework/utilities/RootFileCreationManager.h>
13#include <framework/pcore/ProcHandler.h>
14#include <framework/pcore/RootMergeable.h>
15
16namespace {
17 using namespace Belle2::VariablePersistenceManager;
18
19 std::string typedVariableToLeafName(const TypedVariable& variable)
20 {
21 std::map<VariableDataType, std::string> suffix = {
22 {VariableDataType::c_double, "/D"},
23 {VariableDataType::c_int, "/I"},
24 {VariableDataType::c_bool, "/O"}
25 };
26
27 std::string branchName = Belle2::MakeROOTCompatible::makeROOTCompatible(variable.getName());
28 return branchName + suffix[variable.getDataType()];
29 }
30}
31
32namespace Belle2::VariablePersistenceManager {
38
40 const std::string& treeName,
41 Variables& variables)
42 {
43 m_fileName = fileName;
44 m_treeName = treeName;
45 m_variables = variables;
46
49 }
50
51 void ConcreteVariablesToNtuplePersistenceManager::addEntry(const EvaluatedVariables& evaluatedVariables)
52 {
53 for (const auto& [variableName, value] : evaluatedVariables) {
54 std::string branchName = Belle2::MakeROOTCompatible::makeROOTCompatible(variableName.c_str());
55 updateBranch(branchName, value);
56 }
57 m_tree->get().Fill();
58 }
59
61 {
63 B2INFO("Writing Ntuple: " << m_treeName);
64 TDirectory::TContext directoryGuard{m_file.get()};
65 m_tree->write(m_file.get());
66
67 const bool writeError = m_file->TestBit(TFile::kWriteError);
68 m_file.reset();
69 if (writeError) {
70 B2FATAL("A write error occurred while saving '" << m_fileName << "', please check if enough disk space is available.");
71 }
72 }
73 }
74
76 {
77 if (m_fileName.empty()) {
78 B2FATAL("Output root file name is not set.");
79 }
81
82 if (not m_file) {
83 B2FATAL("Could not create file: " << m_fileName << ".");
84 }
85
86 TDirectory::TContext directoryGuard{m_file.get()};
87
88 if (m_file->Get(m_treeName.c_str())) {
89 B2FATAL("A tree with the name: " << m_treeName << "already exists in the file: " << m_fileName << ".");
90 }
91
93 m_tree.construct(m_treeName.c_str(), "");
94 m_tree->get().SetCacheSize(100000);
95 }
96
98 {
99 for (const auto& variable : m_variables) {
100 std::visit([&](const auto & typedVariable) {
101 using T = std::decay_t<decltype(typedVariable)>;
102 if constexpr(std::is_same_v<T, TypedVariable>) {
103 std::string branchName = Belle2::MakeROOTCompatible::makeROOTCompatible(typedVariable.getName().c_str());
104 std::string leafName = typedVariableToLeafName(typedVariable);
105
106 switch (typedVariable.getDataType()) {
107 case VariableDataType::c_double:
108 m_branchesDouble[branchName] = double{};
109 m_tree->get().Branch(branchName.c_str(), &m_branchesDouble[branchName], leafName.c_str());
110 break;
111 case VariableDataType::c_int:
112 m_branchesInt[branchName] = int{};
113 m_tree->get().Branch(branchName.c_str(), &m_branchesInt[branchName], leafName.c_str());
114 break;
115 case VariableDataType::c_bool:
116 m_branchesBool[branchName] = bool{};
117 m_tree->get().Branch(branchName.c_str(), &m_branchesBool[branchName], leafName.c_str());
118 break;
119 default:
120 break;
121 }
122 } else {
123 B2WARNING("Incompatible variable type. Skipping branch registration.");
124 }
125 }, variable);
126 }
127 m_tree->get().SetBasketSize("*", m_basketSize);
128 }
129
131 const VariableType& evaluatedValue)
132 {
133 std::visit([&](auto&& value) {
134 using T = std::decay_t<decltype(value)>;
135
136 if constexpr(std::is_same_v<T, double>) {
137 m_branchesDouble[branchName] = value;
138 } else if constexpr(std::is_same_v<T, int>) {
139 m_branchesInt[branchName] = value;
140 } else if constexpr(std::is_same_v<T, bool>) {
141 m_branchesBool[branchName] = value;
142 }
143 }, evaluatedValue);
144 }
145}
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.
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.
void updateBranch(const std::string &variableName, const VariableType &value)
Updates the branch buffer for a given variable with a new value.
void addEntry(const EvaluatedVariables &evaluatedVariables) override
Adds a single event's worth of variable data to the TTree.
std::map< std::string, double > m_branchesDouble
Maps of variable names to storage for double, int, and bool TTree branches.
void initialize(const std::string &fileName, const std::string &treeName, Variables &variables) override
Initializes the manager by opening a ROOT file and preparing a TTree.
StoreObjPtr< RootMergeable< TTree > > m_tree
A store pointer to the RootMergeable wrapper for the TTree.
A variable with a specified data type.
Definition Types.h:57
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.