Adding pt and eta control plots, increasing ClassDef version
[u/mrichter/AliRoot.git] / PWG0 / esdTrackCuts / AliCutTask.cxx
1 #include "TChain.h"
2 #include "TTree.h"
3 #include "TH1F.h"
4 #include "TCanvas.h"
5 #include "TList.h"
6 #include "TFile.h"
7 #include "TParticle.h"
8 #include "TParticlePDG.h"
9
10 #include "esdTrackCuts/AliESDtrackCuts.h"
11 #include "AliPWG0Helper.h"
12
13 #include "AliAnalysisTask.h"
14 #include "AliAnalysisManager.h"
15 #include "AliESDEvent.h"
16 #include "AliESDInputHandler.h"
17 #include "AliESDVertex.h"
18
19 #include "AliMCEventHandler.h"
20 #include "AliMCEvent.h"
21 #include "AliStack.h"
22 #include "AliLog.h"
23
24 #include "AliCutTask.h"
25
26 // simple task that runs the esd track cuts to evaluate the basic plots created during the cuts
27
28 ClassImp(AliCutTask)
29
30 //________________________________________________________________________
31 AliCutTask::AliCutTask(const char *name) 
32   : AliAnalysisTask(name, ""), fESD(0), fTrackCuts(0), fTrackCutsPrimaries(0),
33     fTrackCutsSecondaries(0), fVertex(0), fTriggerStats(0), fOutput(0)
34 {
35   // Constructor
36
37   // Define input and output slots here
38   DefineInput(0, TChain::Class());
39   DefineOutput(0, TList::Class());
40 }
41
42 //________________________________________________________________________
43 void AliCutTask::ConnectInputData(Option_t *) 
44 {
45   // Connect ESD or AOD here
46   // Called once
47
48   TTree* tree = dynamic_cast<TTree*> (GetInputData(0));
49   if (!tree) {
50     Printf("ERROR: Could not read chain from input slot 0");
51   } else {
52     // Disable all branches and enable only the needed ones
53     //tree->SetBranchStatus("*", kFALSE);
54     //tree->SetBranchStatus("*Calo*", kFALSE);
55
56     tree->SetBranchStatus("fTracks.*", kTRUE);
57     tree->SetBranchStatus("Tracks.*", kTRUE);
58
59     tree->SetBranchStatus("fTriggerMask", kTRUE);
60     tree->SetBranchStatus("AliESDHeader", kTRUE);
61
62     tree->SetBranchStatus("fSPDVertex*", kTRUE);
63
64     // unclear how to enable vertex branch with new ESD format
65     tree->SetBranchStatus("SPDVertex.*", kTRUE);
66     tree->SetBranchStatus("*fPosition*", kTRUE);
67     tree->SetBranchStatus("*fPosition[3]", kTRUE);
68
69     AliESDInputHandler *esdH = dynamic_cast<AliESDInputHandler*> (AliAnalysisManager::GetAnalysisManager()->GetInputEventHandler());
70
71     if (!esdH) {
72       Printf("ERROR: Could not get ESDInputHandler");
73     } else
74       fESD = esdH->GetEvent();
75   }
76 }
77
78 //________________________________________________________________________
79 void AliCutTask::CreateOutputObjects()
80 {
81   // Create histograms
82   // Called once
83
84   fOutput = new TList;
85   fOutput->SetOwner();
86
87   fOutput->Add(fTrackCuts);
88   if (fTrackCutsPrimaries)
89     fOutput->Add(fTrackCutsPrimaries);
90   if (fTrackCutsSecondaries)
91     fOutput->Add(fTrackCutsSecondaries);
92
93   fVertex = new TH1F("fVertex", "fVertex;z vtx (cm);Count", 201, -20, 20);
94   fOutput->Add(fVertex);
95
96   fTriggerStats = new TH1F("fTriggerStats", "fTriggerStats;trigger;Count", 5, -0.5, 4.5);
97   fTriggerStats->GetXaxis()->SetBinLabel(1, "!MB1 & !MB2");
98   fTriggerStats->GetXaxis()->SetBinLabel(2, "MB1");
99   fTriggerStats->GetXaxis()->SetBinLabel(3, "MB2");
100   fTriggerStats->GetXaxis()->SetBinLabel(4, "ITS_SPD_GFO_L0");
101   fTriggerStats->GetXaxis()->SetBinLabel(5, "VZERO_OR_LEFT | VZERO_OR_RIGHT");
102   fOutput->Add(fTriggerStats);
103
104   // disable info messages of AliMCEvent (per event)
105   AliLog::SetClassDebugLevel("AliMCEvent", AliLog::kWarning - AliLog::kDebug + 1);
106 }
107
108 //________________________________________________________________________
109 void AliCutTask::Exec(Option_t *) 
110 {
111   // Main loop
112   // Called for each event
113
114   if (!fESD) {
115     Printf("ERROR: fESD not available");
116     return;
117   }
118
119   // Post output data.
120   PostData(0, fOutput);
121
122   //fESD->GetVertex()->Print();
123
124   if (!AliPWG0Helper::IsEventTriggered(fESD->GetTriggerMask(), AliPWG0Helper::kMB1) && !AliPWG0Helper::IsEventTriggered(fESD->GetTriggerMask(), AliPWG0Helper::kMB2))
125     fTriggerStats->Fill(0);
126
127   if (AliPWG0Helper::IsEventTriggered(fESD->GetTriggerMask(), AliPWG0Helper::kMB1))
128     fTriggerStats->Fill(1);
129
130   if (AliPWG0Helper::IsEventTriggered(fESD->GetTriggerMask(), AliPWG0Helper::kMB2))
131     fTriggerStats->Fill(2);
132
133   if (fESD->GetTriggerMask() & (0x1 << 14))
134     fTriggerStats->Fill(3);
135
136   if (fESD->GetTriggerMask() & 1 || fESD->GetTriggerMask() & 2)
137     fTriggerStats->Fill(4);
138
139   //if (!AliPWG0Helper::IsEventTriggered(fESD->GetTriggerMask()), AliPWG0Helper::kMB1)
140   if (fESD->GetTriggerMask() & (0x1 << 14) == 0)
141     return;
142
143   if (!AliPWG0Helper::IsVertexReconstructed(fESD->GetVertex()))
144     return;
145
146   //Printf("There are %d tracks in this event", fESD->GetNumberOfTracks());
147   fTrackCuts->CountAcceptedTracks(fESD);
148   
149   if (fTrackCutsPrimaries && fTrackCutsSecondaries)
150   {
151     AliMCEventHandler* eventHandler = dynamic_cast<AliMCEventHandler*> (AliAnalysisManager::GetAnalysisManager()->GetMCtruthEventHandler());
152     if (!eventHandler) {
153       Printf("ERROR: Could not retrieve MC event handler");
154       return;
155     }
156
157     AliMCEvent* mcEvent = eventHandler->MCEvent();
158     if (!mcEvent) {
159       Printf("ERROR: Could not retrieve MC event");
160       return;
161     }
162
163     AliStack* stack = mcEvent->Stack();
164     if (!stack)
165     {
166       Printf("ERROR: Stack not available");
167       return;
168     }
169     
170     for (Int_t trackId = 0; trackId < fESD->GetNumberOfTracks(); trackId++)
171     {
172       AliESDtrack* track = fESD->GetTrack(trackId);
173       if (!track)
174       {
175         Printf("UNEXPECTED: Could not get track with id %d", trackId);
176         continue;
177       }
178       
179       if (track->GetLabel() == 0)
180       {
181         Printf("Track %d has no label. Skipping", trackId);
182         continue;
183       }
184       
185       Int_t label = TMath::Abs(track->GetLabel());
186       if (stack->IsPhysicalPrimary(label) == kTRUE)
187       {
188         fTrackCutsPrimaries->AcceptTrack(track);
189       }
190       else
191       {
192         if (!stack->Particle(label)) {
193           Printf("ERROR: Could not get particle with label %d", label);
194           continue;
195         }
196         
197         // Deuteron is ok, but missing in Pdg particles in root
198         if (stack->Particle(label)->GetPdgCode() != 10010020) {
199           if (!stack->Particle(label)->GetPDG()) {
200             Printf("ERROR: Could not get PDG particle of particle with label %d (pdg code is %d)", label, stack->Particle(label)->GetPdgCode());
201             stack->Particle(label)->Print();
202             continue;
203           }
204           
205           if (stack->Particle(label)->GetPDG()->Charge() == 0) {
206             Printf("Particle %d has 0 charge. Skipping.", label);
207             continue;
208           }
209         }
210           
211         fTrackCutsSecondaries->AcceptTrack(track);
212       }
213     }
214   }
215
216   // get the ESD vertex
217   fVertex->Fill(fESD->GetVertex()->GetZv());
218 }
219
220 //________________________________________________________________________
221 void AliCutTask::Terminate(Option_t *) 
222 {
223   // Draw result to the screen
224   // Called once at the end of the query
225
226   fOutput = dynamic_cast<TList*> (GetOutputData(0));
227   if (!fOutput) {
228     Printf("ERROR: fOutput not available");
229     return;
230   }
231
232   fTrackCuts = dynamic_cast<AliESDtrackCuts*> (fOutput->FindObject("AliESDtrackCuts"));
233   if (!fTrackCuts) {
234     Printf("ERROR: fTrackCuts not available");
235     return;
236   }
237
238   fVertex = dynamic_cast<TH1F*> (fOutput->FindObject("fVertex"));
239   if (!fVertex) {
240     Printf("ERROR: fVertex not available");
241     return;
242   }
243
244   fTriggerStats = dynamic_cast<TH1F*> (fOutput->FindObject("fTriggerStats"));
245   if (!fTriggerStats) {
246     Printf("ERROR: fTriggerStats not available");
247     return;
248   }
249
250   fTrackCutsPrimaries = dynamic_cast<AliESDtrackCuts*> (fOutput->FindObject("fTrackCutsPrimaries"));
251   fTrackCutsSecondaries = dynamic_cast<AliESDtrackCuts*> (fOutput->FindObject("fTrackCutsSecondaries"));
252   
253   TFile* file = TFile::Open("trackCuts.root", "RECREATE");
254
255   fTrackCuts->SaveHistograms();
256   fVertex->Write();
257   fTriggerStats->Write();
258   
259   if (fTrackCutsPrimaries)
260     fTrackCutsPrimaries->SaveHistograms();
261
262   if (fTrackCutsSecondaries)
263     fTrackCutsSecondaries->SaveHistograms();
264   
265   file->Close();
266   
267   Printf("Writting results to trackCuts.root.");
268   
269   return;
270
271         fTrackCuts->DrawHistograms();
272
273   new TCanvas;
274   fVertex->Draw();
275
276   new TCanvas;
277   fTriggerStats->Draw();
278 }
279
280 void AliCutTask::EnableSecondaryStudy()
281
282   // 
283   // clones the fTrackCuts
284   //
285   
286   if (!fTrackCuts) {
287     Printf("ERROR: fTrackCuts 0. Do not call this function before SetTrackCuts");
288     return;
289   }
290   
291   fTrackCutsPrimaries = dynamic_cast<AliESDtrackCuts*> (fTrackCuts->Clone("fTrackCutsPrimaries"));
292   fTrackCutsSecondaries = dynamic_cast<AliESDtrackCuts*> (fTrackCuts->Clone("fTrackCutsSecondaries"));
293 }