New classes for event plane calculation (Johanna)
[u/mrichter/AliRoot.git] / ANALYSIS / AliEPSelectionTask.cxx
1 /**************************************************************************
2  * Copyright(c) 1998-2008, 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 //*****************************************************
17 //   Class AliEPSelectionTask
18 //   Class to determine event plane            
19 //   author: Alberica Toia, Johanna Gramling
20 //*****************************************************
21
22 #include "AliEPSelectionTask.h"
23
24 #include <TTree.h>
25 #include <TList.h>
26 #include <TH1F.h>
27 #include <TH2F.h>
28 #include <TProfile.h>
29 #include <TFile.h>
30 #include <TObjString.h>
31 #include <TString.h>
32 #include <TCanvas.h>
33 #include <TROOT.h>
34 #include <TDirectory.h>
35 #include <TSystem.h>
36 #include <iostream>
37 #include <TRandom2.h>
38 #include <TArrayF.h>
39
40 #include "AliAnalysisManager.h"
41 #include "AliVEvent.h"
42 #include "AliESD.h"
43 #include "AliESDEvent.h"
44 #include "AliMCEvent.h"
45 #include "AliESDtrack.h"
46 #include "AliESDtrackCuts.h"
47 #include "AliESDHeader.h"
48 #include "AliESDInputHandler.h"
49 #include "AliAODHandler.h"
50 #include "AliAODEvent.h"
51 #include "AliHeader.h"
52 #include "AliGenHijingEventHeader.h"
53 #include "AliAnalysisTaskSE.h"
54 #include "AliPhysicsSelectionTask.h"
55 #include "AliPhysicsSelection.h"
56 #include "AliBackgroundSelection.h"
57 #include "AliESDUtils.h"
58
59 #include "AliEventplane.h"
60
61 ClassImp(AliEPSelectionTask)
62
63 //________________________________________________________________________
64 AliEPSelectionTask::AliEPSelectionTask():
65 AliAnalysisTaskSE(),
66   fDebug(0),
67   fAnalysisInput("ESD"),
68   fStatus("TPC"),
69   fUseMCRP(kFALSE),
70   fUsePhiWeight(kFALSE),
71   fUsePtWeight(kFALSE),
72   fSaveTrackContribution(kFALSE),
73   fESDtrackCuts(0),
74   ftracklist(0),
75   fPhiDist(0),
76   fQVector(0),
77   fQContributionX(0),
78   fQContributionY(0),
79   fEventplaneQ(0),
80   fQsub1(0),
81   fQsub2(0),
82   fQsubRes(0),  
83   fOutputList(0),
84   fHOutEventplaneQ(0),
85   fHOutPhi(0),
86   fHOutPhiCorr(0),
87   fHOutsub1sub2(0),
88   fHOutNTEPRes(0),
89   fHOutPTPsi(0),
90   fHOutDiff(0),
91   fHOutleadPTPsi(0)
92 {   
93   // Default constructor
94   AliInfo("Event Plane Selection enabled.");
95 }   
96
97 //________________________________________________________________________
98 AliEPSelectionTask::AliEPSelectionTask(const char *name):
99   AliAnalysisTaskSE(name),
100   fDebug(0),
101   fAnalysisInput("ESD"),
102   fStatus("TPC"),
103   fUseMCRP(kFALSE),
104   fUsePhiWeight(kFALSE),
105   fUsePtWeight(kFALSE),  
106   fSaveTrackContribution(kFALSE),
107   fESDtrackCuts(0),
108   ftracklist(0),
109   fPhiDist(0),
110   fQVector(0),
111   fQContributionX(0),
112   fQContributionY(0),
113   fEventplaneQ(0),
114   fQsub1(0),
115   fQsub2(0),
116   fQsubRes(0),
117   fOutputList(0),
118   fHOutEventplaneQ(0),
119   fHOutPhi(0),
120   fHOutPhiCorr(0),
121   fHOutsub1sub2(0),
122   fHOutNTEPRes(0),
123   fHOutPTPsi(0),
124   fHOutDiff(0),
125   fHOutleadPTPsi(0)
126 {
127   // Default constructor
128   AliInfo("Event Plane Selection enabled.");
129   DefineOutput(1, TList::Class());
130 }
131  
132 //________________________________________________________________________
133 AliEPSelectionTask::~AliEPSelectionTask()
134 {
135   // Destructor  
136   if (fOutputList && !AliAnalysisManager::GetAnalysisManager()->IsProofMode()){
137       delete fOutputList;
138       fOutputList = 0;
139   }
140   if (fESDtrackCuts){
141       delete fESDtrackCuts;
142       fESDtrackCuts = 0;
143   }
144   if (fQVector){
145       delete fQVector;
146       fQVector = 0;
147   }
148   if (fQsub1){
149       delete fQsub1;
150       fQsub1 = 0;
151   }
152     if (fQsub2){
153       delete fQsub2;
154       fQsub2 = 0;
155   }
156 }  
157
158 //________________________________________________________________________
159 void AliEPSelectionTask::UserCreateOutputObjects()
160 {  
161   // Create the output containers
162   if (fDebug>1) printf("AliEPSelectionTask::UserCreateOutputObjects() \n");
163   AliLog::SetClassDebugLevel("AliEPSelectionTask", AliLog::kInfo);
164
165   fOutputList = new TList();
166   fOutputList->SetOwner();
167   fHOutEventplaneQ = new TH1F("fHOutEventplaneQ","fHOutEventplaneQ; Event Plane Q",100,0,TMath::Pi());
168   fHOutPhi = new TH1F("fHOutPhi","fHOutPhi; Phi Distribution",100,0,TMath::TwoPi());
169   fHOutPhiCorr = new TH1F("fHOutPhiCorr","fHOutPhiCorr; Corrected Phi Distribution",100,0,TMath::TwoPi());
170   fHOutsub1sub2 = new TH2F("fHOutsub1sub2","fHOutsub1sub2; EP1; EP2",100,0,TMath::Pi(),100,0,TMath::Pi());
171   fHOutNTEPRes = new TH2F("fHOutNTEPRes","fHOutNTEPRes; Number of Tracks; Event Plane Resolution",100,0,5000,100,-TMath::Pi(),TMath::Pi());
172   fHOutPTPsi = new TH2F("fHOutPTPsi","fHOutPTPsi; PT; Phi-EP",100,0,20,100,0,TMath::Pi());
173   fHOutDiff = new TH2F("fHOutDiff","fHOutDiff; EP; MCEP",100,0,TMath::Pi(),100,0,TMath::Pi());
174   fHOutleadPTPsi = new TH2F("fHOutleadPTPsi","fHOutleadPTPsi; leadPT; EP",100,0,TMath::Pi(),100,0,TMath::Pi());
175
176   fOutputList->Add(fHOutEventplaneQ);
177   fOutputList->Add(fHOutPhi);
178   fOutputList->Add(fHOutPhiCorr);
179   fOutputList->Add(fHOutsub1sub2);
180   fOutputList->Add(fHOutNTEPRes);
181   fOutputList->Add(fHOutPTPsi);
182   fOutputList->Add(fHOutleadPTPsi);
183   fOutputList->Add(fHOutDiff);
184   
185   PostData(1, fOutputList); 
186 }
187
188 //________________________________________________________________________
189 void AliEPSelectionTask::UserExec(Option_t */*option*/)
190
191   // Execute analysis for current event:
192   if (fDebug>1) printf(" **** AliEPSelectionTask::UserExec() \n");
193  
194   AliEventplane* esdEP = 0;
195   TVector2 QQ1;
196   TVector2 QQ2;
197   Double_t fRP = 0.; // the monte carlo reaction plane angle
198     
199   if (fAnalysisInput.CompareTo("ESD")==0){
200     
201     AliVEvent* event = InputEvent();
202     AliESDEvent* esd = dynamic_cast<AliESDEvent*>(event);
203     
204     if (fUseMCRP) {
205       AliMCEvent* mcEvent  = MCEvent();      
206       if (mcEvent && mcEvent->GenEventHeader()) {
207         AliGenHijingEventHeader* headerH = dynamic_cast<AliGenHijingEventHeader*>(mcEvent->GenEventHeader());
208         fRP = headerH->ReactionPlaneAngle();
209       }
210     }
211     if (esd){
212       esdEP = esd->GetEventplane();
213       if (fSaveTrackContribution) {
214         esdEP->GetQContributionXArray()->Set(esd->GetNumberOfTracks());
215         esdEP->GetQContributionYArray()->Set(esd->GetNumberOfTracks());
216       }
217       
218       if (fStatus.CompareTo("GLOBAL")==0) ftracklist = fESDtrackCuts->GetAcceptedTracks(esd,kFALSE);
219       if (fStatus.CompareTo("TPC")==0) ftracklist = fESDtrackCuts->GetAcceptedTracks(esd,kTRUE);
220       const int NT = ftracklist->GetEntries();
221       
222       if (NT>4){
223         fQVector = new TVector2(GetQ(esdEP,ftracklist));
224         fEventplaneQ = fQVector->Phi()/2; 
225         GetQsub(QQ1, QQ2, ftracklist);
226         fQsub1 = new TVector2(QQ1);
227         fQsub2 = new TVector2(QQ2);
228         fQsubRes = (fQsub1->Phi()/2 - fQsub2->Phi()/2);
229   
230         esdEP->SetQVector(fQVector);
231         esdEP->SetEventplaneQ(fEventplaneQ);
232         esdEP->SetQsub(fQsub1,fQsub2);
233         esdEP->SetQsubRes(fQsubRes);
234         
235         fHOutEventplaneQ->Fill(fEventplaneQ);
236         fHOutsub1sub2->Fill(fQsub1->Phi()/2,fQsub2->Phi()/2);
237         fHOutNTEPRes->Fill(NT,fQsubRes);
238
239         if (fUseMCRP) fHOutDiff->Fill(fEventplaneQ, fRP);
240         
241         for (int iter = 0; iter<NT;iter++){
242           AliESDtrack* track = dynamic_cast<AliESDtrack*> (ftracklist->At(iter));
243           float delta = track->Phi()-fEventplaneQ;
244           while (delta < 0) delta += TMath::Pi();
245           while (delta > TMath::Pi()) delta -= TMath::Pi();
246           fHOutPTPsi->Fill(track->Pt(),delta);
247           fHOutPhi->Fill(track->Phi());
248           fHOutPhiCorr->Fill(track->Phi(),GetPhiWeight(track));
249         }
250         
251         AliESDtrack* trmax = esd->GetTrack(0);
252         for (int iter = 1; iter<NT;iter++){
253           AliESDtrack* track = dynamic_cast<AliESDtrack*> (ftracklist->At(iter));
254           if (track->Pt() > trmax->Pt()) trmax = track;
255         }
256         fHOutleadPTPsi->Fill(trmax->Phi(),fEventplaneQ);      
257       }
258       delete ftracklist;
259     }
260   }
261   
262   else if (fAnalysisInput.CompareTo("AOD")==0){
263     //AliAODEvent *aod =  dynamic_cast<AliAODEvent*> (InputEvent());
264     // to be implemented
265     printf("  AOD analysis not yet implemented!!!\n\n");
266     return;
267   }  
268   else {
269     printf(" Analysis Input not known!\n\n ");
270     return;
271   }
272   PostData(1, fOutputList); 
273 }
274
275 //________________________________________________________________________
276 void AliEPSelectionTask::Terminate(Option_t */*option*/)
277 {
278   // Terminate analysis
279 }
280
281 //__________________________________________________________________________
282 TVector2 AliEPSelectionTask::GetQ(AliEventplane* EP, TObjArray* tracklist)
283 {
284   TVector2 mQ;
285   float mQx=0, mQy=0;
286   AliESDtrack* track;
287   Double_t weight;
288   
289   int NT = tracklist->GetEntries();
290
291   for (int i=0; i<NT; i++){
292     weight = 1;
293     track = dynamic_cast<AliESDtrack*> (tracklist->At(i));
294     weight = GetWeight(track);
295     if (fSaveTrackContribution){
296       EP->GetQContributionXArray()->AddAt(weight*cos(2*track->Phi()),track->GetID());
297       EP->GetQContributionYArray()->AddAt(weight*sin(2*track->Phi()),track->GetID());
298     }
299     mQx += (weight*cos(2*track->Phi()));
300     mQy += (weight*sin(2*track->Phi()));
301   }
302   mQ.Set(mQx,mQy);
303   return mQ;
304 }
305   
306   //________________________________________________________________________
307 void AliEPSelectionTask::GetQsub(TVector2 &Q1, TVector2 &Q2, TObjArray* tracklist)
308 {
309   TVector2 mQ[2];
310   float mQx1=0, mQy1=0, mQx2=0, mQy2=0;
311   Double_t weight;
312
313   AliESDtrack* track;
314   TRandom2 RN = 0;
315   
316   int NT = tracklist->GetEntries();
317   int trackcounter1=0, trackcounter2=0;
318   
319   for (Int_t i = 0; i < NT; i++) {
320     weight = 1;
321     track = dynamic_cast<AliESDtrack*> (tracklist->At(i));
322     weight = GetWeight(track);
323     
324     // This loop splits the track set into 2 random subsets
325     if( trackcounter1 < int(NT/2.) && trackcounter2 < int(NT/2.)){
326       float random = RN.Rndm();
327       if(random < .5){
328         mQx1 += (weight*cos(2*track->Phi()));
329         mQy1 += (weight*sin(2*track->Phi()));
330         trackcounter1++;
331       }
332       else {
333         mQx2 += (weight*cos(2*track->Phi()));
334         mQy2 += (weight*sin(2*track->Phi()));
335         trackcounter2++;
336       }
337     }
338     else if( trackcounter1 >= int(NT/2.)){
339       mQx2 += (weight*cos(2*track->Phi()));
340       mQy2 += (weight*sin(2*track->Phi()));
341       trackcounter2++;
342     }
343     else {
344       mQx1 += (weight*cos(2*track->Phi()));
345       mQy1 += (weight*sin(2*track->Phi()));
346       trackcounter1++;
347     }
348   }
349   mQ[0].Set(mQx1,mQy1);
350   mQ[1].Set(mQx2,mQy2);
351   Q1 = mQ[0];
352   Q2 = mQ[1];
353 }
354
355 //________________________________________________________________________
356 void AliEPSelectionTask::SetESDtrackCuts(TString status){
357   
358   fStatus = status;
359   if (fStatus.CompareTo("GLOBAL")==0) fESDtrackCuts = AliESDtrackCuts::GetStandardITSTPCTrackCuts2010(kTRUE);
360   if (fStatus.CompareTo("TPC")==0)    fESDtrackCuts = AliESDtrackCuts::GetStandardTPCOnlyTrackCuts();
361   fESDtrackCuts->SetPtRange(0.15,20);
362   fESDtrackCuts->SetEtaRange(-0.8,0.8);
363 }
364
365 //__________________________________________________________________________
366 void AliEPSelectionTask::SetPhiDistribution(char* infilename, char* listname){
367   TFile f(infilename);
368   TObject* list = f.Get(listname);
369   fPhiDist = (TH1F*)list->FindObject("fHOutPhi");
370   if (!fPhiDist) cout << "Phi Distribution not found!!!" << endl;
371   
372   Bool_t emptybins;
373
374   int iter = 0;  
375   while (iter<3){
376       emptybins = kFALSE;
377    
378       for (int i=1; i<fPhiDist->GetNbinsX(); i++){
379         if (!((fPhiDist->GetBinContent(i))>0)) {
380           emptybins = kTRUE;
381         }
382       }  
383       if (emptybins) {
384         cout << "empty bins - rebinning!" << endl;
385         fPhiDist->Rebin();
386         iter++;
387       }      
388       else iter = 3;
389   }
390   
391   if (emptybins) {
392     AliError("After Maximum of rebinning still empty Phi-bins!!!");
393   }
394   f.Close();
395 }
396
397 //________________________________________________________________________
398 Double_t AliEPSelectionTask::GetWeight(AliESDtrack* track)
399 {
400   Double_t ptweight=1;
401
402   if (fUsePtWeight) {      
403     if (track->Pt()<2) ptweight=track->Pt();
404     else ptweight=2;
405   }
406   return ptweight*GetPhiWeight(track);
407 }
408
409 //________________________________________________________________________
410 Double_t AliEPSelectionTask::GetPhiWeight(AliESDtrack* track)
411 {
412   Double_t phiweight=1;
413   
414   if (fUsePhiWeight) {
415     Double_t nParticles = fPhiDist->Integral();
416     Double_t nPhibins = fPhiDist->GetNbinsX();
417   
418     Double_t Phi = track->Phi();
419     
420     while (Phi<0) Phi += TMath::TwoPi();
421     while (Phi>TMath::TwoPi()) Phi -= TMath::TwoPi();
422       
423     Double_t PhiDistValue = fPhiDist->GetBinContent(1+TMath::Floor((track->Phi())*nPhibins/TMath::TwoPi()));
424     
425     if (PhiDistValue > 0) phiweight = nParticles/nPhibins/PhiDistValue;
426   }
427   return phiweight;
428 }