Logo ROOT  
Reference Guide
 
Loading...
Searching...
No Matches
CodeSquashContext.h
Go to the documentation of this file.
1/*
2 * Project: RooFit
3 * Authors:
4 * Garima Singh, CERN 2023
5 * Jonas Rembser, CERN 2023
6 *
7 * Copyright (c) 2023, CERN
8 *
9 * Redistribution and use in source and binary forms,
10 * with or without modification, are permitted according to the terms
11 * listed in LICENSE (http://roofit.sourceforge.net/license.txt)
12 */
13
14#ifndef RooFit_Detail_CodeSquashContext_h
15#define RooFit_Detail_CodeSquashContext_h
16
17#include <RooAbsCollection.h>
19#include <RooNumber.h>
20
21#include <cstddef>
22#include <map>
23#include <sstream>
24#include <string>
25#include <type_traits>
26#include <unordered_map>
27
28template <class T>
30
31namespace RooFit {
32
33namespace Detail {
34
35/// @brief A class to maintain the context for squashing of RooFit models into code.
37public:
38 CodeSquashContext(std::map<RooFit::Detail::DataKey, std::size_t> const &outputSizes) : _nodeOutputSizes(outputSizes)
39 {
40 }
41
42 void addResult(RooAbsArg const *key, std::string const &value);
43 void addResult(const char *key, std::string const &value);
44
45 std::string const &getResult(RooAbsArg const &arg);
46
47 template <class T>
48 std::string const &getResult(RooTemplateProxy<T> const &key)
49 {
50 return getResult(key.arg());
51 }
52
53 /// @brief Figure out the output size of a node. It is the size of the
54 /// vector observable that it depends on, or 1 if it doesn't depend on any
55 /// or is a reducer node.
56 /// @param key The node to look up the size for.
57 std::size_t outputSize(RooFit::Detail::DataKey key) const
58 {
59 auto found = _nodeOutputSizes.find(key);
60 if (found != _nodeOutputSizes.end())
61 return found->second;
62 return 1;
63 }
64
65 void addToGlobalScope(std::string const &str);
66 std::string assembleCode(std::string const &returnExpr);
67 void addVecObs(const char *key, int idx);
68
69 void addToCodeBody(RooAbsArg const *klass, std::string const &in);
70
71 void addToCodeBody(std::string const &in, bool isScopeIndep = false);
72
73 /// @brief Build the code to call the function with name `funcname`, passing some arguments.
74 /// The arguments can either be doubles or some RooFit arguments whose
75 /// results will be looked up in the context.
76 template <typename... Args_t>
77 std::string buildCall(std::string const &funcname, Args_t const &...args)
78 {
79 std::stringstream ss;
80 ss << funcname << "(" << buildArgs(args...) << ")";
81 return ss.str();
82 }
83
84 /// @brief A class to manage loop scopes using the RAII technique. To wrap your code around a loop,
85 /// simply place it between a brace inclosed scope with a call to beginLoop at the top. For e.g.
86 /// {
87 /// auto scope = ctx.beginLoop({<-set of vector observables to loop over->});
88 /// // your loop body code goes here.
89 /// }
90 class LoopScope {
91 public:
92 LoopScope(CodeSquashContext &ctx, std::vector<TNamed const *> &&vars) : _ctx{ctx}, _vars{vars} {}
93 ~LoopScope() { _ctx.endLoop(*this); }
94
95 std::vector<TNamed const *> const &vars() const { return _vars; }
96
97 private:
99 const std::vector<TNamed const *> _vars;
100 };
101
102 std::unique_ptr<LoopScope> beginLoop(RooAbsArg const *in);
103
104 std::string getTmpVarName();
105
106 std::string buildArg(RooAbsCollection const &x);
107 std::string buildArg(std::span<const double> arr);
108
109private:
110 bool isScopeIndependent(RooAbsArg const *in) const;
111
112 void endLoop(LoopScope const &scope);
113
114 void addResult(TNamed const *key, std::string const &value);
115
116 template <class T, typename std::enable_if<std::is_floating_point<T>{}, bool>::type = true>
117 std::string buildArg(T x)
118 {
119 return RooNumber::toString(x);
120 }
121
122 // If input is integer, we want to print it into the code like one (i.e. avoid the unnecessary '.0000').
123 template <class T, typename std::enable_if<std::is_integral<T>{}, bool>::type = true>
124 std::string buildArg(T x)
125 {
126 return std::to_string(x);
127 }
128
129 std::string buildArg(std::string const &x) { return x; }
130
131 std::string buildArg(std::nullptr_t) { return "nullptr"; }
132
133 std::string buildArg(RooAbsArg const &arg) { return getResult(arg); }
134
135 template <class T>
136 std::string buildArg(RooTemplateProxy<T> const &arg)
137 {
138 return getResult(arg);
139 }
140
141 std::string buildArgs() { return ""; }
142
143 template <class Arg_t>
144 std::string buildArgs(Arg_t const &arg)
145 {
146 return buildArg(arg);
147 }
148
149 template <typename Arg_t, typename... Args_t>
150 std::string buildArgs(Arg_t const &arg, Args_t const &...args)
151 {
152 return buildArg(arg) + ", " + buildArgs(args...);
153 }
154
155 /// @brief Map of node names to their result strings.
156 std::unordered_map<const TNamed *, std::string> _nodeNames;
157 /// @brief Block of code that is placed before the rest of the function body.
158 std::string _globalScope;
159 /// @brief A map to keep track of the observable indices if they are non scalar.
160 std::unordered_map<const TNamed *, int> _vecObsIndices;
161 /// @brief Map of node output sizes.
162 std::map<RooFit::Detail::DataKey, std::size_t> _nodeOutputSizes;
163 /// @brief Stores the squashed code body.
164 std::string _code;
165 /// @brief The current number of for loops the started.
166 int _loopLevel = 0;
167 /// @brief Index to get unique names for temporary variables.
168 int _tmpVarIdx = 0;
169 /// @brief Keeps track of the position to go back and insert code to.
170 int _scopePtr = -1;
171 /// @brief Stores code that eventually gets injected into main code body.
172 /// Mainly used for placing decls outside of loops.
173 std::string _tempScope;
174 /// @brief A map to keep track of list names as assigned by addResult.
175 std::unordered_map<RooFit::UniqueId<RooAbsCollection>::Value_t, std::string> listNames;
176};
177
178} // namespace Detail
179
180} // namespace RooFit
181
182#endif
Option_t Option_t TPoint TPoint const char GetTextMagnitude GetFillStyle GetLineColor GetLineWidth GetMarkerStyle GetTextAlign GetTextColor GetTextSize void value
Common abstract base class for objects that represent a value and a "shape" in RooFit.
Definition RooAbsArg.h:77
Abstract container object that can hold multiple RooAbsArg objects.
A class to manage loop scopes using the RAII technique.
std::vector< TNamed const * > const & vars() const
LoopScope(CodeSquashContext &ctx, std::vector< TNamed const * > &&vars)
const std::vector< TNamed const * > _vars
A class to maintain the context for squashing of RooFit models into code.
std::string assembleCode(std::string const &returnExpr)
Assemble and return the final code with the return expression and global statements.
std::map< RooFit::Detail::DataKey, std::size_t > _nodeOutputSizes
Map of node output sizes.
std::string _tempScope
Stores code that eventually gets injected into main code body.
std::string buildCall(std::string const &funcname, Args_t const &...args)
Build the code to call the function with name funcname, passing some arguments.
void addResult(RooAbsArg const *key, std::string const &value)
A function to save an expression that includes/depends on the result of the input node.
std::string buildArg(std::string const &x)
void endLoop(LoopScope const &scope)
std::unordered_map< const TNamed *, int > _vecObsIndices
A map to keep track of the observable indices if they are non scalar.
int _loopLevel
The current number of for loops the started.
int _tmpVarIdx
Index to get unique names for temporary variables.
std::size_t outputSize(RooFit::Detail::DataKey key) const
Figure out the output size of a node.
std::unordered_map< const TNamed *, std::string > _nodeNames
Map of node names to their result strings.
void addToCodeBody(RooAbsArg const *klass, std::string const &in)
Adds the input string to the squashed code body.
void addVecObs(const char *key, int idx)
Since the squashed code represents all observables as a single flattened array, it is important to ke...
bool isScopeIndependent(RooAbsArg const *in) const
CodeSquashContext(std::map< RooFit::Detail::DataKey, std::size_t > const &outputSizes)
std::string const & getResult(RooTemplateProxy< T > const &key)
std::string const & getResult(RooAbsArg const &arg)
Gets the result for the given node using the node name.
std::string getTmpVarName()
Get a unique variable name to be used in the generated code.
std::string _code
Stores the squashed code body.
void addToGlobalScope(std::string const &str)
Adds the given string to the string block that will be emitted at the top of the squashed function.
std::unordered_map< RooFit::UniqueId< RooAbsCollection >::Value_t, std::string > listNames
A map to keep track of list names as assigned by addResult.
std::string buildArgs(Arg_t const &arg, Args_t const &...args)
std::string buildArg(RooAbsCollection const &x)
Function to save a RooListProxy as an array in the squashed code.
std::string buildArgs(Arg_t const &arg)
std::string buildArg(RooAbsArg const &arg)
int _scopePtr
Keeps track of the position to go back and insert code to.
std::string _globalScope
Block of code that is placed before the rest of the function body.
std::unique_ptr< LoopScope > beginLoop(RooAbsArg const *in)
Create a RAII scope for iterating over vector observables.
std::string buildArg(RooTemplateProxy< T > const &arg)
std::string buildArg(std::nullptr_t)
static std::string toString(double x)
Returns an std::to_string compatible number (i.e.
Definition RooNumber.cxx:31
const T & arg() const
Return reference to object held in proxy.
The TNamed class is the base class for all named ROOT classes.
Definition TNamed.h:29
Double_t x[n]
Definition legend1.C:17
The namespace RooFit contains mostly switches that change the behaviour of functions of PDFs (or othe...
Definition JSONIO.h:26