123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125 |
- #include <qaWriter_analysistree2.h>
- ClassImp(qaWriter_analysistree2);
- qaWriter_analysistree2::qaWriter_analysistree2() : fFile(nullptr),
- fTree(nullptr),
- fDataHeader(nullptr),
- fConfig(nullptr),
- fEvent(nullptr),
- fParticles(nullptr),
- isInit(false),
- iB(qaUtility::GetInstance()->error_code),
- iPhiRp(qaUtility::GetInstance()->error_code),
- icharge(qaUtility::GetInstance()->error_code),
- {
- }
- qaWriter_analysistree2::~qaWriter_analysistree2()
- {
- }
- void qaWriter_analysistree2::Init(std::string filename, std::string treename, std::string system, float sqrtSnn)
- {
- fFile = new TFile(filename.c_str(), "recreate");
- fTree = new TTree(treename.c_str(), "AnalysisTree format, model data");
- // Set up AnalysisTree data header
- fDataHeader = new AnalysisTree::DataHeader;
- if (system != "")
- fDataHeader->SetSystem(system);
- if (sqrtSnn > 0)
- fDataHeader->SetBeamMomentum(qaUtility::GetInstance()->GetBeamP(sqrtSnn));
- // Set up AnalysisTree configureation
- fConfig = new AnalysisTree::Configuration;
- // Set up AnalysisTree configuration
- std::string str_event_branch = "Event";
- std::string str_particles_branch = "Particles";
- AnalysisTree::BranchConfig event_branch(str_event_branch.c_str(), AnalysisTree::DetType::kEventHeader);
- AnalysisTree::BranchConfig particles_branch(str_particles_branch.c_str(), AnalysisTree::DetType::kParticle);
- event_branch.AddField<float>("B");
- event_branch.AddField<float>("PhiRp");
- particles_branch.AddField<bool>("is_charged");
- auto hasher = std::hash<std::string>();
- fConfig->AddBranchConfig(event_branch);
- fEvent = new AnalysisTree::EventHeader(Short_t(hasher(event_branch.GetName())));
- fEvent->Init(event_branch);
- fConfig->AddBranchConfig(particles_branch);
- fParticles = new AnalysisTree::Particles(Short_t(hasher(particles_branch.GetName())));
- iB = fConfig->GetBranchConfig(str_event_branch).GetFieldId("B");
- iPhiRp = fConfig->GetBranchConfig(str_event_branch).GetFieldId("PhiRp");
- icharge = fConfig->GetBranchConfig(str_particles_branch).GetFieldId("is_charged");
- // Create branches in the output tree
- fTree->Branch(str_event_branch.c_str(), "AnalysisTree::EventHeader", &fEvent, 32000, 99);
- fTree->Branch(str_particles_branch.c_str(), "AnalysisTree::Particles", &fParticles, 256000, 99);
- fConfig->Print();
- isInit = true;
- }
- void qaWriter_analysistree2::WriteEvent(qaEvent *event)
- {
- if (!event)
- return;
- if (!isInit)
- return;
- fParticles->ClearChannels();
- fEvent->SetField(float(event->GetB()), iB);
- fEvent->SetField(float(event->GetPhiRP()), iPhiRp);
- }
- void qaWriter_analysistree2::WriteParticle(qaParticle *particle)
- {
- if (!particle)
- return;
- if (!isInit)
- return;
- auto *at_particle = fParticles->AddChannel();
- at_particle->Init(fConfig->GetBranchConfig(str_particles_branch));
- at_particle->SetMomentum(particle->GetPx(), particle->GetPy(), particle->GetPz());
- at_particle->SetPid(particle->GetPdg());
- auto particlePDG = (TParticlePDG *)TDatabasePDG::Instance()->GetParticle(particle->GetPdg());
- if (!particlePDG)
- at_particle->SetMass(qaUtility::GetInstance()->error_code);
- if (particlePDG)
- at_particle->SetMass(particlePDG->Mass());
- bool is_charged = (qaUtility::GetInstance()->GetCharge(particle->GetPdg()) == 0 ||
- qaUtility::GetInstance()->GetCharge(particle->GetPdg()) == qaUtility::GetInstance()->error_code)
- ? false
- : true;
- at_particle->SetField(is_charged, icharge);
- }
- void qaWriter_analysistree2::WriteTree()
- {
- if (!isInit) return;
- fFile->cd();
- fTree->Print();
- fTree->Write();
- fDataHeader->Write("DataHeader");
- fConfig->Write("Configuration");
- fFile->Close();
- }
- void qaWriter_analysistree2::FillTree()
- {
- fTree->Fill();
- }
|