qaWriter_analysistree2.cxx 4.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125
  1. #include <qaWriter_analysistree2.h>
  2. ClassImp(qaWriter_analysistree2);
  3. qaWriter_analysistree2::qaWriter_analysistree2() : fFile(nullptr),
  4. fTree(nullptr),
  5. fDataHeader(nullptr),
  6. fConfig(nullptr),
  7. fEvent(nullptr),
  8. fParticles(nullptr),
  9. isInit(false),
  10. iB(qaUtility::GetInstance()->error_code),
  11. iPhiRp(qaUtility::GetInstance()->error_code),
  12. icharge(qaUtility::GetInstance()->error_code),
  13. str_event_branch("Event"),
  14. str_particles_branch("Particles")
  15. {
  16. }
  17. qaWriter_analysistree2::~qaWriter_analysistree2()
  18. {
  19. }
  20. void qaWriter_analysistree2::Init(std::string filename, std::string treename, std::string system, float sqrtSnn)
  21. {
  22. fFile = new TFile(filename.c_str(), "recreate");
  23. fTree = new TTree(treename.c_str(), "AnalysisTree format, model data");
  24. // Set up AnalysisTree data header
  25. fDataHeader = new AnalysisTree::DataHeader;
  26. if (system != "")
  27. fDataHeader->SetSystem(system);
  28. if (sqrtSnn > 0)
  29. fDataHeader->SetBeamMomentum(qaUtility::GetInstance()->GetBeamP(sqrtSnn));
  30. // Set up AnalysisTree configureation
  31. fConfig = new AnalysisTree::Configuration;
  32. // Set up AnalysisTree configuration
  33. AnalysisTree::BranchConfig event_branch(str_event_branch.c_str(), AnalysisTree::DetType::kEventHeader);
  34. AnalysisTree::BranchConfig particles_branch(str_particles_branch.c_str(), AnalysisTree::DetType::kParticle);
  35. event_branch.AddField<float>("B");
  36. event_branch.AddField<float>("PhiRp");
  37. particles_branch.AddField<bool>("is_charged");
  38. auto hasher = std::hash<std::string>();
  39. fConfig->AddBranchConfig(event_branch);
  40. fEvent = new AnalysisTree::EventHeader(Short_t(hasher(event_branch.GetName())));
  41. fEvent->Init(event_branch);
  42. fConfig->AddBranchConfig(particles_branch);
  43. fParticles = new AnalysisTree::Particles(Short_t(hasher(particles_branch.GetName())));
  44. iB = fConfig->GetBranchConfig(str_event_branch).GetFieldId("B");
  45. iPhiRp = fConfig->GetBranchConfig(str_event_branch).GetFieldId("PhiRp");
  46. icharge = fConfig->GetBranchConfig(str_particles_branch).GetFieldId("is_charged");
  47. // Create branches in the output tree
  48. fTree->Branch(str_event_branch.c_str(), "AnalysisTree::EventHeader", &fEvent, 32000, 99);
  49. fTree->Branch(str_particles_branch.c_str(), "AnalysisTree::Particles", &fParticles, 256000, 99);
  50. fConfig->Print();
  51. isInit = true;
  52. }
  53. void qaWriter_analysistree2::WriteEvent(qaEvent *event)
  54. {
  55. if (!event)
  56. return;
  57. if (!isInit)
  58. return;
  59. fParticles->ClearChannels();
  60. fEvent->SetField(float(event->GetB()), iB);
  61. fEvent->SetField(float(event->GetPhiRP()), iPhiRp);
  62. }
  63. void qaWriter_analysistree2::WriteParticle(qaParticle *particle)
  64. {
  65. if (!particle)
  66. return;
  67. if (!isInit)
  68. return;
  69. auto *at_particle = fParticles->AddChannel();
  70. at_particle->Init(fConfig->GetBranchConfig(str_particles_branch));
  71. at_particle->SetMomentum(particle->GetPx(), particle->GetPy(), particle->GetPz());
  72. at_particle->SetPid(particle->GetPdg());
  73. auto particlePDG = (TParticlePDG *)TDatabasePDG::Instance()->GetParticle(particle->GetPdg());
  74. if (!particlePDG)
  75. at_particle->SetMass(qaUtility::GetInstance()->error_code);
  76. if (particlePDG)
  77. at_particle->SetMass(particlePDG->Mass());
  78. bool is_charged = (qaUtility::GetInstance()->GetCharge(particle->GetPdg()) == 0 ||
  79. qaUtility::GetInstance()->GetCharge(particle->GetPdg()) == qaUtility::GetInstance()->error_code)
  80. ? false
  81. : true;
  82. at_particle->SetField(is_charged, icharge);
  83. }
  84. void qaWriter_analysistree2::WriteTree()
  85. {
  86. if (!isInit) return;
  87. fFile->cd();
  88. fTree->Print();
  89. fTree->Write();
  90. fDataHeader->Write("DataHeader");
  91. fConfig->Write("Configuration");
  92. fFile->Close();
  93. }
  94. void qaWriter_analysistree2::FillTree()
  95. {
  96. fTree->Fill();
  97. }