Implementation of asymmetric nsigma cuts.
[u/mrichter/AliRoot.git] / ANALYSIS / AliESDpidCuts.cxx
1 #include <TCanvas.h>
2 #include <TClass.h>
3 #include <TCollection.h>
4 #include <TDirectory.h>
5 #include <TH1F.h>
6 #include <TH1I.h>
7 #include <TH2I.h>
8 #include <TMath.h>
9 #include <TIterator.h>
10 #include <TString.h>
11
12 #include "AliESDtrack.h"
13 #include "AliLog.h"
14 #include "AliTPCpidESD.h"
15 #include "AliTOFpidESD.h"
16
17 #include "AliESDpidCuts.h"
18
19 ClassImp(AliESDpidCuts)
20
21 const Int_t AliESDpidCuts::kNcuts = 3;
22
23 //_____________________________________________________________________
24 AliESDpidCuts::AliESDpidCuts(const Char_t *name, const Char_t *title):
25     AliAnalysisCuts(name, title)
26   , fTPCpid(NULL)
27   , fTOFpid(NULL)
28   , fTPCsigmaCutRequired(0)
29   , fTOFsigmaCutRequired(0)
30   , fCutTPCclusterRatio(0.)
31   , fMinMomentumTOF(0.5)
32   , fHcutStatistics(NULL)
33   , fHcutCorrelation(NULL)
34 {
35   //
36   // Default constructor
37   //
38   
39   fTPCpid = new AliTPCpidESD;
40   fTOFpid = new AliTOFpidESD;
41   memset(fCutTPCnSigma, 0, sizeof(Float_t) * 2);
42   memset(fCutTOFnSigma, 0, sizeof(Float_t) * 2);
43
44   memset(fHclusterRatio, 0, sizeof(TH1F *) * 2);
45   memset(fHnSigmaTPC, 0, sizeof(TH1F *) * AliPID::kSPECIES * 2);
46   memset(fHnSigmaTOF, 0, sizeof(TH1F *) * AliPID::kSPECIES * 2);
47 }
48
49 //_____________________________________________________________________
50 AliESDpidCuts::AliESDpidCuts(const AliESDpidCuts &ref):
51     AliAnalysisCuts(ref)
52   , fTPCpid(NULL)
53   , fTOFpid(NULL)
54   , fTPCsigmaCutRequired(ref.fTPCsigmaCutRequired)
55   , fTOFsigmaCutRequired(ref.fTOFsigmaCutRequired)
56   , fCutTPCclusterRatio(ref.fCutTPCclusterRatio)
57   , fMinMomentumTOF(ref.fMinMomentumTOF)
58   , fHcutStatistics(NULL)
59   , fHcutCorrelation(NULL)
60 {
61   //
62   // Copy constructor
63   //
64   fTPCpid = new AliTPCpidESD(*ref.fTPCpid);
65   fTOFpid = new AliTOFpidESD(*ref.fTOFpid);
66   memcpy(fCutTPCnSigma, ref.fCutTPCnSigma, sizeof(Float_t) * AliPID::kSPECIES * 2);
67   memcpy(fCutTOFnSigma, ref.fCutTOFnSigma, sizeof(Float_t) * AliPID::kSPECIES * 2);
68   
69   if(ref.fHcutStatistics) fHcutStatistics = dynamic_cast<TH1I *>(ref.fHcutStatistics->Clone());
70   if(ref.fHcutCorrelation) fHcutCorrelation = dynamic_cast<TH2I *>(ref.fHcutCorrelation->Clone());
71   for(Int_t imode = 0; imode < 2; imode++){
72     if(ref.fHclusterRatio[imode]) fHclusterRatio[imode] = dynamic_cast<TH1F *>(ref.fHclusterRatio[imode]->Clone());
73     for(Int_t ispec = 0; ispec < AliPID::kSPECIES; ispec++){
74       if(fHnSigmaTPC[ispec][imode]) fHnSigmaTPC[ispec][imode] = dynamic_cast<TH1F *>(fHnSigmaTPC[ispec][imode]->Clone());
75       if(fHnSigmaTOF[ispec][imode]) fHnSigmaTOF[ispec][imode] = dynamic_cast<TH1F *>(fHnSigmaTPC[ispec][imode]->Clone());
76     }
77   }
78 }
79
80 //_____________________________________________________________________
81 AliESDpidCuts &AliESDpidCuts::operator=(const AliESDpidCuts &ref){
82   //
83   // Assignment operator
84   //
85   if(this != &ref)
86     ref.Copy(*this);
87   return *this;
88 }
89
90 //_____________________________________________________________________
91 AliESDpidCuts::~AliESDpidCuts(){
92   //
93   // Destructor
94   //
95   delete fTPCpid;
96   delete fTOFpid;
97
98   delete fHcutStatistics;
99   delete fHcutCorrelation;
100   for(Int_t imode = 0; imode < 2; imode++){
101     delete fHclusterRatio[imode];
102     for(Int_t ispec = 0; ispec < AliPID::kSPECIES; ispec++){
103       delete fHnSigmaTPC[ispec][imode];
104       delete fHnSigmaTOF[ispec][imode];
105     }
106   }
107 }
108
109 //_____________________________________________________________________
110 Bool_t AliESDpidCuts::IsSelected(TObject *o){
111   //
112   // Select Track
113   // 
114   if(TString(o->IsA()->GetName()).CompareTo("AliESDtrack")){
115     Char_t errormessage[256];
116     sprintf(errormessage, "Provided object not an AliESDtrack: Type %s", o->IsA()->GetName());
117     AliError(errormessage);
118     return kFALSE;
119   }
120   return AcceptTrack(const_cast<const AliESDtrack *>(dynamic_cast<AliESDtrack *>(o)));
121 }
122
123 //_____________________________________________________________________
124 void AliESDpidCuts::Copy(TObject &c) const {
125   //
126   // Copy function
127   //
128   AliESDpidCuts &target = dynamic_cast<AliESDpidCuts &>(c);
129
130   target.fTPCpid = new AliTPCpidESD;
131   target.fTOFpid = new AliTOFpidESD;
132
133   target.fCutTPCclusterRatio = fCutTPCclusterRatio;
134   target.fMinMomentumTOF = fMinMomentumTOF;
135
136   target.fTPCsigmaCutRequired = fTPCsigmaCutRequired;
137   target.fTOFsigmaCutRequired = fTOFsigmaCutRequired;
138   
139   if(fHcutStatistics) target.fHcutStatistics = dynamic_cast<TH1I *>(fHcutStatistics->Clone());
140   if(fHcutCorrelation) target.fHcutCorrelation = dynamic_cast<TH2I *>(fHcutCorrelation->Clone());
141   for(Int_t imode = 0; imode < 2; imode++){
142     if(fHclusterRatio[imode]) target.fHclusterRatio[imode] = dynamic_cast<TH1F *>(fHclusterRatio[imode]->Clone());
143     for(Int_t ispec = 0; ispec < AliPID::kSPECIES; ispec++){
144       if(fHnSigmaTPC[ispec][imode]) target.fHnSigmaTPC[ispec][imode] = dynamic_cast<TH1F *>(fHnSigmaTPC[ispec][imode]->Clone());
145       if(fHnSigmaTOF[ispec][imode]) target.fHnSigmaTOF[ispec][imode] = dynamic_cast<TH1F *>(fHnSigmaTOF[ispec][imode]->Clone());
146     }
147   }
148  
149   memcpy(target.fCutTPCnSigma, fCutTPCnSigma, sizeof(Float_t) * AliPID::kSPECIES * 2);
150   memcpy(target.fCutTOFnSigma, fCutTOFnSigma, sizeof(Float_t) * AliPID::kSPECIES);
151  
152   AliESDpidCuts::Copy(c);
153 }
154
155 //_____________________________________________________________________
156 Long64_t AliESDpidCuts::Merge(TCollection *coll){
157   //
158   // Merge Cut objects
159   //
160   if(coll) return 0;
161   if(coll->IsEmpty()) return 1;
162   if(!HasHistograms())  return 0;
163   
164   TIterator *iter = coll->MakeIterator();
165   TObject *o = NULL; 
166   AliESDpidCuts *ref = NULL;
167   Int_t counter = 0;
168   while((o = iter->Next())){
169     ref = dynamic_cast<AliESDpidCuts *>(o);
170     if(!ref) continue;
171     if(!ref->HasHistograms()) continue;
172
173     fHcutStatistics->Add(ref->fHcutStatistics);
174     fHcutCorrelation->Add(ref->fHcutCorrelation);
175     for(Int_t imode = 0; imode < 2; imode++){
176       fHclusterRatio[imode]->Add(ref->fHclusterRatio[imode]);
177       for(Int_t ispec = 0; ispec < AliPID::kSPECIES; ispec++){
178         fHnSigmaTPC[ispec][imode]->Add(ref->fHnSigmaTPC[ispec][imode]);
179         fHnSigmaTOF[ispec][imode]->Add(ref->fHnSigmaTOF[ispec][imode]);
180       }
181     }
182     ++counter;
183   }
184   return ++counter;
185 }
186
187 //_____________________________________________________________________
188 void AliESDpidCuts::DefineHistograms(Color_t color){
189   //
190   // Swich on QA and create the histograms
191   //
192   SetBit(kHasHistograms, kTRUE);
193   fHcutStatistics = new TH1I("fHcutStatistics", "Cut Statistics", kNcuts, 0, kNcuts);
194   fHcutStatistics->SetLineColor(color);
195   fHcutCorrelation = new TH2I("fHcutCorrelation", "Cut Correlation", kNcuts, 0, kNcuts, kNcuts, 0, kNcuts);
196   TString cutname[kNcuts] = {"TPCclusterRatio", "TPC sigma", "TOF sigma"};
197   for(Int_t icut = 0; icut < kNcuts; icut++){
198     fHcutStatistics->GetXaxis()->SetBinLabel(fHcutStatistics->GetXaxis()->GetFirst() + icut, cutname[icut].Data());
199     fHcutCorrelation->GetXaxis()->SetBinLabel(fHcutCorrelation->GetXaxis()->GetFirst() + icut, cutname[icut].Data());
200     fHcutCorrelation->GetYaxis()->SetBinLabel(fHcutCorrelation->GetYaxis()->GetFirst() + icut, cutname[icut].Data());
201   }
202   Char_t hname[256], htitle[256];
203   for(Int_t imode = 0; imode < 2; imode++){
204     sprintf(hname, "fHclusterRatio%s", imode ? "After" : "Before");
205     sprintf(htitle, "TPC cluster Ratio %s cuts;Ratio;Entries", imode ? "after" : "before");
206     fHclusterRatio[imode] = new TH1F(hname, htitle, 20, 0., 1.);
207     for(Int_t ispec = 0; ispec < AliPID::kSPECIES; ispec++){
208       sprintf(hname, "fHnSigma%sTPC%s", AliPID::ParticleName(ispec), imode ? "after" : "before");
209       sprintf(htitle, "TPC sigma for %s %s cuts;sigma;Entries", AliPID::ParticleName(ispec), imode ? "after" : "before");
210       fHnSigmaTPC[ispec][imode] = new TH1F(hname, htitle, 200, -10., 10.);
211       sprintf(hname, "fHnSigma%sTOF%s", AliPID::ParticleName(ispec), imode ? "after" : "before");
212       sprintf(htitle, "TOF sigma for %s %s cuts;sigma;Entries", AliPID::ParticleName(ispec), imode ? "after" : "before");
213       fHnSigmaTOF[ispec][imode] = new TH1F(hname, htitle, 200, -10., 10.);
214     }
215   }
216 }
217
218 //_____________________________________________________________________
219 Bool_t AliESDpidCuts::AcceptTrack(const AliESDtrack *track){
220   //
221   // Check whether the tracks survived the cuts
222   //
223   enum{
224     kCutClusterRatioTPC,
225     kCutNsigmaTPC,
226     kCutNsigmaTOF
227   };
228   Long64_t cutRequired=0, cutFullfiled = 0;
229   Double_t clusterRatio = track->GetTPCNclsF() ? static_cast<Float_t>(track->GetTPCNcls())/static_cast<Float_t>(track->GetTPCNclsF()) : 1.;
230   if(fCutTPCclusterRatio > 0.){
231     SETBIT(cutRequired, kCutClusterRatioTPC);
232     if(clusterRatio >= fCutTPCclusterRatio) 
233       SETBIT(cutFullfiled, kCutClusterRatioTPC);
234   }
235   // check TPC nSigma cut
236   Float_t nsigmaTPC[AliPID::kSPECIES], nsigma;   // need all sigmas for QA plotting
237   for(Int_t ispec = 0; ispec < AliPID::kSPECIES; ispec++){
238     nsigmaTPC[ispec] = nsigma = fTPCpid->GetNumberOfSigmas(track, static_cast<AliPID::EParticleType>(ispec));
239     if(!(fTPCsigmaCutRequired & 1 << ispec)) continue;
240     SETBIT(cutRequired, kCutNsigmaTPC); // We found at least one species where the n-Sigma Cut is required
241     if(nsigma >= fCutTPCnSigma[2*ispec] && nsigma <= fCutTPCnSigma[2*ispec+1]) SETBIT(cutFullfiled, kCutNsigmaTPC);    // Fullfiled for at least one species
242   }
243   // check TOF nSigma cut
244   Float_t nsigmaTOF[AliPID::kSPECIES];    // see above
245   Bool_t hasTOFpid = track->GetStatus() & AliESDtrack::kTOFpid; // only apply TOF n-sigma cut when PID Status Bit is set
246   for(Int_t ispec = 0; ispec < AliPID::kSPECIES; ispec++){
247     if(hasTOFpid) nsigmaTOF[ispec] = nsigma = fTOFpid->GetNumberOfSigmas(track, static_cast<AliPID::EParticleType>(ispec));
248     if(!(fTOFsigmaCutRequired && 1 << ispec)) continue;
249     SETBIT(cutRequired, kCutNsigmaTOF);
250     if(track->GetOuterParam()->P() >= fMinMomentumTOF){
251       if(hasTOFpid && nsigma <= fCutTOFnSigma[2*ispec] && nsigma >= fCutTOFnSigma[2*ispec+1]) SETBIT(cutFullfiled, kCutNsigmaTOF);
252     }
253   }
254
255   // Fill Histograms before cuts
256   if(HasHistograms()){
257     fHclusterRatio[0]->Fill(clusterRatio);
258     for(Int_t ispec = 0; ispec < AliPID::kSPECIES; ispec++){
259       fHnSigmaTPC[ispec][0]->Fill(nsigmaTPC[ispec]);
260       if(hasTOFpid) fHnSigmaTOF[ispec][0]->Fill(nsigmaTOF[ispec]);
261     }
262   }
263   if(cutRequired != cutFullfiled){
264     // Fill cut statistics
265     if(HasHistograms()){
266       for(Int_t icut = 0; icut < kNcuts; icut++){
267         if(TESTBIT(cutRequired, icut) && !TESTBIT(cutFullfiled, icut)){
268           // cut not fullfiled
269           fHcutStatistics->Fill(icut);
270           for(Int_t jcut = 0; jcut <= icut; jcut++)
271             if(TESTBIT(cutRequired, jcut) && !TESTBIT(cutFullfiled, jcut)) fHcutCorrelation->Fill(jcut, icut);
272         }
273       }
274     }
275     return kFALSE;    // At least one cut is not fullfiled
276   }
277
278   // Fill Histograms after cuts
279   if(HasHistograms()){
280     fHclusterRatio[1]->Fill(clusterRatio);
281     for(Int_t ispec = 0; ispec < AliPID::kSPECIES; ispec++){
282       fHnSigmaTPC[ispec][1]->Fill(nsigmaTPC[ispec]);
283       if(hasTOFpid) fHnSigmaTOF[ispec][1]->Fill(nsigmaTOF[ispec]);
284     }
285   }
286
287   return kTRUE;
288 }
289
290 //_____________________________________________________________________
291 void AliESDpidCuts::SaveHistograms(const Char_t * location){
292   //
293   // Save the histograms to a file
294   //
295   if(!HasHistograms()){
296     AliError("Histograms not on - Exiting");
297     return;
298   }
299   if(!location) location = GetName();
300   gDirectory->mkdir(location);
301   gDirectory->cd(location);
302   fHcutStatistics->Write();
303   fHcutCorrelation->Write();
304
305   gDirectory->mkdir("before_cuts");
306   gDirectory->mkdir("after_cuts");
307
308   gDirectory->cd("before_cuts");
309   fHclusterRatio[0]->Write();
310   for(Int_t ispec = 0; ispec < AliPID::kSPECIES; ispec++){
311     fHnSigmaTPC[ispec][0]->Write();
312     fHnSigmaTOF[ispec][0]->Write();
313   }
314
315   gDirectory->cd("../after_cuts");
316   fHclusterRatio[1]->Write();
317   for(Int_t ispec = 0; ispec < AliPID::kSPECIES; ispec++){
318     fHnSigmaTPC[ispec][1]->Write();
319     fHnSigmaTOF[ispec][1]->Write();
320   }
321
322   gDirectory->cd("..");
323 }
324
325 //_____________________________________________________________________
326 void AliESDpidCuts::DrawHistograms(){
327   //
328   // Draw the Histograms
329   //
330   TCanvas *stat = new TCanvas("cutStat", "Cut Statistics", 640, 480);
331   stat->cd();
332   fHcutStatistics->SetStats(kFALSE);
333   fHcutStatistics->Draw();
334   stat->SaveAs(Form("%s_%s.gif", GetName(), stat->GetName()));
335
336   TCanvas *correl = new TCanvas("cutCorrelation", "Cut Correlation", 640, 480);
337   correl->cd();
338   fHcutCorrelation->SetStats(kFALSE);
339   fHcutCorrelation->Draw("colz");
340   correl->SaveAs(Form("%s_%s.gif", GetName(), correl->GetName()));
341
342   TCanvas *cRatio = new TCanvas("ClusterRatioTPC", "TPC cluster Ratio", 640, 480);
343   cRatio->cd();
344   fHclusterRatio[0]->SetLineColor(kRed);
345   fHclusterRatio[0]->SetStats(kFALSE);
346   fHclusterRatio[0]->Draw();
347   fHclusterRatio[1]->SetLineColor(kBlue);
348   fHclusterRatio[1]->SetStats(kFALSE);
349   fHclusterRatio[1]->Draw("same");
350   cRatio->SaveAs(Form("%s_%s.gif",  GetName(), cRatio->GetName()));
351
352   TCanvas *cNsigTPC = new TCanvas("NsigmaTPC", "TPC n-sigma", 640, 480);
353   cNsigTPC->Divide(3,2);
354   for(Int_t ispec = 0; ispec < AliPID::kSPECIES; ispec++){
355     cNsigTPC->cd(ispec + 1);
356     fHnSigmaTPC[ispec][0]->SetLineColor(kRed);
357     fHnSigmaTPC[ispec][0]->SetStats(kFALSE);
358     fHnSigmaTPC[ispec][0]->Draw();
359     fHnSigmaTPC[ispec][1]->SetLineColor(kBlue);
360     fHnSigmaTPC[ispec][1]->SetStats(kFALSE);
361     fHnSigmaTPC[ispec][1]->Draw("same");
362   }
363   cNsigTPC->SaveAs(Form("%s_%s.gif", GetName(), cNsigTPC->GetName()));
364
365   TCanvas *cNsigTOF = new TCanvas("NsigmaTOF", "TOF n-sigma", 640, 480);
366   cNsigTOF->Divide(3,2);
367   for(Int_t ispec = 0; ispec < AliPID::kSPECIES; ispec++){
368     cNsigTOF->cd(ispec + 1);
369     fHnSigmaTOF[ispec][0]->SetLineColor(kRed);
370     fHnSigmaTOF[ispec][0]->SetStats(kFALSE);
371     fHnSigmaTOF[ispec][0]->Draw();
372     fHnSigmaTOF[ispec][1]->SetLineColor(kBlue);
373     fHnSigmaTOF[ispec][1]->SetStats(kFALSE);
374     fHnSigmaTOF[ispec][1]->Draw("same");
375   }
376   cNsigTOF->SaveAs(Form("%s_%s.gif", GetName(), cNsigTOF->GetName()));
377 }
378