Add codegen to RooParametricHist

interface/CombineCodegenImpl.h

@@ -8,6 +8,7 @@ class FastVerticalInterpHistPdf2;
 class FastVerticalInterpHistPdf2D2;
 class ProcessNormalization;
 class VerticalInterpPdf;
+class RooParametricHist;
 
 namespace RooFit::Experimental {
 
@@ -18,8 +19,10 @@ namespace RooFit::Experimental {
   void codegenImpl(FastVerticalInterpHistPdf2D2& arg, CodegenContext& ctx);
   void codegenImpl(ProcessNormalization& arg, CodegenContext& ctx);
   void codegenImpl(VerticalInterpPdf& arg, CodegenContext& ctx);
+  void codegenImpl(RooParametricHist& arg, CodegenContext& ctx);
 
   std::string codegenIntegralImpl(VerticalInterpPdf& arg, int code, const char* rangeName, CodegenContext& ctx);
+  std::string codegenIntegralImpl(RooParametricHist& arg, int code, const char* rangeName, CodegenContext& ctx);
 
 }  // namespace RooFit::Experimental
 

interface/CombineMathFuncs.h

@@ -306,6 +306,167 @@
    return result > 0. ? result : integralFloorVal;
 }
 
+inline int parametricHistFindBin(const int N_bins, const double* bins, const double x) {
+  if (x < bins[0] || x >= bins[N_bins])
+    return -1;
+
+  // Search for the bin
+  for (int i = 0; i < N_bins; ++i) {
+    if (x >= bins[i] && x < bins[i + 1])
+      return i;
+  }
+  return -1;
+}
+
+inline int parametricHistFindBin(const int N_bins, std::vector<double> const& bins, const double x) {
+  return parametricHistFindBin(N_bins, bins.data(), x);
+}
+
+inline Double_t parametricHistMorphScale(const double parVal,
+                                         const int nMorphs,
+                                         const double* morphCoeffs,
+                                         const double* morphDiffs,
+                                         const double* morphSums,
+                                         double smoothRegion) {
+  double morphScale = 1.0;
+  if (!morphDiffs || !morphSums)
+    return morphScale;
+  for (int i = 0; i < nMorphs; ++i) {
+    double coeff = morphCoeffs[i];
+    double a = 0.5 * coeff;
+    double b = smoothStepFunc(coeff, smoothRegion);
+    morphScale *= 1 + (1.0 / parVal) * a * (morphDiffs[i] + b * morphSums[i]);
+  }
+  return morphScale;
+}
+
+inline Double_t parametricHistEvaluate(const int bin_i,
+                                       const double* parVals,
+                                       const double* bins,
+                                       const int N_bins,
+                                       const double* morphCoeffs,
+                                       const int nMorphs,
+                                       const double* morphDiffs,
+                                       const double* morphSums,
+                                       const double* widths,
+                                       const double smoothRegion) {
+  if (bin_i < 0)
+    return 0.0;
+  // Morphing case
+  if (morphCoeffs != nullptr && nMorphs > 0) {
+    // morphDiffs and morphSums are flattened arrays of size N_bins * nMorphs
+    const double* binMorphDiffs = nullptr;
+    const double* binMorphSums = nullptr;
+    if (morphDiffs) {
+      binMorphDiffs = morphDiffs + bin_i * nMorphs;
+    }
+    if (morphSums) {
+      binMorphSums = morphSums + bin_i * nMorphs;
+    }
+    double parVal = parVals[bin_i];
+    double scale = parametricHistMorphScale(parVal, nMorphs, morphCoeffs, binMorphDiffs, binMorphSums, smoothRegion);
+    return (parVal * scale) / widths[bin_i];
+  }
+  // No morphing case
+  return parVals[bin_i] / widths[bin_i];
+}
+
+inline Double_t parametricMorphFunction(const int j,
+                                        const double parVal,
+                                        const bool hasMorphs,
+                                        const int nMorphs,
+                                        const double* morphCoeffs,
+                                        const double* morphDiffs,
+                                        const double* morphSums,
+                                        double smoothRegion) {
+  double morphScale = 1.0;
+  if (!hasMorphs)
+    return morphScale;
+
+  int ndim = nMorphs;
+  // apply all morphs one by one to the bin
+  // almost certaintly a faster way to do this in a vectorized way ....
+  for (int i = 0; i < ndim; ++i) {
+    double x = morphCoeffs[i];
+    double a = 0.5 * x, b = smoothStepFunc(x, smoothRegion);
+    morphScale *= 1 + (1. / parVal) * a * (morphDiffs[j * nMorphs + i] + b * morphSums[j * nMorphs + i]);
+  }
+  return morphScale;
+}
+
+inline Double_t parametricHistFullSum(const double* parVals,
+                                      const int nBins,
+                                      const bool hasMorphs,
+                                      const int nMorphs,
+                                      const double* morphCoeffs,
+                                      const double* morphDiffs,
+                                      const double* morphSums,
+                                      double smoothRegion) {
+  double sum = 0;
+  for (int i = 0; i < nBins; ++i) {
+    double thisVal = parVals[i];
+    if (hasMorphs) {
+      // Apply morphing to this bin, just like in RooParametricHist::evaluate
+      thisVal *=
+          parametricMorphFunction(i, thisVal, hasMorphs, nMorphs, morphCoeffs, morphDiffs, morphSums, smoothRegion);
+    }
+    sum += thisVal;
+  }
+  return sum;
+}
+
+inline Double_t parametricHistIntegral(const double* parVals,
+                                       const double* bins,
+                                       const int N_bins,
+                                       const double* morphCoeffs,
+                                       const int nMorphs,
+                                       const double* morphDiffs,
+                                       const double* morphSums,
+                                       const double* widths,
+                                       const double smoothRegion,
+                                       const char* rangeName,
+                                       const double xmin,
+                                       const double xmax) {
+  // No ranges
+  if (!rangeName) {
+    return parametricHistFullSum(
+        parVals, N_bins, morphCoeffs != nullptr, nMorphs, morphCoeffs, morphDiffs, morphSums, smoothRegion);
+  }
+
+  // Case with ranges, calculate integral explicitly
+  double sum = 0;
+  int i;
+  for (i = 1; i <= N_bins; i++) {
+    // Get maybe-morphed bin value
+    double binVal = parVals[i - 1] / widths[i - 1];
+    if (morphCoeffs != nullptr) {
+      binVal *= parametricMorphFunction(
+          i - 1, parVals[i - 1], true, nMorphs, morphCoeffs, morphDiffs, morphSums, smoothRegion);
+    }
+
+    if (bins[i - 1] >= xmin && bins[i] <= xmax) {
+      // Bin fully in integration domain
+      sum += (bins[i] - bins[i - 1]) * binVal;
+    } else if (bins[i - 1] < xmin && bins[i] > xmax) {
+      // Domain is fully contained in this bin
+      sum += (xmax - xmin) * binVal;
+      // Exit here, this is the last bin to be processed by construction
+      double fullSum = parametricHistFullSum(
+          parVals, N_bins, morphCoeffs != nullptr, nMorphs, morphCoeffs, morphDiffs, morphSums, smoothRegion);
+      return sum / fullSum;
+    } else if (bins[i - 1] < xmin && bins[i] <= xmax && bins[i] > xmin) {
+      // Lower domain boundary is in bin
+      sum += (bins[i] - xmin) * binVal;
+    } else if (bins[i - 1] >= xmin && bins[i] > xmax && bins[i - 1] < xmax) {
+      // Upper domain boundary is in bin
+      // Exit here, this is the last bin to be processed by construction
+      sum += (xmax - bins[i - 1]) * binVal;
+      return sum;
+    }
+  }
+  return sum;
+}
+
 } // namespace MathFuncs
 } // namespace Detail
 } // namespace RooFit

interface/RooParametricHist.h

@@ -31,15 +31,33 @@
   RooArgList & getAllBinVars() const ;
 
   RooRealVar & getObs() const { return (RooRealVar&)x; };
+  RooAbsReal& observable() const { return const_cast<RooAbsReal&>(static_cast<const RooAbsReal&>(x.arg())); }
   const std::vector<double>  getBins()   const { return bins;   };
+  const int getNBins() const { return N_bins; };
   const std::vector<double>  getWidths() const { return widths; };
 
-  const double quickSum() const {return getFullSum() ;}
+  const RooArgList& getPars() const { return pars; };
+  const RooArgList& getCoeffList() const { return _coeffList; };
+
+  const double quickSum() const;
   //RooAddition & getYieldVar(){return sum;};
 
   // how can we pass this version? is there a Collection object for RooDataHists?
   //void addMorphs(RooArgList &_morphPdfsUp, RooArgList &_morphPdfsDown, RooArgList &_coeffs, double smoothRegion);
   void addMorphs(RooDataHist&, RooDataHist&, RooRealVar&, double );
+  Double_t evaluate() const override;
+
+  // Accessors for evaluation data
+  double getX() const { return x; }
+  double getSmoothRegion() const { return _smoothRegion; }
+  bool hasMorphs() const { return _hasMorphs; }
+
+  double getParVal(int bin_i) const;
+
+  // Utility functions for data extraction
+  const std::vector<double>& getParVals() const;
+  const std::vector<double>& getCoeffs() const;
+  void getFlattenedMorphs(std::vector<double>& diffs_flat, std::vector<double>& sums_flat) const;
 
 protected:
 
@@ -56,19 +74,16 @@
   mutable double _smoothRegion;
   mutable bool   _hasMorphs;
   mutable std::vector<std::vector <double> > _diffs;
-  mutable std::vector<std::vector <double> > _sums;
-  double evaluateMorphFunction(int) const;
+  mutable std::vector<std::vector<double> > _sums;
 
+  mutable std::vector<double> pars_vals_;   //! Don't serialize me
+  mutable std::vector<double> coeffs_;      //! Don't serialize me
+  mutable std::vector<double> diffs_flat_;  //! Don't serialize me
+  mutable std::vector<double> sums_flat_;   //! Don't serialize me
 
   void initializeBins(const TH1&) const;
   //void initializeNorm();
 
-
-  double evaluatePartial() const ;
-  double evaluateFull() const ;
-  Double_t evaluate() const override ;
-  double getFullSum() const ;
-
   mutable double _cval;
   void update_cval(double r){_cval=r;};
 

src/CombineCodegenImpl.cxx

@@ -9,6 +9,7 @@
 #include "../interface/VerticalInterpHistPdf.h"
 #include "../interface/VerticalInterpPdf.h"
 #include "../interface/CombineMathFuncs.h"
+#include "../interface/RooParametricHist.h"
 
 #include <RooUniformBinning.h>
 
@@ -242,6 +243,33 @@
                               arg.quadraticAlgo()));
 }
 
+void RooFit::Experimental::codegenImpl(RooParametricHist& arg, CodegenContext& ctx) {
+  std::vector<double> diffs_flat;
+  std::vector<double> sums_flat;
+  if (arg.hasMorphs()) {
+    arg.getFlattenedMorphs(diffs_flat, sums_flat);
+  }
+
+  std::string xName = ctx.getResult(arg.observable());
+  std::stringstream bin_i;
+  bin_i << ctx.buildCall("RooFit::Detail::MathFuncs::parametricHistFindBin", arg.getNBins(), arg.getBins(), xName);
+  std::stringstream code;
+  code << ctx.buildCall("RooFit::Detail::MathFuncs::parametricHistEvaluate",
+                        bin_i.str(),
+                        arg.getPars(),
+                        arg.getBins(),
+                        arg.getNBins(),
+                        arg.getCoeffList(),
+                        static_cast<int>(arg.getCoeffList().size()),
+                        arg.hasMorphs() ? diffs_flat : std::vector<double>{},
+                        arg.hasMorphs() ? sums_flat : std::vector<double>{},
+                        arg.getWidths(),
+                        arg.getSmoothRegion()) +
+              ";\n";
+  ctx.addToCodeBody(code.str());
+  ctx.addResult(&arg, code.str());
+}
+
 std::string RooFit::Experimental::codegenIntegralImpl(VerticalInterpPdf& arg,
                                                       int code,
                                                       const char* rangeName,
@@ -257,4 +285,31 @@
                        arg.quadraticAlgo());
 }
 
+std::string RooFit::Experimental::codegenIntegralImpl(RooParametricHist& arg,
+                                                      int code,
+                                                      const char* rangeName,
+                                                      CodegenContext& ctx) {
+  std::vector<double> diffs_flat;
+  std::vector<double> sums_flat;
+  if (arg.hasMorphs()) {
+    arg.getFlattenedMorphs(diffs_flat, sums_flat);
+  }
+  double xMin = rangeName ? arg.getObs().getMin(rangeName) : -1;
+  double xMax = rangeName ? arg.getObs().getMax(rangeName) : -1;  // if rangeName is null, these values won't be used
+
+  return ctx.buildCall("RooFit::Detail::MathFuncs::parametricHistIntegral",
+                       arg.getPars(),
+                       arg.getBins(),
+                       arg.getNBins(),
+                       arg.getCoeffList(),
+                       static_cast<int>(arg.getCoeffList().size()),
+                       arg.hasMorphs() ? diffs_flat : std::vector<double>{},
+                       arg.hasMorphs() ? sums_flat : std::vector<double>{},
+                       arg.getWidths(),
+                       arg.getSmoothRegion(),
+                       rangeName ? std::string("\"") + rangeName + "\"" : std::string("nullptr"),
+                       xMin,
+                       xMax);
+}
+
 #endif  // ROOT_VERSION_CODE >= ROOT_VERSION(6,36,0)