[RF] Bring vectorizing EvalBackends to createChi2 / chi2FitTo

roofit/codegen/src/CodegenImpl.cxx

@@ -546,9 +546,65 @@ void codegenImpl(RooLognormal &arg, CodegenContext &ctx)
    ctx.addResult(&arg, ctx.buildCall(mathFunc(funcName), arg.getX(), arg.getShapeK(), arg.getMedian()));
 }
 
+namespace {
+
+void codegenChi2(RooFit::Detail::RooNLLVarNew &arg, CodegenContext &ctx)
+{
+   using FuncMode = RooFit::Detail::RooNLLVarNew::FuncMode;
+
+   std::string resName = RooFit::Detail::makeValidVarName(arg.GetName()) + "Result";
+   ctx.addResult(&arg, resName);
+   ctx.addToGlobalScope("double " + resName + " = 0.0;\n");
+
+   // DataError::None means "no errors": every bin contributes zero.
+   if (arg.chi2ErrorType() == RooDataHist::None) {
+      return;
+   }
+
+   // Compute the per-bin normalization factor (constant with respect to the
+   // loop).
+   std::string normFactor;
+   if (arg.funcMode() == FuncMode::Function) {
+      normFactor = "1.0";
+   } else if (arg.funcMode() == FuncMode::ExtendedPdf) {
+      normFactor = ctx.getResult(*arg.expectedEvents());
+   } else { // Pdf
+      std::string weightSumName = RooFit::Detail::makeValidVarName(arg.GetName()) + "WeightSum";
+      ctx.addToGlobalScope("double " + weightSumName + " = 0.0;\n");
+      {
+         auto scope = ctx.beginLoop(&arg);
+         ctx.addToCodeBody(weightSumName + " += " + ctx.getResult(arg.weightVar()) + ";\n");
+      }
+      normFactor = weightSumName;
+   }
+
+   auto scope = ctx.beginLoop(&arg);
+   const std::string mu =
+      "(" + ctx.getResult(arg.func()) + " * " + ctx.getResult(*arg.binVolumes()) + " * " + normFactor + ")";
+
+   std::string term;
+   switch (arg.chi2ErrorType()) {
+   case RooDataHist::Expected: term = ctx.buildCall(mathFunc("chi2Expected"), mu, arg.weightVar()); break;
+   case RooDataHist::SumW2:
+      term = ctx.buildCall(mathFunc("chi2Symmetric"), mu, arg.weightVar(), arg.weightSquaredVar());
+      break;
+   case RooDataHist::Poisson:
+      term = ctx.buildCall(mathFunc("chi2Asymmetric"), mu, arg.weightVar(), *arg.weightErrLo(), *arg.weightErrHi());
+      break;
+   default: break;
+   }
+   ctx.addToCodeBody(&arg, resName + " += " + term + ";");
+}
+
+} // namespace
+
 void codegenImpl(RooFit::Detail::RooNLLVarNew &arg, CodegenContext &ctx)
 {
-   if (arg.binnedL() && !arg.pdf().getAttribute("BinnedLikelihoodActiveYields")) {
+   if (arg.statistic() == RooFit::Detail::RooNLLVarNew::Statistic::Chi2) {
+      return codegenChi2(arg, ctx);
+   }
+
+   if (arg.binnedL() && !arg.func().getAttribute("BinnedLikelihoodActiveYields")) {
       std::stringstream errorMsg;
       errorMsg << "codegen: binned likelihood optimization is only supported when raw pdf "
                   "values can be interpreted as yields."
@@ -579,7 +635,7 @@ void codegenImpl(RooFit::Detail::RooNLLVarNew &arg, CodegenContext &ctx)
    // brackets of the loop is written at the end of the scopes lifetime.
    {
       auto scope = ctx.beginLoop(&arg);
-      std::string term = ctx.buildCall(mathFunc("nll"), arg.pdf(), arg.weightVar(), arg.binnedL(), 0);
+      std::string term = ctx.buildCall(mathFunc("nll"), arg.func(), arg.weightVar(), arg.binnedL(), 0);
       ctx.addToCodeBody(&arg, resName + " += " + term + ";");
    }
    if (arg.expectedEvents()) {

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

@@ -22,6 +22,7 @@
 
 #include <algorithm>
 #include <cmath>
+#include <limits>
 #include <stdexcept>
 
 namespace RooFit::Detail::MathFuncs {
@@ -401,6 +402,53 @@ inline double effProd(double eff, double pdf)
    return eff * pdf;
 }
 
+/// Chi-squared contribution of one bin with "expected" errors:
+/// \f$ \sigma^2 = \mu \f$. Empty/no-prediction bins contribute zero; bins with
+/// non-positive \f$ \mu \f$ but non-empty data yield NaN (to let the minimizer
+/// recover).
+inline double chi2Expected(double mu, double weight)
+{
+   if (mu == 0.0 && weight == 0.0) {
+      return 0.0;
+   }
+   if (mu <= 0.0) {
+      return std::numeric_limits<double>::quiet_NaN();
+   }
+   const double diff = mu - weight;
+   return diff * diff / mu;
+}
+
+/// Chi-squared contribution of one bin with a user-supplied symmetric error
+/// squared (e.g. `SumW2` weights from the data).
+inline double chi2Symmetric(double mu, double weight, double sigma2)
+{
+   if (sigma2 == 0.0 && mu == 0.0 && weight == 0.0) {
+      return 0.0;
+   }
+   if (sigma2 <= 0.0) {
+      return std::numeric_limits<double>::quiet_NaN();
+   }
+   const double diff = mu - weight;
+   return diff * diff / sigma2;
+}
+
+/// Chi-squared contribution of one bin with asymmetric (Poisson-style) data
+/// errors. The side facing the prediction is used: `errHi` when
+/// \f$ \mu > \mathrm{weight} \f$, otherwise `errLo`.
+inline double chi2Asymmetric(double mu, double weight, double errLo, double errHi)
+{
+   const double diff = mu - weight;
+   const double err = diff > 0.0 ? errHi : errLo;
+   const double sigma2 = err * err;
+   if (sigma2 == 0.0 && mu == 0.0 && weight == 0.0) {
+      return 0.0;
+   }
+   if (sigma2 <= 0.0) {
+      return std::numeric_limits<double>::quiet_NaN();
+   }
+   return diff * diff / sigma2;
+}
+
 inline double nll(double pdf, double weight, int binnedL, int doBinOffset)
 {
    if (binnedL) {

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

@@ -18,6 +18,7 @@
 
 #include <RooAbsPdf.h>
 #include <RooAbsReal.h>
+#include <RooDataHist.h>
 #include <RooGlobalFunc.h>
 #include <RooTemplateProxy.h>
 
@@ -29,20 +30,37 @@ namespace Detail {
 class RooNLLVarNew : public RooAbsReal {
 
 public:
-   // The names for the weight variables that the RooNLLVarNew expects
+   // The names for the special variables that the RooNLLVarNew expects
    static constexpr const char *weightVarName = "_weight";
    static constexpr const char *weightVarNameSumW2 = "_weight_sumW2";
-
-   RooNLLVarNew(const char *name, const char *title, RooAbsPdf &pdf, RooArgSet const &observables, bool isExtended,
-                RooFit::OffsetMode offsetMode);
+   static constexpr const char *binVolumeVarName = "_bin_volume";
+   static constexpr const char *weightErrorLoVarName = "_weight_err_lo";
+   static constexpr const char *weightErrorHiVarName = "_weight_err_hi";
+
+   enum class Statistic {
+      NLL,
+      Chi2
+   };
+
+   /// Configuration struct for the unified constructor. Note that `offsetMode`
+   /// only applies to `Statistic::NLL`, and `chi2ErrorType` only applies to
+   /// `Statistic::Chi2`.
+   struct Config {
+      Statistic statistic = Statistic::NLL;
+      bool extended = false;
+      RooFit::OffsetMode offsetMode = RooFit::OffsetMode::None;
+      RooDataHist::ErrorType chi2ErrorType = RooDataHist::Expected;
+   };
+
+   RooNLLVarNew(const char *name, const char *title, RooAbsReal &func, RooArgSet const &observables, Config const &cfg);
    RooNLLVarNew(const RooNLLVarNew &other, const char *name = nullptr);
    TObject *clone(const char *newname) const override { return new RooNLLVarNew(*this, newname); }
 
    /// Return default level for MINUIT error analysis.
-   double defaultErrorLevel() const override { return 0.5; }
+   double defaultErrorLevel() const override { return _statistic == Statistic::Chi2 ? 1.0 : 0.5; }
 
    void doEval(RooFit::EvalContext &) const override;
-   bool canComputeBatchWithCuda() const override { return !_binnedL; }
+   bool canComputeBatchWithCuda() const override { return _statistic == Statistic::NLL && !_binnedL; }
    bool isReducerNode() const override { return true; }
 
    void setPrefix(std::string const &prefix);
@@ -55,28 +73,49 @@ class RooNLLVarNew : public RooAbsReal {
 
    void setSimCount(int simCount) { _simCount = simCount; }
 
-   RooAbsPdf const &pdf() const { return *_pdf; }
+   enum class FuncMode {
+      Pdf,
+      ExtendedPdf,
+      Function
+   };
+
+   RooAbsReal const &func() const { return *_func; }
    RooAbsReal const &weightVar() const { return *_weightVar; }
+   RooAbsReal const &weightSquaredVar() const { return *_weightSquaredVar; }
    bool binnedL() const { return _binnedL; }
    int simCount() const { return _simCount; }
+   Statistic statistic() const { return _statistic; }
+   FuncMode funcMode() const { return _funcMode; }
+   RooDataHist::ErrorType chi2ErrorType() const { return _chi2ErrorType; }
    RooAbsReal const *expectedEvents() const { return _expectedEvents ? &**_expectedEvents : nullptr; }
+   RooAbsReal const *binVolumes() const { return _binVolumes ? &**_binVolumes : nullptr; }
+   RooAbsReal const *weightErrLo() const { return _weightErrLo ? &**_weightErrLo : nullptr; }
+   RooAbsReal const *weightErrHi() const { return _weightErrHi ? &**_weightErrHi : nullptr; }
 
 private:
    double evaluate() const override { return _value; }
    void resetWeightVarNames();
    void finalizeResult(RooFit::EvalContext &, ROOT::Math::KahanSum<double> result, double weightSum) const;
    void fillBinWidthsFromPdfBoundaries(RooAbsReal const &pdf, RooArgSet const &observables);
    void doEvalBinnedL(RooFit::EvalContext &, std::span<const double> preds, std::span<const double> weights) const;
+   void doEvalChi2(RooFit::EvalContext &, std::span<const double> preds, std::span<const double> weights,
+                   std::span<const double> weightsSumW2) const;
 
-   RooTemplateProxy<RooAbsPdf> _pdf;
+   RooTemplateProxy<RooAbsReal> _func;
    RooTemplateProxy<RooAbsReal> _weightVar;
    RooTemplateProxy<RooAbsReal> _weightSquaredVar;
    std::unique_ptr<RooTemplateProxy<RooAbsReal>> _expectedEvents;
    std::unique_ptr<RooTemplateProxy<RooAbsPdf>> _offsetPdf;
+   std::unique_ptr<RooTemplateProxy<RooAbsReal>> _binVolumes;
+   std::unique_ptr<RooTemplateProxy<RooAbsReal>> _weightErrLo;
+   std::unique_ptr<RooTemplateProxy<RooAbsReal>> _weightErrHi;
    bool _weightSquared = false;
    bool _binnedL = false;
    bool _doOffset = false;
    bool _doBinOffset = false;
+   Statistic _statistic = Statistic::NLL;
+   FuncMode _funcMode = FuncMode::Pdf;
+   RooDataHist::ErrorType _chi2ErrorType = RooDataHist::Expected;
    int _simCount = 1;
    std::string _prefix;
    std::vector<double> _binw;

roofit/roofitcore/src/BatchModeDataHelpers.cxx

@@ -15,6 +15,7 @@
 #include "RooFit/BatchModeDataHelpers.h"
 
 #include <RooAbsData.h>
+#include <RooDataHist.h>
 #include <RooRealVar.h>
 #include <RooSimultaneous.h>
 
@@ -91,6 +92,39 @@ getSingleDataSpans(RooAbsData const &data, std::string_view rangeName, std::stri
       insert(RooFit::Detail::RooNLLVarNew::weightVarNameSumW2, weightSumW2);
    }
 
+   // For RooDataHist datasets, also publish per-bin volumes and asymmetric
+   // Poisson errors. These are consumed by the chi2 evaluation path in
+   // RooNLLVarNew.
+   if (auto const *dataHist = dynamic_cast<RooDataHist const *>(&data)) {
+      auto binVolumes = dataHist->binVolumes(0, static_cast<std::size_t>(data.numEntries()));
+      buffers.emplace();
+      auto &bufferBinVol = buffers.top();
+      bufferBinVol.reserve(nNonZeroWeight);
+      buffers.emplace();
+      auto &bufferErrLo = buffers.top();
+      bufferErrLo.reserve(nNonZeroWeight);
+      buffers.emplace();
+      auto &bufferErrHi = buffers.top();
+      bufferErrHi.reserve(nNonZeroWeight);
+
+      auto *dataHistMutable = const_cast<RooDataHist *>(dataHist);
+      for (std::size_t i = 0; i < binVolumes.size(); ++i) {
+         if (hasZeroWeight[i]) {
+            continue;
+         }
+         bufferBinVol.push_back(binVolumes[i]);
+         dataHistMutable->get(static_cast<int>(i));
+         double lo = 0.0;
+         double hi = 0.0;
+         dataHistMutable->weightError(lo, hi, RooAbsData::Poisson);
+         bufferErrLo.push_back(lo);
+         bufferErrHi.push_back(hi);
+      }
+      insert(RooFit::Detail::RooNLLVarNew::binVolumeVarName, {bufferBinVol.data(), bufferBinVol.size()});
+      insert(RooFit::Detail::RooNLLVarNew::weightErrorLoVarName, {bufferErrLo.data(), bufferErrLo.size()});
+      insert(RooFit::Detail::RooNLLVarNew::weightErrorHiVarName, {bufferErrHi.data(), bufferErrHi.size()});
+   }
+
    // Get the real-valued batches and cast the also to double branches to put in
    // the data map
    for (auto const &item : data.getBatches(0, nEvents)) {