Changing GetZv to GetZ
[u/mrichter/AliRoot.git] / ITS / UPGRADE / macros / QA / AliITSUComparisonPileup.C
1 /****************************************************************************
2  *  This macro calculates the efficiency of pileup reconstruction.          *
3  *  Works only for events generated with the AliGenPileup generator.        *
4  *                                                                          *
5  *  Before running, load the ITSU libraries:                                *
6  *  gSystem->Load("libITSUpgradeBase");gSystem->Load("libITSUpgradeRec");   *
7  *                                                                          *
8  *  Definifions:                                                            *
9  *  1) Reconstructable track: physical primary, charged, pT > pTmin         * 
10  *  2) Reconstructable vertex: has at least nMin reconstructable tracks     *
11  *  3) Associated vertex: has at least nAssMin correctly associated tracks  *
12  *  4) Good associated vertex: the fraction of correctly associated tracks  *
13  *        is at least fracMin                                               *
14  *  5) Fake associated vertex: not a good associated vertex                 *
15  *  6) Efficiency:  the ratio of 4) over 3)                                 *
16  *  7) Fake rate:   the ratio of 5) over 3)                                 *
17  *                                                                          *
18  *           Origin: I.Belikov, IPHC, Iouri.Belikov@iphc.cnrs.fr            *
19  ****************************************************************************/
20
21 #if !defined(__CINT__) || defined(__MAKECINT__)
22   #include <Riostream.h>
23   #include <TMath.h>
24   #include <TTree.h>
25   #include <TParticle.h>
26   #include <TParticlePDG.h>
27   #include <TCanvas.h>
28   #include <TFile.h>
29   #include <TLine.h>
30   #include <TROOT.h>
31   #include <TStyle.h>
32   #include <TLegend.h>
33
34   #include "AliStack.h"
35   #include "AliHeader.h"
36   #include "AliGenCocktailEventHeader.h"
37   #include "AliRunLoader.h"
38   #include "AliRun.h"
39   #include "AliESDEvent.h"
40   #include "AliESDtrack.h"
41 #endif
42
43 //**** Parameters used in the definitions
44 const Float_t pTmin=0.2;   // Minimal pT for a reconstructable track
45 const Int_t nMin=3;        // Minimal N of reconstructable tracks per vertex
46 const Int_t nAssMin=2;     // Minimal number of correctly associated tracks
47 const Float_t fracMin=0.8; // Minimal fraction of correctly associated tracks
48
49 extern AliRun *gAlice;
50 extern TROOT *gROOT;
51 extern TStyle *gStyle;
52
53 Int_t AliITSUComparisonPileup(const Char_t *dir=".") {
54    ::Info("AliITSUComparisonPileup.C","Doing comparison...");
55    Int_t GoodPileupVertices(const Char_t *dir=".");
56
57    // **** Book histogramms   
58    Int_t nb=35;
59    Float_t min=0, max=70.;
60    TH2F *h2spd=(TH2F*)gROOT->FindObject("h2spd");
61    if (!h2spd) 
62      h2spd=new TH2F("h2spd","SPD vertices;Number of good vertices;Number of reconstructed vertices",nb,min,max, nb,min,max);
63    h2spd->SetLineColor(2);
64    TH2F *h2trk=(TH2F*)gROOT->FindObject("h2trk");
65    if (!h2trk) 
66      h2trk=new TH2F("h2trk","TRK vertices;Good vertices;Reconstructed vertices",
67      nb,min,max, nb,min,max);
68    h2trk->SetLineColor(4);
69
70    nb=100;
71    min=-0.03; max=0.03;
72    TH1F *hzspd=(TH1F*)gROOT->FindObject("hzspd");
73    if (!hzspd) 
74      hzspd=new TH1F("hzspd","SPD resolution in Z;#DeltaZ (cm);",nb,min,max);
75    hzspd->SetLineColor(2);
76    TH1F *hztrk=(TH1F*)gROOT->FindObject("hztrk");
77    if (!hztrk) 
78      hztrk=new TH1F("hztrk","TRK resolution in Z;#DeltaZ (cm);",nb,min,max);
79    hztrk->SetLineColor(4);
80
81    
82    nb=30;
83    min=-10.; max=10.; 
84    TH1F *hgood=(TH1F*)gROOT->FindObject("hgood");
85    if (!hgood) 
86      hgood=new TH1F("hgood",";Z (cm);",nb,min,max);
87
88    TH1F *hfoundspd=(TH1F*)gROOT->FindObject("hfoundspd");
89    if (!hfoundspd) 
90     hfoundspd=new TH1F("hfoundspd",";Z (cm);",nb,min,max);
91    TH1F *heffspd=(TH1F*)gROOT->FindObject("heffspd");
92    if (!heffspd) 
93       heffspd=new TH1F("heffspd","SPD efficiency + fake rate;Z position of a prim. vertex (cm);Efficiency",nb,min,max);
94    heffspd->SetLineColor(2);
95
96    TH1F *hfoundtrk=(TH1F*)gROOT->FindObject("hfoundtrk");
97    if (!hfoundtrk) 
98     hfoundtrk=new TH1F("hfoundtrk",";Z (cm);",nb,min,max);
99    TH1F *hefftrk=(TH1F*)gROOT->FindObject("hefftrk");
100    if (!hefftrk) 
101       hefftrk=new TH1F("hefftrk","TRK efficiency;Z (cm);Efficiency",nb,min,max);
102    hefftrk->SetLineColor(4);
103
104    TH1F *hfaketrk=(TH1F*)gROOT->FindObject("hfaketrk");
105    if (!hfaketrk) 
106     hfaketrk=new TH1F("hfaketrk",";Z (cm);",nb,min,max);
107    TH1F *heffaketrk=(TH1F*)gROOT->FindObject("heffaketrk");
108    if (!heffaketrk) 
109       heffaketrk=new TH1F("heffaketrk","TRK fake rate;Z (cm);Fake rate",nb,min,max);
110    heffaketrk->SetLineColor(4);
111    heffaketrk->SetFillColor(590);
112
113
114
115    nb=51;
116    min=-0.5; max=50.5;
117    TH1F *hngood=(TH1F*)gROOT->FindObject("hngood");
118    if (!hngood) 
119       hngood=new TH1F("hngood",";Z (cm);",nb,min,max);
120    TH1F *hnfoundtrk=(TH1F*)gROOT->FindObject("hnfoundtrk");
121    if (!hnfoundtrk) 
122       hnfoundtrk=new TH1F("hnfoundtrk",";Z (cm);",nb,min,max);
123    TH1F *hnefftrk=(TH1F*)gROOT->FindObject("hnefftrk");
124    if (!hnefftrk) 
125       hnefftrk=new TH1F("hnefftrk","TRK efficiency;Number of tracks;Efficiency",nb,min,max);
126    hnefftrk->SetLineColor(4);
127
128    TH1F *hnfaketrk=(TH1F*)gROOT->FindObject("hnfaketrk");
129    if (!hnfaketrk) 
130       hnfaketrk=new TH1F("hnfaketrk",";Z (cm);",nb,min,max);
131    TH1F *hneffaketrk=(TH1F*)gROOT->FindObject("hneffaketrk");
132    if (!hneffaketrk) 
133       hneffaketrk=new TH1F("hneffaketrk","TRK fake rate;Number of tracks;Efficiency",nb,min,max);
134    hneffaketrk->SetLineColor(4);
135    hneffaketrk->SetFillColor(590);
136
137
138    // **** Generate a rerefence file with reconstructable vertices
139    Char_t fname[100];
140    sprintf(fname,"%s/GoodPileupVertices.root",dir);
141    TFile *refFile=TFile::Open(fname,"old");
142    if (!refFile || !refFile->IsOpen()) {
143       ::Info("AliITSUComparisonPileup.C",
144       "Marking good pileup vertices (will take a while)...");
145       if (GoodPileupVertices(dir)) {
146      ::Error("AliITSUComparisonPileup.C","Can't generate the reference file !");
147          return 1;
148       }
149    }
150    refFile=TFile::Open(fname,"old");
151    if (!refFile || !refFile->IsOpen()) {
152      ::Error("AliITSUComparisonPileup.C","Can't open the reference file !");
153      return 1;
154    }   
155    TTree *refTree=(TTree*)refFile->Get("refTree");
156    if (!refTree) {
157      ::Error("AliITSUComparisonPileup.C","Can't get the reference tree !");
158      return 2;
159    }
160    TBranch *branch=refTree->GetBranch("Vertices");
161    if (!branch) {
162      ::Error("AliITSUComparisonPileup.C","Can't get the vertex branch !");
163      return 3;
164    }
165    TClonesArray dummy("AliESDVertex",100), *refs=&dummy;
166    branch->SetAddress(&refs);    
167
168
169    // **** Open the ESD 
170    sprintf(fname,"%s/AliESDs.root",dir);
171    TFile *ef=TFile::Open(fname);
172    if ((!ef)||(!ef->IsOpen())) {
173       ::Error("AliITSUComparisonPileup.C","Can't open AliESDs.root !");
174       return 4;
175    }
176    AliESDEvent* event = new AliESDEvent();
177    TTree* esdTree = (TTree*) ef->Get("esdTree");
178    if (!esdTree) {
179       ::Error("AliITSUComparisonPileup.C", "no ESD tree found");
180       return 6;
181    }
182    event->ReadFromTree(esdTree);
183
184
185    //******* Loop over reconstructed events *********
186    Int_t ntrk=0, ntrkcor=0, ntrkwro=0;
187    Int_t e=0;
188    while (esdTree->GetEvent(e)) {
189      cout<<endl<<endl<<"********* Processing event number: "<<e<<"*******\n";
190      Int_t nn=event->GetNumberOfTracks();
191      ntrk += nn;
192
193      TClonesArray *verticesSPD=event->GetPileupVerticesSPD();
194      Int_t nfoundSPD=verticesSPD->GetEntries(); 
195      TClonesArray *verticesTRK=event->GetPileupVerticesTracks();
196      Int_t nfoundTRK=verticesTRK->GetEntries(); 
197
198      if (refTree->GetEvent(e++)==0) {
199         cerr<<"No reconstructable vertices for this event !\n";
200         continue;
201      }
202      Int_t ngood=refs->GetEntriesFast(); 
203      cout<<"Found SPD vertices: "<<nfoundSPD<<
204            "  Reconstructable vertics: "<<ngood<<endl;
205
206      h2spd->Fill(ngood,nfoundSPD);
207      h2trk->Fill(ngood,nfoundTRK);
208
209      Int_t ncor=0, nwro=0;
210      for (Int_t g=0; g<ngood; g++) {
211          Int_t Associate(const AliESDVertex *g, const AliESDVertex *f, 
212             const AliESDEvent *esd); 
213          const AliESDVertex *vtxg=(AliESDVertex*)refs->UncheckedAt(g);
214          Double_t zg=vtxg->GetZ();
215          Double_t ng=vtxg->GetNIndices();
216          hgood->Fill(zg);
217          hngood->Fill(ng);
218
219          AliESDVertex *vtxf=0;
220          Double_t zf=0.;
221          Int_t f=0;
222          for (; f<nfoundSPD; f++) {
223              vtxf=(AliESDVertex*)verticesSPD->UncheckedAt(f);
224              if (!vtxf->GetStatus()) continue;
225              if (Associate(vtxg,vtxf,event)==0) continue; 
226              break;
227          }
228          if (f>=nfoundSPD) {
229              vtxf=(AliESDVertex *)event->GetPrimaryVertexSPD();
230              if (!vtxf->GetStatus()) goto trk;
231              if (Associate(vtxg,vtxf,event)==0) goto trk; 
232          }
233
234          zf=vtxf->GetZ();
235          hfoundspd->Fill(zg);
236          hzspd->Fill(zf-zg);
237
238      trk:
239          Int_t n=0;
240          for (f=0; f<nfoundTRK; f++) {
241              vtxf=(AliESDVertex*)verticesTRK->UncheckedAt(f);
242              if (!vtxf->GetStatus()) continue;
243              n=Associate(vtxg,vtxf,event);
244              if (n < nAssMin) continue;
245              break;
246          }
247          if (f>=nfoundTRK) {
248              vtxf=(AliESDVertex*)event->GetPrimaryVertexTracks();
249              if (!vtxf->GetStatus()) continue;
250              n=Associate(vtxg,vtxf,event);
251              if (n < nAssMin) continue;
252          }
253
254          ncor+=n;
255          nwro+=(vtxf->GetNIndices()-n); 
256          zf=vtxf->GetZ();
257
258          if (Float_t(n)/vtxf->GetNIndices() > fracMin) {
259             hfoundtrk->Fill(zg);
260             hnfoundtrk->Fill(ng);
261          } else {
262             hfaketrk->Fill(zg);
263             hnfaketrk->Fill(ng);
264          }
265          hztrk->Fill(zf-zg);
266
267          vtxf->SetNContributors(0); // Mark this vertex as already associated
268
269      }
270      // Increase the counter of tracks (not)associated with verices
271      ntrkcor += ncor;
272      ntrkwro += nwro;
273
274    } //***** End of the loop over reconstructed events
275
276    delete event;
277    delete esdTree;
278    ef->Close();
279
280    refFile->Close();
281    
282    cout<<"\nTotal number of found tracks: "<<ntrk<<endl;
283    cout<<"Number of tracks correctly associated with vertices: "<<ntrkcor<<endl;
284    cout<<"Number of tracks wrongly associated with vertices: "<<ntrkwro<<endl;
285    if (ntrk != 0) {
286      cout<<"Correctly associated/Found:\t"<<Float_t(ntrkcor)/ntrk<<endl;
287      cout<<"Wrongly associated/Found:\t"<<Float_t(ntrkwro)/ntrk<<endl;
288      cout<<"Not associated/Found:\t\t"<<1.-Float_t(ntrkwro+ntrkcor)/ntrk<<endl;
289    }
290
291    gStyle->SetOptStat(0);
292    gStyle->SetOptTitle(0);
293    TCanvas *c1=new TCanvas("c1","",0,0,700,1000);
294    c1->Divide(1,2);
295    c1->cd(1);
296    gPad->SetGridx(); gPad->SetGridy();
297    h2spd->Draw("box");
298    h2trk->Draw("boxsame");
299    gPad->BuildLegend(0.13,0.65,0.46,0.86)->SetFillColor(0);
300    TLine *l=new TLine(0,0,70,70);
301    l->Draw("same");
302
303    c1->cd(2);
304    gPad->SetGridx(); gPad->SetGridy();
305    hztrk->Draw();
306    hzspd->Draw("same");
307    gPad->BuildLegend(0.13,0.65,0.46,0.86)->SetFillColor(0);
308
309
310    TCanvas *c2=new TCanvas("c2","",0,0,700,1000);
311    c2->Divide(1,2);
312    c2->cd(1);
313    gPad->SetGridx(); gPad->SetGridy();
314    heffspd->Divide(hfoundspd,hgood,1,1,"b");
315    heffspd->SetMinimum(0.); heffspd->SetMaximum(1.2);
316    heffspd->Draw("hist");
317    hefftrk->Divide(hfoundtrk,hgood,1,1,"b");
318    hefftrk->Draw("histsame");
319    heffaketrk->Divide(hfaketrk,hgood,1,1,"b");
320    heffaketrk->Draw("histsame");
321    gPad->BuildLegend(0.13,0.65,0.46,0.86)->SetFillColor(0);
322
323    c2->cd(2);
324    hnefftrk->Divide(hnfoundtrk,hngood,1,1,"b");
325    hnefftrk->SetMinimum(0.); hnefftrk->SetMaximum(1.2);
326    hneffaketrk->Divide(hnfaketrk,hngood,1,1,"b");
327    gPad->SetGridx(); gPad->SetGridy();
328    hnefftrk->Draw();
329    hneffaketrk->Draw("same");
330    gPad->BuildLegend(0.13,0.65,0.46,0.86)->SetFillColor(0);
331
332
333    TFile fc("AliITSUComparisonPileup.root","RECREATE");
334    c1->Write();
335    c2->Write();
336    fc.Close();
337
338    return 0;
339 }
340
341 Int_t 
342 Associate(const AliESDVertex *g,const AliESDVertex *f,const AliESDEvent *esd) { 
343    UShort_t *idxg=g->GetIndices(); Int_t ng=g->GetNIndices(); 
344    UShort_t *idxf=f->GetIndices(); Int_t nf=f->GetNIndices();
345
346    if (nf==0) { 
347    // SPD vertex
348        Double_t zg=g->GetZ();
349        Double_t zf=f->GetZ();
350        if (TMath::Abs(zf-zg)>2e-2) return 0;
351        return 1;
352    }
353    // TRK vertex
354    Int_t nass=0;
355    for (Int_t i=0; i<ng; i++) {
356        UShort_t labg=idxg[i];
357        for (Int_t j=0; j<nf; j++) {
358            const AliESDtrack *t=esd->GetTrack(idxf[j]);
359            UShort_t labf=TMath::Abs(t->GetLabel());
360            if (labg != labf) continue;
361            nass++;
362            break; 
363        }
364    } 
365
366    return nass;
367 }
368
369 Int_t GoodPileupVertices(const Char_t *dir) {
370   Int_t FindContributors(Float_t tz, AliStack *stack, UShort_t *idx);
371    if (gAlice) { 
372        delete AliRunLoader::Instance();
373        delete gAlice;//if everything was OK here it is already NULL
374        gAlice = 0x0;
375    }
376
377    Char_t fname[100];
378    sprintf(fname,"%s/galice.root",dir);
379
380    AliRunLoader *rl = AliRunLoader::Open(fname,"COMPARISON");
381    if (!rl) {
382       ::Error("GoodPileupVertices","Can't start session !");
383       return 1;
384    }
385
386    rl->LoadgAlice();
387    rl->LoadHeader();
388    rl->LoadKinematics();
389
390
391    Int_t nev=rl->GetNumberOfEvents();
392    ::Info("GoodPileupVertices","Number of events : %d\n",nev);  
393
394    sprintf(fname,"%s/GoodPileupVertices.root",dir);
395    TFile *refFile=TFile::Open(fname,"recreate");
396    TClonesArray dummy("AliESDVertex",100), *refs=&dummy;
397    TTree refTree("refTree","Tree with the reconstructable vertices");
398    refTree.Branch("Vertices",&refs);
399
400    //********  Loop over generated events 
401    for (Int_t e=0; e<nev; e++) {
402      rl->GetEvent(e);  refFile->cd();
403      AliStack* stack = rl->Stack();
404
405      AliHeader *ah=rl->GetHeader();
406      AliGenCocktailEventHeader *cock=
407             (AliGenCocktailEventHeader*)ah->GenEventHeader();
408      TList *headers=cock->GetHeaders();
409      const Int_t nvtx=headers->GetEntries();
410      const Int_t np=ah->GetNtrack();
411      cout<<"Event "<<e<<" Number of vertices, particles: "
412          <<nvtx<<' '<<np<<endl;
413
414      Int_t nv=0;
415      for (Int_t v=0; v<nvtx; v++) {
416          AliGenEventHeader *h=(AliGenEventHeader *)headers->At(v);
417          TArrayF vtx(3); h->PrimaryVertex(vtx);
418          Float_t t=h->InteractionTime();
419          UShort_t *idx=new UShort_t[np];
420          Int_t ntrk=FindContributors(t,stack,idx);
421          if (ntrk < nMin) continue;
422          AliESDVertex *vertex=new ((*refs)[nv]) AliESDVertex();
423          vertex->SetXv(vtx[0]);
424          vertex->SetYv(vtx[1]);
425          vertex->SetZv(vtx[2]);
426          vertex->SetNContributors(ntrk);
427          vertex->SetIndices(ntrk,idx);
428          delete idx;
429          nv++;
430      }
431      refTree.Fill();
432      refs->Clear();
433    } //*** end of the loop over generated events
434
435    refTree.Write();
436    refFile->Close();
437
438    delete rl;
439    return 0;
440 }
441
442 Int_t FindContributors(Float_t tz, AliStack *stack, UShort_t *idx) {
443   Int_t ntrk=0;
444   Int_t np=stack->GetNtrack();
445   for (Int_t i=0; i<np; i++) {
446       if (!stack->IsPhysicalPrimary(i)) continue;
447       TParticle *part=stack->Particle(i);
448       if (!part) continue;
449       TParticlePDG *partPDG = part->GetPDG();
450       if (!partPDG) continue;
451       if (TMath::Abs(partPDG->Charge())<1e-10) continue;
452       if (part->Pt() < pTmin) continue;
453       Float_t dt=0.5*(tz-part->T())/(tz+part->T());
454       if (TMath::Abs(dt)>1e-5) continue;
455       idx[ntrk]=i;
456       ntrk++;
457   }
458   return ntrk;
459 }