def getAttributeFloatAst(call_node: ast.Call):
r'''
Return an attribute on one of the xAOD objects.
'''
# Get the name of the moment out
if len(call_node.args) != 1:
raise Exception(
"Calling getAttributeFloat - only one argument is allowed")
if not isinstance(call_node.args[0], ast.Str):
raise Exception(
"Calling getAttributeFloat - only acceptable argument is a string")
r = cpp_ast.CPPCodeValue()
# We don't need include files for this - just quick access
r.args = [
'moment_name',
]
r.replacement_instance_obj = ('obj_j', call_node.func.value.id
) # type: ignore
r.running_code += [
'float result = obj_j->getAttribute<float>(moment_name);'
]
r.result = 'result'
r.result_rep = lambda sc: crep.cpp_variable(
unique_name("jet_attrib"), scope=sc, cpp_type=ctyp.terminal('float'))
# Replace it as the function that is going to get called.
call_node.func = r # type: ignore
return call_node
def DeltaRAst(call_node):
r'''
User is trying to call DeltaR (eta1, phi1, eta2, phi2). We turn this into a call
into a call into ROOT that does the phi subtraction (I can never get that crap right).
'''
if len(call_node.args) != 4:
raise ValueError(
"Calling DeltaR(eta1, phi1, eta2, phi2) has incorrect number of arguments"
)
# Create an AST to hold onto all of this.
r = cpp_ast.CPPCodeValue()
# We need TVector2 included here
r.include_files += ['TVector2.h', 'math.h']
# We need all four arguments pushed through.
r.args = ['eta1', 'phi1', 'eta2', 'phi2']
# The code is three steps
r.running_code += ['auto d_eta = eta1 - eta2;']
r.running_code += ['auto d_phi = TVector2::Phi_mpi_pi(phi1-phi2);']
r.running_code += ['auto result = sqrt(d_eta*d_eta + d_phi*d_phi);']
r.result = 'result'
r.result_rep = lambda scope: cpp_variable(
unique_name('delta_r'), scope=scope, cpp_type='double')
call_node.func = r
return call_node
def visit_BoolOp(self, node):
'''A bool op like And or Or on a set of values
This is a bit more complex than just "anding" things as we want to make sure to short-circuit the
evaluation if we need to.
'''
# The result of this test
result = crep.cpp_variable(unique_name('bool_op'), self._gc.current_scope(), cpp_type='bool')
self._gc.declare_variable(result)
# How we check and short-circuit depends on if we are doing and or or.
check_expr = result.as_cpp() if type(node.op) == ast.And else '!{0}'.format(result.as_cpp())
check = crep.cpp_value(check_expr, self._gc.current_scope(), cpp_type=ctyp.terminal('bool'))
first = True
scope = self._gc.current_scope()
for v in node.values:
if not first:
self._gc.add_statement(statement.iftest(check))
rep_v = self.get_rep(v)
self._gc.add_statement(statement.set_var(result, rep_v))
if not first:
self._gc.set_scope(scope)
first = False
# Cache result variable so those above us have something to use.
self._result = result
node.rep = result
def getCollection(info, call_node):
r'''
Return a cpp ast for accessing the jet collection
'''
# Get the name jet collection to look at.
if len(call_node.args) != 1:
raise ValueError(f"Calling {info['function_name']} - only one argument is allowed")
if not isinstance(call_node.args[0], ast.Str):
raise ValueError(f"Calling {info['function_name']} - only acceptable argument is a string")
# Fill in the CPP block next.
r = cpp_ast.CPPCodeValue()
r.args = ['collection_name', ]
r.include_files += info['include_files']
r.running_code += ['{0} result = 0;'.format(info['container_type']),
'ANA_CHECK (evtStore()->retrieve(result, collection_name));']
r.result = 'result'
is_collection = info['is_collection'] if 'is_collection' in info else True
if is_collection:
r.result_rep = lambda scope: crep.cpp_collection(unique_name(info['function_name'].lower()), scope=scope, collection_type=info['container_type']) # type: ignore
else:
r.result_rep = lambda scope: crep.cpp_variable(unique_name(info['function_name'].lower()), scope=scope, cpp_type=info['container_type'])
# Replace it as the function that is going to get called.
call_node.func = r
return call_node
def visit_IfExp(self, node):
r'''
We'd like to be able to use the "?" operator in C++, but the
problem is lazy evaluation. It could be when we look at one or the
other item, a bunch of prep work has to be done - and that will
show up in separate statements. So we have to use if/then/else with
a result value.
'''
# The result we'll store everything in.
result = crep.cpp_variable(unique_name("if_else_result"), self._gc.current_scope(), cpp_type=ctyp.terminal("double"))
self._gc.declare_variable(result)
# We always have to evaluate the test.
current_scope = self._gc.current_scope()
test_expr = self.get_rep(node.test)
self._gc.add_statement(statement.iftest(test_expr))
if_scope = self._gc.current_scope()
# Next, we do the true and false if statement.
self._gc.add_statement(statement.set_var(result, self.get_rep(node.body)))
self._gc.set_scope(if_scope)
self._gc.pop_scope()
self._gc.add_statement(statement.elsephrase())
self._gc.add_statement(statement.set_var(result, self.get_rep(node.orelse)))
self._gc.set_scope(current_scope)
# Done, the result is the rep of this node!
node.rep = result
self._result = result
def _create_accumulator(self, seq: crep.cpp_sequence, acc_type: ctyp.terminal, initial_value=None):
'Helper to create an accumulator for the Aggregate function'
accumulator_type = acc_type
# When we implement other types of aggregate, this code will need to
# be back in.
# if accumulator_type is None:
# sv = seq.sequence_value()
# if not isinstance(sv, crep.cpp_value):
# raise Exception("Do not know how to accumulate a sequence!")
# accumulator_type = sv.cpp_type()
if not check_accumulator_type(accumulator_type):
raise ValueError(f"Aggregate over a sequence of type '{str(accumulator_type)}' is not supported.")
# Getting the scope level right is tricky. If this is a straight sequence of items, then we want the sequence level.
# But if this is a sequence of sequences, we are aggregating over the sequence itself. So we need to do it one level
# up from where the iterator is running on the interior sequence.
seq_val = seq.sequence_value()
if isinstance(seq_val, crep.cpp_sequence):
accumulator_scope = seq_val.iterator_value().scope()[-1]
else:
accumulator_scope = seq.iterator_value().scope()[-1]
accumulator = crep.cpp_variable(unique_name("aggResult"),
accumulator_scope,
accumulator_type,
initial_value=initial_value if initial_value is not None else crep.cpp_value(accumulator_type.default_value(), self._gc.current_scope(), accumulator_type))
accumulator_scope.declare_variable(accumulator)
return accumulator, accumulator_scope
def test_variable_type_update():
tc = gc_scope_top_level()
expr = "a"
ctype = ctyp.terminal('int', False)
v = crep.cpp_variable(expr, tc, ctype)
v.update_type(ctyp.terminal('float', False))
assert v.cpp_type().type == 'float'
def write_cpp_files(self, ast: ast.AST,
output_path: Path) -> xAODExecutionInfo:
r"""
Given the AST generate the C++ files that need to run. Return them along with
the input files.
"""
# Find the base file dataset and mark it.
from func_adl import find_EventDataset
file = find_EventDataset(ast)
iterator = crep.cpp_variable("bogus-do-not-use",
top_level_scope(),
cpp_type=None)
file.rep = crep.cpp_sequence(iterator, iterator,
top_level_scope()) # type: ignore
# Visit the AST to generate the code structure and find out what the
# result is going to be.
qv = query_ast_visitor()
result_rep = qv.get_rep(ast) if _is_format_request(ast) \
else qv.get_as_ROOT(ast)
# Emit the C++ code into our dictionaries to be used in template generation below.
query_code = _cpp_source_emitter()
qv.emit_query(query_code)
book_code = _cpp_source_emitter()
qv.emit_book(book_code)
class_dec_code = qv.class_declaration_code()
includes = qv.include_files()
# The replacement dict to pass to the template generator can now be filled
info = {}
info['query_code'] = query_code.lines_of_query_code()
info['book_code'] = book_code.lines_of_query_code()
info['class_dec'] = class_dec_code
info['include_files'] = includes
# We use jinja2 templates. Write out everything.
template_dir = _find_dir("func_adl_xAOD/R21Code")
j2_env = jinja2.Environment(
loader=jinja2.FileSystemLoader(template_dir))
self._copy_template_file(j2_env, info, 'ATestRun_eljob.py',
output_path)
self._copy_template_file(j2_env, info, 'package_CMakeLists.txt',
output_path)
self._copy_template_file(j2_env, info, 'query.cxx', output_path)
self._copy_template_file(j2_env, info, 'query.h', output_path)
self._copy_template_file(j2_env, info, 'runner.sh', output_path)
(output_path / 'runner.sh').chmod(0o755)
# Build the return object.
return xAODExecutionInfo(result_rep, output_path, 'runner.sh', [
'ATestRun_eljob.py', 'package_CMakeLists.txt', 'query.cxx',
'query.h', 'runner.sh'
])
def test_variable_type__with_initial_update():
tc = gc_scope_top_level()
expr = "a"
c_type = ctyp.terminal('int', False)
c_init = crep.cpp_value('0.0', tc, ctyp.terminal('int', False))
v = crep.cpp_variable(expr, tc, c_type, c_init)
v.update_type(ctyp.terminal('float', False))
assert v.cpp_type().type == 'float'
assert v.initial_value().cpp_type().type == 'float'
def fill_collection_levels(seq: crep.cpp_sequence, accumulator: crep.cpp_value):
inner = seq.sequence_value()
scope = seq.scope()
if isinstance(inner, crep.cpp_sequence):
scope = seq.iterator_value().scope()
storage = crep.cpp_variable(unique_name('ntuple'), scope, cpp_type=inner.cpp_type())
assert not isinstance(scope, gc_scope_top_level)
scope.declare_variable(storage)
fill_collection_levels(inner, storage)
inner = storage
set_scope(scope, scope_fill)
self._gc.add_statement(statement.push_back(accumulator, inner))
async def exe_from_qastle(q: str):
'Dummy executor that will return the ast properly rendered. If qastle_roundtrip is true, then we will round trip the ast via qastle first.'
# Round trip qastle if requested.
import qastle
a = qastle.text_ast_to_python_ast(q).body[0].value
# Setup the rep for this filter
from func_adl import find_EventDataset
file = find_EventDataset(a)
iterator = cpp_variable("bogus-do-not-use",
top_level_scope(),
cpp_type=None)
file.rep = cpp_sequence(iterator, iterator,
top_level_scope()) # type: ignore
# Use the dummy executor to process this, and return it.
exe = dummy_executor()
exe.evaluate(a)
return exe
def call_First(self, node: ast.AST, args: List[ast.AST]) -> Any:
'We are in a sequence. Take the first element of the sequence and use that for future things.'
# Unpack the source here
assert len(args) == 1
source = args[0]
# Make sure we are in a loop.
seq = self.as_sequence(source)
# The First terminal works by protecting the code with a if (first_time) {} block.
# We need to declare the first_time variable outside the block where the thing we are
# looping over here is defined. This is a little tricky, so we delegate to another method.
loop_scope = seq.iterator_value().scope()
outside_block_scope = loop_scope[-1]
# Define the variable to track this outside that block.
is_first = crep.cpp_variable(unique_name('is_first'),
outside_block_scope,
cpp_type=ctyp.terminal('bool'),
initial_value=crep.cpp_value('true', self._gc.current_scope(), ctyp.terminal('bool')))
outside_block_scope.declare_variable(is_first)
# Now, as long as is_first is true, we can execute things inside this statement.
# The trick is putting the if statement in the right place. We need to locate it just one level
# below where we defined the scope above.
s = statement.iftest(is_first)
s.add_statement(statement.set_var(is_first, crep.cpp_value('false', top_level_scope(), cpp_type=ctyp.terminal('bool'))))
sv = seq.sequence_value()
if isinstance(sv, crep.cpp_sequence):
self._gc.set_scope(sv.iterator_value().scope()[-1])
else:
self._gc.set_scope(sv.scope())
self._gc.add_statement(s)
# If we just found the first sequence in a sequence, return that.
# Otherwise return a new version of the value.
first_value = sv if isinstance(sv, crep.cpp_sequence) else sv.copy_with_new_scope(self._gc.current_scope())
node.rep = first_value # type: ignore
self._result = first_value
async def execute_result_async(self, a: ast.AST) -> Any:
'Dummy executor that will return the ast properly rendered. If qastle_roundtrip is true, then we will round trip the ast via qastle first.'
# Round trip qastle if requested.
if self._q_roundtrip:
import qastle
print(f'before: {ast.dump(a)}')
a_text = qastle.python_ast_to_text_ast(a)
a = qastle.text_ast_to_python_ast(a_text).body[0].value
print(f'after: {ast.dump(a)}')
# Setup the rep for this dataset
from func_adl import find_EventDataset
file = find_EventDataset(a)
iterator = cpp_variable("bogus-do-not-use",
top_level_scope(),
cpp_type=None)
file.rep = cpp_sequence(iterator, iterator,
top_level_scope()) # type: ignore
# Use the dummy executor to process this, and return it.
exe = dummy_executor()
rnr = atlas_xaod_executor()
exe.evaluate(a)
return exe
def call_ResultTTree(self, node: ast.Call, args: List[ast.AST]):
'''This AST means we are taking an iterable and converting it to a ROOT file.
'''
# Unpack the variables.
assert len(args) == 4
source = args[0]
column_names = _extract_column_names(args[1])
tree_name = ast.literal_eval(args[2])
assert isinstance(tree_name, str)
# root_filename = args[3]
# Get the representations for each variable. We expect some sort of structure
# for the variables - or perhaps a single variable.
self.generic_visit(source)
v_rep_not_norm = self.as_sequence(source)
# What we have is a sequence of the data values we want to fill. The iterator at play
# here is the scope we want to use to run our Fill() calls to the TTree.
scope_fill = v_rep_not_norm.iterator_value().scope()
# Clean the data up so it is uniform and the next bit can proceed smoothly.
# If we don't have a tuple of data to log, turn it into a tuple.
seq_values = v_rep_not_norm.sequence_value()
if not isinstance(seq_values, crep.cpp_tuple):
seq_values = crep.cpp_tuple((v_rep_not_norm.sequence_value(),), scope_fill)
# Make sure the number of items is the same as the number of columns specified.
if len(seq_values.values()) != len(column_names):
raise Exception("Number of columns ({0}) is not the same as labels ({1}) in TTree creation".format(len(seq_values.values()), len(column_names)))
# Next, look at each on in turn to decide if it is a vector or a simple variable.
# Create a variable that we will fill for each one.
var_names = [(name, crep.cpp_variable(unique_name(name, is_class_var=True), self._gc.current_scope(), cpp_type=get_ttree_type(rep)))
for name, rep in zip(column_names, seq_values.values())]
# For each incoming variable, we need to declare something we are going to write.
for cv in var_names:
self._gc.declare_class_variable(cv[1])
# Next, emit the booking code
self._gc.add_book_statement(statement.book_ttree(tree_name, var_names))
# Note that the output file and tree are what we are going to return.
# The output filename is fixed - the hose code in AnalysisBase has that hard coded.
# To allow it to be different we have to modify that template too, and pass the
# information there. If more than one tree is written, the current code would
# lead to a bug.
node.rep = rh.cpp_ttree_rep("ANALYSIS.root", tree_name, self._gc.current_scope()) # type: ignore
# For each varable we need to save, cache it or push it back, depending.
# Make sure that it happens at the proper scope, where what we are after is defined!
s_orig = self._gc.current_scope()
for e_rep, e_name in zip(seq_values.values(), var_names):
scope_fill = self.code_fill_ttree(e_rep, e_name[1], scope_fill)
# The fill statement. This should happen at the scope where the tuple was defined.
# The scope where this should be done is a bit tricky (note the update above):
# - If a sequence, you want it where the sequence iterator is defined - or outside that scope
# - If a value, you want it at the level where the value is set.
self._gc.set_scope(scope_fill)
self._gc.add_statement(statement.ttree_fill(tree_name))
for e in zip(seq_values.values(), var_names):
if rep_is_collection(e[0]):
self._gc.add_statement(statement.container_clear(e[1][1]))
# And we are a terminal, so pop off the block.
self._gc.set_scope(s_orig)
self._gc.pop_scope()
return node.rep # type: ignore