Adding the SetOwner
[u/mrichter/AliRoot.git] / PWG2 / EBYE / AliAnalysisTaskBF.cxx
1 #include "TChain.h"
2 #include "TList.h"
3 #include "TCanvas.h"
4 #include "TLorentzVector.h"
5 #include "TGraphErrors.h"
6 #include "TH1F.h"
7
8 #include "AliAnalysisTaskSE.h"
9 #include "AliAnalysisManager.h"
10
11 #include "AliESDVertex.h"
12 #include "AliESDEvent.h"
13 #include "AliESDInputHandler.h"
14 #include "AliAODEvent.h"
15 #include "AliAODTrack.h"
16 #include "AliAODInputHandler.h"
17 #include "AliMCEventHandler.h"
18 #include "AliMCEvent.h"
19 #include "AliStack.h"
20 #include "AliESDtrackCuts.h"
21
22 #include "AliBalance.h"
23
24 #include "AliAnalysisTaskBF.h"
25
26 // Analysis task for the BF code
27 // Authors: Panos Cristakoglou@cern.ch
28
29 ClassImp(AliAnalysisTaskBF)
30
31 //________________________________________________________________________
32 AliAnalysisTaskBF::AliAnalysisTaskBF(const char *name) 
33 : AliAnalysisTaskSE(name), 
34   fBalance(0),
35   fList(0),
36   fHistEventStats(0),
37   fHistVx(0),
38   fHistVy(0),
39   fHistVz(0),
40   fESDtrackCuts(0),
41   fUseOfflineTrigger(kFALSE),
42   fVxMax(0.3),
43   fVyMax(0.3),
44   fVzMax(10.) {
45   // Constructor
46
47   // Define input and output slots here
48   // Input slot #0 works with a TChain
49   DefineInput(0, TChain::Class());
50   // Output slot #0 writes into a TH1 container
51   DefineOutput(1, AliBalance::Class());
52   DefineOutput(2, TList::Class());
53 }
54
55 //________________________________________________________________________
56 void AliAnalysisTaskBF::UserCreateOutputObjects() {
57   // Create histograms
58   // Called once
59   if(!fBalance) {
60     fBalance = new AliBalance();
61     fBalance->SetAnalysisLevel("ESD");
62     fBalance->SetAnalysisType(1);
63     fBalance->SetNumberOfBins(18);
64     fBalance->SetInterval(-0.9,0.9);
65   }
66
67   //QA list
68   fList = new TList();
69   fList->SetName("listQA");
70   fList->SetOwner();
71
72   //Event stats.
73   TString gCutName[4] = {"Total","Offline trigger",
74                          "Vertex","Analyzed"};
75   fHistEventStats = new TH1F("fHistEventStats",
76                              "Event statistics;;N_{events}",
77                              4,0.5,4.5);
78   for(Int_t i = 1; i <= 4; i++)
79     fHistEventStats->GetXaxis()->SetBinLabel(i,gCutName[i-1].Data());
80   fList->Add(fHistEventStats);
81
82   //Vertex distributions
83   fHistVx = new TH1F("fHistVx","Primary vertex distribution - x coordinate;V_{x} (cm);Entries",100,-0.5,0.5);
84   fList->Add(fHistVx);
85   fHistVy = new TH1F("fHistVy","Primary vertex distribution - y coordinate;V_{y} (cm);Entries",100,-0.5,0.5);
86   fList->Add(fHistVy);
87   fHistVz = new TH1F("fHistVz","Primary vertex distribution - z coordinate;V_{z} (cm);Entries",100,-20.,20.);
88   fList->Add(fHistVz);
89
90   if(fESDtrackCuts) fList->Add(fESDtrackCuts);
91
92   // Post output data.
93   PostData(1, fBalance);
94   PostData(2, fList);
95 }
96
97 //________________________________________________________________________
98 void AliAnalysisTaskBF::UserExec(Option_t *) {
99   // Main loop
100   // Called for each event
101   TString gAnalysisLevel = fBalance->GetAnalysisLevel();
102
103   TObjArray *array = new TObjArray();
104   
105   //ESD analysis
106   if(gAnalysisLevel == "ESD") {
107     AliESDEvent* gESD = dynamic_cast<AliESDEvent*>(InputEvent()); // from TaskSE
108     if (!gESD) {
109       Printf("ERROR: gESD not available");
110       return;
111     }
112
113     fHistEventStats->Fill(1); //all events
114     Bool_t isSelected = kTRUE;
115     if(fUseOfflineTrigger)
116       isSelected = ((AliInputEventHandler*)(AliAnalysisManager::GetAnalysisManager()->GetInputEventHandler()))->IsEventSelected();
117     if(isSelected) {
118       fHistEventStats->Fill(2); //triggered events
119
120       const AliESDVertex *vertex = gESD->GetPrimaryVertex();
121       if(vertex) {
122         if(vertex->GetNContributors() > 0) {
123           if(vertex->GetZRes() != 0) {
124             fHistEventStats->Fill(3); //events with a proper vertex
125             if(TMath::Abs(vertex->GetXv()) < fVxMax) {
126               if(TMath::Abs(vertex->GetYv()) < fVyMax) {
127                 if(TMath::Abs(vertex->GetZv()) < fVzMax) {
128                   fHistEventStats->Fill(4); //analayzed events
129                   fHistVx->Fill(vertex->GetXv());
130                   fHistVy->Fill(vertex->GetYv());
131                   fHistVz->Fill(vertex->GetZv());
132
133                   //Printf("There are %d tracks in this event", gESD->GetNumberOfTracks());
134                   for (Int_t iTracks = 0; iTracks < gESD->GetNumberOfTracks(); iTracks++) {
135                     AliESDtrack* track = gESD->GetTrack(iTracks);
136                     if (!track) {
137                       Printf("ERROR: Could not receive track %d", iTracks);
138                       continue;
139                     }
140                     
141                     //ESD track cuts
142                     if(fESDtrackCuts) 
143                       if(!fESDtrackCuts->AcceptTrack(track)) continue;
144                     array->Add(track);
145                   } //track loop
146                 }//Vz cut
147               }//Vy cut
148             }//Vx cut
149           }//proper vertex resolution
150         }//proper number of contributors
151       }//vertex object valid
152     }//triggered event 
153   }//ESD analysis
154   //AOD analysis
155   else if(gAnalysisLevel == "AOD") {
156     AliAODEvent* gAOD = dynamic_cast<AliAODEvent*>(InputEvent()); // from TaskSE
157     if(!gAOD) {
158       Printf("ERROR: gAOD not available");
159       return;
160     }
161
162     Printf("There are %d tracks in this event", gAOD->GetNumberOfTracks());
163     for (Int_t iTracks = 0; iTracks < gAOD->GetNumberOfTracks(); iTracks++) {
164       AliAODTrack* track = gAOD->GetTrack(iTracks);
165       if (!track) {
166         Printf("ERROR: Could not receive track %d", iTracks);
167         continue;
168       }
169       array->Add(track);
170     } //track loop
171   }//AOD analysis
172   //MC analysis
173   else if(gAnalysisLevel == "MC") {
174
175     AliMCEvent*  mcEvent = MCEvent(); 
176     if (!mcEvent) {
177       Printf("ERROR: mcEvent not available");
178       return;
179     }
180     
181     Printf("There are %d tracks in this event", mcEvent->GetNumberOfPrimaries());
182     for (Int_t iTracks = 0; iTracks < mcEvent->GetNumberOfPrimaries(); iTracks++) {
183       AliMCParticle* track = dynamic_cast<AliMCParticle *>(mcEvent->GetTrack(iTracks));
184       if (!track) {
185         Printf("ERROR: Could not receive particle %d", iTracks);
186         continue;
187       }
188       array->Add(track);
189     } //track loop
190   }//MC analysis
191
192   fBalance->CalculateBalance(array);
193   
194   delete array;
195   
196 }      
197
198 //________________________________________________________________________
199 void AliAnalysisTaskBF::Terminate(Option_t *) {
200   // Draw result to the screen
201   // Called once at the end of the query
202   fBalance = dynamic_cast<AliBalance*> (GetOutputData(1));
203   if (!fBalance) {
204     Printf("ERROR: fBalance not available");
205     return;
206   }
207   
208   TGraphErrors *gr = fBalance->DrawBalance();
209   gr->SetMarkerStyle(20);
210   gr->Draw("AP");
211
212   fBalance->PrintResults();
213 }