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