UserTaskWrite.cpp 7.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188
  1. #include "UserTaskWrite.hpp"
  2. #include "TaskManager.hpp"
  3. using namespace AnalysisTree;
  4. void UserTaskWrite::Init() { // Preparations before reading events
  5. // Read input chain with AnalysisTree
  6. auto man = TaskManager::GetInstance();
  7. auto chain = man->GetChain();
  8. AddInputBranch("RecoEvent");
  9. reco_evt_ = chain->GetBranch("RecoEvent");
  10. reco_evt_.Freeze();
  11. AddInputBranch("McEvent");
  12. mc_evt_ = chain->GetBranch("McEvent");
  13. mc_evt_.Freeze();
  14. AddInputBranch("TpcTracks");
  15. tracks_ = chain->GetBranch("TpcTracks");
  16. tracks_.Freeze();
  17. AddInputBranch("McTracks");
  18. particles_ = chain->GetBranch("McTracks");
  19. particles_.Freeze();
  20. match_ = chain->GetMatching("TpcTracks", "McTracks");
  21. // Set up branch configurations (copy them from the input and then modify)
  22. auto br_conf_reco_evt = chain->GetConfiguration()->GetBranchConfig("RecoEvent");
  23. auto br_conf_mc_evt = chain->GetConfiguration()->GetBranchConfig("McEvent");
  24. auto new_conf_reco_evt = br_conf_reco_evt.Clone("RecoEventExt", DetType::kEventHeader);
  25. new_conf_reco_evt.AddFields<float>({"bcent"}, "Extended parameter");
  26. auto br_conf_tracks = chain->GetConfiguration()->GetBranchConfig("TpcTracks");
  27. auto new_conf_tracks = br_conf_tracks.Clone("TpcTracksExt", DetType::kTrack);
  28. new_conf_tracks.AddFields<float>({"mc_E", "mc_pT", "mc_eta", "mc_phi", "mc_y", "rapidity"}, "Extended parameter");
  29. new_conf_tracks.AddFields<int>({"eta_sign", "mc_pid", "mc_mother_id"}, "Extended parameter");
  30. auto br_conf_particles = chain->GetConfiguration()->GetBranchConfig("McTracks");
  31. auto new_conf_particles = br_conf_particles.Clone("McTracksExt", DetType::kParticle);
  32. new_conf_particles.AddFields<bool>({"is_charged"}, "Extended parameter");
  33. new_conf_particles.AddFields<int>({"eta_sign"}, "Extended parameter");
  34. // Prepare output branches and then add them to the output file via task manager
  35. new_reco_evt_ = Branch(new_conf_reco_evt);
  36. new_reco_evt_.SetMutable();
  37. new_tracks_ = Branch(new_conf_tracks);
  38. new_tracks_.SetMutable();
  39. new_particles_ = Branch(new_conf_particles);
  40. new_particles_.SetMutable();
  41. man->AddBranch(&new_reco_evt_);
  42. man->AddBranch(&new_tracks_);
  43. man->AddBranch(&new_particles_);
  44. }
  45. void UserTaskWrite::Exec() { // Procedures to do iside loop over events
  46. // Get links to variables from the input mc event header and output reco event header
  47. auto imc_event_b = mc_evt_.GetField("B");
  48. auto ireco_evt_bcent = new_reco_evt_.GetField("bcent");
  49. // Get links to variables from the input and output reco tracks
  50. auto itpc_pt = tracks_.GetField("pT");
  51. auto itpc_p = tracks_.GetField("p");
  52. auto itpc_pz = tracks_.GetField("pz");
  53. auto itpc_eta = tracks_.GetField("eta");
  54. auto itpc_track_mc_E = new_tracks_.GetField("mc_E");
  55. auto itpc_track_mc_pT = new_tracks_.GetField("mc_pT");
  56. auto itpc_track_mc_eta = new_tracks_.GetField("mc_eta");
  57. auto itpc_track_mc_phi = new_tracks_.GetField("mc_phi");
  58. auto itpc_track_mc_y = new_tracks_.GetField("mc_y");
  59. auto itpc_track_rapidity = new_tracks_.GetField("rapidity");
  60. auto itpc_track_eta_sign = new_tracks_.GetField("eta_sign");
  61. auto itpc_track_mc_pid = new_tracks_.GetField("mc_pid");
  62. auto itpc_track_mc_mother_id = new_tracks_.GetField("mc_mother_id");
  63. // Get links to variables from the input and output mc tracks
  64. auto imc_mass = particles_.GetField("mass");
  65. auto imc_p = particles_.GetField("p");
  66. auto imc_pT = particles_.GetField("pT");
  67. auto imc_eta = particles_.GetField("eta");
  68. auto imc_phi = particles_.GetField("phi");
  69. auto imc_rapidity = particles_.GetField("rapidity");
  70. auto imc_pid = particles_.GetField("pid");
  71. auto imc_mother_id = particles_.GetField("mother_id");
  72. auto imc_particle_is_charged = new_particles_.GetField("is_charged");
  73. auto imc_particle_eta_sign = new_particles_.GetField("eta_sign");
  74. // Processing reco & mc event
  75. for(size_t i=0; i<1; ++i){
  76. auto mc_evt_var = mc_evt_[i];
  77. auto new_reco_evt_var = new_reco_evt_[i];
  78. float bcent = GetBCent(mc_evt_var.Value(imc_event_b));
  79. new_reco_evt_var.SetValue(ireco_evt_bcent, bcent);
  80. }
  81. // Processing reco tracks
  82. new_tracks_.ClearChannels();
  83. for(size_t i=0; i<tracks_.size(); ++i){ // Loop over reco tracks
  84. auto track = tracks_[i];
  85. auto match_id = match_->GetMatch(track.GetId());
  86. auto new_track = new_tracks_.NewChannel();
  87. new_track.CopyContent(track);
  88. int tpc_eta_sign = (track.Value(itpc_eta) < 0) ? -1 : 1;
  89. new_track.SetValue(itpc_track_eta_sign, tpc_eta_sign);
  90. // Fill extended fields with info from associated mc tracks
  91. if(match_id >= 0){
  92. auto particle = particles_[match_id];
  93. new_track.SetValue(itpc_track_mc_pid, particle.Value(imc_pid));
  94. new_track.SetValue(itpc_track_mc_pT, particle.Value(imc_pT));
  95. new_track.SetValue(itpc_track_mc_eta, particle.Value(imc_eta));
  96. new_track.SetValue(itpc_track_mc_phi, particle.Value(imc_phi));
  97. new_track.SetValue(itpc_track_mc_y, particle.Value(imc_rapidity));
  98. new_track.SetValue(itpc_track_mc_mother_id, particle.Value(imc_mother_id));
  99. double Ene = sqrt(pow(particle.Value(imc_p),2) + pow(particle.Value(imc_mass),2));
  100. new_track.SetValue(itpc_track_mc_E, Ene);
  101. double rapidityPDG = GetRapidityPDG((float)track.Value(itpc_p), (float)track.Value(itpc_pz), (int)particle.Value(imc_pid));
  102. new_track.SetValue(itpc_track_rapidity, rapidityPDG);
  103. }
  104. }
  105. // Processing mc tracks
  106. new_particles_.ClearChannels();
  107. for(size_t i=0; i<particles_.size(); ++i){ // Loop over mc tracks
  108. auto particle = particles_[i];
  109. auto new_particle = new_particles_.NewChannel();
  110. new_particle.CopyContent(particle);
  111. int mc_eta_sign = (particle.Value(itpc_eta) < 0) ? -1 : 1;
  112. bool is_charged = IsCharged(particle.Value(imc_pid));
  113. new_particle.SetValue(imc_particle_eta_sign, mc_eta_sign);
  114. new_particle.SetValue(imc_particle_is_charged, is_charged);
  115. }
  116. }
  117. UserTaskWrite::~UserTaskWrite() { // Destructor. Make sure to delete branch pointers for each output branch
  118. delete new_reco_evt_ptr_;
  119. delete new_particles_ptr_;
  120. delete new_tracks_ptr_;
  121. }
  122. bool UserTaskWrite::IsCharged(Int_t pdg) { // Additional function returning bool value wheather the particle is charged or not
  123. auto particle = (TParticlePDG *)TDatabasePDG::Instance()->GetParticle(pdg);
  124. if (!particle)
  125. return false;
  126. return (particle->Charge() != 0);
  127. }
  128. float UserTaskWrite::GetRapidityPDG(Float_t p, Float_t pz, Int_t pdg) { // Additional function returning rapidity value based on particle's pdg code
  129. Float_t mass;
  130. auto pdgParticle = (TParticlePDG *)TDatabasePDG::Instance()->GetParticle(pdg);
  131. if (!pdgParticle)
  132. return -999.;
  133. mass = pdgParticle->Mass();
  134. Float_t E = TMath::Sqrt(p * p + mass * mass);
  135. return 0.5 * TMath::Log((E + pz) / (E - pz));
  136. }
  137. float UserTaskWrite::GetBCent(Float_t bimp) { // Additional function returning centrality class based on fixed impact parameter ranges
  138. // Impact parameter values for each centrlity class
  139. const std::vector<float> v_bimp = {0., 4.06, 5.84, 7.17, 8.27, 9.26, 10.17, 11.04, 11.99, 14.};
  140. // Centrality class values
  141. const std::vector<float> v_cent = {0., 10., 20., 30., 40., 50., 60., 70., 80., 100.};
  142. Float_t cent=-1;
  143. for (int i = 0; i<v_bimp.size()-1; i++)
  144. {
  145. if (bimp >= v_bimp.at(i) && bimp < v_bimp.at(i+1))
  146. {
  147. cent = v_cent.at(i) + (v_cent.at(i+1) - v_cent.at(i))/2.;
  148. }
  149. }
  150. return cent;
  151. }