 |
ROOT
Reference Guide |
|
| |
ROOT
Reference Guide |
|
Developer guide on how to add support for Automatic Differentiation via code generation.
RooFit is a statistical data analysis tool, widely used in scientific research, especially in the high-energy physics (HEP) field. It is an extension of the ROOT framework, a C++ based data analysis framework that provides tools for data storage, analysis, and visualization. RooFit provides a set of tools/classes to define and evaluate probability density functions (PDFs), perform maximum likelihood fits, perform statistical tests, etc.
RooFit is used to reduce statistical models (functions) to find a set of parameters that minimize the value of the function. This minimization happens via one of several methods relying heavily on the computation of derivatives of the function with respect to its free parameters. Currently, the computation of Numerical Derivatives is the most time consuming component of RooFit [^1]. On the other hand, derivatives computed using the Automatic Differentiation tool Clad have been shown to be far more efficient [^2].
Main Advantage of using AD with RooFit: efficient and more precise derivatives. It computes derivatives with high precision, avoiding the errors that may arise from approximating derivatives using finite differences.
RooFit is an extensive toolkit. The initiative to add AD support has been started, but has not yet achieved full coverage for the models defined/maintained in RooFit.
Clad is a C++ plugin for Clang. It implements a technique called Source Code Transformation to enable AD support.
Source Code Transformation takes the source code (that needs to be differentiated) as the input and generates an output code that represents the derivative of the input. This output code can be used instead of the input code for more efficient compilation.
For more technical details, please see the following paper:
Automatic Differentiation of Binned Likelihoods With RooFit and Clad
In RooFit jargon, what is meant by a "RooFit class" is a class inheriting from RooAbsArg that represents a mathematical function, a PDF, or any other transformation of inputs that are also represented by RooAbsArg objects. Almost all final classes deriving from RooAbsArg have RooAbsReal as an intermediate base class, which is the base class for all RooAbsArg that represent real-valued nodes in the computation graph. As such RooFit objects are so prevalent in practice, the names RooAbsArg and RooAbsReal are used interchangeably in this guide.
Users take these classes to build a computational graph that represents the PDF (also called "model") that they want to use for fitting the data. The user then passes his final PDF and a RooAbsData object to the RooAbsPdf::fitTo() method, which implicitly creates a negative-log likelihood (NLL) that RooFit minimizes for parameter estimation. The NLL object, internally created by RooAbsPdf::createNLL(), is a RooAbsArg itself. In technical terms, it's another larger computation graph that emcompasses the computation graph representing the PDF.
To enable source code transformation AD for RooFit NLLs with Clad, RooFit has a routine that can traverse a computation graph made of RooAbsArg objects and transform it to much simpler C++ code that mathematically represents the same computation, but without any overhead that is hard to digest by the AD tool.
On a high level, this code generation is implemented as follows:
In summary, the important ingredient to enable AD in RooFit is to support the C++ code generation from RooFit classes.
Clad requires the code for models to be within a single translation unit. This has been implemented in RooFit by moving the computational aspects of the RooFit class to free functions in a single header file named EvaluateFuncs.
Furthermore, the RooAbsArg::translate() function needs to be overridden to specify how the class is translating to C++ code that is using the aforementioned free function.
To add AD support to an existing RooFit class, following changes are required.
1. Extract logic into a separate file Implement what your class is supposed to do as a free function in EvaluateFuncs. This implementation must be compatible with the syntax supported by Clad.
2. Refactor evaluate(): Refactor the existing evaluate() function to use the EvaluateFuncs.h implementation. This is optional, but can reduce code duplication and potential for bugs. This may require some effort if an extensive caching infrastructure is used in your model.
3. Add translate(): RooFit classes are extended using a (typically) simple translate() function that extracts the mathematically differentiable properties out of the RooFit classes that make up the statistical model.
The translate() function helps implement the Code Squashing logic that is used to optimize numerical evaluations. It accomplishes this by using a small subset of helper functions that are available in the RooFit::Detail::CodeSquashContext and RooFuncWrapper classes (see Appendix B). It converts a RooFit expression into a form that can be efficiently evaluated by Clad.
The translate() function returns an std::string representing the underlying mathematical notation of the class as code, that can later be concatenated into a single string representing the entire model. This string of code is then just-in-time compiled by Cling (a C++ interpreter for Root).
4. analyticalIntegral() Use Case: If your class includes (or should include) the analyticalIntegral() function, then a simple buildCallToAnalyticIntegral() function needs to be created to help call the analyticalIntegral() function.
Let us take the RooPoisson.cxx class as an example.
First step is to locate the evaluate() function. Most RooFit classes implement this function.
RooFit internally calls the
evaluate()function to evaluate a single node in a compute graph.
Following is a code snippet from RooPoisson before it had AD support.
TMath::Poisson() is a simple mathematical function. For this example, the relevant part is return TMath::Poisson(k,mean);. This needs to be extracted into the EvaluateFuncs.h file and the fully qualified name of the function referencing that file should be used here instead.
Following is a code snippet from RooPoisson after it has AD support.
Note that the evaluate() function was refactored in such a way that the mathematical parts were moved to an inline function in a separate header file named EvaluateFuncs, so that Clad could see and differentiate that function. The rest of the contents of the function remain unchanged.
All contents of the
evaluate()function don't always need to be pulled
out, only the required parts (mathematical logic) should be moved to EvaluateFuncs.
What is EvaluateFuncs?
Moving away from the class-based hierarchy design, EvaluateFuncs.h a simply a flat file of function implementations.
This file is required since Clad will not be able to see anything that is not inlined and explicitly available to it during compilation (since it has to be in the same translation). So other than of generating these functions on the fly, your only other option is to place these functions in a separate header file and make them inline.
Theoretically, multiple header files can also be used and then mashed together.
Directory path: roofit/roofitcore/inc/RooFit/Detail/EvaluateFuncs.h
translate() Example 1: Continuing our RooPoisson example:
To translate the RooPoisson class, create a translate function and in it include a call to the updated function.
Here we can see that the name of the variable x (remember that "x" is a member of RooPoisson) is retrieved and stored in the xName variable. Next, there's an if condition that does an operation on x (may or may not round it to the nearest integer, depending on the condition).
The important part is where the addResult() function helps add the result of evaluating the Poisson function to the context (ctx). It uses the buildCall() method to construct a function call to the fully qualified name of poissonEvaluate (which now resides in the EvaluateFuncs file), with arguments xName and mean.
Essentially, the translate() function constructs a function call to evaluate the Poisson function using 'x' and 'mean' variables, and adds the result to the context.
Helper Functions:
getResult() helps lookup the result of a child node (the string that the child node previously saved in a variable using the addResult() function).addResult() It may include a function call, an expression, or something more complicated. For a specific class, it will add whatever is represented on the right hand side to the result of that class, which can then be propagated in the rest of the compute graph.translate() function, it is important to call addResult() since this is what enables the squashing to happen.translate() Example 2: Following is a code snippet from RooGaussian.cxx after it has AD support.
Here we can see that the translate() function constructs a function call using the buildCall() method. It specifies the fully qualified name of the gaussianEvalaute function (which is now part of the EvaluateFuncs file), and includes the x, mean, and sigma variables as arguments to this function call.
Helper Function:
buildCall() helps build a function call. Requires the fully qualified name (RooFit::Detail::EvaluateFuncs::gaussianEvaluate) of the function. When this external buildCall() function is called, internally, the getResult() function is called on the input RooFit objects (e.g., x, mean, sigma). That's the only way to propagate these upwards into the compute graph.translate() Example 3: A more complicated example of a translate() function can be seen here:
Source: - RooNLLVarNew
The complexity of the translate() function in this example can be attributed to the more complex scenarios/operations specific to the computation of negative log-likelihood (NLL) values for probability density functions (PDFs) in RooFit, especially for simultaneous fits (multiple simultaneous PDFs being considered) and binned likelihoods (adding further complexity).
In this example, the translate() function generates code to compute the Negative Log likelihood (NLL). We can see that the intermediate result variable resName is added to the context so that it can be accessed and used in the generated code. This variable is made available globally (using addToGlobalScope()).
If a weight sum is needed, then it creates a loop, and weightSumName is accumulated with the weight variable. Otherwise, if there are multiple simultaneous PDFs, then it adds a term to the result that scales with the logarithm of the count of simultaneous PDFs. The rest of the function body (including the loop scope with NLL computation) has omitted from this example to keep it brief.
Helper functions:
makeValidVarName() helps get a valid name from the name of the respective RooFit class. It then helps save it to the variable that represents the result of this class (the squashed code/ C++ function that will be created).addToGlobalScope() helps declare and initialize the results variable, so that it can be available globally (throughout the function body). For local variables, the addToCodeBody() function can be used to keep the variables in the respective scope (for example, within a loop).beginLoop() helps build the start and the end of a For loop for your class. Simply place this function in the scope and place the contents of the For loop below this statement. The code squashing task will automatically build a loop around the statements that follow it. There's no need to worry about the index of these loops, because they get propagated. For example, if you want to iterate over a vector of RooFit objects using a loop, you don't have to think about indexing them properly because the beginLoop() function takes care of that. Simply call this function, place your function call in a scope and after the scope ends, the loop will also end.addToCodeBody() helps add things to the body of the C++ function that you're creating. It takes whatever string is computed in its arguments and adds it to the overall function string (which will later be just-in-time compiled). The addToCodeBody() function is important since not everything can be added in-line and this function helps split the code into multiple lines.Besides the
evaluate()function, this tutorial illustrates how the
analyticalIntegral() can be updated. This highly dependent on the class that is being transformed for AD support, but will be necessary in those specific instances.
Let's consider a fictional class RooFoo, that performs some arbitrary mathematical operations called 'Foo' (as seen in doFoo() function below).
Note that doFoo is a simplified example, in many cases the mathematical
operations are not limited to a single function, so they need to be spotted within the evaluate() function.
Note how the evaluate() function overrides the RooAbsReal for the RooFoo class. Similarly, the analyticalIntegral() function has also been overridden from the RooAbsReal class.
The evaluate() function includes some bookkeeping steps (commented out in above example) that are not relevant to AD. The important part is that it calls a specific function (doFoo() in this example), and returns the results.
Similarly, the analyticalIntegral() function calls a specific function ( integralFoo() in this example), and returns the results. It may also include some code that may need to be looked at, but for simplicity, its contents are commented out in this example.
Before creating the translate() function, move the mathematical logic ( doFoo() function in this example) out of the source class (RooFoo in this example) and into a separate header file called EvaluateFuncs.h. Also note that the parameters a and b have been defined as inputs, instead of them just being class members.
Directory path: roofit/roofitcore/inc/RooFit/Detail/EvaluateFuncs.h
So now that the doFoo() function exists in the EvaluateFuncs namespace, we need to comment out its original function definition in the RooFoo class and also add the namespace EvaluateFuncs to wherever doFoo() it is referenced (and also define input parameters for it).
Next, create the translate function. Most translate functions include a buildCall() function, that includes the fully qualified name (including 'EvaluateFuncs') of the function to be called along with the input parameters as they appear in the function (a,b in the following example).
Also, each translate() function requires the addResult() function. It will add whatever is represented on the right hand side to the result (saved in the res variable in the following example) of this class, which can then be propagated in the rest of the compute graph.
Besides creating the translate() function, the buildCallToAnalyticIntegral() function also needs to be added when analyticalIntegral() is found in your class. Depending on the code, you can call one or more integral functions using the code parameter. Our RooFoo example above only contains one integral function (integralFoo()).
Similar to doFoo(), comment out ‘integralFoo()’ in the original file and move it to a separate file called 'AnalyticalIntegrals.h' (this is not the same file as the EvaluateFuncs.h, but it works in a similar manner).
As with doFoo(). add the relevant inputs (a,b) as parameters, instead of just class members.
Directory path: hist/hist/src/AnalyticalIntegrals.h
Next, in the original RooFoo class, update all references to the integralFoo() function with its new fully qualified path ( EvaluateFunc::integralFoo) and include the input parameters as well ( EvaluateFunc::integralFoo(a, b)).
Next, in the buildCallToAnalyticIntegral() function, simply return the output using the buildCall() function.
buildCallToAnalyticIntegral() function is quiet similar to the translate() function, except that in translate(), you have to add to the result (using addResult()), while for buildCallToAnalyticIntegral(), you only have to return the string (using buildCall()).Consolidated Code changes in RooFoo example
Final RooFoo code:
Mathematical code moved to EvaluateFuncs.h file.
Integrals moved to the 'AnalyticalIntegrals.h' file.
Remember, as long as your code is supported by Clad (e.g., meaning there are custom derivatives defined for all external Math library functions used in
your code), it should work for AD support efforts. Please view Clad documentation for more details.
Both scenarios are possible:
1 - where analyticalIntegral() will be able to support AD
2 - where analyticalIntegral() will not be able to support AD
This requires further research.
In some cases. the evaluate() function is written in a piece-wise format (multiple evaluations based on multiple chunks of code). You can review the EvaluateFuncs.h file to find AD support for several piece-wise (if code==1 {...} else if code==2 {...} ) code snippets.
However, there may still be some cases where AD support may not be possible due to the way that evaluate() function works in that instance.
For simple caching, the caching logic can be separated from the mathematical code that is being moved to EvaluateFuncs.h, so that it can retained in the original file.
For more complicated scenarios, the code variable can be used to identify use cases (parts of the mathematical code in evaluate()) that should be supported, while other parts that are explicitly not be supported (e.g., using if code==1 {...} else if code==2 {...}).
So far, no. This needs further exploration. Hint: classes using Numerical Integration can be identified with the absence of the analyticalIntegral() function.
If you call in to an external Math library, and you use a function that has a customized variant with an already defined custom derivative, then you may see a warning like "falling back to Numeric Differentiation". In most such cases, your derivative should still work, since Numeric Differentiation is already well-tested in Clad.
To handle this, either define a custom derivative for that external function, or find a way to expose it to Clad.
An example of this can be seen with gamma_cdf() in AnalyticalIntegrals.h, for which the custom derivative is not supported, but in this specific instance, it falls back to Numeric Differentiation and works fine, since gamma_cdf()` doesn't have a lot of parameters.
In such cases, Numeric Differentiation fallback is only used for that
specific function. In above example, gamma_cdf() falls back to Numeric Differentiation but other functions in AnalyticalIntegrals.h will still be able to use AD. This is because Clad is going to assume that you have a derivative for this gamma_cdf() function, and the remaining functions will use AD as expected. In the end, the remaining functions (including gamma_cdf()) will try to fall back to Numeric Differentiation.
However, if you want to add pure AD support, you need to make sure that all your external functions are supported by Clad (meaning there is a custom derivative defined for each of them).
Please look at the test classes that test the derivatives, evaluates, fixtures, etc. (defined in 'roofit/roofitcore/test'). You can clone and adapt these tests to your class as needed. For example:
Tip: Tests like above can be referenced to see which parts of RooFit already support AD.
This is an area of research that still needs some work. In most cases, the compile times are reasonable, but with an increase in the level of complexity, higher compile times may be encountered.
Following classes provide several Helper Functions to translate existing logic into AD-supported logic.
a - RooFit::Detail::CodeSquashContext
b - RooFuncWrapper
roofit/roofitcore/inc/RooFit/Detail/RooFit::Detail::CodeSquashContext.h
It handles how to create a C++ function out of the compute graph (which is created with different RooFit classes). This C++ function will be independent of these RooFit classes.
RooFit::Detail::CodeSquashContext helps traverse the compute graph received from RooFit and then it translates that into a single piece of code (a C++ function), that can then be differentiated using Clad. It also helps evaluate the model.
In RooFit, evaluation is done using the 'evaluate()' function. It also performs a lot of book-keeping, caching, etc. that is required for RooFit (but not necessarily for AD).
A new translate() function is added to RooFit classes that includes a call to this evaluate() function. translate() helps implement the Code Squashing logic. All RooFit classes that should support AD need to use this function. It creates a string of code, which is then just-in-time compiled using Cling (C++ interpreter for ROOT). For each of the translate() functions, it is important to call addResult() since this is what enables the squashing to happen.
beginloop() and endloop() are used to create a scope for iterating over vector observables (collections of data). This is especially useful when dealing with data that comes in sets or arrays.For each
translate()function, it is important to calladdResult()since
this is what enables the squashing to happen.
translate().These functions will appear again in this document with more contextual examples. For detailed in-line documentation (code comments), please see:
roofit/roofitcore/src/RooFit/Detail/RooFit::Detail::CodeSquashContext.cxx
This class wraps the generated C++ code in a RooFit object, so that it can be used like other RooFit objects.
It takes a function body as input and creates a callable function from it. This allows users to evaluate the function and its derivatives efficiently.
These functions will appear again in this document with more contextual examples. For detailed in-line documentation (code comments), please see:
addResult()must be included in translate() function.key (the name of the node to add the result for), value (the new name to assign/overwrite).addResult() function).key (the node to get the result string for).klass (the class requesting this addition, usually 'this'), in (string to add to the squashed code).str (the string to add to the global scope).returnExpr (he string representation of what the squashed function should return, usually the head node).in (a pointer to the calling class, used to determine the loop dependent variables).in (the list to convert to array).funcname, passing some arguments.in (the input string).head (starting mathematical expression).funcName (name of the function being differentiated), funcBody (actual mathematical formula or equation).out (array where the computed gradient values will be stored).out array with gradient values.funcName (function name), head (structure of the function), paramSet (input function's parameters), data (optional data points).[^1]: For more details, please see this paper on [Automatic Differentiation of Binned Likelihoods With Roofit and Clad]
[^2]: For more details, please see this paper on GPU Accelerated Automatic Differentiation With Clad
