Corrected protection.
[u/mrichter/AliRoot.git] / STEER / AliKFParticleTest.C
1
2 //---------------------------------------------------------------------------------
3 // The example of usage of AliKFParticle & AliKFVertex classes for V0 analysis
4 // .
5 // @author  S.Gorbunov, I.Kisel
6 // @version 1.0
7 // @since   13.05.07
8 // 
9 // The AliKFParticleTest macro contains a toy V0 finder for ESD tracks.
10 // At the first step, event primary vertex is reconstructed. 
11 // At the second step, ideal PID hypothesis are assigned to all the particles.
12 // At the third step, V0 candidates are constructed for each pair 
13 // of positive and negative particles.
14 // V0 candidate considered as good when it passes Chi^2 cut and 
15 // it is placed >= 3 Sigma away from the primary vertex.
16 // Invariant mass distribution for all good V0 candidates is plotted.
17 //
18 //  -= Copyright &copy ALICE HLT Group =-
19 //_________________________________________________________________________________
20
21
22
23 #if !defined( __CINT__) || defined(__MAKECINT__)
24
25   //gSystem->AddIncludePath("-I\"$ALICE_ROOT/include\"");
26
27   #include <Riostream.h>
28   #include <TTree.h>
29   #include <TFile.h>
30   #include <TCanvas.h>
31   #include <TH1D.h>
32   #include <TLatex.h>
33   #include <TDatabasePDG.h>
34   #include <TParticle.h>
35
36   #include "AliRun.h"
37   #include "AliStack.h"
38   #include "AliESDEvent.h"
39   #include "AliTracker.h"
40   #include "AliKFParticle.h"
41   #include "AliKFVertex.h"
42
43 #endif
44
45
46 Bool_t old_file = 1;
47 TCanvas *canvas=0;
48 TH1D  *histoMass[3]={0,0,0};
49 Int_t histoPDG [4]= {22,310,3122,421};
50 char* histoName[4]= {"#gamma","K^{0}_{s}","#Lambda","D^{0}"};
51 TLatex *histoText[4]={0,0,0,0};
52
53 void DrawV0() 
54 {
55   //* Draw the invariant mass histogram
56
57   if( !canvas ) return;
58   histoMass[0]->Draw();
59   for( Int_t i=1; i<3; i++ ) histoMass[i]->Draw("same");
60   for( Int_t i=0; i<4; i++ ){
61     Double_t mass = TDatabasePDG::Instance()->GetParticle(histoPDG[i])->Mass();    
62     Int_t bin = histoMass[2]->FindBin(mass) -5;
63     if( bin<0 ) bin =0;
64     Double_t max = 0;
65     for( Int_t j=bin; j<bin+10; j++ ) 
66     if( max<histoMass[2]->GetBinContent(j) ) max = histoMass[2]->GetBinContent(j);
67     histoText[i]->SetY(max);
68     histoText[i]->SetX(mass+.05);
69     if(max>0) histoText[i]->Draw();
70   }
71   if( canvas ) canvas->Update();
72 }
73
74 void StartV0()
75 {
76   //* (Re)create histograms and fill them from temporary file
77   
78   TDirectory *curr = gDirectory;
79   canvas = new TCanvas();
80   histoMass[0] = new TH1D("massAll","AliKFParticle test", 500,0,3);
81   histoMass[0]->SetXTitle("V^{0} invariant mass [GeV]");
82   histoMass[1] = new TH1D("massMulti","V^{+-} contributions", 500,0,3);
83   histoMass[2] = new TH1D("massV0","V^{0} signal", 500,0,3);
84   histoMass[0]->SetLineColor(8);
85   histoMass[0]->SetFillColor(8);
86   histoMass[1]->SetLineColor(kMagenta);
87   histoMass[1]->SetFillColor(kMagenta);
88   histoMass[2]->SetLineColor(kBlue);
89   histoMass[2]->SetFillColor(kBlue);
90   for( Int_t i=0; i<4; i++ ){
91     histoText[i] = new TLatex(0,0,histoName[i]);
92     histoText[i]->SetTextColor(kBlue);
93   }
94   TFile* file = new TFile("AliKFParticleTest.root", "READ" );
95   if( !old_file && file ){
96     file->cd();
97     for( Int_t i=0; i<3; i++ ){
98       TObject *old = file->FindObjectAny(histoMass[i]->GetName());
99       if( old && histoMass[i] ) histoMass[i]->Add( (TH1D*) old);
100     }
101     file->Close();
102   }
103   curr->cd();
104   DrawV0();
105 }
106
107 void EndV0()
108 {
109   //* End of gAlice -> store all histos in temporary file, clean memory
110
111   TDirectory *curr = gDirectory;
112   TFile* file = new TFile("AliKFParticleTest.root", "RECREATE" );
113   old_file = 0;
114   file->cd();
115   for( Int_t i=0; i<3; i++ ) histoMass[i]->Write();
116   file->Close();
117   for( Int_t i=0; i<3; i++ ) delete histoMass[i];
118   for( Int_t i=0; i<4; i++ ) delete histoText[i];
119   delete canvas;
120   curr->cd();
121 }
122
123 class TESDTrackInfo
124 {
125 public:
126   TESDTrackInfo(){}
127   AliKFParticle fParticle; //* assigned KFParticle
128   Bool_t fPrimUsedFlag;    //* flag shows that the particle was used for primary vertex fit
129   Bool_t fOK;              //* is the track good enough
130   Int_t mcPDG;             //* Monte Carlo PDG code
131   Int_t mcMotherID;        //* Monte Carlo ID of its mother
132 };
133
134 void RunV0(  AliESDEvent *event )
135 {
136   //* V0 finder
137
138   static Int_t iEvent = 0;
139   cout<<"Event "<<iEvent++<<endl;
140
141   if( !gAlice ) return;
142   AliRunLoader *rl = AliRunLoader::Instance(); 
143  
144   AliStack *stack = rl->Stack();
145   if( !stack ) return;  
146
147   Int_t nESDTracks=event->GetNumberOfTracks();
148   if( nESDTracks>1000 ) nESDTracks=1000;
149
150
151   TESDTrackInfo ESDTrackInfo[1000]; //* parallel to ESD tracks
152
153   //* Fill ESDTrackInfo array
154
155   for (Int_t iTr=0; iTr<nESDTracks; iTr++)
156     {   
157       TESDTrackInfo &info = ESDTrackInfo[iTr];
158       info.fOK = 0;
159       info.fPrimUsedFlag = 0;
160       info.mcPDG = -1;
161       info.mcMotherID = -1;
162
163       //* track quality check
164
165       AliESDtrack *pTrack = event->GetTrack(iTr);    
166       if( !pTrack  ) continue;
167       if (pTrack->GetKinkIndex(0)>0) continue;
168       if ( !( pTrack->GetStatus()&AliESDtrack::kITSrefit ) ) continue;
169       Int_t indi[12];
170       if( pTrack->GetITSclusters(indi) <5 ) continue;
171       Int_t PDG = ( pTrack->GetSigned1Pt() <0 ) ?321 :211;
172
173       //* take MC PDG  
174       { 
175         Int_t mcID = TMath::Abs(pTrack->GetLabel());
176         TParticle * part = stack->Particle(TMath::Abs(mcID));
177         if( part ){
178           info.mcPDG = part->GetPdgCode();
179           PDG = info.mcPDG;
180           if( mcID>=0 ) info.mcMotherID = part->GetFirstMother();
181         }    
182       }
183
184       //* Construct KFParticle for the track
185
186       info.fParticle = AliKFParticle( *pTrack, PDG );
187       info.fOK = 1;   
188     }
189
190   //* Find event primary vertex
191   
192   AliKFVertex primVtx;  
193   {
194     const AliKFParticle * vSelected[1000]; //* Selected particles for vertex fit
195     Int_t vIndex[1000];                    //* Indices of selected particles
196     Bool_t vFlag[1000];                    //* Flags returned by the vertex finder
197
198     Int_t nSelected = 0;
199     for( Int_t i = 0; i<nESDTracks; i++){ 
200       if(ESDTrackInfo[i].fOK ){
201         vSelected[nSelected] = &(ESDTrackInfo[i].fParticle);
202         vIndex[nSelected] = i;
203         nSelected++;
204       }
205     }
206     primVtx.ConstructPrimaryVertex( vSelected, nSelected, vFlag, 3. );
207     for( Int_t i = 0; i<nSelected; i++){ 
208       if( vFlag[i] ) ESDTrackInfo[vIndex[i]].fPrimUsedFlag = 1;
209     }
210     if( primVtx.GetNDF() <1 ) return; //* Less then two tracks in primary vertex 
211   }
212
213   //* V0 finder
214
215   for( Int_t iTr = 0; iTr<nESDTracks; iTr++ ){ //* first daughter
216
217     TESDTrackInfo &info = ESDTrackInfo[iTr];
218     if( !info.fOK ) continue;    
219
220     for( Int_t jTr = iTr+1; jTr<nESDTracks; jTr++ ){  //* second daughter
221       TESDTrackInfo &jnfo = ESDTrackInfo[jTr];
222       if( !jnfo.fOK ) continue;
223       
224       //* check for different charge
225
226       if( info.fParticle.GetQ() == jnfo.fParticle.GetQ() ) continue;      
227
228       //* construct V0 mother
229
230       AliKFParticle V0( info.fParticle, jnfo.fParticle );     
231
232       //* check V0 Chi^2
233
234       if( V0.GetNDF()<1 ) continue;
235       if( TMath::Sqrt(TMath::Abs(V0.GetChi2()/V0.GetNDF())) >3. ) continue;
236
237       //* subtruct daughters from primary vertex 
238
239       AliKFVertex primVtxCopy = primVtx;
240        
241       if( info.fPrimUsedFlag ) primVtxCopy -= info.fParticle;
242       if( jnfo.fPrimUsedFlag ) primVtxCopy -= jnfo.fParticle;
243
244       //* Check V0 Chi^2 deviation from primary vertex 
245
246       if( V0.GetDeviationFromVertex( primVtxCopy ) >3. ) continue;
247
248       //* Add V0 to primary vertex to improve the primary vertex resolution
249
250       primVtxCopy += V0;      
251
252       //* Set production vertex for V0
253
254       V0.SetProductionVertex( primVtxCopy );
255
256       //* Check chi^2 for a case
257
258       if( TMath::Sqrt( TMath::Abs(V0.GetChi2()/V0.GetNDF()) >3. )) continue;
259
260       //* Get V0 decay length with estimated error
261
262       Double_t length, sigmaLength;
263       if( V0.GetDecayLength( length, sigmaLength ) ) continue;
264
265       //* Reject V0 if it decays too close to the primary vertex
266
267       if( length  <3.*sigmaLength ) continue;
268
269       //* Get V0 invariant mass and plot it
270
271       Double_t mass, sigmaMass;
272       if( V0.GetMass( mass, sigmaMass ) ) continue;
273       histoMass[0]->Fill(mass);            
274
275       //* Fill signal histograms using MC information
276
277       if( info.mcMotherID==jnfo.mcMotherID && info.mcMotherID>=0 ){
278         TParticle *mother = stack->Particle(info.mcMotherID);
279         if( mother && TMath::Abs(mother->GetPdgCode())!=21 ){
280           histoMass[1]->Fill(mass);
281           if( mother->GetNDaughters()==2 ){
282             histoMass[2]->Fill(mass);   
283           }
284           cout<<"PDG V0,pi,pj, ndaughters, mc mass, reco mass = "<<mother->GetPdgCode()<<","<<info.mcPDG<<","<<jnfo.mcPDG<<", "
285               << mother->GetNDaughters()<<", "<<mother->GetMass()<<", "<<mass<<endl;
286         }
287       }
288     }
289   }
290   if( iEvent %100 == 0 || (iEvent<100 && iEvent %10==0)) DrawV0();     
291 }
292
293
294
295 Int_t AliKFParticleTest(Int_t n1=0,Int_t n2=1000,char *dire="/d/alice10/sma/my_v4-05-Release/pp/"){
296   //* Main macro
297     
298   //  LOOP  OVER  SERIES  OF  DIRECTORIES
299   
300   for (Int_t ifi=n1; ifi<=n2; ifi++) {
301    
302     char nstring[5], filename[100], esdfile[100];
303     sprintf(nstring,"%3.3d",ifi);
304     sprintf(filename,"%s%s/galice.root",dire,nstring);
305     sprintf(esdfile,"%s%s/AliESDs.root",dire,nstring);
306     
307     cout <<"  Opening "<<filename<<"\nand ESD "<<esdfile<<endl;
308     
309     if (gAlice) {
310       delete AliRunLoader::Instance();
311       delete gAlice;
312       gAlice=0;
313     }
314     
315     AliRunLoader *rl = AliRunLoader::Open(filename);
316     if ( !rl ) { 
317       ::Error("AliKFParticleTest.C","Can not open session !");
318       continue;
319     }
320     if (rl->LoadgAlice()) {
321       ::Error("AliKFParticleTest.C","LoadgAlice returned error !");
322       delete rl;
323       continue;
324     }
325     if (rl->LoadHeader()) {
326       ::Error("AliKFParticleTest.C","LoadHeader returned error !");
327       delete rl;
328       continue;
329     }
330     rl->LoadKinematics();
331     AliTracker::SetFieldMap(gAlice->Field(),1);
332     AliKFParticle::SetField( AliTracker::GetBz() );
333
334     //---------------------------------------//
335     //                                       //
336     //               ESD  file               //
337     //                                       //
338     //---------------------------------------//
339     
340     // Open file with the ESD
341     TFile *ef=TFile::Open(esdfile);
342
343     //Check if the file could be opened
344     if (!ef || !ef->IsOpen()) {cerr<<"Error open AliESDs.root !\n"; continue ;}
345
346     //create event object
347     AliESD *event = new AliESDEvent;
348
349     //Set pointer to the esd tree in the file
350     TTree* tree = (TTree*) ef->Get("esdTree");
351     
352     //check if the tree exists
353     if (!tree) {cerr<<"no ESD tree found\n"; continue;};
354     
355     //Set pointer to the esd object in the tree
356     event->ReadFromTree(event);
357     
358     //Number of events
359     Int_t nevents=tree->GetEntriesFast();
360     cout << "Number of events: " << nevents << endl;
361
362     StartV0();
363     for (Int_t iev=0; iev<nevents; iev++){
364       tree->GetEvent(iev);
365       rl->GetEvent(iev);
366       RunV0(event);
367     }
368     EndV0();
369     delete event;
370     ef->Close();
371   }
372   StartV0();
373   return 0;
374 }