PID info added
[u/mrichter/AliRoot.git] / PWGCF / Correlations / Base / AliAnalysisTaskCFTree.cxx
1 /**************************************************************************
2  * Copyright(c) 1998-1999, ALICE Experiment at CERN, All rights reserved. *
3  *                                                                        *
4  * Author: The ALICE Off-line Project.                                    *
5  * Contributors are mentioned in the code where appropriate.              *
6  *                                                                        *
7  * Permission to use, copy, modify and distribute this software and its   *
8  * documentation strictly for non-commercial purposes is hereby granted   *
9  * without fee, provided that the above copyright notice appears in all   *
10  * copies and that both the copyright notice and this permission notice   *
11  * appear in the supporting documentation. The authors make no claims     *
12  * about the suitability of this software for any purpose. It is          *
13  * provided "as is" without express or implied warranty.                  *
14  **************************************************************************/
15
16 // Analysis task to produce trees of lightweight events
17 // evgeny.kryshen@cern.ch
18
19 // aliroot
20 #include "AliAnalysisTaskCFTree.h"
21 #include "AliCFParticle.h"
22 #include "AliAnalysisManager.h"
23 #include "AliInputEventHandler.h"
24 #include "AliESDEvent.h"
25 #include "AliAODEvent.h"
26 #include "AliMCEvent.h"
27 #include "AliHeader.h"
28 #include "AliESDHeader.h"
29 #include "AliAODHeader.h"
30 #include "AliAODMCHeader.h"
31 #include "AliGenEventHeader.h"
32 #include "AliGenCocktailEventHeader.h"
33 #include "AliAODTrack.h"
34 #include "AliESDtrack.h"
35 #include "AliMCParticle.h"
36 #include "AliAODMCParticle.h"
37 #include "AliExternalTrackParam.h"
38 #include "AliCentrality.h"
39 #include "AliAnalysisFilter.h"
40 #include "AliPIDResponse.h"
41 #include "AliTPCPIDResponse.h"
42 // root
43 #include "TMath.h"
44 #include "TFile.h"
45 #include "TList.h"
46 #include "TH1I.h"
47 #include "TChain.h"
48 #include "TTree.h"
49 #include "TClonesArray.h"
50 ClassImp(AliAnalysisTaskCFTree)
51
52 //-----------------------------------------------------------------------------
53 AliAnalysisTaskCFTree::AliAnalysisTaskCFTree(const char* name) :
54   AliAnalysisTask(name,""),
55   fDebug(0),
56   fMode(0),
57   fInputHandler(0x0),
58   fMcHandler(0x0),
59   fTrackFilter(0x0),
60   fHybridConstrainedMask(0),
61   fTPConlyConstrainedMask(0),
62   fPIDResponse(0x0),
63   fListOfHistos(0x0),
64   fEventStatistics(0x0),
65   fTree(0x0),
66   fParticles(0x0),
67   fField(0),
68   fCentrality(0),
69   fZvtx(0),
70   fRunNumber(0),
71   fPeriod(0),
72   fOrbit(),
73   fBc(),
74   fSelectMask(0),
75   fSelectBit(AliVEvent::kMB),
76   fZVertexCut(10.),
77   fnTracksVertex(1),
78   fCentralityMethod("V0M"),
79   fTrackFilterBit(0xffffffff),
80   fTrackEtaCut(1.0),
81   fPtMin(0.15),
82   fSharedClusterCut(0.4),
83   fCrossedRowsCut(100),
84   fFoundFractionCut(0.8)
85 {
86   Info("AliAnalysisTaskCFTree","Calling Constructor");
87   DefineInput(0,TChain::Class());
88   DefineOutput(0,TList::Class());
89   DefineOutput(1,TTree::Class());
90 }
91 //-----------------------------------------------------------------------------
92
93
94 //-----------------------------------------------------------------------------
95 void AliAnalysisTaskCFTree::ConnectInputData(Option_t* /*option*/){
96   fInputHandler = dynamic_cast<AliInputEventHandler*> (AliAnalysisManager::GetAnalysisManager()->GetInputEventHandler());
97   fMcHandler    = dynamic_cast<AliInputEventHandler*> (AliAnalysisManager::GetAnalysisManager()->GetMCtruthEventHandler());
98 }
99 //-----------------------------------------------------------------------------
100
101
102 //-----------------------------------------------------------------------------
103 void AliAnalysisTaskCFTree::CreateOutputObjects(){
104   fListOfHistos = new TList();
105   fListOfHistos->SetOwner();
106   fEventStatistics = new TH1I("fEventStatistics","",10,0,10);
107   fEventStatistics->SetBit(TH1::kCanRebin);
108
109   fListOfHistos->Add(fEventStatistics);
110
111   fParticles = new TClonesArray("AliCFParticle",2000);
112   // create file-resident tree
113   TDirectory *owd = gDirectory;
114   OpenFile(1);
115   fTree = new TTree("events","events");
116   owd->cd();
117   fTree->Branch("cent",&fCentrality);
118   fTree->Branch("zvtx",&fZvtx);
119   fTree->Branch("field",&fField);
120   fTree->Branch("run",&fRunNumber);
121   fTree->Branch("period",&fPeriod);
122   fTree->Branch("orbit",&fOrbit);
123   fTree->Branch("bc",&fBc);
124   fTree->Branch("mask",&fSelectMask);
125   fTree->Branch("particles",&fParticles);
126   PostData(0,fListOfHistos);
127   PostData(1,fTree);
128 }
129 //-----------------------------------------------------------------------------
130
131
132 //-----------------------------------------------------------------------------
133 void AliAnalysisTaskCFTree::Exec(Option_t *){
134   fParticles->Clear();
135   
136   fEventStatistics->Fill("before cuts",1);
137   
138   AliVEvent* event = fInputHandler->GetEvent();
139   if (!event) return;
140   fEventStatistics->Fill("after event check",1);
141   
142   fPIDResponse = fInputHandler->GetPIDResponse();
143   
144   fSelectMask = fInputHandler->IsEventSelected();
145   if (!(fSelectMask & fSelectBit)) return;
146   fEventStatistics->Fill("after trigger selection",1);
147   fRunNumber  = event->GetRunNumber();
148   fPeriod     = event->GetPeriodNumber();
149   fOrbit      = event->GetOrbitNumber();
150   fBc         = event->GetBunchCrossNumber();
151   fField      = event->GetMagneticField();
152   fCentrality = event->GetCentrality()->GetCentralityPercentile(fCentralityMethod);
153   
154   if (fMode==0){ // Data analysis
155     const AliVVertex* vertex  = event->GetPrimaryVertex();
156     if (!vertex) return;
157     fZvtx  = vertex->GetZ();
158     TString name(vertex->GetName());
159     if (name.CompareTo("PrimaryVertex") && name.CompareTo("SPDVertex")) return; // Reject TPC only vertex
160     fEventStatistics->Fill("after rejection of TPC only vertex",1);
161     if (TMath::Abs(fZvtx) >= fZVertexCut || vertex->GetNContributors()<fnTracksVertex)  return;
162     fEventStatistics->Fill("after vertex selection",1);
163     
164     for (Int_t ipart=0;ipart<event->GetNumberOfTracks();ipart++){
165       AliVTrack* track = (AliVTrack*) event->GetTrack(ipart);
166       if (!track) continue;
167       UInt_t mask = GetFilterMap(track);
168       if (!(mask & fTrackFilterBit)) continue;
169
170       if (track->InheritsFrom("AliAODTrack")) AddTrack(track,mask,0);
171       else if (track->InheritsFrom("AliESDtrack")) {
172         if (mask)                           AddTrack(track,mask,1);
173         if (mask & fHybridConstrainedMask)  AddTrack(track,mask,2);
174         if (mask & fTPConlyConstrainedMask) AddTrack(track,mask,3);
175       }
176     }
177   }
178   else { // MC analysis
179     AliMCEvent* mcEvent = fMcHandler ? fMcHandler->MCEvent() : 0;
180     TClonesArray* mcTracks = dynamic_cast<TClonesArray*>(event->FindListObject(AliAODMCParticle::StdBranchName()));
181     if (!mcEvent && !mcTracks) { printf("No mc object found\n"); return; }
182     fEventStatistics->Fill("after mc objects check",1);
183
184     Int_t nPrimGen = 0;
185     Int_t nProduced = 0;
186     if (mcEvent) { // ESD
187       AliHeader* header = (AliHeader*) mcEvent->Header();
188       AliGenCocktailEventHeader* cocktailHeader = dynamic_cast<AliGenCocktailEventHeader*> (header->GenEventHeader());
189       AliGenEventHeader* mcHeader = dynamic_cast<AliGenEventHeader*> (cocktailHeader ? cocktailHeader->GetHeaders()->First() : header->GenEventHeader());
190       if (!mcHeader) { printf("mc header not found\n"); return; }
191       nProduced = mcEvent->GetNumberOfTracks();
192       nPrimGen = mcHeader->NProduced();
193       fZvtx = mcEvent->GetPrimaryVertex()->GetZ();
194     } else { // AOD
195       AliAODMCHeader* mcHeader = (AliAODMCHeader*) event->FindListObject(AliAODMCHeader::StdBranchName());
196       if (!mcHeader) { printf("AliAODMCHeader not found\n"); return; }
197       if (mcHeader->GetCocktailHeaders()) {
198         nProduced = mcTracks->GetEntriesFast();
199         AliGenEventHeader* header0 =  mcHeader->GetCocktailHeader(0);
200         if (!header0) { printf("first header expected but not found\n"); return; }
201         nPrimGen = header0->NProduced();
202       } else  nPrimGen = nProduced;
203       fZvtx = mcHeader->GetVtxZ();
204     }
205     fEventStatistics->Fill("after mc header check",1);
206
207     if (TMath::Abs(fZvtx)>fZVertexCut) return;
208     fEventStatistics->Fill("after MC vertex cut",1);
209
210     const AliVVertex* vertex = event->GetPrimaryVertex();
211     if (!vertex) return;
212     fEventStatistics->Fill("after check for vertex object",1);
213     TString name(vertex->GetName());
214     if (name.CompareTo("PrimaryVertex") && name.CompareTo("SPDVertex")) fZvtx=1000; // Reject TPC only vertex
215     fEventStatistics->Fill("after rejection of TPC only vertex",1);
216     if (vertex->GetNContributors()<fnTracksVertex)  fZvtx=-1000;
217     fEventStatistics->Fill("after check for vertex contributors",1);
218     
219     UInt_t* masks = new UInt_t[nProduced];
220     for (Int_t i=0;i<nProduced;i++) masks[i]=0;
221     
222     // Loop over reconstructed tracks to set masks
223     for (Int_t ipart=0;ipart<event->GetNumberOfTracks();ipart++){
224       AliVTrack* part = (AliVTrack*) event->GetTrack(ipart);
225       if (!part) continue;
226       Int_t label = TMath::Abs(part->GetLabel());
227       if (label>=nProduced) continue;
228       masks[label] |= GetFilterMap(part);
229     }
230     
231     // Loop over mc tracks to be stored
232     for (Int_t ipart=0;ipart<nProduced;ipart++){
233       AliVParticle* part = 0;
234       Bool_t isPrimary = 0;
235       if (mcTracks) { // AOD analysis
236         AliAODMCParticle* particle = (AliAODMCParticle*) mcTracks->At(ipart);
237         if (!particle) continue;
238         // skip injected signals
239         AliAODMCParticle* mother = particle;
240         while (!mother->IsPrimary()) mother = (AliAODMCParticle*) mcTracks->At(mother->GetMother());
241         if (mother->GetLabel()>=nPrimGen) continue;
242         part = particle;
243         // check for primary
244         isPrimary = particle->IsPhysicalPrimary();
245       } else { // ESD analysis
246         AliMCParticle* particle = (AliMCParticle*) mcEvent->GetTrack(ipart);
247         if (!particle) continue;
248         // skip injected signals
249         AliMCParticle* mother = particle;
250         while (mother->GetMother()>=0) mother = (AliMCParticle*) mcEvent->GetTrack(mother->GetMother());
251         if (mother->GetLabel()>=nPrimGen) continue;
252         part = particle;
253         // check for primary
254         isPrimary = mcEvent->IsPhysicalPrimary(ipart);
255       }
256       
257       // store only primaries and all reconstructed (non-zero mask)
258       Int_t mask = masks[ipart];
259       if (isPrimary) mask |= (1 << 30);
260       if (!mask) continue; 
261       AddTrack(part,mask);
262     }
263     delete[] masks;
264   }
265   fTree->Fill();
266   PostData(0,fListOfHistos);
267   PostData(1,fTree);
268 }
269 //-----------------------------------------------------------------------------
270
271
272 //-----------------------------------------------------------------------------
273 UInt_t AliAnalysisTaskCFTree::GetFilterMap(AliVParticle* particle){
274   UInt_t mask = 0;
275   if (particle->InheritsFrom("AliAODTrack")) {
276     AliAODTrack* part = (AliAODTrack*) particle;
277     mask = part->GetFilterMap();
278     Double_t nCrossedRaws      = part->GetTPCNCrossedRows();
279     Double_t nFindableClusters = part->GetTPCNclsF();
280     Double_t nSharedClusters   = part->GetTPCnclsS();
281     Double_t nClusters         = part->GetTPCncls();
282     Bool_t itsRefit            = part->GetStatus() & AliVTrack::kITSrefit;
283     if (nCrossedRaws/nFindableClusters > fFoundFractionCut) mask |= (1 << 26);
284     if (nCrossedRaws>fCrossedRowsCut)                       mask |= (1 << 27);
285     if (itsRefit)                                           mask |= (1 << 28);
286     if (nSharedClusters/nClusters<=fSharedClusterCut)       mask |= (1 << 29);
287     if (part->GetLabel()<0)                                 mask |= (1 << 31);
288   } else if (particle->InheritsFrom("AliESDtrack")){
289     AliESDtrack* part = (AliESDtrack*) particle;
290     if (!fTrackFilter) AliFatal("Track filter undefined");
291     mask |= fTrackFilter->IsSelected(part);
292   }
293   
294   return mask;
295 }
296 //-----------------------------------------------------------------------------
297
298
299 //-----------------------------------------------------------------------------
300 AliCFParticle* AliAnalysisTaskCFTree::AddTrack(AliVParticle* part, UInt_t mask, UInt_t flag){
301
302   // skip neutral mc particles
303   Char_t charge = part->Charge();
304   if (charge==0) return NULL;
305   
306   // set pt,eta,phi
307   Float_t pt=0,eta=0,phi=0;
308   if (flag==0 || flag==1){ // AOD, MC or Global ESD tracks
309     pt  = part->Pt();
310     eta = part->Eta();
311     phi = part->Phi();
312     if  (flag==1) mask &= (~fHybridConstrainedMask) & (~fTPConlyConstrainedMask);
313   } 
314   else if (flag==2) { // Hybrid constrained tracks (ESD)
315     AliESDtrack* track = (AliESDtrack*) part;
316     const AliExternalTrackParam* param = track->GetConstrainedParam();
317     pt  = param->Pt();
318     eta = param->Eta();
319     phi = param->Phi();
320     mask &= fHybridConstrainedMask;
321   } 
322   else if (flag==3) { // TPC only constrained tracks (ESD)
323     AliESDtrack* track = (AliESDtrack*) part;
324     AliESDtrack tpcTrack;
325     if (!track->FillTPCOnlyTrack(tpcTrack)) return NULL;
326     AliExternalTrackParam param;
327     const AliESDVertex* vtxSPD = ((AliESDEvent*) fInputHandler->GetEvent())->GetPrimaryVertexSPD();
328     if (!tpcTrack.RelateToVertexTPC(vtxSPD,fField,1.e30,&param)) return NULL;
329     pt  = param.Pt();
330     eta = param.Eta();
331     phi = param.Phi();
332     mask &= fTPConlyConstrainedMask;
333   }
334
335   // kinematic cuts
336   if (pt < fPtMin || TMath::Abs(eta) > fTrackEtaCut) return NULL;
337
338   
339   AliCFParticle* cfpart = new ((*fParticles)[fParticles->GetEntriesFast()]) AliCFParticle(pt,eta,phi,charge,mask,3);
340
341   AliVTrack* track = dynamic_cast<AliVTrack*> (part);
342   if (!track) return cfpart;
343   Float_t ncl  = track->GetTPCsignalN();
344   Float_t dedx = track->GetTPCsignalTunedOnData(); if (dedx<=0) dedx = track->GetTPCsignal();
345   Float_t beta = -99;
346   if (track->GetStatus()&AliESDtrack::kTOFpid){
347     Double_t tof[5];
348     track->GetIntegratedTimes(tof);
349     beta = tof[0]/track->GetTOFsignal();
350   }
351   cfpart->SetAt(ncl,0);
352   cfpart->SetAt(dedx,1);
353   cfpart->SetAt(beta,2);
354   return cfpart;
355 }