[RF] Let RooGenericPdf/RooFormulaVar declare a piecewise-flat binning

roofit/roofitcore/inc/RooFormulaVar.h

@@ -20,12 +20,16 @@
 #include "RooArgList.h"
 #include "RooListProxy.h"
 #include "RooTrace.h"
+#include "RooAbsBinning.h"
 
 #include <memory>
 #include <list>
+#include <map>
+#include <string>
 
 class RooArgSet ;
 class RooFormula ;
+class RooAbsRealLValue;
 
 class RooFormulaVar : public RooAbsReal {
 public:
@@ -68,6 +72,14 @@ class RooFormulaVar : public RooAbsReal {
 
   double defaultErrorLevel() const override ;
 
+  // Declare this function to be piecewise constant (flat) within the bins of
+  // the given `binning` of observable `obs`. This lets integration use the fast
+  // bin integrator instead of the generic numeric integrator. A binning can be
+  // set for more than one observable. Use a RooUniformBinning to describe many
+  // uniform bins compactly.
+  void setBinBoundaries(RooAbsRealLValue &obs, const RooAbsBinning &binning, bool checkFlatness = true);
+
+  bool isBinnedDistribution(const RooArgSet &obs) const override;
   std::list<double>* binBoundaries(RooAbsRealLValue& /*obs*/, double /*xlo*/, double /*xhi*/) const override ;
   std::list<double>* plotSamplingHint(RooAbsRealLValue& /*obs*/, double /*xlo*/, double /*xhi*/) const override ;
 
@@ -94,7 +106,10 @@ class RooFormulaVar : public RooAbsReal {
   mutable RooArgSet* _nset{nullptr}; ///<! Normalization set to be passed along to contents
   TString _formExpr ;            ///< Formula expression string
 
-  ClassDefOverride(RooFormulaVar,1) // Real-valued function of other RooAbsArgs calculated by a TFormula expression
+  std::map<int, std::unique_ptr<RooAbsBinning>> _binnings; ///< User-defined binnings, keyed by the observable's index
+                                                           ///< in _actualVars, for a piecewise-flat distribution
+
+  ClassDefOverride(RooFormulaVar, 2) // Real-valued function of other RooAbsArgs calculated by a TFormula expression
 };
 
 #endif

roofit/roofitcore/inc/RooGenericPdf.h

@@ -18,9 +18,15 @@
 
 #include "RooAbsPdf.h"
 #include "RooListProxy.h"
+#include "RooAbsBinning.h"
+
+#include <map>
+#include <memory>
+#include <string>
 
 class RooArgList ;
 class RooFormula ;
+class RooAbsRealLValue;
 
 class RooGenericPdf : public RooAbsPdf {
 public:
@@ -61,8 +67,18 @@ class RooGenericPdf : public RooAbsPdf {
 
   std::string getUniqueFuncName() const;
 
-protected:
+  // Declare this pdf to be piecewise constant (flat) within the bins of the
+  // given `binning` of observable `obs`. This lets the integration use the fast
+  // bin integrator instead of the generic numeric integrator. A binning can be
+  // set for more than one observable. Use a RooUniformBinning to describe many
+  // uniform bins compactly.
+  void setBinBoundaries(RooAbsRealLValue &obs, const RooAbsBinning &binning, bool checkFlatness = true);
+
+  bool isBinnedDistribution(const RooArgSet &obs) const override;
+  std::list<double> *binBoundaries(RooAbsRealLValue &obs, double xlo, double xhi) const override;
+  std::list<double> *plotSamplingHint(RooAbsRealLValue &obs, double xlo, double xhi) const override;
 
+  protected:
   RooFormula& formula() const ;
 
   // Function evaluation
@@ -78,7 +94,10 @@ class RooGenericPdf : public RooAbsPdf {
   mutable RooFormula * _formula = nullptr; ///<! Formula engine
   TString _formExpr ;            ///< Formula expression string
 
-  ClassDefOverride(RooGenericPdf,1) // Generic PDF defined by string expression and list of variables
+  std::map<int, std::unique_ptr<RooAbsBinning>> _binnings; ///< User-defined binnings, keyed by the observable's index
+                                                           ///< in _actualVars, for a piecewise-flat distribution
+
+  ClassDefOverride(RooGenericPdf, 2) // Generic PDF defined by string expression and list of variables
 };
 
 #endif

roofit/roofitcore/inc/RooHelpers.h

@@ -17,13 +17,17 @@
 #include <RooAbsArg.h>
 #include <RooAbsReal.h>
 
+#include <ROOT/RSpan.hxx>
+
 #include <sstream>
+#include <list>
 #include <vector>
 #include <string>
 #include <utility>
 
 class RooAbsPdf;
 class RooAbsData;
+class RooAbsRealLValue;
 
 namespace RooHelpers {
 
@@ -91,6 +95,22 @@ void checkRangeOfParameters(const RooAbsReal *callingClass, std::initializer_lis
 /// set all RooRealVars to constants. return true if at least one changed status
 bool setAllConstant(const RooAbsCollection &coll, bool constant = true);
 
+/// Check that `function` is constant (flat) inside each bin defined by the
+/// sorted `boundaries` when scanning the observable `obs`. Several interior
+/// points are sampled per bin and compared to the bin's first sample; if any
+/// of them deviates by more than `relTol` (relative to the value scale), the
+/// function is not flat and false is returned. The value of `obs` is restored
+/// on return.
+bool isFunctionFlatInBins(const RooAbsReal &function, RooAbsRealLValue &obs, std::span<const double> boundaries,
+                          double relTol = 1e-9);
+
+/// Return a newly allocated list with the subset of `boundaries` that lies
+/// strictly inside [`xlo`, `xhi`], with `xlo` and `xhi` added as the first and
+/// last entries. This is the form expected by RooFit's binBoundaries()
+/// interface, so the bin integrator covers exactly the integration range.
+/// The caller takes ownership of the returned list.
+std::list<double> *binBoundariesInRange(std::span<const double> boundaries, double xlo, double xhi);
+
 } // namespace RooHelpers
 
 #endif

roofit/roofitcore/src/RooFormulaVar.cxx

@@ -52,6 +52,10 @@
 #include "RooMsgService.h"
 #include "RooTrace.h"
 #include "RooFormula.h"
+#include "RooAbsRealLValue.h"
+#include "RooAbsBinning.h"
+#include "RooCurve.h"
+#include "RooHelpers.h"
 
 #ifdef ROOFIT_LEGACY_EVAL_BACKEND
 #include "RooNLLVar.h"
@@ -123,6 +127,9 @@ RooFormulaVar::RooFormulaVar(const RooFormulaVar& other, const char* name) :
   _actualVars("actualVars",this,other._actualVars),
   _formExpr(other._formExpr)
 {
+   for (auto const &item : other._binnings) {
+      _binnings[item.first] = std::unique_ptr<RooAbsBinning>{item.second->clone()};
+   }
   if (other._formula && other._formula->ok()) {
     _formula = new RooFormula(*other._formula);
     _formExpr = _formula->formulaString().c_str();
@@ -228,13 +235,74 @@ void RooFormulaVar::writeToStream(ostream& os, bool compact) const
   }
 }
 
+////////////////////////////////////////////////////////////////////////////////
+/// Declare that this function is piecewise constant (flat) within the bins of
+/// the given `binning` of the observable `obs`, which must be one of the formula
+/// variables. The method can be called several times to set a binning for more
+/// than one observable. See RooGenericPdf::setBinBoundaries() for details.
+
+void RooFormulaVar::setBinBoundaries(RooAbsRealLValue &obs, const RooAbsBinning &binning, bool checkFlatness)
+{
+   // Match the observable to a formula variable by name, so that a same-named
+   // stand-in for the actual server is accepted too.
+   const int idx = _actualVars.index(obs.GetName());
+   if (idx < 0) {
+      coutE(InputArguments) << "RooFormulaVar::setBinBoundaries(" << GetName() << ") the observable " << obs.GetName()
+                            << " is not one of the formula variables of this function, nothing done." << std::endl;
+      return;
+   }
+
+   if (checkFlatness) {
+      // Vary the actual formula variable (the server), which may be a different
+      // object than `obs` if `obs` is just a same-named stand-in.
+      auto *serverObs = dynamic_cast<RooAbsRealLValue *>(_actualVars.at(idx));
+      RooAbsRealLValue &flatObs = serverObs ? *serverObs : obs;
+      std::span<const double> boundaries{binning.array(), static_cast<std::size_t>(binning.numBoundaries())};
+      if (!RooHelpers::isFunctionFlatInBins(*this, flatObs, boundaries)) {
+         coutE(InputArguments) << "RooFormulaVar::setBinBoundaries(" << GetName() << ") the expression \"" << _formExpr
+                               << "\" is not flat within the given bins of " << obs.GetName()
+                               << ". The binning is not set. Pass checkFlatness=false to override this check."
+                               << std::endl;
+         return;
+      }
+   }
+
+   // Key the binning by the observable's index in _actualVars (not its name), so
+   // that it survives a renaming of the variable or a server redirection.
+   _binnings[idx] = std::unique_ptr<RooAbsBinning>{binning.clone()};
+}
+
+////////////////////////////////////////////////////////////////////////////////
+/// Return true if a binning was set with setBinBoundaries() for every
+/// observable in the integration set `obs`.
 
+bool RooFormulaVar::isBinnedDistribution(const RooArgSet &obs) const
+{
+   if (obs.empty() || _binnings.empty()) {
+      return false;
+   }
+   for (RooAbsArg *o : obs) {
+      if (_binnings.find(_actualVars.index(o->GetName())) == _binnings.end()) {
+         return false;
+      }
+   }
+   return true;
+}
 
 ////////////////////////////////////////////////////////////////////////////////
-/// Forward the plot sampling hint from the p.d.f. that defines the observable obs
+/// Return the boundaries of the binning set with setBinBoundaries() that fall
+/// within [xlo, xhi]. If no binning was set for this observable, forward the bin
+/// boundaries from the server that defines the observable obs.
 
 std::list<double>* RooFormulaVar::binBoundaries(RooAbsRealLValue& obs, double xlo, double xhi) const
 {
+   auto found = _binnings.find(_actualVars.index(obs.GetName()));
+   if (found != _binnings.end()) {
+      const RooAbsBinning &binning = *found->second;
+      return RooHelpers::binBoundariesInRange({binning.array(), static_cast<std::size_t>(binning.numBoundaries())}, xlo,
+                                              xhi);
+   }
+
   for (const auto par : _actualVars) {
     auto func = static_cast<const RooAbsReal*>(par);
     list<double>* binb = nullptr;
@@ -247,13 +315,20 @@ std::list<double>* RooFormulaVar::binBoundaries(RooAbsRealLValue& obs, double xl
   return nullptr;
 }
 
-
-
 ////////////////////////////////////////////////////////////////////////////////
-/// Forward the plot sampling hint from the p.d.f. that defines the observable obs
+/// Return sampling hints that draw the piecewise-flat shape exactly if a binning
+/// was set for this observable. Otherwise, forward the plot sampling hint from
+/// the server that defines the observable obs.
 
 std::list<double>* RooFormulaVar::plotSamplingHint(RooAbsRealLValue& obs, double xlo, double xhi) const
 {
+   auto found = _binnings.find(_actualVars.index(obs.GetName()));
+   if (found != _binnings.end()) {
+      const RooAbsBinning &binning = *found->second;
+      return RooCurve::plotSamplingHintForBinBoundaries(
+         {binning.array(), static_cast<std::size_t>(binning.numBoundaries())}, xlo, xhi);
+   }
+
   for (const auto par : _actualVars) {
     auto func = dynamic_cast<const RooAbsReal*>(par);
     list<double>* hint = nullptr;

roofit/roofitcore/src/RooGenericPdf.cxx

@@ -48,6 +48,10 @@ the names of the arguments are not hard coded.
 #include "RooMsgService.h"
 #include "RooArgList.h"
 #include "RooFormula.h"
+#include "RooAbsRealLValue.h"
+#include "RooAbsBinning.h"
+#include "RooCurve.h"
+#include "RooHelpers.h"
 
 using std::istream, std::ostream, std::endl;
 
@@ -107,6 +111,9 @@ RooGenericPdf::RooGenericPdf(const RooGenericPdf& other, const char* name) :
   _actualVars("actualVars",this,other._actualVars),
   _formExpr(other._formExpr)
 {
+   for (auto const &item : other._binnings) {
+      _binnings[item.first] = std::unique_ptr<RooAbsBinning>{item.second->clone()};
+   }
   formula();
 }
 
@@ -122,7 +129,99 @@ RooFormula& RooGenericPdf::formula() const
   return *_formula ;
 }
 
+////////////////////////////////////////////////////////////////////////////////
+/// Declare that this pdf is piecewise constant (flat) within the bins of the
+/// given `binning` of the observable `obs`, which must be one of the formula
+/// variables. The method can be called several times to set a binning for more
+/// than one observable. Use a RooUniformBinning to describe many uniform bins
+/// compactly.
+///
+/// Once set, integrals over `obs` use the fast bin integrator (which sums the
+/// central value of each bin times the bin width) instead of the generic
+/// numeric integrator, and plotting samples the step shape exactly.
+///
+/// If `checkFlatness` is true (the default), the function is sampled at several
+/// points inside each bin to verify that it is indeed flat; if it is not, an
+/// error is issued and the binning is not stored.
+
+void RooGenericPdf::setBinBoundaries(RooAbsRealLValue &obs, const RooAbsBinning &binning, bool checkFlatness)
+{
+   // Match the observable to a formula variable by name, so that a same-named
+   // stand-in for the actual server is accepted too.
+   const int idx = _actualVars.index(obs.GetName());
+   if (idx < 0) {
+      coutE(InputArguments) << "RooGenericPdf::setBinBoundaries(" << GetName() << ") the observable " << obs.GetName()
+                            << " is not one of the formula variables of this pdf, nothing done." << std::endl;
+      return;
+   }
+
+   if (checkFlatness) {
+      // Vary the actual formula variable (the server), which may be a different
+      // object than `obs` if `obs` is just a same-named stand-in.
+      auto *serverObs = dynamic_cast<RooAbsRealLValue *>(_actualVars.at(idx));
+      RooAbsRealLValue &flatObs = serverObs ? *serverObs : obs;
+      std::span<const double> boundaries{binning.array(), static_cast<std::size_t>(binning.numBoundaries())};
+      if (!RooHelpers::isFunctionFlatInBins(*this, flatObs, boundaries)) {
+         coutE(InputArguments) << "RooGenericPdf::setBinBoundaries(" << GetName() << ") the expression \"" << _formExpr
+                               << "\" is not flat within the given bins of " << obs.GetName()
+                               << ". The binning is not set. Pass checkFlatness=false to override this check."
+                               << std::endl;
+         return;
+      }
+   }
+
+   // Key the binning by the observable's index in _actualVars (not its name), so
+   // that it survives a renaming of the variable or a server redirection.
+   _binnings[idx] = std::unique_ptr<RooAbsBinning>{binning.clone()};
+}
+
+////////////////////////////////////////////////////////////////////////////////
+/// Return true if a binning was set with setBinBoundaries() for every
+/// observable in the integration set `obs`.
 
+bool RooGenericPdf::isBinnedDistribution(const RooArgSet &obs) const
+{
+   if (obs.empty() || _binnings.empty()) {
+      return false;
+   }
+   for (RooAbsArg *o : obs) {
+      if (_binnings.find(_actualVars.index(o->GetName())) == _binnings.end()) {
+         return false;
+      }
+   }
+   return true;
+}
+
+////////////////////////////////////////////////////////////////////////////////
+/// Return the boundaries of the binning set with setBinBoundaries() that fall
+/// within [xlo, xhi], or a null pointer if no binning was set for this observable.
+
+std::list<double> *RooGenericPdf::binBoundaries(RooAbsRealLValue &obs, double xlo, double xhi) const
+{
+   auto found = _binnings.find(_actualVars.index(obs.GetName()));
+   if (found == _binnings.end()) {
+      return nullptr;
+   }
+   const RooAbsBinning &binning = *found->second;
+   return RooHelpers::binBoundariesInRange({binning.array(), static_cast<std::size_t>(binning.numBoundaries())}, xlo,
+                                           xhi);
+}
+
+////////////////////////////////////////////////////////////////////////////////
+/// Return sampling hints that draw the piecewise-flat shape exactly (a pair of
+/// points just left and right of every bin boundary), or a null pointer if no
+/// binning was set for this observable.
+
+std::list<double> *RooGenericPdf::plotSamplingHint(RooAbsRealLValue &obs, double xlo, double xhi) const
+{
+   auto found = _binnings.find(_actualVars.index(obs.GetName()));
+   if (found == _binnings.end()) {
+      return nullptr;
+   }
+   const RooAbsBinning &binning = *found->second;
+   return RooCurve::plotSamplingHintForBinBoundaries(
+      {binning.array(), static_cast<std::size_t>(binning.numBoundaries())}, xlo, xhi);
+}
 
 ////////////////////////////////////////////////////////////////////////////////
 /// Calculate current value of this object