removal of a few debug lines
[u/mrichter/AliRoot.git] / PWG0 / multiplicity / AliMultiplicityTask.cxx
1 /* $Id$ */
2
3 #include "AliMultiplicityTask.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 <TH2F.h>
12 #include <TH3F.h>
13 #include <TParticle.h>
14 #include <TRandom.h>
15 #include <TNtuple.h>
16 #include <TObjString.h>
17 #include <TF1.h>
18
19 #include <AliLog.h>
20 #include <AliESDVertex.h>
21 #include <AliESDEvent.h>
22 #include <AliStack.h>
23 #include <AliHeader.h>
24 #include <AliGenEventHeader.h>
25 #include <AliMultiplicity.h>
26 #include <AliAnalysisManager.h>
27 #include <AliMCEventHandler.h>
28 #include <AliMCEvent.h>
29 #include <AliESDInputHandler.h>
30
31 #include "esdTrackCuts/AliESDtrackCuts.h"
32 #include "AliPWG0Helper.h"
33 #include "dNdEta/AliMultiplicityCorrection.h"
34 #include "AliCorrection.h"
35 #include "AliCorrectionMatrix3D.h"
36
37 //#define TPCMEASUREMENT
38 #define ITSMEASUREMENT
39
40 ClassImp(AliMultiplicityTask)
41
42 AliMultiplicityTask::AliMultiplicityTask(const char* opt) :
43   AliAnalysisTask("AliMultiplicityTask", ""),
44   fESD(0),
45   fOption(opt),
46   fMultiplicity(0),
47   fEsdTrackCuts(0),
48   fSystSkipParticles(kFALSE),
49   fSelectProcessType(0),
50   fParticleSpecies(0),
51   fPtSpectrum(0),
52   fOutput(0)
53 {
54   //
55   // Constructor. Initialization of pointers
56   //
57
58   for (Int_t i = 0; i<4; i++)
59     fParticleCorrection[i] = 0;
60
61   // Define input and output slots here
62   DefineInput(0, TChain::Class());
63   DefineOutput(0, TList::Class());
64 }
65
66 AliMultiplicityTask::~AliMultiplicityTask()
67 {
68   //
69   // Destructor
70   //
71
72   // histograms are in the output list and deleted when the output
73   // list is deleted by the TSelector dtor
74
75   if (fOutput) {
76     delete fOutput;
77     fOutput = 0;
78   }
79 }
80
81 //________________________________________________________________________
82 void AliMultiplicityTask::ConnectInputData(Option_t *)
83 {
84   // Connect ESD
85   // Called once
86
87   TTree* tree = dynamic_cast<TTree*> (GetInputData(0));
88   if (!tree) {
89     Printf("ERROR: Could not read tree from input slot 0");
90   } else {
91     // Disable all branches and enable only the needed ones
92     tree->SetBranchStatus("*", 0);
93     tree->SetBranchStatus("fTriggerMask", 1);
94     tree->SetBranchStatus("fSPDVertex*", 1);
95
96     #ifdef ITSMEASUREMENT
97       tree->SetBranchStatus("fSPDMult*", 1);
98     #endif
99
100     #ifdef TPCMEASUREMENT
101       AliESDtrackCuts::EnableNeededBranches(tree);
102       tree->SetBranchStatus("fTracks.fLabel", 1);
103     #endif
104
105     tree->SetCacheSize(0);
106
107     AliESDInputHandler *esdH = dynamic_cast<AliESDInputHandler*> (AliAnalysisManager::GetAnalysisManager()->GetInputEventHandler());
108
109     if (!esdH) {
110       Printf("ERROR: Could not get ESDInputHandler");
111     } else
112       fESD = esdH->GetEvent();
113   }
114 }
115
116 void AliMultiplicityTask::CreateOutputObjects()
117 {
118   // create result objects and add to output list
119
120   fOutput = new TList;
121   fOutput->SetOwner();
122
123   fMultiplicity = new AliMultiplicityCorrection("Multiplicity", "Multiplicity");
124   fOutput->Add(fMultiplicity);
125
126   if (fOption.Contains("skip-particles"))
127   {
128     fSystSkipParticles = kTRUE;
129     AliInfo("WARNING: Systematic study enabled. Particles will be skipped.");
130   }
131
132   if (fOption.Contains("particle-efficiency"))
133     for (Int_t i = 0; i<4; i++)
134     {
135       fParticleCorrection[i] = new AliCorrection(Form("correction_%d", i), Form("correction_%d", i));
136       fOutput->Add(fParticleCorrection[i]);
137     }
138
139   if (fOption.Contains("only-process-type-nd"))
140     fSelectProcessType = 1;
141
142   if (fOption.Contains("only-process-type-sd"))
143     fSelectProcessType = 2;
144
145   if (fOption.Contains("only-process-type-dd"))
146     fSelectProcessType = 3;
147
148   if (fSelectProcessType != 0)
149     AliInfo(Form("WARNING: Systematic study enabled. Only considering process type %d", fSelectProcessType));
150
151   if (fOption.Contains("pt-spectrum-hist"))
152   {
153     TFile* file = TFile::Open("ptspectrum_fit.root");
154     if (file)
155     {
156       TString subStr(fOption(fOption.Index("pt-spectrum")+17, 3));
157       TString histName(Form("ptspectrum_%s", subStr.Data()));
158       AliInfo(Form("Pt-Spectrum modification. Using %s.", histName.Data()));
159       fPtSpectrum = (TH1*) file->Get(histName);
160       if (!fPtSpectrum)
161         AliError("Histogram not found");
162     }
163     else
164       AliError("Could not open ptspectrum_fit.root. Pt Spectrum study could not be enabled.");
165
166     if (fPtSpectrum)
167       AliInfo("WARNING: Systematic study enabled. Pt spectrum will be modified");
168   }
169
170   if (fOption.Contains("pt-spectrum-func"))
171   {
172     if (fPtSpectrum)
173     {
174       Printf("Using function from input list for systematic p_t study");
175     }
176     else
177     {
178       fPtSpectrum = new TH1F("fPtSpectrum", "fPtSpectrum", 1, 0, 100);
179       fPtSpectrum->SetBinContent(1, 1);
180     }
181
182     if (fPtSpectrum)
183       AliInfo("WARNING: Systematic study enabled. Pt spectrum will be modified");
184   }
185
186   if (fOption.Contains("particle-species")) {
187     fParticleSpecies = new TNtuple("fParticleSpecies", "fParticleSpecies", "vtx:Pi_True:K_True:p_True:o_True:Pi_Rec:K_Rec:p_Rec:o_Rec:nolabel:doublePrim:doubleCount");
188     fOutput->Add(fParticleSpecies);
189   }
190
191   // TODO set seed for random generator
192 }
193
194 void AliMultiplicityTask::Exec(Option_t*)
195 {
196   // process the event
197
198   // Check prerequisites
199   if (!fESD)
200   {
201     AliDebug(AliLog::kError, "ESD branch not available");
202     return;
203   }
204
205   if (!fEsdTrackCuts)
206   {
207     AliDebug(AliLog::kError, "fESDTrackCuts not available");
208     return;
209   }
210
211   AliAnalysisManager* mgr = AliAnalysisManager::GetAnalysisManager();
212   if (!mgr) {
213      Printf("ERROR: Could not retrieve manager");
214      return;
215   }
216
217   AliMCEventHandler* eventHandler = dynamic_cast<AliMCEventHandler*> (mgr->GetMCtruthEventHandler());
218   if (!eventHandler) {
219     Printf("ERROR: Could not retrieve MC event handler");
220     return;
221   }
222
223   AliMCEvent* mcEvent = eventHandler->MCEvent();
224   if (!mcEvent) {
225      Printf("ERROR: Could not retrieve MC event");
226      return;
227   }
228
229   AliStack* stack = mcEvent->Stack();
230   if (!stack)
231   {
232     AliDebug(AliLog::kError, "Stack not available");
233     return;
234   }
235
236   AliHeader* header = mcEvent->Header();
237   if (!header)
238   {
239     AliDebug(AliLog::kError, "Header not available");
240     return;
241   }
242
243   // post the output already here, because there are several exit possibilities below
244   PostData(0, fOutput);
245
246   if (fSelectProcessType > 0)
247   {
248     // getting process information; NB: this only works for Pythia
249     Int_t processtype = AliPWG0Helper::GetPythiaEventProcessType(header);
250
251     Bool_t processEvent = kFALSE;
252
253     // non diffractive
254     if (fSelectProcessType == 1 && processtype !=92 && processtype !=93 && processtype != 94)
255       processEvent = kTRUE;
256
257     // single diffractive
258     if (fSelectProcessType == 2 && (processtype == 92 || processtype == 93))
259       processEvent = kTRUE;
260
261     // double diffractive
262     if (fSelectProcessType == 3 && processtype == 94)
263       processEvent = kTRUE;
264
265     if (!processEvent)
266     {
267       AliDebug(AliLog::kDebug, Form("Skipping this event, because it is not of the requested process type (%d)", processtype));
268       return;
269     }
270   }
271
272   Bool_t eventTriggered = AliPWG0Helper::IsEventTriggered(fESD->GetTriggerMask());
273   eventTriggered = kTRUE; // no generated trigger for the simulated events
274
275   Bool_t eventVertex = AliPWG0Helper::IsVertexReconstructed(fESD->GetVertex());
276
277   // get the ESD vertex
278   const AliESDVertex* vtxESD = fESD->GetVertex();
279   Double_t vtx[3];
280   vtxESD->GetXYZ(vtx);
281
282   // get the MC vertex
283   AliGenEventHeader* genHeader = header->GenEventHeader();
284   TArrayF vtxMC(3);
285   genHeader->PrimaryVertex(vtxMC);
286
287   // get tracklets
288   const AliMultiplicity* mult = fESD->GetMultiplicity();
289   if (!mult)
290   {
291     AliDebug(AliLog::kError, "AliMultiplicity not available");
292     return;
293   }
294
295   //const Float_t kPtCut = 0.3;
296
297   // get number of tracks from MC
298   // loop over mc particles
299   Int_t nPrim  = stack->GetNprimary();
300   Int_t nMCPart = stack->GetNtrack();
301
302   // tracks in different eta ranges
303   Int_t nMCTracks05 = 0;
304   Int_t nMCTracks09 = 0;
305   Int_t nMCTracks15 = 0;
306   Int_t nMCTracks20 = 0;
307   Int_t nMCTracksAll = 0;
308
309   // tracks per particle species, in |eta| < 2 (systematic study)
310   Int_t nMCTracksSpecies[4]; // (pi, K, p, other)
311   for (Int_t i = 0; i<4; ++i)
312     nMCTracksSpecies[i] = 0;
313   // eta range for nMCTracksSpecies, nESDTracksSpecies
314 #ifdef ITSMEASUREMENT
315   const Float_t etaRange = 2.0;
316 #else
317   const Float_t etaRange = 0.9;
318 #endif
319
320   for (Int_t iMc = 0; iMc < nPrim; ++iMc)
321   {
322     AliDebug(AliLog::kDebug+1, Form("MC Loop: Processing particle %d.", iMc));
323
324     TParticle* particle = stack->Particle(iMc);
325
326     if (!particle)
327     {
328       AliDebug(AliLog::kError, Form("UNEXPECTED: particle with label %d not found in stack (mc loop).", iMc));
329       continue;
330     }
331
332     Bool_t debug = kFALSE;
333     if (AliPWG0Helper::IsPrimaryCharged(particle, nPrim, debug) == kFALSE)
334     {
335       //printf("%d) DROPPED ", iMC);
336       //particle->Print();
337       continue;
338     }
339
340     //printf("%d) OK      ", iMC);
341     //particle->Print();
342
343     //if (particle->Pt() < kPtCut)
344     //  continue;
345
346     Int_t particleWeight = 1;
347
348     Float_t pt = particle->Pt();
349
350     // in case of systematic study, weight according to the change of the pt spectrum
351     // it cannot be just multiplied because we cannot count "half of a particle"
352     // instead a random generator decides if the particle is counted twice (if value > 1)
353     // or not (if value < 0)
354     if (fPtSpectrum)
355     {
356       Int_t bin = fPtSpectrum->FindBin(pt);
357       if (bin > 0 && bin <= fPtSpectrum->GetNbinsX())
358       {
359         Float_t factor = fPtSpectrum->GetBinContent(bin);
360         if (factor > 0)
361         {
362           Float_t random = gRandom->Uniform();
363           if (factor > 1 && random < factor - 1)
364           {
365             particleWeight = 2;
366           }
367           else if (factor < 1 && random < 1 - factor)
368             particleWeight = 0;
369         }
370       }
371     }
372
373     //Printf("MC weight is: %d", particleWeight);
374
375     if (TMath::Abs(particle->Eta()) < 0.5)
376       nMCTracks05 += particleWeight;
377
378     if (TMath::Abs(particle->Eta()) < 0.9)
379       nMCTracks09 += particleWeight;
380
381     if (TMath::Abs(particle->Eta()) < 1.5)
382       nMCTracks15 += particleWeight;
383
384     if (TMath::Abs(particle->Eta()) < 2.0)
385       nMCTracks20 += particleWeight;
386
387     nMCTracksAll += particleWeight;
388
389     if (fParticleCorrection[0] || fParticleSpecies)
390     {
391       Int_t id = -1;
392       switch (TMath::Abs(particle->GetPdgCode()))
393       {
394         case 211: id = 0; break;
395         case 321: id = 1; break;
396         case 2212: id = 2; break;
397         default: id = 3; break;
398       }
399       if (TMath::Abs(particle->Eta()) < etaRange)
400         nMCTracksSpecies[id]++;
401       if (fParticleCorrection[id])
402         fParticleCorrection[id]->GetTrackCorrection()->FillGene(vtxMC[2], particle->Eta(), particle->Pt());
403     }
404   }// end of mc particle
405
406   fMultiplicity->FillGenerated(vtxMC[2], eventTriggered, eventVertex, (Int_t) nMCTracks05, (Int_t) nMCTracks09, (Int_t) nMCTracks15, (Int_t) nMCTracks20, (Int_t) nMCTracksAll);
407
408   if (!eventTriggered || !eventVertex)
409     return;
410
411   Int_t nESDTracks05 = 0;
412   Int_t nESDTracks09 = 0;
413   Int_t nESDTracks15 = 0;
414   Int_t nESDTracks20 = 0;
415
416   // tracks per particle species, in |eta| < 2 (systematic study)
417   Int_t nESDTracksSpecies[7]; // (pi, K, p, other, nolabel, doublecount_prim, doublecount_all)
418   for (Int_t i = 0; i<7; ++i)
419     nESDTracksSpecies[i] = 0;
420
421   Bool_t* foundPrimaries = new Bool_t[nPrim];   // to prevent double counting
422   for (Int_t i=0; i<nPrim; i++)
423     foundPrimaries[i] = kFALSE;
424
425   Bool_t* foundTracks = new Bool_t[nMCPart];    // to prevent double counting
426   for (Int_t i=0; i<nMCPart; i++)
427     foundTracks[i] = kFALSE;
428
429 #ifdef ITSMEASUREMENT
430   // get multiplicity from ITS tracklets
431   for (Int_t i=0; i<mult->GetNumberOfTracklets(); ++i)
432   {
433     //printf("%d %f %f %f\n", i, mult->GetTheta(i), mult->GetPhi(i), mult->GetDeltaPhi(i));
434
435     // This removes non-tracklets in PDC06 data. Very bad solution. New solution is implemented for newer data. Keeping this for compatibility.
436     if (mult->GetDeltaPhi(i) < -1000)
437       continue;
438
439     // systematic study: 10% lower efficiency
440     if (fSystSkipParticles && (gRandom->Uniform() < 0.1))
441       continue;
442
443     Float_t theta = mult->GetTheta(i);
444     Float_t eta   = -TMath::Log(TMath::Tan(theta/2.));
445
446     // TODO define needed, because this only works with new AliRoot
447     Int_t label = mult->GetLabel(i);
448 #endif
449 #ifdef TPCMEASUREMENT
450   // get multiplicity from ESD tracks
451   TObjArray* list = fEsdTrackCuts->GetAcceptedTracks(fESD);
452   Int_t nGoodTracks = list->GetEntries();
453   // loop over esd tracks
454   for (Int_t i=0; i<nGoodTracks; i++)
455   {
456     AliESDtrack* esdTrack = dynamic_cast<AliESDtrack*> (list->At(i));
457     if (!esdTrack)
458     {
459       AliDebug(AliLog::kError, Form("ERROR: Could not retrieve track %d.", i));
460       continue;
461     }
462
463     Double_t p[3];
464     esdTrack->GetConstrainedPxPyPz(p); 
465     TVector3 vector(p);
466
467     Float_t theta = vector.Theta();
468     Float_t eta   = -TMath::Log(TMath::Tan(theta/2.));
469
470     Int_t label = esdTrack->GetLabel();
471
472     //Float_t pt = vector.Pt();
473     //if (pt < kPtCut)
474     //  continue;
475 #endif
476
477     Int_t particleWeight = 1;
478
479     // in case of systematic study, weight according to the change of the pt spectrum
480     if (fPtSpectrum)
481     {
482       TParticle* mother = 0;
483
484       // preserve label for later
485       Int_t labelCopy = label;
486       if (labelCopy >= 0)
487         labelCopy = AliPWG0Helper::FindPrimaryMotherLabel(stack, labelCopy);
488       if (labelCopy >= 0)
489         mother = stack->Particle(labelCopy);
490
491       // in case of pt study we do not count particles w/o label, because they cannot be scaled
492       if (!mother)
493         continue;
494
495       // it cannot be just multiplied because we cannot count "half of a particle"
496       // instead a random generator decides if the particle is counted twice (if value > 1) 
497       // or not (if value < 0)
498       Int_t bin = fPtSpectrum->FindBin(mother->Pt());
499       if (bin > 0 && bin <= fPtSpectrum->GetNbinsX())
500       {
501         Float_t factor = fPtSpectrum->GetBinContent(bin);
502         if (factor > 0)
503         {
504           Float_t random = gRandom->Uniform();
505           if (factor > 1 && random < factor - 1)
506           {
507             particleWeight = 2;
508           }
509           else if (factor < 1 && random < 1 - factor)
510             particleWeight = 0;
511         }
512       }
513     }
514
515     //Printf("ESD weight is: %d", particleWeight);
516
517     if (TMath::Abs(eta) < 0.5)
518       nESDTracks05 += particleWeight;
519
520     if (TMath::Abs(eta) < 0.9)
521       nESDTracks09 += particleWeight;
522
523     if (TMath::Abs(eta) < 1.5)
524       nESDTracks15 += particleWeight;
525
526     if (TMath::Abs(eta) < 2.0)
527       nESDTracks20 += particleWeight;
528
529
530     if (fParticleCorrection[0] || fParticleSpecies)
531     {
532       Int_t motherLabel = -1;
533       TParticle* mother = 0;
534
535       // find mother
536       if (label >= 0)
537         motherLabel = AliPWG0Helper::FindPrimaryMotherLabel(stack, label);
538       if (motherLabel >= 0)
539         mother = stack->Particle(motherLabel);
540
541       if (!mother)
542       {
543         // count tracks that did not have a label
544         if (TMath::Abs(eta) < etaRange)
545           nESDTracksSpecies[4]++;
546         continue;
547       }
548
549       // get particle type (pion, proton, kaon, other)
550       Int_t id = -1;
551       switch (TMath::Abs(mother->GetPdgCode()))
552       {
553         case 211: id = 0; break;
554         case 321: id = 1; break;
555         case 2212: id = 2; break;
556         default: id = 3; break;
557       }
558
559       // double counting is ok for particle ratio study
560       if (TMath::Abs(eta) < etaRange)
561         nESDTracksSpecies[id]++;
562
563       // double counting is not ok for efficiency study
564
565       // check if we already counted this particle, this way distinguishes double counted particles (bug/artefact in tracking) or double counted primaries due to secondaries (physics)
566       if (foundTracks[label])
567       {
568         if (TMath::Abs(eta) < etaRange)
569           nESDTracksSpecies[6]++;
570         continue;
571       }
572       foundTracks[label] = kTRUE;
573
574       // particle (primary) already counted?
575       if (foundPrimaries[motherLabel])
576       {
577         if (TMath::Abs(eta) < etaRange)
578           nESDTracksSpecies[5]++;
579         continue;
580       }
581       foundPrimaries[motherLabel] = kTRUE;
582
583       if (fParticleCorrection[id])
584         fParticleCorrection[id]->GetTrackCorrection()->FillMeas(vtxMC[2], mother->Eta(), mother->Pt());
585     }
586   }
587
588   delete[] foundTracks;
589   delete[] foundPrimaries;
590
591   if ((Int_t) nMCTracks20 == 0 && nESDTracks20 > 3)
592     printf("WARNING: Event has %d generated and %d reconstructed...\n", nMCTracks20, nESDTracks20);
593
594   //printf("%.2f generated and %.2f reconstructed\n", nMCTracks20, nESDTracks20);
595
596   //Printf("%f %f %f %f %f %f %f %f %f %f", vtxMC[2], nMCTracks05, nMCTracks09, nMCTracks15, nMCTracks20, nMCTracksAll, nESDTracks05, nESDTracks09, nESDTracks15, nESDTracks20);
597
598   //Printf("%f %f %f %f %f %f %f %f %f %f", vtxMC[2], nMCTracks05, nMCTracks09, nMCTracks15, nMCTracks20, nMCTracksAll, nESDTracks05, nESDTracks09, nESDTracks15, nESDTracks20);
599
600   fMultiplicity->FillMeasured(vtx[2], nESDTracks05, nESDTracks09, nESDTracks15, nESDTracks20);
601
602   // fill response matrix using vtxMC (best guess)
603   fMultiplicity->FillCorrection(vtxMC[2],  nMCTracks05,  nMCTracks09,  nMCTracks15,  nMCTracks20,  nMCTracksAll,  nESDTracks05,  nESDTracks09,  nESDTracks15,  nESDTracks20);
604
605   if (fParticleSpecies)
606     fParticleSpecies->Fill(vtxMC[2], nMCTracksSpecies[0], nMCTracksSpecies[1], nMCTracksSpecies[2], nMCTracksSpecies[3], nESDTracksSpecies[0], nESDTracksSpecies[1], nESDTracksSpecies[2], nESDTracksSpecies[3], nESDTracksSpecies[4], nESDTracksSpecies[5], nESDTracksSpecies[6]);
607
608   //Printf("%d = %d %d %d %d %d %d %d", (Int_t) nESDTracks20, nESDTracksSpecies[0], nESDTracksSpecies[1], nESDTracksSpecies[2], nESDTracksSpecies[3], nESDTracksSpecies[4], nESDTracksSpecies[5], nESDTracksSpecies[6]);
609 }
610
611 void AliMultiplicityTask::Terminate(Option_t *)
612 {
613   // The Terminate() function is the last function to be called during
614   // a query. It always runs on the client, it can be used to present
615   // the results graphically or save the results to file.
616
617   fOutput = dynamic_cast<TList*> (GetOutputData(0));
618   if (!fOutput) {
619     Printf("ERROR: fOutput not available");
620     return;
621   }
622
623   fMultiplicity = dynamic_cast<AliMultiplicityCorrection*> (fOutput->FindObject("Multiplicity"));
624   for (Int_t i = 0; i < 4; ++i)
625     fParticleCorrection[i] = dynamic_cast<AliCorrection*> (fOutput->FindObject(Form("correction_%d", i)));
626   fParticleSpecies = dynamic_cast<TNtuple*> (fOutput->FindObject("fParticleSpecies"));
627
628   if (!fMultiplicity)
629   {
630     AliDebug(AliLog::kError, Form("ERROR: Histograms not available %p", (void*) fMultiplicity));
631     return;
632   }
633
634   TFile* file = TFile::Open("multiplicityMC.root", "RECREATE");
635
636   fMultiplicity->SaveHistograms();
637   for (Int_t i = 0; i < 4; ++i)
638     if (fParticleCorrection[i])
639       fParticleCorrection[i]->SaveHistograms();
640   if (fParticleSpecies)
641     fParticleSpecies->Write();
642
643   TObjString option(fOption);
644   option.Write();
645
646   file->Close();
647
648   Printf("Writting result to multiplicityMC.root");
649 }