Refactoring of dN/deta task to common base class
[u/mrichter/AliRoot.git] / PWG2 / FORWARD / analysis2 / AliBasedNdetaTask.cxx
1 //====================================================================
2 #include "AliBasedNdetaTask.h"
3 #include <TMath.h>
4 #include <TH2D.h>
5 #include <TH1D.h>
6 #include <THStack.h>
7 #include <TList.h>
8 #include <AliAnalysisManager.h>
9 #include <AliAODEvent.h>
10 #include <AliAODHandler.h>
11 #include <AliAODInputHandler.h>
12 #include "AliForwardUtil.h"
13 #include "AliAODForwardMult.h"
14
15 //____________________________________________________________________
16 AliBasedNdetaTask::AliBasedNdetaTask()
17   : AliAnalysisTaskSE(), 
18     fSum(0),            //  Sum of histograms 
19     fSumMC(0),          //  Sum of MC histograms (if any)
20     fSums(0),           // Container of sums 
21     fOutput(0),         // Container of outputs 
22     fTriggers(0),       // Histogram of triggers 
23     fVtxMin(0),         // Minimum v_z
24     fVtxMax(0),         // Maximum v_z
25     fTriggerMask(0),    // Trigger mask 
26     fRebin(0),          // Rebinning factor 
27     fCutEdges(false), 
28     fSymmetrice(true),
29     fCorrEmpty(true)
30 {}
31
32 //____________________________________________________________________
33 AliBasedNdetaTask::AliBasedNdetaTask(const char* name)
34   : AliAnalysisTaskSE(name), 
35     fSum(0),            // Sum of histograms 
36     fSumMC(0),          // Sum of MC histograms (if any)
37     fSums(0),           // Container of sums 
38     fOutput(0),         // Container of outputs 
39     fTriggers(0),       // Histogram of triggers 
40     fVtxMin(-10),       // Minimum v_z
41     fVtxMax(10),        // Maximum v_z
42     fTriggerMask(AliAODForwardMult::kInel), 
43     fRebin(5),          // Rebinning factor 
44     fCutEdges(false), 
45     fSymmetrice(true),
46     fCorrEmpty(true)
47 {
48   // Output slot #1 writes into a TH1 container
49   DefineOutput(1, TList::Class()); 
50   DefineOutput(2, TList::Class()); 
51 }
52
53 //____________________________________________________________________
54 AliBasedNdetaTask::AliBasedNdetaTask(const AliBasedNdetaTask& o)
55   : AliAnalysisTaskSE(o), 
56     fSum(o.fSum),       // TH2D* -  Sum of histograms 
57     fSumMC(o.fSumMC),// TH2D* -  Sum of MC histograms (if any)
58     fSums(o.fSums),             // TList* - Container of sums 
59     fOutput(o.fOutput),         // TList* - Container of outputs 
60     fTriggers(o.fTriggers),     // TH1D* - Histogram of triggers 
61     fVtxMin(o.fVtxMin),         // Double_t - Minimum v_z
62     fVtxMax(o.fVtxMax),         // Double_t - Maximum v_z
63     fTriggerMask(o.fTriggerMask),// Int_t - Trigger mask 
64     fRebin(o.fRebin),           // Int_t - Rebinning factor 
65     fCutEdges(o.fCutEdges),     // Bool_t - Whether to cut edges when rebinning
66     fSymmetrice(true),
67     fCorrEmpty(true)
68 {}
69
70 //____________________________________________________________________
71 AliBasedNdetaTask::~AliBasedNdetaTask()
72 {
73   if (fSums) { 
74     fSums->Delete();
75     delete fSums;
76     fSums = 0;
77   }
78   if (fOutputs) { 
79     fOutputs->Delete();
80     delete fOutputs;
81     fOutputs = 0;
82   }
83 }
84
85 //________________________________________________________________________
86 void 
87 AliBasedNdetaTask::SetTriggerMask(const char* mask)
88 {
89   UShort_t    trgMask = 0;
90   TString     trgs(mask);
91   trgs.ToUpper();
92   TObjString* trg;
93   TIter       next(trgs.Tokenize(" ,|"));
94   while ((trg = static_cast<TObjString*>(next()))) { 
95     TString s(trg->GetString());
96     if      (s.IsNull()) continue;
97     if      (s.CompareTo("INEL")  == 0) trgMask = AliAODForwardMult::kInel;
98     else if (s.CompareTo("INEL>0")== 0) trgMask = AliAODForwardMult::kInelGt0;
99     else if (s.CompareTo("NSD")   == 0) trgMask = AliAODForwardMult::kNSD;
100     else 
101       Warning("SetTriggerMask", "Unknown trigger %s", s.Data());
102   }
103   if (trgMask == 0) trgMask = 1;
104   SetTriggerMask(trgMask);
105 }
106
107 //________________________________________________________________________
108 void 
109 AliBasedNdetaTask::UserCreateOutputObjects()
110 {
111   // Create histograms
112   // Called once (on the worker node)
113
114   fOutput = new TList;
115   fOutput->SetName(Form("%s_result", GetName()));
116   fOutput->SetOwner();
117
118   fSums = new TList;
119   fSums->SetName(Form("%s_sums", GetName()));
120   fSums->SetOwner();
121
122
123   fTriggers = new TH1D("triggers", "Number of triggers", 
124                        kAccepted, 1, kAccepted);
125   fTriggers->SetYTitle("# of events");
126   fTriggers->GetXaxis()->SetBinLabel(kAll,         "All events");
127   fTriggers->GetXaxis()->SetBinLabel(kB,           "w/B trigger");
128   fTriggers->GetXaxis()->SetBinLabel(kA,           "w/A trigger");
129   fTriggers->GetXaxis()->SetBinLabel(kC,           "w/C trigger");
130   fTriggers->GetXaxis()->SetBinLabel(kE,           "w/E trigger");
131   fTriggers->GetXaxis()->SetBinLabel(kMB,          "w/Collision trigger");
132   fTriggers->GetXaxis()->SetBinLabel(kWithVertex,  "w/Vertex");
133   fTriggers->GetXaxis()->SetBinLabel(kWithTrigger, "w/Selected trigger");
134   fTriggers->GetXaxis()->SetBinLabel(kAccepted,    "Accepted by cut");
135   fTriggers->GetXaxis()->SetNdivisions(kAccepted, false);
136   fTriggers->SetFillColor(kRed+1);
137   fTriggers->SetFillStyle(3001);
138   fTriggers->SetStats(0);
139   fSums->Add(fTriggers);
140
141   // Check that we have an AOD input handler 
142   AliAnalysisManager* am = AliAnalysisManager::GetAnalysisManager();
143   AliAODInputHandler* ah = 
144     dynamic_cast<AliAODInputHandler*>(am->GetInputEventHandler());
145   if (!ah) AliFatal("No AOD input handler set in analysis manager");
146
147   // Post data for ALL output slots >0 here, to get at least an empty histogram
148   PostData(1, fSums); 
149 }
150
151 //____________________________________________________________________
152 TH2D*
153 AliBasedNdetaTask::CloneHist(TH2D* in, const char* name) 
154 {
155   if (!in) return 0;
156   TH2D* ret = static_cast<TH2D*>(in->Clone(name));
157   ret->SetDirectory(0);
158   ret->Sumw2();
159   fSums->Add(ret);
160
161   return ret;
162 }
163
164 //____________________________________________________________________
165 Bool_t
166 AliBasedNdetaTask::CheckEvent(AliAODEvent* aod)
167 {
168   TObject*           oForward   = aod->FindListObject("Forward");
169   if (!oForward) { 
170     AliWarning("No forward object found");
171     return false;
172   }
173   AliAODForwardMult* forward   = static_cast<AliAODForwardMult*>(oForward);
174   
175   // Count event 
176   fTriggers->AddBinContent(kAll);
177   if (forward->IsTriggerBits(AliAODForwardMult::kB)) 
178     fTriggers->AddBinContent(kB);
179   if (forward->IsTriggerBits(AliAODForwardMult::kA)) 
180     fTriggers->AddBinContent(kA);
181   if (forward->IsTriggerBits(AliAODForwardMult::kC)) 
182     fTriggers->AddBinContent(kC);
183   if (forward->IsTriggerBits(AliAODForwardMult::kE)) 
184     fTriggers->AddBinContent(kE);
185   if (forward->IsTriggerBits(AliAODForwardMult::kInel)) 
186     fTriggers->AddBinContent(kMB);
187
188   // Check if we have an event of interest. 
189   if (!forward->IsTriggerBits(fTriggerMask)) return false;
190   fTriggers->AddBinContent(kWithTrigger);
191
192   // Check that we have a valid vertex
193   if (!forward->HasIpZ()) return false;
194   fTriggers->AddBinContent(kWithVertex);
195
196   // Check that vertex is within cuts 
197   if (!forward->InRange(fVtxMin, fVtxMax)) return false;
198   fTriggers->AddBinContent(kAccepted);
199
200   return true;
201 }
202 //____________________________________________________________________
203 void 
204 AliBasedNdetaTask::UserExec(Option_t *) 
205 {
206   // Main loop
207   AliAODEvent* aod = dynamic_cast<AliAODEvent*>(InputEvent());
208   if (!aod) {
209     AliError("Cannot get the AOD event");
210     return;
211   }  
212
213   // Get the histogram(s) 
214   TH2D* data   = GetHistogram(aod);
215   TH2D* dataMC = GetHistogram(aod, true);
216
217   // We should have a base object at least 
218   if (!data) { 
219     AliWarning("No data object found in AOD");
220     return;
221   }
222   
223   // Create our sum histograms 
224   if (!fSum)             fSum   = CloneHist(data,  GetName());
225   if (!fSumMC && dataMC) fSumMC = CloneHist(dataMC,Form("%sMC", GetName()));
226
227   // Check the event 
228   if (!CheckEvent(aod)) return;
229
230   // Add contribution 
231   fSum->Add((data));
232   if (fSumMC) fSumMC->Add((dataMC));
233
234   PostData(1, fSums);
235 }
236
237 //________________________________________________________________________
238 void 
239 AliBasedNdetaTask::SetHistogramAttributes(TH1D* h, Int_t colour, Int_t marker,
240                                           const char* title, const char* ytitle)
241 {
242   h->SetTitle(title);
243   h->SetMarkerColor(colour);
244   h->SetMarkerStyle(marker);
245   h->SetMarkerSize(1);
246   h->SetFillStyle(0);
247   h->SetYTitle(ytitle);
248   h->GetXaxis()->SetTitleFont(132);
249   h->GetXaxis()->SetLabelFont(132);
250   h->GetXaxis()->SetNdivisions(10);
251   h->GetYaxis()->SetTitleFont(132);
252   h->GetYaxis()->SetLabelFont(132);
253   h->GetYaxis()->SetNdivisions(10);
254   h->GetYaxis()->SetDecimals();
255   h->SetStats(0);
256 }
257
258 //________________________________________________________________________
259 TH1D*
260 AliBasedNdetaTask::ProjectX(const TH2D* h, 
261                             const char* name,
262                             Int_t firstbin, 
263                             Int_t lastbin, 
264                             bool  corr,
265                             bool  error) const
266 {
267   if (!h) return 0;
268 #if USE_ROOT_PROJECT
269   return h->ProjectionX(name, firstbin, lastbin, (error ? "e" : ""));
270 #endif
271   
272   TAxis* xaxis = h->GetXaxis();
273   TAxis* yaxis = h->GetYaxis();
274   TH1D*  ret   = new TH1D(name, h->GetTitle(), xaxis->GetNbins(), 
275                           xaxis->GetXmin(), xaxis->GetXmax());
276   static_cast<const TAttLine*>(h)->Copy(*ret);
277   static_cast<const TAttFill*>(h)->Copy(*ret);
278   static_cast<const TAttMarker*>(h)->Copy(*ret);
279   ret->GetXaxis()->ImportAttributes(xaxis);
280
281   Int_t first = firstbin; 
282   Int_t last  = lastbin;
283   if (first < 0)                         first = 0;
284   else if (first >= yaxis->GetNbins()+1) first = yaxis->GetNbins();
285   if      (last  < 0)                    last  = yaxis->GetNbins();
286   else if (last  >  yaxis->GetNbins()+1) last  = yaxis->GetNbins();
287   if (last-first < 0) { 
288     AliWarning(Form("Nothing to project [%d,%d]", first, last));
289     return 0;
290     
291   }
292
293   // Loop over X bins 
294   // AliInfo(Form("Projecting bins [%d,%d] of %s", first, last, h->GetName()));
295   Int_t ybins = (last-first+1);
296   for (Int_t xbin = 0; xbin <= xaxis->GetNbins()+1; xbin++) { 
297     Double_t content = 0;
298     Double_t error2  = 0;
299     Int_t    nbins   = 0;
300     
301     
302     for (Int_t ybin = first; ybin <= last; ybin++) { 
303       Double_t c1 = h->GetCellContent(xbin, ybin);
304       Double_t e1 = h->GetCellError(xbin, ybin);
305
306       // Ignore empty bins 
307       if (c1 < 1e-12) continue;
308       if (e1 < 1e-12) {
309         if (error) continue; 
310         e1 = 1;
311       }
312
313       content    += c1;
314       error2     += e1*e1;
315       nbins++;
316     } // for (ybin)
317     if(content > 0 && nbins > 0) {
318       Double_t factor = (corr ? Double_t(ybins) / nbins : 1);
319       if (error) { 
320         // Calculate weighted average
321         ret->SetBinContent(xbin, content * factor);
322         ret->SetBinError(xbin, factor * TMath::Sqrt(error2));
323       }
324       else 
325         ret->SetBinContent(xbin, factor * content);
326     }
327   } // for (xbin)
328   
329   return ret;
330 }
331   
332 //________________________________________________________________________
333 void 
334 AliBasedNdetaTask::Terminate(Option_t *) 
335 {
336   // Draw result to screen, or perform fitting, normalizations Called
337   // once at the end of the query
338         
339   fSums = dynamic_cast<TList*> (GetOutputData(1));
340   if(!fSums) {
341     AliError("Could not retrieve TList fSums"); 
342     return; 
343   }
344   
345   if (!fOutput) { 
346     fOutput = new TList;
347     fOutput->SetName(Form("%s_result", GetName()));
348     fOutput->SetOwner();
349   }
350
351   fSum     = static_cast<TH2D*>(fSums->FindObject(GetName()));
352   fSumMC   = static_cast<TH2D*>(fSums->FindObject(Form("%sMC", GetName())));
353   fTriggers       = static_cast<TH1D*>(fSums->FindObject("triggers"));
354
355   if (!fTriggers) { 
356     AliError("Couldn't find histogram 'triggers' in list");
357     return;
358   }
359   if (!fSum) { 
360     AliError("Couldn't find histogram 'base' in list");
361     return;
362   }
363
364   Int_t nAll        = Int_t(fTriggers->GetBinContent(kAll));
365   Int_t nB          = Int_t(fTriggers->GetBinContent(kB));
366   Int_t nA          = Int_t(fTriggers->GetBinContent(kA));
367   Int_t nC          = Int_t(fTriggers->GetBinContent(kC));
368   Int_t nE          = Int_t(fTriggers->GetBinContent(kE));
369   Int_t nMB         = Int_t(fTriggers->GetBinContent(kMB));
370   Int_t nTriggered  = Int_t(fTriggers->GetBinContent(kWithTrigger));
371   Int_t nWithVertex = Int_t(fTriggers->GetBinContent(kWithVertex));
372   Int_t nAccepted   = Int_t(fTriggers->GetBinContent(kAccepted));
373   Int_t nGood       = nB - nA - nC + 2 * nE;
374   if (nTriggered <= 0) { 
375     AliError("Number of triggered events <= 0");
376     return;
377   }
378   if (nGood <= 0) { 
379     AliWarning(Form("Number of good events=%d=%d-%d-%d+2*%d<=0",
380                     nGood, nB, nA, nC, nE));
381     nGood = nMB;
382   }
383   Double_t vtxEff   = Double_t(nMB) / nTriggered * Double_t(nAccepted) / nGood;
384   Double_t vNorm    = Double_t(nAccepted) / nGood;
385   AliInfo(Form("Total of %9d events\n"
386                "                   of these %9d are minimum bias\n"
387                "                   of these %9d has a %s trigger\n" 
388                "                   of these %9d has a vertex\n" 
389                "                   of these %9d were in [%+4.1f,%+4.1f]cm\n"
390                "                   Triggers by type:\n"
391                "                     B   = %9d\n"
392                "                     A|C = %9d (%9d+%-9d)\n"
393                "                     E   = %9d\n"
394                "                   Implies %9d good triggers\n"
395                "                   Vertex efficiency: %f (%f)",
396                nAll, nMB, nTriggered, 
397                AliAODForwardMult::GetTriggerString(fTriggerMask),
398                nWithVertex, nAccepted,
399                fVtxMin, fVtxMax, 
400                nB, nA+nC, nA, nC, nE, nGood, vtxEff, vNorm));
401   
402   const char* name = GetName();
403
404   // Get acceptance normalisation from underflow bins 
405   TH1D* norm   = ProjectX(fSum, Form("norm%s",name), 0, 1, fCorrEmpty, false);
406   // Project onto eta axis - _ignoring_underflow_bins_!
407   TH1D* dndeta = ProjectX(fSum, Form("dndeta%s",name), 1, fSum->GetNbinsY(),
408                           fCorrEmpty);
409   // Normalize to the acceptance 
410   dndeta->Divide(norm);
411   // Scale by the vertex efficiency 
412   dndeta->Scale(vNorm, "width");
413   norm->Scale(1. / nAccepted);
414
415   SetHistogramAttributes(dndeta, kRed+1, 20, Form("ALICE %s", name));
416   SetHistogramAttributes(norm, kRed+1,20,Form("ALICE %s normalisation", name)); 
417
418   fOutput->Add(fTriggers->Clone());
419   if (fSymmetrice)   fOutput->Add(Symmetrice(dndeta));
420   fOutput->Add(dndeta);
421   fOutput->Add(norm);
422   fOutput->Add(Rebin(dndeta));
423   if (fSymmetrice)   fOutput->Add(Symmetrice(Rebin(dndeta)));
424
425   if (fSumMC) { 
426     // Get acceptance normalisation from underflow bins 
427     TH1D* normMC   = ProjectX(fSumMC,Form("norm%sMC", name), 0, 1, 
428                               fCorrEmpty, false);
429     // Project onto eta axis - _ignoring_underflow_bins_!
430     TH1D* dndetaMC = ProjectX(fSumMC,Form("dndeta%sMC", name),1,
431                               fSum->GetNbinsY(), fCorrEmpty);
432     // Normalize to the acceptance 
433     dndetaMC->Divide(normMC);
434     // Scale by the vertex efficiency 
435     dndetaMC->Scale(vNorm, "width");
436     normMC->Scale(1. / nAccepted);
437
438     SetHistogramAttributes(dndetaMC, kRed+3, 21, Form("ALICE %s (MC)",name));
439     SetHistogramAttributes(normMC, kRed+3, 21, Form("ALICE %s (MC) "
440                                                     "normalisation",name)); 
441
442     fOutput->Add(dndetaMC);
443     if (fSymmetrice)   fOutput->Add(Symmetrice(dndetaMC));    
444     fOutput->Add(normMC);
445     fOutput->Add(Rebin(dndetaMC));
446     if (fSymmetrice)   fOutput->Add(Symmetrice(Rebin(dndetaMC)));
447   }
448
449   TNamed* trigString = 
450     new TNamed("trigString", AliAODForwardMult::GetTriggerString(fTriggerMask));
451   trigString->SetUniqueID(fTriggerMask);
452   fOutput->Add(trigString);
453
454   TAxis* vtxAxis = new TAxis(1,fVtxMin,fVtxMax);
455   vtxAxis->SetName("vtxAxis");
456   vtxAxis->SetTitle(Form("v_{z}#in[%+5.1f,%+5.1f]cm", fVtxMin,fVtxMax));
457   fOutput->Add(vtxAxis);
458
459   PostData(2, fOutput);
460 }
461
462 //________________________________________________________________________
463 TH1D*
464 AliBasedNdetaTask::Rebin(const TH1D* h) const
465 {
466   if (fRebin <= 1) return 0;
467
468   Int_t nBins = h->GetNbinsX();
469   if(nBins % fRebin != 0) {
470     Warning("Rebin", "Rebin factor %d is not a devisor of current number "
471             "of bins %d in the histogram %s", fRebin, nBins, h->GetName());
472     return 0;
473   }
474     
475   // Make a copy 
476   TH1D* tmp = static_cast<TH1D*>(h->Clone(Form("%s_rebin%02d", 
477                                                h->GetName(), fRebin)));
478   tmp->Rebin(fRebin);
479   tmp->SetDirectory(0);
480
481   // The new number of bins 
482   Int_t nBinsNew = nBins / fRebin;
483   for(Int_t i = 1;i<= nBinsNew; i++) {
484     Double_t content = 0;
485     Double_t sumw    = 0;
486     Double_t wsum    = 0;
487     Int_t    nbins   = 0;
488     for(Int_t j = 1; j<=fRebin;j++) {
489       Int_t    bin = (i-1)*fRebin + j;
490       Double_t c   =  h->GetBinContent(bin);
491
492       if (c <= 0) continue;
493
494       if (fCutEdges) {
495         if (h->GetBinContent(bin+1)<=0 || 
496             h->GetBinContent(bin-1)) {
497           Warning("Rebin", "removing bin %d=%f of %s (%d=%f,%d=%f)", 
498                   bin, c, h->GetName(), 
499                   bin+1, h->GetBinContent(bin+1), 
500                   bin-1, h->GetBinContent(bin-1));
501           continue;
502         }       
503       }
504       Double_t e =  h->GetBinError(bin);
505       Double_t w =  1 / (e*e); // 1/c/c
506       content    += c;
507       sumw       += w;
508       wsum       += w * c;
509       nbins++;
510     }
511       
512     if(content > 0 && nbins > 0) {
513       tmp->SetBinContent(i, wsum / sumw);
514       tmp->SetBinError(i,1./TMath::Sqrt(sumw));
515     }
516   }
517   
518   return tmp;
519 }
520
521 //__________________________________________________________________
522 /** 
523  * Make an extension of @a h to make it symmetric about 0 
524  * 
525  * @param h Histogram to symmertrice 
526  * 
527  * @return Symmetric extension of @a h 
528  */
529 TH1* 
530 AliBasedNdetaTask::Symmetrice(const TH1* h) const
531 {
532   Int_t nBins = h->GetNbinsX();
533   TH1* s = static_cast<TH1*>(h->Clone(Form("%s_mirror", h->GetName())));
534   s->SetTitle(Form("%s (mirrored)", h->GetTitle()));
535   s->Reset();
536   s->SetBins(nBins, -h->GetXaxis()->GetXmax(), -h->GetXaxis()->GetXmin());
537   s->SetMarkerColor(h->GetMarkerColor());
538   s->SetMarkerSize(h->GetMarkerSize());
539   s->SetMarkerStyle(h->GetMarkerStyle()+4);
540   s->SetFillColor(h->GetFillColor());
541   s->SetFillStyle(h->GetFillStyle());
542   s->SetDirectory(0);
543
544   // Find the first and last bin with data 
545   Int_t first = nBins+1;
546   Int_t last  = 0;
547   for (Int_t i = 1; i <= nBins; i++) { 
548     if (h->GetBinContent(i) <= 0) continue; 
549     first = TMath::Min(first, i);
550     last  = TMath::Max(last,  i);
551   }
552     
553   Double_t xfirst = h->GetBinCenter(first-1);
554   Int_t    f1     = h->GetXaxis()->FindBin(-xfirst);
555   Int_t    l2     = s->GetXaxis()->FindBin(xfirst);
556   for (Int_t i = f1, j=l2; i <= last; i++,j--) { 
557     s->SetBinContent(j, h->GetBinContent(i));
558     s->SetBinError(j, h->GetBinError(i));
559   }
560   // Fill in overlap bin 
561   s->SetBinContent(l2+1, h->GetBinContent(first));
562   s->SetBinError(l2+1, h->GetBinError(first));
563   return s;
564 }