Skip to content

Commit

Permalink
[RF] Support RooBernstein in code generation for AD
Browse files Browse the repository at this point in the history 
Adding code generation and AD support for the final class in upstream
RooFit that is used in the CMS Higgs discovery workspace: the
RooBernstein.

Also:

  * fix missing variable warning in `RooCBShape.cxx`

  * fix AD support for RooLandau by using code that clad can digest

  * implement test for RooLandau fits with AD
  • Loading branch information
@guitargeek
guitargeek committed May 13, 2024
1 parent 44cff59 commit 50a1bdd
Show file tree
Hide file tree
Showing 7 changed files with 245 additions and 183 deletions.
78 changes: 40 additions & 38 deletions roofit/roofit/inc/RooBernstein.h
View file
@@ -1,54 +1,56 @@
/*****************************************************************************
* Project: RooFit *
* Package: RooFitModels *
* File: $Id$
* Authors: *
* Kyle Cranmer
* *
* *
* Redistribution and use in source and binary forms, *
* with or without modification, are permitted according to the terms *
* listed in LICENSE (http://roofit.sourceforge.net/license.txt) *
*****************************************************************************/
#ifndef ROO_BERNSTEIN
#define ROO_BERNSTEIN

#include "RooAbsPdf.h"
#include "RooTemplateProxy.h"
#include "RooRealVar.h"
#include "RooListProxy.h"
#include "RooAbsRealLValue.h"
#include <string>
/*
* Project: RooFit
*
* Copyright (c) 2024, CERN
*
* Redistribution and use in source and binary forms,
* with or without modification, are permitted according to the terms
* listed in LICENSE (http://roofit.sourceforge.net/license.txt)
*/

#ifndef RooFit_RooBernstein_h
#define RooFit_RooBernstein_h

#include <RooAbsPdf.h>
#include <RooAbsRealLValue.h>
#include <RooListProxy.h>
#include <RooRealVar.h>
#include <RooTemplateProxy.h>

class RooRealVar;
class RooArgList;
#include <string>

class RooBernstein : public RooAbsPdf {
public:
RooBernstein() = default;
RooBernstein(const char *name, const char *title, RooAbsRealLValue &_x, const RooArgList &_coefList);

RooBernstein() {}
RooBernstein(const char *name, const char *title,
RooAbsRealLValue& _x, const RooArgList& _coefList) ;
RooBernstein(const RooBernstein &other, const char *name = nullptr);

RooBernstein(const RooBernstein &other, const char *name = nullptr);
TObject *clone(const char *newname) const override { return new RooBernstein(*this, newname); }

TObject* clone(const char* newname) const override { return new RooBernstein(*this, newname); }
Int_t getAnalyticalIntegral(RooArgSet &allVars, RooArgSet &analVars, const char *rangeName = nullptr) const override;
double analyticalIntegral(Int_t code, const char *rangeName = nullptr) const override;
void selectNormalizationRange(const char *rangeName = nullptr, bool force = false) override;

Int_t getAnalyticalIntegral(RooArgSet& allVars, RooArgSet& analVars, const char* rangeName=nullptr) const override ;
double analyticalIntegral(Int_t code, const char* rangeName=nullptr) const override ;
void selectNormalizationRange(const char* rangeName=nullptr, bool force=false) override ;
void translate(RooFit::Detail::CodeSquashContext &ctx) const override;
std::string buildCallToAnalyticIntegral(Int_t code, const char *rangeName,
RooFit::Detail::CodeSquashContext &ctx) const override;

private:
void fillBuffer() const;
inline double xmin() const { return _buffer[_coefList.size()]; }
inline double xmax() const { return _buffer[_coefList.size() + 1]; }

RooTemplateProxy<RooAbsRealLValue> _x ;
RooListProxy _coefList ;
std::string _refRangeName ;
RooTemplateProxy<RooAbsRealLValue> _x;
RooListProxy _coefList;
std::string _refRangeName;
mutable std::vector<double> _buffer; ///<!

double evaluate() const override;
void doEval(RooFit::EvalContext &) const override;
inline bool canComputeBatchWithCuda() const override { return true; }
double evaluate() const override;
void doEval(RooFit::EvalContext &) const override;
inline bool canComputeBatchWithCuda() const override { return true; }

ClassDefOverride(RooBernstein,2) // Bernstein polynomial PDF
ClassDefOverride(RooBernstein, 2) // Bernstein polynomial PDF
};

#endif
174 changes: 56 additions & 118 deletions roofit/roofit/src/RooBernstein.cxx
View file
@@ -1,11 +1,12 @@
/*****************************************************************************
* Project: RooFit *
* Package: RooFitModels *
* @(#)root/roofit:$Id$
* Authors: *
* Kyle Cranmer *
* *
*****************************************************************************/
/*
* Project: RooFit
*
* Copyright (c) 2024, CERN
*
* Redistribution and use in source and binary forms,
* with or without modification, are permitted according to the terms
* listed in LICENSE (http://roofit.sourceforge.net/license.txt)
*/

/** \class RooBernstein
\ingroup Roofit
Expand All @@ -32,148 +33,85 @@ See also
http://www.idav.ucdavis.edu/education/CAGDNotes/Bernstein-Polynomials.pdf
**/

#include "RooBernstein.h"
#include "RooRealVar.h"
#include "RooArgList.h"
#include "RooBatchCompute.h"
#include <RooBernstein.h>
#include <RooRealVar.h>
#include <RooBatchCompute.h>

#include "TMath.h"

#include <cmath>
#include <RooFit/Detail/AnalyticalIntegrals.h>
#include <RooFit/Detail/EvaluateFuncs.h>

ClassImp(RooBernstein);

/// Constructor
////////////////////////////////////////////////////////////////////////////////

RooBernstein::RooBernstein(const char* name, const char* title,
RooAbsRealLValue& x, const RooArgList& coefList):
RooAbsPdf(name, title),
_x("x", "Dependent", this, x),
_coefList("coefficients","List of coefficients",this)
RooBernstein::RooBernstein(const char *name, const char *title, RooAbsRealLValue &x, const RooArgList &coefList)
: RooAbsPdf(name, title), _x("x", "Dependent", this, x), _coefList("coefficients", "List of coefficients", this)
{
_coefList.addTyped<RooAbsReal>(coefList);
_coefList.addTyped<RooAbsReal>(coefList);
_buffer.resize(_coefList.size() + 2);
}

////////////////////////////////////////////////////////////////////////////////

RooBernstein::RooBernstein(const RooBernstein &other, const char *name)
: RooAbsPdf(other, name),
_x("x", this, other._x),
_coefList("coefList", this, other._coefList),
_refRangeName{other._refRangeName}
_refRangeName{other._refRangeName},
_buffer{other._buffer}
{
}

////////////////////////////////////////////////////////////////////////////////

/// Force use of a given normalisation range.
/// Needed for functions or PDFs (e.g. RooAddPdf) whose shape depends on the choice of normalisation.
void RooBernstein::selectNormalizationRange(const char* rangeName, bool force)
void RooBernstein::selectNormalizationRange(const char *rangeName, bool force)
{
if (rangeName && (force || !_refRangeName.empty())) {
_refRangeName = rangeName;
}
if (rangeName && (force || !_refRangeName.empty())) {
_refRangeName = rangeName;
}
}

////////////////////////////////////////////////////////////////////////////////
void RooBernstein::fillBuffer() const
{
std::size_t n = _coefList.size();
for (std::size_t i = 0; i < n; ++i) {
_buffer[i] = static_cast<RooAbsReal &>(_coefList[i]).getVal();
}
std::tie(_buffer[n], _buffer[n + 1]) = _x->getRange(_refRangeName.empty() ? nullptr : _refRangeName.c_str());
}

double RooBernstein::evaluate() const
{
double xmax;
double xmin;
std::tie(xmin, xmax) = _x->getRange(_refRangeName.empty() ? nullptr : _refRangeName.c_str());
double x = (_x - xmin) / (xmax - xmin); // rescale to [0,1]
Int_t degree = _coefList.size() - 1; // n+1 polys of degree n

if(degree == 0) {

return static_cast<RooAbsReal &>(_coefList[0]).getVal();

} else if(degree == 1) {

double a0 = static_cast<RooAbsReal &>(_coefList[0]).getVal(); // c0
double a1 = static_cast<RooAbsReal &>(_coefList[1]).getVal() - a0; // c1 - c0
return a1 * x + a0;

} else if(degree == 2) {

double a0 = static_cast<RooAbsReal &>(_coefList[0]).getVal(); // c0
double a1 = 2 * (static_cast<RooAbsReal &>(_coefList[1]).getVal() - a0); // 2 * (c1 - c0)
double a2 = static_cast<RooAbsReal &>(_coefList[2]).getVal() - a1 - a0; // c0 - 2 * c1 + c2
return (a2 * x + a1) * x + a0;

} else if(degree > 2) {

double t = x;
double s = 1 - x;

double result = static_cast<RooAbsReal &>(_coefList[0]).getVal() * s;
for(Int_t i = 1; i < degree; i++) {
result = (result + t * TMath::Binomial(degree, i)
* static_cast<RooAbsReal &>(_coefList[i]).getVal()) * s;
t *= x;
}
result += t * static_cast<RooAbsReal &>(_coefList[degree]).getVal();

return result;
}
fillBuffer();
return RooFit::Detail::EvaluateFuncs::bernsteinEvaluate(_x, xmin(), xmax(), _buffer.data(), _coefList.size());
}

// in case list of arguments passed is empty
return TMath::SignalingNaN();
void RooBernstein::translate(RooFit::Detail::CodeSquashContext &ctx) const
{
fillBuffer();
ctx.addResult(this, ctx.buildCall("RooFit::Detail::EvaluateFuncs::bernsteinEvaluate", _x, xmin(), xmax(), _coefList,
_coefList.size()));
}

////////////////////////////////////////////////////////////////////////////////
/// Compute multiple values of Bernstein distribution.
void RooBernstein::doEval(RooFit::EvalContext & ctx) const
void RooBernstein::doEval(RooFit::EvalContext &ctx) const
{
const int nCoef = _coefList.size();
std::vector<double> extraArgs(nCoef+2);
for (int i=0; i<nCoef; i++)
extraArgs[i] = static_cast<RooAbsReal&>(_coefList[i]).getVal();
extraArgs[nCoef] = _x.min();
extraArgs[nCoef+1] = _x.max();

RooBatchCompute::compute(ctx.config(this), RooBatchCompute::Bernstein, ctx.output(), {ctx.at(_x)}, extraArgs);
fillBuffer();
RooBatchCompute::compute(ctx.config(this), RooBatchCompute::Bernstein, ctx.output(), {ctx.at(_x)}, _buffer);
}

////////////////////////////////////////////////////////////////////////////////

Int_t RooBernstein::getAnalyticalIntegral(RooArgSet& allVars, RooArgSet& analVars, const char* /*rangeName*/) const
Int_t RooBernstein::getAnalyticalIntegral(RooArgSet &allVars, RooArgSet &analVars, const char * /*rangeName*/) const
{

if (matchArgs(allVars, analVars, _x)) return 1;
return 0;
return matchArgs(allVars, analVars, _x) ? 1 : 0;
}

////////////////////////////////////////////////////////////////////////////////
double RooBernstein::analyticalIntegral(Int_t /*code*/, const char *rangeName) const
{
fillBuffer();
return RooFit::Detail::AnalyticalIntegrals::bernsteinIntegral(_x.min(rangeName), _x.max(rangeName), xmin(), xmax(),
_buffer.data(), _coefList.size());
}

double RooBernstein::analyticalIntegral(Int_t code, const char* rangeName) const
std::string RooBernstein::buildCallToAnalyticIntegral(Int_t /*code*/, const char *rangeName,
RooFit::Detail::CodeSquashContext &ctx) const
{
R__ASSERT(code==1) ;

double xmax;
double xmin;
std::tie(xmin, xmax) = _x->getRange(_refRangeName.empty() ? nullptr : _refRangeName.c_str());

const double xlo = (_x.min(rangeName) - xmin) / (xmax - xmin);
const double xhi = (_x.max(rangeName) - xmin) / (xmax - xmin);

Int_t degree= _coefList.size()-1; // n+1 polys of degree n
double norm(0) ;

double temp=0;
for (int i=0; i<=degree; ++i){
// for each of the i Bernstein basis polynomials
// represent it in the 'power basis' (the naive polynomial basis)
// where the integral is straight forward.
temp = 0;
for (int j=i; j<=degree; ++j){ // power basisŧ
temp += pow(-1.,j-i) * TMath::Binomial(degree, j) * TMath::Binomial(j,i) * (TMath::Power(xhi,j+1) - TMath::Power(xlo,j+1)) / (j+1);
}
temp *= static_cast<RooAbsReal &>(_coefList[i]).getVal(); // include coeff
norm += temp; // add this basis's contribution to total
}

return norm * (xmax - xmin);
fillBuffer(); // to get the right xmin() and xmax()
return ctx.buildCall("RooFit::Detail::AnalyticalIntegrals::bernsteinIntegral", _x.min(rangeName), _x.max(rangeName),
xmin(), xmax(), _coefList, _coefList.size());
}
13 changes: 6 additions & 7 deletions roofit/roofit/src/RooLandau.cxx
View file
Expand Up @@ -26,6 +26,8 @@ Landau distribution p.d.f
#include "RooRandom.h"
#include "RooBatchCompute.h"

#include "RooFit/Detail/EvaluateFuncs.h"

#include "TMath.h"
#include "Math/ProbFunc.h"

Expand Down Expand Up @@ -56,14 +58,14 @@ RooLandau::RooLandau(const RooLandau& other, const char* name) :

double RooLandau::evaluate() const
{
return TMath::Landau(x, mean, sigma);
return RooFit::Detail::EvaluateFuncs::landauEvaluate(x, mean, sigma);
}

////////////////////////////////////////////////////////////////////////////////

void RooLandau::translate(RooFit::Detail::CodeSquashContext &ctx) const
{
ctx.addResult(this, ctx.buildCall("TMath::Landau", x, mean, sigma));
ctx.addResult(this, ctx.buildCall("RooFit::Detail::EvaluateFuncs::landauEvaluate", x, mean, sigma));
}

////////////////////////////////////////////////////////////////////////////////
Expand All @@ -78,15 +80,12 @@ void RooLandau::doEval(RooFit::EvalContext &ctx) const

Int_t RooLandau::getGenerator(const RooArgSet& directVars, RooArgSet &generateVars, bool /*staticInitOK*/) const
{
if (matchArgs(directVars,generateVars,x)) return 1 ;
return 0;
return matchArgs(directVars,generateVars,x) ? 1 : 0;
}

Int_t RooLandau::getAnalyticalIntegral(RooArgSet &allVars, RooArgSet &analVars, const char * /*rangeName*/) const
{
if (matchArgs(allVars, analVars, x))
return 1;
return 0;
return matchArgs(allVars, analVars, x) ? 1 : 0;
}

Double_t RooLandau::analyticalIntegral(Int_t /*code*/, const char *rangeName) const
Expand Down

0 comments on commit 50a1bdd

Please sign in to comment.