MiniTreeDataSet.hpp 7.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185
  1. /**
  2. * @file
  3. * @author Caleb Fangmeier <caleb@fangmeier.tech>
  4. * @version 0.1
  5. *
  6. * @section LICENSE
  7. *
  8. *
  9. * MIT License
  10. *
  11. * Copyright (c) 2017 Caleb Fangmeier
  12. *
  13. * Permission is hereby granted, free of charge, to any person obtaining a copy
  14. * of this software and associated documentation files (the "Software"), to deal
  15. * in the Software without restriction, including without limitation the rights
  16. * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
  17. * copies of the Software, and to permit persons to whom the Software is
  18. * furnished to do so, subject to the following conditions:
  19. *
  20. * The above copyright notice and this permission notice shall be included in all
  21. * copies or substantial portions of the Software.
  22. *
  23. * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
  24. * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
  25. * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
  26. * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
  27. * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
  28. * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
  29. * SOFTWARE.
  30. */
  31. #ifndef minitreedataset_h
  32. #define minitreedataset_h
  33. #include <string>
  34. #include <tuple>
  35. #include "filval/filval.hpp"
  36. #include "filval_root/filval_root.hpp"
  37. #include "MiniTree.hpp"
  38. using namespace std;
  39. using namespace fv;
  40. using namespace fv::root;
  41. class MiniTreeDataSet : public DataSet,
  42. public MiniTree{
  43. private:
  44. // Maps filenames to data category. Either "signal" or "background"
  45. std::map<std::string,std::string> input_categories;
  46. // Maps filenames to data label, eg. "TTTT", or "TTZ"
  47. std::map<std::string,std::string> input_labels;
  48. std::string output_filename;
  49. TFile* output_file;
  50. long next_entry;
  51. long nentries;
  52. bool load_next(){
  53. if (next_entry >= nentries) return false;
  54. fChain->GetEntry(next_entry);
  55. ++next_entry;
  56. return true;
  57. }
  58. int get_events(){
  59. return nentries;
  60. }
  61. int get_current_event(){
  62. return next_entry-1;
  63. }
  64. void save_event_count_and_xsection(){
  65. std::map<std::string,int> event_counts;
  66. std::map<std::string,float> xsecs;
  67. string fname, label;
  68. for(auto& p : input_labels){
  69. std::tie(fname, label) = p;
  70. TFile f(fname.c_str());
  71. TH1D* count = (TH1D*)f.Get("Count");
  72. event_counts[label] = (int)count->GetBinContent(1);
  73. TTree* tree = (TTree*)f.Get("tree");
  74. TBranch* b = tree->GetBranch("xsec");
  75. float xsec;
  76. b->SetAddress(&xsec);
  77. b->GetEntry(1);
  78. xsecs[label] = xsec;
  79. }
  80. output_file->cd();
  81. gDirectory->WriteObjectAny(&event_counts, "std::map<std::string,int>", "_event_counts");
  82. gDirectory->WriteObjectAny(&xsecs, "std::map<std::string,float>", "_xsecs");
  83. }
  84. public:
  85. MiniTreeDataSet(const std::string& output_filename, const std::string input_filename, const std::string data_label)
  86. :DataSet(),
  87. input_categories({ {input_filename, "signal"} }),
  88. input_labels({ {input_filename, data_label} }),
  89. output_filename(output_filename),
  90. next_entry(0) {
  91. TChain* chain = new TChain("tree");
  92. chain->Add(input_filename.c_str());
  93. Init(chain);
  94. nentries = fChain->GetEntries();
  95. output_file = TFile::Open(output_filename.c_str(), "RECREATE");
  96. this->fChain->SetBranchStatus("*", false);
  97. }
  98. MiniTreeDataSet(const std::string& output_filename, const std::map<std::string,std::string>& filenames_with_labels)
  99. :DataSet(),
  100. input_categories(filenames_with_labels),
  101. output_filename(output_filename),
  102. next_entry(0) {
  103. TChain* chain = new TChain("tree");
  104. for(auto& p : filenames_with_labels){
  105. std::string filename;
  106. std::tie(filename, std::ignore) = p;
  107. chain->Add(filename.c_str());
  108. }
  109. Init(chain);
  110. nentries = fChain->GetEntries();
  111. output_file = TFile::Open(output_filename.c_str(), "RECREATE");
  112. this->fChain->SetBranchStatus("*", false);
  113. }
  114. ~MiniTreeDataSet(){
  115. save_event_count_and_xsection();
  116. output_file->Close();
  117. }
  118. const std::string& get_current_event_label() const{
  119. TFile* file = fChain->GetFile();
  120. std::string filename = file->GetName();
  121. return input_categories.at(filename);
  122. }
  123. template <typename T>
  124. Value<T>* track_branch(const std::string& bname){
  125. TBranch* branch = fChain->GetBranch(bname.c_str());
  126. if (branch == nullptr){
  127. CRITICAL("Branch: " << bname << " does not exist in input tree.", -1);
  128. }
  129. T* bref = (T*) branch->GetAddress();
  130. fChain->SetBranchStatus(bname.c_str(), true);
  131. INFO("Registering branch \"" << bname
  132. << "\" with address " << bref
  133. << " and type " << typeid(bref).name());
  134. return new ObservedValue<T>(bname, bref);
  135. }
  136. template <typename T>
  137. Value<T*>* track_branch_ptr(const std::string& bname){
  138. TBranch* branch = fChain->GetBranch(bname.c_str());
  139. if (branch == nullptr){
  140. CRITICAL("Branch: " << bname << " does not exist in input tree.", -1);
  141. }
  142. T* bref = (T*) branch->GetAddress();
  143. fChain->SetBranchStatus(bname.c_str(), true);
  144. INFO("Registering pointer branch \"" << bname
  145. << "\" with address " << bref
  146. << " and type " << typeid(bref).name());
  147. return new PointerValue<T>(bname, bref);
  148. }
  149. template <typename T>
  150. decltype(auto) track_branch_vec(const std::string& size_bname, const std::string& value_bname){
  151. track_branch_ptr<T>(value_bname);
  152. return wrapper_vector<T>(lookup<int>(size_bname), lookup<T*>(value_bname), value_bname);
  153. }
  154. void save_all(){
  155. output_file->cd();
  156. // Save the value names for each container to enable looking up
  157. // what values are plotted
  158. std::map<string,string> value_lookup = this->get_container_name_value_map();
  159. gDirectory->WriteObjectAny(&value_lookup, "std::map<std::string,std::string>", "_value_lookup");
  160. std::map<string,string> fn_impl_lookup = this->get_function_name_impl_map();
  161. gDirectory->WriteObjectAny(&fn_impl_lookup, "std::map<std::string,std::string>", "_function_impl_lookup");
  162. for(auto container : containers){
  163. container.second->save_as("outfile", SaveOption::ROOT);
  164. }
  165. }
  166. };
  167. #endif // minitreedataset_h