first checkin of Claus' code with a few modifications
[u/mrichter/AliRoot.git] / PWG0 / dNdEta / dNdEtaCorrection.cxx
1 #include "dNdEtaCorrection.h"
2
3 #include <TCanvas.h>
4
5 //____________________________________________________________________
6 ClassImp(dNdEtaCorrection);
7
8 //____________________________________________________________________
9 dNdEtaCorrection::dNdEtaCorrection(Char_t* name) {
10
11   fName = TString(name);
12   
13   hEtaVsVtx_meas  = new TH2F("etaVsVtx_meas", "etaVsVtx_meas" ,80,-20,20,120,-6,6);
14   hEtaVsVtx_gene  = new TH2F("etaVsVtx_gene", "etaVsVtx_gene" ,80,-20,20,120,-6,6);
15   hEtaVsVtx_corr  = new TH2F("etaVsVtx_corr", "etaVsVtx_corr", 80,-20,20,120,-6,6);
16
17   hEtaVsVtx_ratio = new TH2F("etaVsVtx_ratio","etaVsVtx_ratio",80,-20,20,120,-6,6);
18
19   hEtaVsVtx_meas ->SetXTitle("vtx z [cm]");  hEtaVsVtx_meas ->SetYTitle("#eta"); 
20   hEtaVsVtx_gene ->SetXTitle("vtx z [cm]");  hEtaVsVtx_gene ->SetYTitle("#eta"); 
21   hEtaVsVtx_corr ->SetXTitle("vtx z [cm]");  hEtaVsVtx_corr ->SetYTitle("#eta"); 
22   hEtaVsVtx_ratio->SetXTitle("vtx z [cm]");  hEtaVsVtx_ratio->SetYTitle("#eta"); 
23
24 }
25
26 //____________________________________________________________________
27 void
28 dNdEtaCorrection::Finish() {  
29
30   hEtaVsVtx_ratio->Divide(hEtaVsVtx_meas, hEtaVsVtx_gene, 1,1,"B");
31   hEtaVsVtx_corr->Divide(hEtaVsVtx_gene, hEtaVsVtx_meas, 1,1,"B");
32
33   Int_t nBinsVtx = hEtaVsVtx_corr->GetNbinsX();
34   Int_t nBinsEta = hEtaVsVtx_corr->GetNbinsY();
35
36   TH2F* tmp = (TH2F*)hEtaVsVtx_corr->Clone("tmp");
37
38   // cut at 0.2
39   for (Int_t bx=0; bx<=nBinsVtx; bx++) {
40     for (Int_t by=0; by<=nBinsEta; by++) {
41       if (tmp->GetBinContent(bx,by)<0.2) {
42         hEtaVsVtx_corr->SetBinContent(bx,by,0);
43         hEtaVsVtx_corr->SetBinError(bx,by,0);
44         
45         tmp->SetBinContent(bx,by,0);
46       }
47       else 
48         tmp->SetBinContent(bx,by,1);
49     }
50   }
51   
52 }
53
54 //____________________________________________________________________
55 void
56 dNdEtaCorrection::RemoveEdges(Float_t cut, Int_t nBinsVtx, Int_t nBinsEta) {
57
58   // remove edges of correction histogram by removing 
59   // - bins with content less than cut
60   // - bins next to bins with zero bin content
61   
62   Int_t nBinsX = hEtaVsVtx_corr->GetNbinsX();
63   Int_t nBinsY = hEtaVsVtx_corr->GetNbinsY();
64
65   // set bin content to zero for bins with content smaller cut
66   for (Int_t bx=0; bx<=nBinsX; bx++) {
67     for (Int_t by=0; by<=nBinsY; by++) {
68       if (hEtaVsVtx_corr->GetBinContent(bx,by)>cut) {
69         hEtaVsVtx_corr->SetBinContent(bx,by,0);
70         hEtaVsVtx_corr->SetBinError(bx,by,0);
71       }
72     }
73   }
74
75   // set bin content to zero for bins next to bins with zero
76   TH2F* tmp = (TH2F*)hEtaVsVtx_corr->Clone("tmp");
77   tmp->Reset();
78   
79   Bool_t done = kFALSE;
80   Int_t nBinsVtxCount = 0;
81   Int_t nBinsEtaCount = 0;
82   while (!done) {    
83     if (nBinsVtxCount<nBinsVtx) 
84       for (Int_t bx=0; bx<=nBinsX; bx++) {
85         for (Int_t by=0; by<=nBinsY; by++) {
86           if ((hEtaVsVtx_corr->GetBinContent(bx+1,by)==0)|| 
87               (hEtaVsVtx_corr->GetBinContent(bx-1,by)==0))
88             tmp->SetBinContent(bx,by,1);        
89           
90         }
91       }
92     if (nBinsEtaCount<nBinsEta) 
93       for (Int_t bx=0; bx<=nBinsX; bx++) {
94         for (Int_t by=0; by<=nBinsY; by++) {
95           if ((hEtaVsVtx_corr->GetBinContent(bx,by+1)==0)|| 
96               (hEtaVsVtx_corr->GetBinContent(bx,by-1)==0))
97             tmp->SetBinContent(bx,by,1);        
98         }
99       }    
100     for (Int_t bx=0; bx<=nBinsX; bx++) {
101       for (Int_t by=0; by<=nBinsY; by++) {
102         if (tmp->GetBinContent(bx,by)==1) {
103           hEtaVsVtx_corr->SetBinContent(bx,by,0);
104           hEtaVsVtx_corr->SetBinError(bx,by,0);
105         }
106       }
107     }
108     nBinsVtxCount++;
109     nBinsEtaCount++;
110     if ((nBinsVtxCount>=nBinsVtx)&&(nBinsEtaCount>=nBinsEta)) done=kTRUE;
111   }
112   tmp->Delete();  
113
114 }
115
116 //____________________________________________________________________
117 Bool_t
118 dNdEtaCorrection::LoadHistograms(Char_t* fileName, Char_t* dir) {
119
120   TFile* fin = TFile::Open(fileName);  
121   
122   // add test of file
123   // return kFALSE
124
125   hEtaVsVtx_meas  = (TH2F*)fin->Get(Form("%s/etaVsVtx_meas",dir));
126   hEtaVsVtx_gene  = (TH2F*)fin->Get(Form("%s/etaVsVtx_gene",dir));
127   hEtaVsVtx_corr  = (TH2F*)fin->Get(Form("%s/etaVsVtx_corr",dir));
128
129   hEtaVsVtx_ratio = (TH2F*)fin->Get(Form("%s/etaVsVtx_ratio",dir));
130
131   return kTRUE;
132 }
133
134
135 //____________________________________________________________________
136 void
137 dNdEtaCorrection::SaveHistograms() {
138
139   gDirectory->mkdir(fName.Data());
140   gDirectory->cd(fName.Data());
141   
142   hEtaVsVtx_meas ->Write();
143   hEtaVsVtx_gene ->Write();
144
145   if (hEtaVsVtx_corr)
146     hEtaVsVtx_corr->Write();
147
148   if (hEtaVsVtx_ratio)
149     hEtaVsVtx_ratio->Write();
150
151
152   gDirectory->cd("../");
153 }
154
155 //____________________________________________________________________
156 void dNdEtaCorrection::DrawHistograms()
157 {
158   TCanvas* canvas = new TCanvas("dNdEtaCorrection", "dNdEtaCorrection", 800, 800);
159   canvas->Divide(2, 2);
160   
161   canvas->cd(1);
162   if (hEtaVsVtx_meas)
163     hEtaVsVtx_meas->Draw("COLZ");
164
165   canvas->cd(2);
166   if (hEtaVsVtx_gene)
167     hEtaVsVtx_gene->Draw("COLZ");
168
169   canvas->cd(3);
170   if (hEtaVsVtx_ratio)
171     hEtaVsVtx_ratio->Draw("COLZ");
172
173   canvas->cd(4);
174   if (hEtaVsVtx_corr)
175     hEtaVsVtx_corr->Draw("COLZ");
176 }