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