New directory for the VMC tests (Ivana, Eva)
[u/mrichter/AliRoot.git] / TRD / qaRec / AliTRDpidRefMaker.cxx
1 #include "TPDGCode.h"
2 #include "TFile.h"
3 #include "TTree.h"
4 #include "TEventList.h"
5 #include "TObjArray.h"
6 #include "TH2.h"
7
8 #include "AliLog.h"
9 #include "AliESDtrack.h"
10 #include "AliTrackReference.h"
11
12 #include "AliTRDReconstructor.h"
13 #include "AliTRDtrackV1.h"
14 #include "AliTRDseedV1.h"
15 #include "AliTRDpidRefMaker.h"
16 #include "info/AliTRDv0Info.h"
17
18
19 // Defines and implements basic functionality for building reference data for TRD PID. 
20 // 
21 // Here is the list of functionality provided by this class
22 // 1.
23 // 2.
24 // 
25 // Authors:
26 //   Alex Bercuci <A.Bercuci@gsi.de>
27 //   Alex Wilk    <wilka@uni-muenster.de>
28 //   Markus Fasel <mfasel@gsi.de>
29 //   Markus Heide <mheide@uni-muenster.de>
30 // 
31
32 ClassImp(AliTRDpidRefMaker)
33 ClassImp(AliTRDpidRefMaker::AliTRDpidRefData)
34 ClassImp(AliTRDpidRefMaker::AliTRDpidRefDataArray)
35
36 //________________________________________________________________________
37 AliTRDpidRefMaker::AliTRDpidRefDataArray::AliTRDpidRefDataArray() :
38   fNtracklets(0)
39   ,fData(0x0)
40 {
41   // Constructor of data array
42   fData = new AliTRDpidRefData[AliTRDgeometry::kNlayer];
43 }
44
45 //________________________________________________________________________
46 AliTRDpidRefMaker::AliTRDpidRefDataArray::~AliTRDpidRefDataArray()
47 {
48   // Destructor
49   delete [] fData;
50 }
51
52 //________________________________________________________________________
53 void AliTRDpidRefMaker::AliTRDpidRefDataArray::PushBack(Int_t ly, Int_t p, Float_t *dedx)
54 {
55 // Add PID data to the end of the array 
56   fData[fNtracklets].fPLbin= (ly<<4) | (p&0xf);
57   memcpy(fData[fNtracklets].fdEdx, dedx, 8*sizeof(Float_t));
58   fNtracklets++;
59 }
60
61 //________________________________________________________________________
62 void AliTRDpidRefMaker::AliTRDpidRefDataArray::Reset()
63 {
64 // Reset content
65
66   if(!fNtracklets) return;
67   while(fNtracklets--){
68     fData[fNtracklets].fPLbin = 0xff;
69     memset(fData[fNtracklets].fdEdx, 0, 8*sizeof(Float_t));
70   }
71   fNtracklets=0;
72 }
73
74 //________________________________________________________________________
75 AliTRDpidRefMaker::AliTRDpidRefMaker(const char *name, const char *title) 
76   :AliTRDrecoTask(name, title)
77   ,fReconstructor(0x0)
78   ,fV0s(0x0)
79   ,fData(0x0)
80   ,fPIDdataArray(0x0)
81   ,fRefPID(kMC)
82   ,fRefP(kMC)
83   ,fPIDbin(0xff)
84   ,fFreq(1.)
85   ,fP(-1.)
86   ,fPthreshold(0.5)
87 {
88   //
89   // Default constructor
90   //
91
92   fReconstructor = new AliTRDReconstructor();
93   fReconstructor->SetRecoParam(AliTRDrecoParam::GetLowFluxParam());
94   memset(fdEdx, 0, 10*sizeof(Float_t));
95   memset(fPID, 0, AliPID::kSPECIES*sizeof(Float_t));
96
97   DefineInput(1, TObjArray::Class());
98   DefineOutput(1, TTree::Class());
99 }
100
101
102 //________________________________________________________________________
103 AliTRDpidRefMaker::~AliTRDpidRefMaker() 
104 {
105   if(fPIDdataArray) delete fPIDdataArray;
106   if(fReconstructor) delete fReconstructor;
107 }
108
109 //________________________________________________________________________
110 Bool_t AliTRDpidRefMaker::CheckQuality(AliTRDseedV1* /*trklt*/)
111 {
112 // Place holder for checking tracklet quality for PID.
113   return kTRUE;  
114 }
115
116
117 //________________________________________________________________________
118 Float_t* AliTRDpidRefMaker::CookdEdx(AliTRDseedV1 *trklt)
119 {
120   trklt->CookdEdx(AliTRDpidUtil::kNNslices);
121   memcpy(fdEdx, trklt->GetdEdx(), AliTRDpidUtil::kNNslices*sizeof(Float_t));
122   return fdEdx;
123 }
124
125
126 //________________________________________________________________________
127 void AliTRDpidRefMaker::ConnectInputData(Option_t *opt)
128 {
129   AliTRDrecoTask::ConnectInputData(opt);
130   fV0s = dynamic_cast<TObjArray*>(GetInputData(1));
131 }
132
133 //________________________________________________________________________
134 void AliTRDpidRefMaker::CreateOutputObjects()
135 {
136   // Create histograms
137   // Called once
138
139   OpenFile(0, "RECREATE");
140   fContainer = new TObjArray();
141   fContainer->SetName(Form("Moni%s", GetName()));
142
143   TH2 *h2 = new TH2I("hPDG","Particle abundance", AliPID::kSPECIES, -0.5, 4.5, AliTRDCalPID::kNMom, -0.5, AliTRDCalPID::kNMom-0.5);
144   TAxis *ax = h2->GetXaxis();
145   ax->SetNdivisions(505);
146   ax->SetTitle("Particle species");
147   for(Int_t is=AliPID::kSPECIES; is--;) ax->SetBinLabel(is+1, AliPID::ParticleShortName(is));
148   h2->GetYaxis()->SetTitle("P bins");
149   h2->GetYaxis()->SetNdivisions(511);
150
151   fContainer->AddAt(h2, 0);
152
153   fData = new TTree(GetName(), Form("Reference data for %s", GetName()));
154   fData->Branch("s", &fPIDbin, "s/b");
155   fData->Branch("data", &fPIDdataArray);
156 }
157
158 //________________________________________________________________________
159 void AliTRDpidRefMaker::Exec(Option_t *) 
160 {
161   // Main loop
162   // Called for each event
163
164   AliTRDtrackInfo     *track = 0x0;
165   AliTRDtrackV1    *trackTRD = 0x0;
166   AliTrackReference     *ref = 0x0;
167   const AliTRDtrackInfo::AliESDinfo *infoESD = 0x0;
168   for(Int_t itrk=0; itrk<fTracks->GetEntriesFast(); itrk++){
169     track = (AliTRDtrackInfo*)fTracks->UncheckedAt(itrk);
170     if(!track->HasESDtrack()) continue;
171     trackTRD = track->GetTrack();
172     infoESD  = track->GetESDinfo();
173     Double32_t *infoPID = infoESD->GetSliceIter();
174     Int_t n = infoESD->GetNSlices() - AliTRDgeometry::kNlayer;
175     Double32_t *p = &infoPID[n];
176
177
178     ULong_t status = track->GetStatus();
179     if(!(status&AliESDtrack::kTPCout)) continue;
180
181     // fill the pid information
182     SetRefPID(fRefPID, track, fPID);
183
184     // prepare PID data array
185     if(!fPIDdataArray){ 
186       fPIDdataArray = new AliTRDpidRefDataArray();
187     } else fPIDdataArray->Reset();
188
189     // fill PID information
190     for(Int_t ily = 0; ily < AliTRDgeometry::kNlayer; ily++){
191
192       // fill P & dE/dx information
193       if(HasFriends()){ // from TRD track
194         if(!trackTRD) continue;
195         AliTRDseedV1 *trackletTRD(0x0);
196         trackTRD -> SetReconstructor(fReconstructor);
197         if(!(trackletTRD = trackTRD -> GetTracklet(ily))) continue;
198         if(!CheckQuality(trackletTRD)) continue;
199         if(!CookdEdx(trackletTRD)) continue;
200
201         // fill momentum information
202         fP = 0.;
203         switch(fRefP){
204         case kMC:
205           if(!(ref = track->GetTrackRef(trackletTRD))) continue;
206           fP = ref->P();
207           break;
208         case kRec:
209           fP = trackletTRD->GetMomentum();
210           break;
211         default: continue;
212         }
213       } else { // from ESD track
214         // fill momentum information
215         switch(fRefP){
216         case kMC:
217           if(!(ref = track->GetTrackRef(ily))) continue;
218           fP = ref->P();
219           break;
220         case kRec:
221           fP = p[ily];
222           break;
223         default: continue;
224         } 
225         Double32_t *it = &infoPID[ily*AliTRDCalPID::kNSlicesNN];
226         for(Int_t is=AliTRDCalPID::kNSlicesNN; is--; it++) fdEdx[is] = (*it);
227       }
228
229       // momentum threshold
230       if(fP < fPthreshold) continue;
231
232       // store information
233       fPIDdataArray->PushBack(ily, AliTRDpidUtil::GetMomentumBin(fP), fdEdx);
234     }
235
236     Fill();
237   }
238
239   PostData(0, fContainer);
240   PostData(1, fData);
241 }
242
243
244 //________________________________________________________________________
245 void AliTRDpidRefMaker::Fill() 
246 {
247 // Fill data tree
248
249   if(!fPIDdataArray->fNtracklets) return;
250   fPIDbin = TMath::LocMax(AliPID::kSPECIES, fPID); // get particle type
251 // Fill data tree
252   fData->Fill();
253
254   
255   // fill monitor
256   for(Int_t ily=fPIDdataArray->fNtracklets; ily--;){
257     Int_t pBin = fPIDdataArray->fData[ily].fPLbin & 0xf;
258     ((TH2*)fContainer->At(0))->Fill(fPIDbin, pBin);
259   }
260 }
261
262 //________________________________________________________________________
263 void AliTRDpidRefMaker::LinkPIDdata() 
264 {
265 // Link data tree to data members
266   fData->SetBranchAddress("s", &fPIDbin);
267   fData->SetBranchAddress("data", &fPIDdataArray);
268 }
269
270 //________________________________________________________________________
271 void AliTRDpidRefMaker::SetRefPID(ETRDpidRefMakerSource select, void *source, Float_t *pid) 
272 {
273 // Fill the reference PID values "pid" from "source" object
274 // according to the option "select". Possible options are
275 // - kV0  - v0 based PID
276 // - kMC  - MC truth [default]
277 // - kRec - outside detectors
278
279
280   memset(fPID, 0, AliPID::kSPECIES*sizeof(Float_t));
281   switch(select){ 
282   case kV0:
283     {
284       AliTRDtrackInfo *track = static_cast<AliTRDtrackInfo*>(source);
285       if(!track){
286         AliError("No trackInfo found");
287         return;
288       }
289
290       //Get V0 PID decisions from the AliTRDv0Info for all particle species (implemented so far : electrons from conversions, pions from K0s and protons from Lambdas) :
291       AliTRDv0Info v0;
292       for(Int_t is=AliPID::kSPECIES; is--;) fPID[is] = v0.GetV0PID(is, track);
293     }
294     break;
295   case kMC:
296     if(!HasMCdata()){
297       AliError("Could not retrive reference PID from MC");
298       return;
299     }
300     {
301       AliTRDtrackInfo *track = static_cast<AliTRDtrackInfo*>(source);
302       switch(track->GetPDG()){
303       case kElectron:
304       case kPositron:
305         fPID[AliPID::kElectron] = 1.;
306         break;
307       case kMuonPlus:
308       case kMuonMinus:
309         fPID[AliPID::kMuon] = 1.;
310         break;
311       case kPiPlus:
312       case kPiMinus:
313         fPID[AliPID::kPion] = 1.;
314         break;
315       case kKPlus:
316       case kKMinus:
317         fPID[AliPID::kKaon] = 1.;
318         break;
319       case kProton:
320       case kProtonBar:
321         fPID[AliPID::kProton] = 1.;
322         break;
323       }
324     }
325     break;
326   case kRec:
327     { 
328       AliTRDtrackInfo *track = static_cast<AliTRDtrackInfo*>(source);
329       AliTRDtrackV1 *trackTRD = track->GetTrack();
330       trackTRD -> SetReconstructor(fReconstructor);
331       //fReconstructor -> SetOption("nn");
332       trackTRD -> CookPID();
333       for(Int_t iPart = 0; iPart < AliPID::kSPECIES; iPart++){
334         pid[iPart] = trackTRD -> GetPID(iPart);
335         AliDebug(4, Form("PDG is (in V0info) %d %f", iPart, pid[iPart]));
336       }
337     }
338     break;
339   default:
340     AliWarning("PID reference source not implemented");
341     return;
342   }
343   AliDebug(4, Form("Ref PID [%] : %s[%5.2f] %s[%5.2f] %s[%5.2f] %s[%5.2f] %s[%5.2f]"
344     ,AliPID::ParticleShortName(0), 1.e2*fPID[0]
345     ,AliPID::ParticleShortName(1), 1.e2*fPID[1]
346     ,AliPID::ParticleShortName(2), 1.e2*fPID[2]
347     ,AliPID::ParticleShortName(3), 1.e2*fPID[3]
348     ,AliPID::ParticleShortName(4), 1.e2*fPID[4]
349   ));
350 }
351
352 //________________________________________________________________________
353 void AliTRDpidRefMaker::SetAbundance(Float_t train) 
354 {
355 // Split data sample between trainning and testing
356
357   if(train<0. || train >1.){
358     AliWarning("The input data should be in the interval [0, 1]");
359     return;
360   }
361
362   fFreq = train; 
363 }
364