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 BaseException( "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 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 BaseException( "Calling getMomentFloat - only one argument is allowed") if not isinstance(call_node.args[0], ast.Str): raise BaseException( "Calling getMomentFloat - 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) 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 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='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 BaseException("Calling {0} - only one argument is allowed".format(info['function_name'])) if not isinstance(call_node.args[0], ast.Str): raise BaseException("Calling {0} - only acceptable argument is a string".format(info['function_name'])) # 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']) 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 visit_Call_Aggregate_only(self, node: ast.Call): ''' - (acc lambda): the accumulator is set to the first element, and the lambda is called to update it after that. This is called `agg_only`. ''' agg_lambda = node.args[0] # Get the sequence we are calling against and the accumulator if not isinstance(node.func, ast.Attribute): raise BaseException("Wrong type of function") seq = self.as_sequence(node.func.value) accumulator, accumulator_scope = self.create_accumulator(seq) # We have to do a simple if statement here so that the first time through we can set the # accumulator, and the second time we can add to it. is_first_iter = crep.cpp_variable(unique_name("is_first"), self._gc.current_scope(), cpp_type=ctyp.terminal('bool'), initial_value=crep.cpp_value( 'true', self._gc.current_scope(), ctyp.terminal('bool'))) accumulator_scope.declare_variable(is_first_iter) # Set the scope where we will be doing the accumulation 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()) # Code up if statement to select out the first element. if_first = statement.iftest(is_first_iter) self._gc.add_statement(if_first) self._gc.add_statement( statement.set_var( is_first_iter, crep.cpp_value("false", self._gc.current_scope(), ctyp.terminal('bool')))) # Set the accumulator self._gc.add_statement( statement.set_var(accumulator, seq.sequence_value())) self._gc.pop_scope() # Now do the if statement and make the call to calculate the accumulation. self._gc.add_statement(statement.elsephrase()) call = ast.Call( func=agg_lambda, args=[accumulator.as_ast(), seq.sequence_value().as_ast()]) self._gc.add_statement( statement.set_var(accumulator, self.get_rep(call))) # Finally, since this is a terminal, we need to pop off the top. self._gc.set_scope(accumulator_scope) # Cache the results in our result in case we are skipping nodes in the AST. node.rep = accumulator self._result = accumulator
def create_accumulator(self, seq: crep.cpp_sequence, initial_value=None, acc_type=None): 'Helper to create an accumulator for the Aggregate function' accumulator_type = acc_type if accumulator_type is None: sv = seq.sequence_value() if not isinstance(sv, crep.cpp_value): raise BaseException( "Do not know how to accumulate a sequence!") accumulator_type = sv.cpp_type() if not check_accumulator_type(accumulator_type): raise BaseException( "Aggregate over a sequence of type '{0}' is not supported.". format(str(accumulator_type))) # 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 write_cpp_files(self, ast: ast.AST, output_path: str) -> 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.xAOD.backend.util_LINQ import find_dataset file = find_dataset(ast) iterator = crep.cpp_variable("bogus-do-not-use", top_level_scope(), cpp_type=None) file.rep = crep.cpp_sequence(iterator, iterator) # 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) # 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/backend/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) os.chmod(os.path.join(str(output_path), 'runner.sh'), 0o755) # Build the return object. return xAODExecutionInfo( file.url, result_rep, output_path, 'runner.sh', [ 'ATestRun_eljob.py', 'package_CMakeLists.txt', 'query.cxx', 'query.h', 'runner.sh' ])
async def exe_for_test(a: ast.AST): 'Dummy executor that will return the ast properly rendered' # Setup the rep for this filter file = find_dataset(a) iterator = cpp_variable("bogus-do-not-use", top_level_scope(), cpp_type=None) file.rep = cpp_sequence(iterator, iterator) # Use the dummy executor to process this, and return it. exe = dummy_executor() rnr = atlas_xaod_executor() exe.evaluate(rnr.apply_ast_transformations(a)) return exe
def visit_First(self, node): 'We are in a sequence. Take the first element of the sequence and use that for future things.' # Make sure we are in a loop. seq = self.as_sequence(node.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 self._result = first_value
def visit_ResultTTree(self, node: ResultTTree): '''This AST means we are taking an iterable and converting it to a ROOT file. ''' # Get the representations for each variable. We expect some sort of structure # for the variables - or perhaps a single variable. self.generic_visit(node) v_rep_not_norm = self.as_sequence(node.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(node.column_names): raise BaseException( "Number of columns ({0}) is not the same as labels ({1}) in TTree creation" .format(len(seq_values.values()), len(node.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(node.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 tree_name = unique_name(node.tree_name) 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()) # 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): # Set the scope. Normally we want to do it where the variable was calculated # (think of cases when you have to calculate something with a `push_back`), # but if the variable was already calculated, we want to make sure we are at least # in the same scope as the tree fill. e_rep_scope = e_rep.scope() if not isinstance( e_rep, crep.cpp_sequence) else e_rep.sequence_value().scope() if e_rep_scope.starts_with(scope_fill): self._gc.set_scope(e_rep_scope) else: self._gc.set_scope(scope_fill) # If the variable is something we are iterating over, then fill it, otherwise, # just set it. if rep_is_collection(e_rep): self._gc.add_statement( statement.push_back(e_name[1], e_rep.sequence_value())) else: self._gc.add_statement(statement.set_var(e_name[1], e_rep)) cs = self._gc.current_scope() if cs.starts_with(scope_fill): scope_fill = cs # 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()