(martin) 2D correction matrix calculation class
[u/mrichter/AliRoot.git] / PWG0 / CorrectionMatrix2D.cxx
1 #include "CorrectionMatrix2D.h"
2
3 //____________________________________________________________________
4 ClassImp(CorrectionMatrix2D);
5
6 //____________________________________________________________________
7 CorrectionMatrix2D::CorrectionMatrix2D(Char_t* name) {
8
9   fName = TString(name);
10
11 }
12 //________________________________________________________________________
13 void CorrectionMatrix2D::SetHist(Char_t* title ,Int_t nBinX, Float_t Xmin, Float_t Xmax,
14                                   Int_t nBinY, Float_t Ymin, Float_t Ymax) 
15 {
16 //Char_t* name , Chat_t* title ,   
17
18   hmeas  = new TH2F("meas",  Form("meas_%s",title), nBinX,Xmin,Xmax,nBinY,Ymin,Ymax);
19   hgene  = new TH2F("gene",  Form("gene_%s",title), nBinX,Xmin,Xmax,nBinY,Ymin,Ymax);
20   hcorr  = new TH2F("corr",  Form("corr_%s",title), nBinX,Xmin,Xmax,nBinY,Ymin,Ymax);
21   hratio = new TH2F("ratio", Form("ratio_%s",title), nBinX,Xmin,Xmax,nBinY,Ymin,Ymax);
22 }
23
24
25 void CorrectionMatrix2D::SetHist(Char_t* title ,Int_t nBinX, Float_t *X,
26                                                 Int_t nBinY, Float_t *Y) 
27 {
28
29   hmeas  = new TH2F("meas",  Form("meas_%s",title), nBinX,X,nBinY,Y);
30   hgene  = new TH2F("gene",  Form("gene_%s",title), nBinX,X,nBinY,Y);
31   hcorr  = new TH2F("corr",  Form("corr_%s",title), nBinX,X,nBinY,Y);
32   hratio = new TH2F("ratio", Form("ratio_%s",title),nBinX,X,nBinY,Y);
33   
34 }
35
36
37
38 //________________________________________________________________________
39 void CorrectionMatrix2D::SetHistTitle(Char_t* titleX, Char_t* titleY) 
40 {
41  
42   hmeas ->SetXTitle(titleX);  hmeas ->SetYTitle(titleY); 
43   hgene ->SetXTitle(titleX);  hgene ->SetYTitle(titleY); 
44   hcorr ->SetXTitle(titleX);  hcorr ->SetYTitle(titleY); 
45   hratio->SetXTitle(titleX);  hratio->SetYTitle(titleY); 
46
47 }
48
49
50 //____________________________________________________________________
51 void CorrectionMatrix2D::Finish() {  
52
53 if (!hmeas || !hgene)  return;
54  
55 TCanvas *c1 = new TCanvas();
56 c1->Divide(2,1);
57 c1->cd(1);
58 hgene->Draw();
59 c1->cd(2);
60 hmeas->Draw();
61
62   hratio->Divide(hmeas, hgene, 1,1,"B");
63   hcorr->Divide(hgene, hmeas, 1,1,"B");
64
65   Int_t nBinsVtx = hcorr->GetNbinsX();
66   Int_t nBinsEta = hcorr->GetNbinsY();
67
68   TH2F* tmp = (TH2F*)hcorr->Clone("tmp");
69
70   // cut at 0.2
71   for (Int_t bx=0; bx<=nBinsVtx; bx++) {
72     for (Int_t by=0; by<=nBinsEta; by++) {
73       if (tmp->GetBinContent(bx,by)<0.2) {
74         hcorr->SetBinContent(bx,by,0);
75         hcorr->SetBinError(bx,by,0);
76         
77         tmp->SetBinContent(bx,by,0);
78       }
79       else 
80         tmp->SetBinContent(bx,by,1);
81     }
82   }
83   
84 }
85
86 //____________________________________________________________________
87 void
88 CorrectionMatrix2D::RemoveEdges(Float_t cut, Int_t nBinsXedge, Int_t nBinsYedge) {
89
90   // remove edges of correction histogram by removing 
91   // - bins with content less than cut
92   // - bins next to bins with zero bin content
93   
94   Int_t nBinsX = hcorr->GetNbinsX();
95   Int_t nBinsY = hcorr->GetNbinsY();
96
97   // set bin content to zero for bins with content smaller cut
98   for (Int_t bx=0; bx<=nBinsX; bx++) {
99     for (Int_t by=0; by<=nBinsY; by++) {
100       if (hcorr->GetBinContent(bx,by)>cut) {
101           hcorr->SetBinContent(bx,by,0);
102           hcorr->SetBinError(bx,by,0);
103       }
104     }
105   }
106
107   // set bin content to zero for bins next to bins with zero
108   TH2F* tmp = (TH2F*)hcorr->Clone("tmp");
109   tmp->Reset();
110   
111   Bool_t done = kFALSE;
112   Int_t nBinsXCount = 0;
113   Int_t nBinsYCount = 0;
114   while (!done) {    
115     if (nBinsXCount<nBinsXedge) 
116       for (Int_t bx=0; bx<=nBinsX; bx++) {
117         for (Int_t by=0; by<=nBinsY; by++) {
118           if ((hcorr->GetBinContent(bx+1,by)==0)|| 
119               (hcorr->GetBinContent(bx-1,by)==0))
120             tmp->SetBinContent(bx,by,1);        
121           
122         }
123       }
124     if (nBinsYCount<nBinsYedge) 
125       for (Int_t bx=0; bx<=nBinsX; bx++) {
126         for (Int_t by=0; by<=nBinsY; by++) {
127           if ((hcorr->GetBinContent(bx,by+1)==0)|| 
128               (hcorr->GetBinContent(bx,by-1)==0))
129             tmp->SetBinContent(bx,by,1);        
130         }
131       }    
132     for (Int_t bx=0; bx<=nBinsX; bx++) {
133       for (Int_t by=0; by<=nBinsY; by++) {
134         if (tmp->GetBinContent(bx,by)==1) {
135           hcorr->SetBinContent(bx,by,0);
136           hcorr->SetBinError(bx,by,0);
137         }
138       }
139     }
140     nBinsXCount++;
141     nBinsYCount++;
142     if ((nBinsXCount>=nBinsXedge)&&(nBinsYCount>=nBinsYedge)) done=kTRUE;
143   }
144   tmp->Delete();  
145
146 }
147
148 //____________________________________________________________________
149 Bool_t CorrectionMatrix2D::LoadHistograms(Char_t* fileName, Char_t* dir) {
150
151   TFile* fin = TFile::Open(fileName);  
152   
153   if(!fin) 
154   {
155     //Info("LoadHistograms",Form(" %s file does not exist",fileName));
156   return kFALSE;
157   }
158   
159   if(hgene)  {delete hgene;  hgene=0;}
160   if(hcorr)  {delete hcorr;  hcorr=0;}
161   if(hratio) {delete hratio; hratio=0;}
162   if(hmeas)  {delete hmeas;  hmeas=0;}
163   
164   hmeas  = (TH2F*)fin->Get(Form("%s/meas",dir));
165       if(!hmeas)  Info("LoadHistograms","No meas  hist available");
166   hgene  = (TH2F*)fin->Get(Form("%s/gene",dir));
167       if(!hgene)  Info("LoadHistograms","No gene  hist available");
168   hratio = (TH2F*)fin->Get(Form("%s/ratio",dir));
169       if(!hratio) Info("LoadHistograms","No ratio hist available");
170   hcorr  = (TH2F*)fin->Get(Form("%s/corr",dir));
171       if(!hcorr) {Info("LoadHistograms","No corr  hist available");
172       return kFALSE;}
173       
174   return kTRUE;
175 }
176
177
178 //____________________________________________________________________
179 void
180 CorrectionMatrix2D::SaveHistograms() {
181
182   gDirectory->mkdir(fName.Data());
183   gDirectory->cd(fName.Data());
184   
185   hmeas ->Write();
186   hgene ->Write();
187
188   if (hcorr)
189     hcorr->Write();
190
191   if (hratio)
192     hratio->Write();
193
194   gDirectory->cd("../");
195 }
196
197 //____________________________________________________________________
198 void CorrectionMatrix2D::DrawHistograms()
199 {
200   TCanvas* canvas = new TCanvas("Correction", "Correction", 800, 800);
201   canvas->Divide(2, 2);
202
203   canvas->cd(1);
204   if (hmeas)
205     hmeas->Draw("COLZ");
206
207   canvas->cd(2);
208   if (hgene)
209     hgene->Draw("COLZ");
210
211   canvas->cd(3);
212   if (hratio)
213     hratio->Draw("COLZ");
214
215   canvas->cd(4);
216   if (hcorr)
217     hcorr->Draw("COLZ");
218 }
219
220