Coding rules (Markus)
[u/mrichter/AliRoot.git] / PWG3 / hfe / AliHFEpid.cxx
1 /**************************************************************************
2 * Copyright(c) 1998-1999, ALICE Experiment at CERN, All rights reserved. *
3 *                                                                        *
4 * Author: The ALICE Off-line Project.                                    *
5 * Contributors are mentioned in the code where appropriate.              *
6 *                                                                        *
7 * Permission to use, copy, modify and distribute this software and its   *
8 * documentation strictly for non-commercial purposes is hereby granted   *
9 * without fee, provided that the above copyright notice appears in all   *
10 * copies and that both the copyright notice and this permission notice   *
11 * appear in the supporting documentation. The authors make no claims     *
12 * about the suitability of this software for any purpose. It is          *
13 * provided "as is" without express or implied warranty.                  *
14 **************************************************************************/
15 //
16 // PID Steering Class 
17 // Interface to the user task
18 // 
19 // Authors:
20 //   Markus Fasel <M.Fasel@gsi.de>
21 //
22 #include <TAxis.h>
23 #include <TClass.h>
24 #include <THnSparse.h>
25 #include <TIterator.h>
26 #include <TList.h>
27 #include <TObjArray.h>
28 #include <TObjString.h>
29 #include <TString.h>
30
31 #include "AliESDtrack.h"
32 #include "AliLog.h"
33 #include "AliPID.h"
34
35 #include "AliHFEpid.h"
36 #include "AliHFEpidBase.h"
37 #include "AliHFEpidITS.h"
38 #include "AliHFEpidTPC.h"
39 #include "AliHFEpidTRD.h"
40 #include "AliHFEpidTOF.h"
41 #include "AliHFEpidMC.h"
42
43 ClassImp(AliHFEpid)
44
45 //____________________________________________________________
46 AliHFEpid::AliHFEpid():
47   fEnabledDetectors(0),
48   fPIDstrategy(0),
49   fQAlist(0x0),
50   fDebugLevel(0)
51 {
52   //
53   // Default constructor
54   //
55   memset(fDetectorPID, 0, sizeof(AliHFEpidBase *) * kNdetectorPID);
56 }
57
58 //____________________________________________________________
59 AliHFEpid::AliHFEpid(const AliHFEpid &c):
60   TObject(c),
61   fEnabledDetectors(c.fEnabledDetectors),
62   fPIDstrategy(0),
63   fQAlist(0x0),
64   fDebugLevel(c.fDebugLevel)
65 {
66   //
67   // Copy Constructor
68   //
69   memset(fDetectorPID, 0, sizeof(AliHFEpidBase *) * kNdetectorPID);
70   if(c.fDetectorPID[kMCpid])
71     fDetectorPID[kMCpid] = new AliHFEpidMC(*(dynamic_cast<AliHFEpidMC *>(c.fDetectorPID[kMCpid])));
72   if(c.fDetectorPID[kTPCpid])
73     fDetectorPID[kTPCpid] = new AliHFEpidTPC(*(dynamic_cast<AliHFEpidTPC *>(c.fDetectorPID[kTPCpid])));
74   if(c.fDetectorPID[kTRDpid])
75     fDetectorPID[kTRDpid] = new AliHFEpidTRD(*(dynamic_cast<AliHFEpidTRD *>(c.fDetectorPID[kTOFpid])));
76   if(c.fDetectorPID[kTOFpid])
77     fDetectorPID[kTOFpid] = new AliHFEpidTOF(*(dynamic_cast<AliHFEpidTOF *>(c.fDetectorPID[kTOFpid])));
78   if(c.IsQAOn()) SetQAOn();
79   if(c.HasMCData()) SetHasMCData(kTRUE);
80   for(Int_t idet = 0; idet < kNdetectorPID; idet++){
81     if(c.IsQAOn() && fDetectorPID[idet]) fDetectorPID[idet]->SetQAOn(fQAlist);
82     if(c.HasMCData() && fDetectorPID[idet]) fDetectorPID[idet]->SetHasMCData(kTRUE);
83   }
84 }
85
86 //____________________________________________________________
87 AliHFEpid& AliHFEpid::operator=(const AliHFEpid &c){
88   //
89   // Assignment operator
90   //
91   TObject::operator=(c);
92
93   if(this != &c){
94     fEnabledDetectors = c.fEnabledDetectors;
95     fPIDstrategy = c.fPIDstrategy;
96     fQAlist = 0x0;
97     fDebugLevel = c.fDebugLevel;
98   
99     memset(fDetectorPID, 0, sizeof(AliHFEpidBase *) * kNdetectorPID);
100     if(c.fDetectorPID[kMCpid])
101       fDetectorPID[kMCpid] = new AliHFEpidMC(*(dynamic_cast<AliHFEpidMC *>(c.fDetectorPID[kMCpid])));
102     if(c.fDetectorPID[kTPCpid])
103       fDetectorPID[kTPCpid] = new AliHFEpidTPC(*(dynamic_cast<AliHFEpidTPC *>(c.fDetectorPID[kTPCpid])));
104     if(c.fDetectorPID[kTRDpid])
105       fDetectorPID[kTRDpid] = new AliHFEpidTRD(*(dynamic_cast<AliHFEpidTRD *>(c.fDetectorPID[kTOFpid])));
106     if(c.fDetectorPID[kTOFpid])
107       fDetectorPID[kTOFpid] = new AliHFEpidTOF(*(dynamic_cast<AliHFEpidTOF *>(c.fDetectorPID[kTOFpid])));
108     if(c.IsQAOn()) SetQAOn();
109     if(c.HasMCData()) SetHasMCData(kTRUE);
110     for(Int_t idet = 0; idet < kNdetectorPID; idet++){
111       if(c.IsQAOn() && fDetectorPID[idet]) fDetectorPID[idet]->SetQAOn(fQAlist);
112       if(c.HasMCData() && fDetectorPID[idet]) fDetectorPID[idet]->SetHasMCData();
113     }
114   }
115   return *this; 
116 }
117
118 //____________________________________________________________
119 AliHFEpid::~AliHFEpid(){
120   //
121   // Destructor
122   //
123   for(Int_t idet = 0; idet < kNdetectorPID; idet++){
124     if(fDetectorPID[idet])
125       delete fDetectorPID[idet];
126   } 
127   if(fQAlist) delete fQAlist; fQAlist = 0x0;  // Each detector has to care about its Histograms
128 }
129
130 //____________________________________________________________
131 Bool_t AliHFEpid::InitializePID(TString arg){
132   //
133   // Initializes PID Object:
134   // + Defines which detectors to use
135   // + Initializes Detector PID objects
136   // + Handles QA
137   //
138   fDetectorPID[kMCpid] = new AliHFEpidMC("Monte Carlo PID"); // Always there
139   SETBIT(fEnabledDetectors, kMCpid);
140   // Initialize detector PIDs according to PID Strategies
141   if(arg.BeginsWith("Strategy")){
142     arg.ReplaceAll("Strategy", "");
143     fPIDstrategy = arg.Atoi();
144     AliInfo(Form("PID Strategy %d enabled", fPIDstrategy));
145     Int_t strategyStatus = kTRUE;
146     switch(fPIDstrategy){
147       case 1: InitStrategy1(); break;
148       case 2: InitStrategy2(); break;
149       case 3: InitStrategy3(); break;
150       case 4: InitStrategy4(); break;
151       case 5: InitStrategy5(); break;
152       default: strategyStatus = kFALSE;
153     }
154     return strategyStatus;
155   }
156   // No Strategy defined, Initialize according to detectors specified
157   AliInfo(Form("Doing InitializePID for Detectors %s end", arg.Data()));
158   fDetectorPID[kITSpid] = new AliHFEpidITS("ITS development PID");  // Development version of the ITS pid, for the moment always there
159   SETBIT(fEnabledDetectors, kITSpid);
160   
161   TObjArray *detsEnabled = arg.Tokenize(":");
162   TIterator *detIterator = detsEnabled->MakeIterator();
163   TObjString *det = 0x0;
164   while((det = dynamic_cast<TObjString *>(detIterator->Next()))){
165     if(det->String().CompareTo("TPC") == 0){
166       AliInfo("Doing TPC PID");
167       fDetectorPID[kTPCpid] = new AliHFEpidTPC("TPC PID");
168       SETBIT(fEnabledDetectors, kTPCpid);
169     } else if(det->String().CompareTo("TRD") == 0){
170       fDetectorPID[kTRDpid] = new AliHFEpidTRD("TRD PID");
171       SETBIT(fEnabledDetectors, kTRDpid);
172     } else if(det->String().CompareTo("TOF") == 0){
173       fDetectorPID[kTOFpid] = new AliHFEpidTOF("TOF PID");
174       SETBIT(fEnabledDetectors, kTOFpid);
175     }
176     // Here is still space for ESD PID
177   }
178   // Initialize PID Objects
179   Bool_t status = kTRUE;
180   for(Int_t idet = 0; idet < kNdetectorPID; idet++){
181     if(fDetectorPID[idet]){ 
182       status &= fDetectorPID[idet]->InitializePID();
183       if(IsQAOn() && status) fDetectorPID[idet]->SetQAOn(fQAlist);
184       if(HasMCData() && status) fDetectorPID[idet]->SetHasMCData();
185     }
186   }
187   return status;
188 }
189
190 //____________________________________________________________
191 Bool_t AliHFEpid::IsSelected(AliHFEpidObject *track){
192   //
193   // Steers PID decision for single detectors respectively combined
194   // PID decision
195   //
196   if(!track->fRecTrack){
197     // MC Event
198     return (TMath::Abs(fDetectorPID[kMCpid]->IsSelected(track)) == 11);
199   }
200   if(fPIDstrategy > 0 && fPIDstrategy < 6){
201     Int_t pid = 0;
202     switch(fPIDstrategy){
203       case 1: pid = IdentifyStrategy1(track); break;
204       case 2: pid = IdentifyStrategy2(track); break;
205       case 3: pid = IdentifyStrategy3(track); break;
206       case 4: pid = IdentifyStrategy4(track); break;
207       case 5: pid = IdentifyStrategy5(track); break;
208       default: break;
209     }
210     return pid;
211   }
212   if(TESTBIT(fEnabledDetectors, kTPCpid)){
213     if(IsQAOn() && fDebugLevel > 1){ 
214       AliInfo("Filling QA plots");
215       MakePlotsItsTpc(track);  // First fill the QA histograms
216     }
217     if(TESTBIT(fEnabledDetectors, kTOFpid)){
218       // case TPC-TOF
219       return MakePidTpcTof(track);
220     } else if(TESTBIT(fEnabledDetectors, kTRDpid)){
221       // case TPC-TRD with low level detector Signals
222       return MakePidTpcTrd(track);
223     } else
224       return (TMath::Abs(fDetectorPID[kTPCpid]->IsSelected(track)) ==11);
225   } else if(TESTBIT(fEnabledDetectors, kTRDpid)){
226     return (TMath::Abs(fDetectorPID[kTRDpid]->IsSelected(track)) ==11);
227   } else if(TESTBIT(fEnabledDetectors, kTOFpid)){
228     return (TMath::Abs(fDetectorPID[kTOFpid]->IsSelected(track)) ==11);
229   }
230   
231   return kFALSE;
232 }
233
234 //____________________________________________________________
235 Bool_t AliHFEpid::MakePidTpcTof(AliHFEpidObject *track){
236   //
237   // Combines TPC and TOF PID decision
238   //
239   if(fDetectorPID[kTOFpid]->IsSelected(track)) return fDetectorPID[kTPCpid]->IsSelected(track);
240   return kFALSE;
241 }
242
243 //____________________________________________________________
244 Bool_t AliHFEpid::MakePidTpcTrd(AliHFEpidObject *track){
245   //
246   // Combination of TPC and TRD PID
247   // Fills Histograms TPC Signal vs. TRD signal for different
248   // momentum slices
249   //
250   if(track->fAnalysisType != AliHFEpidObject::kESDanalysis) return kFALSE; //AOD based detector PID combination not yet implemented
251   AliDebug(1, "Analysis Type OK, do PID");
252   AliESDtrack *esdTrack = dynamic_cast<AliESDtrack *>(track->fRecTrack);
253   AliHFEpidTRD *trdPid = dynamic_cast<AliHFEpidTRD *>(fDetectorPID[kTRDpid]);
254   Int_t pdg = TMath::Abs(fDetectorPID[kMCpid]->IsSelected(track));
255   Int_t pid = -1;
256   switch(pdg){
257     case 11:    pid = AliPID::kElectron; break;
258     case 13:    pid = AliPID::kMuon; break;
259     case 211:   pid = AliPID::kPion; break;
260     case 321:   pid = AliPID::kKaon; break;
261     case 2212:  pid = AliPID::kProton; break;
262     default:    pid = -1;
263   };
264   Double_t content[10];
265   content[0] = pid;
266   content[1] = esdTrack->P();
267   content[2] = esdTrack->GetTPCsignal();
268   content[3] = trdPid->GetTRDSignalV1(esdTrack, pid);
269   if(IsQAOn() && fDebugLevel > 1){
270     content[4] = trdPid->GetTRDSignalV2(esdTrack, pid);
271     AliDebug(1, Form("Momentum: %f, TRD Signal: Method 1[%f], Method 2[%f]", content[1], content[3], content[4]));
272     (dynamic_cast<THnSparseF *>(fQAlist->At(kTRDSignal)))->Fill(content);
273   }
274   if(content[1] > 2){ // perform combined
275     AliDebug(1, "Momentum bigger 2 GeV/c, doing combined PID"); 
276     if(content[2] > 65 && content[3] > 500) return kTRUE;
277     else return kFALSE;
278   }
279   else {
280     AliDebug(1, "Momentum smaller 2GeV/c, doing TPC alone PID");
281     return fDetectorPID[kTPCpid]->IsSelected(track) == 11;
282   }
283 }
284
285 //____________________________________________________________
286 void AliHFEpid::MakePlotsItsTpc(AliHFEpidObject *track){
287   //
288   // Make a plot ITS signal - TPC signal for several momentum bins
289   //
290   if(track->fAnalysisType != AliHFEpidObject::kESDanalysis) return; //AOD based detector PID combination not yet implemented
291   AliESDtrack * esdTrack = dynamic_cast<AliESDtrack *>(track->fRecTrack);
292    // Fill My Histograms for MC PID
293   Int_t pdg = TMath::Abs(fDetectorPID[kMCpid]->IsSelected(track));
294   Int_t pid = -1;
295   switch(pdg){
296     case 11:    pid = AliPID::kElectron; break;
297     case 13:    pid = AliPID::kMuon; break;
298     case 211:   pid = AliPID::kPion; break;
299     case 321:   pid = AliPID::kKaon; break;
300     case 2212:  pid = AliPID::kProton; break;
301     default:    pid = -1;
302   };
303   if(IsQAOn() && fDebugLevel > 0){
304     Double_t content[10];
305     content[0] = pid;
306     content[1] = esdTrack->GetTPCInnerParam() ? esdTrack->GetTPCInnerParam()->P() : esdTrack->P();
307     content[2] = (dynamic_cast<AliHFEpidITS *>(fDetectorPID[kITSpid]))->GetITSSignalV1(track->fRecTrack, pid);
308     content[3] = esdTrack->GetTPCsignal();
309     AliDebug(1, Form("Momentum %f, TPC Signal %f, ITS Signal %f", content[1], content[2], content[3]));
310     (dynamic_cast<THnSparseF *>(fQAlist->At(kITSSignal)))->Fill(content);
311   }
312 }
313
314 //____________________________________________________________
315 void AliHFEpid::SetQAOn(){
316   //
317   // Switch on QA
318   //
319   SetBit(kIsQAOn, kTRUE);
320   AliInfo("QA switched on");
321   if(fQAlist) return;
322   fQAlist = new TList;
323   fQAlist->SetName("PIDqa");
324   THnSparseF *histo = NULL;
325
326   // Prepare axis for QA histograms
327   const Int_t kMomentumBins = 41;
328   const Double_t kPtMin = 0.1;
329   const Double_t kPtMax = 10.;
330   Double_t momentumBins[kMomentumBins];
331   for(Int_t ibin = 0; ibin < kMomentumBins; ibin++)
332     momentumBins[ibin] = static_cast<Double_t>(TMath::Power(10,TMath::Log10(kPtMin) + (TMath::Log10(kPtMax)-TMath::Log10(kPtMin))/(kMomentumBins-1)*static_cast<Double_t>(ibin)));
333
334   // Add Histogram for combined TPC-TRD PID
335   if(fDebugLevel > 1){
336     AliDebug(1, "Adding histogram for ITS-TPC investigation");
337     const Int_t kDimensionsTRDsig = 5;
338     Int_t kNbinsTRDsig[kDimensionsTRDsig] = {AliPID::kSPECIES + 1, kMomentumBins - 1, 200, 3000, 3000};
339     Double_t binMinTRDsig[kDimensionsTRDsig] = {-1., 0.1, 0, 0, 0};
340     Double_t binMaxTRDsig[kDimensionsTRDsig] = {AliPID::kSPECIES, 10., 200., 3000., 3000.};
341     fQAlist->AddAt((histo = new THnSparseF("fCombTPCTRDpid", "Combined TPC-TRD PID", kDimensionsTRDsig, kNbinsTRDsig, binMinTRDsig, binMaxTRDsig)), kTRDSignal);
342     histo->GetAxis(1)->Set(kMomentumBins - 1, momentumBins);
343     histo->GetAxis(0)->SetTitle("Particle Species");
344     histo->GetAxis(1)->SetTitle("p / GeV/c");
345     histo->GetAxis(2)->SetTitle("TPC Signal / a.u.");
346     histo->GetAxis(3)->SetTitle("TRD Signal / a.u.");
347     histo->GetAxis(4)->SetTitle("TRD Signal / a.u.");
348   }
349
350   // Add Histogram for combined TPC-ITS PID
351   if(fDebugLevel > 0){
352     AliDebug(1, "Adding histogram for TPC-TRD investigation");
353     const Int_t kDimensionsITSsig = 4;
354     Int_t kNbinsITSsig[kDimensionsITSsig] = {AliPID::kSPECIES + 1, kMomentumBins - 1, 300, 3000};
355     Double_t binMinITSsig[kDimensionsITSsig] = {-1., 0.1, 0., 0.};
356     Double_t binMaxITSsig[kDimensionsITSsig] = {AliPID::kSPECIES, 10., 300., 300.};
357     fQAlist->AddAt((histo = new THnSparseF("fCombTPCITSpid", "Combined TPC-ITS PID", kDimensionsITSsig, kNbinsITSsig, binMinITSsig, binMaxITSsig)), kITSSignal);
358     histo->GetAxis(1)->Set(kMomentumBins - 1, momentumBins);
359     histo->GetAxis(0)->SetTitle("Particle Species");
360     histo->GetAxis(1)->SetTitle("p / GeV/c");
361     histo->GetAxis(2)->SetTitle("ITS Signal / a.u.");
362     histo->GetAxis(3)->SetTitle("TPC Signal / a.u.");
363   }
364 }
365
366 //____________________________________________________________
367 void AliHFEpid::InitStrategy1(){
368   //
369   // TPC alone, 3-sigma cut
370   //
371   AliHFEpidTPC *pid = new AliHFEpidTPC("strat1TPCpid");
372   pid->SetTPCnSigma(3);
373   Bool_t status = pid->InitializePID();
374   if(IsQAOn() && status) pid->SetQAOn(fQAlist);
375   if(HasMCData() && status) pid->SetHasMCData();
376   fDetectorPID[kTPCpid] = pid;
377 }
378
379 //____________________________________________________________
380 void AliHFEpid::InitStrategy2(){
381   //
382   // TPC alone, symmetric 3 sigma cut and asymmetric sigma cut in the momentum region between 2GeV/c and 10 GeV/c and sigma between -1 and 100
383   //
384   AliHFEpidTPC *pid = new AliHFEpidTPC("strat2TPCpid");
385   pid->SetTPCnSigma(3);
386   pid->SetAsymmetricTPCsigmaCut(2., 10., 0., 4.);
387   Bool_t status = pid->InitializePID();
388   if(IsQAOn() && status) pid->SetQAOn(fQAlist);
389   if(HasMCData() && status) pid->SetHasMCData();
390   fDetectorPID[kTPCpid] = pid;
391 }
392
393 //____________________________________________________________
394 void AliHFEpid::InitStrategy3(){
395   //
396   // TPC alone, symmetric 3 sigma cut and 2 - -100 sigma pion rejection
397   //   
398   AliHFEpidTPC *pid = new AliHFEpidTPC("strat3TPCpid");
399   pid->SetTPCnSigma(3);
400   pid->SetRejectParticle(AliPID::kPion, 0., -100., 10., 3.);
401   Bool_t status = pid->InitializePID();
402   if(IsQAOn() && status) pid->SetQAOn(fQAlist);
403   if(HasMCData() && status) pid->SetHasMCData();
404   fDetectorPID[kTPCpid] = pid;
405 }
406
407 //____________________________________________________________
408 void AliHFEpid::InitStrategy4(){
409   //
410   // TPC and TRD combined, TPC 3 sigma cut and TRD NN 90% el efficiency level above 2 GeV/c
411   //
412   InitStrategy1();
413   AliHFEpidTRD *trdpid = new AliHFEpidTRD("strat4TRDpid");
414   Bool_t status = trdpid->InitializePID();
415   if(IsQAOn() && status) trdpid->SetQAOn(fQAlist);
416   if(HasMCData() && status) trdpid->SetHasMCData();
417   fDetectorPID[kTRDpid] = trdpid;
418 }
419
420 //____________________________________________________________
421 void AliHFEpid::InitStrategy5(){
422   //
423   // TPC and TRD combined, TPC 3 sigma cut and TRD NN 90% el efficiency level above 2 GeV/c
424   //
425   InitStrategy1();
426   AliHFEpidTRD *trdpid = new AliHFEpidTRD("strat5TRDpid");
427   Bool_t status = trdpid->InitializePID();
428   if(IsQAOn() && status) trdpid->SetQAOn(fQAlist);
429   if(HasMCData() && status) trdpid->SetHasMCData();
430   fDetectorPID[kTRDpid] = trdpid;
431 }
432
433 //____________________________________________________________
434 Bool_t AliHFEpid::IdentifyStrategy1(AliHFEpidObject *track){
435   return TMath::Abs(fDetectorPID[kTPCpid]->IsSelected(track)) == 11;
436 }
437
438 //____________________________________________________________
439 Bool_t AliHFEpid::IdentifyStrategy2(AliHFEpidObject *track){
440   return TMath::Abs(fDetectorPID[kTPCpid]->IsSelected(track)) == 11;
441 }
442
443 //____________________________________________________________
444 Bool_t AliHFEpid::IdentifyStrategy3(AliHFEpidObject *track){
445   return TMath::Abs(fDetectorPID[kTPCpid]->IsSelected(track)) == 11;
446 }
447
448 //____________________________________________________________
449 Bool_t AliHFEpid::IdentifyStrategy4(AliHFEpidObject *track){
450   Int_t trdpid = TMath::Abs(fDetectorPID[kTRDpid]->IsSelected(track));
451   return (trdpid == 0 || trdpid == 11) && (TMath::Abs(fDetectorPID[kTPCpid]->IsSelected(track)) == 11);
452 }
453
454 //____________________________________________________________
455 Bool_t AliHFEpid::IdentifyStrategy5(AliHFEpidObject *track){
456   return MakePidTpcTrd(track);
457 }