-
Notifications
You must be signed in to change notification settings - Fork 0
/
TupleReader.cc
72 lines (64 loc) · 2.03 KB
/
TupleReader.cc
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
#include <iostream>
#include <utility>
#include <exception>
#include "TupleReader.h"
//ROOT Headers
#include <TBranch.h>
using std::cout;
using std::endl;
using std::vector;
using std::string;
using std::map;
using std::make_pair;
TupleReader::TupleReader(vector<string> var_types,
map<string, vector<string>> var_names,
string root_filename,
string root_treename)
: var_types_(var_types), var_names_(var_names) {
root_file_ = new TFile(root_filename.c_str());
root_tree_ = (TTree*) root_file_->Get(root_treename.c_str());
current_event_idx_ = 0;
num_events_ = root_tree_->GetEntries();
SetAddresses();
}
TupleReader::~TupleReader() {
root_file_->Close();
delete root_file_;
}
void TupleReader::SetAddresses() {
for (auto t : var_types_) {
if (t != "int" && t != "float")
throw std::invalid_argument("Var_types error: Only supports int and float");
for (auto v : var_names_[t]) {
if (t == "int") {
var_values_int_.insert(make_pair(v, 0));
root_tree_->SetBranchAddress(v.c_str(), &var_values_int_[v]);
}
else { // float
var_values_float_.insert(make_pair(v, 0.));
root_tree_->SetBranchAddress(v.c_str(), &var_values_float_[v]);
}
}
}
}
bool TupleReader::next_record() {
if (current_event_idx_ < num_events_) {
if (root_tree_->GetEntry(current_event_idx_) > 0) {
++current_event_idx_;
return true;
}
else
throw std::out_of_range("GetEntry error: either entry doesn't exist or I/O error.");
}
return false;
}
int TupleReader::GetVarInt(string var_name) const {
if (var_values_int_.count(var_name) == 0)
throw std::domain_error("No such variable, only those in var_names are allowed.");
return var_values_int_.at(var_name);
}
float TupleReader::GetVarFloat(string var_name) const {
if (var_values_float_.count(var_name) == 0)
throw std::domain_error("No such variable, only those in var_names are allowed.");
return var_values_float_.at(var_name);
}