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