test/Macros/fit2Dprojection.cc

0001 #include "TROOT.h"
0002 #include "TH1.h"
0003 #include "TH2.h"
0004 #include "TF1.h"
0005 #include "TGraphErrors.h"
0006 #include "TLegend.h"
0007 #include "TCanvas.h"
0008 #include "TSystem.h"
0009 #include "TStyle.h"
0010 #include "TFile.h"
0011 #include "TPaveText.h"
0012 #include <vector>
0013 #include <iostream>
0014 #include <sstream>
0015 #include <cmath>
0016
0017 void setTDRStyle();
0018
0019 /// Small function to simplify the creation of text in the TPaveText
0020 TString setText(const char * text, const double & num1, const char * divider = "", const double & num2 = 0) {
0021   std::stringstream numString;
0022   TString textString(text);
0023   numString << num1;
0024   textString += numString.str();
0025   if( num2 != 0 ) {
0026     textString += divider;
0027     numString.str("");
0028     numString << num2;
0029     textString += numString.str();
0030   }
0031   return textString;
0032 }
0033
0034 /// This function sets up the TPaveText with chi2/ndf and parameters values +- errors
0035 void setTPaveText(const TF1 * fit, TPaveText * paveText) {
0036   Double_t chi2 = fit->GetChisquare();
0037   Int_t ndf = fit->GetNDF();
0038   paveText->AddText(setText("#chi^2 / ndf = ", chi2,  " / ", ndf));
0039
0040   for( int iPar=0; iPar<fit->GetNpar(); ++iPar ) {
0041     TString parName(fit->GetParName(iPar));
0042     Double_t parValue = fit->GetParameter(iPar); // value of Nth parameter
0043     Double_t parError = fit->GetParError(iPar);  // error on Nth parameter
0044     paveText->AddText(setText(parName + " = ", parValue, " #pm ", parError));
0045   }
0046 }
0047
0048 // Structure to hold the fit results
0049 struct FitResult
0050 {
0051   double top;
0052   double topError;
0053   double width;
0054   double widthError;
0055   double mean;
0056   double meanError;
0057   double xCenter;
0058   double xCenterError;
0059   double chi2;
0060 };
0061
0062 class ProjectionFitter
0063 {
0064  public:
0065   TGraphErrors * fit2Dprojection(const TString & inputFileName, const TString & histoName,
0066                                  const unsigned int rebinX, const unsigned int rebinY,
0067                                  const double & fitXmin, const double & fitXmax,
0068                                  const TString & fitName, const TString & append,
0069                                  TFile * outputFile, const unsigned int minEntries = 100);
0070  protected:
0071   /// Take the projection and give it a suitable name and title
0072   TH1 * takeProjectionY( const TH2 * histo, const int index ) {
0073     // Name
0074     std::stringstream number;
0075     number << index;
0076     TString numberString(number.str());
0077     TString name = TString(histo->GetName()) + "_" + numberString;
0078     // Title
0079     double xBin = histo->GetXaxis()->GetBinCenter(index);
0080     std::stringstream xBinString;
0081     xBinString << xBin;
0082     TString title("Projection of x = ");
0083     title += xBinString.str();
0084
0085     TH1 * histoY = histo->ProjectionY(name, index, index);
0086     histoY->SetName(title);
0087     return histoY;
0088   }
0089 };
0090
0091 /// This method fits TH2F slices with a selected fitType function (gaussian, lorentz, ...).
0092 TGraphErrors * ProjectionFitter::fit2Dprojection(const TString & inputFileName, const TString & histoName,
0093                                                  const unsigned int rebinX, const unsigned int rebinY,
0094                                                  const double & fitXmin, const double & fitXmax,
0095                                                  const TString & fitName, const TString & append,
0096                                                  TFile * outputFile, const unsigned int minEntries)
0097 {
0098   // Read the TH2 from file
0099   // std::cout << "inputFileName = " << inputFileName << std::endl;
0100   TFile *inputFile = new TFile(inputFileName);
0101   // FindObjectAny finds the object of the given name looking also in the subdirectories.
0102   // It does not change the current directory if it finds a matching (for that use FindKeyAny).
0103   TH2 * histo = (TH2*) inputFile->FindObjectAny(histoName);
0104   if( rebinX > 0 ) histo->RebinX(rebinX);
0105   if( rebinY > 0 ) histo->RebinY(rebinY);
0106
0107   std::vector<FitResult> fitResults;
0108
0109   // Output file
0110   // TString outputFileName("fitCompare2"+histoName);
0111   // outputFileName += append;
0112   // outputFileName += ".root";
0113   // TFile * outputFile = new TFile(outputFileName,"RECREATE");
0114
0115   // For each bin in X take the projection on Y and fit.
0116   // All the fits are saved in a single canvas.
0117   unsigned int nBins = histo->GetXaxis()->GetNbins();
0118
0119   std::vector<TH1*> projections;
0120   std::vector<TF1*> projectionsFits;
0121
0122   for( unsigned int i=1; i<=nBins; ++i ) {
0123
0124     // Project on Y
0125     // ------------
0126     TH1 * histoY = takeProjectionY( histo, i );
0127
0128     // Require a minimum number of entries to do the fit
0129     if( histoY->GetEntries() > minEntries ) {
0130
0131       // Gaussian fit
0132       std::stringstream ss;
0133       ss << i;
0134       TString fitFunctionName(fitName+"Fit_"+ss.str());
0135       TF1 * fitFunction = new TF1(fitFunctionName.Data(), fitName.Data(), fitXmin, fitXmax);
0136       fitFunction->SetLineColor(kRed);
0137       projectionsFits.push_back(fitFunction);
0138       // Options: M = "more": improve fit results; Q = "quiet"
0139       // If not using the N options it will try to save the histogram and corrupt the memory.
0140       // We save also draw functions and write the later.
0141       histoY->Fit(fitFunctionName, "MQN", "", fitXmin, fitXmax);
0142
0143       projections.push_back(histoY);
0144
0145       double *par = fitFunction->GetParameters();
0146       double *err = fitFunction->GetParErrors();
0147
0148       // Store the fit results
0149       FitResult fitResult;
0150       if( par[0] == par[0] ) {
0151         fitResult.top = par[0];
0152         fitResult.topError = err[0];
0153         // sometimes the gaussian has negative width (checked with Rene Brun)
0154         fitResult.mean = fabs(par[1]);
0155         fitResult.meanError = err[1];
0156         fitResult.width = par[2];
0157         fitResult.widthError = err[2];
0158
0159         fitResult.xCenter = histo->GetXaxis()->GetBinCenter(i);
0160         fitResult.xCenterError = 0;
0161
0162         double chi2 = fitFunction->GetChisquare()/fitFunction->GetNDF();
0163         if( chi2 == chi2 ) fitResult.chi2 = fitFunction->GetChisquare()/fitFunction->GetNDF();
0164         else fitResult.chi2 = 100000;
0165       }
0166       else {
0167         // Skip nan
0168         fitResult.top = 0;
0169         fitResult.topError = 1;
0170         // sometimes the gaussian has negative width (checked with Rene Brun)
0171         fitResult.mean = 0;
0172         fitResult.meanError = 1;
0173         fitResult.width = 0;
0174         fitResult.widthError = 1;
0175
0176         fitResult.xCenter = histo->GetXaxis()->GetBinCenter(i);
0177         fitResult.xCenterError = 0;
0178
0179         fitResult.chi2 = 100000;
0180       }
0181       fitResults.push_back(fitResult);
0182     }
0183   }
0184
0185   TCanvas * canvasYcheck = new TCanvas(histoName+"_canvas_check"+append, histoName+" fits check", 1000, 800);
0186   int sizeCheck = projections.size();
0187   int x = int(sqrt(sizeCheck));
0188   int y = x;
0189   if( x*y < sizeCheck ) y += 1;
0190   if( x*y < sizeCheck ) x += 1;
0191   std::cout << "sizeCheck = " << sizeCheck << std::endl;
0192   std::cout << "x*y = " << x*y << std::endl;
0193   canvasYcheck->Divide(x, y);
0194
0195   std::vector<TH1*>::const_iterator it = projections.begin();
0196   std::vector<TF1*>::const_iterator fit = projectionsFits.begin();
0197   for( int i=1; it != projections.end(); ++it, ++fit, ++i ) {
0198     canvasYcheck->cd(i);
0199     (*it)->Draw();
0200     (*it)->GetXaxis()->SetRangeUser(fitXmin, fitXmax);
0201     (*fit)->Draw("same");
0202     (*fit)->SetLineColor(kRed);
0203   }
0204
0205   // Plot the fit results in TGraphs
0206   const unsigned int fitsNumber = fitResults.size();
0207   double xCenter[fitsNumber], xCenterError[fitsNumber],
0208     mean[fitsNumber], meanError[fitsNumber],
0209     width[fitsNumber], widthError[fitsNumber],
0210     chi2[fitsNumber];
0211
0212   double xDisplace = 0;
0213
0214   for( unsigned int i=0; i<fitsNumber; ++i ) {
0215
0216     FitResult * fitResult = &(fitResults[i]);
0217
0218     xCenter[i]      = fitResult->xCenter + xDisplace;
0219     xCenterError[i] = fitResult->xCenterError;
0220     mean[i]         = fitResult->mean;
0221     meanError[i]    = fitResult->meanError;
0222     width[i]        = fitResult->width;
0223     width[i]        = fitResult->widthError;
0224     chi2[i]         = fitResult->chi2;
0225   }
0226
0227   TGraphErrors *grMean = new TGraphErrors(fitsNumber, xCenter, mean, xCenterError, meanError);
0228   grMean->SetTitle(histoName+"_Mean"+append);
0229   grMean->SetName(histoName+"_Mean"+append);
0230   TGraphErrors *grWidth = new TGraphErrors(fitsNumber, xCenter, width, xCenterError, widthError);
0231   grWidth->SetTitle(histoName+"_Width"+append);
0232   grWidth->SetName(histoName+"_Width"+append);
0233   TGraphErrors *grChi2 = new TGraphErrors(fitsNumber, xCenter, chi2, xCenterError, xCenterError);
0234   grChi2->SetTitle(histoName+"_chi2"+append);
0235   grChi2->SetName(histoName+"_chi2"+append);
0236
0237   grMean->SetMarkerColor(4);
0238   grMean->SetMarkerStyle(20);
0239   grWidth->SetMarkerColor(4);
0240   grWidth->SetMarkerStyle(20);
0241   grChi2->SetMarkerColor(4);
0242   grChi2->SetMarkerStyle(20);
0243
0244   outputFile->cd();
0245   // Draw and save the graphs
0246   TCanvas * c1 = new TCanvas(histoName+"_Width"+append, histoName+"_Width");
0247   c1->cd();
0248   grWidth->Draw("AP");
0249   c1->Write();
0250   TCanvas * c2 = new TCanvas(histoName+"_Mean"+append, histoName+"_Mean");
0251   c2->cd();
0252   grMean->Draw("AP");
0253   c2->Write();
0254   TCanvas * c3 = new TCanvas(histoName+"_Chi2"+append, histoName+"_Chi2");
0255   c3->cd();
0256   grChi2->Draw("AP");
0257   c3->Write();
0258
0259   // Write the canvas with the fits
0260   canvasYcheck->Write();
0261
0262   // outputFile->Close();
0263
0264   return grMean;
0265 }
0266
0267 /****************************************************************************************/
0268 void macroPlot( TString name, const TString & nameFile1, const TString & nameFile2, const TString & title,
0269                 const TString & resonanceType, const int rebinX, const int rebinY, const TString & outputFileName,
0270                 TFile * externalOutputFile = 0)
0271 {
0272   gROOT->SetBatch(true);
0273
0274   //Save the graphs in a file
0275   TFile *outputFile = externalOutputFile;
0276   if( outputFile == 0 ) {
0277     outputFile = new TFile(outputFileName,"RECREATE");
0278   }
0279
0280   setTDRStyle();
0281
0282   ProjectionFitter projectionFitter;
0283
0284   std::cout << "File 1 = " << nameFile1 << std::endl;
0285   std::cout << "File 2 = " << nameFile2 << std::endl;
0286
0287   double y[2];
0288   if( resonanceType == "JPsi" || resonanceType == "Psi2S" ) { y[0]=2.9; y[1]=3.3; }
0289   else if( resonanceType.Contains("Upsilon1S") ) { y[0]=9.25; y[1]=9.85; }
0290   else if( resonanceType.Contains("Upsilon2S") ) { y[0]=9.85; y[1]=10.2; }
0291   else if( resonanceType.Contains("Upsilon3S") ) { y[0]=10.2; y[1]=10.7; }
0292   else if( resonanceType == "Z" ) { y[0]=70; y[1]=110; }
0293   else if( resonanceType == "Resolution" ) { y[0]=-0.3; y[1]=0.3; }
0294
0295   TGraphErrors *grM_1 = projectionFitter.fit2Dprojection( nameFile1, name, rebinX, rebinY, y[0], y[1], "gaus", "_1", outputFile );
0296   TGraphErrors *grM_2 = projectionFitter.fit2Dprojection( nameFile2, name, rebinX, rebinY, y[0], y[1], "gaus", "_2", outputFile );
0297
0298   TCanvas *c = new TCanvas(name,name);
0299   c->SetGridx();
0300   c->SetGridy();
0301
0302   grM_1->SetMarkerColor(1);
0303   grM_2->SetMarkerColor(2);
0304
0305   TString xAxisTitle;
0306
0307   double x[2];
0308
0309   if( name.Contains("Eta") ) {
0310     x[0]=-3; x[1]=3;
0311     xAxisTitle = "#eta";
0312   }
0313   else if( name.Contains("PhiPlus") || name.Contains("PhiMinus") ) {
0314     x[0] = -3.2; x[1] = 3.2;
0315     xAxisTitle = "#phi(rad)";
0316   }
0317   else {
0318     x[0] = 0.; x[1] = 200;
0319     xAxisTitle = "pt(GeV)";
0320   }
0321
0322   // This is used to have a canvas containing both histogram points
0323   TGraph *gr = new TGraph(2,x,y);
0324   gr->SetMarkerStyle(8);
0325   gr->SetMarkerColor(108);
0326   gr->GetYaxis()->SetTitle("Mass(GeV)   ");
0327   gr->GetYaxis()->SetTitleOffset(1);
0328   gr->GetXaxis()->SetTitle(xAxisTitle);
0329   gr->SetTitle(title);
0330
0331   // Text for the fits
0332   TPaveText * paveText1 = new TPaveText(0.20,0.15,0.49,0.35,"NDC");
0333   paveText1->SetFillColor(0);
0334   paveText1->SetTextColor(1);
0335   paveText1->SetTextSize(0.02);
0336   paveText1->SetBorderSize(1);
0337   TPaveText * paveText2 = new TPaveText(0.59,0.15,0.88,0.35,"NDC");
0338   paveText2->SetFillColor(0);
0339   paveText2->SetTextColor(2);
0340   paveText2->SetTextSize(0.02);
0341   paveText2->SetBorderSize(1);
0342
0343   c->cd();
0344   gr->Draw("AP");
0345   grM_1->Draw("P");
0346   grM_2->Draw("P");
0347
0348   paveText1->Draw("same");
0349   paveText2->Draw("same");
0350
0351   TLegend *leg = new TLegend(0.65,0.85,1,1);
0352   leg->SetFillColor(0);
0353   leg->AddEntry(grM_1,"before calibration","P");
0354   leg->AddEntry(grM_2,"after calibration","P");
0355   leg->Draw("same");
0356
0357   outputFile->cd();
0358   c->Write();
0359   if( externalOutputFile == 0 ) {
0360     outputFile->Close();
0361   }
0362 }
0363
0364 /****************************************************************************************/
0365
0366 void setTDRStyle() {
0367   TStyle *tdrStyle = new TStyle("tdrStyle","Style for P-TDR");
0368
0369   // For the canvas:
0370   tdrStyle->SetCanvasBorderMode(0);
0371   tdrStyle->SetCanvasColor(kWhite);
0372   tdrStyle->SetCanvasDefH(600); //Height of canvas
0373   tdrStyle->SetCanvasDefW(600); //Width of canvas
0374   tdrStyle->SetCanvasDefX(0);   //POsition on screen
0375   tdrStyle->SetCanvasDefY(0);
0376
0377   // For the Pad:
0378   tdrStyle->SetPadBorderMode(0);
0379   tdrStyle->SetPadColor(kWhite);
0380   tdrStyle->SetPadGridX(false);
0381   tdrStyle->SetPadGridY(false);
0382   tdrStyle->SetGridColor(0);
0383   tdrStyle->SetGridStyle(3);
0384   tdrStyle->SetGridWidth(1);
0385
0386   // For the frame:
0387   tdrStyle->SetFrameBorderMode(0);
0388   tdrStyle->SetFrameBorderSize(1);
0389   tdrStyle->SetFrameFillColor(0);
0390   tdrStyle->SetFrameFillStyle(0);
0391   tdrStyle->SetFrameLineColor(1);
0392   tdrStyle->SetFrameLineStyle(1);
0393   tdrStyle->SetFrameLineWidth(1);
0394
0395   // For the histo:
0396   tdrStyle->SetHistLineColor(1);
0397   tdrStyle->SetHistLineStyle(0);
0398   tdrStyle->SetHistLineWidth(1);
0399
0400   tdrStyle->SetEndErrorSize(2);
0401   tdrStyle->SetErrorX(0.);
0402
0403   tdrStyle->SetMarkerStyle(20);
0404
0405   //For the fit/function:
0406   tdrStyle->SetOptFit(1);
0407   tdrStyle->SetFitFormat("5.4g");
0408   tdrStyle->SetFuncColor(2);
0409   tdrStyle->SetFuncStyle(1);
0410   tdrStyle->SetFuncWidth(1);
0411
0412   //For the date:
0413   tdrStyle->SetOptDate(0);
0414
0415   // For the statistics box:
0416   tdrStyle->SetOptFile(0);
0417   tdrStyle->SetOptStat(0); // To display the mean and RMS:   SetOptStat("mr");
0418   tdrStyle->SetStatColor(kWhite);
0419   tdrStyle->SetStatFont(42);
0420   tdrStyle->SetStatFontSize(0.025);
0421   tdrStyle->SetStatTextColor(1);
0422   tdrStyle->SetStatFormat("6.4g");
0423   tdrStyle->SetStatBorderSize(1);
0424   tdrStyle->SetStatH(0.1);
0425   tdrStyle->SetStatW(0.15);
0426
0427   // Margins:
0428   tdrStyle->SetPadTopMargin(0.05);
0429   tdrStyle->SetPadBottomMargin(0.13);
0430   tdrStyle->SetPadLeftMargin(0.13);
0431   tdrStyle->SetPadRightMargin(0.05);
0432
0433   // For the Global title:
0434   tdrStyle->SetOptTitle(0);
0435
0436   // For the axis titles:
0437   tdrStyle->SetTitleColor(1, "XYZ");
0438   tdrStyle->SetTitleFont(42, "XYZ");
0439   tdrStyle->SetTitleSize(0.06, "XYZ");
0440   tdrStyle->SetTitleXOffset(0.9);
0441   tdrStyle->SetTitleYOffset(1.05);
0442
0443   // For the axis labels:
0444   tdrStyle->SetLabelColor(1, "XYZ");
0445   tdrStyle->SetLabelFont(42, "XYZ");
0446   tdrStyle->SetLabelOffset(0.007, "XYZ");
0447   tdrStyle->SetLabelSize(0.05, "XYZ");
0448
0449   // For the axis:
0450   tdrStyle->SetAxisColor(1, "XYZ");
0451   tdrStyle->SetStripDecimals(kTRUE);
0452   tdrStyle->SetTickLength(0.03, "XYZ");
0453   tdrStyle->SetNdivisions(510, "XYZ");
0454   tdrStyle->SetPadTickX(1);  // To get tick marks on the opposite side of the frame
0455   tdrStyle->SetPadTickY(1);
0456
0457   // Change for log plots:
0458   tdrStyle->SetOptLogx(0);
0459   tdrStyle->SetOptLogy(0);
0460   tdrStyle->SetOptLogz(0);
0461
0462   //tdrStyle->SetOptFit(00000);
0463   tdrStyle->cd();
0464 }