debug message removed
[u/mrichter/AliRoot.git] / STEER / AliPIDComparison.C
1 //*************************************************************************
2 //                Example of using the PID classes
3 //  which calculates the efficiency and contamination of the combined PID.
4 //
5 //      Origin: Iouri Belikov, CERN, Jouri.Belikov@cern.ch
6 //*************************************************************************
7
8 #if !defined( __CINT__) || defined(__MAKECINT__)
9   #include <TMath.h>
10   #include <TError.h>
11   #include <Riostream.h>
12   #include <TF1.h>
13   #include <TH1F.h>
14   #include <TH2F.h>
15   #include <TParticle.h>
16   #include <TCanvas.h>
17   #include <TBenchmark.h>
18   #include <TFile.h>
19   #include <TTree.h>
20   #include <TROOT.h>
21
22   #include <AliStack.h>
23   #include <AliRunLoader.h>
24   #include <AliRun.h>
25   #include <AliTPCpidESD.h>
26   #include <AliESDEvent.h>
27 #endif
28
29 extern AliRun *gAlice;
30 extern TBenchmark *gBenchmark;
31 extern TROOT *gROOT;
32
33 static Int_t allpisel=0;
34 static Int_t allkasel=0;
35 static Int_t allprsel=0;
36 static Int_t allnsel=0;
37
38
39 Double_t tpcBethe(Double_t *x, Double_t *par) {
40   //
41   // This is the TPC Bethe-Bloch function normalised to 1 at the minimum
42   //
43   Double_t p=x[0];   //momentum
44   Double_t m=par[0]; //mass
45   Double_t bg=p/m;
46   return AliTPCpidESD::Bethe(bg);  
47 }
48
49
50 Int_t AliPIDComparison(const Char_t *dir=".") { 
51    gBenchmark->Start("AliPIDComparison");
52
53    ::Info("AliPIDComparison.C","Doing comparison...");
54
55    Double_t pi=0.2,pa=3;
56
57    TH2F *tpcHist=(TH2F*)gROOT->FindObject("tpcHist");
58    if (!tpcHist)
59      tpcHist=new TH2F("tpcHist","TPC dE/dX vs momentum",100,pi,pa,100,0.,6.);
60    tpcHist->SetXTitle("p (GeV/c)"); tpcHist->SetYTitle("dE/dx (Arb. Units)");
61    tpcHist->SetMarkerStyle(8); 
62    tpcHist->SetMarkerSize(0.3);
63  
64    const Char_t *hname[]={
65     "piG","piR","piF","piGood","piFake",
66     "kaG","kaR","kaF","kaGood","kaFake",
67     "prG","prR","prF","prGood","prFake"
68    };
69    Int_t nh=sizeof(hname)/sizeof(const Char_t *);
70    TH1F **hprt=new TH1F*[nh]; 
71
72    for (Int_t i=0; i<nh; i++) {
73      hprt[i]=(TH1F*)gROOT->FindObject(hname[i]);
74      if (hprt[i]==0) {hprt[i]=new TH1F(hname[i],"",20,pi,pa);hprt[i]->Sumw2();}
75    }
76    TH1F *piG=hprt[0];
77    TH1F *piR=hprt[1];
78    TH1F *piF=hprt[2];
79    TH1F *piGood=hprt[3]; 
80         piGood->SetTitle("Combined PID for pions");
81         piGood->SetLineColor(4); piGood->SetXTitle("p (GeV/c)");
82    TH1F *piFake=hprt[4]; 
83         piFake->SetLineColor(2);
84    TH1F *kaG=hprt[5];
85    TH1F *kaR=hprt[6];
86    TH1F *kaF=hprt[7];
87    TH1F *kaGood=hprt[8];
88         kaGood->SetTitle("Combined PID for kaons"); 
89         kaGood->SetLineColor(4); kaGood->SetXTitle("p (GeV/c)");
90    TH1F *kaFake=hprt[9]; 
91         kaFake->SetLineColor(2);
92    TH1F *prG=hprt[10];
93    TH1F *prR=hprt[11];
94    TH1F *prF=hprt[12];
95    TH1F *prGood=hprt[13];
96         prGood->SetTitle("Combined PID for protons"); 
97         prGood->SetLineColor(4); prGood->SetXTitle("p (GeV/c)");
98    TH1F *prFake=hprt[14]; 
99         prFake->SetLineColor(2);
100
101    delete[] hprt;
102
103    Char_t fname[100];
104
105    if (gAlice) {
106       delete AliRunLoader::Instance();
107       delete gAlice;
108       gAlice=0;
109    }
110    sprintf(fname,"%s/galice.root",dir);
111    AliRunLoader *rl = AliRunLoader::Open(fname);
112    if (rl == 0x0) {
113       ::Error("AliPIDComparison.C","Can not open session !");
114       return 1;
115    }
116    if (rl->LoadgAlice()) {
117       ::Error("AliPIDComparison.C","LoadgAlice returned error !");
118       delete rl;
119       return 1;
120    }
121    if (rl->LoadHeader()) {
122       ::Error("AliPIDComparison.C","LoadHeader returned error !");
123       delete rl;
124       return 1;
125    }
126    rl->LoadKinematics();
127
128    sprintf(fname,"%s/AliESDs.root",dir);
129    TFile *ef=TFile::Open(fname);
130    if (!ef || !ef->IsOpen()) {
131       ::Error("AliPIDComparison.C","Can't AliESDs.root !"); 
132       delete rl;
133       return 1;
134    }
135    AliESDEvent* event = new AliESDEvent();
136    TTree* tree = (TTree*) ef->Get("esdTree");
137    if (!tree) {
138       ::Error("AliPIDComparison.C", "no ESD tree found");
139       delete rl;
140       return 1;
141    }
142    event->ReadFromTree(tree);
143
144
145    //****** Tentative particle type "concentrations"
146    Double_t c[5]={0.01, 0.01, 0.85, 0.10, 0.05};
147    AliPID::SetPriors(c);
148
149
150    //******* The loop over events
151    Int_t e=0;
152    while (tree->GetEvent(e)) {
153       cout<<endl<<endl<<"********* Processing event number: "<<e<<"*******\n";
154
155       rl->GetEvent(e);
156  
157       e++;
158
159       Int_t ntrk=event->GetNumberOfTracks();
160       cerr<<"Number of ESD tracks : "<<ntrk<<endl; 
161
162       Int_t pisel=0,kasel=0,prsel=0,nsel=0;
163
164       AliStack *stack = rl->Stack();
165
166       while (ntrk--) {
167         AliESDtrack *t=event->GetTrack(ntrk);
168
169         //*** Some track quality cuts ****
170         if (t->GetITSclusters(0) < 6 ) continue;
171         if (t->GetTPCclusters(0) < 60) continue;
172         //if (t->GetTRDclusters(0) < 60) continue;
173
174         if (!t->IsOn(AliESDtrack::kITSpid)) continue;
175         if (!t->IsOn(AliESDtrack::kTPCpid)) continue;
176         //if (!t->IsOn(AliESDtrack::kTRDpid)) continue;
177         if (!t->IsOn(AliESDtrack::kTOFpid)) continue;
178         {
179            nsel++;
180
181            Double_t p=t->GetInnerParam()->GetP();
182            Double_t dedx=t->GetTPCsignal()/47.;
183            tpcHist->Fill(p,dedx,1);
184
185            Int_t lab=TMath::Abs(t->GetLabel());
186            TParticle *part=stack->Particle(lab);
187            Int_t code=part->GetPdgCode();
188
189            Double_t r[10]; t->GetESDpid(r);
190
191            AliPID pid(r);
192
193            Double_t w[10];
194            w[0]=pid.GetProbability(AliPID::kElectron);
195            w[1]=pid.GetProbability(AliPID::kMuon);
196            w[2]=pid.GetProbability(AliPID::kPion);
197            w[3]=pid.GetProbability(AliPID::kKaon);
198            w[4]=pid.GetProbability(AliPID::kProton);
199
200            if (TMath::Abs(code)==2212) prR->Fill(p);
201            if (w[4]>w[3] && w[4]>w[2] && w[4]>w[1] && w[4]>w[0]) {//proton
202               prsel++;
203               prG->Fill(p);
204               if (TMath::Abs(code)!=2212) prF->Fill(p);
205            }
206
207            if (TMath::Abs(code)==321) kaR->Fill(p);
208            if (w[3]>w[4] && w[3]>w[2] && w[3]>w[1] && w[3]>w[0]) {//kaon
209               kasel++;
210               kaG->Fill(p);
211               if (TMath::Abs(code)!=321) kaF->Fill(p);
212            }
213
214            if (TMath::Abs(code)==211) piR->Fill(p);
215            if (w[2]>w[4] && w[2]>w[3] && w[2]>w[0] && w[2]>w[1]) {//pion
216               pisel++;
217               piG->Fill(p);
218               if (TMath::Abs(code)!=211) piF->Fill(p);
219            }
220            /*
221            if (TMath::Abs(code)==11) piR->Fill(p);
222            if (w[0]>w[4] && w[0]>w[3] && w[0]>w[3] && w[0]>w[1]) {//electron
223               pisel++;
224               piG->Fill(p);
225               if (TMath::Abs(code)!=11) piF->Fill(p);
226            }
227            */
228         }
229       }
230       cout<<"Number of selected ESD tracks : "<<nsel<<endl;
231       cout<<"Number of selected pion ESD tracks : "<<pisel<<endl;
232       cout<<"Number of selected kaon ESD tracks : "<<kasel<<endl;
233       cout<<"Number of selected proton ESD tracks : "<<prsel<<endl;
234
235       allnsel+=nsel; allpisel+=pisel; allkasel+=kasel; allprsel+=prsel;
236
237    } // ***** End of the loop over events
238
239    delete event;
240    delete tree;
241    ef->Close();
242
243    TCanvas *c1=(TCanvas*)gROOT->FindObject("c1");
244    if (!c1) {
245       c1=new TCanvas("c1","",0,0,600,1000);
246       c1->Divide(1,4);
247    }
248
249    c1->cd(1);
250    tpcHist->Draw();
251
252    TF1 *fpi=(TF1*)gROOT->FindObject("piBethe");
253    if (!fpi) {
254       fpi=new TF1("piBethe",tpcBethe,pi,pa,1);
255       fpi->SetLineColor(2);
256       fpi->SetParameter(0,AliPID::ParticleMass(AliPID::kPion));
257    }
258    fpi->Draw("same");
259
260    TF1 *fka=(TF1*)gROOT->FindObject("kaBethe");
261    if (!fka) {
262       fka=new TF1("kaBethe",tpcBethe,pi,pa,1);
263       fka->SetLineColor(3);
264       fka->SetParameter(0,AliPID::ParticleMass(AliPID::kKaon));
265    }
266    fka->Draw("same");
267
268    TF1 *fpr=(TF1*)gROOT->FindObject("prBethe");
269    if (!fpr) {
270       fpr=new TF1("prBethe",tpcBethe,pi,pa,1);
271       fpr->SetLineColor(4);
272       fpr->SetParameter(0,AliPID::ParticleMass(AliPID::kProton));
273    }
274    fpr->Draw("same");
275
276
277    c1->cd(2);
278    //piG->Sumw2(); piF->Sumw2(); piR->Sumw2();
279    piGood->Add(piG,piF,1,-1);
280    piGood->Divide(piGood,piR,1,1,"b");
281    piFake->Divide(piF,piG,1,1,"b");
282    piGood->SetMinimum(0);
283    piGood->Draw("hist");
284    piFake->Draw("same");
285
286    c1->cd(3);
287    //kaG->Sumw2(); kaF->Sumw2(); kaR->Sumw2();
288    kaGood->Add(kaG,kaF,1,-1);
289    kaGood->Divide(kaGood,kaR,1,1,"b");
290    kaFake->Divide(kaF,kaG,1,1,"b");
291    kaGood->SetMinimum(0);
292    kaGood->Draw("hist");
293    kaFake->Draw("same");
294
295    c1->cd(4);
296    //prG->Sumw2(); prF->Sumw2(); prR->Sumw2();
297    prGood->Add(prG,prF,1,-1);
298    prGood->Divide(prGood,prR,1,1,"b");
299    prFake->Divide(prF,prG,1,1,"b");
300    prGood->SetMinimum(0);
301    prGood->Draw("hist");
302    prFake->Draw("same");
303
304    c1->Update();
305
306    cout<<endl;
307    cout<<endl;
308    e=(Int_t)piR->GetEntries();
309    Int_t o=(Int_t)piG->GetEntries();
310    if (e*o) cout<<"Efficiency (contamination) for pions : "<<
311       piGood->GetEntries()/e<<'('<<piF->GetEntries()/o<<')'<<endl; 
312    e=(Int_t)kaR->GetEntries();
313    o=(Int_t)kaG->GetEntries();
314    if (e*o) cout<<"Efficiency (contamination) for kaons : "<<
315       kaGood->GetEntries()/e<<'('<<kaF->GetEntries()/o<<')'<<endl; 
316    e=(Int_t)prR->GetEntries();
317    o=(Int_t)prG->GetEntries();
318    if (e*o) cout<<"Efficiency (contamination) for protons : "<<
319       prGood->GetEntries()/e<<'('<<prF->GetEntries()/o<<')'<<endl;
320    cout<<endl; 
321
322    cout<<"Total number of selected ESD tracks : "<<allnsel<<endl;
323    cout<<"Total number of selected pion ESD tracks : "<<allpisel<<endl;
324    cout<<"Total number of selected kaon ESD tracks : "<<allkasel<<endl;
325    cout<<"Total number of selected proton ESD tracks : "<<allprsel<<endl;
326
327    ef->Close();
328    TFile fc("AliPIDComparison.root","RECREATE");
329    c1->Write();
330    fc.Close();
331
332    gBenchmark->Stop("AliPIDComparison");
333    gBenchmark->Show("AliPIDComparison");
334
335    delete rl;
336    return 0;
337 }