9eb8f984a735768f7848f06348c56af47fe1e29f
[u/mrichter/AliRoot.git] / PWG0 / dNdEta / AlidNdEtaTask.cxx
1 /* $Id$ */
2
3 #include "AlidNdEtaTask.h"
4
5 #include <TStyle.h>
6 #include <TSystem.h>
7 #include <TCanvas.h>
8 #include <TVector3.h>
9 #include <TChain.h>
10 #include <TFile.h>
11 #include <TH1F.h>
12 #include <TH2F.h>
13 #include <TH3F.h>
14 #include <TParticle.h>
15 #include <TRandom.h>
16 #include <TNtuple.h>
17 #include <TObjString.h>
18 #include <TF1.h>
19
20 #include <AliLog.h>
21 #include <AliESDVertex.h>
22 #include <AliESDEvent.h>
23 #include <AliStack.h>
24 #include <AliHeader.h>
25 #include <AliGenEventHeader.h>
26 #include <AliMultiplicity.h>
27 #include <AliAnalysisManager.h>
28 #include <AliMCEventHandler.h>
29 #include <AliMCEvent.h>
30 #include <AliESDInputHandler.h>
31
32 #include "esdTrackCuts/AliESDtrackCuts.h"
33 #include "AliPWG0Helper.h"
34 #include "AliCorrection.h"
35 #include "AliCorrectionMatrix3D.h"
36 #include "dNdEta/dNdEtaAnalysis.h"
37
38 ClassImp(AlidNdEtaTask)
39
40 AlidNdEtaTask::AlidNdEtaTask(const char* opt) :
41   AliAnalysisTask("AlidNdEtaTask", ""),
42   fESD(0),
43   fOutput(0),
44   fOption(opt),
45   fAnalysisMode(kTPC),
46   fReadMC(kFALSE),
47   fEsdTrackCuts(0),
48   fdNdEtaAnalysisESD(0),
49   fMult(0),
50   fMultVtx(0),
51   fEvents(0),
52   fdNdEtaAnalysis(0),
53   fdNdEtaAnalysisTr(0),
54   fdNdEtaAnalysisTrVtx(0),
55   fdNdEtaAnalysisTracks(0),
56   fVertex(0),
57   fPartPt(0)
58 {
59   //
60   // Constructor. Initialization of pointers
61   //
62
63   // Define input and output slots here
64   DefineInput(0, TChain::Class());
65   DefineOutput(0, TList::Class());
66 }
67
68 AlidNdEtaTask::~AlidNdEtaTask()
69 {
70   //
71   // Destructor
72   //
73
74   // histograms are in the output list and deleted when the output
75   // list is deleted by the TSelector dtor
76
77   if (fOutput) {
78     delete fOutput;
79     fOutput = 0;
80   }
81 }
82
83 //________________________________________________________________________
84 void AlidNdEtaTask::ConnectInputData(Option_t *)
85 {
86   // Connect ESD
87   // Called once
88
89   Printf("AlidNdEtaTask::ConnectInputData called");
90
91   TTree* tree = dynamic_cast<TTree*> (GetInputData(0));
92   if (!tree) {
93     Printf("ERROR: Could not read tree from input slot 0");
94   } else {
95     // Disable all branches and enable only the needed ones
96     //tree->SetBranchStatus("*", 0);
97
98     tree->SetBranchStatus("fTriggerMask", 1);
99     tree->SetBranchStatus("fSPDVertex*", 1);
100
101     if (fAnalysisMode == kSPD)
102       tree->SetBranchStatus("fSPDMult*", 1);
103
104     if (fAnalysisMode == kTPC) {
105       AliESDtrackCuts::EnableNeededBranches(tree);
106       tree->SetBranchStatus("fTracks.fLabel", 1);
107     }
108
109     AliESDInputHandler *esdH = dynamic_cast<AliESDInputHandler*> (AliAnalysisManager::GetAnalysisManager()->GetInputEventHandler());
110
111     if (!esdH) {
112       Printf("ERROR: Could not get ESDInputHandler");
113     } else
114       fESD = esdH->GetEvent();
115   }
116 }
117
118 void AlidNdEtaTask::CreateOutputObjects()
119 {
120   // create result objects and add to output list
121
122   fOutput = new TList;
123   fOutput->SetOwner();
124
125   TString detector;
126   if (fAnalysisMode == kTPC)
127   {
128     detector = "TPC";
129   }
130   else if (fAnalysisMode == kSPD)
131     detector = "SPD";
132
133   fdNdEtaAnalysisESD = new dNdEtaAnalysis("fdNdEtaAnalysisESD", "fdNdEtaAnalysisESD", detector);
134   fOutput->Add(fdNdEtaAnalysisESD);
135
136   fMult = new TH1F("fMult", "fMult;Ntracks;Count", 201, -0.5, 200.5);
137   fOutput->Add(fMult);
138
139   fMultVtx = new TH1F("fMultVtx", "fMultVtx;Ntracks;Count", 201, -0.5, 200.5);
140   fOutput->Add(fMultVtx);
141
142   for (Int_t i=0; i<3; ++i)
143   {
144     fPartEta[i] = new TH1F(Form("dndeta_check_%d", i), Form("dndeta_check_%d", i), 60, -6, 6);
145     fPartEta[i]->Sumw2();
146     fOutput->Add(fPartEta[i]);
147   }
148
149   fEvents = new TH1F("dndeta_check_vertex", "dndeta_check_vertex", 40, -20, 20);
150   fOutput->Add(fEvents);
151
152   if (fReadMC)
153   {
154     fdNdEtaAnalysis = new dNdEtaAnalysis("dndeta", "dndeta", detector);
155     fOutput->Add(fdNdEtaAnalysis);
156
157     fdNdEtaAnalysisTr = new dNdEtaAnalysis("dndetaTr", "dndetaTr", detector);
158     fOutput->Add(fdNdEtaAnalysisTr);
159
160     fdNdEtaAnalysisTrVtx = new dNdEtaAnalysis("dndetaTrVtx", "dndetaTrVtx", detector);
161     fOutput->Add(fdNdEtaAnalysisTrVtx);
162
163     fdNdEtaAnalysisTracks = new dNdEtaAnalysis("dndetaTracks", "dndetaTracks", detector);
164     fOutput->Add(fdNdEtaAnalysisTracks);
165
166     fVertex = new TH3F("vertex_check", "vertex_check", 50, -50, 50, 50, -50, 50, 50, -50, 50);
167     fOutput->Add(fVertex);
168
169     fPartPt =  new TH1F("dndeta_check_pt", "dndeta_check_pt", 1000, 0, 10);
170     fPartPt->Sumw2();
171     fOutput->Add(fPartPt);
172   }
173 }
174
175 void AlidNdEtaTask::Exec(Option_t*)
176 {
177   // process the event
178
179   // Check prerequisites
180   if (!fESD)
181   {
182     AliDebug(AliLog::kError, "ESD branch not available");
183     return;
184   }
185
186   // trigger definition
187   Bool_t eventTriggered = AliPWG0Helper::IsEventTriggered(fESD->GetTriggerMask(), AliPWG0Helper::kMB2);
188
189   Bool_t eventVertex = AliPWG0Helper::IsVertexReconstructed(fESD->GetVertex());
190
191   // get the ESD vertex
192   const AliESDVertex* vtxESD = fESD->GetVertex();
193   Double_t vtx[3];
194   vtxESD->GetXYZ(vtx);
195
196   // post the data already here
197   PostData(0, fOutput);
198
199   // create list of (label, eta, pt) tuples
200   Int_t inputCount = 0;
201   Int_t* labelArr = 0;
202   Float_t* etaArr = 0;
203   Float_t* ptArr = 0;
204   if (fAnalysisMode == kSPD)
205   {
206     // get tracklets
207     const AliMultiplicity* mult = fESD->GetMultiplicity();
208     if (!mult)
209     {
210       AliDebug(AliLog::kError, "AliMultiplicity not available");
211       return;
212     }
213
214     labelArr = new Int_t[mult->GetNumberOfTracklets()];
215     etaArr = new Float_t[mult->GetNumberOfTracklets()];
216     ptArr = new Float_t[mult->GetNumberOfTracklets()];
217
218     // get multiplicity from ITS tracklets
219     for (Int_t i=0; i<mult->GetNumberOfTracklets(); ++i)
220     {
221       //printf("%d %f %f %f\n", i, mult->GetTheta(i), mult->GetPhi(i), mult->GetDeltaPhi(i));
222
223       // This removes non-tracklets in PDC06 data. Very bad solution. New solution is implemented for newer data. Keeping this for compatibility.
224       if (mult->GetDeltaPhi(i) < -1000)
225         continue;
226
227       etaArr[inputCount] = mult->GetEta(i);
228       labelArr[inputCount] = mult->GetLabel(i);
229       ptArr[inputCount] = 0; // no pt for tracklets
230       ++inputCount;
231     }
232   }
233   else if (fAnalysisMode == kTPC)
234   {
235     if (!fEsdTrackCuts)
236     {
237       AliDebug(AliLog::kError, "fESDTrackCuts not available");
238       return;
239     }
240
241     // get multiplicity from ESD tracks
242     TObjArray* list = fEsdTrackCuts->GetAcceptedTracks(fESD);
243     Int_t nGoodTracks = list->GetEntries();
244
245     labelArr = new Int_t[nGoodTracks];
246     etaArr = new Float_t[nGoodTracks];
247     ptArr = new Float_t[nGoodTracks];
248
249     // loop over esd tracks
250     for (Int_t i=0; i<nGoodTracks; i++)
251     {
252       AliESDtrack* esdTrack = dynamic_cast<AliESDtrack*> (list->At(i));
253       if (!esdTrack)
254       {
255         AliDebug(AliLog::kError, Form("ERROR: Could not retrieve track %d.", i));
256         continue;
257       }
258
259       etaArr[inputCount] = esdTrack->Eta();
260       labelArr[inputCount] = TMath::Abs(esdTrack->GetLabel());
261       ptArr[inputCount] = esdTrack->Pt();
262       ++inputCount;
263
264       Printf("%f", esdTrack->Pt());
265     }
266   }
267   else
268     return;
269
270   // Processing of ESD information (always)
271   if (eventTriggered)
272   {
273     // control hists
274     fMult->Fill(inputCount);
275     if (eventVertex)
276       fMultVtx->Fill(inputCount);
277
278     // count all events
279     if (!eventVertex)
280       vtx[2] = 0;
281
282     for (Int_t i=0; i<inputCount; ++i)
283     {
284       Float_t eta = etaArr[i];
285       Float_t pt  = ptArr[i];
286
287       fdNdEtaAnalysisESD->FillTrack(vtx[2], eta, pt);
288
289       if (TMath::Abs(vtx[2]) < 20)
290       {
291         fPartEta[0]->Fill(eta);
292
293         if (vtx[2] < 0)
294           fPartEta[1]->Fill(eta);
295         else
296           fPartEta[2]->Fill(eta);
297       }
298     }
299
300     // for event count per vertex
301     fdNdEtaAnalysisESD->FillEvent(vtx[2], inputCount);
302
303     // control hists
304     fEvents->Fill(vtx[2]);
305   }
306
307   if (fReadMC)   // Processing of MC information (optional)
308   {
309     AliMCEventHandler* eventHandler = dynamic_cast<AliMCEventHandler*> (AliAnalysisManager::GetAnalysisManager()->GetMCtruthEventHandler());
310     if (!eventHandler) {
311       Printf("ERROR: Could not retrieve MC event handler");
312       return;
313     }
314
315     AliMCEvent* mcEvent = eventHandler->MCEvent();
316     if (!mcEvent) {
317       Printf("ERROR: Could not retrieve MC event");
318       return;
319     }
320
321     AliStack* stack = mcEvent->Stack();
322     if (!stack)
323     {
324       AliDebug(AliLog::kError, "Stack not available");
325       return;
326     }
327
328     AliHeader* header = mcEvent->Header();
329     if (!header)
330     {
331       AliDebug(AliLog::kError, "Header not available");
332       return;
333     }
334
335     // get the MC vertex
336     AliGenEventHeader* genHeader = header->GenEventHeader();
337     TArrayF vtxMC(3);
338     genHeader->PrimaryVertex(vtxMC);
339
340     // loop over mc particles
341     Int_t nPrim  = stack->GetNprimary();
342
343     Int_t nAcceptedParticles = 0;
344
345     for (Int_t iMc = 0; iMc < nPrim; ++iMc)
346     {
347       //Printf("Getting particle %d", iMc);
348       TParticle* particle = stack->Particle(iMc);
349
350       if (!particle)
351         continue;
352
353       if (AliPWG0Helper::IsPrimaryCharged(particle, nPrim) == kFALSE)
354         continue;
355
356       AliDebug(AliLog::kDebug+1, Form("Accepted primary %d, unique ID: %d", iMc, particle->GetUniqueID()));
357       ++nAcceptedParticles;
358
359       Float_t eta = particle->Eta();
360       Float_t pt = particle->Pt();
361
362       fdNdEtaAnalysis->FillTrack(vtxMC[2], eta, pt);
363       fVertex->Fill(particle->Vx(), particle->Vy(), particle->Vz());
364
365       if (eventTriggered)
366       {
367         fdNdEtaAnalysisTr->FillTrack(vtxMC[2], eta, pt);
368         if (eventVertex)
369           fdNdEtaAnalysisTrVtx->FillTrack(vtxMC[2], eta, pt);
370       }
371
372       if (TMath::Abs(eta) < 0.8)
373         fPartPt->Fill(pt);
374     }
375
376     fdNdEtaAnalysis->FillEvent(vtxMC[2], nAcceptedParticles);
377     if (eventTriggered)
378     {
379       fdNdEtaAnalysisTr->FillEvent(vtxMC[2], nAcceptedParticles);
380       if (eventVertex)
381         fdNdEtaAnalysisTrVtx->FillEvent(vtxMC[2], nAcceptedParticles);
382     }
383
384     if (eventTriggered && eventVertex)
385     {
386       // from tracks is only done for triggered and vertex reconstructed events
387
388       for (Int_t i=0; i<inputCount; ++i)
389       {
390         Int_t label = labelArr[i];
391
392         if (label < 0)
393           continue;
394
395         //Printf("Getting particle of track %d", label);
396         TParticle* particle = stack->Particle(label);
397
398         if (!particle)
399         {
400           AliDebug(AliLog::kError, Form("ERROR: Could not retrieve particle %d.", label));
401           continue;
402         }
403
404         fdNdEtaAnalysisTracks->FillTrack(vtxMC[2], particle->Eta(), particle->Pt());
405       } // end of track loop
406
407       // for event count per vertex
408       fdNdEtaAnalysisTracks->FillEvent(vtxMC[2], inputCount);
409     }
410   }
411
412   delete[] etaArr;
413   delete[] labelArr;
414   delete[] ptArr;
415 }
416
417 void AlidNdEtaTask::Terminate(Option_t *)
418 {
419   // The Terminate() function is the last function to be called during
420   // a query. It always runs on the client, it can be used to present
421   // the results graphically or save the results to file.
422
423   fOutput = dynamic_cast<TList*> (GetOutputData(0));
424   if (!fOutput) {
425     Printf("ERROR: fOutput not available");
426     return;
427   }
428
429   fdNdEtaAnalysisESD = dynamic_cast<dNdEtaAnalysis*> (fOutput->FindObject("fdNdEtaAnalysisESD"));
430   fMult = dynamic_cast<TH1F*> (fOutput->FindObject("fMult"));
431   fMultVtx = dynamic_cast<TH1F*> (fOutput->FindObject("fMultVtx"));
432   fEvents = dynamic_cast<TH1F*> (fOutput->FindObject("dndeta_check_vertex"));
433   for (Int_t i=0; i<3; ++i)
434     fPartEta[i] = dynamic_cast<TH1F*> (fOutput->FindObject(Form("dndeta_check_%d", i)));
435
436   if (!fdNdEtaAnalysisESD)
437   {
438     AliDebug(AliLog::kError, "ERROR: fdNdEtaAnalysisESD not available");
439     return;
440   }
441
442   if (fMult && fMultVtx)
443   {
444     new TCanvas;
445     fMult->Draw();
446     fMultVtx->SetLineColor(2);
447     fMultVtx->DrawCopy("SAME");
448   }
449
450   if (fPartEta[0])
451   {
452     Int_t events1 = (Int_t) fEvents->Integral(fEvents->GetXaxis()->FindBin(-19.9), fEvents->GetXaxis()->FindBin(-0.001));
453     Int_t events2 = (Int_t) fEvents->Integral(fEvents->GetXaxis()->FindBin(0.001), fEvents->GetXaxis()->FindBin(19.9));
454
455     fPartEta[0]->Scale(1.0 / (events1 + events2));
456     fPartEta[1]->Scale(1.0 / events1);
457     fPartEta[2]->Scale(1.0 / events2);
458
459     for (Int_t i=0; i<3; ++i)
460       fPartEta[i]->Scale(1.0 / fPartEta[i]->GetBinWidth(1));
461
462     new TCanvas("control", "control", 500, 500);
463     for (Int_t i=0; i<3; ++i)
464     {
465       fPartEta[i]->SetLineColor(i+1);
466       fPartEta[i]->Draw((i==0) ? "" : "SAME");
467     }
468   }
469
470   if (fEvents)
471   {
472     new TCanvas("control3", "control3", 500, 500);
473     fEvents->DrawCopy();
474   }
475
476   TFile* fout = new TFile("analysis_esd_raw.root", "RECREATE");
477
478   if (fdNdEtaAnalysisESD)
479     fdNdEtaAnalysisESD->SaveHistograms();
480
481   if (fEsdTrackCuts)
482     fEsdTrackCuts->SaveHistograms("esd_tracks_cuts");
483
484   if (fMult)
485     fMult->Write();
486
487   if (fMultVtx)
488     fMultVtx->Write();
489
490   for (Int_t i=0; i<3; ++i)
491     if (fPartEta[i])
492       fPartEta[i]->Write();
493
494   if (fEvents)
495     fEvents->Write();
496
497   fout->Write();
498   fout->Close();
499
500   Printf("Writting result to analysis_esd_raw.root");
501
502   if (fReadMC)
503   {
504     fdNdEtaAnalysis = dynamic_cast<dNdEtaAnalysis*> (fOutput->FindObject("dndeta"));
505     fdNdEtaAnalysisTr = dynamic_cast<dNdEtaAnalysis*> (fOutput->FindObject("dndetaTr"));
506     fdNdEtaAnalysisTrVtx = dynamic_cast<dNdEtaAnalysis*> (fOutput->FindObject("dndetaTrVtx"));
507     fdNdEtaAnalysisTracks = dynamic_cast<dNdEtaAnalysis*> (fOutput->FindObject("dndetaTracks"));
508     fPartPt = dynamic_cast<TH1F*> (fOutput->FindObject("dndeta_check_pt"));
509
510     if (!fdNdEtaAnalysis || !fdNdEtaAnalysisTr || !fdNdEtaAnalysisTrVtx || !fPartPt)
511     {
512       AliDebug(AliLog::kError, Form("ERROR: Histograms not available %p %p", (void*) fdNdEtaAnalysis, (void*) fPartPt));
513       return;
514     }
515
516     fdNdEtaAnalysis->Finish(0, -1, AlidNdEtaCorrection::kNone);
517     fdNdEtaAnalysisTr->Finish(0, -1, AlidNdEtaCorrection::kNone);
518     fdNdEtaAnalysisTrVtx->Finish(0, -1, AlidNdEtaCorrection::kNone);
519     fdNdEtaAnalysisTracks->Finish(0, -1, AlidNdEtaCorrection::kNone);
520
521     Int_t events = (Int_t) fdNdEtaAnalysis->GetData()->GetEventCorrection()->GetMeasuredHistogram()->Integral();
522     fPartPt->Scale(1.0/events);
523     fPartPt->Scale(1.0/fPartPt->GetBinWidth(1));
524
525     TFile* fout = new TFile("analysis_mc.root","RECREATE");
526
527     fdNdEtaAnalysis->SaveHistograms();
528     fdNdEtaAnalysisTr->SaveHistograms();
529     fdNdEtaAnalysisTrVtx->SaveHistograms();
530     fdNdEtaAnalysisTracks->SaveHistograms();
531     fPartPt->Write();
532
533     fout->Write();
534     fout->Close();
535
536     Printf("Writting result to analysis_mc.root");
537
538     if (fPartPt)
539     {
540       new TCanvas("control2", "control2", 500, 500);
541       fPartPt->DrawCopy();
542     }
543   }
544 }