FileReader.cpp 10.6 KB
Newer Older
Simon Spannagel's avatar
Simon Spannagel committed
1
2
3
4
5
6
7
/**
 * @file
 * @brief Implementation of ROOT data file reader module
 * @copyright Copyright (c) 2017-2019 CERN and the Corryvreckan authors.
 * This software is distributed under the terms of the MIT License, copied verbatim in the file "LICENSE.md".
 * In applying this license, CERN does not waive the privileges and immunities granted to it by virtue of its status as an
 * Intergovernmental Organization or submit itself to any jurisdiction.
8
 * @remarks The implementation of this module is based on the ROOTObjectReader module of the Allpix Squared project
Simon Spannagel's avatar
Simon Spannagel committed
9
10
 */

11
12
#include "FileReader.h"

Simon Spannagel's avatar
Simon Spannagel committed
13
14
15
#include <climits>
#include <string>
#include <utility>
16

Simon Spannagel's avatar
Simon Spannagel committed
17
18
19
20
21
#include <TBranch.h>
#include <TKey.h>
#include <TObjArray.h>
#include <TProcessID.h>
#include <TTree.h>
22

Simon Spannagel's avatar
Simon Spannagel committed
23
24
25
#include "core/utils/file.h"
#include "core/utils/log.h"
#include "core/utils/type.h"
26

Simon Spannagel's avatar
Simon Spannagel committed
27
28
#include "objects/Object.hpp"
#include "objects/objects.h"
29

Simon Spannagel's avatar
Simon Spannagel committed
30
using namespace corryvreckan;
31

Simon Spannagel's avatar
Simon Spannagel committed
32
33
FileReader::FileReader(Configuration config, std::vector<std::shared_ptr<Detector>> detectors)
    : Module(std::move(config), std::move(detectors)) {}
34

Simon Spannagel's avatar
Simon Spannagel committed
35
36
37
38
39
40
/**
 * @note Objects cannot be stored in smart pointers due to internal ROOT logic
 */
FileReader::~FileReader() {
    for(auto object_inf : object_info_array_) {
        delete object_inf.objects;
41
    }
Simon Spannagel's avatar
Simon Spannagel committed
42
}
43

Simon Spannagel's avatar
Simon Spannagel committed
44
45
46
47
48
/**
 * Adds lambda function map to convert a vector of generic objects to a templated vector of objects containing this
 * particular type of object from its typeid.
 */
template <typename T> static void add_creator(FileReader::ObjectCreatorMap& map) {
49
    map[typeid(T)] = [&](std::vector<Object*> objects, std::string detector, std::shared_ptr<Clipboard> clipboard) {
Simon Spannagel's avatar
Simon Spannagel committed
50
        std::vector<T*> data;
Simon Spannagel's avatar
Simon Spannagel committed
51
52
        // Copy the objects to data vector
        for(auto& object : objects) {
53
            data.push_back(new T(*static_cast<T*>(object)));
Simon Spannagel's avatar
Simon Spannagel committed
54
        }
55

Simon Spannagel's avatar
Simon Spannagel committed
56
57
58
59
        // Fix the object references (NOTE: we do this after insertion as otherwise the objects could have been relocated)
        for(size_t i = 0; i < objects.size(); ++i) {
            auto& prev_obj = *objects[i];
            auto& new_obj = data[i];
60

Simon Spannagel's avatar
Simon Spannagel committed
61
62
            // Only update the reference for objects that have been referenced before
            if(prev_obj.TestBit(kIsReferenced)) {
Simon Spannagel's avatar
Simon Spannagel committed
63
                auto pid = TProcessID::GetProcessWithUID(new_obj);
Simon Spannagel's avatar
Simon Spannagel committed
64
65
                if(pid->GetObjectWithID(prev_obj.GetUniqueID()) != &prev_obj) {
                    LOG(ERROR) << "Duplicate object IDs, cannot correctly resolve previous history!";
66
                }
Simon Spannagel's avatar
Simon Spannagel committed
67
                prev_obj.ResetBit(kIsReferenced);
Simon Spannagel's avatar
Simon Spannagel committed
68
69
                new_obj->SetBit(kIsReferenced);
                pid->PutObjectWithID(new_obj);
Simon Spannagel's avatar
Simon Spannagel committed
70
71
            }
        }
72

73
74
75
76
77
78
        // Store the ojects on the clipboard:
        if(detector.empty()) {
            clipboard->put(std::make_shared<std::vector<T*>>(std::move(data)));
        } else {
            clipboard->put(std::make_shared<std::vector<T*>>(std::move(data)), detector);
        }
Simon Spannagel's avatar
Simon Spannagel committed
79
    };
80
81
}

Simon Spannagel's avatar
Simon Spannagel committed
82
83
84
85
86
87
88
89
90
/**
 * Uses SFINAE trick to call the add_creator function for all template arguments of a container class. Used to add creators
 * for every object in a tuple of objects.
 */
template <template <typename...> class T, typename... Args>
static void gen_creator_map_from_tag(FileReader::ObjectCreatorMap& map, type_tag<T<Args...>>) {
    std::initializer_list<int> value{(add_creator<Args>(map), 0)...};
    (void)value;
}
Simon Spannagel's avatar
Simon Spannagel committed
91

Simon Spannagel's avatar
Simon Spannagel committed
92
93
94
95
96
97
98
99
/**
 * Wrapper function to make the SFINAE trick in \ref gen_creator_map_from_tag work.
 */
template <typename T> static FileReader::ObjectCreatorMap gen_creator_map() {
    FileReader::ObjectCreatorMap ret_map;
    gen_creator_map_from_tag(ret_map, type_tag<T>());
    return ret_map;
}
Simon Spannagel's avatar
Simon Spannagel committed
100

Simon Spannagel's avatar
Simon Spannagel committed
101
102
103
104
105
106
107
108
109
110
111
112
void FileReader::initialise() {
    // Read include and exclude list
    if(m_config.has("include") && m_config.has("exclude")) {
        throw InvalidCombinationError(
            m_config, {"exclude", "include"}, "include and exclude parameter are mutually exclusive");
    } else if(m_config.has("include")) {
        auto inc_arr = m_config.getArray<std::string>("include");
        include_.insert(inc_arr.begin(), inc_arr.end());
    } else if(m_config.has("exclude")) {
        auto exc_arr = m_config.getArray<std::string>("exclude");
        exclude_.insert(exc_arr.begin(), exc_arr.end());
    }
Simon Spannagel's avatar
Simon Spannagel committed
113

Simon Spannagel's avatar
Simon Spannagel committed
114
115
    // Initialize the call map from the tuple of available objects
    object_creator_map_ = gen_creator_map<corryvreckan::OBJECTS>();
Simon Spannagel's avatar
Simon Spannagel committed
116

Simon Spannagel's avatar
Simon Spannagel committed
117
118
    // Open the file with the objects
    input_file_ = std::make_unique<TFile>(m_config.getPath("file_name", true).c_str());
Simon Spannagel's avatar
Simon Spannagel committed
119

Simon Spannagel's avatar
Simon Spannagel committed
120
121
122
    // Read all the trees in the file
    TList* keys = input_file_->GetListOfKeys();
    std::set<std::string> tree_names;
Simon Spannagel's avatar
Simon Spannagel committed
123

Simon Spannagel's avatar
Simon Spannagel committed
124
125
126
127
    for(auto&& object : *keys) {
        auto& key = dynamic_cast<TKey&>(*object);
        if(std::string(key.GetClassName()) == "TTree") {
            auto tree = static_cast<TTree*>(key.ReadObjectAny(nullptr));
Simon Spannagel's avatar
Simon Spannagel committed
128

129
130
131
132
133
134
            if(tree->GetName() == std::string("Event")) {
                LOG(DEBUG) << "Found Event object tree";
                event_tree_ = tree;
                continue;
            }

Simon Spannagel's avatar
Simon Spannagel committed
135
136
137
138
139
140
141
            // Check if a version of this tree has already been read
            if(tree_names.find(tree->GetName()) != tree_names.end()) {
                LOG(TRACE) << "Skipping copy of tree with name " << tree->GetName()
                           << " because one with identical name has already been processed";
                continue;
            }
            tree_names.insert(tree->GetName());
Simon Spannagel's avatar
Simon Spannagel committed
142

Simon Spannagel's avatar
Simon Spannagel committed
143
144
145
146
147
148
149
            // Check if this tree should be used
            if((!include_.empty() && include_.find(tree->GetName()) == include_.end()) ||
               (!exclude_.empty() && exclude_.find(tree->GetName()) != exclude_.end())) {
                LOG(TRACE) << "Ignoring tree with " << tree->GetName()
                           << " objects because it has been excluded or not explicitly included";
                continue;
            }
Simon Spannagel's avatar
Simon Spannagel committed
150

Simon Spannagel's avatar
Simon Spannagel committed
151
152
153
154
155
            trees_.push_back(tree);
        }
    }

    if(trees_.empty()) {
156
        throw ModuleError("Provided ROOT file does not contain any trees, module cannot read any data");
Simon Spannagel's avatar
Simon Spannagel committed
157
    }
Simon Spannagel's avatar
Simon Spannagel committed
158

159
160
161
162
163
164
165
166
167
168
169
170
    // Prepare event branch:
    if(event_tree_ == nullptr) {
        throw ModuleError("Could not find \"Event\" tree to read event definitions from");
    }

    // Loop over the list of branches and create the set of receiver objects
    TObjArray* event_branches = event_tree_->GetListOfBranches();
    if(event_branches->GetEntries() != 1) {
        throw ModuleError("\"Event\" tree invalid, cannot read event data from file");
    }

    auto* event_branch = static_cast<TBranch*>(event_branches->At(0));
171
172
    event_ = new Event();
    event_branch->SetAddress(&event_);
173

Simon Spannagel's avatar
Simon Spannagel committed
174
175
176
177
    // Loop over all found trees
    for(auto& tree : trees_) {
        // Loop over the list of branches and create the set of receiver objects
        TObjArray* branches = tree->GetListOfBranches();
178
        LOG(TRACE) << "Tree \"" << tree->GetName() << "\" has " << branches->GetEntries() << " branches";
Simon Spannagel's avatar
Simon Spannagel committed
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
        for(int i = 0; i < branches->GetEntries(); i++) {
            auto* branch = static_cast<TBranch*>(branches->At(i));

            // Add a new vector of objects and bind it to the branch
            object_info object_inf;
            object_inf.objects = new std::vector<Object*>;
            object_info_array_.emplace_back(object_inf);
            branch->SetAddress(&(object_info_array_.back().objects));

            // Check tree structure and if object type matches name
            auto split_type = corryvreckan::split<std::string>(branch->GetClassName(), "<>");
            if(split_type.size() != 2 || split_type[1].size() <= 2) {
                throw ModuleError("Tree is malformed and cannot be used for creating objetcs");
            }
            std::string class_name = split_type[1].substr(0, split_type[1].size() - 1);
            std::string corry_namespace = "corryvreckan::";
            size_t corry_idx = class_name.find(corry_namespace);
            if(corry_idx != std::string::npos) {
                class_name.replace(corry_idx, corry_namespace.size(), "");
            }
            if(class_name != tree->GetName()) {
                throw ModuleError("Tree contains objects of the wrong type");
Simon Spannagel's avatar
Simon Spannagel committed
201
202
            }

Simon Spannagel's avatar
Simon Spannagel committed
203
204
205
206
207
208
209
210
211
            std::string branch_name = branch->GetName();
            if(branch_name != "global") {
                // Check if detector is registered by fetching it:
                auto detector = get_detector(branch_name);
                object_info_array_.back().detector = branch_name;
            }
        }
    }
}
Simon Spannagel's avatar
Simon Spannagel committed
212

Simon Spannagel's avatar
Simon Spannagel committed
213
StatusCode FileReader::run(std::shared_ptr<Clipboard> clipboard) {
Simon Spannagel's avatar
Simon Spannagel committed
214

Simon Spannagel's avatar
Simon Spannagel committed
215
216
217
    if(clipboard->event_defined()) {
        ModuleError("Clipboard event already defined, cannot continue");
    }
Simon Spannagel's avatar
Simon Spannagel committed
218

219
220
    // Read event object from tree and store it on the clipboard:
    event_tree_->GetEntry(event_num_);
221
    clipboard->put_event(std::make_shared<Event>(*event_));
222
    read_cnt_++;
223

Simon Spannagel's avatar
Simon Spannagel committed
224
    for(auto& tree : trees_) {
225
        LOG(TRACE) << "Reading tree \"" << tree->GetName() << "\"";
Simon Spannagel's avatar
Simon Spannagel committed
226
227
        tree->GetEntry(event_num_);
    }
Simon Spannagel's avatar
Simon Spannagel committed
228

Simon Spannagel's avatar
Simon Spannagel committed
229
230
231
232
233
234
235
236
    LOG(TRACE) << "Putting stored objects on the clipboard";

    // Loop through all branches
    for(auto object_inf : object_info_array_) {
        auto objects = object_inf.objects;
        // Skip empty objects in current event
        if(objects->empty()) {
            continue;
237
        }
238

Simon Spannagel's avatar
Simon Spannagel committed
239
240
241
242
        // Check if a pointer to a dispatcher method exist
        auto first_object = (*objects)[0];
        auto iter = object_creator_map_.find(typeid(*first_object));
        if(iter == object_creator_map_.end()) {
243
244
            LOG(INFO) << "Cannot create object with type " << corryvreckan::demangle(typeid(*first_object).name())
                      << " because it not registered for clipboard storage";
Simon Spannagel's avatar
Simon Spannagel committed
245
246
247
            continue;
        }

248
249
        LOG(TRACE) << "- " << objects->size() << " " << corryvreckan::demangle(typeid(*first_object).name()) << ", detector "
                   << object_inf.detector;
250
251
        // Create the object vector and store it on the clipboard:
        iter->second(*objects, object_inf.detector, clipboard);
Simon Spannagel's avatar
Simon Spannagel committed
252
253
254

        // Update statistics
        read_cnt_ += objects->size();
Simon Spannagel's avatar
Simon Spannagel committed
255
    }
256

Simon Spannagel's avatar
Simon Spannagel committed
257
    event_num_++;
258

259
260
261
262
263
264
    if(event_num_ >= event_tree_->GetEntries()) {
        LOG(INFO) << "Requesting end of run because TTree only contains data for " << (event_num_ + 1) << " events";
        return StatusCode::EndRun;
    } else {
        return StatusCode::Success;
    }
265
266
}

Simon Spannagel's avatar
Simon Spannagel committed
267
void FileReader::finalise() {
Simon Spannagel's avatar
Simon Spannagel committed
268
269
270
271
    int branch_count = 0;
    for(auto& tree : trees_) {
        branch_count += tree->GetListOfBranches()->GetEntries();
    }
272
    branch_count += event_tree_->GetListOfBranches()->GetEntries();
273

Simon Spannagel's avatar
Simon Spannagel committed
274
275
    // Print statistics
    LOG(INFO) << "Read " << read_cnt_ << " objects from " << branch_count << " branches";
276

Simon Spannagel's avatar
Simon Spannagel committed
277
278
    // Close the file
    input_file_->Close();
279
}