ae2c2e39e9627b09620b9b3aea03bd11b51b47a6
[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    // **** Generate a rerefence file with reconstructable vertices
116    Char_t fname[100];
117    sprintf(fname,"%s/GoodPileupVertices.root",dir);
118    TFile *refFile=TFile::Open(fname,"old");
119    if (!refFile || !refFile->IsOpen()) {
120       ::Info("AliITSUComparisonPileup.C",
121       "Marking good pileup vertices (will take a while)...");
122       if (GoodPileupVertices(dir)) {
123      ::Error("AliITSUComparisonPileup.C","Can't generate the reference file !");
124          return 1;
125       }
126    }
127    refFile=TFile::Open(fname,"old");
128    if (!refFile || !refFile->IsOpen()) {
129      ::Error("AliITSUComparisonPileup.C","Can't open the reference file !");
130      return 1;
131    }   
132    TTree *refTree=(TTree*)refFile->Get("refTree");
133    if (!refTree) {
134      ::Error("AliITSUComparisonPileup.C","Can't get the reference tree !");
135      return 2;
136    }
137    TBranch *branch=refTree->GetBranch("Vertices");
138    if (!branch) {
139      ::Error("AliITSUComparisonPileup.C","Can't get the vertex branch !");
140      return 3;
141    }
142    TClonesArray dummy("AliESDVertex",100), *refs=&dummy;
143    branch->SetAddress(&refs);    
144
145
146    // **** Open the ESD 
147    sprintf(fname,"%s/AliESDs.root",dir);
148    TFile *ef=TFile::Open(fname);
149    if ((!ef)||(!ef->IsOpen())) {
150       ::Error("AliITSUComparisonPileup.C","Can't open AliESDs.root !");
151       return 4;
152    }
153    AliESDEvent* event = new AliESDEvent();
154    TTree* esdTree = (TTree*) ef->Get("esdTree");
155    if (!esdTree) {
156       ::Error("AliITSUComparisonPileup.C", "no ESD tree found");
157       return 6;
158    }
159    event->ReadFromTree(esdTree);
160
161
162    //******* Loop over reconstructed events *********
163    Int_t ntrk=0, ntrkcor=0, ntrkwro=0;
164    Int_t e=0;
165    while (esdTree->GetEvent(e)) {
166      cout<<endl<<endl<<"********* Processing event number: "<<e<<"*******\n";
167      Int_t nn=event->GetNumberOfTracks();
168      ntrk += nn;
169
170      TClonesArray *verticesSPD=event->GetPileupVerticesSPD();
171      Int_t nfoundSPD=verticesSPD->GetEntries(); 
172      TClonesArray *verticesTRK=event->GetPileupVerticesTracks();
173      Int_t nfoundTRK=verticesTRK->GetEntries(); 
174
175      if (refTree->GetEvent(e++)==0) {
176         cerr<<"No reconstructable vertices for this event !\n";
177         continue;
178      }
179      Int_t ngood=refs->GetEntriesFast(); 
180      cout<<"Found SPD vertices: "<<nfoundSPD<<
181            "  Reconstructable vertics: "<<ngood<<endl;
182
183      h2spd->Fill(ngood,nfoundSPD);
184      h2trk->Fill(ngood,nfoundTRK);
185
186      Int_t ncor=0, nwro=0;
187      for (Int_t g=0; g<ngood; g++) {
188          Int_t Associate(const AliESDVertex *g, const AliESDVertex *f, 
189             const AliESDEvent *esd); 
190          const AliESDVertex *vtxg=(AliESDVertex*)refs->UncheckedAt(g);
191          Double_t zg=vtxg->GetZv();
192          hgood->Fill(zg);
193
194          AliESDVertex *vtxf=0;
195          Double_t zf=0.;
196          Int_t f=0;
197          for (; f<nfoundSPD; f++) {
198              vtxf=(AliESDVertex*)verticesSPD->UncheckedAt(f);
199              if (!vtxf->GetStatus()) continue;
200              if (Associate(vtxg,vtxf,event)==0) continue; 
201              break;
202          }
203          if (f>=nfoundSPD) {
204              vtxf=(AliESDVertex *)event->GetPrimaryVertexSPD();
205              if (!vtxf->GetStatus()) goto trk;
206              if (Associate(vtxg,vtxf,event)==0) goto trk; 
207          }
208
209          zf=vtxf->GetZv();
210          hfoundspd->Fill(zg);
211          hzspd->Fill(zf-zg);
212
213      trk:
214          Int_t n=0;
215          for (f=0; f<nfoundTRK; f++) {
216              vtxf=(AliESDVertex*)verticesTRK->UncheckedAt(f);
217              if (!vtxf->GetStatus()) continue;
218              n=Associate(vtxg,vtxf,event);
219              if (n < nAssMin) continue;
220              break;
221          }
222          if (f>=nfoundTRK) {
223              vtxf=(AliESDVertex*)event->GetPrimaryVertexTracks();
224              if (!vtxf->GetStatus()) continue;
225              n=Associate(vtxg,vtxf,event);
226              if (n < nAssMin) continue;
227          }
228
229          ncor+=n;
230          nwro+=(vtxf->GetNIndices()-n); 
231          zf=vtxf->GetZv();
232
233          if (Float_t(n)/vtxf->GetNIndices() > fracMin) {
234             hfoundtrk->Fill(zg);
235          } else {
236             hfaketrk->Fill(zg);
237          }
238          hztrk->Fill(zf-zg);
239
240          vtxf->SetNContributors(0); // Mark this vertex as already associated
241
242      }
243      // Increase the counter of tracks (not)associated with verices
244      ntrkcor += ncor;
245      ntrkwro += nwro;
246
247    } //***** End of the loop over reconstructed events
248
249    delete event;
250    delete esdTree;
251    ef->Close();
252
253    refFile->Close();
254    
255    cout<<"\nTotal number of found tracks: "<<ntrk<<endl;
256    cout<<"Number of tracks correctly associated with vertices: "<<ntrkcor<<endl;
257    cout<<"Number of tracks wrongly associated with vertices: "<<ntrkwro<<endl;
258    if (ntrk != 0) {
259      cout<<"Correctly associated/Found:\t"<<Float_t(ntrkcor)/ntrk<<endl;
260      cout<<"Wrongly associated/Found:\t"<<Float_t(ntrkwro)/ntrk<<endl;
261      cout<<"Not associated/Found:\t\t"<<1.-Float_t(ntrkwro+ntrkcor)/ntrk<<endl;
262    }
263
264    gStyle->SetOptStat(0);
265    gStyle->SetOptTitle(0);
266    TCanvas *c1=new TCanvas("c1","",0,0,700,1000);
267    c1->Divide(1,3);
268    c1->cd(1);
269    gPad->SetGridx(); gPad->SetGridy();
270    h2spd->Draw("box");
271    h2trk->Draw("boxsame");
272    gPad->BuildLegend(0.13,0.65,0.46,0.86)->SetFillColor(0);
273    TLine *l=new TLine(0,0,70,70);
274    l->Draw("same");
275
276    c1->cd(2);
277    gPad->SetGridx(); gPad->SetGridy();
278    heffspd->Divide(hfoundspd,hgood,1,1,"b");
279    heffspd->SetMinimum(0.); heffspd->SetMaximum(1.2);
280    heffspd->Draw("hist");
281    hefftrk->Divide(hfoundtrk,hgood,1,1,"b");
282    hefftrk->Draw("histsame");
283    heffaketrk->Divide(hfaketrk,hgood,1,1,"b");
284    heffaketrk->Draw("histsame");
285    gPad->BuildLegend(0.13,0.65,0.46,0.86)->SetFillColor(0);
286
287    c1->cd(3);
288    gPad->SetGridx(); gPad->SetGridy();
289    hztrk->Draw();
290    hzspd->Draw("same");
291    gPad->BuildLegend(0.13,0.65,0.46,0.86)->SetFillColor(0);
292
293    TFile fc("AliITSUComparisonPileup.root","RECREATE");
294    c1->Write();
295    fc.Close();
296
297    return 0;
298 }
299
300 Int_t 
301 Associate(const AliESDVertex *g,const AliESDVertex *f,const AliESDEvent *esd) { 
302    UShort_t *idxg=g->GetIndices(); Int_t ng=g->GetNIndices(); 
303    UShort_t *idxf=f->GetIndices(); Int_t nf=f->GetNIndices();
304
305    if (nf==0) { 
306    // SPD vertex
307        Double_t zg=g->GetZv();
308        Double_t zf=f->GetZv();
309        if (TMath::Abs(zf-zg)>2e-2) return 0;
310        return 1;
311    }
312    // TRK vertex
313    Int_t nass=0;
314    for (Int_t i=0; i<ng; i++) {
315        UShort_t labg=idxg[i];
316        for (Int_t j=0; j<nf; j++) {
317            const AliESDtrack *t=esd->GetTrack(idxf[j]);
318            UShort_t labf=TMath::Abs(t->GetLabel());
319            if (labg != labf) continue;
320            nass++;
321            break; 
322        }
323    } 
324
325    return nass;
326 }
327
328 Int_t GoodPileupVertices(const Char_t *dir) {
329   Int_t FindContributors(Float_t tz, AliStack *stack, UShort_t *idx);
330    if (gAlice) { 
331        delete AliRunLoader::Instance();
332        delete gAlice;//if everything was OK here it is already NULL
333        gAlice = 0x0;
334    }
335
336    Char_t fname[100];
337    sprintf(fname,"%s/galice.root",dir);
338
339    AliRunLoader *rl = AliRunLoader::Open(fname,"COMPARISON");
340    if (!rl) {
341       ::Error("GoodPileupVertices","Can't start session !");
342       return 1;
343    }
344
345    rl->LoadgAlice();
346    rl->LoadHeader();
347    rl->LoadKinematics();
348
349
350    Int_t nev=rl->GetNumberOfEvents();
351    ::Info("GoodPileupVertices","Number of events : %d\n",nev);  
352
353    sprintf(fname,"%s/GoodPileupVertices.root",dir);
354    TFile *refFile=TFile::Open(fname,"recreate");
355    TClonesArray dummy("AliESDVertex",100), *refs=&dummy;
356    TTree refTree("refTree","Tree with the reconstructable vertices");
357    refTree.Branch("Vertices",&refs);
358
359    //********  Loop over generated events 
360    for (Int_t e=0; e<nev; e++) {
361      rl->GetEvent(e);  refFile->cd();
362      AliStack* stack = rl->Stack();
363
364      AliHeader *ah=rl->GetHeader();
365      AliGenCocktailEventHeader *cock=
366             (AliGenCocktailEventHeader*)ah->GenEventHeader();
367      TList *headers=cock->GetHeaders();
368      const Int_t nvtx=headers->GetEntries();
369      const Int_t np=ah->GetNtrack();
370      cout<<"Event "<<e<<" Number of vertices, particles: "
371          <<nvtx<<' '<<np<<endl;
372
373      Int_t nv=0;
374      for (Int_t v=0; v<nvtx; v++) {
375          AliGenEventHeader *h=(AliGenEventHeader *)headers->At(v);
376          TArrayF vtx(3); h->PrimaryVertex(vtx);
377          Float_t t=h->InteractionTime();
378          UShort_t *idx=new UShort_t[np];
379          Int_t ntrk=FindContributors(t,stack,idx);
380          if (ntrk < nMin) continue;
381          AliESDVertex *vertex=new ((*refs)[nv]) AliESDVertex();
382          vertex->SetXv(vtx[0]);
383          vertex->SetYv(vtx[1]);
384          vertex->SetZv(vtx[2]);
385          vertex->SetNContributors(ntrk);
386          vertex->SetIndices(ntrk,idx);
387          delete idx;
388          nv++;
389      }
390      refTree.Fill();
391      refs->Clear();
392    } //*** end of the loop over generated events
393
394    refTree.Write();
395    refFile->Close();
396
397    delete rl;
398    return 0;
399 }
400
401 Int_t FindContributors(Float_t tz, AliStack *stack, UShort_t *idx) {
402   Int_t ntrk=0;
403   Int_t np=stack->GetNtrack();
404   for (Int_t i=0; i<np; i++) {
405       if (!stack->IsPhysicalPrimary(i)) continue;
406       TParticle *part=stack->Particle(i);
407       if (!part) continue;
408       TParticlePDG *partPDG = part->GetPDG();
409       if (!partPDG) continue;
410       if (TMath::Abs(partPDG->Charge())<1e-10) continue;
411       if (part->Pt() < pTmin) continue;
412       Float_t dt=0.5*(tz-part->T())/(tz+part->T());
413       if (TMath::Abs(dt)>1e-5) continue;
414       idx[ntrk]=i;
415       ntrk++;
416   }
417   return ntrk;
418 }