PicoDstConverter.cpp 9.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307
  1. #include <iostream>
  2. #include <set>
  3. #include <map>
  4. #include <TROOT.h>
  5. #include <TChain.h>
  6. #include <TH1.h>
  7. #include <TF1.h>
  8. #include <TFile.h>
  9. #include <TTree.h>
  10. #include <TRandom1.h>
  11. #include <TClonesArray.h>
  12. #include <TFile.h>
  13. #include <TTree.h>
  14. #include <TString.h>
  15. #include <TMath.h>
  16. #include <TSystem.h>
  17. #include <TVector3.h>
  18. #include <TStopwatch.h>
  19. #include "FairMCEventHeader.h"
  20. #ifdef _OLD_MCSTACK_
  21. #include "FairMCTrack.h"
  22. #endif
  23. #ifdef _NEW_MCSTACK_
  24. #include "MpdMCTrack.h"
  25. #endif
  26. #include "MpdEvent.h"
  27. #include "MpdZdcDigi.h"
  28. #include "MpdPid.h"
  29. #include "MpdTrack.h"
  30. #include "MpdKalmanTrack.h"
  31. #include "MpdVertex.h"
  32. #include "PicoDstMCEvent.h"
  33. #include "PicoDstRecoEvent.h"
  34. #include "PicoDstMCTrack.h"
  35. #include "PicoDstRecoTrack.h"
  36. #include "PicoDstFHCal.h"
  37. int main(int argc, char **argv)
  38. {
  39. TString iFileName, oFileName;
  40. Bool_t writeAll = false;
  41. if (argc < 5)
  42. {
  43. std::cerr << "./PicoDstConverter -i inputfile -o outputfile [optional: --all. Writes all events (event with no reco tracks). False by default.]" << std::endl;
  44. return 1;
  45. }
  46. for (int i = 1; i < argc; i++)
  47. {
  48. if (std::string(argv[i]) != "-i" &&
  49. std::string(argv[i]) != "-o" &&
  50. std::string(argv[i]) != "--all")
  51. {
  52. std::cerr << "\n[ERROR]: Unknown parameter " << i << ": " << argv[i] << std::endl;
  53. return 1;
  54. }
  55. else
  56. {
  57. if (std::string(argv[i]) == "-i" && i != argc - 1)
  58. {
  59. iFileName = argv[++i];
  60. continue;
  61. }
  62. if (std::string(argv[i]) == "-i" && i == argc - 1)
  63. {
  64. std::cerr << "\n[ERROR]: Input file name was not specified " << std::endl;
  65. return 1;
  66. }
  67. if (std::string(argv[i]) == "-o" && i != argc - 1)
  68. {
  69. oFileName = argv[++i];
  70. continue;
  71. }
  72. if (std::string(argv[i]) == "-o" && i == argc - 1)
  73. {
  74. std::cerr << "\n[ERROR]: Output file name was not specified " << std::endl;
  75. return 1;
  76. }
  77. if (std::string(argv[i]) == "--all")
  78. {
  79. writeAll = true;
  80. continue;
  81. }
  82. }
  83. }
  84. TStopwatch timer;
  85. timer.Start();
  86. ///////////////////////////////////////////
  87. //TChain *dstTree;
  88. TFile *fi = new TFile(iFileName.Data(),"read");
  89. if (!fi || fi->IsZombie())
  90. {
  91. std::cerr << "ERROR: Input file probably is empty. Exit the program." << std::endl;
  92. return 100;
  93. }
  94. TTree *dstTree = (TTree*) fi->Get("mpdsim");
  95. TFile *outFile = new TFile(oFileName.Data(), "RECREATE");
  96. // TTree connfiguration
  97. TTree *outTree = new TTree("picodst","Pico Dst for flow calculations at MPD");
  98. const Double_t PIDsigM = 4.0;
  99. const Double_t PIDsigE = 4.0;
  100. const Double_t PIDenergy = 11.;
  101. const Double_t PIDkoeff = 1.;
  102. const TString PIDgenerator = "URQMD";
  103. const TString PIDtracking = "CF";
  104. const TString PIDparticles = "pikapr";
  105. const Int_t Num_Of_Modules = 90;
  106. MpdPid *pid = new MpdPid(PIDsigM, PIDsigE, PIDenergy, PIDkoeff, PIDgenerator, PIDtracking, PIDparticles);
  107. // PicoDst-related variables
  108. PicoDstMCEvent *mcEvent = new PicoDstMCEvent();
  109. PicoDstRecoEvent *recoEvent = new PicoDstRecoEvent();
  110. TClonesArray *mcTracks = new TClonesArray("PicoDstMCTrack");
  111. TClonesArray *recoTracks = new TClonesArray("PicoDstRecoTrack");
  112. TClonesArray *fhcalModules = new TClonesArray("PicoDstFHCal");
  113. outTree->Branch("mcevent.", &mcEvent, 32000, 99);
  114. outTree->Branch("recoevent.", &recoEvent, 32000, 99);
  115. outTree->Branch("mctracks", &mcTracks, 256000, 99);
  116. outTree->Branch("recotracks", &recoTracks, 256000, 99);
  117. outTree->Branch("FHCalModules", &fhcalModules, 128000, 99);
  118. mcTracks->BypassStreamer();
  119. recoTracks->BypassStreamer();
  120. fhcalModules->BypassStreamer();
  121. FairMCEventHeader *MCHeader;
  122. TClonesArray *MCTracks;
  123. MpdEvent *MPDEvent;
  124. TClonesArray *FHCalHits;
  125. TClonesArray *MpdGlobalTracks;
  126. MpdZdcDigi *FHCalHit;
  127. //TClonesArray *mpdKalmanTracks;
  128. TClonesArray *vertexes;
  129. MCHeader = nullptr;
  130. MCTracks = nullptr;
  131. MPDEvent = nullptr;
  132. FHCalHits = nullptr;
  133. //mpdKalmanTracks = nullptr;
  134. vertexes = nullptr;
  135. dstTree->SetBranchAddress("MCEventHeader.", &MCHeader);
  136. dstTree->SetBranchAddress("MCTrack", &MCTracks);
  137. dstTree->SetBranchAddress("MPDEvent.", &MPDEvent);
  138. dstTree->SetBranchAddress("ZdcDigi", &FHCalHits);
  139. //dstTree->SetBranchAddress("TpcKalmanTrack", &mpdKalmanTracks);
  140. dstTree->SetBranchAddress("Vertex", &vertexes);
  141. // Starting event loop
  142. TVector3 primaryVertex;
  143. std::set <Int_t> UsedMCTracks;
  144. std::map <Int_t,Int_t> InitMcNewMcId; // map[old-mc-id] = new-mc-id
  145. Float_t FHCalSumEnergy[Num_Of_Modules];
  146. Int_t n_entries = dstTree->GetEntriesFast();
  147. Bool_t isGoodPID;
  148. Short_t charge_mpd;
  149. for (int iEv = 0; iEv < n_entries; iEv++)
  150. {
  151. std::cout << "EVENT N " << iEv << std::endl;
  152. dstTree->GetEntry(iEv);
  153. mcTracks->Clear();
  154. recoTracks->Clear();
  155. UsedMCTracks.clear();
  156. InitMcNewMcId.clear();
  157. for (int i=0; i<Num_Of_Modules; i++)
  158. {
  159. FHCalSumEnergy[i] = 0.;
  160. }
  161. // Read MC Event
  162. mcEvent->SetB(MCHeader->GetB());
  163. mcEvent->SetPhiRP(MCHeader->GetRotZ());
  164. mcEvent->SetVertex(MCHeader->GetX(), MCHeader->GetY(), MCHeader->GetZ());
  165. // Reading Reco Event
  166. MpdVertex *vertex = (MpdVertex *)vertexes->First();
  167. vertex->Position(primaryVertex);
  168. recoEvent->SetVertex(primaryVertex.X(), primaryVertex.Y(), primaryVertex.Z());
  169. // Read energy reconstructed in the FHCal modules
  170. Int_t number_of_FHCal_hits = FHCalHits->GetEntriesFast();
  171. for(int ihit=0; ihit<number_of_FHCal_hits; ihit++)
  172. {
  173. FHCalHit = (MpdZdcDigi*) FHCalHits->UncheckedAt(ihit);
  174. Int_t DetId = FHCalHit->GetDetectorID();
  175. FHCalSumEnergy[(FHCalHit->GetModuleID()-1)+(Num_Of_Modules/2)*(DetId-1)] += FHCalHit->GetELoss();
  176. }
  177. // Fill fDetectorID = 1 part
  178. for (int imodule=0; imodule<Num_Of_Modules; imodule++)
  179. {
  180. PicoDstFHCal *fhcalModule = (PicoDstFHCal*)fhcalModules->ConstructedAt(imodule);
  181. fhcalModule->SetEnergy(FHCalSumEnergy[imodule]);
  182. }
  183. // Reading Reco Tracks
  184. MpdGlobalTracks = (TClonesArray*)MPDEvent->GetGlobalTracks();
  185. for (int itrack=0; itrack<MpdGlobalTracks->GetEntriesFast(); itrack++)
  186. {
  187. MpdTrack* mpdtrack = (MpdTrack*) MpdGlobalTracks->UncheckedAt(itrack);
  188. UsedMCTracks.insert(mpdtrack->GetID());
  189. PicoDstRecoTrack *recoTrack = (PicoDstRecoTrack*)recoTracks->ConstructedAt(itrack);
  190. if (writeAll) recoTrack->SetInitialMcId(mpdtrack->GetID());
  191. recoTrack->SetPx(mpdtrack->GetPx());
  192. recoTrack->SetPy(mpdtrack->GetPy());
  193. recoTrack->SetPz(mpdtrack->GetPz());
  194. recoTrack->SetNhits(mpdtrack->GetNofHits());
  195. recoTrack->SetNhitsPoss(mpdtrack->GetNofHitsPossTpc());
  196. recoTrack->SetTofMass2(mpdtrack->GetTofMass2());
  197. recoTrack->SetTofFlag(mpdtrack->GetTofFlag());
  198. recoTrack->SetTpcdEdx(mpdtrack->GetdEdXTPC());
  199. recoTrack->SetChi2(mpdtrack->GetChi2());
  200. recoTrack->SetDCAx(mpdtrack->GetDCAX());
  201. recoTrack->SetDCAy(mpdtrack->GetDCAY());
  202. recoTrack->SetDCAz(mpdtrack->GetDCAZ());
  203. charge_mpd = (mpdtrack->GetPt() < 0) ? 1 : -1;
  204. recoTrack->SetCharge(charge_mpd);
  205. // PID
  206. if (mpdtrack->GetTofFlag() == 2 || mpdtrack->GetTofFlag() == 6)
  207. {
  208. isGoodPID = pid->FillProbs(TMath::Abs(mpdtrack->GetPt()) * TMath::CosH(mpdtrack->GetEta()),
  209. mpdtrack->GetdEdXTPC(), mpdtrack->GetTofMass2(), charge_mpd);
  210. }
  211. else
  212. {
  213. isGoodPID = pid->FillProbs(TMath::Abs(mpdtrack->GetPt()) * TMath::CosH(mpdtrack->GetEta()),
  214. mpdtrack->GetdEdXTPC(), charge_mpd);
  215. }
  216. if (isGoodPID)
  217. {
  218. recoTrack->SetPidProbPion(pid->GetProbPi()); //mpdtrack->GetTPCPidProbPion();
  219. recoTrack->SetPidProbKaon(pid->GetProbKa()); //mpdtrack->GetTPCPidProbKaon();
  220. recoTrack->SetPidProbProton(pid->GetProbPr()); //mpdtrack->GetTPCPidProbProton();
  221. }
  222. else
  223. {
  224. recoTrack->SetPidProbPion(-999.);
  225. recoTrack->SetPidProbKaon(-999.);
  226. recoTrack->SetPidProbProton(-999.);
  227. }
  228. }
  229. // Reading MC Tracks
  230. int mcTrackCounter = 0;
  231. for (int imctrack=0; imctrack<MCTracks->GetEntriesFast(); imctrack++)
  232. {
  233. #ifdef _OLD_MCSTACK_
  234. FairMCTrack *mctrack = (FairMCTrack*) MCTracks->UncheckedAt(imctrack);
  235. #endif
  236. #ifdef _NEW_MCSTACK_
  237. MpdMCTrack *mctrack = (MpdMCTrack*) MCTracks->UncheckedAt(imctrack);
  238. #endif
  239. bool isUsed = (UsedMCTracks.count(imctrack));
  240. // If motherId != 1 and mc track doesn't relate to any reco track - skip
  241. if (mctrack->GetMotherId() != -1 && !isUsed) continue;
  242. if (isUsed)
  243. {
  244. InitMcNewMcId[imctrack] = mcTrackCounter;
  245. }
  246. PicoDstMCTrack *mcPicoTrack = (PicoDstMCTrack*)mcTracks->ConstructedAt(mcTrackCounter);
  247. if (writeAll) mcPicoTrack->SetInitialId(imctrack);
  248. mcPicoTrack->SetPx(mctrack->GetPx());
  249. mcPicoTrack->SetPy(mctrack->GetPy());
  250. mcPicoTrack->SetPz(mctrack->GetPz());
  251. mcPicoTrack->SetMotherId(mctrack->GetMotherId());
  252. mcPicoTrack->SetPdg(mctrack->GetPdgCode());
  253. mcPicoTrack->SetEnergy(mctrack->GetEnergy());
  254. mcTrackCounter++;
  255. }
  256. for (int itrack=0; itrack<recoTracks->GetEntriesFast(); itrack++)
  257. {
  258. PicoDstRecoTrack *recoTrack = (PicoDstRecoTrack*) recoTracks->UncheckedAt(itrack);
  259. MpdTrack* mpdtrack = (MpdTrack*) MpdGlobalTracks->UncheckedAt(itrack);
  260. recoTrack->SetMcId(InitMcNewMcId[mpdtrack->GetID()]);
  261. }
  262. // Skip if the event has no recontructed tracks
  263. if (recoTracks->GetEntriesFast() == 0 && !(writeAll)) continue;
  264. outTree->Fill();
  265. } // End of the Event loop
  266. outFile->cd();
  267. outTree->Print();
  268. outTree->Write();
  269. outFile->Close();
  270. timer.Stop();
  271. timer.Print();
  272. return 0;
  273. }