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