0b06b09175b01a8191e0b22a256324e273ae2452
[u/mrichter/AliRoot.git] / PWG0 / AliPWG0Helper.cxx
1 /* $Id$ */
2
3 #include <AliPWG0Helper.h>
4
5 #include <TParticle.h>
6 #include <TParticlePDG.h>
7 #include <TH1.h>
8 #include <TH2.h>
9 #include <TH3.h>
10 #include <TList.h>
11 #include <TTree.h>
12 #include <TBranch.h>
13 #include <TLeaf.h>
14
15 #include <AliHeader.h>
16 #include <AliStack.h>
17 #include <AliLog.h>
18
19 #include <AliLog.h>
20 #include <AliESD.h>
21 #include <AliESDEvent.h>
22 #include <AliESDVertex.h>
23
24 #include <AliGenEventHeader.h>
25 #include <AliGenPythiaEventHeader.h>
26 #include <AliGenCocktailEventHeader.h>
27
28 //____________________________________________________________________
29 ClassImp(AliPWG0Helper)
30
31 //____________________________________________________________________
32 Bool_t AliPWG0Helper::IsEventTriggered(const AliESD* aEsd, Trigger trigger)
33 {
34   // see function with ULong64_t argument
35
36   ULong64_t triggerMask = aEsd->GetTriggerMask();
37   return IsEventTriggered(triggerMask, trigger);
38 }
39
40 //____________________________________________________________________
41 Bool_t AliPWG0Helper::IsEventTriggered(ULong64_t triggerMask, Trigger trigger)
42 {
43   // check if the event was triggered
44   //
45   // this function needs the branch fTriggerMask
46   
47
48   // definitions from p-p.cfg
49   ULong64_t spdFO = (1 << 14);
50   ULong64_t v0left = (1 << 11);
51   ULong64_t v0right = (1 << 12);
52
53   switch (trigger)
54   {
55     case kMB1:
56     {
57       if (triggerMask & spdFO || ((triggerMask & v0left) || (triggerMask & v0right)))
58         return kTRUE;
59       break;
60     }
61     case kMB2:
62     {
63       if (triggerMask & spdFO && ((triggerMask & v0left) || (triggerMask & v0right)))
64         return kTRUE;
65       break;
66     }
67   }
68
69   return kFALSE;
70 }
71
72 //____________________________________________________________________
73 Bool_t AliPWG0Helper::IsVertexReconstructed(const AliESD* aEsd)
74 {
75   // see function with AliESDVertex argument
76
77   const AliESDVertex* vtxESD = aEsd->GetVertex();
78   return IsVertexReconstructed(vtxESD);
79 }
80
81 //____________________________________________________________________
82 Bool_t AliPWG0Helper::IsVertexReconstructed(const AliESDVertex* vtxESD)
83 {
84   // checks if the vertex is reasonable
85   //
86   // this function needs the branches fSPDVertex*
87
88   if (!vtxESD)
89     return kFALSE;
90
91   // the vertex should be reconstructed
92   if (strcmp(vtxESD->GetName(), "default")==0)
93     return kFALSE;
94
95   // check Ncontributors
96   if (vtxESD->GetNContributors() <= 0)
97     return kFALSE;
98
99   Double_t vtx_res[3];
100   vtx_res[0] = vtxESD->GetXRes();
101   vtx_res[1] = vtxESD->GetYRes();
102   vtx_res[2] = vtxESD->GetZRes();
103
104   if (vtx_res[2]==0 || vtx_res[2]>0.1)
105     return kFALSE;
106
107   return kTRUE;
108 }
109
110 //____________________________________________________________________
111 const AliESDVertex* AliPWG0Helper::GetVertex(const AliESDEvent* aEsd, AnalysisMode analysisMode)
112 {
113   // Get the vertex from the ESD and returns it if the vertex is valid
114   //
115   // Second argument decides which vertex is used (this selects
116   // also the quality criteria that are applied)
117
118   const AliESDVertex* vertex = 0;
119   Float_t requiredZResolution = -1;
120   if (analysisMode == kSPD || analysisMode == kTPCITS)
121   {
122     vertex = aEsd->GetVertex();
123     requiredZResolution = 0.1;
124   }
125   else if (analysisMode == kTPC) 
126   {
127     vertex = aEsd->GetPrimaryVertex();
128     requiredZResolution = 0.6;
129   }
130   else
131     Printf("AliPWG0Helper::GetVertex: ERROR: Invalid second argument %d", analysisMode);
132
133   if (!vertex)
134     return 0;
135
136   // check Ncontributors
137   if (vertex->GetNContributors() <= 0)
138     return 0;
139
140   // check resolution
141   Double_t zRes = vertex->GetZRes();
142
143   if (zRes == 0 || zRes > requiredZResolution)
144     return 0;
145
146   return vertex;
147 }
148
149 //____________________________________________________________________
150 Bool_t AliPWG0Helper::IsPrimaryCharged(TParticle* aParticle, Int_t aTotalPrimaries, Bool_t adebug)
151 {
152   //
153   // this function checks if a particle from the event generator (i.e. among the nPrim particles in the stack)
154   // shall be counted as a primary particle
155   //
156   // This function or a equivalent should be available in some common place of AliRoot
157   //
158   // WARNING: Call this function only for particles that are among the particles from the event generator!
159   // --> stack->Particle(id) with id < stack->GetNprimary()
160
161   // if the particle has a daughter primary, we do not want to count it
162   if (aParticle->GetFirstDaughter() != -1 && aParticle->GetFirstDaughter() < aTotalPrimaries)
163   {
164     if (adebug)
165       printf("Dropping particle because it has a daughter among the primaries.\n");
166     return kFALSE;
167   }
168
169   Int_t pdgCode = TMath::Abs(aParticle->GetPdgCode());
170
171   // skip quarks and gluon
172   if (pdgCode <= 10 || pdgCode == 21)
173   {
174     if (adebug)
175       printf("Dropping particle because it is a quark or gluon.\n");
176     return kFALSE;
177   }
178
179   if (strcmp(aParticle->GetName(),"XXX") == 0)
180   {
181     if (adebug)
182       printf("WARNING: There is a particle named XXX.\n");
183     return kFALSE;
184   }
185
186   TParticlePDG* pdgPart = aParticle->GetPDG();
187
188   if (strcmp(pdgPart->ParticleClass(),"Unknown") == 0)
189   {
190     if (adebug)
191       printf("WARNING: There is a particle with an unknown particle class (pdg code %d).\n", pdgCode);
192     return kFALSE;
193   }
194
195   if (pdgPart->Charge() == 0)
196   {
197     if (adebug)
198       printf("Dropping particle because it is not charged.\n");
199     return kFALSE;
200   }
201
202   return kTRUE;
203 }
204
205 //____________________________________________________________________
206 void AliPWG0Helper::CreateProjections(TH3* hist, Bool_t save)
207 {
208   // create projections of 3d hists to all 2d combinations
209   // the histograms are not returned, just use them from memory or use this to create them in a file
210
211   TH1* proj = hist->Project3D("yx");
212   proj->SetXTitle(hist->GetXaxis()->GetTitle());
213   proj->SetYTitle(hist->GetYaxis()->GetTitle());
214   if (save)
215     proj->Write();
216
217   proj = hist->Project3D("zx");
218   proj->SetXTitle(hist->GetXaxis()->GetTitle());
219   proj->SetYTitle(hist->GetZaxis()->GetTitle());
220   if (save)
221     proj->Write();
222
223   proj = hist->Project3D("zy");
224   proj->SetXTitle(hist->GetYaxis()->GetTitle());
225   proj->SetYTitle(hist->GetZaxis()->GetTitle());
226   if (save)
227     proj->Write();
228 }
229
230 //____________________________________________________________________
231 void AliPWG0Helper::CreateDividedProjections(TH3* hist, TH3* hist2, const char* axis, Bool_t putErrors, Bool_t save)
232 {
233   // create projections of the 3d hists divides them
234   // axis decides to which plane, if axis is 0 to all planes
235   // the histograms are not returned, just use them from memory or use this to create them in a file
236
237   if (axis == 0)
238   {
239     CreateDividedProjections(hist, hist2, "yx", putErrors, save);
240     CreateDividedProjections(hist, hist2, "zx", putErrors, save);
241     CreateDividedProjections(hist, hist2, "zy", putErrors, save);
242
243     return;
244   }
245
246   TH1* proj = hist->Project3D(axis);
247
248   if (strlen(axis) == 2)
249   {
250     proj->SetYTitle(GetAxisTitle(hist, axis[0]));
251     proj->SetXTitle(GetAxisTitle(hist, axis[1]));
252   }
253   else if (strlen(axis) == 1)
254     proj->SetXTitle(GetAxisTitle(hist, axis[0]));
255
256   TH1* proj2 = hist2->Project3D(axis);
257   if (strlen(axis) == 2)
258   {
259     proj2->SetYTitle(GetAxisTitle(hist2, axis[0]));
260     proj2->SetXTitle(GetAxisTitle(hist2, axis[1]));
261   }
262   else if (strlen(axis) == 1)
263     proj2->SetXTitle(GetAxisTitle(hist2, axis[0]));
264
265   TH1* division = dynamic_cast<TH1*> (proj->Clone(Form("%s_div_%s", proj->GetName(), proj2->GetName())));
266   //printf("doing axis: %s, x axis has %d %d bins, min %f %f max %f %f\n", axis, division->GetNbinsX(), proj2->GetNbinsX(), division->GetXaxis()->GetBinLowEdge(1), proj2->GetXaxis()->GetBinLowEdge(1), division->GetXaxis()->GetBinUpEdge(division->GetNbinsX()), proj2->GetXaxis()->GetBinUpEdge(proj2->GetNbinsX()));
267   //printf("doing axis: %s, y axis has %d %d bins, min %f %f max %f %f\n", axis, division->GetNbinsY(), proj2->GetNbinsY(), division->GetYaxis()->GetBinLowEdge(1), proj2->GetYaxis()->GetBinLowEdge(1), division->GetYaxis()->GetBinUpEdge(division->GetNbinsY()), proj2->GetYaxis()->GetBinUpEdge(proj2->GetNbinsY()));
268   division->Divide(proj, proj2, 1, 1, "B");
269   division->SetTitle(Form("%s divided %s", proj->GetTitle(), proj2->GetTitle()));
270
271   if (putErrors)
272   {
273     division->Sumw2();
274     if (division->GetDimension() == 1)
275     {
276       Int_t nBins = division->GetNbinsX();
277       for (Int_t i = 1; i <= nBins; ++i)
278         if (proj2->GetBinContent(i) != 0)
279           division->SetBinError(i, TMath::Sqrt(proj->GetBinContent(i)) / proj2->GetBinContent(i));
280     }
281     else if (division->GetDimension() == 2)
282     {
283       Int_t nBinsX = division->GetNbinsX();
284       Int_t nBinsY = division->GetNbinsY();
285       for (Int_t i = 1; i <= nBinsX; ++i)
286         for (Int_t j = 1; j <= nBinsY; ++j)
287           if (proj2->GetBinContent(i, j) != 0)
288             division->SetBinError(i, j, TMath::Sqrt(proj->GetBinContent(i, j)) / proj2->GetBinContent(i, j));
289     }
290   }
291
292   if (save)
293   {
294     proj->Write();
295     proj2->Write();
296     division->Write();
297   }
298 }
299
300 //____________________________________________________________________
301 const char* AliPWG0Helper::GetAxisTitle(TH3* hist, const char axis)
302 {
303   // returns the title of the axis given in axis (x, y, z)
304
305   if (axis == 'x')
306     return hist->GetXaxis()->GetTitle();
307   else if (axis == 'y')
308     return hist->GetYaxis()->GetTitle();
309   else if (axis == 'z')
310     return hist->GetZaxis()->GetTitle();
311
312   return 0;
313 }
314
315 //____________________________________________________________________
316 Int_t AliPWG0Helper::GetPythiaEventProcessType(AliHeader* aHeader, Bool_t adebug) {
317   //
318   // get the process type of the event.
319   //
320
321   // can only read pythia headers, either directly or from cocktalil header
322   AliGenPythiaEventHeader* pythiaGenHeader = dynamic_cast<AliGenPythiaEventHeader*>(aHeader->GenEventHeader());
323
324   if (!pythiaGenHeader) {
325
326     AliGenCocktailEventHeader* genCocktailHeader = dynamic_cast<AliGenCocktailEventHeader*>(aHeader->GenEventHeader());
327     if (!genCocktailHeader) {
328       printf("AliPWG0Helper::GetProcessType : Unknown header type (not Pythia or Cocktail). \n");
329       return -1;
330     }
331
332     TList* headerList = genCocktailHeader->GetHeaders();
333     if (!headerList) {
334       return -1;
335     }
336
337     for (Int_t i=0; i<headerList->GetEntries(); i++) {
338       pythiaGenHeader = dynamic_cast<AliGenPythiaEventHeader*>(headerList->At(i));
339       if (pythiaGenHeader)
340         break;
341     }
342
343     if (!pythiaGenHeader) {
344       printf("AliPWG0Helper::GetProcessType : Could not find Pythia header. \n");
345       return -1;
346     }
347   }
348
349   if (adebug) {
350     printf("AliPWG0Helper::GetProcessType : Pythia process type found: %d \n",pythiaGenHeader->ProcessType());
351   }
352
353   return pythiaGenHeader->ProcessType();
354 }
355
356 //____________________________________________________________________
357 TParticle* AliPWG0Helper::FindPrimaryMother(AliStack* stack, Int_t label)
358 {
359   //
360   // Finds the first mother among the primary particles of the particle identified by <label>,
361   // i.e. the primary that "caused" this particle
362   //
363
364   Int_t motherLabel = FindPrimaryMotherLabel(stack, label);
365   if (motherLabel < 0)
366     return 0;
367
368   return stack->Particle(motherLabel);
369 }
370
371 //____________________________________________________________________
372 Int_t AliPWG0Helper::FindPrimaryMotherLabel(AliStack* stack, Int_t label)
373 {
374   //
375   // Finds the first mother among the primary particles of the particle identified by <label>,
376   // i.e. the primary that "caused" this particle
377   //
378   // returns its label
379   //
380
381   Int_t nPrim  = stack->GetNprimary();
382
383   while (label >= nPrim)
384   {
385     //printf("Particle %d (pdg %d) is not a primary. Let's check its mother %d\n", label, mother->GetPdgCode(), mother->GetMother(0));
386
387     TParticle* particle = stack->Particle(label);
388     if (!particle)
389     {
390       AliDebugGeneral("FindPrimaryMother", AliLog::kError, Form("UNEXPECTED: particle with label %d not found in stack.", label));
391       return -1;
392     }
393  
394     // find mother
395     if (particle->GetMother(0) < 0)
396     {
397       AliDebugGeneral("FindPrimaryMother", AliLog::kError, Form("UNEXPECTED: Could not find mother of secondary particle %d.", label));
398       return -1;
399     }
400
401     label = particle->GetMother(0);
402   }
403
404   return label;
405 }
406
407 void AliPWG0Helper::SetBranchStatusRecursive(TTree* tree, char *bname, Bool_t status, Bool_t debug)
408 {
409   // Function  to switch on/off all data members of a top level branch
410   // this is needed for branches without a trailing dot ".", for those
411   // the root functionality with regular expressions does not work.
412   // Usage e.g.
413   // chain->SetBranchStatus("*", 0); 
414   // SetBranchStatusRecursive(chain,"SPDVertex",1);
415   // You need to give the full name of the top level branch zou want to access
416   //==========================================================================
417   // Author Christian.Klein-Boesing@cern.ch
418
419   if (!tree)
420     return;
421
422   TBranch *br = tree->GetBranch(bname);
423   if(!br) {
424     Printf("AliPWG0Helper::SetBranchStatusRecursive: Branch %s not found", bname);
425   }
426
427   TObjArray *leaves = tree->GetListOfLeaves();
428   Int_t nleaves = leaves->GetEntries();
429   TLeaf *leaf = 0;
430   TBranch *branch = 0;
431   TBranch *mother = 0;
432   for (Int_t i=0;i<nleaves;i++)  {
433     // the matched entry e.g. SPDVertex is its own Mother
434     leaf = (TLeaf*)leaves->UncheckedAt(i);
435     branch = (TBranch*)leaf->GetBranch();
436     mother = branch->GetMother();
437     if (mother==br){
438       if (debug)
439         Printf(">>>> AliPWG0Helper::SetBranchStatusRecursive: Setting status %s %s to %d", mother->GetName(), leaf->GetName(), status);
440       if (status) branch->ResetBit(kDoNotProcess);
441       else        branch->SetBit(kDoNotProcess);
442     }
443   }
444 }
445
446 //____________________________________________________________________
447 void AliPWG0Helper::NormalizeToBinWidth(TH1* hist)
448 {
449   //
450   // normalizes a 1-d histogram to its bin width
451   //
452
453   for (Int_t i=1; i<=hist->GetNbinsX(); ++i)
454   {
455     hist->SetBinContent(i, hist->GetBinContent(i) / hist->GetBinWidth(i));
456     hist->SetBinError(i, hist->GetBinError(i) / hist->GetBinWidth(i));
457   }
458 }
459
460 //____________________________________________________________________
461 void AliPWG0Helper::NormalizeToBinWidth(TH2* hist)
462 {
463   //
464   // normalizes a 2-d histogram to its bin width (x width * y width)
465   //
466
467   for (Int_t i=1; i<=hist->GetNbinsX(); ++i)
468     for (Int_t j=1; j<=hist->GetNbinsY(); ++j)
469     {
470       Double_t factor = hist->GetXaxis()->GetBinWidth(i) * hist->GetYaxis()->GetBinWidth(j);
471       hist->SetBinContent(i, j, hist->GetBinContent(i, j) / factor);
472       hist->SetBinError(i, j, hist->GetBinError(i, j) / factor);
473     }
474 }