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