[RF] Added test and correction factor in RooMultiPdf

GalinBistrev2 · guitargeek · commit 06cb8baab343 · 2025-07-21T14:14:44.000+02:00
diff --git a/roofit/roofit/CMakeLists.txt b/roofit/roofit/CMakeLists.txt
@@ -16,7 +16,6 @@ endif()
 
 ROOT_STANDARD_LIBRARY_PACKAGE(RooFit
   HEADERS
-    RooMultiPdf.h
     Roo2DKeysPdf.h
     RooArgusBG.h
     RooBCPEffDecay.h
@@ -82,7 +81,6 @@ ROOT_STANDARD_LIBRARY_PACKAGE(RooFit
     RooVoigtian.h
     RooJohnson.h
   SOURCES
-    src/RooMultiPdf.cxx
     src/Roo2DKeysPdf.cxx
     src/RooArgusBG.cxx
     src/RooBCPEffDecay.cxx
diff --git a/roofit/roofit/inc/LinkDef1.h b/roofit/roofit/inc/LinkDef1.h
@@ -3,7 +3,6 @@
 #pragma link off all classes;
 #pragma link off all functions;
 #pragma link C++ class Roo2DKeysPdf+ ;
-#pragma link C++ class RooMultiPdf + ;
 #pragma link C++ class RooArgusBG+ ;
 #pragma link C++ class RooBCPEffDecay+ ;
 #pragma link C++ class RooBCPGenDecay+ ;
@@ -59,6 +58,7 @@
 #pragma link C++ class RooSpline+ ;
 #pragma link C++ class RooStepFunction+ ;
 #pragma link C++ class RooMultiBinomial+ ;
+
 /* #pragma link C++ class std::vector< TVector2 >; */
 /* #pragma link C++ class std::vector< TVector2 >::iterator ; */
 /* #pragma link C++ class RooPolyMorph2D+ ; */
diff --git a/roofit/roofit/test/CMakeLists.txt b/roofit/roofit/test/CMakeLists.txt
@@ -20,7 +20,7 @@ ROOT_ADD_GTEST(testRooLandau testRooLandau.cxx LIBRARIES RooFit)
 ROOT_ADD_GTEST(testRooParamHistFunc testRooParamHistFunc.cxx LIBRARIES Gpad RooFit)
 ROOT_ADD_GTEST(testRooPoisson testRooPoisson.cxx LIBRARIES RooFit)
 ROOT_ADD_GTEST(testRooStepFunc testRooStepFunc.cxx LIBRARIES RooFit)
-ROOT_ADD_GTEST(testRooMultiPdf testRooMultiPdf.cxx LIBRARIES RooFit)
+
 if(mathmore)
   ROOT_EXECUTABLE(testRooFit testRooFit.cxx LIBRARIES RooFit MathMore)
   ROOT_ADD_TEST(test-fit-testRooFit COMMAND testRooFit)
diff --git a/roofit/roofit/test/testRooMultiPdf b/roofit/roofit/test/testRooMultiPdf
diff --git a/roofit/roofitcore/CMakeLists.txt b/roofit/roofitcore/CMakeLists.txt
@@ -171,6 +171,7 @@ ROOT_STANDARD_LIBRARY_PACKAGE(RooFitCore
     RooMsgService.h
     RooMultiCategory.h
     RooMultiVarGaussian.h
+    RooMultiPdf.h
     RooNameReg.h
     RooNormSetCache.h
     RooNumCdf.h
@@ -366,6 +367,7 @@ ROOT_STANDARD_LIBRARY_PACKAGE(RooFitCore
     src/RooMsgService.cxx
     src/RooMultiCategory.cxx
     src/RooMultiVarGaussian.cxx
+    src/RooMultiPdf.cxx
     src/RooNLLVarNew.cxx
     src/RooNameReg.cxx
     src/RooNormSetCache.cxx
diff --git a/roofit/roofitcore/inc/LinkDef.h b/roofit/roofitcore/inc/LinkDef.h
@@ -133,6 +133,7 @@
 #pragma read sourceClass="RooMappedCategory::Entry" targetClass="RooMappedCategory::Entry" version="[1]" include="RooFitLegacy/RooCatTypeLegacy.h" \
     source="RooCatType _cat" target="_catIdx" code="{ _catIdx = onfile._cat.getVal(); }"
 #pragma link C++ class RooMultiCategory+ ;
+#pragma link C++ class RooMultiPdf+ ;
 #pragma link off class RooNameReg+ ;
 #pragma link C++ class RooNumConvolution+ ;
 #pragma link C++ class RooNumConvPdf+ ;
diff --git a/roofit/roofitcore/inc/RooAbsPdf.h b/roofit/roofitcore/inc/RooAbsPdf.h
@@ -177,6 +177,8 @@ class RooAbsPdf : public RooAbsReal {
   // Project p.d.f into lower dimensional p.d.f
   virtual RooAbsPdf* createProjection(const RooArgSet& iset) ;
 
+  virtual double getCorrection() const;
+
   // Create cumulative density function from p.d.f
   RooFit::OwningPtr<RooAbsReal> createCdf(const RooArgSet& iset, const RooArgSet& nset=RooArgSet()) ;
   RooFit::OwningPtr<RooAbsReal> createCdf(const RooArgSet& iset, const RooCmdArg& arg1, const RooCmdArg& arg2={},
diff --git a/roofit/roofitcore/inc/RooFit/Detail/RooNormalizedPdf.h b/roofit/roofitcore/inc/RooFit/Detail/RooNormalizedPdf.h
@@ -47,6 +47,8 @@ class RooNormalizedPdf : public RooAbsPdf {
 
    bool selfNormalized() const override { return true; }
 
+   inline double getCorrection() const override { return _pdf->getCorrection(); }
+
    bool forceAnalyticalInt(const RooAbsArg & /*dep*/) const override { return true; }
    /// Forward determination of analytical integration capabilities to input p.d.f
    Int_t getAnalyticalIntegralWN(RooArgSet &allVars, RooArgSet &analVars, const RooArgSet * /*normSet*/,
diff --git a/roofit/roofitcore/inc/RooMultiPdf.h b/roofit/roofitcore/inc/RooMultiPdf.h
@@ -18,7 +18,7 @@ class RooMultiPdf : public RooAbsPdf {
    TObject *clone(const char *newname) const override { return new RooMultiPdf(*this, newname); }
 
    inline bool checkIndexDirty() const { return _oldIndex != x; }
-   inline double getCorrection() const { return cFactor * static_cast<RooAbsReal *>(corr.at(x))->getVal(); }
+   inline double getCorrection() const override { return cFactor * static_cast<RooAbsReal *>(corr.at(x))->getVal(); }
    inline RooAbsPdf *getCurrentPdf() const { return getPdf(getCurrentIndex()); }
    int getNumPdfs() const { return c.size(); }
 
diff --git a/roofit/roofitcore/src/FitHelpers.cxx b/roofit/roofitcore/src/FitHelpers.cxx
@@ -33,6 +33,7 @@
 #include <RooFuncWrapper.h>
 #include <RooLinkedList.h>
 #include <RooMinimizer.h>
+#include <RooConstVar.h>
 #include <RooRealVar.h>
 #include <RooSimultaneous.h>
 #include <RooFormulaVar.h>
@@ -810,8 +811,26 @@ std::unique_ptr<RooAbsReal> createNLL(RooAbsPdf &pdf, RooAbsData &data, const Ro
             takeGlobalObservablesFromData);
       }
 
+      const double correction = pdfClone->getCorrection();
       nllWrapper->addOwnedComponents(std::move(nll));
       nllWrapper->addOwnedComponents(std::move(pdfClone));
+
+      if (correction > 0) {
+         oocoutI(&pdf, Fitting) << "[FitHelpers] Detected correction term from RooAbsPdf::getCorrection(). "
+                                << "Adding penalty to NLL." << std::endl;
+
+         // Convert the multiplicative correction to an additive term in -log L
+         auto penaltyTerm = std::make_unique<RooConstVar>((baseName + "_Penalty").c_str(),
+                                                          "Penalty term from getCorrection()", correction);
+
+         auto correctedNLL = std::make_unique<RooAddition>(
+            // add penalty and NLL
+            (baseName + "_corrected").c_str(), "NLL + penalty", RooArgSet(*nllWrapper, *penaltyTerm));
+
+         // transfer ownership of terms
+         correctedNLL->addOwnedComponents(std::move(nllWrapper), std::move(penaltyTerm));
+         nllWrapper = std::move(correctedNLL);
+      }
       return nllWrapper;
    }
 
@@ -888,6 +907,23 @@ std::unique_ptr<RooAbsReal> createNLL(RooAbsPdf &pdf, RooAbsData &data, const Ro
    throw std::runtime_error("RooFit was not built with the legacy evaluation backend");
 #endif
 
+   if (const double correction = pdf.getCorrection(); correction > 0) {
+      oocoutI(&pdf, Fitting) << "[FitHelpers] Detected correction term from RooAbsPdf::getCorrection(). "
+                             << "Adding penalty to NLL." << std::endl;
+
+      // Convert the multiplicative correction to an additive term in -log L
+      auto penaltyTerm = std::make_unique<RooConstVar>((baseName + "_Penalty").c_str(),
+                                                       "Penalty term from getCorrection()", correction);
+
+      auto correctedNLL = std::make_unique<RooAddition>(
+         // add penalty and NLL
+         (baseName + "_corrected").c_str(), "NLL + penalty", RooArgSet(*nll, *penaltyTerm));
+
+      // transfer ownership of terms
+      correctedNLL->addOwnedComponents(std::move(nll), std::move(penaltyTerm));
+      nll = std::move(correctedNLL);
+   }
+
    return nll;
 }
 
diff --git a/roofit/roofitcore/src/RooAbsPdf.cxx b/roofit/roofitcore/src/RooAbsPdf.cxx
@@ -620,6 +620,12 @@ double RooAbsPdf::getLogVal(const RooArgSet* nset) const
   return getLog(getVal(nset), this);
 }
 
+/// Return correction factor
+
+double RooAbsPdf::getCorrection() const
+{
+   return 0;
+}
 
 ////////////////////////////////////////////////////////////////////////////////
 /// Check for infinity or NaN.
diff --git a/roofit/roofitcore/src/RooFuncWrapper.cxx b/roofit/roofitcore/src/RooFuncWrapper.cxx
@@ -142,7 +142,6 @@ RooFuncWrapper::loadParamsAndData(RooArgSet const &paramSet, const RooAbsData *d
       idx += n;
    }
 
-   // Extract parameters
    for (auto *param : paramSet) {
       if (!dynamic_cast<RooAbsReal *>(param)) {
          if (param->isConstant()) {
diff --git a/roofit/roofitcore/src/RooMultiPdf.cxx b/roofit/roofitcore/src/RooMultiPdf.cxx
diff --git a/roofit/roofitcore/test/CMakeLists.txt b/roofit/roofitcore/test/CMakeLists.txt
@@ -20,6 +20,7 @@ ROOT_ADD_GTEST(testRooWrapperPdf testRooWrapperPdf.cxx LIBRARIES Gpad RooFitCore
 ROOT_ADD_GTEST(testGenericPdf testGenericPdf.cxx LIBRARIES RooFitCore)
 ROOT_ADD_GTEST(testRooAbsPdf testRooAbsPdf.cxx LIBRARIES RooFitCore)
 ROOT_ADD_GTEST(testRooAbsCollection testRooAbsCollection.cxx LIBRARIES RooFitCore)
+
 ROOT_ADD_GTEST(testRooDataSet testRooDataSet.cxx LIBRARIES Tree RooFitCore
   COPY_TO_BUILDDIR ${CMAKE_CURRENT_SOURCE_DIR}/dataSet_with_errors_6_26_10.root)
 ROOT_ADD_GTEST(testRooFormula testRooFormula.cxx LIBRARIES RooFitCore ROOT::TestSupport)
@@ -59,6 +60,7 @@ ROOT_ADD_GTEST(testRooExtendedBinding testRooExtendedBinding.cxx LIBRARIES RooFi
 ROOT_ADD_GTEST(testRooHist testRooHist.cxx LIBRARIES RooFitCore)
 ROOT_ADD_GTEST(testRooHistPdf testRooHistPdf.cxx LIBRARIES RooFitCore)
 ROOT_ADD_GTEST(testRooMinimizer testRooMinimizer.cxx LIBRARIES RooFitCore RooFit)
+ROOT_ADD_GTEST(testRooMultiPdf testRooMultiPdf.cxx LIBRARIES RooFitCore RooFit)
 ROOT_ADD_GTEST(testRooPolyFunc testRooPolyFunc.cxx LIBRARIES Gpad RooFitCore)
 ROOT_ADD_GTEST(testRooRealL TestStatistics/testRooRealL.cxx LIBRARIES RooFitCore)
 ROOT_ADD_GTEST(testRooRombergIntegrator testRooRombergIntegrator.cxx LIBRARIES MathCore RooFitCore)
@@ -67,6 +69,7 @@ ROOT_ADD_GTEST(testRooSimultaneous testRooSimultaneous.cxx LIBRARIES RooFitCore)
 ROOT_ADD_GTEST(testRooTruthModel testRooTruthModel.cxx LIBRARIES RooFitCore RooFit)
 ROOT_ADD_GTEST(testSumW2Error testSumW2Error.cxx LIBRARIES Gpad RooFitCore)
 ROOT_ADD_GTEST(testTestStatistics testTestStatistics.cxx LIBRARIES RooFitCore)
+
 if (roofit_multiprocess)
   ROOT_ADD_GTEST(testTestStatisticsPlot TestStatistics/testPlot.cxx LIBRARIES RooFitMultiProcess RooFitCore RooFit
                    COPY_TO_BUILDDIR ${CMAKE_CURRENT_SOURCE_DIR}/TestStatistics/TestStatistics_ref.root)
diff --git a/roofit/roofitcore/test/testRooMultiPdf.cxx b/roofit/roofitcore/test/testRooMultiPdf.cxx
@@ -1,4 +1,5 @@
 #include <RooCategory.h>
+#include <RooConstVar.h>
 #include <RooGaussian.h>
 #include <RooMultiPdf.h>
 #include <RooRealVar.h>
@@ -8,6 +9,20 @@
 #include <gtest/gtest.h>
 
 #include <memory>
+// Helper function to count parameters including x
+int countFloatingParametersIncludingObservable(const RooAbsPdf &pdf)
+{
+   std::unique_ptr<RooArgSet> params{pdf.getParameters(RooArgSet())};
+   int count = 0;
+   for (auto *obj : *params) {
+      auto *var = dynamic_cast<RooRealVar *>(obj);
+      if (var && !var->isConstant()) {
+         count++;
+      }
+   }
+
+   return count;
+}
 
 TEST(RooMultiPdf, SelectsCorrectPdf)
 {
@@ -24,7 +39,6 @@ TEST(RooMultiPdf, SelectsCorrectPdf)
 
    RooCategory indx("my_special_index", "my_index");
 
-
    RooArgList list{gaus1, gaus2};
 
    RooMultiPdf pdf("mult", "multi_pdf", indx, list);
@@ -39,9 +53,10 @@ TEST(RooMultiPdf, SelectsCorrectPdf)
 }
 
 TEST(RooMultiPdfTest, FitConvergesAndReturnsReasonableResult)
+
 {
    RooRealVar x("x", "x", -10, 10);
-   RooRealVar m1("mean1", "mean1", 0.);
+   RooRealVar m1("mean1", "mean1", 0., -10, 10);
    RooRealVar s1("sigma1", "sigma1", 1., 0.001, 10.);
    RooGaussian gaus1("gaus1", "gaus1", x, m1, s1);
 
@@ -83,15 +98,20 @@ TEST(RooMultiPdfTest, FitConvergesAndReturnsReasonableResult)
    // Fit 2 - Reference fit
    std::unique_ptr<RooAbsReal> nll1{gaus1.createNLL(*data, RooFit::EvalBackend("codegen"))};
 
-   RooMinimizer minim1{*nll};
+   RooMinimizer minim1{*nll1};
    minim1.setStrategy(0);
    int status1 = minim1.minimize("Minuit2", "");
 
+   int n_param_gaus1 = countFloatingParametersIncludingObservable(gaus1);
+
+   double first_fit = nll->getVal();
+   double ref_fit = nll1->getVal() + 0.5 * n_param_gaus1; // 1.5 because the gaussian has 3 param*0.5
    // Now test the results
    EXPECT_EQ(status, 0) << "Fit 1 did not converge.";
    EXPECT_EQ(status1, 0) << "Fit 2 did not converge.";
-   EXPECT_TRUE(std::isfinite(nll->getVal())) << "NLL is not finite.";
-   EXPECT_DOUBLE_EQ(nll->getVal(), nll1->getVal());
+   EXPECT_TRUE(std::isfinite(first_fit)) << "NLL is not finite.";
+   EXPECT_TRUE(std::isfinite(ref_fit)) << "NLL1 is not finite.";
+   EXPECT_DOUBLE_EQ(first_fit, ref_fit);
 
    // Check that the correct number of PDFs are present
    EXPECT_EQ(multipdf.getNumPdfs(), 3);
@@ -101,6 +121,7 @@ TEST(RooMultiPdfTest, FitConvergesAndReturnsReasonableResult)
    EXPECT_NEAR(s1.getVal(), 1.0, 0.2);
 
    // Check whether RooMultiPdf chooses the correct index
+
    for (int i = 0; i < multipdf.getNumPdfs(); ++i) {
 
       indx.setIndex(i);
@@ -110,9 +131,58 @@ TEST(RooMultiPdfTest, FitConvergesAndReturnsReasonableResult)
       RooAbsPdf *selectedPdf = multipdf.getPdf(i);
       std::unique_ptr<RooAbsReal> nll_direct{selectedPdf->createNLL(*data, RooFit::EvalBackend("codegen"))};
 
+      int n_param = countFloatingParametersIncludingObservable(*selectedPdf);
+
       double multi = nll_multi->getVal();
-      double direct = nll_direct->getVal();
+      double direct = nll_direct->getVal() + 0.5 * n_param;
 
-      EXPECT_NEAR(multi, direct, 1e-6) << "Mismatch" << i;
+      std::cout << "PDF index " << i << ": n_param = " << n_param << ", direct+penalty = " << direct
+                << ", multipdf = " << multi << std::endl;
+
+      EXPECT_NEAR(multi, direct, 1e-6) << "Mismatch at index " << i;
    }
+}
+TEST(RooMultiPdfTest, PenaltyTermIsAppliedCorrectly)
+{
+   using namespace RooFit;
+
+   RooRealVar x("x", "x", -10, 10);
+
+   RooRealVar mean("mean", "mean", 0, -5, 5);
+   RooRealVar sigma("sigma", "sigma", 1, 0.1, 10);
+
+   RooGaussian gauss1("gauss1", "gauss1", x, mean, sigma);
+   RooGaussian gauss2("gauss2", "gauss2", x, mean, sigma);
+
+   RooCategory index("index", "index");
+
+   RooArgList pdfList(gauss1, gauss2);
+
+   RooMultiPdf multiPdf("multiPdf", "multiPdf", index, pdfList);
+   index.setConstant();
+
+   index.setIndex(0);
+
+   std::unique_ptr<RooDataSet> data{gauss1.generate(x, 100)};
+
+   std::unique_ptr<RooAbsReal> nll_gauss1{gauss1.createNLL(*data, EvalBackend("codegen"))};
+
+   std::unique_ptr<RooAbsReal> nll_multi{multiPdf.createNLL(*data, EvalBackend("codegen"))};
+
+   double val_gauss1 = nll_gauss1->getVal();
+   double val_multi = nll_multi->getVal();
+   int n_params = countFloatingParametersIncludingObservable(gauss1);
+
+   const double expected_penalty = 0.5 * n_params;
+   const double delta = val_multi - val_gauss1;
+
+   std::cout << "NLL(gauss1):     " << val_gauss1 << std::endl;
+   std::cout << "NLL(multiPdf):   " << val_multi << std::endl;
+   std::cout << "Expected penalty: " << expected_penalty << std::endl;
+   std::cout << "Delta:           " << delta << std::endl;
+
+   EXPECT_TRUE(std::isfinite(val_gauss1));
+   EXPECT_TRUE(std::isfinite(val_multi));
+
+   EXPECT_NEAR(delta, expected_penalty, 1e-6) << "Penalty term not correctly applied.";
 }
