[RF] Implement SumW2 correction in new BatchMode with RooFitDriver

This has not been implemented so far. This commit also includes a unit
test for it.

For easier toggling of squared weights, a new virtual function
`RooAbsArg::applyWeightsSquared` was introduced such that one doesn't
have to pick up manully the likelihood classes from the computation
graph when applying the weights squared correction.

Author: guitargeek

roofit/roofitcore/inc/RooAbsArg.h

@@ -599,6 +599,8 @@ class RooAbsArg : public TNamed, public RooPrintable {
   virtual bool canComputeBatchWithCuda() const { return false; }
   virtual bool isReducerNode() const { return false; }
 
+  virtual void applyWeightSquared(bool flag);
+
   operator RooBatchCompute::DataKey() const { return RooBatchCompute::DataKey::create(this->namePtr()); }
 
 protected:

roofit/roofitcore/inc/RooAbsPdf.h

@@ -406,7 +406,7 @@ class RooAbsPdf : public RooAbsReal {
 
 private:
   int calcAsymptoticCorrectedCovariance(RooMinimizer& minimizer, RooAbsData const& data);
-  int calcSumW2CorrectedCovariance(RooMinimizer& minimizer, RooAbsReal const& nll) const;
+  int calcSumW2CorrectedCovariance(RooMinimizer& minimizer, RooAbsReal & nll) const;
 
   ClassDefOverride(RooAbsPdf,5) // Abstract PDF with normalization support
 };

roofit/roofitcore/inc/RooNLLVar.h

@@ -53,7 +53,7 @@ class RooNLLVar : public RooAbsOptTestStatistic {
 
   ~RooNLLVar() override;
 
-  void applyWeightSquared(Bool_t flag) ;
+  void applyWeightSquared(bool flag) override;
 
   Double_t defaultErrorLevel() const override { return 0.5 ; }
 

roofit/roofitcore/res/RooFitDriver.h

@@ -127,6 +127,11 @@ class RooFitDriver {
 
       double getValV(const RooArgSet *) const override { return evaluate(); }
 
+      void applyWeightSquared(bool flag) override
+      {
+         const_cast<RooAbsReal &>(_driver->topNode()).applyWeightSquared(flag);
+      }
+
    protected:
       double evaluate() const override { return _driver ? _driver->getVal() : 0.0; }
 

roofit/roofitcore/res/RooNLLVarNew.h

@@ -26,6 +26,11 @@ namespace Experimental {
 class RooNLLVarNew : public RooAbsReal {
 
 public:
+   // The names for the weight variables that the RooNLLVarNew expects
+   static constexpr auto weightVarName = "_weight";
+   static constexpr auto weightVarNameSumW2 = "_weight_sumW2";
+   static constexpr auto weightVarNameSumW2Suffix = "_sumW2";
+
    RooNLLVarNew(){};
    RooNLLVarNew(const char *name, const char *title, RooAbsPdf &pdf, RooArgSet const &observables, RooAbsReal *weight,
                 bool isExtended, std::string const &rangeName);
@@ -41,13 +46,17 @@ class RooNLLVarNew : public RooAbsReal {
    void computeBatch(cudaStream_t *, double *output, size_t nOut, RooBatchCompute::DataMap &) const override;
    inline bool isReducerNode() const override { return true; }
 
+   RooArgSet prefixObservableAndWeightNames(std::string const &prefix);
+
+   void applyWeightSquared(bool flag) override;
+
+protected:
    void setObservables(RooArgSet const &observables)
    {
       _observables.clear();
       _observables.add(observables);
    }
 
-protected:
    RooTemplateProxy<RooAbsPdf> _pdf;
    RooArgSet _observables;
    std::unique_ptr<RooTemplateProxy<RooAbsReal>> _weight;

roofit/roofitcore/src/BatchModeHelpers.cxx

@@ -27,6 +27,9 @@
 
 #include <string>
 
+using ROOT::Experimental::RooFitDriver;
+using ROOT::Experimental::RooNLLVarNew;
+
 namespace {
 
 std::unique_ptr<RooAbsArg> prepareSimultaneousModelForBatchMode(RooSimultaneous &simPdf, RooArgSet &observables,
@@ -40,8 +43,8 @@ std::unique_ptr<RooAbsArg> prepareSimultaneousModelForBatchMode(RooSimultaneous
       auto const &catName = catItem.first;
       auto *pdf = simPdf.getPdf(catName.c_str());
       auto nllName = std::string("nll_") + pdf->GetName();
-      nllTerms.add(*new ROOT::Experimental::RooNLLVarNew(nllName.c_str(), nllName.c_str(), *pdf, observables, weight,
-                                                         isExtended, rangeName));
+      nllTerms.add(
+         *new RooNLLVarNew(nllName.c_str(), nllName.c_str(), *pdf, observables, weight, isExtended, rangeName));
    }
 
    RooArgSet newObservables;
@@ -50,22 +53,8 @@ std::unique_ptr<RooAbsArg> prepareSimultaneousModelForBatchMode(RooSimultaneous
    std::size_t iNLL = 0;
    for (auto const &catItem : simPdf.indexCat()) {
       auto const &catName = catItem.first;
-      auto &nll = nllTerms[iNLL];
-      RooArgSet pdfObs;
-      nll.getObservables(&observables, pdfObs);
-      if (weight)
-         pdfObs.add(*weight);
-      RooArgSet obsClones;
-      pdfObs.snapshot(obsClones);
-      for (RooAbsArg *arg : obsClones) {
-         auto newName = std::string("_") + catName + "_" + arg->GetName();
-         arg->setAttribute((std::string("ORIGNAME:") + arg->GetName()).c_str());
-         arg->SetName(newName.c_str());
-      }
-      nll.recursiveRedirectServers(obsClones, false, true);
-      newObservables.add(obsClones);
-      static_cast<ROOT::Experimental::RooNLLVarNew &>(nll).setObservables(obsClones);
-      nll.addOwnedComponents(std::move(obsClones));
+      auto &nll = static_cast<RooNLLVarNew &>(nllTerms[iNLL]);
+      newObservables.add(nll.prefixObservableAndWeightNames(std::string("_") + catName + "_"));
       ++iNLL;
    }
 
@@ -86,7 +75,7 @@ RooFit::BatchModeHelpers::createNLL(RooAbsPdf &pdf, RooAbsData &data, std::uniqu
 {
    std::unique_ptr<RooRealVar> weightVar;
 
-   std::unique_ptr<ROOT::Experimental::RooFitDriver> driver;
+   std::unique_ptr<RooFitDriver> driver;
 
    RooArgSet observables;
    pdf.getObservables(data.get(), observables);
@@ -104,16 +93,9 @@ RooFit::BatchModeHelpers::createNLL(RooAbsPdf &pdf, RooAbsData &data, std::uniqu
    }
 
    if (data.isWeighted()) {
-      std::string weightVarName = "_weight";
-      if (auto *dataSet = dynamic_cast<RooDataSet const *>(&data)) {
-         if (dataSet->weightVar())
-            weightVarName = dataSet->weightVar()->GetName();
-      }
-
-      // make a clone of the weight variable (or an initial instance, if it doesn't exist)
-      // the clone will hold the weight value (or values as a batch) and will participate
-      // in the computation graph of the RooFit driver.
-      weightVar = std::make_unique<RooRealVar>(weightVarName.c_str(), "Weight(s) of events", data.weight());
+      // RooRealVar for the weight value (or values as a batch) that will
+      // participate in the computation graph of the RooFit driver.
+      weightVar = std::make_unique<RooRealVar>(RooNLLVarNew::weightVarName, "Weight(s) of events", data.weight());
    }
 
    // Deal with the IntegrateBins argument
@@ -135,8 +117,8 @@ RooFit::BatchModeHelpers::createNLL(RooAbsPdf &pdf, RooAbsData &data, std::uniqu
       nllTerms.addOwned(
          prepareSimultaneousModelForBatchMode(*simPdfClone, observables, weightVar.get(), isExtended, rangeName));
    } else {
-      nllTerms.addOwned(std::make_unique<ROOT::Experimental::RooNLLVarNew>(
-         "RooNLLVarNew", "RooNLLVarNew", finalPdf, observables, weightVar.get(), isExtended, rangeName));
+      nllTerms.addOwned(std::make_unique<RooNLLVarNew>("RooNLLVarNew", "RooNLLVarNew", finalPdf, observables,
+                                                       weightVar.get(), isExtended, rangeName));
    }
    if (constraints) {
       nllTerms.addOwned(std::move(constraints));
@@ -151,10 +133,9 @@ RooFit::BatchModeHelpers::createNLL(RooAbsPdf &pdf, RooAbsData &data, std::uniqu
       RooArgSet parameters;
       pdf.getParameters(data.get(), parameters);
       nll->recursiveRedirectServers(parameters);
-      driver = std::make_unique<ROOT::Experimental::RooFitDriver>(data, *nll, observables, batchMode, rangeName,
-                                                                  &simPdf->indexCat());
+      driver = std::make_unique<RooFitDriver>(data, *nll, observables, batchMode, rangeName, &simPdf->indexCat());
    } else {
-      driver = std::make_unique<ROOT::Experimental::RooFitDriver>(data, *nll, observables, batchMode, rangeName);
+      driver = std::make_unique<RooFitDriver>(data, *nll, observables, batchMode, rangeName);
    }
 
    // Set the fitrange attribute so that RooPlot can automatically plot the fitting range by default
@@ -179,7 +160,7 @@ RooFit::BatchModeHelpers::createNLL(RooAbsPdf &pdf, RooAbsData &data, std::uniqu
       pdf.setStringAttribute("fitrange", fitrangeValue.c_str());
    }
 
-   auto driverWrapper = ROOT::Experimental::RooFitDriver::makeAbsRealWrapper(std::move(driver));
+   auto driverWrapper = RooFitDriver::makeAbsRealWrapper(std::move(driver));
    driverWrapper->addOwnedComponents(std::move(nll));
    if (weightVar)
       driverWrapper->addOwnedComponents(std::move(weightVar));

roofit/roofitcore/src/RooAbsArg.cxx

@@ -2545,3 +2545,12 @@ std::string printValue(RooAbsArg *raa)
    return s.str();
 }
 } // namespace cling
+
+
+/// Disables or enables the usage of squared weights. Needs to be overloaded in
+/// the likelihood classes for which this is relevant.
+void RooAbsArg::applyWeightSquared(bool flag) {
+   for(auto * server : servers()) {
+      server->applyWeightSquared(flag);
+   }
+}

roofit/roofitcore/src/RooAbsPdf.cxx

@@ -1309,31 +1309,17 @@ int RooAbsPdf::calcAsymptoticCorrectedCovariance(RooMinimizer &minimizer, RooAbs
 ///            matrix caltulated here will be applied to it via
 ///            RooMinimizer::applyCovarianceMatrix().
 /// \param[in] nll The NLL object that was used for the fit.
-int RooAbsPdf::calcSumW2CorrectedCovariance(RooMinimizer &minimizer, RooAbsReal const &nll) const
+int RooAbsPdf::calcSumW2CorrectedCovariance(RooMinimizer &minimizer, RooAbsReal &nll) const
 {
-
-   // Make list of RooNLLVar components of FCN
-   std::vector<RooNLLVar *> nllComponents;
-   std::unique_ptr<RooArgSet> comps{nll.getComponents()};
-   for (auto const &arg : *comps) {
-      if (RooNLLVar *nllComp = dynamic_cast<RooNLLVar *>(arg)) {
-         nllComponents.push_back(nllComp);
-      }
-   }
-
    // Calculated corrected errors for weighted likelihood fits
    std::unique_ptr<RooFitResult> rw{minimizer.save()};
-   for (auto &comp : nllComponents) {
-      comp->applyWeightSquared(true);
-   }
+   nll.applyWeightSquared(true);
    coutI(Fitting) << "RooAbsPdf::fitTo(" << this->GetName()
                   << ") Calculating sum-of-weights-squared correction matrix for covariance matrix"
                   << std::endl;
    minimizer.hesse();
    std::unique_ptr<RooFitResult> rw2{minimizer.save()};
-   for (auto &comp : nllComponents) {
-      comp->applyWeightSquared(false);
-   }
+   nll.applyWeightSquared(false);
 
    // Apply correction matrix
    const TMatrixDSym &matV = rw->covarianceMatrix();

roofit/roofitcore/src/RooFitDriver.cxx

@@ -40,6 +40,7 @@ and gets destroyed when the fitting ends.
 #include <RooRealVar.h>
 #include <RooSimultaneous.h>
 #include <RooBatchCompute/Initialisation.h>
+#include <RooBatchCompute/DataKey.h>
 
 #include <ROOT/StringUtils.hxx>
 
@@ -48,6 +49,19 @@ and gets destroyed when the fitting ends.
 #include <thread>
 #include <unordered_set>
 
+namespace {
+
+// Little wrapper to use a TNamed directly as a RooBatchCompute DataKey
+class NamePtrWrapper {
+public:
+   NamePtrWrapper(TNamed const *namePtr) : _namePtr(namePtr) {}
+   operator RooBatchCompute::DataKey() const { return RooBatchCompute::DataKey::create(_namePtr); }
+
+private:
+   TNamed const *_namePtr;
+};
+} // namespace
+
 namespace ROOT {
 namespace Experimental {
 
@@ -66,6 +80,7 @@ RooFitDriver::Dataset::Dataset(RooAbsData const &data, RooArgSet const &observab
                                RooAbsCategory const *indexCat)
    : _nEvents{static_cast<size_t>(data.numEntries())}
 {
+   auto &nameReg = RooNameReg::instance();
 
    // fill the RunContext with the observable data and map the observables
    // by namePtr in order to replace their memory addresses later, with
@@ -80,7 +95,7 @@ RooFitDriver::Dataset::Dataset(RooAbsData const &data, RooArgSet const &observab
    // the data map
    for (auto const &item : data.getCategoryBatches(0, _nEvents)) {
 
-      const TNamed *namePtr = RooNameReg::instance().constPtr(item.first.c_str());
+      const TNamed *namePtr = nameReg.constPtr(item.first.c_str());
       RooSpan<const RooAbsCategory::value_type> intSpan{item.second};
 
       _buffers.emplace(_nEvents);
@@ -92,19 +107,16 @@ RooFitDriver::Dataset::Dataset(RooAbsData const &data, RooArgSet const &observab
       _dataSpans[namePtr] = RooSpan<const double>(buffer, _nEvents);
    }
 
-   // Check if there is a batch for weights and if it's already in the dataMap.
-   // If not, we need to put the batch and give as a key a RooRealVar* that has
-   // the same name as RooNLLVarNew's _weight proxy, so that it gets renamed like
-   // every other observable.
-   RooSpan<const double> weights = data.getWeightBatch(0, _nEvents);
-   if (!weights.empty()) {
-      std::string weightVarName = "_weight";
-      if (auto *dataSet = dynamic_cast<RooDataSet const *>(&data)) {
-         if (dataSet->weightVar())
-            weightVarName = dataSet->weightVar()->GetName();
+   // Add weights to the datamap. They should have the names expected by the
+   // RooNLLVarNew. We also add the sumW2 weights here under a different name,
+   // so we can apply the sumW2 correction by easily swapping the spans.
+   {
+      auto weight = data.getWeightBatch(0, _nEvents, /*sumW2=*/false);
+      auto weightSumW2 = data.getWeightBatch(0, _nEvents, /*sumW2=*/true);
+      if (!weight.empty()) {
+         _dataSpans[nameReg.constPtr(RooNLLVarNew::weightVarName)] = weight;
+         _dataSpans[nameReg.constPtr(RooNLLVarNew::weightVarNameSumW2)] = weightSumW2;
       }
-      const TNamed *pTNamed = RooNameReg::instance().constPtr(weightVarName.c_str());
-      _dataSpans[pTNamed] = weights;
    }
 
    // Now we have do do the range selection
@@ -210,8 +222,8 @@ RooFitDriver::RooFitDriver(const RooAbsData &data, const RooAbsReal &topNode, Ro
 }
 
 RooFitDriver::RooFitDriver(RooBatchCompute::RunContext const &data, const RooAbsReal &topNode, RooArgSet const &normSet)
-   : _name{topNode.GetName()}, _title{topNode.GetTitle()},
-     _batchMode{RooFit::BatchModeOption::Cpu}, _dataset{data}, _topNode{topNode}, _normSet{std::make_unique<RooArgSet>(normSet)}
+   : _name{topNode.GetName()}, _title{topNode.GetTitle()}, _batchMode{RooFit::BatchModeOption::Cpu}, _dataset{data},
+     _topNode{topNode}, _normSet{std::make_unique<RooArgSet>(normSet)}
 {
    init();
 }
@@ -263,12 +275,15 @@ void RooFitDriver::init()
       serverSet.add(*arg);
    }
 
+   for (auto const &span : _dataset.spans()) {
+      _dataMapCPU[NamePtrWrapper(span.first)] = span.second;
+   }
+
    for (RooAbsArg *arg : serverSet) {
       _orderedNodes.push_back(arg);
       auto &argInfo = _nodeInfos[arg];
-      if (_dataset.contains(arg)) {
-         _dataMapCPU[arg] = _dataset.span(arg);
-         argInfo.outputSize = _dataset.span(arg).size();
+      if (_dataMapCPU.count(arg) > 0) {
+         argInfo.outputSize = _dataMapCPU[arg].size();
       }
 
       for (auto *client : arg->clients()) {

roofit/roofitcore/src/RooNLLVarNew.cxx

@@ -221,3 +221,66 @@ void RooNLLVarNew::getParametersHook(const RooArgSet * /*nset*/, RooArgSet *para
    if (_weight)
       params->remove(**_weight, true, true);
 }
+
+////////////////////////////////////////////////////////////////////////////////
+/// Replaces all observables and the weight variable of this NLL with clones
+/// that only differ by a prefix added to the names. Used for simultaneous fits.
+/// \return A RooArgSet with the new observable args.
+/// \param[in] prefix The prefix to add to the observables and weight names.
+RooArgSet RooNLLVarNew::prefixObservableAndWeightNames(std::string const &prefix)
+{
+   RooArgSet obsSet{_observables};
+   if (_weight)
+      obsSet.add(**_weight);
+   RooArgSet obsClones;
+   obsSet.snapshot(obsClones);
+   for (RooAbsArg *arg : obsClones) {
+      arg->setAttribute((std::string("ORIGNAME:") + arg->GetName()).c_str());
+      arg->SetName((prefix + arg->GetName()).c_str());
+   }
+   recursiveRedirectServers(obsClones, false, true);
+
+   RooArgSet newObservables{obsClones};
+   if (_weight) {
+      newObservables.remove(**_weight);
+   }
+
+   setObservables(obsClones);
+   addOwnedComponents(std::move(obsClones));
+
+   return newObservables;
+}
+
+namespace {
+
+inline bool endsWith(std::string const &value, std::string_view ending)
+{
+   if (ending.size() > value.size())
+      return false;
+   return std::equal(ending.rbegin(), ending.rend(), value.rbegin());
+}
+
+} // namespace
+
+////////////////////////////////////////////////////////////////////////////////
+/// Toggles the weight square correction by changing the name of the weight
+/// variable to get different weights from the RooFitDriver data map. The
+/// RooNLLVarNew::weightVarNameSumW2Suffix is either added or removed from the
+/// weight variable name.
+void RooNLLVarNew::applyWeightSquared(bool flag)
+{
+   if (!_weight)
+      return;
+   auto &w = **_weight;
+
+   const std::string name = w.GetName();
+   const std::string suffix = weightVarNameSumW2Suffix;
+
+   bool isAlreadyWeightSquared = endsWith(name, suffix);
+
+   if (isAlreadyWeightSquared && !flag) {
+      w.SetName(name.substr(0, name.length() - suffix.length()).c_str());
+   } else if (!isAlreadyWeightSquared && flag) {
+      w.SetName((name + suffix).c_str());
+   }
+}