[RF] Implement channel masking for simultaneous likelihoods

Upstreaming a feature from CMS combine.

Draft for now.

roofit/codegen/inc/RooFit/CodegenImpl.h

@@ -67,6 +67,7 @@ namespace Detail {
 class RooFixedProdPdf;
 class RooNLLVarNew;
 class RooNormalizedPdf;
+class RooSimNLL;
 } // namespace Detail
 
 namespace Experimental {
@@ -76,6 +77,7 @@ class CodegenContext;
 void codegenImpl(RooFit::Detail::RooFixedProdPdf &arg, CodegenContext &ctx);
 void codegenImpl(RooFit::Detail::RooNLLVarNew &arg, CodegenContext &ctx);
 void codegenImpl(RooFit::Detail::RooNormalizedPdf &arg, CodegenContext &ctx);
+void codegenImpl(RooFit::Detail::RooSimNLL &arg, CodegenContext &ctx);
 void codegenImpl(ParamHistFunc &arg, CodegenContext &ctx);
 void codegenImpl(PiecewiseInterpolation &arg, CodegenContext &ctx);
 void codegenImpl(RooAbsArg &arg, CodegenContext &ctx);

roofit/codegen/src/CodegenImpl.cxx

@@ -133,6 +133,28 @@ void codegenImpl(RooFit::Detail::RooFixedProdPdf &arg, CodegenContext &ctx)
    }
 }
 
+void codegenImpl(RooFit::Detail::RooSimNLL &arg, CodegenContext &ctx)
+{
+   if (arg.terms().empty()) {
+      ctx.addResult(&arg, "0.0");
+   }
+
+   std::string resName = RooFit::Detail::makeValidVarName(arg.GetName()) + "Result";
+   ctx.addResult(&arg, resName);
+   ctx.addToGlobalScope("double " + resName + " = 0.0;\n");
+
+   std::stringstream ss;
+
+   std::size_t i = 0;
+   for (auto *component : static_range_cast<RooAbsReal *>(arg.terms())) {
+
+      // TODO: support channel masking here
+      ss << resName << " += " << ctx.buildFunction(*component, ctx.outputSizes()) << "(params, obs, xlArr);\n";
+      ++i;
+   }
+   ctx.addToGlobalScope(ss.str());
+}
+
 void codegenImpl(ParamHistFunc &arg, CodegenContext &ctx)
 {
    std::string const &idx = arg.dataHist().calculateTreeIndexForCodeSquash(&arg, ctx, arg.dataVars(), true);
@@ -251,15 +273,7 @@ void codegenImpl(RooAddition &arg, CodegenContext &ctx)
 
    std::size_t i = 0;
    for (auto *component : static_range_cast<RooAbsReal *>(arg.list())) {
-
-      if (!dynamic_cast<RooFit::Detail::RooNLLVarNew *>(component) || arg.list().size() == 1) {
-         result += ctx.getResult(*component);
-         ++i;
-         if (i < arg.list().size())
-            result += '+';
-         continue;
-      }
-      result += ctx.buildFunction(*component, ctx.outputSizes()) + "(params, obs, xlArr)";
+      result += ctx.getResult(*component);
       ++i;
       if (i < arg.list().size())
          result += '+';

roofit/roofitcore/inc/LinkDef.h

@@ -337,5 +337,6 @@
 #pragma link C++ class RooBinWidthFunction+;
 #pragma link C++ class RooFit::Detail::RooNLLVarNew+;
 #pragma link C++ class RooFit::Detail::RooNormalizedPdf+ ;
+#pragma link C++ class RooFit::Detail::RooSimNLL+;
 
 #endif

roofit/roofitcore/inc/RooFit/Detail/RooNLLVarNew.h

@@ -19,10 +19,13 @@
 #include <RooAbsPdf.h>
 #include <RooAbsReal.h>
 #include <RooGlobalFunc.h>
+#include <RooListProxy.h>
 #include <RooTemplateProxy.h>
 
 #include <Math/Util.h>
 
+class RooAbsCategory;
+
 namespace RooFit {
 namespace Detail {
 
@@ -87,6 +90,30 @@ class RooNLLVarNew : public RooAbsReal {
    ClassDefOverride(RooFit::Detail::RooNLLVarNew, 0);
 };
 
+class RooSimNLL : public RooAbsReal {
+public:
+   RooSimNLL(const char *name, const char *title, const RooArgSet &terms, RooAbsCategoryLValue const &indexCat,
+             bool channelMasking);
+
+   RooSimNLL(const RooSimNLL &other, const char *name = nullptr);
+   TObject *clone(const char *newname) const override { return new RooSimNLL(*this, newname); }
+
+   double defaultErrorLevel() const override;
+
+   const RooArgSet &terms() const { return _set; }
+   const RooArgSet &masks() const { return _mask; }
+
+   void doEval(RooFit::EvalContext &) const override;
+
+protected:
+   double evaluate() const override;
+
+   RooSetProxy _set; ///< set of terms to be summed
+   RooSetProxy _mask;
+
+   ClassDefOverride(RooFit::Detail::RooSimNLL, 0) // Sum of RooNLLVarNew instances
+};
+
 } // namespace Detail
 } // namespace RooFit
 

roofit/roofitcore/src/FitHelpers.cxx

@@ -51,6 +51,7 @@
 #endif
 
 using RooFit::Detail::RooNLLVarNew;
+using RooFit::Detail::RooSimNLL;
 
 namespace {
 
@@ -357,7 +358,7 @@ std::unique_ptr<RooAbsArg> createSimultaneousNLL(RooSimultaneous const &simPdf,
    }
 
    // Time to sum the NLLs
-   auto nll = std::make_unique<RooAddition>("mynll", "mynll", nllTerms);
+   auto nll = std::make_unique<RooSimNLL>("mynll", "mynll", nllTerms, simCat, true);
    nll->addOwnedComponents(std::move(nllTerms));
    return nll;
 }

roofit/roofitcore/src/RooNLLVarNew.cxx

@@ -25,14 +25,15 @@ computation times.
 
 #include "RooFit/Detail/RooNLLVarNew.h"
 
-#include <RooHistPdf.h>
+#include <RooAbsCategoryLValue.h>
 #include <RooBatchCompute.h>
+#include <RooConstVar.h>
 #include <RooDataHist.h>
+#include <RooFit/Detail/MathFuncs.h>
+#include <RooHistPdf.h>
 #include <RooNaNPacker.h>
-#include <RooConstVar.h>
 #include <RooRealVar.h>
 #include <RooSetProxy.h>
-#include <RooFit/Detail/MathFuncs.h>
 
 #include "RooFitImplHelpers.h"
 
@@ -47,6 +48,7 @@ computation times.
 #include <vector>
 
 ClassImp(RooFit::Detail::RooNLLVarNew);
+ClassImp(RooFit::Detail::RooSimNLL);
 
 namespace RooFit {
 namespace Detail {
@@ -336,6 +338,57 @@ void RooNLLVarNew::finalizeResult(RooFit::EvalContext &ctx, ROOT::Math::KahanSum
    ctx.setOutputWithOffset(this, result, _offset);
 }
 
+RooSimNLL::RooSimNLL(const char *name, const char *title, const RooArgSet &terms, RooAbsCategoryLValue const &indexCat,
+                     bool channelMasking)
+   : RooAbsReal(name, title), _set("!set", "set of components", this), _mask("!mask", "set of masks", this)
+{
+   _set.addTyped<RooAbsReal>(terms);
+
+   if (channelMasking) {
+      for (auto const &catState : indexCat) {
+         std::string const &catName = catState.first;
+         std::string maskName = "mask_" + catName;
+         _mask.addOwned(std::make_unique<RooRealVar>(maskName.c_str(), maskName.c_str(), 0.0));
+      }
+   }
+}
+
+RooSimNLL::RooSimNLL(const RooSimNLL &other, const char *name)
+   : RooAbsReal(other, name), _set("!set", this, other._set), _mask("!mask", this, other._set)
+{
+}
+
+double RooSimNLL::evaluate() const
+{
+   double sum(0);
+   const RooArgSet *nset = _set.nset();
+
+   std::size_t i = 0;
+   for (auto *comp : static_range_cast<RooAbsReal *>(_set)) {
+      if (_mask.empty() || static_cast<RooAbsReal const *>(_mask[i])->getVal() == 0.0) {
+         sum += comp->getVal(nset);
+      }
+      ++i;
+   }
+   return sum;
+}
+
+void RooSimNLL::doEval(RooFit::EvalContext &ctx) const
+{
+   double result = 0.;
+   for (std::size_t i = 0; i < _set.size(); ++i) {
+      if (_mask.empty() || ctx.at(_mask[i])[0] == 0.0) {
+         result += ctx.at(_set[i])[0];
+      }
+   }
+   ctx.output()[0] = result;
+}
+
+double RooSimNLL::defaultErrorLevel() const
+{
+   return static_cast<RooNLLVarNew *>(_set[0])->defaultErrorLevel();
+}
+
 } // namespace Detail
 } // namespace RooFit