o update dielectron package
[u/mrichter/AliRoot.git] / PWGDQ / dielectron / AliDielectronPID.cxx
1 /*************************************************************************
2 * Copyright(c) 1998-2009, 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 ///////////////////////////////////////////////////////////////////////////
17 //                Dielectron PID                                  //
18 //                                                                       //
19 //                                                                       //
20 /*
21 Detailed description
22
23
24 */
25 //                                                                       //
26 ///////////////////////////////////////////////////////////////////////////
27
28 #include <TMath.h>
29 #include <TF1.h>
30 #include <TGraph.h>
31
32 #include <AliVTrack.h>
33 #include <AliVCluster.h>
34 #include <AliLog.h>
35 #include <AliExternalTrackParam.h>
36 #include <AliPIDResponse.h>
37 #include <AliESDtrack.h> //!!!!! Remove once Eta correction is treated in the tender
38
39 #include "AliDielectronVarManager.h"
40
41 #include "AliDielectronPID.h"
42
43 ClassImp(AliDielectronPID)
44
45 TGraph *AliDielectronPID::fgFitCorr=0x0;
46 Double_t AliDielectronPID::fgCorr=0.0;
47 TF1 *AliDielectronPID::fgFunEtaCorr=0x0;
48
49 AliDielectronPID::AliDielectronPID() :
50   AliAnalysisCuts(),
51   fNcuts(0),
52   fPIDResponse(0x0)
53 {
54   //
55   // Default Constructor
56   //
57   for (Int_t icut=0; icut<kNmaxPID; ++icut){
58     fDetType[icut]=kTPC;
59     fPartType[icut]=AliPID::kPion;
60     fNsigmaLow[icut]=0;
61     fNsigmaUp[icut]=0;
62     fmin[icut]=0;
63     fmax[icut]=0;
64     fExclude[icut]=kFALSE;
65     fFunUpperCut[icut]=0x0;
66     fFunLowerCut[icut]=0x0;
67     fRequirePIDbit[icut]=0;
68     fActiveCuts[icut]=-1;
69   }
70 }
71
72 //______________________________________________
73 AliDielectronPID::AliDielectronPID(const char* name, const char* title) :
74   AliAnalysisCuts(name, title),
75   fNcuts(0),
76   fPIDResponse(0x0)
77 {
78   //
79   // Named Constructor
80   //
81   for (Int_t icut=0; icut<kNmaxPID; ++icut){
82     fDetType[icut]=kTPC;
83     fPartType[icut]=AliPID::kPion;
84     fNsigmaLow[icut]=0;
85     fNsigmaUp[icut]=0;
86     fmin[icut]=0;
87     fmax[icut]=0;
88     fExclude[icut]=kFALSE;
89     fFunUpperCut[icut]=0x0;
90     fFunLowerCut[icut]=0x0;
91     fRequirePIDbit[icut]=0;
92     fActiveCuts[icut]=0;
93   }
94 }
95
96 //______________________________________________
97 AliDielectronPID::~AliDielectronPID()
98 {
99   //
100   // Default Destructor
101   //
102 }
103
104 //______________________________________________
105 void AliDielectronPID::AddCut(DetType det, AliPID::EParticleType type, Double_t nSigmaLow, Double_t nSigmaUp/*=-99999.*/,
106                               Double_t min/*=0*/, Double_t max/*=0*/, Bool_t exclude/*=kFALSE*/,
107                               UInt_t pidBitType/*=AliDielectronPID::kRequire*/, Int_t var/*=-1*/)
108 {
109   //
110   // Add a pid nsigma cut
111   // use response of detector 'det' in the range ['min'] to ['max'] for var
112   // use a sigma band between 'nSigmaLow' and 'nSigmaUp'
113   // if nSigmaUp==-99999. then nSigmaLow will be uesd as a symmetric band +- nSigmaLow
114   // specify whether to 'exclude' the given band
115   //
116
117   if (fNcuts>=kNmaxPID){
118     AliError(Form("only %d pid cut ranges allowed",kNmaxPID));
119     return;
120   }
121   if (TMath::Abs(nSigmaUp+99999.)<1e-20){
122     nSigmaUp=TMath::Abs(nSigmaLow);
123     nSigmaLow=-1.*nSigmaUp;
124   }
125   fDetType[fNcuts]=det;
126   fPartType[fNcuts]=type;
127   fNsigmaLow[fNcuts]=nSigmaLow;
128   fNsigmaUp[fNcuts]=nSigmaUp;
129   fmin[fNcuts]=min;
130   fmax[fNcuts]=max;
131   fExclude[fNcuts]=exclude;
132   fRequirePIDbit[fNcuts]=pidBitType;
133   fActiveCuts[fNcuts]=(var==-1 ? AliDielectronVarManager::kP : var);
134
135   AliInfo(Form("Add PID cut %d: sigma [% .1f,% .1f] \t cut [% .1f,% .f] \t var %d->%s \n",
136                fNcuts,nSigmaLow,nSigmaUp,min,max,fActiveCuts[fNcuts],AliDielectronVarManager::GetValueName(fActiveCuts[fNcuts])));  
137   
138   ++fNcuts;
139
140 }
141
142 //______________________________________________
143 void AliDielectronPID::AddCut(DetType det, AliPID::EParticleType type, Double_t nSigmaLow, TF1 * const funUp,
144                               Double_t min/*=0*/, Double_t max/*=0*/, Bool_t exclude/*=kFALSE*/,
145                               UInt_t pidBitType/*=AliDielectronPID::kRequire*/, Int_t var/*=-1*/)
146 {
147   //
148   // cut using a TF1 as upper cut
149   //
150   if (funUp==0x0){
151     AliError("A valid function is required for the upper cut. Not adding the cut!");
152     return;
153   }
154   fFunUpperCut[fNcuts]=funUp;
155   AddCut(det,type,nSigmaLow,0.,min,max,exclude,pidBitType,var);
156 }
157
158 //______________________________________________
159 void AliDielectronPID::AddCut(DetType det, AliPID::EParticleType type, TF1 * const funLow, Double_t nSigmaUp,
160                               Double_t min/*=0*/, Double_t max/*=0*/, Bool_t exclude/*=kFALSE*/,
161                               UInt_t pidBitType/*=AliDielectronPID::kRequire*/, Int_t var/*=-1*/)
162 {
163   //
164   // cut using a TF1 as lower cut
165   //
166   if (funLow==0x0){
167     AliError("A valid function is required for the lower cut. Not adding the cut!");
168     return;
169   }
170   fFunLowerCut[fNcuts]=funLow;
171   AddCut(det,type,0.,nSigmaUp,min,max,exclude,pidBitType,var);
172 }
173
174 //______________________________________________
175 void AliDielectronPID::AddCut(DetType det, AliPID::EParticleType type, TF1 * const funLow, TF1 * const funUp,
176                               Double_t min/*=0*/, Double_t max/*=0*/, Bool_t exclude/*=kFALSE*/,
177                               UInt_t pidBitType/*=AliDielectronPID::kRequire*/, Int_t var/*=-1*/)
178 {
179   //
180   // cut using a TF1 as lower and upper cut
181   //
182   if ( (funUp==0x0) || (funLow==0x0) ){
183     AliError("A valid function is required for upper and lower cut. Not adding the cut!");
184     return;
185   }
186   fFunUpperCut[fNcuts]=funUp;
187   fFunLowerCut[fNcuts]=funLow;
188   AddCut(det,type,0.,0.,min,max,exclude,pidBitType,var);
189 }
190
191 //______________________________________________
192 Bool_t AliDielectronPID::IsSelected(TObject* track)
193 {
194   //
195   // perform PID cuts
196   //
197
198   //loop over all cuts
199   AliVTrack *part=static_cast<AliVTrack*>(track);
200   AliESDtrack *esdTrack=0x0;
201   Double_t origdEdx=-1;
202   
203   // apply ETa correction, remove once this is in the tender
204   if( (part->IsA() == AliESDtrack::Class()) && fgFunEtaCorr ){
205     esdTrack=static_cast<AliESDtrack*>(part);
206     origdEdx=esdTrack->GetTPCsignal();
207     esdTrack->SetTPCsignal(origdEdx/GetEtaCorr(esdTrack),esdTrack->GetTPCsignalSigma(),esdTrack->GetTPCsignalN());
208   }
209
210   
211   //Fill values
212   Double_t values[AliDielectronVarManager::kNMaxValues];
213   AliDielectronVarManager::Fill(track,values);
214
215   Bool_t selected=kFALSE;
216   fPIDResponse=AliDielectronVarManager::GetPIDResponse();
217   for (UChar_t icut=0; icut<fNcuts; ++icut){
218     Double_t min=fmin[icut];
219     Double_t max=fmax[icut];
220     Double_t val=values[fActiveCuts[icut]];
221     
222     // test var range. In case min==max do not cut
223     if ( (TMath::Abs(min-max)>1e-20) && (val<min || val>=max) ) continue;
224     
225     switch (fDetType[icut]){
226     case kITS:
227       selected = IsSelectedITS(part,icut);
228       break;
229     case kTPC:
230       selected = IsSelectedTPC(part,icut);
231       break;
232     case kTRD:
233       selected = IsSelectedTRD(part,icut);
234       break;
235     case kTRDeleEff:
236       selected = IsSelectedTRDeleEff(part,icut);
237       break;
238     case kTOF:
239       selected = IsSelectedTOF(part,icut);
240       break;
241     case kEMCAL:
242       selected = IsSelectedEMCAL(part,icut);
243       break;
244     }
245     if (!selected) {
246       if (esdTrack) esdTrack->SetTPCsignal(origdEdx,esdTrack->GetTPCsignalSigma(),esdTrack->GetTPCsignalN());
247
248       return kFALSE;
249     }
250   }
251
252   if (esdTrack) esdTrack->SetTPCsignal(origdEdx,esdTrack->GetTPCsignalSigma(),esdTrack->GetTPCsignalN());
253   return selected;
254 }
255
256 //______________________________________________
257 Bool_t AliDielectronPID::IsSelectedITS(AliVTrack * const part, Int_t icut)
258 {
259   //
260   // ITS part of the PID check
261   // Don't accept the track if there was no pid bit set
262   //
263   
264   if (fRequirePIDbit[icut]==AliDielectronPID::kRequire&&!(part->GetStatus()&AliVTrack::kITSpid)) return kFALSE;
265   if (fRequirePIDbit[icut]==AliDielectronPID::kIfAvailable&&!(part->GetStatus()&AliVTrack::kITSpid)) return kTRUE;
266
267   Double_t mom=part->P();
268   
269   Float_t numberOfSigmas=fPIDResponse->NumberOfSigmasITS(part, fPartType[icut]);
270   
271   // test if we are supposed to use a function for the cut
272   if (fFunUpperCut[icut]) fNsigmaUp[icut] =fFunUpperCut[icut]->Eval(mom);
273   if (fFunLowerCut[icut]) fNsigmaLow[icut]=fFunLowerCut[icut]->Eval(mom);
274   
275   Bool_t selected=((numberOfSigmas>=fNsigmaLow[icut])&&(numberOfSigmas<=fNsigmaUp[icut]))^fExclude[icut];
276   return selected;
277 }
278
279 //______________________________________________
280 Bool_t AliDielectronPID::IsSelectedTPC(AliVTrack * const part, Int_t icut)
281 {
282   //
283   // TPC part of the PID check
284   // Don't accept the track if there was no pid bit set
285   //
286   if (fRequirePIDbit[icut]==AliDielectronPID::kRequire&&!(part->GetStatus()&AliVTrack::kTPCpid)) return kFALSE;
287   if (fRequirePIDbit[icut]==AliDielectronPID::kIfAvailable&&!(part->GetStatus()&AliVTrack::kTPCpid)) return kTRUE;
288
289   Double_t mom=part->GetTPCmomentum();
290   
291   Float_t numberOfSigmas=fPIDResponse->NumberOfSigmasTPC(part, fPartType[icut]);
292
293   if (fPartType[icut]==AliPID::kElectron){
294     numberOfSigmas-=fgCorr;
295   }
296   
297   // test if we are supposed to use a function for the cut
298   if (fFunUpperCut[icut]) fNsigmaUp[icut] =fFunUpperCut[icut]->Eval(mom);
299   if (fFunLowerCut[icut]) fNsigmaLow[icut]=fFunLowerCut[icut]->Eval(mom);
300   
301   Bool_t selected=((numberOfSigmas>=fNsigmaLow[icut])&&(numberOfSigmas<=fNsigmaUp[icut]))^fExclude[icut];
302   return selected;
303 }
304
305 //______________________________________________
306 Bool_t AliDielectronPID::IsSelectedTRD(AliVTrack * const part, Int_t icut)
307 {
308   //   
309   // TRD part of the pid check
310   // the TRD checks on the probabilities.
311   //
312
313   if (fRequirePIDbit[icut]==AliDielectronPID::kRequire&&!(part->GetStatus()&AliVTrack::kTRDpid)) return kFALSE;
314   if (fRequirePIDbit[icut]==AliDielectronPID::kIfAvailable&&!(part->GetStatus()&AliVTrack::kTRDpid)) return kTRUE;
315   if (fRequirePIDbit[icut]==AliDielectronPID::kIfAvailable && (part->GetTRDntrackletsPID()<4)) return kTRUE;
316
317   Double_t p[AliPID::kSPECIES]={0.};
318   fPIDResponse->ComputeTRDProbability(part,AliPID::kSPECIES,p);
319   Float_t particleProb=p[fPartType[icut]];
320
321   Bool_t selected=((particleProb>=fNsigmaLow[icut])&&(particleProb<=fNsigmaUp[icut]))^fExclude[icut];
322   return selected;
323 }
324
325 //______________________________________________
326 Bool_t AliDielectronPID::IsSelectedTRDeleEff(AliVTrack * const part, Int_t icut)
327 {
328   //
329   // TRD part of the pid check using electron efficiency requirement
330   // in this case the upper limit as well as the particle specie is ignored 
331   //   and the lower limit regarded as the requested electron efficiency
332   //
333
334   if (fRequirePIDbit[icut]==AliDielectronPID::kRequire&&!(part->GetStatus()&AliVTrack::kTRDpid)) return kFALSE;
335   if (fRequirePIDbit[icut]==AliDielectronPID::kIfAvailable&&!(part->GetStatus()&AliVTrack::kTRDpid)) return kTRUE;
336
337   Bool_t selected=fPIDResponse->IdentifiedAsElectronTRD(part,fNsigmaLow[icut])^fExclude[icut];
338   return selected;
339 }
340
341 //______________________________________________
342 Bool_t AliDielectronPID::IsSelectedTOF(AliVTrack * const part, Int_t icut)
343 {
344   //
345   // TOF part of the PID check
346   // Don't accept the track if there was no pid bit set
347   //
348   
349   if (fRequirePIDbit[icut]==AliDielectronPID::kRequire&&!(part->GetStatus()&AliVTrack::kTOFpid)) return kFALSE;
350   if (fRequirePIDbit[icut]==AliDielectronPID::kIfAvailable&&!(part->GetStatus()&AliVTrack::kTOFpid)) return kTRUE;
351
352   Float_t numberOfSigmas=fPIDResponse->NumberOfSigmasTOF(part, fPartType[icut]);
353   
354   Bool_t selected=((numberOfSigmas>=fNsigmaLow[icut])&&(numberOfSigmas<=fNsigmaUp[icut]))^fExclude[icut];
355   return selected;
356 }
357
358 //______________________________________________
359 Bool_t AliDielectronPID::IsSelectedEMCAL(AliVTrack * const part, Int_t icut)
360 {
361   //
362   // emcal pid selecttion
363   //
364
365   //TODO: correct way to check for emcal pid?
366   Float_t numberOfSigmas=fPIDResponse->NumberOfSigmasEMCAL(part, fPartType[icut]);
367
368   Bool_t hasPID=numberOfSigmas>-998.;
369
370   if (fRequirePIDbit[icut]==AliDielectronPID::kRequire&&!hasPID) return kFALSE;
371   if (fRequirePIDbit[icut]==AliDielectronPID::kIfAvailable&&!hasPID) return kTRUE;
372
373
374   Bool_t selected=((numberOfSigmas>=fNsigmaLow[icut])&&(numberOfSigmas<=fNsigmaUp[icut]))^fExclude[icut];
375   return selected;
376 }
377
378 //______________________________________________
379 void AliDielectronPID::SetDefaults(Int_t def){
380   //
381   // initialise default pid strategies
382   //
383
384   if (def==0){
385     // 2sigma bands TPC:
386     // - include e
387     // - exclude mu,K,pi,p
388     // -complete p range
389     AddCut(kTPC,AliPID::kElectron,2);
390     AddCut(kTPC,AliPID::kMuon,-2.,2.,0.,0.,kTRUE);
391     AddCut(kTPC,AliPID::kPion,-2.,2.,0.,0.,kTRUE);
392     AddCut(kTPC,AliPID::kKaon,-2.,2.,0.,0.,kTRUE);
393     AddCut(kTPC,AliPID::kProton,-2.,2.,0.,0.,kTRUE);
394   } else if (def==1) {
395     // 2sigma bands TPC:
396     // - include e 0<p<inf
397     // - exclude mu,K,pi,p 0<p<2
398     AddCut(kTPC,AliPID::kElectron,2.);
399     AddCut(kTPC,AliPID::kMuon,-2.,2.,0.,2.,kTRUE);
400     AddCut(kTPC,AliPID::kPion,-2.,2.,0.,2.,kTRUE);
401     AddCut(kTPC,AliPID::kKaon,-2.,2.,0.,2.,kTRUE);
402     AddCut(kTPC,AliPID::kProton,-2.,2.,0.,2.,kTRUE);
403     
404   } else if (def==2) {
405     // include 2sigma e TPC
406     // 3sigma bands TOF
407     // - exclude K,P
408     AddCut(kTPC,AliPID::kElectron,-3.,3.);
409     AddCut(kTPC,AliPID::kPion,-3.,3.,0.,0.,kTRUE);
410     AddCut(kTPC,AliPID::kProton,-3.,3.,0.,0.,kTRUE);
411
412   } else if (def==3 || def==4) { // def3 and def4 are the same??
413     // include 2sigma e TPC
414     // 3sigma bands TOF
415     // - exclude K 0<p<1
416     // - exclude P 0<p<2
417     AddCut(kTPC,AliPID::kElectron,2);
418     AddCut(kTOF,AliPID::kKaon,-3.,3.,0.,1.,kTRUE);
419     AddCut(kTOF,AliPID::kProton,-6.,6.,0.,1.,kTRUE);
420     AddCut(kTOF,AliPID::kProton,-3.,3.,1.,2.,kTRUE);
421     
422   } else if (def==5) {
423     AddCut(kTPC,AliPID::kElectron,-0.5,3);
424     AddCut(kTOF,AliPID::kElectron,-3,3,0,1.5);
425   } else if (def==6) {
426     // lower cut TPC: parametrisation by HFE
427     // upper cut TPC: 3 sigma
428     // TOF ele band 3sigma 0<p<1.5GeV
429     TF1 *lowerCut=new TF1("lowerCut", "[0] * TMath::Exp([1]*x)", 0, 100);
430     lowerCut->SetParameters(-2.7,-0.4357);
431     AddCut(kTPC,AliPID::kElectron,lowerCut,3.);
432     AddCut(kTOF,AliPID::kElectron,-3,3,0,1.5);
433   } else if (def==7) {
434     // wide TPC cut
435     // TOF ele band 3sigma 0<p<1.5GeV
436     AddCut(kTPC,AliPID::kElectron,10.);
437     AddCut(kTOF,AliPID::kElectron,-3,3,0,1.5);
438   } else if (def==8) {
439     // TOF 5 sigma inclusion if TOFpid available
440     // this should reduce K,p,Pi to a large extent
441     AddCut(kTOF,AliPID::kElectron,-5,5,0,200,kFALSE,AliDielectronPID::kIfAvailable);
442   } else if (def==9) {
443     // lower cut TPC: parametrisation by HFE
444     // upper cut TPC: 3 sigma
445     // TOF 5 sigma inclusion if TOFpid available
446     // this should reduce K,p,Pi to a large extent
447     TF1 *lowerCut=new TF1("lowerCut", "[0] * TMath::Exp([1]*x)", 0, 100);
448     lowerCut->SetParameters(-2.65,-0.6757);
449     AddCut(kTPC,AliPID::kElectron,lowerCut,4.);
450     AddCut(kTOF,AliPID::kElectron,-5,5,0,200,kFALSE,AliDielectronPID::kIfAvailable);
451
452   } else if (def==10) {
453     AddCut(kTOF,AliPID::kElectron,-5,5,0,200,kFALSE,AliDielectronPID::kIfAvailable);
454     AddCut(kTPC,AliPID::kElectron,3.);
455     AddCut(kTPC,AliPID::kPion,-3.,3.,0.,0.,kTRUE);
456     AddCut(kTPC,AliPID::kProton,-3.,3.,0.,0.,kTRUE);
457     
458   } else if (def==11) {
459     // lower cut TPC: parametrisation by HFE
460     // only use from period d on !!
461     // upper cut TPC: 3 sigma
462     // TOF ele band 3sigma 0<p<2.0GeV
463     TF1 *lowerCut=new TF1("lowerCut", "[0] * TMath::Exp([1]*x)+[2]", 0, 100);
464     lowerCut->SetParameters(-3.718,-0.4,0.27);
465     AddCut(kTPC,AliPID::kElectron,lowerCut,3.);
466     AddCut(kTOF,AliPID::kElectron,-3,3,0,2.);
467   } else if (def==12) {
468     // lower cut TPC: parametrisation by HFE
469     // only use from period d on !!
470     // upper cut TPC: 3 sigma
471     // TOF 5 sigma inclusion if TOFpid available
472     // this should reduce K,p,Pi to a large extent
473     TF1 *lowerCut=new TF1("lowerCut", "[0] * TMath::Exp([1]*x)+[2]", 0, 100);
474     lowerCut->SetParameters(-3.718,-0.4,0.27);
475     AddCut(kTPC,AliPID::kElectron,lowerCut,4.);
476     AddCut(kTOF,AliPID::kElectron,-5,5,0,200,kFALSE,AliDielectronPID::kIfAvailable);
477   }
478
479 }
480
481 //______________________________________________
482 void AliDielectronPID::SetCorrVal(Double_t run)
483 {
484   //
485   // set correction value for run
486   //
487   fgCorr=0;
488   if (!fgFitCorr) return;
489   fgCorr=fgFitCorr->Eval(run);
490   if (run<fgFitCorr->GetX()[0]) fgCorr=fgFitCorr->GetY()[0];
491   if (run>fgFitCorr->GetX()[fgFitCorr->GetN()-1]) fgCorr=fgFitCorr->GetY()[fgFitCorr->GetN()-1];
492 //   printf("Corr: %f\n",fgCorr);
493 }
494
495 //______________________________________________
496 Double_t AliDielectronPID::GetEtaCorr(AliVTrack *track)
497 {
498   //
499   // return eta correction
500   //
501   if (!fgFunEtaCorr) return 1;
502   return fgFunEtaCorr->Eval(track->Eta());
503 }