MiniTreeDataSet.hpp 4.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125
  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. std::string input_filename;
  45. std::string output_filename;
  46. TFile* input_file;
  47. TFile* output_file;
  48. long next_entry;
  49. long nentries;
  50. bool load_next(){
  51. if (next_entry >= nentries) return false;
  52. fChain->GetEntry(next_entry);
  53. ++next_entry;
  54. return true;
  55. }
  56. int get_events(){
  57. return nentries;
  58. }
  59. int get_current_event(){
  60. return next_entry-1;
  61. }
  62. public:
  63. MiniTreeDataSet(const std::string& input_filename, const std::string output_filename)
  64. :input_filename(input_filename),
  65. output_filename(output_filename),
  66. next_entry(0) {
  67. input_file = TFile::Open(input_filename.c_str());
  68. Init((TTree*) input_file->Get("tree"));
  69. nentries = fChain->GetEntriesFast();
  70. output_file = TFile::Open(output_filename.c_str(), "RECREATE");
  71. }
  72. ~MiniTreeDataSet(){
  73. input_file->Close();
  74. output_file->Close();
  75. }
  76. template <typename T>
  77. Value<T>* track_branch(const std::string& bname){
  78. TBranch* branch = fChain->GetBranch(bname.c_str());
  79. if (branch == nullptr){
  80. CRITICAL("Branch: " << bname << " does not exist in input tree.", -1);
  81. }
  82. T* bref = (T*) branch->GetAddress();
  83. fChain->SetBranchStatus(bname.c_str(), true);
  84. INFO("Registering branch \"" << bname
  85. << "\" with address " << bref
  86. << " and type " << typeid(bref).name());
  87. return new ObservedValue<T>(bname, bref);
  88. }
  89. template <typename T>
  90. Value<T*>* track_branch_ptr(const std::string& bname){
  91. TBranch* branch = fChain->GetBranch(bname.c_str());
  92. if (branch == nullptr){
  93. CRITICAL("Branch: " << bname << " does not exist in input tree.", -1);
  94. }
  95. T* bref = (T*) branch->GetAddress();
  96. fChain->SetBranchStatus(bname.c_str(), true);
  97. INFO("Registering pointer branch \"" << bname
  98. << "\" with address " << bref
  99. << " and type " << typeid(bref).name());
  100. return new PointerValue<T>(bname, bref);
  101. }
  102. template <typename T>
  103. WrapperVector<T>* track_branch_vec(const std::string& size_bname, const std::string& bname){
  104. track_branch_ptr<T>(bname);
  105. return new WrapperVector<T>(lookup<int>(size_bname), lookup<T*>(bname), bname);
  106. }
  107. void save_all(){
  108. output_file->cd();
  109. for(auto container : containers){
  110. container.second->save_as("outfile", SaveOption::ROOT);
  111. }
  112. }
  113. };
  114. #endif // minitreedataset_h