MpdDst2AT.C 21 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548
  1. /*// Standard c++ headers
  2. #include <iostream>
  3. #include <string>
  4. #include <utility>
  5. #include <set>
  6. #include <map>
  7. #include <functional>
  8. // Standard ROOT headers
  9. #include <Rtypes.h>
  10. #include <TFile.h>
  11. #include <TTree.h>
  12. #include <TString.h>
  13. #include <TStopwatch.h>
  14. #include <TClonesArray.h>
  15. #include <TObject.h>
  16. #include <TMath.h>
  17. #include <TDatabasePDG.h>
  18. #include <TParticlePDG.h>
  19. // FairRoot/MpdRoot headers
  20. #include "FairMCEventHeader.h"
  21. #ifdef _MCSTACK_
  22. #include "FairMCTrack.h"
  23. #endif
  24. #ifdef _MPDMCSTACK_
  25. #include "MpdMCTrack.h"
  26. #endif
  27. #include "MpdEvent.h"
  28. #include "MpdZdcDigi.h"
  29. #include "MpdPid.h"
  30. #include "MpdTrack.h"
  31. #include "MpdKalmanTrack.h"
  32. #include "MpdVertex.h"
  33. // AnalysisTree headers
  34. #include "AnalysisTree/Configuration.hpp"
  35. #include "AnalysisTree/DataHeader.hpp"
  36. #include "AnalysisTree/EventHeader.hpp"
  37. #include "AnalysisTree/Detector.hpp"
  38. #include "AnalysisTree/Matching.hpp"
  39. */
  40. #define _MPDMCSTACK_
  41. float get_beamP(float sqrtSnn, float m_target = 0.938, float m_beam = 0.938);
  42. float GetFHCalPhi(int iModule);
  43. TVector3 GetFHCalPos(int iModule);
  44. Bool_t IsCharged(Int_t pdg);
  45. Float_t GetRapidity(Float_t p, Float_t pz, Int_t pid_type);
  46. Float_t GetRapidityPDG(Float_t p, Float_t pz, Int_t pdg);
  47. void MpdDst2AT(TString iFileName, TString oFileName, std::string system = "", float sqrtSnn = -1.)
  48. {
  49. bool isDataHeaderDefined = false;
  50. if (system != "" && sqrtSnn != -1) isDataHeaderDefined = true;
  51. TStopwatch timer;
  52. timer.Start();
  53. // Open input dst tree
  54. TFile *fi = new TFile(iFileName.Data(),"read");
  55. if (!fi || fi->IsZombie())
  56. {
  57. std::cerr << "ERROR: Input file probably is empty. Exit the program." << std::endl;
  58. return 100;
  59. }
  60. TTree *dstTree = (TTree*) fi->Get("mpdsim");
  61. // PID related parameters
  62. const Double_t PIDsigM = 4.0;
  63. const Double_t PIDsigE = 4.0;
  64. const Double_t PIDenergy = 11.;
  65. const Double_t PIDkoeff = 1.;
  66. const TString PIDgenerator = "URQMD";
  67. const TString PIDtracking = "CF";
  68. const TString PIDparticles = "pikapr";
  69. const Int_t Num_Of_Modules = 90;
  70. MpdPid *pid = new MpdPid(PIDsigM, PIDsigE, PIDenergy, PIDkoeff, PIDgenerator, PIDtracking, PIDparticles);
  71. // Prepare to read input dst
  72. FairMCEventHeader *MCHeader;
  73. TClonesArray *MCTracks;
  74. MpdEvent *MPDEvent;
  75. TClonesArray *FHCalHits;
  76. TClonesArray *MpdGlobalTracks;
  77. MpdZdcDigi *FHCalHit;
  78. //TClonesArray *mpdKalmanTracks;
  79. TClonesArray *vertexes;
  80. MCHeader = nullptr;
  81. MCTracks = nullptr;
  82. MPDEvent = nullptr;
  83. FHCalHits = nullptr;
  84. //mpdKalmanTracks = nullptr;
  85. vertexes = nullptr;
  86. dstTree->SetBranchAddress("MCEventHeader.", &MCHeader);
  87. dstTree->SetBranchAddress("MCTrack", &MCTracks);
  88. dstTree->SetBranchAddress("MPDEvent.", &MPDEvent);
  89. dstTree->SetBranchAddress("ZdcDigi", &FHCalHits);
  90. //dstTree->SetBranchAddress("TpcKalmanTrack", &mpdKalmanTracks);
  91. dstTree->SetBranchAddress("Vertex", &vertexes);
  92. // Set up output dst
  93. TFile *outFile = new TFile(oFileName.Data(), "RECREATE");
  94. TTree *outTree = new TTree("aTree","AnalysisTree Dst at MPD");
  95. // Set up AnalysisTree data header
  96. AnalysisTree::DataHeader *dataHeader = new AnalysisTree::DataHeader;
  97. if (system != "") dataHeader->SetSystem(system);
  98. if (sqrtSnn > 0) dataHeader->SetBeamMomentum(get_beamP(sqrtSnn));
  99. auto &fhcal_mod_pos = dataHeader->AddDetector();
  100. TVector3 modulePos;
  101. for (int imodule=0; imodule<Num_Of_Modules; imodule++)
  102. {
  103. auto *module = fhcal_mod_pos.AddChannel();
  104. modulePos = GetFHCalPos(imodule);
  105. module->SetPosition(modulePos);
  106. }
  107. // Set up AnalysisTree configureation
  108. AnalysisTree::Configuration *out_config = new AnalysisTree::Configuration;
  109. // Set up AnalysisTree configuration
  110. std::string str_reco_event_branch = "RecoEvent";
  111. std::string str_mc_event_branch = "McEvent";
  112. std::string str_tpc_tracks_branch = "TpcTracks";
  113. std::string str_fhcal_branch = "FHCalModules";
  114. std::string str_mc_tracks_branch = "McTracks";
  115. std::string str_tpc2mc_tracks_branch = "TpcTracks2McTracks";
  116. AnalysisTree::BranchConfig reco_event_branch(str_reco_event_branch.c_str(), AnalysisTree::DetType::kEventHeader);
  117. AnalysisTree::BranchConfig mc_event_branch(str_mc_event_branch.c_str(), AnalysisTree::DetType::kEventHeader);
  118. AnalysisTree::BranchConfig fhcal_branch(str_fhcal_branch.c_str(), AnalysisTree::DetType::kModule);
  119. AnalysisTree::BranchConfig tpc_tracks_branch(str_tpc_tracks_branch.c_str(), AnalysisTree::DetType::kTrack);
  120. AnalysisTree::BranchConfig mc_tracks_branch(str_mc_tracks_branch.c_str(), AnalysisTree::DetType::kParticle);
  121. mc_event_branch.AddField<float>("B");
  122. mc_event_branch.AddField<float>("PhiRp");
  123. // mc_event_branch.AddField<float>("Centrality_B");
  124. // reco_event_branch.AddField<float>("refMult");
  125. // reco_event_branch.AddField<float>("Centrality_refMult");
  126. tpc_tracks_branch.AddField<int>("nhits");
  127. tpc_tracks_branch.AddField<int>("nhits_poss");
  128. tpc_tracks_branch.AddField<int>("charge");
  129. tpc_tracks_branch.AddFields<float>({"dca_x","dca_y","dca_z"});
  130. tpc_tracks_branch.AddField<float>("chi2");
  131. tpc_tracks_branch.AddField<float>("tof_mass2");
  132. tpc_tracks_branch.AddField<float>("tof_flag");
  133. tpc_tracks_branch.AddField<float>("dedx");
  134. tpc_tracks_branch.AddField<float>("pid_prob_pion");
  135. tpc_tracks_branch.AddField<float>("pid_prob_kaon");
  136. tpc_tracks_branch.AddField<float>("pid_prob_proton");
  137. tpc_tracks_branch.AddField<float>("mc_E");
  138. tpc_tracks_branch.AddField<float>("mc_pT");
  139. tpc_tracks_branch.AddField<float>("mc_eta");
  140. tpc_tracks_branch.AddField<float>("mc_phi");
  141. tpc_tracks_branch.AddField<float>("mc_rapidity");
  142. tpc_tracks_branch.AddField<int>("mc_mother_id");
  143. tpc_tracks_branch.AddField<int>("mc_pdg");
  144. tpc_tracks_branch.AddField<int>("eta_sign");
  145. tpc_tracks_branch.AddField<float>("rapidity_pdg");
  146. tpc_tracks_branch.AddField<float>("rapidity_pion");
  147. tpc_tracks_branch.AddField<float>("rapidity_kaon");
  148. tpc_tracks_branch.AddField<float>("rapidity_proton");
  149. fhcal_branch.AddField<float>("phi");
  150. fhcal_branch.AddField<float>("signal_eta_signed");
  151. mc_tracks_branch.AddField<int>("mother_id");
  152. mc_tracks_branch.AddField<bool>("is_charged");
  153. mc_tracks_branch.AddField<int>("eta_sign");
  154. auto hasher = std::hash<std::string>();
  155. // Initialize AnalysisTree Dst components
  156. out_config->AddBranchConfig(reco_event_branch);
  157. AnalysisTree::EventHeader *reco_event = new AnalysisTree::EventHeader( Short_t(hasher(reco_event_branch.GetName())) );
  158. out_config->AddBranchConfig(mc_event_branch);
  159. AnalysisTree::EventHeader *mc_event = new AnalysisTree::EventHeader( Short_t(hasher(mc_event_branch.GetName())) );
  160. out_config->AddBranchConfig(tpc_tracks_branch);
  161. AnalysisTree::TrackDetector *tpc_tracks = new AnalysisTree::TrackDetector( Short_t(hasher(tpc_tracks_branch.GetName())) );
  162. out_config->AddBranchConfig(fhcal_branch);
  163. AnalysisTree::ModuleDetector *fhcal_modules = new AnalysisTree::ModuleDetector( Short_t(hasher(fhcal_branch.GetName())) );
  164. out_config->AddBranchConfig(mc_tracks_branch);
  165. AnalysisTree::Particles *mc_tracks = new AnalysisTree::Particles( Short_t(hasher(mc_tracks_branch.GetName())) );
  166. AnalysisTree::Matching *tpc2mc_tracks = new AnalysisTree::Matching(out_config->GetBranchConfig(str_tpc_tracks_branch).GetId(), out_config->GetBranchConfig(str_mc_tracks_branch).GetId());
  167. out_config->AddMatch(tpc2mc_tracks);
  168. reco_event->Init(reco_event_branch);
  169. mc_event->Init(mc_event_branch);
  170. // mc_event's additional field ids
  171. //const int mceventid = mc_event->GetId();
  172. const int iB = out_config->GetBranchConfig(str_mc_event_branch).GetFieldId("B");
  173. const int iPhiRp = out_config->GetBranchConfig(str_mc_event_branch).GetFieldId("PhiRp");
  174. // fhcal_modules' additional field ids
  175. //const int fhcalid = fhcal_modules->GetId();
  176. const int ifhcalphi = out_config->GetBranchConfig(str_fhcal_branch).GetFieldId("phi");
  177. const int ifhcalsign = out_config->GetBranchConfig(str_fhcal_branch).GetFieldId("signal_eta_signed");
  178. // tpc_tracks' additional field ids
  179. const int inhits = out_config->GetBranchConfig(str_tpc_tracks_branch).GetFieldId("nhits");
  180. const int inhits_poss = out_config->GetBranchConfig(str_tpc_tracks_branch).GetFieldId("nhits_poss");
  181. const int icharge = out_config->GetBranchConfig(str_tpc_tracks_branch).GetFieldId("charge");
  182. const int idcax = out_config->GetBranchConfig(str_tpc_tracks_branch).GetFieldId("dca_x");
  183. const int idcay = out_config->GetBranchConfig(str_tpc_tracks_branch).GetFieldId("dca_y");
  184. const int idcaz = out_config->GetBranchConfig(str_tpc_tracks_branch).GetFieldId("dca_z");
  185. const int ichi2 = out_config->GetBranchConfig(str_tpc_tracks_branch).GetFieldId("chi2");
  186. const int itof_mass2 = out_config->GetBranchConfig(str_tpc_tracks_branch).GetFieldId("tof_mass2");
  187. const int itof_flag = out_config->GetBranchConfig(str_tpc_tracks_branch).GetFieldId("tof_flag");
  188. const int idedx = out_config->GetBranchConfig(str_tpc_tracks_branch).GetFieldId("dedx");
  189. const int ipid_prob_pion = out_config->GetBranchConfig(str_tpc_tracks_branch).GetFieldId("pid_prob_pion");
  190. const int ipid_prob_kaon = out_config->GetBranchConfig(str_tpc_tracks_branch).GetFieldId("pid_prob_kaon");
  191. const int ipid_prob_proton = out_config->GetBranchConfig(str_tpc_tracks_branch).GetFieldId("pid_prob_proton");
  192. const int imc_E = out_config->GetBranchConfig(str_tpc_tracks_branch).GetFieldId("mc_E");
  193. const int imc_pT = out_config->GetBranchConfig(str_tpc_tracks_branch).GetFieldId("mc_pT");
  194. const int imc_eta = out_config->GetBranchConfig(str_tpc_tracks_branch).GetFieldId("mc_eta");
  195. const int imc_phi = out_config->GetBranchConfig(str_tpc_tracks_branch).GetFieldId("mc_phi");
  196. const int imc_rapidity = out_config->GetBranchConfig(str_tpc_tracks_branch).GetFieldId("mc_rapidity");
  197. const int imc_mother_id = out_config->GetBranchConfig(str_tpc_tracks_branch).GetFieldId("mc_mother_id");
  198. const int imc_pdg = out_config->GetBranchConfig(str_tpc_tracks_branch).GetFieldId("mc_pdg");
  199. const int ietasign = out_config->GetBranchConfig(str_tpc_tracks_branch).GetFieldId("eta_sign");
  200. const int irapidity_pion = out_config->GetBranchConfig(str_tpc_tracks_branch).GetFieldId("rapidity_pion");
  201. const int irapidity_pdg = out_config->GetBranchConfig(str_tpc_tracks_branch).GetFieldId("rapidity_pdg");
  202. const int irapidity_kaon = out_config->GetBranchConfig(str_tpc_tracks_branch).GetFieldId("rapidity_kaon");
  203. const int irapidity_proton = out_config->GetBranchConfig(str_tpc_tracks_branch).GetFieldId("rapidity_proton");
  204. // mc_tracks' additional field ids
  205. const int imother_id = out_config->GetBranchConfig(str_mc_tracks_branch).GetFieldId("mother_id");
  206. const int icharge_mc = out_config->GetBranchConfig(str_mc_tracks_branch).GetFieldId("is_charged");
  207. const int ietasign_mc = out_config->GetBranchConfig(str_mc_tracks_branch).GetFieldId("eta_sign");
  208. // Create branches in the output tree
  209. outTree->Branch(str_reco_event_branch.c_str(), "AnalysisTree::EventHeader", &reco_event, 32000, 99);
  210. outTree->Branch(str_mc_event_branch.c_str(), "AnalysisTree::EventHeader", &mc_event, 32000, 99);
  211. outTree->Branch(str_tpc_tracks_branch.c_str(), "AnalysisTree::TrackDetector", &tpc_tracks, 256000, 99);
  212. outTree->Branch(str_fhcal_branch.c_str(), "AnalysisTree::ModuleDetector", &fhcal_modules, 128000, 99);
  213. outTree->Branch(str_mc_tracks_branch.c_str(), "AnalysisTree::Particles", &mc_tracks, 256000, 99);
  214. outTree->Branch(str_tpc2mc_tracks_branch.c_str(), "AnalysisTree::Matching", &tpc2mc_tracks, 32000, 99);
  215. // Printout basic configuration info
  216. out_config->Print();
  217. // Starting event loop
  218. TVector3 primaryVertex, mc_assoc_mom;
  219. std::set <Int_t> UsedMCTracks; // using to remap mc-reco track matching
  220. std::map <Int_t,Int_t> InitMcNewMcId; // map[old-mc-id] = new-mc-id
  221. Float_t FHCalSumEnergy[Num_Of_Modules];
  222. Int_t FHCalNumOfHits[Num_Of_Modules];
  223. Int_t n_entries = dstTree->GetEntriesFast();
  224. bool isGoodPID;
  225. int charge;
  226. for (int iEv = 0; iEv < n_entries; iEv++)
  227. {
  228. std::cout << "Event [" << iEv << "/" << n_entries << "]" << std::endl;
  229. dstTree->GetEntry(iEv);
  230. UsedMCTracks.clear();
  231. InitMcNewMcId.clear();
  232. tpc2mc_tracks->Clear();
  233. for (int i=0; i<Num_Of_Modules; i++)
  234. {
  235. FHCalSumEnergy[i] = 0.;
  236. FHCalNumOfHits[i] = 0;
  237. }
  238. tpc_tracks->ClearChannels();
  239. fhcal_modules->ClearChannels();
  240. mc_tracks->ClearChannels();
  241. // Reading Reco Event
  242. MpdVertex *vertex = (MpdVertex *)vertexes->First();
  243. vertex->Position(primaryVertex);
  244. reco_event->SetVertexPosition3(primaryVertex);
  245. // Read MC Event
  246. TVector3 vtx(MCHeader->GetX(),MCHeader->GetY(),MCHeader->GetZ());
  247. mc_event->SetVertexPosition3(vtx);
  248. mc_event->SetField(float(MCHeader->GetB()), iB);
  249. mc_event->SetField(float(MCHeader->GetRotZ()), iPhiRp);
  250. // Read energy in FHCal modules
  251. fhcal_modules->Reserve(Num_Of_Modules);
  252. for (int imodule=0; imodule<Num_Of_Modules; imodule++)
  253. {
  254. auto *module = fhcal_modules->AddChannel();
  255. module->Init(out_config->GetBranchConfig(str_fhcal_branch));
  256. module->SetSignal(0.f);
  257. }
  258. Int_t number_of_FHCal_hits = FHCalHits->GetEntriesFast();
  259. for(int ihit=0; ihit<number_of_FHCal_hits; ihit++)
  260. {
  261. FHCalHit = (MpdZdcDigi*) FHCalHits->UncheckedAt(ihit);
  262. Int_t DetId = FHCalHit->GetDetectorID();
  263. Int_t ModId = FHCalHit->GetModuleID()-1;
  264. Int_t ModNumber = ModId + (Num_Of_Modules/2) * (DetId-1);
  265. FHCalSumEnergy[ModNumber] += FHCalHit->GetELoss();
  266. FHCalNumOfHits[ModNumber]++;
  267. }
  268. for (int imodule=0; imodule<Num_Of_Modules; imodule++)
  269. {
  270. auto& module = fhcal_modules->Channel(imodule);
  271. int fhcal_sign = (imodule < 45) ? -1 : 1;
  272. module.SetNumber(FHCalNumOfHits[imodule]); // Number of hits that got in the module
  273. module.SetSignal(FHCalSumEnergy[imodule]); // Total energy from hits in the module
  274. module.SetField(float(GetFHCalPhi(imodule)), ifhcalphi);
  275. module.SetField(float(FHCalSumEnergy[imodule]*(float)fhcal_sign), ifhcalsign);
  276. }
  277. // Reading Reco Tracks
  278. MpdGlobalTracks = (TClonesArray*)MPDEvent->GetGlobalTracks();
  279. Int_t Num_of_tpc_tracks = MpdGlobalTracks->GetEntriesFast();
  280. tpc_tracks->Reserve(Num_of_tpc_tracks);
  281. for (int itrack=0; itrack<Num_of_tpc_tracks; itrack++)
  282. {
  283. MpdTrack* mpdtrack = (MpdTrack*) MpdGlobalTracks->UncheckedAt(itrack);
  284. UsedMCTracks.insert(mpdtrack->GetID());
  285. auto *track = tpc_tracks->AddChannel();
  286. track->Init(out_config->GetBranchConfig(str_tpc_tracks_branch));
  287. int tpc_eta_sign = (mpdtrack->GetEta() < 0) ? -1 : 1;
  288. track->SetMomentum(mpdtrack->GetPx(),mpdtrack->GetPy(),mpdtrack->GetPz());
  289. track->SetField(int(mpdtrack->GetNofHits()), inhits);
  290. track->SetField(int(mpdtrack->GetNofHitsPossTpc()), inhits_poss);
  291. track->SetField(float(mpdtrack->GetTofMass2()), itof_mass2);
  292. track->SetField(float(mpdtrack->GetTofFlag()), itof_flag);
  293. track->SetField(float(mpdtrack->GetdEdXTPC()), idedx);
  294. track->SetField(float(mpdtrack->GetChi2()), ichi2);
  295. track->SetField(float(mpdtrack->GetDCAX()), idcax);
  296. track->SetField(float(mpdtrack->GetDCAY()), idcay);
  297. track->SetField(float(mpdtrack->GetDCAZ()), idcaz);
  298. charge = (mpdtrack->GetPt() < 0) ? 1 : -1;
  299. track->SetField(int(charge), icharge);
  300. track->SetField(int(tpc_eta_sign), ietasign);
  301. Float_t mpdtrack_p = TMath::Sqrt(TMath::Power(mpdtrack->GetPx(),2) + TMath::Power(mpdtrack->GetPy(),2) + TMath::Power(mpdtrack->GetPz(),2));
  302. track->SetField(float(GetRapidity(mpdtrack_p, mpdtrack->GetPz(), 0)), irapidity_pion);
  303. track->SetField(float(GetRapidity(mpdtrack_p, mpdtrack->GetPz(), 1)), irapidity_kaon);
  304. track->SetField(float(GetRapidity(mpdtrack_p, mpdtrack->GetPz(), 2)), irapidity_proton);
  305. // PID
  306. if (mpdtrack->GetTofFlag() == 2 || mpdtrack->GetTofFlag() == 6)
  307. {
  308. // TOF + TPC
  309. isGoodPID = pid->FillProbs(TMath::Abs(mpdtrack->GetPt()) * TMath::CosH(mpdtrack->GetEta()),
  310. mpdtrack->GetdEdXTPC(), mpdtrack->GetTofMass2(), charge);
  311. }
  312. else
  313. {
  314. // TPC only
  315. isGoodPID = pid->FillProbs(TMath::Abs(mpdtrack->GetPt()) * TMath::CosH(mpdtrack->GetEta()),
  316. mpdtrack->GetdEdXTPC(), charge);
  317. }
  318. if (isGoodPID)
  319. {
  320. track->SetField(float(pid->GetProbPi()), ipid_prob_pion);
  321. track->SetField(float(pid->GetProbKa()), ipid_prob_kaon);
  322. track->SetField(float(pid->GetProbPr()), ipid_prob_proton);
  323. }
  324. else
  325. {
  326. track->SetField(float(-999.), ipid_prob_pion);
  327. track->SetField(float(-999.), ipid_prob_kaon);
  328. track->SetField(float(-999.), ipid_prob_proton);
  329. }
  330. #ifdef _MCSTACK_
  331. FairMCTrack *mctrack = (FairMCTrack*) MCTracks->UncheckedAt(mpdtrack->GetID());
  332. #endif
  333. #ifdef _MPDMCSTACK_
  334. MpdMCTrack *mctrack = (MpdMCTrack*) MCTracks->UncheckedAt(mpdtrack->GetID());
  335. #endif
  336. mc_assoc_mom.SetXYZ(mctrack->GetPx(), mctrack->GetPy(), mctrack->GetPz());
  337. track->SetField(float(mctrack->GetEnergy()), imc_E);
  338. track->SetField(float(mc_assoc_mom.Pt()), imc_pT);
  339. track->SetField(float(mc_assoc_mom.Eta()), imc_eta);
  340. track->SetField(float(mc_assoc_mom.Phi()), imc_phi);
  341. track->SetField(float(mctrack->GetRapidity()), imc_rapidity);
  342. track->SetField(int(mctrack->GetMotherId()), imc_mother_id);
  343. track->SetField(int(mctrack->GetPdgCode()), imc_pdg);
  344. track->SetField(float(GetRapidityPDG(mpdtrack_p, mpdtrack->GetPz(), mctrack->GetPdgCode())), irapidity_pdg);
  345. } // End of the tpc track loop
  346. // Read Mc tracks
  347. Int_t Num_of_mc_tracks = MCTracks->GetEntriesFast();
  348. mc_tracks->Reserve(Num_of_mc_tracks);
  349. for (int imctrack=0; imctrack<Num_of_mc_tracks; imctrack++)
  350. {
  351. #ifdef _MCSTACK_
  352. FairMCTrack *mctrack = (FairMCTrack*) MCTracks->UncheckedAt(imctrack);
  353. #endif
  354. #ifdef _MPDMCSTACK_
  355. MpdMCTrack *mctrack = (MpdMCTrack*) MCTracks->UncheckedAt(imctrack);
  356. #endif
  357. bool isUsed = (UsedMCTracks.count(imctrack));
  358. // If motherId != 1 and mc track doesn't relate to any reco track - skip
  359. if (mctrack->GetMotherId() != -1 && !isUsed) continue;
  360. auto *track = mc_tracks->AddChannel();
  361. track->Init(out_config->GetBranchConfig(str_mc_tracks_branch));
  362. // Collect new Mc Ids
  363. InitMcNewMcId[imctrack] = track->GetId();
  364. int mc_eta_sign = (mctrack->GetRapidity() < 0) ? -1 : 1;
  365. track->SetMomentum(mctrack->GetPx(), mctrack->GetPy(), mctrack->GetPz());
  366. track->SetPid(int(mctrack->GetPdgCode()));
  367. track->SetMass(float(mctrack->GetMass()));
  368. track->SetField(int(mctrack->GetMotherId()), imother_id);
  369. track->SetField(bool(IsCharged(mctrack->GetPdgCode())), icharge_mc);
  370. track->SetField(int(mc_eta_sign), ietasign_mc);
  371. } // End of the mc track loop
  372. // reco-mc tracks matching
  373. for (int itrack=0; itrack<Num_of_tpc_tracks; itrack++)
  374. {
  375. MpdTrack* mpdtrack = (MpdTrack*) MpdGlobalTracks->UncheckedAt(itrack);
  376. const int tpc_track_id = itrack;
  377. const int mc_track_id = InitMcNewMcId[mpdtrack->GetID()];
  378. tpc2mc_tracks->AddMatch(tpc_track_id, mc_track_id);
  379. }
  380. outTree->Fill();
  381. } // End of the event loop
  382. outFile->cd();
  383. outTree->Print();
  384. outTree->Write();
  385. if (isDataHeaderDefined) dataHeader->Write("DataHeader");
  386. out_config->Write("Configuration");
  387. outFile->Close();
  388. timer.Stop();
  389. timer.Print();
  390. }
  391. float get_beamP(float sqrtSnn, float m_target, float m_beam)
  392. {
  393. return sqrt( pow((pow(sqrtSnn,2) - pow(m_target,2) - pow(m_beam,2))/(2*m_target), 2) - pow(m_beam,2) );
  394. }
  395. float GetFHCalPhi(int iModule)
  396. {
  397. const int Nmodules = 45;
  398. int xAxisSwitch = (iModule < Nmodules) ? 1 : -1;
  399. int module = (iModule < Nmodules) ? iModule : iModule - Nmodules;
  400. float x, y, phi = -999.;
  401. if (module >= 0 && module <= 4)
  402. {
  403. y = 45.;
  404. x = (module - 2) * 15.;
  405. phi = TMath::ATan2(y, x * xAxisSwitch);
  406. }
  407. else if ((module >= 5) && (module <= 39))
  408. {
  409. y = (3 - (module + 2) / 7) * 15.;
  410. x = (3 - (module + 2) % 7) * 15.;
  411. phi = TMath::ATan2(y, x * xAxisSwitch);
  412. }
  413. else if ((module >= 40) && (module <= 44))
  414. {
  415. y = -45.;
  416. x = (module - 42) * 15.;
  417. phi = TMath::ATan2(y, x * xAxisSwitch);
  418. }
  419. return phi;
  420. }
  421. TVector3 GetFHCalPos(int iModule)
  422. {
  423. const int Nmodules = 45;
  424. int xAxisSwitch = (iModule < Nmodules) ? 1 : -1;
  425. int module = (iModule < Nmodules) ? iModule : iModule - Nmodules;
  426. float x, y, z;
  427. z = (iModule < Nmodules) ? -320. : 320.;
  428. if (module >= 0 && module <= 4)
  429. {
  430. y = 45.;
  431. x = (module - 2) * 15.;
  432. }
  433. else if ((module >= 5) && (module <= 39))
  434. {
  435. y = (3 - (module + 2) / 7) * 15.;
  436. x = (3 - (module + 2) % 7) * 15.;
  437. }
  438. else if ((module >= 40) && (module <= 44))
  439. {
  440. y = -45.;
  441. x = (module - 42) * 15.;
  442. }
  443. TVector3 vec(x * xAxisSwitch, y, z);
  444. return vec;
  445. }
  446. Bool_t IsCharged(Int_t pdg)
  447. {
  448. auto particle = (TParticlePDG *)TDatabasePDG::Instance()->GetParticle(pdg);
  449. if (!particle)
  450. return false;
  451. return ( particle->Charge() != 0 );
  452. }
  453. Float_t GetRapidity(Float_t p, Float_t pz, Int_t pid_type)
  454. {
  455. Float_t mass; // in GeV/c^2
  456. if (pid_type == 0) mass = 0.13957; // pions
  457. if (pid_type == 1) mass = 0.493677; // kaons
  458. if (pid_type == 2) mass = 0.938272; // protons
  459. if (pid_type != 0 && pid_type != 1 && pid_type != 2) return -999.;
  460. Float_t E = TMath::Sqrt(p*p + mass*mass);
  461. return 0.5 * TMath::Log((E + pz)/(E-pz));
  462. }
  463. Float_t GetRapidityPDG(Float_t p, Float_t pz, Int_t pdg)
  464. {
  465. Float_t mass;
  466. auto pdgParticle = (TParticlePDG*)TDatabasePDG::Instance()->GetParticle(pdg);
  467. if (!pdgParticle) return -999.;
  468. mass = pdgParticle->Mass();
  469. Float_t E = TMath::Sqrt(p*p + mass*mass);
  470. return 0.5 * TMath::Log((E + pz)/(E - pz));
  471. }