MiniTreeDataSet.hpp 2.3 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374
  1. #ifndef minitreedataset_h
  2. #define minitreedataset_h
  3. #include <string>
  4. #include <tuple>
  5. #include "filval/filval.hpp"
  6. #include "filval_root/filval_root.hpp"
  7. #include "MiniTree.hpp"
  8. using namespace std;
  9. using namespace fv;
  10. using namespace fv::root;
  11. class MiniTreeDataSet : public DataSet,
  12. public MiniTree{
  13. private:
  14. long next_entry;
  15. long nentries;
  16. bool load_next(){
  17. if (next_entry >= nentries) return false;
  18. fChain->GetEntry(next_entry);
  19. ++next_entry;
  20. return true;
  21. }
  22. int get_events(){
  23. return nentries;
  24. }
  25. int get_current_event(){
  26. return next_entry-1;
  27. }
  28. public:
  29. MiniTreeDataSet(TTree *tree)
  30. :MiniTree(tree){
  31. next_entry = 0;
  32. nentries = fChain->GetEntriesFast();
  33. }
  34. template <typename T>
  35. Value<T>* track_branch(const std::string& bname){
  36. TBranch* branch = fChain->GetBranch(bname.c_str());
  37. if (branch == nullptr){
  38. CRITICAL("Branch: " << bname << " does not exist in input tree.", -1);
  39. }
  40. T* bref = (T*) branch->GetAddress();
  41. fChain->SetBranchStatus(bname.c_str(), true);
  42. INFO("Registering branch \"" << bname
  43. << "\" with address " << bref
  44. << " and type " << typeid(bref).name());
  45. return new ObservedValue<T>(bname, bref);
  46. }
  47. template <typename T>
  48. Value<T*>* track_branch_ptr(const std::string& bname){
  49. TBranch* branch = fChain->GetBranch(bname.c_str());
  50. if (branch == nullptr){
  51. CRITICAL("Branch: " << bname << " does not exist in input tree.", -1);
  52. }
  53. T* bref = (T*) branch->GetAddress();
  54. fChain->SetBranchStatus(bname.c_str(), true);
  55. INFO("Registering pointer branch \"" << bname
  56. << "\" with address " << bref
  57. << " and type " << typeid(bref).name());
  58. return new PointerValue<T>(bname, bref);
  59. }
  60. void save_all(){
  61. for(auto container : containers){
  62. container.second->save_as("outfile", SaveOption::ROOT);
  63. }
  64. }
  65. };
  66. #endif // minitreedataset_h