4ba3692bce222fe5b845df2c2401b3506508aa62
[u/mrichter/AliRoot.git] / PWG0 / dNdEta / AlidNdEtaCorrectionTask.cxx
1 /* $Id$ */
2
3 #include "AlidNdEtaCorrectionTask.h"
4
5 #include <TROOT.h>
6 #include <TCanvas.h>
7 #include <TChain.h>
8 #include <TFile.h>
9 #include <TH1F.h>
10 #include <TH2F.h>
11 #include <TProfile.h>
12 #include <TParticle.h>
13
14 #include <AliLog.h>
15 #include <AliESDVertex.h>
16 #include <AliESDEvent.h>
17 #include <AliStack.h>
18 #include <AliHeader.h>
19 #include <AliGenEventHeader.h>
20 #include <AliMultiplicity.h>
21 #include <AliAnalysisManager.h>
22 #include <AliMCEventHandler.h>
23 #include <AliMCEvent.h>
24 #include <AliESDInputHandler.h>
25
26 #include "esdTrackCuts/AliESDtrackCuts.h"
27 #include "AliPWG0Helper.h"
28 //#include "AliCorrection.h"
29 //#include "AliCorrectionMatrix3D.h"
30 #include "dNdEta/dNdEtaAnalysis.h"
31 #include "dNdEta/AlidNdEtaCorrection.h"
32 #include "AliVertexerTracks.h"
33
34 ClassImp(AlidNdEtaCorrectionTask)
35
36 AlidNdEtaCorrectionTask::AlidNdEtaCorrectionTask(const char* opt) :
37   AliAnalysisTask("AlidNdEtaCorrectionTask", ""),
38   fESD(0),
39   fOutput(0),
40   fOption(opt),
41   fAnalysisMode(AliPWG0Helper::kTPC),
42   fSignMode(0),
43   fEsdTrackCuts(0),
44   fdNdEtaCorrection(0),
45   fdNdEtaAnalysisMC(0),
46   fdNdEtaAnalysisESD(0),
47   fPIDParticles(0),
48   fPIDTracks(0),
49   fVertexCorrelation(0),
50   fVertexProfile(0),
51   fVertexShiftNorm(0),
52   fSigmaVertexTracks(0),
53   fSigmaVertexPrim(0)
54 {
55   //
56   // Constructor. Initialization of pointers
57   //
58
59   // Define input and output slots here
60   DefineInput(0, TChain::Class());
61   DefineOutput(0, TList::Class());
62
63   for (Int_t i=0; i<2; i++)
64     fdNdEtaCorrectionProcessType[i] = 0;
65 }
66
67 AlidNdEtaCorrectionTask::~AlidNdEtaCorrectionTask()
68 {
69   //
70   // Destructor
71   //
72
73   // histograms are in the output list and deleted when the output
74   // list is deleted by the TSelector dtor
75
76   if (fOutput) {
77     delete fOutput;
78     fOutput = 0;
79   }
80 }
81
82 //________________________________________________________________________
83 void AlidNdEtaCorrectionTask::ConnectInputData(Option_t *)
84 {
85   // Connect ESD
86   // Called once
87
88   Printf("AlidNdEtaCorrectionTask::ConnectInputData called");
89
90   TTree* tree = dynamic_cast<TTree*> (GetInputData(0));
91   if (!tree) {
92     Printf("ERROR: Could not read tree from input slot 0");
93   } else {
94     // Disable all branches and enable only the needed ones
95     //tree->SetBranchStatus("*", 0);
96
97     tree->SetBranchStatus("fTriggerMask", 1);
98     tree->SetBranchStatus("fSPDVertex*", 1);
99     // PrimaryVertex
100
101     if (fAnalysisMode == AliPWG0Helper::kSPD)
102       tree->SetBranchStatus("fSPDMult*", 1);
103
104     if (fAnalysisMode == AliPWG0Helper::kTPC || fAnalysisMode == AliPWG0Helper::kTPCITS) {
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   // disable info messages of AliMCEvent (per event)
118   AliLog::SetClassDebugLevel("AliMCEvent", AliLog::kWarning - AliLog::kDebug + 1);
119 }
120
121 void AlidNdEtaCorrectionTask::CreateOutputObjects()
122 {
123   // create result objects and add to output list
124
125   if (fOption.Contains("only-positive"))
126   {
127     Printf("INFO: Processing only positive particles.");
128     fSignMode = 1;
129   }
130   else if (fOption.Contains("only-negative"))
131   {
132     Printf("INFO: Processing only negative particles.");
133     fSignMode = -1;
134   }
135
136   fOutput = new TList;
137   fOutput->SetOwner();
138
139   fdNdEtaCorrection = new AlidNdEtaCorrection("dndeta_correction", "dndeta_correction", fAnalysisMode);
140   fOutput->Add(fdNdEtaCorrection);
141
142   fPIDParticles = new TH1F("pid_particles", "PID of generated primary particles", 10001, -5000.5, 5000.5);
143   fOutput->Add(fPIDParticles);
144
145   fPIDTracks = new TH1F("pid_tracks", "MC PID of reconstructed tracks", 10001, -5000.5, 5000.5);
146   fOutput->Add(fPIDTracks);
147
148   fdNdEtaAnalysisMC = new dNdEtaAnalysis("dndetaMC", "dndetaMC", fAnalysisMode);
149   fOutput->Add(fdNdEtaAnalysisMC);
150
151   fdNdEtaAnalysisESD = new dNdEtaAnalysis("dndetaESD", "dndetaESD", fAnalysisMode);
152   fOutput->Add(fdNdEtaAnalysisESD);
153
154   if (fOption.Contains("process-types")) {
155     fdNdEtaCorrectionProcessType[0] = new AlidNdEtaCorrection("dndeta_correction_ND", "dndeta_correction_ND", fAnalysisMode);
156     fdNdEtaCorrectionProcessType[1] = new AlidNdEtaCorrection("dndeta_correction_SD", "dndeta_correction_SD", fAnalysisMode);
157     fdNdEtaCorrectionProcessType[2] = new AlidNdEtaCorrection("dndeta_correction_DD", "dndeta_correction_DD", fAnalysisMode);
158
159     fOutput->Add(fdNdEtaCorrectionProcessType[0]);
160     fOutput->Add(fdNdEtaCorrectionProcessType[1]);
161     fOutput->Add(fdNdEtaCorrectionProcessType[2]);
162   }
163
164   if (fOption.Contains("sigma-vertex"))
165   {
166     fSigmaVertexTracks = new TH1F("fSigmaVertexTracks", "fSigmaVertexTracks;Nsigma2vertex;NacceptedTracks", 60, 0.05, 6.05);
167     fSigmaVertexPrim = new TH1F("fSigmaVertexPrim", "fSigmaVertexPrim;Nsigma2vertex;NacceptedPrimaries", 60, 0.05, 6.05);
168     fOutput->Add(fSigmaVertexTracks);
169     fOutput->Add(fSigmaVertexPrim);
170     Printf("WARNING: sigma-vertex analysis enabled. This will produce weird results in the AliESDtrackCuts histograms");
171   }
172
173   fVertexCorrelation = new TH2F("fVertexCorrelation", "fVertexCorrelation;MC z-vtx;ESD z-vtx", 80, -20, 20, 80, -20, 20);
174   fVertexProfile = new TProfile("fVertexProfile", "fVertexProfile;MC z-vtx;MC z-vtx - ESD z-vtx", 40, -20, 20);
175   fVertexShiftNorm = new TH1F("fVertexShiftNorm", "fVertexShiftNorm;(MC z-vtx - ESD z-vtx) / #sigma_{ESD z-vtx};Entries", 200, -100, 100);
176 }
177
178 void AlidNdEtaCorrectionTask::Exec(Option_t*)
179 {
180   // process the event
181
182   // Check prerequisites
183   if (!fESD)
184   {
185     AliDebug(AliLog::kError, "ESD branch not available");
186     return;
187   }
188
189   // trigger definition
190   Bool_t eventTriggered = AliPWG0Helper::IsEventTriggered(fESD->GetTriggerMask(), AliPWG0Helper::kMB2);
191
192   // post the data already here
193   PostData(0, fOutput);
194
195   // MC info
196   AliMCEventHandler* eventHandler = dynamic_cast<AliMCEventHandler*> (AliAnalysisManager::GetAnalysisManager()->GetMCtruthEventHandler());
197   if (!eventHandler) {
198     Printf("ERROR: Could not retrieve MC event handler");
199     return;
200   }
201
202   AliMCEvent* mcEvent = eventHandler->MCEvent();
203   if (!mcEvent) {
204     Printf("ERROR: Could not retrieve MC event");
205     return;
206   }
207
208   AliStack* stack = mcEvent->Stack();
209   if (!stack)
210   {
211     AliDebug(AliLog::kError, "Stack not available");
212     return;
213   }
214
215   AliHeader* header = mcEvent->Header();
216   if (!header)
217   {
218     AliDebug(AliLog::kError, "Header not available");
219     return;
220   }
221
222   // get process type; NB: this only works for Pythia
223   Int_t processType = AliPWG0Helper::GetPythiaEventProcessType(header);
224   AliDebug(AliLog::kDebug+1, Form("Found pythia process type %d", processType));
225
226   if (processType<0)
227     AliDebug(AliLog::kError, Form("Unknown Pythia process type %d.", processType));
228
229   // get the MC vertex
230   AliGenEventHeader* genHeader = header->GenEventHeader();
231   TArrayF vtxMC(3);
232   genHeader->PrimaryVertex(vtxMC);
233
234   // get the ESD vertex
235   const AliESDVertex* vtxESD = AliPWG0Helper::GetVertex(fESD, fAnalysisMode);
236
237   Bool_t eventVertex = kFALSE;
238   Double_t vtx[3];
239   if (vtxESD) 
240   {
241     vtxESD->GetXYZ(vtx);
242     eventVertex = kTRUE;
243     
244     Double_t diff = vtxMC[2] - vtx[2];
245     if (vtxESD->GetZRes() > 0) 
246       fVertexShiftNorm->Fill(diff / vtxESD->GetZRes());
247   } 
248   else
249     Printf("No vertex found");
250
251
252   // create list of (label, eta, pt) tuples
253   Int_t inputCount = 0;
254   Int_t* labelArr = 0;
255   Float_t* etaArr = 0;
256   Float_t* ptArr = 0;
257   if (fAnalysisMode == AliPWG0Helper::kSPD)
258   {
259     // get tracklets
260     const AliMultiplicity* mult = fESD->GetMultiplicity();
261     if (!mult)
262     {
263       AliDebug(AliLog::kError, "AliMultiplicity not available");
264       return;
265     }
266
267     labelArr = new Int_t[mult->GetNumberOfTracklets()];
268     etaArr = new Float_t[mult->GetNumberOfTracklets()];
269     ptArr = new Float_t[mult->GetNumberOfTracklets()];
270
271     // get multiplicity from ITS tracklets
272     for (Int_t i=0; i<mult->GetNumberOfTracklets(); ++i)
273     {
274       //printf("%d %f %f %f\n", i, mult->GetTheta(i), mult->GetPhi(i), mult->GetDeltaPhi(i));
275
276       // This removes non-tracklets in PDC06 data. Very bad solution. New solution is implemented for newer data. Keeping this for compatibility.
277       if (mult->GetDeltaPhi(i) < -1000)
278         continue;
279
280       etaArr[inputCount] = mult->GetEta(i);
281       labelArr[inputCount] = mult->GetLabel(i);
282       ptArr[inputCount] = 0; // no pt for tracklets
283       ++inputCount;
284     }
285   }
286   else if (fAnalysisMode == AliPWG0Helper::kTPC || fAnalysisMode == AliPWG0Helper::kTPCITS)
287   {
288     if (!fEsdTrackCuts)
289     {
290       AliDebug(AliLog::kError, "fESDTrackCuts not available");
291       return;
292     }
293
294     // get multiplicity from ESD tracks
295     TObjArray* list = fEsdTrackCuts->GetAcceptedTracks(fESD);
296     Int_t nGoodTracks = list->GetEntries();
297
298     Printf("Accepted %d tracks", nGoodTracks);
299
300     labelArr = new Int_t[nGoodTracks];
301     etaArr = new Float_t[nGoodTracks];
302     ptArr = new Float_t[nGoodTracks];
303
304     // loop over esd tracks
305     for (Int_t i=0; i<nGoodTracks; i++)
306     {
307       AliESDtrack* esdTrack = dynamic_cast<AliESDtrack*> (list->At(i));
308       if (!esdTrack)
309       {
310         AliDebug(AliLog::kError, Form("ERROR: Could not retrieve track %d.", i));
311         continue;
312       }
313
314       etaArr[inputCount] = esdTrack->Eta();
315       labelArr[inputCount] = TMath::Abs(esdTrack->GetLabel());
316       ptArr[inputCount] = esdTrack->Pt();
317       ++inputCount;
318     }
319
320     delete list;
321   }
322   else
323     return;
324
325   if (eventTriggered && eventVertex)
326   {
327     fVertexCorrelation->Fill(vtxMC[2], vtx[2]);
328     fVertexProfile->Fill(vtxMC[2], vtxMC[2] - vtx[2]);
329   }
330
331
332   // loop over mc particles
333   Int_t nPrim  = stack->GetNprimary();
334
335   for (Int_t iMc = 0; iMc < nPrim; ++iMc)
336   {
337     //Printf("Getting particle %d", iMc);
338     TParticle* particle = stack->Particle(iMc);
339
340     if (!particle)
341       continue;
342
343     if (AliPWG0Helper::IsPrimaryCharged(particle, nPrim) == kFALSE)
344       continue;
345
346     if (SignOK(particle->GetPDG()) == kFALSE)
347       continue;
348
349     Float_t eta = particle->Eta();
350     Float_t pt = particle->Pt();
351
352     fdNdEtaCorrection->FillMCParticle(vtxMC[2], eta, pt, eventTriggered, eventVertex, processType);
353
354     if (fdNdEtaCorrectionProcessType[0])
355     {
356       // non diffractive
357       if (processType!=92 && processType!=93 && processType!=94)
358         fdNdEtaCorrectionProcessType[0]->FillMCParticle(vtxMC[2], eta, pt, eventTriggered, eventVertex, processType);
359
360       // single diffractive
361       if (processType==92 || processType==93)
362         fdNdEtaCorrectionProcessType[1]->FillMCParticle(vtxMC[2], eta, pt, eventTriggered, eventVertex, processType);
363
364       // double diffractive
365       if (processType==94)
366         fdNdEtaCorrectionProcessType[2]->FillMCParticle(vtxMC[2], eta, pt, eventTriggered, eventVertex, processType);
367     }
368
369     if (eventTriggered)
370       if (eventVertex)
371         fdNdEtaAnalysisMC->FillTrack(vtxMC[2], eta, pt);
372   }
373
374   for (Int_t i=0; i<inputCount; ++i)
375   {
376     Int_t label = labelArr[i];
377
378     if (label < 0)
379     {
380       Printf("WARNING: cannot find corresponding mc part for track(let) %d with label %d.", i, label);
381       continue;
382     }
383
384     TParticle* particle = stack->Particle(label);
385     if (!particle)
386     {
387       AliDebug(AliLog::kError, Form("UNEXPECTED: particle with label %d not found in stack (track loop).", label));
388       continue;
389     }
390
391     if (SignOK(particle->GetPDG()) == kFALSE)
392         continue;
393
394     if (eventTriggered && eventVertex)
395     {
396       fdNdEtaCorrection->FillTrackedParticle(vtxMC[2], particle->Eta(), particle->Pt());
397       fdNdEtaAnalysisESD->FillTrack(vtxMC[2], particle->Eta(), particle->Pt());
398       if (particle->Pt() > 0.1 && particle->Pt() < 0.2)
399       {
400         fPIDTracks->Fill(particle->GetPdgCode());
401       }
402
403       if (fdNdEtaCorrectionProcessType[0])
404       {
405         // non diffractive
406         if (processType!=92 && processType!=93 && processType!=94)
407           fdNdEtaCorrectionProcessType[0]->FillTrackedParticle(vtxMC[2], particle->Eta(), particle->Pt());
408
409         // single diffractive
410         if (processType==92 || processType==93)
411           fdNdEtaCorrectionProcessType[1]->FillTrackedParticle(vtxMC[2], particle->Eta(), particle->Pt());
412
413         // double diffractive
414         if (processType==94)
415           fdNdEtaCorrectionProcessType[2]->FillTrackedParticle(vtxMC[2], particle->Eta(), particle->Pt());
416       }
417     }
418   }
419
420   if (eventTriggered && eventVertex)
421   {
422     fdNdEtaAnalysisMC->FillEvent(vtxMC[2], inputCount);
423     fdNdEtaAnalysisESD->FillEvent(vtxMC[2], inputCount);
424   }
425
426    // stuff regarding the vertex reco correction and trigger bias correction
427   fdNdEtaCorrection->FillEvent(vtxMC[2], inputCount, eventTriggered, eventVertex, processType);
428
429   if (fdNdEtaCorrectionProcessType[0])
430   {
431     // non diffractive
432     if (processType!=92 && processType!=93 && processType!=94)
433       fdNdEtaCorrectionProcessType[0]->FillEvent(vtxMC[2], inputCount, eventTriggered, eventVertex, processType);
434
435     // single diffractive
436     if (processType==92 || processType==93)
437       fdNdEtaCorrectionProcessType[1]->FillEvent(vtxMC[2], inputCount, eventTriggered, eventVertex, processType);
438
439     // double diffractive
440     if (processType==94)
441       fdNdEtaCorrectionProcessType[2]->FillEvent(vtxMC[2], inputCount, eventTriggered, eventVertex, processType);
442   }
443
444   delete[] etaArr;
445   delete[] labelArr;
446   delete[] ptArr;
447
448   // fills the fSigmaVertex histogram (systematic study)
449   if (fSigmaVertexTracks)
450   {
451     // save the old value
452     Float_t oldSigmaVertex = fEsdTrackCuts->GetMinNsigmaToVertex();
453
454     // set to maximum
455     fEsdTrackCuts->SetMinNsigmaToVertex(6);
456
457     TObjArray* list = fEsdTrackCuts->GetAcceptedTracks(fESD);
458     Int_t nGoodTracks = list->GetEntries();
459
460     // loop over esd tracks
461     for (Int_t i=0; i<nGoodTracks; i++)
462     {
463       AliESDtrack* esdTrack = dynamic_cast<AliESDtrack*> (list->At(i));
464       if (!esdTrack)
465       {
466         AliError(Form("ERROR: Could not retrieve track %d.", i));
467         continue;
468       }
469
470       Float_t sigma2Vertex = fEsdTrackCuts->GetSigmaToVertex(esdTrack);
471
472       for (Double_t nSigma = 0.1; nSigma < 6.05; nSigma += 0.1)
473       {
474         if (sigma2Vertex < nSigma)
475         {
476           fSigmaVertexTracks->Fill(nSigma);
477
478           Int_t label = TMath::Abs(esdTrack->GetLabel());
479           TParticle* particle = stack->Particle(label);
480           if (!particle || label >= nPrim)
481             continue;
482
483           if (AliPWG0Helper::IsPrimaryCharged(particle, nPrim))
484             fSigmaVertexPrim->Fill(nSigma);
485         }
486       }
487     }
488
489     delete list;
490     list = 0;
491
492     // set back the old value
493     fEsdTrackCuts->SetMinNsigmaToVertex(oldSigmaVertex);
494   }
495 }
496
497 void AlidNdEtaCorrectionTask::Terminate(Option_t *)
498 {
499   // The Terminate() function is the last function to be called during
500   // a query. It always runs on the client, it can be used to present
501   // the results graphically or save the results to file.
502
503   fOutput = dynamic_cast<TList*> (GetOutputData(0));
504   if (!fOutput) {
505     Printf("ERROR: fOutput not available");
506     return;
507   }
508
509   fdNdEtaCorrection = dynamic_cast<AlidNdEtaCorrection*> (fOutput->FindObject("dndeta_correction"));
510   fdNdEtaAnalysisMC = dynamic_cast<dNdEtaAnalysis*> (fOutput->FindObject("dndetaMC"));
511   fdNdEtaAnalysisESD = dynamic_cast<dNdEtaAnalysis*> (fOutput->FindObject("dndetaESD"));
512   if (!fdNdEtaCorrection || !fdNdEtaAnalysisMC || !fdNdEtaAnalysisESD)
513   {
514     AliDebug(AliLog::kError, "Could not read object from output list");
515     return;
516   }
517
518   fdNdEtaCorrection->Finish();
519
520   TString fileName;
521   fileName.Form("correction_map%s.root", fOption.Data());
522   TFile* fout = new TFile(fileName, "RECREATE");
523
524   if (fEsdTrackCuts)
525     fEsdTrackCuts->SaveHistograms("esd_track_cuts");
526   fdNdEtaCorrection->SaveHistograms();
527   fdNdEtaAnalysisMC->SaveHistograms();
528   fdNdEtaAnalysisESD->SaveHistograms();
529
530   fdNdEtaCorrectionProcessType[0] = dynamic_cast<AlidNdEtaCorrection*> (fOutput->FindObject("dndeta_correction_ND"));
531   fdNdEtaCorrectionProcessType[1] = dynamic_cast<AlidNdEtaCorrection*> (fOutput->FindObject("dndeta_correction_SD"));
532   fdNdEtaCorrectionProcessType[2] = dynamic_cast<AlidNdEtaCorrection*> (fOutput->FindObject("dndeta_correction_DD"));
533   for (Int_t i=0; i<3; ++i)
534     if (fdNdEtaCorrectionProcessType[i])
535       fdNdEtaCorrectionProcessType[i]->SaveHistograms();
536
537   fSigmaVertexTracks = dynamic_cast<TH1F*> (fOutput->FindObject("fSigmaVertexTracks"));
538   if (fSigmaVertexTracks)
539     fSigmaVertexTracks->Write();
540   fSigmaVertexPrim = dynamic_cast<TH1F*> (fOutput->FindObject("fSigmaVertexPrim"));
541   if (fSigmaVertexPrim)
542     fSigmaVertexPrim->Write();
543
544   if (fVertexCorrelation)
545     fVertexCorrelation->Write();
546   if (fVertexProfile)
547     fVertexProfile->Write();
548   if (fVertexShiftNorm)
549     fVertexShiftNorm->Write();
550
551   fout->Write();
552   fout->Close();
553
554   //fdNdEtaCorrection->DrawHistograms();
555
556   Printf("Writting result to %s", fileName.Data());
557
558   if (fPIDParticles && fPIDTracks)
559   {
560     new TCanvas("pidcanvas", "pidcanvas", 500, 500);
561
562     fPIDParticles->Draw();
563     fPIDTracks->SetLineColor(2);
564     fPIDTracks->Draw("SAME");
565
566     TDatabasePDG* pdgDB = new TDatabasePDG;
567
568     for (Int_t i=0; i <= fPIDParticles->GetNbinsX()+1; ++i)
569       if (fPIDParticles->GetBinContent(i) > 0)
570         printf("PDG = %d (%s): generated: %d, reconstructed: %d, ratio: %f\n", (Int_t) fPIDParticles->GetBinCenter(i), pdgDB->GetParticle((Int_t) fPIDParticles->GetBinCenter(i))->GetName(), (Int_t) fPIDParticles->GetBinContent(i), (Int_t) fPIDTracks->GetBinContent(i), ((fPIDTracks->GetBinContent(i) > 0) ? fPIDParticles->GetBinContent(i) / fPIDTracks->GetBinContent(i) : -1));
571
572     delete pdgDB;
573     pdgDB = 0;
574   }
575 }
576
577 Bool_t AlidNdEtaCorrectionTask::SignOK(TParticlePDG* particle)
578 {
579   // returns if a particle with this sign should be counted
580   // this is determined by the value of fSignMode, which should have the same sign
581   // as the charge
582   // if fSignMode is 0 all particles are counted
583
584   if (fSignMode == 0)
585     return kTRUE;
586
587   if (!particle)
588   {
589     printf("WARNING: not counting a particle that does not have a pdg particle\n");
590     return kFALSE;
591   }
592
593   Double_t charge = particle->Charge();
594
595   if (fSignMode > 0)
596     if (charge < 0)
597       return kFALSE;
598
599   if (fSignMode < 0)
600     if (charge > 0)
601       return kFALSE;
602
603   return kTRUE;
604 }
605