Major dielectron framework update; includes "alignment" to updates in
[u/mrichter/AliRoot.git] / PWG3 / 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 <AliLog.h>
34 #include <AliExternalTrackParam.h>
35 #include <AliESDtrack.h>
36 #include <AliESDpid.h>
37 #include <AliAODpidUtil.h>
38 #include <AliAODPid.h>
39
40 #include "AliDielectronVarManager.h"
41
42 #include "AliDielectronPID.h"
43
44 ClassImp(AliDielectronPID)
45
46 TGraph *AliDielectronPID::fgFitCorr=0x0;
47 Double_t AliDielectronPID::fgCorr=0.0;
48
49 AliDielectronPID::AliDielectronPID() :
50   AliAnalysisCuts(),
51   fNcuts(0),
52   fESDpid(0x0),
53   fAODpidUtil(0x0)
54 {
55   //
56   // Default Constructor
57   //
58   for (Int_t icut=0; icut<kNmaxPID; ++icut){
59     fDetType[icut]=kTPC;
60     fPartType[icut]=AliPID::kPion;
61     fNsigmaLow[icut]=0;
62     fNsigmaUp[icut]=0;
63     fPmin[icut]=0;
64     fPmax[icut]=0;
65     fExclude[icut]=kFALSE;
66     fFunUpperCut[icut]=0x0;
67     fFunLowerCut[icut]=0x0;
68     fRequirePIDbit[icut]=0;
69   }
70 }
71
72 //______________________________________________
73 AliDielectronPID::AliDielectronPID(const char* name, const char* title) :
74   AliAnalysisCuts(name, title),
75   fNcuts(0),
76   fESDpid(0x0),
77   fAODpidUtil(0x0)
78 {
79   //
80   // Named Constructor
81   //
82   for (Int_t icut=0; icut<kNmaxPID; ++icut){
83     fDetType[icut]=kTPC;
84     fPartType[icut]=AliPID::kPion;
85     fNsigmaLow[icut]=0;
86     fNsigmaUp[icut]=0;
87     fPmin[icut]=0;
88     fPmax[icut]=0;
89     fExclude[icut]=kFALSE;
90     fFunUpperCut[icut]=0x0;
91     fFunLowerCut[icut]=0x0;
92     fRequirePIDbit[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 pMin/*=0*/, Double_t pMax/*=0*/, Bool_t exclude/*=kFALSE*/,
107                               UInt_t pidBitType/*=AliDielectronPID::kRequire*/)
108 {
109   //
110   // Add a pid nsigma cut
111   // use response of detector 'det' in the momentum range ('pMin') to ['pMax']
112   // use a sigma band betwee '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   }
120   if (TMath::Abs(nSigmaUp+99999.)<1e-20){
121     nSigmaUp=TMath::Abs(nSigmaLow);
122     nSigmaLow=-1.*nSigmaUp;
123   }
124   fDetType[fNcuts]=det;
125   fPartType[fNcuts]=type;
126   fNsigmaLow[fNcuts]=nSigmaLow;
127   fNsigmaUp[fNcuts]=nSigmaUp;
128   fPmin[fNcuts]=pMin;
129   fPmax[fNcuts]=pMax;
130   fExclude[fNcuts]=exclude;
131   fRequirePIDbit[fNcuts]=pidBitType;
132   ++fNcuts;
133   
134 }
135
136 //______________________________________________
137 void AliDielectronPID::AddCut(DetType det, AliPID::EParticleType type, Double_t nSigmaLow, TF1 * const funUp,
138                               Double_t pMin/*=0*/, Double_t pMax/*=0*/, Bool_t exclude/*=kFALSE*/,
139                               UInt_t pidBitType/*=AliDielectronPID::kRequire*/)
140 {
141   //
142   // cut using a TF1 as upper cut
143   //
144   if (funUp==0x0){
145     AliError("A valid function is required for the upper cut. Not adding the cut!");
146     return;
147   }
148   fFunUpperCut[fNcuts]=funUp;
149   AddCut(det,type,nSigmaLow,0.,pMin,pMax,exclude,pidBitType);
150 }
151
152 //______________________________________________
153 void AliDielectronPID::AddCut(DetType det, AliPID::EParticleType type, TF1 * const funLow, Double_t nSigmaUp,
154                               Double_t pMin/*=0*/, Double_t pMax/*=0*/, Bool_t exclude/*=kFALSE*/,
155                               UInt_t pidBitType/*=AliDielectronPID::kRequire*/)
156 {
157   //
158   // cut using a TF1 as lower cut
159   //
160   if (funLow==0x0){
161     AliError("A valid function is required for the lower cut. Not adding the cut!");
162     return;
163   }
164   fFunLowerCut[fNcuts]=funLow;
165   AddCut(det,type,0.,nSigmaUp,pMin,pMax,exclude,pidBitType);
166 }
167
168 //______________________________________________
169 void AliDielectronPID::AddCut(DetType det, AliPID::EParticleType type, TF1 * const funLow, TF1 * const funUp,
170                               Double_t pMin/*=0*/, Double_t pMax/*=0*/, Bool_t exclude/*=kFALSE*/,
171                               UInt_t pidBitType/*=AliDielectronPID::kRequire*/)
172 {
173   //
174   // cut using a TF1 as lower and upper cut
175   //
176   if ( (funUp==0x0) || (funLow==0x0) ){
177     AliError("A valid function is required for upper and lower cut. Not adding the cut!");
178     return;
179   }
180   fFunUpperCut[fNcuts]=funUp;
181   fFunLowerCut[fNcuts]=funLow;
182   AddCut(det,type,0.,0.,pMin,pMax,exclude,pidBitType);
183 }
184
185 //______________________________________________
186 Bool_t AliDielectronPID::IsSelected(TObject* track)
187 {
188   //
189   // perform PID cuts
190   //
191
192   //loop over all cuts
193   AliVTrack *part=static_cast<AliVTrack*>(track);
194   //TODO: Which momentum to use?
195   //      Different momenta for different detectors?
196   Double_t mom=part->P();
197   
198   Bool_t selected=kFALSE;
199   fESDpid=AliDielectronVarManager::GetESDpid();
200   fAODpidUtil=AliDielectronVarManager::GetAODpidUtil();
201   
202   for (UChar_t icut=0; icut<fNcuts; ++icut){
203     Double_t pMin=fPmin[icut];
204     Double_t pMax=fPmax[icut];
205
206     // test momentum range. In case pMin==pMax use all momenta
207     if ( (TMath::Abs(pMin-pMax)>1e-20) && (mom<=pMin || mom>pMax) ) continue;
208
209
210     switch (fDetType[icut]){
211     case kITS:
212       selected = IsSelectedITS(part,icut);
213       break;
214     case kTPC:
215       selected = IsSelectedTPC(part,icut);
216       break;
217     case kTRD:
218       selected = IsSelectedTRD(part,icut);
219       break;
220     case kTOF:
221       selected = IsSelectedTOF(part,icut);
222       break;
223     }
224     if (!selected) return kFALSE;
225   }
226
227   return selected;
228 }
229
230 //______________________________________________
231 Bool_t AliDielectronPID::IsSelectedITS(AliVTrack * const part, Int_t icut)
232 {
233   //
234   // ITS part of the PID check
235   // Don't accept the track if there was no pid bit set
236   //
237   Float_t numberOfSigmas=-1000.;
238   
239   if (fRequirePIDbit[icut]==AliDielectronPID::kRequire&&!(part->GetStatus()&AliESDtrack::kITSpid)) return kFALSE;
240   if (fRequirePIDbit[icut]==AliDielectronPID::kIfAvailable&&!(part->GetStatus()&AliESDtrack::kITSpid)) return kTRUE;
241
242   Double_t mom=part->P();
243   
244   if (part->IsA()==AliESDtrack::Class()){
245     // ESD case in case the PID bit is not set, don't use this track!
246     AliESDtrack *track=static_cast<AliESDtrack*>(part);
247     numberOfSigmas=fESDpid->NumberOfSigmasITS(track, fPartType[icut]);
248   }else if(part->IsA()==AliAODTrack::Class()){
249     // AOD case
250     // FIXME: Is there a place to check whether the PID is was set in ESD???
251     AliAODTrack *track=static_cast<AliAODTrack*>(part);
252     numberOfSigmas=fAODpidUtil->NumberOfSigmasITS(track, fPartType[icut]);
253   }
254   
255   // test if we are supposed to use a function for the cut
256   if (fFunUpperCut[icut]) fNsigmaUp[icut] =fFunUpperCut[icut]->Eval(mom);
257   if (fFunLowerCut[icut]) fNsigmaLow[icut]=fFunLowerCut[icut]->Eval(mom);
258   
259   Bool_t selected=((numberOfSigmas>=fNsigmaLow[icut])&&(numberOfSigmas<=fNsigmaUp[icut]))^fExclude[icut];
260   return selected;
261 }
262
263 //______________________________________________
264 Bool_t AliDielectronPID::IsSelectedTPC(AliVTrack * const part, Int_t icut)
265 {
266   //
267   // TPC part of the PID check
268   // Don't accept the track if there was no pid bit set
269   //
270   Float_t numberOfSigmas=-1000.;
271   
272   if (fRequirePIDbit[icut]==AliDielectronPID::kRequire&&!(part->GetStatus()&AliESDtrack::kTPCpid)) return kFALSE;
273   if (fRequirePIDbit[icut]==AliDielectronPID::kIfAvailable&&!(part->GetStatus()&AliESDtrack::kTPCpid)) return kTRUE;
274
275   Double_t mom=part->P();
276   
277   if (part->IsA()==AliESDtrack::Class()){
278     // ESD case in case the PID bit is not set, don't use this track!
279     AliESDtrack *track=static_cast<AliESDtrack*>(part);
280     const AliExternalTrackParam *in = track->GetInnerParam();
281     if (in) mom = in->GetP();
282     numberOfSigmas=fESDpid->NumberOfSigmasTPC(track, fPartType[icut]);
283   }else if(part->IsA()==AliAODTrack::Class()){
284     // AOD case
285     // FIXME: Is there a place to check whether the PID is was set in ESD???
286     AliAODTrack *track=static_cast<AliAODTrack*>(part);
287     const AliAODPid *pidObj = track->GetDetPid();
288     if (pidObj) mom = pidObj->GetTPCmomentum();
289     numberOfSigmas=fAODpidUtil->NumberOfSigmasTPC(track, fPartType[icut]);
290   }
291   if (fPartType[icut]==AliPID::kElectron){
292     numberOfSigmas-=fgCorr;
293   }
294   
295   // test if we are supposed to use a function for the cut
296   if (fFunUpperCut[icut]) fNsigmaUp[icut] =fFunUpperCut[icut]->Eval(mom);
297   if (fFunLowerCut[icut]) fNsigmaLow[icut]=fFunLowerCut[icut]->Eval(mom);
298   
299   Bool_t selected=((numberOfSigmas>=fNsigmaLow[icut])&&(numberOfSigmas<=fNsigmaUp[icut]))^fExclude[icut];
300   return selected;
301 }
302
303 //______________________________________________
304 Bool_t AliDielectronPID::IsSelectedTRD(AliVTrack * const /*part*/, Int_t /*icut*/)
305 {
306   //   
307   // TRD part of the pid check
308   //
309   return kFALSE;
310 }
311
312 //______________________________________________
313 Bool_t AliDielectronPID::IsSelectedTOF(AliVTrack * const part, Int_t icut)
314 {
315   //
316   // TOF part of the PID check
317   // Don't accept the track if there was no pid bit set
318   //
319   Float_t numberOfSigmas=-1000.;
320   
321   if (fRequirePIDbit[icut]==AliDielectronPID::kRequire&&!(part->GetStatus()&AliESDtrack::kTOFpid)) return kFALSE;
322   if (fRequirePIDbit[icut]==AliDielectronPID::kIfAvailable&&!(part->GetStatus()&AliESDtrack::kTOFpid)) return kTRUE;
323
324   if (part->IsA()==AliESDtrack::Class()){
325     // ESD case in case the PID bit is not set, don't use this track!
326     AliESDtrack *track=static_cast<AliESDtrack*>(part);
327     numberOfSigmas=fESDpid->NumberOfSigmasTOF(track, fPartType[icut], fESDpid->GetTOFResponse().GetTimeZero());
328   }else if(part->IsA()==AliAODTrack::Class()){
329     // AOD case
330     // FIXME: Is there a place to check whether the PID is was set in ESD???
331     AliAODTrack *track=static_cast<AliAODTrack*>(part);
332     numberOfSigmas=fAODpidUtil->NumberOfSigmasTOF(track, fPartType[icut]);
333   }
334   
335   Bool_t selected=((numberOfSigmas>=fNsigmaLow[icut])&&(numberOfSigmas<=fNsigmaUp[icut]))^fExclude[icut];
336   return selected;
337 }
338
339 //______________________________________________
340 void AliDielectronPID::SetDefaults(Int_t def){
341   //
342   // initialise default pid strategies
343   //
344
345   if (def==0){
346     // 2sigma bands TPC:
347     // - include e
348     // - exclude mu,K,pi,p
349     // -complete p range
350     AddCut(kTPC,AliPID::kElectron,2);
351     AddCut(kTPC,AliPID::kMuon,-2.,2.,0.,0.,kTRUE);
352     AddCut(kTPC,AliPID::kPion,-2.,2.,0.,0.,kTRUE);
353     AddCut(kTPC,AliPID::kKaon,-2.,2.,0.,0.,kTRUE);
354     AddCut(kTPC,AliPID::kProton,-2.,2.,0.,0.,kTRUE);
355   } else if (def==1) {
356     // 2sigma bands TPC:
357     // - include e 0<p<inf
358     // - exclude mu,K,pi,p 0<p<2
359
360     AddCut(kTPC,AliPID::kElectron,2.);
361     AddCut(kTPC,AliPID::kMuon,-2.,2.,0.,2.,kTRUE);
362     AddCut(kTPC,AliPID::kPion,-2.,2.,0.,2.,kTRUE);
363     AddCut(kTPC,AliPID::kKaon,-2.,2.,0.,2.,kTRUE);
364     AddCut(kTPC,AliPID::kProton,-2.,2.,0.,2.,kTRUE);
365     
366   } else if (def==2) {
367     // include 2sigma e TPC
368     // 3sigma bands TOF
369     // - exclude K,P
370     AddCut(kTPC,AliPID::kElectron,-3.,3.);
371     AddCut(kTPC,AliPID::kPion,-3.,3.,0.,0.,kTRUE);
372     AddCut(kTPC,AliPID::kProton,-3.,3.,0.,0.,kTRUE);
373
374   } else if (def==3) {
375     // include 2sigma e TPC
376     // 3sigma bands TOF
377     // - exclude K 0<p<1
378     // - exclude P 0<p<2
379     AddCut(kTPC,AliPID::kElectron,2);
380     AddCut(kTOF,AliPID::kKaon,-3.,3.,0.,1.,kTRUE);
381     AddCut(kTOF,AliPID::kProton,-6.,6.,0.,1.,kTRUE);
382     AddCut(kTOF,AliPID::kProton,-3.,3.,1.,2.,kTRUE);
383     
384   } else if (def==4) {
385     // include 2sigma e TPC
386     // 3sigma bands TOF
387     // - exclude K 0<p<1
388     // - exclude P 0<p<2
389     AddCut(kTPC,AliPID::kElectron,2.);
390     AddCut(kTOF,AliPID::kKaon,-3.,3.,0.,1.,kTRUE);
391     AddCut(kTOF,AliPID::kProton,-6.,6.,0.,1.,kTRUE);
392     AddCut(kTOF,AliPID::kProton,-3.,3.,1.,2.,kTRUE);
393   } else if (def==5) {
394     AddCut(kTPC,AliPID::kElectron,-0.5,3);
395     AddCut(kTOF,AliPID::kElectron,-3,3,0,1.5);
396   } else if (def==6) {
397     // lower cut TPC: parametrisation by HFE
398     // upper cut TPC: 3 sigma
399     // TOF ele band 3sigma 0<p<1.5GeV
400     TF1 *lowerCut=new TF1("lowerCut", "[0] * TMath::Exp([1]*x)", 0, 100);
401     lowerCut->SetParameters(-2.7,-0.4357);
402     AddCut(kTPC,AliPID::kElectron,lowerCut,3.);
403     AddCut(kTOF,AliPID::kElectron,-3,3,0,1.5);
404   } else if (def==7) {
405     // wide TPC cut
406     // TOF ele band 3sigma 0<p<1.5GeV
407     AddCut(kTPC,AliPID::kElectron,10.);
408     AddCut(kTOF,AliPID::kElectron,-3,3,0,1.5);
409   } else if (def==8) {
410     // TOF 5 sigma inclusion if TOFpid available
411     // this should reduce K,p,Pi to a large extent
412     AddCut(kTOF,AliPID::kElectron,-5,5,0,200,kFALSE,AliDielectronPID::kIfAvailable);
413   } else if (def==9) {
414     // lower cut TPC: parametrisation by HFE
415     // upper cut TPC: 3 sigma
416     // TOF 5 sigma inclusion if TOFpid available
417     // this should reduce K,p,Pi to a large extent
418     TF1 *lowerCut=new TF1("lowerCut", "[0] * TMath::Exp([1]*x)", 0, 100);
419     lowerCut->SetParameters(-2.65,-0.6757);
420     AddCut(kTPC,AliPID::kElectron,lowerCut,4.);
421     AddCut(kTOF,AliPID::kElectron,-5,5,0,200,kFALSE,AliDielectronPID::kIfAvailable);
422   } else if (def==10) {
423     AddCut(kTOF,AliPID::kElectron,-5,5,0,200,kFALSE,AliDielectronPID::kIfAvailable);
424     AddCut(kTPC,AliPID::kElectron,3.);
425     AddCut(kTPC,AliPID::kPion,-3.,3.,0.,0.,kTRUE);
426     AddCut(kTPC,AliPID::kProton,-3.,3.,0.,0.,kTRUE);
427     
428   }
429 }
430
431 //______________________________________________
432 void AliDielectronPID::SetCorrVal(Double_t run)
433 {
434   //
435   // set correction value for run
436   //
437   fgCorr=0;
438   if (!fgFitCorr) return;
439   fgCorr=fgFitCorr->Eval(run);
440   if (run<fgFitCorr->GetX()[0]) fgCorr=fgFitCorr->GetY()[0];
441   if (run>fgFitCorr->GetX()[fgFitCorr->GetN()-1]) fgCorr=fgFitCorr->GetY()[fgFitCorr->GetN()-1];
442 //   printf("Corr: %f\n",fgCorr);
443 }
444