Implemented a remapper for parent identification in the particle tree when adding...
[u/mrichter/AliRoot.git] / PWGPP / AliKFParticleTest.h
1 //---------------------------------------------------------------------------------
2 // The example of usage of AliKFParticle & AliKFVertex classes for V0 analysis
3 // .
4 // @author  S.Gorbunov, I.Kisel
5 // @version 1.0
6 // @since   13.05.07
7 // 
8 // The AliKFParticleTest macro contains a toy V0 finder for ESD tracks.
9 // At the first step, event primary vertex is reconstructed. 
10 // At the second step, ideal PID hypothesis are assigned to all the particles.
11 // At the third step, V0 candidates are constructed for each pair 
12 // of positive and negative particles.
13 // V0 candidate considered as good when it passes Chi^2 cut and 
14 // it is placed >= 3 Sigma away from the primary vertex.
15 // Invariant mass distribution for all good V0 candidates is plotted.
16 //
17 //  -= Copyright &copy ALICE HLT Group =-
18 //_________________________________________________________________________________
19
20
21 #ifndef ALIKFPARTICLETEST_H
22 #define ALIKFPARTICLETEST_H
23
24 #include "TH1.h"
25 #include "TCanvas.h"
26 #include "AliESD.h"
27 #include "AliAnalysisTaskRL.h"
28
29 class AliKFParticleTest: public AliAnalysisTaskRL {
30  public:
31   AliKFParticleTest(const char *name);
32   virtual ~AliKFParticleTest() {}
33   
34   virtual void   ConnectInputData(Option_t *);
35   virtual void   CreateOutputObjects();
36   virtual void   Exec(Option_t *option);
37   virtual void   Terminate(Option_t *);
38   
39  private:
40   AliESD *fESD; //ESD object
41   TH1D   *fHistoMassAll;
42   TH1D   *fHistoMassSignal;
43   TCanvas *fCanvas;
44   void DrawV0(); 
45
46   TObjArray * fOutputContainer; // ! output data container
47
48   ClassDef(AliKFParticleTest, 0); // example of analysis
49 };
50
51
52 #include "Riostream.h"
53 #include "TChain.h"
54 #include "TSystem.h"
55 #include "TROOT.h"
56 #include "TParticle.h"
57 #include "TRandom.h"
58 #include "AliESD.h"
59 #include "AliLog.h"
60 #include "AliStack.h"
61 #include "AliTracker.h"
62 #include "TLatex.h"
63 #include "TDatabasePDG.h"
64 #include "AliKFParticle.h"
65 #include "AliKFVertex.h"
66
67
68 ClassImp(AliKFParticleTest);
69
70
71 //________________________________________________________________________
72 AliKFParticleTest::AliKFParticleTest(const char *name) :AliAnalysisTaskRL(name,""), fESD(0) {
73   // Constructor.
74   // Input slot #0 works with an Ntuple
75   DefineInput(0, TChain::Class());
76  // Output slot #0 writes into a TH1 container
77   DefineOutput(0, TObjArray::Class());
78 }
79
80 //___________________________________________________________________________
81 void AliKFParticleTest::ConnectInputData(Option_t *) {
82   // Initialize branches.
83   printf("   ConnectInputData of task %s\n", GetName());
84   if (!fESD) {
85     char ** address = (char **)GetBranchAddress(0, "ESD");
86     if (address) fESD = (AliESD*)(*address);
87     if (!fESD) {
88       fESD = new AliESD();
89       SetBranchAddress(0, "ESD", &fESD);
90     }
91   }
92
93 }
94
95 //___________________________________________________________________________
96 void AliKFParticleTest::CreateOutputObjects() {
97   printf("   CreateOutputObjects of task %s\n", GetName());
98
99   if ( !gROOT->IsBatch() ) fCanvas = new TCanvas();
100   else fCanvas = 0;
101
102   fHistoMassAll = new TH1D("massAll","AliKFParticle test", 500,0,3);
103   fHistoMassAll->SetXTitle("V^{0} invariant mass [GeV]");
104   fHistoMassAll->SetLineColor(kGreen);
105   fHistoMassAll->SetFillColor(kGreen);
106
107   fHistoMassSignal = new TH1D("massSignal","V^{0} signal", 500,0,3);
108   fHistoMassSignal->SetXTitle("V^{0} invariant mass [GeV]");
109   fHistoMassSignal->SetLineColor(kBlue);
110   fHistoMassSignal->SetFillColor(kBlue);
111
112   fOutputContainer = new TObjArray(1);
113   fOutputContainer->SetName(GetName()) ;
114   fOutputContainer->Add(fHistoMassAll);
115   fOutputContainer->Add(fHistoMassSignal);
116 }
117
118
119 void AliKFParticleTest::DrawV0() 
120 {
121   //* Draw the invariant mass histogram
122
123   if ( gROOT->IsBatch() ) return;
124
125   const Int_t histoPDG [4]= {22,310,3122,421};
126   const Char_t* histoName[4]= {"#gamma","K^{0}_{s}","#Lambda","D^{0}"};
127   TLatex histoText[4];
128
129   if( !fCanvas ) return;
130   fCanvas->Clear();
131   fCanvas->cd();
132
133   for( Int_t i=0; i<4; i++ ){
134     histoText[i].SetTextColor(kBlue);
135   }
136
137   fHistoMassAll->Draw();
138   fHistoMassSignal->Draw("same");
139
140   for( Int_t i=0; i<4; i++ ){
141     Double_t mass = TDatabasePDG::Instance()->GetParticle(histoPDG[i])->Mass();    
142     Int_t bin = fHistoMassSignal->FindBin(mass) -5;
143     if( bin<0 ) bin =0;
144     Double_t max = 0;
145     for( Int_t j=bin; j<bin+10; j++ ) 
146     if( max<fHistoMassSignal->GetBinContent(j) ) max = fHistoMassSignal->GetBinContent(j);
147     if(max>0) histoText[i].DrawLatex( mass+.05, max, histoName[i] );
148   }
149   fCanvas->Update();
150 }
151
152 //________________________________________________________________________
153 void AliKFParticleTest::Exec(Option_t *) {
154
155   static Int_t iEvent = 0;
156   if( ++iEvent%100 ==0 ) cout<<"Event "<<iEvent<<endl;
157
158   // Get input data
159   TTree *tinput = (TTree*)GetInputData(0);
160   Long64_t ientry = tinput->GetReadEntry();
161   if (AliAnalysisTaskRL::GetEntry(ientry) == kFALSE) {
162     printf("Couldn't get event from the runLoader\n");
163     return;
164   }  
165   if (!fESD) return;
166
167   Double_t Bz = fESD->GetMagneticField();
168   AliKFParticle::SetField( Bz );
169
170   if (ientry==0) Notify ();
171
172  //cout <<"Event "<<ientry<<endl;
173
174   AliStack* stack = GetStack();
175   if (!stack) {
176     AliDebug(AliLog::kError, "Stack not available");
177     return;
178   }
179
180
181   class TESDTrackInfo
182   {
183   public:
184     TESDTrackInfo(){}
185     AliKFParticle fParticle; //* assigned KFParticle
186     Bool_t fPrimUsedFlag;    //* flag shows that the particle was used for primary vertex fit
187     Bool_t fOK;              //* is the track good enough
188     Int_t mcPDG;             //* Monte Carlo PDG code
189     Int_t mcMotherID;        //* Monte Carlo ID of its mother
190   };
191
192   Int_t nESDTracks=fESD->GetNumberOfTracks();
193   if( nESDTracks>1000 ) nESDTracks=1000;
194
195   TESDTrackInfo ESDTrackInfo[1000]; //* parallel to ESD tracks
196
197   //* Fill ESDTrackInfo array
198
199   for (Int_t iTr=0; iTr<nESDTracks; iTr++)
200     {   
201       TESDTrackInfo &info = ESDTrackInfo[iTr];
202       info.fOK = 0;
203       info.fPrimUsedFlag = 0;
204       info.mcPDG = -1;
205       info.mcMotherID = -1;
206
207       //* track quality check
208
209       AliESDtrack *pTrack = fESD->GetTrack(iTr);    
210       if( !pTrack  ) continue;
211       if (pTrack->GetKinkIndex(0)>0) continue;
212       if ( !( pTrack->GetStatus()&AliESDtrack::kITSrefit ) ) continue;
213       Int_t indi[12];
214       if( pTrack->GetITSclusters(indi) <5 ) continue;
215       Int_t PDG = ( pTrack->Get1Pt() <0 ) ?321 :211;
216
217       //* take MC PDG  
218       { 
219         Int_t mcID = TMath::Abs(pTrack->GetLabel());
220         TParticle * part = stack->Particle(TMath::Abs(mcID));
221         if( part ){
222           info.mcPDG = part->GetPdgCode();
223           PDG = info.mcPDG;
224           if( mcID>=0 ) info.mcMotherID = part->GetFirstMother();
225         }    
226       }
227
228       //* Construct KFParticle for the track
229
230       info.fParticle = AliKFParticle( *pTrack, PDG );
231       info.fOK = 1;   
232     }
233
234   //* Find event primary vertex
235   
236   AliKFVertex primVtx;  
237   {
238     const AliKFParticle * vSelected[1000]; //* Selected particles for vertex fit
239     Int_t vIndex[1000];                    //* Indices of selected particles
240     Bool_t vFlag[1000];                    //* Flags returned by the vertex finder
241
242     Int_t nSelected = 0;
243     for( Int_t i = 0; i<nESDTracks; i++){ 
244       if(ESDTrackInfo[i].fOK ){
245         vSelected[nSelected] = &(ESDTrackInfo[i].fParticle);
246         vIndex[nSelected] = i;
247         nSelected++;
248       }
249     }
250     primVtx.ConstructPrimaryVertex( vSelected, nSelected, vFlag, 3. );
251     for( Int_t i = 0; i<nSelected; i++){ 
252       if( vFlag[i] ) ESDTrackInfo[vIndex[i]].fPrimUsedFlag = 1;
253     }
254     if( primVtx.GetNDF() <1 ) return; //* Less then two tracks in primary vertex 
255   }
256
257   //* V0 finder
258
259   for( Int_t iTr = 0; iTr<nESDTracks; iTr++ ){ //* first daughter
260
261     TESDTrackInfo &info = ESDTrackInfo[iTr];
262     if( !info.fOK ) continue;    
263
264     for( Int_t jTr = iTr+1; jTr<nESDTracks; jTr++ ){  //* second daughter
265       TESDTrackInfo &jnfo = ESDTrackInfo[jTr];
266       if( !jnfo.fOK ) continue;
267       
268       //* check for different charge
269
270       if( info.fParticle.GetQ() == jnfo.fParticle.GetQ() ) continue;      
271
272       //* construct V0 mother
273
274       AliKFParticle V0( info.fParticle, jnfo.fParticle );     
275
276       //* check V0 Chi^2
277
278       if( V0.GetNDF()<1 ) continue;
279       if( TMath::Sqrt(TMath::Abs(V0.GetChi2()/V0.GetNDF())) >3. ) continue;
280
281       //* subtruct daughters from primary vertex 
282
283       AliKFVertex primVtxCopy = primVtx;
284        
285       if( info.fPrimUsedFlag ) primVtxCopy -= info.fParticle;
286       if( jnfo.fPrimUsedFlag ) primVtxCopy -= jnfo.fParticle;
287
288       //* Check V0 Chi^2 deviation from primary vertex 
289
290       if( V0.GetDeviationFromVertex( primVtxCopy ) >3. ) continue;
291
292       //* Add V0 to primary vertex to improve the primary vertex resolution
293
294       primVtxCopy += V0;      
295
296       //* Set production vertex for V0
297
298       V0.SetProductionVertex( primVtxCopy );
299
300       //* Check chi^2 for a case
301
302       if( TMath::Sqrt( TMath::Abs(V0.GetChi2()/V0.GetNDF()) >3. )) continue;
303
304       //* Get V0 decay length with estimated error
305
306       Double_t length, sigmaLength;
307       if( V0.GetDecayLength( length, sigmaLength ) ) continue;
308
309       //* Reject V0 if it decays too close to the primary vertex
310
311       if( length  <3.*sigmaLength ) continue;
312
313       //* Get V0 invariant mass and plot it
314
315       Double_t mass, sigmaMass;
316       if( V0.GetMass( mass, sigmaMass ) ) continue;
317       fHistoMassAll->Fill(mass);            
318
319       //* Fill signal histograms using MC information
320
321       if( info.mcMotherID==jnfo.mcMotherID && info.mcMotherID>=0 ){
322         TParticle *mother = stack->Particle(info.mcMotherID);
323         if( mother && TMath::Abs(mother->GetPdgCode())!=21 ){    
324           if( mother->GetNDaughters()==2 ){
325             fHistoMassSignal->Fill(mass);       
326           }
327           cout<<"PDG V0,pi,pj, ndaughters, mc mass, reco mass = "<<mother->GetPdgCode()<<","<<info.mcPDG<<","<<jnfo.mcPDG<<", "
328               << mother->GetNDaughters()<<", "<<mother->GetMass()<<", "<<mass<<endl;
329         }
330       }
331     }
332   }
333   if( iEvent %1000 == 0 || (iEvent ==200)) DrawV0();     
334
335   // Post final data. It will be written to a file with option "RECREATE"
336   PostData(0, fOutputContainer);
337 }      
338
339
340 //________________________________________________________________________
341 void AliKFParticleTest::Terminate(Option_t *) 
342 {
343   DrawV0();
344   fOutputContainer = (TObjArray*)GetOutputData(0);
345   //fHistoMassAll=(TH1D*)fOutputContainer->At(0) ;
346 }
347
348 #endif