f0a59f6a1850005e38463b6c109db12184c8f539
[u/mrichter/AliRoot.git] / STEER / AliESDComparison.C
1 //********************************************************************
2 //     Example (very naive for the moment) of using the ESD classes
3 //        It demonstrates the idea of the "combined PID".
4 //      Origin: Iouri Belikov, CERN, Jouri.Belikov@cern.ch
5 //********************************************************************
6
7 #if !defined( __CINT__) || defined(__MAKECINT__)
8   #include <Riostream.h>
9   #include "TKey.h"
10   #include "TFile.h"
11   #include "TH1F.h"
12   #include "TH2F.h"
13   #include "TCanvas.h"
14   #include "TStopwatch.h"
15   #include "TParticle.h"
16   #include "TROOT.h"
17
18   #include "AliRun.h"
19   #include "AliStack.h"
20   #include "AliRunLoader.h"
21   #include "AliLoader.h"
22
23   #include "AliESD.h"
24 #endif
25
26 extern AliRun *gAlice;
27 extern TROOT *gROOT;
28
29 Int_t AliESDComparison(const Char_t *dir=".") { 
30    Double_t pi=0.2,pa=2;
31
32    TH2F *tpcHist=(TH2F*)gROOT->FindObject("tpcHist");
33    if (!tpcHist)
34      tpcHist=new TH2F("tpcHist","TPC dE/dX vs momentum",100,pi,pa,100,0.,300.);
35    tpcHist->SetXTitle("p (GeV/c)"); tpcHist->SetYTitle("dE/dx (Arb. Units)");
36    tpcHist->SetMarkerStyle(8); 
37    tpcHist->SetMarkerSize(0.3);
38  
39    const Char_t *hname[]={
40     "piG","piR","piF","piGood","piFake",
41     "kaG","kaR","kaF","kaGood","kaFake",
42     "prG","prR","prF","prGood","prFake"
43    };
44    Int_t nh=sizeof(hname)/sizeof(const Char_t *);
45    TH1F **hprt=new TH1F*[nh]; 
46
47    for (Int_t i=0; i<nh; i++) {
48      hprt[i]=(TH1F*)gROOT->FindObject(hname[i]);
49      if (hprt[i]==0) {hprt[i]=new TH1F(hname[i],"",20,pi,pa);hprt[i]->Sumw2();}
50    }
51    TH1F *piG=hprt[0];
52    TH1F *piR=hprt[1];
53    TH1F *piF=hprt[2];
54    TH1F *piGood=hprt[3]; 
55         piGood->SetTitle("Combined PID for pions"); 
56         piGood->SetLineColor(4); piGood->SetXTitle("p (GeV/c)");
57    TH1F *piFake=hprt[4]; 
58         piFake->SetLineColor(2);
59    TH1F *kaG=hprt[5];
60    TH1F *kaR=hprt[6];
61    TH1F *kaF=hprt[7];
62    TH1F *kaGood=hprt[8];
63         kaGood->SetTitle("Combined PID for kaons"); 
64         kaGood->SetLineColor(4); kaGood->SetXTitle("p (GeV/c)");
65    TH1F *kaFake=hprt[9]; 
66         kaFake->SetLineColor(2);
67    TH1F *prG=hprt[10];
68    TH1F *prR=hprt[11];
69    TH1F *prF=hprt[12];
70    TH1F *prGood=hprt[13];
71         prGood->SetTitle("Combined PID for protons"); 
72         prGood->SetLineColor(4); prGood->SetXTitle("p (GeV/c)");
73    TH1F *prFake=hprt[14]; 
74         prFake->SetLineColor(2);
75
76    delete[] hprt;
77
78    Char_t fname[100];
79
80    if (gAlice) {
81       delete gAlice->GetRunLoader();
82       delete gAlice;
83       gAlice=0;
84    }
85    sprintf(fname,"%s/galice.root",dir);
86    AliRunLoader *rl = AliRunLoader::Open(fname);
87    if (rl == 0x0) {
88       cerr<<"Can not open session"<<endl;
89       return 1;
90    }
91    if (rl->LoadgAlice()) {
92       cerr<<"LoadgAlice returned error"<<endl;
93       delete rl;
94       return 1;
95    }
96    if (rl->LoadHeader()) {
97       cerr<<"LoadHeader returned error"<<endl;
98       delete rl;
99       return 1;
100    }
101    rl->LoadKinematics();
102
103    sprintf(fname,"%s/AliESDs.root",dir);
104    TFile *ef=TFile::Open(fname);
105    if (!ef->IsOpen()) {cerr<<"Can't AliESDs.root !\n"; return 1;}
106
107    TStopwatch timer;
108    Int_t rc=0,n=0;
109    TKey *key=0;
110    TIter next(ef->GetListOfKeys());
111
112    //****** Tentative particle type "concentrations"
113    Double_t c[5]={0.01, 0.01, 0.85, 0.10, 0.05};
114
115    //******* The loop over events
116    while ((key=(TKey*)next())!=0) {
117      rl->GetEvent(n);
118      AliStack *stack = rl->Stack();
119
120      cerr<<"Processing event number : "<<n++<<endl;
121
122      ef->cd();
123
124      AliESD *event=(AliESD*)key->ReadObj();
125
126      Int_t ntrk=event->GetNumberOfTracks();
127      cerr<<"Number of ESD tracks : "<<ntrk<<endl; 
128
129      // ****** The loop over tracks
130
131      Int_t pisel=0,kasel=0,prsel=0,nsel=0;
132
133      while (ntrk--) {
134        AliESDtrack *t=event->GetTrack(ntrk);
135
136        UInt_t status=AliESDtrack::kESDpid;
137        status|=AliESDtrack::kITSpid; 
138        status|=AliESDtrack::kTPCpid; 
139        status|=AliESDtrack::kTOFpid; 
140
141        if ((t->GetStatus()&status) == status) {
142          nsel++;
143
144          Double_t p=t->GetP();
145          Double_t dedx=t->GetTPCsignal();
146          tpcHist->Fill(p,dedx,1);
147
148          Int_t lab=TMath::Abs(t->GetLabel());
149          TParticle *part=stack->Particle(lab);
150          Int_t code=part->GetPdgCode();
151
152          Double_t r[10]; t->GetESDpid(r);
153
154          Double_t rcc=0.;
155          Int_t i;
156          for (i=0; i<AliESDtrack::kSPECIES; i++) rcc+=(c[i]*r[i]);
157          if (rcc==0.) continue;
158
159          //Here we apply Bayes' formula
160          Double_t w[10];
161          for (i=0; i<AliESDtrack::kSPECIES; i++) w[i]=c[i]*r[i]/rcc;
162
163          if (TMath::Abs(code)==2212) prR->Fill(p);
164          if (w[4]>w[3] && w[4]>w[2] && w[4]>w[1] && w[4]>w[0]) {//proton
165             prsel++;
166             prG->Fill(p);
167             if (TMath::Abs(code)!=2212) prF->Fill(p);
168          }
169
170          if (TMath::Abs(code)==321) kaR->Fill(p);
171          if (w[3]>w[4] && w[3]>w[2] && w[3]>w[1] && w[3]>w[0]) {//kaon
172             kasel++;
173             kaG->Fill(p);
174             if (TMath::Abs(code)!=321) kaF->Fill(p);
175          }
176
177          if (TMath::Abs(code)==211) piR->Fill(p);
178          if (w[2]>w[3] && w[2]>w[4] && w[2]>w[1] && w[2]>w[0]) {//pion
179             pisel++;
180             piG->Fill(p);
181             if (TMath::Abs(code)!=211) piF->Fill(p);
182          }
183
184        }
185
186      } 
187      delete event;
188      cerr<<"Number of selected ESD tracks : "<<nsel<<endl;
189      cerr<<"Number of selected pion ESD tracks : "<<pisel<<endl;
190      cerr<<"Number of selected kaon ESD tracks : "<<kasel<<endl;
191      cerr<<"Number of selected proton ESD tracks : "<<prsel<<endl;
192
193    }
194
195    timer.Stop(); timer.Print();
196
197    TCanvas *c1=(TCanvas*)gROOT->FindObject("c1");
198    if (c1) delete c1; 
199    c1=new TCanvas("c1","",0,0,600,1200);
200    c1->Divide(1,4);
201
202    c1->cd(1);
203    tpcHist->Draw();
204
205    c1->cd(2);
206    //piG->Sumw2(); piF->Sumw2(); piR->Sumw2();
207    piGood->Add(piG,piF,1,-1);
208    piGood->Divide(piGood,piR,1,1,"b");
209    piFake->Divide(piF,piG,1,1,"b");
210    piGood->Draw("hist");
211    piFake->Draw("same");
212
213    c1->cd(3);
214    //kaG->Sumw2(); kaF->Sumw2(); kaR->Sumw2();
215    kaGood->Add(kaG,kaF,1,-1);
216    kaGood->Divide(kaGood,kaR,1,1,"b");
217    kaFake->Divide(kaF,kaG,1,1,"b");
218    kaGood->Draw("hist");
219    kaFake->Draw("same");
220
221    c1->cd(4);
222    //prG->Sumw2(); prF->Sumw2(); prR->Sumw2();
223    prGood->Add(prG,prF,1,-1);
224    prGood->Divide(prGood,prR,1,1,"b");
225    prFake->Divide(prF,prG,1,1,"b");
226    prGood->Draw("hist");
227    prFake->Draw("same");
228
229    ef->Close();
230    TFile fc("AliESDComparison.root","RECREATE");
231    c1->Write();
232    fc.Close();
233
234    delete rl;
235
236    return rc;
237 }