HFInvMassFitter: Fix canvas division, add ratio histograms

Author: pstahlhu

PWGHF/D2H/Macros/HFInvMassFitter.cxx

@@ -359,6 +359,7 @@ void HFInvMassFitter::doFit()
       mTotalPdf->plotOn(mInvMassFrame, Name("Tot_c"), LineColor(kBlue));
       mSgnPdf->plotOn(mInvMassFrame, Normalization(1.0, RooAbsReal::RelativeExpected), DrawOption("F"), FillColor(TColor::GetColorTransparent(kBlue, 0.2)), VLines());
       mChiSquareOverNdfTotal = mInvMassFrame->chiSquare("Tot_c", "data_c"); // calculate reduced chi2 / DNF
+
       // plot residual distribution
       mResidualFrame = mass->frame(Title("Residual Distribution"));
       RooHist* residualHistogram = mInvMassFrame->residHist("data_c", "Bkg_c");
@@ -375,6 +376,7 @@ void HFInvMassFitter::doFit()
     calculateSignal(mRawYield, mRawYieldErr);
     countSignal(mRawYieldCounted, mRawYieldCountedErr);
     calculateSignificance(mSignificance, mSignificanceErr);
+    calculateFitToDataRatio();
   }
 }
 
@@ -640,6 +642,37 @@ void HFInvMassFitter::drawResidual(TVirtualPad* pad)
   highlightPeakRegion(mResidualFrame);
 }
 
+// draw ratio on canvas
+void HFInvMassFitter::drawRatio(TVirtualPad* pad)
+{
+  pad->cd();
+  mRatioFrame->GetYaxis()->SetTitle("");
+  TPaveText* textInfo = new TPaveText(0.12, 0.65, 0.47, .89, "NDC");
+  textInfo->SetBorderSize(0);
+  textInfo->SetFillStyle(0);
+  textInfo->SetTextColor(kBlue);
+  textInfo->AddText(Form("S = %.0f #pm %.0f ", mRawYield, mRawYieldErr));
+  textInfo->AddText(Form("S_{count} = %.0f #pm %.0f ", mRawYieldCounted, mRawYieldCountedErr));
+  textInfo->AddText(Form("mean = %.3f #pm %.3f", mRooMeanSgn->getVal(), mRooMeanSgn->getError()));
+  if (mTypeOfSgnPdf == DoubleGaus) {
+    auto const& baseSigmaSgn = mWorkspace->var("sigma");
+    textInfo->AddText(Form("sigma = %.3f #pm %.3f", baseSigmaSgn->getVal(), baseSigmaSgn->getError()));
+    textInfo->AddText(Form("sigma 2 = %.3f #pm %.3f", mRooSigmaSgn->getVal(), mRooSigmaSgn->getError()));
+  } else {
+    textInfo->AddText(Form("sigma = %.3f #pm %.3f", mRooSigmaSgn->getVal(), mRooSigmaSgn->getError()));
+  }
+  mRatioFrame->addObject(textInfo);
+  double xMin = mRatioFrame->GetXaxis()->GetXmin();
+  double xMax = mRatioFrame->GetXaxis()->GetXmax();
+  TLine* line = new TLine(xMin, 1.0, xMax, 1.0);
+  line->SetLineColor(kGray);
+  line->SetLineStyle(2);
+  line->SetLineWidth(2);
+  mRatioFrame->addObject(line);
+  mRatioFrame->Draw();
+  highlightPeakRegion(mRatioFrame);
+}
+
 // draw peak region with vertical lines
 void HFInvMassFitter::highlightPeakRegion(const RooPlot* plot, Color_t color, Width_t width, Style_t style) const
 {
@@ -820,6 +853,42 @@ void HFInvMassFitter::plotRefl(RooAbsPdf* pdf)
   pdf->plotOn(mInvMassFrame, Components(namesOfReflPdf.at(mTypeOfReflPdf).c_str()), Name("Refl_c"), LineColor(kGreen));
 }
 
+// Calculate fit to data ratio
+void HFInvMassFitter::calculateFitToDataRatio()
+{
+  if (!mInvMassFrame)
+    return;
+  // Create a new RooPlot for the ratio
+  mRatioFrame = mWorkspace->var("mass")->frame(Title("Fit/Data Ratio"));
+
+  // Get the data and fit curves from the frame
+  auto* dataHist = dynamic_cast<RooHist*>(mInvMassFrame->findObject("data_c"));
+  auto* fitCurve = dynamic_cast<RooCurve*>(mInvMassFrame->findObject("Tot_c")); // or the relevant fit curve
+
+  if (!dataHist || !fitCurve)
+    return;
+
+  RooHist* ratioHist = new RooHist();
+
+  for (int i = 0; i < dataHist->GetN(); ++i) {
+    double x, dataY, dataErr;
+    dataHist->GetPoint(i, x, dataY);
+    dataErr = dataHist->GetErrorY(i);
+
+    double fitY = fitCurve->Eval(x);
+
+    double ratio = dataY != 0 ? fitY / dataY : 0;
+    double err = dataY != 0 ? ratio * dataErr / dataY : 0;
+
+    ratioHist->SetPoint(i, x, ratio);
+    ratioHist->SetPointError(i, 0, 0, err, err);
+  }
+
+  mRatioFrame->SetMinimum(0.0);
+  mRatioFrame->SetMaximum(2.0);
+  mRatioFrame->addPlotable(ratioHist, "P");
+}
+
 // Fix reflection pdf
 void HFInvMassFitter::setReflFuncFixed()
 {

PWGHF/D2H/Macros/HFInvMassFitter.h

@@ -224,8 +224,10 @@ class HFInvMassFitter : public TNamed
   void calculateBackground(Double_t& bkg, Double_t& bkgErr) const;
   void calculateSignificance(Double_t& significance, Double_t& significanceErr) const;
   void checkForSignal(Double_t& estimatedSignal);
+  void calculateFitToDataRatio();
   void drawFit(TVirtualPad* c, Int_t writeFitInfo = 2);
   void drawResidual(TVirtualPad* c);
+  void drawRatio(TVirtualPad* c);
   void drawReflection(TVirtualPad* c);
 
  private:
@@ -293,6 +295,7 @@ class HFInvMassFitter : public TNamed
   RooPlot* mReflFrame;               /// reflection frame
   RooPlot* mReflOnlyFrame;           /// reflection frame plot on reflection only
   RooPlot* mResidualFrame;           /// residual frame
+  RooPlot* mRatioFrame;              /// fit/data ratio frame
   RooPlot* mResidualFrameForCalculation;
   RooWorkspace* mWorkspace;    /// workspace
   Double_t mIntegralHisto;     /// integral of histogram to fit

PWGHF/D2H/Macros/runMassFitter.C

@@ -346,6 +346,7 @@ int runMassFitter(const TString& configFileName)
   const Int_t nCanvases = std::ceil(static_cast<float>(nSliceVarBins) / nCanvasesMax);
   std::vector<TCanvas*> canvasMass(nCanvases);
   std::vector<TCanvas*> canvasResiduals(nCanvases);
+  std::vector<TCanvas*> canvasRatio(nCanvases);
   std::vector<TCanvas*> canvasRefl(nCanvases);
   for (int iCanvas = 0; iCanvas < nCanvases; iCanvas++) {
     const int nPads = (nCanvases == 1) ? nSliceVarBins : nCanvasesMax;
@@ -356,9 +357,16 @@ int runMassFitter(const TString& configFileName)
     canvasResiduals[iCanvas] =
       new TCanvas(Form("canvasResiduals%d", iCanvas), Form("canvasResiduals%d", iCanvas), canvasSize[0], canvasSize[1]);
     divideCanvas(canvasResiduals[iCanvas], nPads);
-    canvasRefl[iCanvas] = new TCanvas(Form("canvasRefl%d", iCanvas), Form("canvasRefl%d", iCanvas),
-                                      canvasSize[0], canvasSize[1]);
-    divideCanvas(canvasRefl[iCanvas], nPads);
+
+    canvasRatio[iCanvas] = new TCanvas(Form("canvasRatio%d", iCanvas), Form("canvasRatio%d", iCanvas),
+                                       canvasSize[0], canvasSize[1]);
+    divideCanvas(canvasRatio[iCanvas], nPads);
+
+    if (enableRefl) {
+      canvasRefl[iCanvas] = new TCanvas(Form("canvasRefl%d", iCanvas), Form("canvasRefl%d", iCanvas),
+                                        canvasSize[0], canvasSize[1]);
+      divideCanvas(canvasRefl[iCanvas], nPads);
+    }
   }
 
   for (unsigned int iSliceVar = 0; iSliceVar < nSliceVarBins; iSliceVar++) {
@@ -533,6 +541,15 @@ int runMassFitter(const TString& configFileName)
       massFitter->drawResidual(gPad);
       canvasResiduals[iCanvas]->Modified();
       canvasResiduals[iCanvas]->Update();
+
+      if (nSliceVarBins > 1) {
+        canvasRatio[iCanvas]->cd(iSliceVar - nCanvasesMax * iCanvas + 1);
+      } else {
+        canvasRatio[iCanvas]->cd();
+      }
+      massFitter->drawRatio(gPad);
+      canvasRatio[iCanvas]->Modified();
+      canvasRatio[iCanvas]->Update();
     }
 
     hFitConfig->SetBinContent(1, iSliceVar + 1, massMin[iSliceVar]);
@@ -555,7 +572,10 @@ int runMassFitter(const TString& configFileName)
     canvasMass[iCanvas]->Write();
     if (!isMc) {
       canvasResiduals[iCanvas]->Write();
-      canvasRefl[iCanvas]->Write();
+      canvasRatio[iCanvas]->Write();
+      if (enableRefl) {
+        canvasRefl[iCanvas]->Write();
+      }
     }
   }
 
@@ -614,13 +634,9 @@ void setHistoStyle(TH1* histo, Color_t color, Size_t markerSize)
 
 void divideCanvas(TCanvas* canvas, int nSliceVarBins)
 {
-  const int rectangularSideMin = std::floor(std::sqrt(nSliceVarBins));
-  constexpr int RectangularSidesDiffMax = 2;
-  for (int rectangularSidesDiff = 0; rectangularSidesDiff < RectangularSidesDiffMax; ++rectangularSidesDiff) {
-    if (rectangularSideMin * (rectangularSideMin + rectangularSidesDiff) >= nSliceVarBins) {
-      canvas->Divide(rectangularSideMin + rectangularSidesDiff, rectangularSideMin);
-    }
-  }
+  int nCols = std::ceil(std::sqrt(nSliceVarBins));
+  int nRows = std::ceil(double(nSliceVarBins) / nCols);
+  canvas->Divide(nCols, nRows);
 }
 
 int main(int argc, char* argv[])