added correction for events with vertex but 0 tracks
[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), 
33   fAnalysisMode(AliPWG0Helper::kTPCITS),
34   fTrackCutsPrimaries(0),
35   fTrackCutsSecondaries(0), fVertex(0), fTriggerStats(0), fOutput(0),
36   fPrimStats(0)
37 {
38   // Constructor
39
40   // Define input and output slots here
41   DefineInput(0, TChain::Class());
42   DefineOutput(0, TList::Class());
43 }
44
45 //________________________________________________________________________
46 void AliCutTask::ConnectInputData(Option_t *) 
47 {
48   // Connect ESD or AOD here
49   // Called once
50
51   TTree* tree = dynamic_cast<TTree*> (GetInputData(0));
52   if (!tree) {
53     Printf("ERROR: Could not read chain from input slot 0");
54   } else {
55     // Disable all branches and enable only the needed ones
56     //tree->SetBranchStatus("*", kFALSE);
57
58     // old esd
59     tree->SetBranchStatus("fTriggerMask", kTRUE);
60     tree->SetBranchStatus("fTracks.*", kTRUE);
61     tree->SetBranchStatus("fSPDVertex*", kTRUE);
62
63     // new esd
64     tree->SetBranchStatus("TriggerMask", kTRUE);
65     tree->SetBranchStatus("AliESDHeader", kTRUE);
66     //AliPWG0Helper::SetBranchStatusRecursive(tree, "Tracks", kTRUE);
67     //AliPWG0Helper::SetBranchStatusRecursive(tree, "SPDVertex", 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   fPrimStats = new TH1F("fPrimStats", "fPrimStats", 8, 0.5, 8.5);
105   fPrimStats->GetXaxis()->SetBinLabel(1, "Primary accepted once");
106   fPrimStats->GetXaxis()->SetBinLabel(2, "Primary accepted more than once");
107   fPrimStats->GetXaxis()->SetBinLabel(3, "Primary accepted more than once (count)");
108   fPrimStats->GetXaxis()->SetBinLabel(4, "Primary track rejected");
109   fPrimStats->GetXaxis()->SetBinLabel(5, "Primary track rejected (count)");
110   fPrimStats->GetXaxis()->SetBinLabel(6, "Primary track rejected, but other track accepted");
111   fPrimStats->GetXaxis()->SetBinLabel(7, "Primary track rejected, but other track accepted (count)");
112   fOutput->Add(fPrimStats);
113
114   // disable info messages of AliMCEvent (per event)
115   AliLog::SetClassDebugLevel("AliMCEvent", AliLog::kWarning - AliLog::kDebug + 1);
116 }
117
118 //________________________________________________________________________
119 void AliCutTask::Exec(Option_t *) 
120 {
121   // Main loop
122   // Called for each event
123
124   if (!fESD) {
125     Printf("ERROR: fESD not available");
126     return;
127   }
128
129   // Post output data.
130   PostData(0, fOutput);
131
132   //fESD->GetVertex()->Print();
133
134   if (!AliPWG0Helper::IsEventTriggered(fESD->GetTriggerMask(), AliPWG0Helper::kMB1) && !AliPWG0Helper::IsEventTriggered(fESD->GetTriggerMask(), AliPWG0Helper::kMB2))
135     fTriggerStats->Fill(0);
136
137   if (AliPWG0Helper::IsEventTriggered(fESD->GetTriggerMask(), AliPWG0Helper::kMB1))
138     fTriggerStats->Fill(1);
139
140   if (AliPWG0Helper::IsEventTriggered(fESD->GetTriggerMask(), AliPWG0Helper::kMB2))
141     fTriggerStats->Fill(2);
142
143   if (fESD->GetTriggerMask() & (0x1 << 14))
144     fTriggerStats->Fill(3);
145
146   if (fESD->GetTriggerMask() & 1 || fESD->GetTriggerMask() & 2)
147     fTriggerStats->Fill(4);
148
149   //if (fESD->GetTriggerMask() & (0x1 << 14) == 0)
150   if (!AliPWG0Helper::IsEventTriggered(fESD->GetTriggerMask(), AliPWG0Helper::kMB2))
151     return;
152
153   // get the ESD vertex
154   const AliESDVertex* vtxESD = AliPWG0Helper::GetVertex(fESD, fAnalysisMode);
155   
156   if (!vtxESD)
157   {
158     Printf("No vertex. Skipping event");
159     return;
160   }
161
162   Printf("There are %d tracks in this event", fESD->GetNumberOfTracks());
163   Int_t acceptedTracks = fTrackCuts->CountAcceptedTracks(fESD);
164   Printf("Accepted %d tracks", acceptedTracks);
165   
166   if (fTrackCutsPrimaries && fTrackCutsSecondaries)
167   {
168     AliMCEventHandler* eventHandler = dynamic_cast<AliMCEventHandler*> (AliAnalysisManager::GetAnalysisManager()->GetMCtruthEventHandler());
169     if (!eventHandler) {
170       Printf("ERROR: Could not retrieve MC event handler");
171       return;
172     }
173
174     AliMCEvent* mcEvent = eventHandler->MCEvent();
175     if (!mcEvent) {
176       Printf("ERROR: Could not retrieve MC event");
177       return;
178     }
179
180     AliStack* stack = mcEvent->Stack();
181     if (!stack)
182     {
183       Printf("ERROR: Stack not available");
184       return;
185     }
186
187     // count if primaries are counted several times
188     Int_t max = stack->GetNprimary();
189     Int_t* primAcc = new Int_t[max];
190     Int_t* primRej = new Int_t[max];
191     for (Int_t i=0; i<max; i++)
192     {
193       primAcc[i] = 0;
194       primRej[i] = 0;
195     }
196
197     for (Int_t trackId = 0; trackId < fESD->GetNumberOfTracks(); trackId++)
198     {
199       AliESDtrack* track = fESD->GetTrack(trackId);
200       if (!track)
201       {
202         Printf("UNEXPECTED: Could not get track with id %d", trackId);
203         continue;
204       }
205       
206       if (track->GetLabel() == 0)
207       {
208         Printf("Track %d has no label. Skipping", trackId);
209         continue;
210       }
211       
212       Int_t label = TMath::Abs(track->GetLabel());
213       if (stack->IsPhysicalPrimary(label) == kTRUE)
214       {
215         if (label >= max)
216         {
217           Printf("Warning label %d is higher than number of primaries %d", label, max);
218           continue;
219         }
220
221         if (fTrackCutsPrimaries->AcceptTrack(track))
222         {
223           primAcc[label]++;
224         }
225         else
226           primRej[label]++;
227
228       }
229       else
230       {
231         if (!stack->Particle(label)) {
232           Printf("ERROR: Could not get particle with label %d", label);
233           continue;
234         }
235         
236         // Deuteron is ok, but missing in Pdg particles in root
237         if (stack->Particle(label)->GetPdgCode() != 10010020) {
238           if (!stack->Particle(label)->GetPDG()) {
239             Printf("ERROR: Could not get PDG particle of particle with label %d (pdg code is %d)", label, stack->Particle(label)->GetPdgCode());
240             stack->Particle(label)->Print();
241             continue;
242           }
243           
244           if (stack->Particle(label)->GetPDG()->Charge() == 0) {
245             Printf("Particle %d has 0 charge. Skipping.", label);
246             continue;
247           }
248         }
249           
250         fTrackCutsSecondaries->AcceptTrack(track);
251       }
252     }
253
254     for (Int_t i=0; i<max; i++) {
255       if (primAcc[i] == 1) {
256         fPrimStats->Fill(1);
257       } else if (primAcc[i] > 1) {
258         fPrimStats->Fill(2);
259         fPrimStats->Fill(3, primAcc[i]);
260       }
261
262       if (primRej[i] > 0) {
263         if (primAcc[i] == 0) {
264           fPrimStats->Fill(4);
265           fPrimStats->Fill(5, primRej[i]);
266         } else if (primAcc[i] > 0) {
267           fPrimStats->Fill(6);
268           fPrimStats->Fill(7, primRej[i]);
269         }
270       }
271     }
272
273     delete[] primAcc;
274     delete[] primRej;
275   }
276
277   // get the ESD vertex
278   fVertex->Fill(vtxESD->GetZv());
279 }
280
281 //________________________________________________________________________
282 void AliCutTask::Terminate(Option_t *) 
283 {
284   // Draw result to the screen
285   // Called once at the end of the query
286
287   fOutput = dynamic_cast<TList*> (GetOutputData(0));
288   if (!fOutput) {
289     Printf("ERROR: fOutput not available");
290     return;
291   }
292
293   fTrackCuts = dynamic_cast<AliESDtrackCuts*> (fOutput->FindObject("AliESDtrackCuts"));
294   if (!fTrackCuts) {
295     Printf("ERROR: fTrackCuts not available");
296     return;
297   }
298
299   fVertex = dynamic_cast<TH1F*> (fOutput->FindObject("fVertex"));
300   if (!fVertex) {
301     Printf("ERROR: fVertex not available");
302     return;
303   }
304
305   fTriggerStats = dynamic_cast<TH1F*> (fOutput->FindObject("fTriggerStats"));
306   if (!fTriggerStats) {
307     Printf("ERROR: fTriggerStats not available");
308     return;
309   }
310
311   fTrackCutsPrimaries = dynamic_cast<AliESDtrackCuts*> (fOutput->FindObject("fTrackCutsPrimaries"));
312   fTrackCutsSecondaries = dynamic_cast<AliESDtrackCuts*> (fOutput->FindObject("fTrackCutsSecondaries"));
313   
314   TFile* file = TFile::Open("trackCuts.root", "RECREATE");
315
316   fTrackCuts->SaveHistograms();
317   fVertex->Write();
318   fTriggerStats->Write();
319   
320   if (fTrackCutsPrimaries)
321     fTrackCutsPrimaries->SaveHistograms();
322
323   if (fTrackCutsSecondaries)
324     fTrackCutsSecondaries->SaveHistograms();
325
326   if (fPrimStats)
327     fPrimStats->Write();
328   
329   file->Close();
330   
331   Printf("Writting results to trackCuts.root.");
332   
333   fTrackCuts->DrawHistograms();
334
335   new TCanvas;
336   fVertex->Draw();
337
338   new TCanvas;
339   fTriggerStats->Draw();
340 }
341
342 void AliCutTask::EnableSecondaryStudy()
343
344   // 
345   // clones the fTrackCuts
346   //
347   
348   if (!fTrackCuts) {
349     Printf("ERROR: fTrackCuts 0. Do not call this function before SetTrackCuts");
350     return;
351   }
352   
353   fTrackCutsPrimaries = dynamic_cast<AliESDtrackCuts*> (fTrackCuts->Clone("fTrackCutsPrimaries"));
354   fTrackCutsSecondaries = dynamic_cast<AliESDtrackCuts*> (fTrackCuts->Clone("fTrackCutsSecondaries"));
355 }