[RF] Implement codegen support for RooFunctorBinding and friends

This enables easily plugging in user-defined functions including their
gradients into RooFit computation graphs, which also works from Python
thanks to the `std::function` pythonization.

Author: guitargeek

roofit/codegen/inc/RooFit/CodegenImpl.h

@@ -25,8 +25,8 @@ class RooAddPdf;
 class RooAddition;
 class RooBernstein;
 class RooBifurGauss;
-class RooCategory;
 class RooCBShape;
+class RooCategory;
 class RooChebychev;
 class RooConstVar;
 class RooConstraintSum;
@@ -35,13 +35,18 @@ class RooEfficiency;
 class RooExponential;
 class RooExtendPdf;
 class RooFormulaVar;
+class RooFunctor1DBinding;
+class RooFunctor1DPdfBinding;
+class RooFunctorBinding;
+class RooFunctorPdfBinding;
 class RooGamma;
 class RooGaussian;
 class RooGenericPdf;
 class RooHistFunc;
 class RooHistPdf;
 class RooLandau;
 class RooLognormal;
+class RooMultiPdf;
 class RooMultiVarGaussian;
 class RooParamHistFunc;
 class RooPoisson;
@@ -55,7 +60,6 @@ class RooRealSumPdf;
 class RooRealVar;
 class RooRecursiveFraction;
 class RooUniform;
-class RooMultiPdf;
 class RooWrapperPdf;
 
 namespace RooStats {
@@ -96,6 +100,10 @@ void codegenImpl(RooEfficiency &arg, CodegenContext &ctx);
 void codegenImpl(RooExponential &arg, CodegenContext &ctx);
 void codegenImpl(RooExtendPdf &arg, CodegenContext &ctx);
 void codegenImpl(RooFormulaVar &arg, CodegenContext &ctx);
+void codegenImpl(RooFunctor1DBinding &arg, CodegenContext &ctx);
+void codegenImpl(RooFunctor1DPdfBinding &arg, CodegenContext &ctx);
+void codegenImpl(RooFunctorBinding &arg, CodegenContext &ctx);
+void codegenImpl(RooFunctorPdfBinding &arg, CodegenContext &ctx);
 void codegenImpl(RooGamma &arg, CodegenContext &ctx);
 void codegenImpl(RooGaussian &arg, CodegenContext &ctx);
 void codegenImpl(RooGenericPdf &arg, CodegenContext &ctx);

roofit/codegen/src/CodegenImpl.cxx

@@ -31,6 +31,8 @@
 #include <RooFit/Detail/RooNLLVarNew.h>
 #include <RooFit/Detail/RooNormalizedPdf.h>
 #include <RooFormulaVar.h>
+#include <RooFunctor1DBinding.h>
+#include <RooFunctorBinding.h>
 #include <RooGamma.h>
 #include <RooGaussian.h>
 #include <RooGenericPdf.h>
@@ -62,8 +64,9 @@
 
 #include <TInterpreter.h>
 
-namespace RooFit {
-namespace Experimental {
+#include <unordered_set>
+
+namespace RooFit::Experimental {
 
 namespace {
 
@@ -371,6 +374,92 @@ void codegenImpl(RooConstraintSum &arg, CodegenContext &ctx)
    ctx.addResult(&arg, ctx.buildCall(mathFunc("constraintSum"), arg.list(), arg.list().size()));
 }
 
+// Generate RooFit codegen wrappers for RooFunctorBinding and similar objects,
+// emitting both the primal function call and its gradient pullback for
+// Clad-based AD.
+template <class RooArg_t>
+void functorCodegenImpl(RooArg_t &arg, RooArgList const &variables, CodegenContext &ctx)
+{
+   if (!arg.function()->HasGradient()) {
+      std::stringstream errorMsg;
+      errorMsg << "Functor wrapped by \"" << arg.GetName() << "\" doesn't provide a gradient function."
+               << " RooFit codegen is therefore not supported.";
+      oocoutE(&arg, InputArguments) << errorMsg.str() << std::endl;
+      throw std::runtime_error(errorMsg.str());
+   }
+
+   std::string funcAddrStr = TString::Format("0x%zx", reinterpret_cast<std::size_t>(arg.function())).Data();
+   std::string wrapperName = "roo_functor_" + funcAddrStr;
+
+   static std::unordered_set<std::string> wrapperNames;
+
+   if (wrapperNames.find(wrapperName) == wrapperNames.end()) {
+
+      wrapperNames.insert(wrapperName);
+
+      std::string pullbackName = wrapperName + "_pullback";
+      std::string nStr = std::to_string(std::size(variables));
+
+      std::string type;
+      if constexpr (std::is_same_v<RooArg_t, RooFunctor1DBinding> || std::is_same_v<RooArg_t, RooFunctor1DPdfBinding>)
+         type = "::ROOT::Math::IGradientFunctionOneDim";
+      else
+         type = "::ROOT::Math::IGradientFunctionMultiDim";
+
+      std::string funcAddrCasted = "reinterpret_cast<" + type + " const *>(" + funcAddrStr + ")";
+
+      std::string code;
+
+      code += "double " + wrapperName +
+              "(double const *x) {\n"
+              "   return " +
+              funcAddrCasted +
+              "->operator()(x);\n"
+              "}\n\n"
+              "namespace clad::custom_derivatives {\n\n"
+              "void " +
+              pullbackName +
+              "(double const* x, double d_y, double *d_x) {\n"
+              "   double output[" +
+              nStr +
+              "]{};\n"
+              "   " +
+              funcAddrCasted +
+              "->Gradient(x, output);\n"
+              "   for (int i = 0; i < " +
+              nStr +
+              "; ++i) {\n"
+              "      d_x[i] += output[i] * d_y;\n"
+              "   }\n"
+              "}\n"
+              "} // namespace clad::custom_derivatives\n";
+
+      gInterpreter->Declare(code.c_str());
+   }
+
+   ctx.addResult(&arg, ctx.buildCall(wrapperName, variables));
+}
+
+void codegenImpl(RooFunctor1DBinding &arg, CodegenContext &ctx)
+{
+   functorCodegenImpl(arg, arg.variable(), ctx);
+}
+
+void codegenImpl(RooFunctor1DPdfBinding &arg, CodegenContext &ctx)
+{
+   functorCodegenImpl(arg, arg.variable(), ctx);
+}
+
+void codegenImpl(RooFunctorBinding &arg, CodegenContext &ctx)
+{
+   functorCodegenImpl(arg, arg.variables(), ctx);
+}
+
+void codegenImpl(RooFunctorPdfBinding &arg, CodegenContext &ctx)
+{
+   functorCodegenImpl(arg, arg.variables(), ctx);
+}
+
 void codegenImpl(RooGamma &arg, CodegenContext &ctx)
 {
    ctx.addResult(&arg, ctx.buildCall("TMath::GammaDist", arg.getX(), arg.getGamma(), arg.getMu(), arg.getBeta()));
@@ -909,5 +998,4 @@ std::string codegenIntegralImpl(RooUniform &arg, int code, const char *rangeName
    return doubleToString(arg.analyticalIntegral(code, rangeName));
 }
 
-} // namespace Experimental
-} // namespace RooFit
+} // namespace RooFit::Experimental

roofit/roofit/inc/RooFunctor1DBinding.h

@@ -18,7 +18,6 @@
 #include "RooListProxy.h"
 #include "RooAbsPdf.h"
 #include "RooRealProxy.h"
-#include "RooMsgService.h"
 #include "Math/IFunction.h"
 
 
@@ -32,49 +31,45 @@ RooAbsPdf*  bindPdf(const char* name, const ROOT::Math::IBaseFunctionOneDim& fto
 
 class RooFunctor1DBinding : public RooAbsReal {
 public:
-  RooFunctor1DBinding() : func(nullptr) {
-    // Default constructor
-  } ;
+  RooFunctor1DBinding() = default;
   RooFunctor1DBinding(const char *name, const char *title, const ROOT::Math::IBaseFunctionOneDim& ftor, RooAbsReal& var);
   RooFunctor1DBinding(const RooFunctor1DBinding& other, const char* name=nullptr) ;
   TObject* clone(const char* newname=nullptr) const override { return new RooFunctor1DBinding(*this,newname); }
   void printArgs(std::ostream& os) const override ;
 
+  ROOT::Math::IBaseFunctionOneDim const *function() const { return func; }
+  RooAbsReal const &variable() const { return *var; }
+
 protected:
 
   double evaluate() const override ;
 
-  const ROOT::Math::IBaseFunctionOneDim* func ;    // Functor
+  const ROOT::Math::IBaseFunctionOneDim *func = nullptr; // Functor
   RooRealProxy                       var ;    // Argument reference
 
-
-private:
-
   ClassDefOverride(RooFunctor1DBinding,1) // RooAbsReal binding to a ROOT::Math::IBaseFunctionOneDim
 };
 
 
 
 class RooFunctor1DPdfBinding : public RooAbsPdf {
 public:
-  RooFunctor1DPdfBinding() : func(nullptr) {
-    // Default constructor
-  } ;
+  RooFunctor1DPdfBinding() = default;
   RooFunctor1DPdfBinding(const char *name, const char *title, const ROOT::Math::IBaseFunctionOneDim& ftor, RooAbsReal& vars);
   RooFunctor1DPdfBinding(const RooFunctor1DPdfBinding& other, const char* name=nullptr) ;
   TObject* clone(const char* newname=nullptr) const override { return new RooFunctor1DPdfBinding(*this,newname); }
   void printArgs(std::ostream& os) const override ;
 
+  ROOT::Math::IBaseFunctionOneDim const *function() const { return func; }
+  RooAbsReal const &variable() const { return *var; }
+
 protected:
 
   double evaluate() const override ;
 
-  const ROOT::Math::IBaseFunctionOneDim* func ;    // Functor
+  ROOT::Math::IBaseFunctionOneDim const *func = nullptr; // Functor
   RooRealProxy                           var ;    // Argument reference
 
-
-private:
-
   ClassDefOverride(RooFunctor1DPdfBinding,1) // RooAbsPdf binding to a ROOT::Math::IBaseFunctionOneDim
 };
 

roofit/roofit/inc/RooFunctorBinding.h

@@ -18,7 +18,6 @@
 #include "RooListProxy.h"
 #include "RooAbsPdf.h"
 #include "RooRealProxy.h"
-#include "RooMsgService.h"
 #include "Math/IFunction.h"
 
 namespace RooFit {
@@ -37,6 +36,9 @@ class RooFunctorBinding : public RooAbsReal {
   inline ~RooFunctorBinding() override { delete[] x ; }
   void printArgs(std::ostream& os) const override ;
 
+  ROOT::Math::IBaseFunctionMultiDim const *function() const { return func; }
+  RooArgList const &variables() const { return vars; }
+
 protected:
 
   double evaluate() const override ;
@@ -45,9 +47,6 @@ class RooFunctorBinding : public RooAbsReal {
   RooListProxy vars; // Argument reference
   double *x = nullptr; // Argument value array
 
-
-private:
-
   ClassDefOverride(RooFunctorBinding,1) // RooAbsReal binding to a ROOT::Math::IBaseFunctionMultiDim
 };
 
@@ -62,6 +61,9 @@ class RooFunctorPdfBinding : public RooAbsPdf {
   inline ~RooFunctorPdfBinding() override { delete[] x ; }
   void printArgs(std::ostream& os) const override ;
 
+  ROOT::Math::IBaseFunctionMultiDim const *function() const { return func; }
+  RooArgList const &variables() const { return vars; }
+
 protected:
 
   double evaluate() const override ;
@@ -70,9 +72,6 @@ class RooFunctorPdfBinding : public RooAbsPdf {
   RooListProxy vars ;    // Argument reference
   double *x = nullptr; // Argument value array
 
-
-private:
-
   ClassDefOverride(RooFunctorPdfBinding,1) // RooAbsPdf binding to a ROOT::Math::IBaseFunctionMultiDim
 };