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