MpdDst2AT.C 16 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446
  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. // _MPDMCSTACK_ - for newer versions of mpdroot (since 2020), _MCSTACK_ - for the older mpdroot versions
  41. #define _MPDMCSTACK_
  42. float get_beamP(float sqrtSnn, float m_target = 0.938, float m_beam = 0.938);
  43. float GetFHCalPhi(int iModule);
  44. TVector3 GetFHCalPos(int iModule);
  45. Bool_t IsCharged(Int_t pdg);
  46. Float_t GetRapidity(Float_t p, Float_t pz, Int_t pid_type);
  47. Float_t GetRapidityPDG(Float_t p, Float_t pz, Int_t pdg);
  48. void MpdDst2AT(TString iFileName, TString oFileName, std::string system = "", float sqrtSnn = -1.)
  49. {
  50. bool isDataHeaderDefined = false;
  51. if (system != "" && sqrtSnn != -1) isDataHeaderDefined = true;
  52. TStopwatch timer;
  53. timer.Start();
  54. // Open input dst tree
  55. TFile *fi = new TFile(iFileName.Data(),"read");
  56. if (!fi || fi->IsZombie())
  57. {
  58. std::cerr << "ERROR: Input file probably is empty. Exit the program." << std::endl;
  59. return 100;
  60. }
  61. TTree *dstTree = (TTree*) fi->Get("mpdsim");
  62. // PID related parameters
  63. const Double_t PIDsigM = 4.0;
  64. const Double_t PIDsigE = 4.0;
  65. const Double_t PIDenergy = 11.;
  66. const Double_t PIDkoeff = 1.;
  67. const TString PIDgenerator = "URQMD";
  68. const TString PIDtracking = "CF";
  69. const TString PIDparticles = "pikapr";
  70. const Int_t Num_Of_Modules = 90;
  71. MpdPid *pid = new MpdPid(PIDsigM, PIDsigE, PIDenergy, PIDkoeff, PIDgenerator, PIDtracking, PIDparticles);
  72. // Prepare to read input dst
  73. FairMCEventHeader *MCHeader;
  74. TClonesArray *MCTracks;
  75. MpdEvent *MPDEvent;
  76. TClonesArray *FHCalHits;
  77. TClonesArray *MpdGlobalTracks;
  78. MpdZdcDigi *FHCalHit;
  79. //TClonesArray *mpdKalmanTracks;
  80. TClonesArray *vertexes;
  81. MCHeader = nullptr;
  82. MCTracks = nullptr;
  83. MPDEvent = nullptr;
  84. FHCalHits = nullptr;
  85. //mpdKalmanTracks = nullptr;
  86. vertexes = nullptr;
  87. dstTree->SetBranchAddress("MCEventHeader.", &MCHeader);
  88. dstTree->SetBranchAddress("MCTrack", &MCTracks);
  89. dstTree->SetBranchAddress("MPDEvent.", &MPDEvent);
  90. dstTree->SetBranchAddress("ZdcDigi", &FHCalHits);
  91. //dstTree->SetBranchAddress("TpcKalmanTrack", &mpdKalmanTracks);
  92. dstTree->SetBranchAddress("Vertex", &vertexes);
  93. // Set up output dst
  94. TFile *outFile = new TFile(oFileName.Data(), "RECREATE");
  95. TTree *outTree = new TTree("aTree","AnalysisTree Dst at MPD");
  96. // Set up AnalysisTree data header
  97. AnalysisTree::DataHeader *dataHeader = new AnalysisTree::DataHeader;
  98. if (system != "") dataHeader->SetSystem(system);
  99. if (sqrtSnn > 0) dataHeader->SetBeamMomentum(get_beamP(sqrtSnn));
  100. auto &fhcal_mod_pos = dataHeader->AddDetector();
  101. TVector3 modulePos;
  102. for (int imodule=0; imodule<Num_Of_Modules; imodule++)
  103. {
  104. auto *module = fhcal_mod_pos.AddChannel();
  105. modulePos = GetFHCalPos(imodule);
  106. module->SetPosition(modulePos);
  107. }
  108. // Set up AnalysisTree configureation
  109. AnalysisTree::Configuration *out_config = new AnalysisTree::Configuration;
  110. // Set up AnalysisTree configuration
  111. std::string str_reco_event_branch = "RecoEvent";
  112. std::string str_mc_event_branch = "McEvent";
  113. std::string str_tpc_tracks_branch = "TpcTracks";
  114. std::string str_fhcal_branch = "FHCalModules";
  115. std::string str_mc_tracks_branch = "McTracks";
  116. std::string str_tpc2mc_tracks_branch = "TpcTracks2McTracks";
  117. AnalysisTree::BranchConfig reco_event_branch(str_reco_event_branch.c_str(), AnalysisTree::DetType::kEventHeader);
  118. AnalysisTree::BranchConfig mc_event_branch(str_mc_event_branch.c_str(), AnalysisTree::DetType::kEventHeader);
  119. AnalysisTree::BranchConfig fhcal_branch(str_fhcal_branch.c_str(), AnalysisTree::DetType::kModule);
  120. AnalysisTree::BranchConfig tpc_tracks_branch(str_tpc_tracks_branch.c_str(), AnalysisTree::DetType::kTrack);
  121. AnalysisTree::BranchConfig mc_tracks_branch(str_mc_tracks_branch.c_str(), AnalysisTree::DetType::kParticle);
  122. mc_event_branch.AddField<float>("B");
  123. mc_event_branch.AddField<float>("PhiRp");
  124. tpc_tracks_branch.AddField<int>("nhits");
  125. tpc_tracks_branch.AddField<int>("nhits_poss");
  126. tpc_tracks_branch.AddField<int>("charge");
  127. tpc_tracks_branch.AddFields<float>({"dca_x","dca_y","dca_z"});
  128. tpc_tracks_branch.AddField<float>("chi2");
  129. tpc_tracks_branch.AddField<float>("tof_mass2");
  130. tpc_tracks_branch.AddField<float>("tof_flag");
  131. tpc_tracks_branch.AddField<float>("dedx");
  132. mc_tracks_branch.AddField<int>("mother_id");
  133. auto hasher = std::hash<std::string>();
  134. // Initialize AnalysisTree Dst components
  135. out_config->AddBranchConfig(reco_event_branch);
  136. AnalysisTree::EventHeader *reco_event = new AnalysisTree::EventHeader( Short_t(hasher(reco_event_branch.GetName())) );
  137. out_config->AddBranchConfig(mc_event_branch);
  138. AnalysisTree::EventHeader *mc_event = new AnalysisTree::EventHeader( Short_t(hasher(mc_event_branch.GetName())) );
  139. out_config->AddBranchConfig(tpc_tracks_branch);
  140. AnalysisTree::TrackDetector *tpc_tracks = new AnalysisTree::TrackDetector( Short_t(hasher(tpc_tracks_branch.GetName())) );
  141. out_config->AddBranchConfig(fhcal_branch);
  142. AnalysisTree::ModuleDetector *fhcal_modules = new AnalysisTree::ModuleDetector( Short_t(hasher(fhcal_branch.GetName())) );
  143. out_config->AddBranchConfig(mc_tracks_branch);
  144. AnalysisTree::Particles *mc_tracks = new AnalysisTree::Particles( Short_t(hasher(mc_tracks_branch.GetName())) );
  145. AnalysisTree::Matching *tpc2mc_tracks = new AnalysisTree::Matching(out_config->GetBranchConfig(str_tpc_tracks_branch).GetId(), out_config->GetBranchConfig(str_mc_tracks_branch).GetId());
  146. out_config->AddMatch(tpc2mc_tracks);
  147. reco_event->Init(reco_event_branch);
  148. mc_event->Init(mc_event_branch);
  149. // mc_event's additional field ids
  150. const int iB = out_config->GetBranchConfig(str_mc_event_branch).GetFieldId("B");
  151. const int iPhiRp = out_config->GetBranchConfig(str_mc_event_branch).GetFieldId("PhiRp");
  152. // tpc_tracks' additional field ids
  153. const int inhits = out_config->GetBranchConfig(str_tpc_tracks_branch).GetFieldId("nhits");
  154. const int inhits_poss = out_config->GetBranchConfig(str_tpc_tracks_branch).GetFieldId("nhits_poss");
  155. const int icharge = out_config->GetBranchConfig(str_tpc_tracks_branch).GetFieldId("charge");
  156. const int idcax = out_config->GetBranchConfig(str_tpc_tracks_branch).GetFieldId("dca_x");
  157. const int idcay = out_config->GetBranchConfig(str_tpc_tracks_branch).GetFieldId("dca_y");
  158. const int idcaz = out_config->GetBranchConfig(str_tpc_tracks_branch).GetFieldId("dca_z");
  159. const int ichi2 = out_config->GetBranchConfig(str_tpc_tracks_branch).GetFieldId("chi2");
  160. const int itof_mass2 = out_config->GetBranchConfig(str_tpc_tracks_branch).GetFieldId("tof_mass2");
  161. const int itof_flag = out_config->GetBranchConfig(str_tpc_tracks_branch).GetFieldId("tof_flag");
  162. const int idedx = out_config->GetBranchConfig(str_tpc_tracks_branch).GetFieldId("dedx");
  163. // mc_tracks' additional field ids
  164. const int imother_id = out_config->GetBranchConfig(str_mc_tracks_branch).GetFieldId("mother_id");
  165. // Create branches in the output tree
  166. outTree->Branch(str_reco_event_branch.c_str(), "AnalysisTree::EventHeader", &reco_event, 32000, 99);
  167. outTree->Branch(str_mc_event_branch.c_str(), "AnalysisTree::EventHeader", &mc_event, 32000, 99);
  168. outTree->Branch(str_tpc_tracks_branch.c_str(), "AnalysisTree::TrackDetector", &tpc_tracks, 256000, 99);
  169. outTree->Branch(str_fhcal_branch.c_str(), "AnalysisTree::ModuleDetector", &fhcal_modules, 128000, 99);
  170. outTree->Branch(str_mc_tracks_branch.c_str(), "AnalysisTree::Particles", &mc_tracks, 256000, 99);
  171. outTree->Branch(str_tpc2mc_tracks_branch.c_str(), "AnalysisTree::Matching", &tpc2mc_tracks, 32000, 99);
  172. // Printout basic configuration info
  173. out_config->Print();
  174. // Starting event loop
  175. TVector3 primaryVertex, mc_assoc_mom;
  176. std::set <Int_t> UsedMCTracks; // using to remap mc-reco track matching
  177. std::map <Int_t,Int_t> InitMcNewMcId; // map[old-mc-id] = new-mc-id
  178. Float_t FHCalSumEnergy[Num_Of_Modules];
  179. Int_t FHCalNumOfHits[Num_Of_Modules];
  180. Int_t n_entries = dstTree->GetEntriesFast();
  181. bool isGoodPID;
  182. int charge;
  183. for (int iEv = 0; iEv < n_entries; iEv++)
  184. {
  185. std::cout << "Event [" << iEv << "/" << n_entries << "]" << std::endl;
  186. dstTree->GetEntry(iEv);
  187. UsedMCTracks.clear();
  188. InitMcNewMcId.clear();
  189. tpc2mc_tracks->Clear();
  190. for (int i=0; i<Num_Of_Modules; i++)
  191. {
  192. FHCalSumEnergy[i] = 0.;
  193. FHCalNumOfHits[i] = 0;
  194. }
  195. tpc_tracks->ClearChannels();
  196. fhcal_modules->ClearChannels();
  197. mc_tracks->ClearChannels();
  198. // Reading Reco Event
  199. MpdVertex *vertex = (MpdVertex *)vertexes->First();
  200. vertex->Position(primaryVertex);
  201. reco_event->SetVertexPosition3(primaryVertex);
  202. // Read MC Event
  203. TVector3 vtx(MCHeader->GetX(),MCHeader->GetY(),MCHeader->GetZ());
  204. mc_event->SetVertexPosition3(vtx);
  205. mc_event->SetField(float(MCHeader->GetB()), iB);
  206. mc_event->SetField(float(MCHeader->GetRotZ()), iPhiRp);
  207. // Read energy in FHCal modules
  208. fhcal_modules->Reserve(Num_Of_Modules);
  209. for (int imodule=0; imodule<Num_Of_Modules; imodule++)
  210. {
  211. auto *module = fhcal_modules->AddChannel();
  212. module->Init(out_config->GetBranchConfig(str_fhcal_branch));
  213. module->SetSignal(0.f);
  214. }
  215. Int_t number_of_FHCal_hits = FHCalHits->GetEntriesFast();
  216. for(int ihit=0; ihit<number_of_FHCal_hits; ihit++)
  217. {
  218. FHCalHit = (MpdZdcDigi*) FHCalHits->UncheckedAt(ihit);
  219. Int_t DetId = FHCalHit->GetDetectorID();
  220. Int_t ModId = FHCalHit->GetModuleID()-1;
  221. Int_t ModNumber = ModId + (Num_Of_Modules/2) * (DetId-1);
  222. FHCalSumEnergy[ModNumber] += FHCalHit->GetELoss();
  223. FHCalNumOfHits[ModNumber]++;
  224. }
  225. for (int imodule=0; imodule<Num_Of_Modules; imodule++)
  226. {
  227. auto& module = fhcal_modules->Channel(imodule);
  228. // module.SetNumber(FHCalNumOfHits[imodule]); // Number of hits that got in the module
  229. module.SetNumber(imodule); // Number of the FHCal module
  230. module.SetSignal(FHCalSumEnergy[imodule]); // Total energy from hits in the module
  231. }
  232. // Reading Reco Tracks
  233. MpdGlobalTracks = (TClonesArray*)MPDEvent->GetGlobalTracks();
  234. Int_t Num_of_tpc_tracks = MpdGlobalTracks->GetEntriesFast();
  235. tpc_tracks->Reserve(Num_of_tpc_tracks);
  236. for (int itrack=0; itrack<Num_of_tpc_tracks; itrack++)
  237. {
  238. MpdTrack* mpdtrack = (MpdTrack*) MpdGlobalTracks->UncheckedAt(itrack);
  239. UsedMCTracks.insert(mpdtrack->GetID());
  240. auto *track = tpc_tracks->AddChannel();
  241. track->Init(out_config->GetBranchConfig(str_tpc_tracks_branch));
  242. track->SetMomentum(mpdtrack->GetPx(),mpdtrack->GetPy(),mpdtrack->GetPz());
  243. track->SetField(int(mpdtrack->GetNofHits()), inhits);
  244. track->SetField(int(mpdtrack->GetNofHitsPossTpc()), inhits_poss);
  245. track->SetField(float(mpdtrack->GetTofMass2()), itof_mass2);
  246. track->SetField(float(mpdtrack->GetTofFlag()), itof_flag);
  247. track->SetField(float(mpdtrack->GetdEdXTPC()), idedx);
  248. track->SetField(float(mpdtrack->GetChi2()), ichi2);
  249. track->SetField(float(mpdtrack->GetDCAX()), idcax);
  250. track->SetField(float(mpdtrack->GetDCAY()), idcay);
  251. track->SetField(float(mpdtrack->GetDCAZ()), idcaz);
  252. charge = (mpdtrack->GetPt() < 0) ? 1 : -1;
  253. track->SetField(int(charge), icharge);
  254. } // End of the tpc track loop
  255. // Read Mc tracks
  256. Int_t Num_of_mc_tracks = MCTracks->GetEntriesFast();
  257. mc_tracks->Reserve(Num_of_mc_tracks);
  258. for (int imctrack=0; imctrack<Num_of_mc_tracks; imctrack++)
  259. {
  260. #ifdef _MCSTACK_
  261. FairMCTrack *mctrack = (FairMCTrack*) MCTracks->UncheckedAt(imctrack);
  262. #endif
  263. #ifdef _MPDMCSTACK_
  264. MpdMCTrack *mctrack = (MpdMCTrack*) MCTracks->UncheckedAt(imctrack);
  265. #endif
  266. bool isUsed = (UsedMCTracks.count(imctrack));
  267. // If motherId != 1 and mc track doesn't relate to any reco track - skip
  268. if (mctrack->GetMotherId() != -1 && !isUsed) continue;
  269. auto *track = mc_tracks->AddChannel();
  270. track->Init(out_config->GetBranchConfig(str_mc_tracks_branch));
  271. // Collect new Mc Ids
  272. InitMcNewMcId[imctrack] = track->GetId();
  273. track->SetMomentum(mctrack->GetPx(), mctrack->GetPy(), mctrack->GetPz());
  274. track->SetPid(int(mctrack->GetPdgCode()));
  275. track->SetMass(float(mctrack->GetMass()));
  276. track->SetField(int(mctrack->GetMotherId()), imother_id);
  277. } // End of the mc track loop
  278. // reco-mc tracks matching
  279. for (int itrack=0; itrack<Num_of_tpc_tracks; itrack++)
  280. {
  281. MpdTrack* mpdtrack = (MpdTrack*) MpdGlobalTracks->UncheckedAt(itrack);
  282. const int tpc_track_id = itrack;
  283. const int mc_track_id = InitMcNewMcId[mpdtrack->GetID()];
  284. tpc2mc_tracks->AddMatch(tpc_track_id, mc_track_id);
  285. }
  286. outTree->Fill();
  287. } // End of the event loop
  288. outFile->cd();
  289. outTree->Print();
  290. outTree->Write();
  291. if (isDataHeaderDefined) dataHeader->Write("DataHeader");
  292. out_config->Write("Configuration");
  293. outFile->Close();
  294. timer.Stop();
  295. timer.Print();
  296. }
  297. // Additional functions
  298. float get_beamP(float sqrtSnn, float m_target, float m_beam)
  299. {
  300. return sqrt( pow((pow(sqrtSnn,2) - pow(m_target,2) - pow(m_beam,2))/(2*m_target), 2) - pow(m_beam,2) );
  301. }
  302. float GetFHCalPhi(int iModule)
  303. {
  304. const int Nmodules = 45;
  305. int xAxisSwitch = (iModule < Nmodules) ? 1 : -1;
  306. int module = (iModule < Nmodules) ? iModule : iModule - Nmodules;
  307. float x, y, phi = -999.;
  308. if (module >= 0 && module <= 4)
  309. {
  310. y = 45.;
  311. x = -1. * (module - 2) * 15.;
  312. phi = TMath::ATan2(y, x * xAxisSwitch);
  313. }
  314. else if ((module >= 5) && (module <= 39))
  315. {
  316. y = (3 - (module + 2) / 7) * 15.;
  317. x = (3 - (module + 2) % 7) * 15.;
  318. phi = TMath::ATan2(y, x * xAxisSwitch);
  319. }
  320. else if ((module >= 40) && (module <= 44))
  321. {
  322. y = -45.;
  323. x = -1. * (module - 42) * 15.;
  324. phi = TMath::ATan2(y, x * xAxisSwitch);
  325. }
  326. return phi;
  327. }
  328. TVector3 GetFHCalPos(int iModule)
  329. {
  330. const int Nmodules = 45;
  331. int xAxisSwitch = (iModule < Nmodules) ? 1 : -1;
  332. int module = (iModule < Nmodules) ? iModule : iModule - Nmodules;
  333. float x, y, z;
  334. z = (iModule < Nmodules) ? -320. : 320.;
  335. if (module >= 0 && module <= 4)
  336. {
  337. y = 45.;
  338. x = -1. * (module - 2) * 15.;
  339. }
  340. else if ((module >= 5) && (module <= 39))
  341. {
  342. y = (3 - (module + 2) / 7) * 15.;
  343. x = (3 - (module + 2) % 7) * 15.;
  344. }
  345. else if ((module >= 40) && (module <= 44))
  346. {
  347. y = -45.;
  348. x = -1. * (module - 42) * 15.;
  349. }
  350. TVector3 vec(x * xAxisSwitch, y, z);
  351. return vec;
  352. }
  353. Bool_t IsCharged(Int_t pdg)
  354. {
  355. auto particle = (TParticlePDG *)TDatabasePDG::Instance()->GetParticle(pdg);
  356. if (!particle)
  357. return false;
  358. return ( particle->Charge() != 0 );
  359. }
  360. Float_t GetRapidity(Float_t p, Float_t pz, Int_t pid_type)
  361. {
  362. Float_t mass; // in GeV/c^2
  363. if (pid_type == 0) mass = 0.13957; // pions
  364. if (pid_type == 1) mass = 0.493677; // kaons
  365. if (pid_type == 2) mass = 0.938272; // protons
  366. if (pid_type != 0 && pid_type != 1 && pid_type != 2) return -999.;
  367. Float_t E = TMath::Sqrt(p*p + mass*mass);
  368. return 0.5 * TMath::Log((E + pz)/(E-pz));
  369. }
  370. Float_t GetRapidityPDG(Float_t p, Float_t pz, Int_t pdg)
  371. {
  372. Float_t mass;
  373. auto pdgParticle = (TParticlePDG*)TDatabasePDG::Instance()->GetParticle(pdg);
  374. if (!pdgParticle) return -999.;
  375. mass = pdgParticle->Mass();
  376. Float_t E = TMath::Sqrt(p*p + mass*mass);
  377. return 0.5 * TMath::Log((E + pz)/(E - pz));
  378. }