FileWriter.cpp 7.1 KB
Newer Older
1
2
3
4
5
6
7
/**
 * @file
 * @brief Implementation of ROOT data file writer module
 * @copyright Copyright (c) 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 ROOTObjectWriter module of the Allpix Squared project
9
10
 */

11
12
#include "FileWriter.h"

13
14
15
#include <fstream>
#include <string>
#include <utility>
16

17
18
#include <TBranchElement.h>
#include <TClass.h>
19

20
21
22
#include "core/utils/file.h"
#include "core/utils/log.h"
#include "core/utils/type.h"
23

24
#include "objects/Object.hpp"
25

26
using namespace corryvreckan;
27

28
29
30
31
FileWriter::FileWriter(Configuration config, std::vector<std::shared_ptr<Detector>> detectors)
    : Module(std::move(config), std::move(detectors)) {}
/**
 * @note Objects cannot be stored in smart pointers due to internal ROOT logic
32
 */
33
34
35
36
37
38
FileWriter::~FileWriter() {
    // Delete all object pointers
    for(auto& index_data : write_list_) {
        delete index_data.second;
    }
}
39

40
void FileWriter::initialise() {
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
    // Create output file
    output_file_name_ =
        createOutputFile(corryvreckan::add_file_extension(m_config.get<std::string>("file_name", "data"), "root"), true);
    output_file_ = std::make_unique<TFile>(output_file_name_.c_str(), "RECREATE");
    output_file_->cd();

    // Read include and exclude list
    if(m_config.has("include") && m_config.has("exclude")) {
        throw InvalidValueError(m_config, "exclude", "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());
56
    }
57

58
59
    // Create event tree:
    event_tree_ = std::make_unique<TTree>("Event", (std::string("Tree of Events").c_str()));
60
    event_tree_->Bronch("global", "corryvreckan::Event", &event_);
61
62
}

63
StatusCode FileWriter::run(std::shared_ptr<Clipboard> clipboard) {
Simon Spannagel's avatar
Simon Spannagel committed
64

65
66
67
    if(!clipboard->event_defined()) {
        ModuleError("No Clipboard event defined, cannot continue");
    }
Simon Spannagel's avatar
Simon Spannagel committed
68

69
70
    // Read event from clipboard and write to tree:
    event_ = clipboard->get_event().get();
71
    event_tree_->Fill();
72
    write_cnt_++;
73
74
75
76
77
78
79
80
81

    auto data = clipboard->get_all();
    LOG(DEBUG) << "Clipboard has " << data.size() << " different object types.";

    for(auto& block : data) {
        try {
            auto type_idx = block.first;
            auto class_name = corryvreckan::demangle(type_idx.name());
            auto class_name_full = corryvreckan::demangle(type_idx.name(), true);
82
            LOG(TRACE) << "Received objects of type \"" << class_name << "\" in " << block.second.size() << " blocks";
83
84
85
86
87
88
89
90

            // Check if these objects should be stored
            if((!include_.empty() && include_.find(class_name) == include_.end()) ||
               (!exclude_.empty() && exclude_.find(class_name) != exclude_.end())) {
                LOG(TRACE) << "Ignoring object " << corryvreckan::demangle(type_idx.name())
                           << " because it has been excluded or not explicitly included";
                continue;
            }
Simon Spannagel's avatar
Simon Spannagel committed
91

92
93
94
95
96
97
            for(auto& detector_block : block.second) {
                // Get the detector name
                std::string detector_name;
                if(!detector_block.first.empty()) {
                    detector_name = detector_block.first;
                }
Simon Spannagel's avatar
Simon Spannagel committed
98

99
                auto objects = std::static_pointer_cast<ObjectVector>(detector_block.second);
100
                LOG(TRACE) << " - " << detector_name << ": " << objects->size();
Simon Spannagel's avatar
Simon Spannagel committed
101

102
103
104
                // Create a new branch of the correct type if this object has not been received before
                auto index_tuple = std::make_tuple(type_idx, detector_name);
                if(write_list_.find(index_tuple) == write_list_.end()) {
Simon Spannagel's avatar
Simon Spannagel committed
105

106
107
108
                    // Add vector of objects to write to the write list
                    write_list_[index_tuple] = new std::vector<Object*>();
                    auto addr = &write_list_[index_tuple];
Simon Spannagel's avatar
Simon Spannagel committed
109

110
111
                    auto new_tree = (trees_.find(class_name) == trees_.end());
                    if(new_tree) {
112
113
114
115
116
117
                        // Create new tree
                        output_file_->cd();
                        trees_.emplace(
                            class_name,
                            std::make_unique<TTree>(class_name.c_str(), (std::string("Tree of ") + class_name).c_str()));
                    }
Simon Spannagel's avatar
Simon Spannagel committed
118

119
                    std::string branch_name = detector_name.empty() ? "global" : detector_name;
Simon Spannagel's avatar
Simon Spannagel committed
120

121
122
                    trees_[class_name]->Bronch(
                        branch_name.c_str(), (std::string("std::vector<") + class_name_full + "*>").c_str(), addr);
123

124
125
126
127
128
129
130
131
132
133
134
135
                    if(new_tree) {
                        LOG(DEBUG) << "Pre-filling new tree of " << class_name << " with " << last_event_ << " empty events";
                        for(unsigned int i = 0; i < last_event_; ++i) {
                            trees_[class_name]->Fill();
                        }
                    } else {
                        LOG(DEBUG) << "Pre-filling new branch " << branch_name << " of " << class_name << " with "
                                   << last_event_ << " empty events";
                        auto* branch = trees_[class_name]->GetBranch(branch_name.c_str());
                        for(unsigned int i = 0; i < last_event_; ++i) {
                            branch->Fill();
                        }
136
                    }
Simon Spannagel's avatar
Simon Spannagel committed
137
138
                }

139
140
141
142
143
                // Fill the branch vector
                for(auto& object : *objects) {
                    ++write_cnt_;
                    write_list_[index_tuple]->push_back(object);
                }
Simon Spannagel's avatar
Simon Spannagel committed
144
            }
145
146
147
        } catch(...) {
            LOG(WARNING) << "Cannot process object of type" << corryvreckan::demangle(block.first.name());
            return StatusCode::NoData;
148
149
150
        }
    }

151
152
153
    LOG(TRACE) << "Writing new objects to tree";
    output_file_->cd();

154
155
    last_event_++;

156
157
158
159
160
161
162
163
164
    // Fill the tree with the current received messages
    for(auto& tree : trees_) {
        tree.second->Fill();
    }

    // Clear the current message list
    for(auto& index_data : write_list_) {
        index_data.second->clear();
    }
165

166
    return StatusCode::Success;
167
168
}

Simon Spannagel's avatar
Simon Spannagel committed
169
void FileWriter::finalise() {
170
171
    LOG(TRACE) << "Writing objects to file";
    output_file_->cd();
172

173
174
175
176
    int branch_count = 0;
    for(auto& tree : trees_) {
        // Update statistics
        branch_count += tree.second->GetListOfBranches()->GetEntries();
177
    }
178
    branch_count += event_tree_->GetListOfBranches()->GetEntries();
179

180
181
182
183
184
185
    // Finish writing to output file
    output_file_->Write();

    // Print statistics
    LOG(STATUS) << "Wrote " << write_cnt_ << " objects to " << branch_count << " branches in file:" << std::endl
                << output_file_name_;
186
}