def test_first_after_where(): # Part of testing that First puts its outer settings in the right place. # This also tests First on a collection of objects that hasn't been pulled a part # in a select. dataset_for_testing() \ .Select('lambda e: e.Jets("AntiKt4EMTopoJets").Where(lambda j: j.pt() > 10).First().pt()') \ .value()
def test_first_object_in_each_event(): # Part of testing that First puts its outer settings in the right place. # This also tests First on a collection of objects that hasn't been pulled a part # in a select. dataset_for_testing() \ .Select('lambda e: e.Jets("AntiKt4EMTopoJets").First().pt()/1000.0') \ .value()
def test_get_attribute(): with pytest.raises(Exception) as e: dataset_for_testing() \ .SelectMany('lambda e: e.Jets("AntiKt4EMTopoJets").Select(lambda j: j.getAttribute("emf"))') \ .value() assert "getAttributeFloat" in str(e.value)
def test_get_attribute_float_wrong_arg_type(): with pytest.raises(Exception) as e: dataset_for_testing() \ .SelectMany('lambda e: e.Jets("AntiKt4EMTopoJets").Select(lambda j: j.getAttributeFloat(1))') \ .value() assert 'getAttribute' in str(e.value)
def test_get_attribute_vector_float_wrong_args(): with pytest.raises(Exception) as e: dataset_for_testing() \ .SelectMany('lambda e: e.Jets("AntiKt4EMTopoJets")') \ .SelectMany('lambda j: j.getAttributeVectorFloat("emf", 22)') \ .value() assert 'getAttributeVectorFloat' in str(e.value)
def test_Select_of_2D_array_with_tuple(): # We do not support structured output - so array or array(array), but not array(array, array), # at least not yet. Make sure error is reasonable. with pytest.raises(Exception) as e: dataset_for_testing() \ .Select('lambda e: e.Jets("AntiKt4EMTopoJets").Select(lambda j: (j.pt()/1000.0, j.eta()))') \ .value() assert "data structures" in str(e.value)
def test_cant_call_double(): msg = "" try: dataset_for_testing("file://root.root") \ .Select("lambda e: e.Jets('AntiKt4EMTopoJets').Select(lambda j: j.pt().eta()).Sum()") \ .value() except xAODTranslationError as e: msg = str(e) assert 'Unable to call method eta on type double' in msg
def test_generate_unary_not(): r = dataset_for_testing() \ .SelectMany('lambda e: e.Jets("AntiKt4EMTopoJets").Select(lambda j: not (j.pt() > 50.0))') \ .value() lines = get_lines_of_code(r) print_lines(lines) _ = find_line_with(f"!(", lines)
def test_per_event_item(): r=dataset_for_testing() \ .Select('lambda e: e.EventInfo("EventInfo").runNumber()') \ .value() vs = r.QueryVisitor._gc._class_vars assert 1 == len(vs) assert "double" == str(vs[0].cpp_type())
def test_dict_output_fail_expansion(): my_old_dict = {1: 'hi'} with pytest.raises(ValueError): r=dataset_for_testing() \ .Select(lambda e: e.EventInfo("EventInfo").runNumber()) \ .Select(lambda e: {'run_number': e, **my_old_dict}) \ .value()
def test_generate_unary_operations(): ops = ['+', '-'] for o in ops: r = dataset_for_testing() \ .SelectMany('lambda e: e.Jets("AntiKt4EMTopoJets").Select(lambda j: j.pt()+({0}1))'.format(o)) \ .value() lines = get_lines_of_code(r) print_lines(lines) _ = find_line_with(f"pt()+({o}(1))", lines)
def test_dict_output(): 'This is integration testing - making sure the dict to root conversion works' r=dataset_for_testing() \ .Select(lambda e: e.EventInfo("EventInfo").runNumber()) \ .Select(lambda e: {'run_number': e}) \ .value() vs = r.QueryVisitor._gc._class_vars assert 1 == len(vs) assert "double" == str(vs[0].cpp_type())
def test_builtin_sin_function_no_math_import(): # The following statement should be a straight sequence, not an array. r = dataset_for_testing() \ .SelectMany('lambda e: e.Jets("AntiKt4EMTopoJets").Select(lambda j: sin(j.pt()))') \ .value() # Check to see if there mention of push_back anywhere. lines = get_lines_of_code(r) print_lines(lines) l_abs = find_line_with("std::sin", lines) assert "->pt()" in lines[l_abs]
def test_generate_binary_operators(): # Make sure the binary operators work correctly - that they don't cause a crash in generation. ops = ['+', '-', '*', '/', '%'] for o in ops: r = dataset_for_testing() \ .SelectMany('lambda e: e.Jets("AntiKt4EMTopoJets").Select(lambda j: j.pt(){0}1)'.format(o)) \ .value() lines = get_lines_of_code(r) print_lines(lines) _ = find_line_with(f"pt(){o}1", lines)
def test_nested_lambda_argument_name_with_monad(): # Need both the monad and the "e" reused to get this error! r = dataset_for_testing() \ .Select('lambda e: (e.Electrons("Electrons"), e.Muons("Muons"))') \ .Select('lambda e: e[0].Select(lambda e: e.E())') \ .value() lines = get_lines_of_code(r) print_lines(lines) l_push = find_line_with('push_back', lines) assert "->E()" in lines[l_push]
def test_ifexpr(): r = dataset_for_testing(qastle_roundtrip=True) \ .SelectMany('lambda e: e.Jets("AntiKt4EMTopoJets").Select(lambda j: 1.0 if j.pt() > 10.0 else 2.0)') \ .value() # Make sure that a test around 10.0 occurs. lines = get_lines_of_code(r) print_lines(lines) lines = [l for l in lines if '10.0' in l] assert len(lines) == 1 assert 'if ' in lines[0]
def test_get_attribute_float(): # The following statement should be a straight sequence, not an array. r = dataset_for_testing() \ .SelectMany('lambda e: e.Jets("AntiKt4EMTopoJets").Select(lambda j: j.getAttributeFloat("emf"))') \ .value() # Check to see if there mention of push_back anywhere. lines = get_lines_of_code(r) print_lines(lines) l_attribute = find_line_with("getAttribute", lines) assert 'getAttribute<float>("emf")' in lines[l_attribute]
def test_First_Of_Select_After_Where_is_in_right_place(): # Make sure that we have the "First" predicate after if Where's if statement. r = dataset_for_testing() \ .Select('lambda e: e.Jets("AntiKt4EMTopoJets").Select(lambda j: j.pt()/1000.0).Where(lambda jpt: jpt > 10.0).First()') \ .value() lines = get_lines_of_code(r) print_lines(lines) l = find_line_with(">10.0", lines) # Look for the "false" that First uses to remember it has gone by one. assert find_line_with("false", lines[l:], throw_if_not_found=False) > 0
def test_electron_and_muon_with_list_qastle(): # See if we can re-create a bug we are seeing with # Marc's long query. r = dataset_for_testing(qastle_roundtrip=True) \ .Select('lambda e: [e.Electrons("Electrons"), e.Muons("Muons")]') \ .Select('lambda e: [e[0].Select(lambda ele: ele.E()), e[0].Select(lambda ele: ele.pt()), e[0].Select(lambda ele: ele.phi()), e[0].Select(lambda ele: ele.eta()), e[1].Select(lambda mu: mu.E()), e[1].Select(lambda mu: mu.pt()), e[1].Select(lambda mu: mu.phi()), e[1].Select(lambda mu: mu.eta())]') \ .value() lines = get_lines_of_code(r) print_lines(lines) assert find_line_with("->Fill()", lines) != 0
def test_SelectMany_of_tuple_is_not_array(): # The following statement should be a straight sequence, not an array. r = dataset_for_testing() \ .SelectMany('lambda e: e.Jets("AntiKt4EMTopoJets").Select(lambda j: (j.pt()/1000.0, j.eta()))') \ .value() lines = get_lines_of_code(r) print_lines(lines) assert 0 == ["push_back" in l for l in lines].count(True) l_push_back = find_line_with("Fill()", lines) active_blocks = find_open_blocks(lines[:l_push_back]) assert 1 == ["for" in a for a in active_blocks].count(True)
def test_result_awkward(): # The following statement should be a straight sequence, not an array. r = dataset_for_testing() \ .SelectMany('lambda e: e.Jets("AntiKt4EMTopoJets")') \ .Select("lambda j: j.pt()") \ .value() # Make sure that the tree Fill is at the same level as the _JetPts2 getting set. lines = get_lines_of_code(r) print_lines(lines) l_jetpt = find_line_with("_col1", lines) assert "Fill()" in lines[l_jetpt + 1]
def test_per_jet_item(): # The following statement should be a straight sequence, not an array. r = dataset_for_testing() \ .SelectMany('lambda e: e.Jets("AntiKt4EMTopoJets").Select(lambda j: j.pt())') \ .value() # Check to see if there mention of push_back anywhere. lines = get_lines_of_code(r) print_lines(lines) assert 0 == ["push_back" in l for l in lines].count(True) l_push_back = find_line_with("Fill()", lines) active_blocks = find_open_blocks(lines[:l_push_back]) assert 1 == ["for" in a for a in active_blocks].count(True)
def test_and_clause_in_where(): # The following statement should be a straight sequence, not an array. r = dataset_for_testing() \ .SelectMany('lambda e: e.Jets("AntiKt4EMTopoJets")') \ .Where("lambda j: j.pt()>40.0 and j.eta()<2.5") \ .Select("lambda j: j.pt()") \ .value() # Make sure that the tree Fill is at the same level as the _JetPts2 getting set. lines = get_lines_of_code(r) print_lines(lines) l_if = [l for l in lines if "if (" in l] assert len(l_if) == 2 assert l_if[0] == l_if[1]
def test_per_jet_with_delta(): # Trying to repro a bug we saw in the wild r = dataset_for_testing() \ .Select('lambda e: (e.Jets("AntiKt4EMTopoJets"),e.TruthParticles("TruthParticles").Where(lambda tp1: tp1.pdgId() == 35))') \ .SelectMany('lambda ev: ev[0].Select(lambda j1: (j1, ev[1].Where(lambda tp2: DeltaR(tp2.eta(), tp2.phi(), j1.eta(), j1.phi()) < 0.4)))') \ .Select('lambda ji: (ji[0].pt(), 0 if ji[1].Count()==0 else abs(ji[1].First().prodVtx().x()-ji[1].First().decayVtx().x()))') \ .Where('lambda jall: jall[0] > 40.0') \ .value() lines = get_lines_of_code(r) print_lines(lines) l_numbers = find_line_numbers_with("if (i_obj", lines) for line in [lines[ln] for ln in l_numbers]: assert "x()" not in line
def test_per_jet_item_with_event_level(): r = dataset_for_testing() \ .Select('lambda e: (e.Jets("AntiKt4EMTopoJets").Select(lambda j: j.pt()), e.EventInfo("EventInfo").runNumber())') \ .SelectMany(lambda ji: ji[0].Select(lambda pt: { 'JetPt': pt, 'runNumber': ji[1]} )) \ .value() lines = get_lines_of_code(r) print_lines(lines) l_jetpt = find_line_with("_JetPt", lines) l_runnum = find_line_with("_runNumber", lines) l_fill = find_line_with("->Fill()", lines) assert l_jetpt + 1 == l_runnum assert l_runnum + 1 == l_fill
def test_Select_of_2D_with_where(): # This should generate a 2D array. r = dataset_for_testing() \ .Select('lambda e: e.Jets("AntiKt4EMTopoJets").Select(lambda j: e.Electrons("Electrons").Where(lambda ele: ele.pt() > 10).Select(lambda e: e.pt()))') \ .value() lines = get_lines_of_code(r) print_lines(lines) l_vector_decl = find_line_with("vector<double>", lines) l_vector_active = len(find_open_blocks(lines[:l_vector_decl])) l_first_push = find_line_with("push_back", lines) l_first_push_active = len(find_open_blocks(lines[:l_first_push])) assert ( l_vector_active + 2 ) == l_first_push_active # +2 because it is inside the for loop and the if block
def test_Select_of_3D_array(): # This should generate a 2D array. r = dataset_for_testing() \ .Select('lambda e: e.Jets("AntiKt4EMTopoJets").Select(lambda j: e.Electrons("Electrons").Select(lambda e: e.Jets("AntiKt4EMTopoJets").Select(lambda j: j.pt())))') \ .value() lines = get_lines_of_code(r) print_lines(lines) l_vector_decl = find_line_with("vector<double> ", lines) l_vector_active = len(find_open_blocks(lines[:l_vector_decl])) l_vector_double_decl = find_line_with("vector<std::vector<double>>", lines) l_vector_double_active = len(find_open_blocks( lines[:l_vector_double_decl])) assert l_vector_active == (l_vector_double_active + 1)
def test_per_jet_with_matching_and_zeros_and_sum(): # Trying to repro a bug we saw in the wild r = dataset_for_testing() \ .Select('lambda e: (e.Jets("AntiKt4EMTopoJets"),e.TruthParticles("TruthParticles").Where(lambda tp1: tp1.pdgId() == 35))') \ .SelectMany('lambda ev: ev[0].Select(lambda j1: (j1, ev[1].Where(lambda tp2: DeltaR(tp2.eta(), tp2.phi(), j1.eta(), j1.phi()) < 0.4)))') \ .Select(lambda ji: { 'JetPts': ji[0].pt(), 'NumLLPs': 0 if ji[1].Count() == 0 else (ji[1].First().pt()-ji[1].First().pt()), 'sums': ji[0].getAttributeVectorFloat("EnergyPerSampling").Sum()}) \ .value() lines = get_lines_of_code(r) print_lines(lines) l_jetpt = find_line_with("_JetPts", lines) l_nllp = find_line_with("_NumLLPs", lines) l_fill = find_line_with("->Fill()", lines) assert l_jetpt + 1 == l_nllp assert l_nllp + 2 == l_fill
def test_Select_Multiple_arrays_2_step(): # The following statement should be a straight sequence, not an array. r = dataset_for_testing() \ .Select('lambda e: e.Jets("AntiKt4EMTopoJets")') \ .Select('lambda jets: (jets.Select(lambda j: j.pt()/1000.0),jets.Select(lambda j: j.eta()))') \ .value() # Check to see if there mention of push_back anywhere. lines = get_lines_of_code(r) print_lines(lines) l_push_back = find_line_numbers_with("push_back", lines) assert all([ len([l for l in find_open_blocks(lines[:pb]) if "for" in l]) == 1 for pb in l_push_back ]) assert 2 == ["push_back" in l for l in lines].count(True) l_push_back = find_line_with("Fill()", lines) active_blocks = find_open_blocks(lines[:l_push_back]) assert 0 == ["for" in a for a in active_blocks].count(True)
def test_per_jet_with_Count_matching(): # Trying to repro a bug we saw in the wild # The problem is with the "Where" below, it gets moved way up to the top. If it is put near the top then the # generated code is fine. In this case, where it is currently located, the code generated to look at the DeltaR particles # is missed when calculating the y() component (for some reason). This bug may not be in the executor, but, rather, may # be in the function simplifier. # Also, if the "else" doesn't include a "first" thing, then things seem to work just fine too. # .Where('lambda jall: jall[0].pt() > 40.0') \ r = dataset_for_testing() \ .Select('lambda e: (e.Jets("AntiKt4EMTopoJets"),e.TruthParticles("TruthParticles").Where(lambda tp1: tp1.pdgId() == 35))') \ .SelectMany('lambda ev: ev[0].Select(lambda j1: (j1, ev[1].Where(lambda tp2: DeltaR(tp2.eta(), tp2.phi(), j1.eta(), j1.phi()) < 0.4)))') \ .Select('lambda ji: (ji[0].pt(), 0 if ji[1].Count()==0 else ji[1].First().prodVtx().y())') \ .Where('lambda jall: jall[0] > 40.0') \ .value() lines = get_lines_of_code(r) print_lines(lines) l = find_line_numbers_with("if (0)", lines) assert len(l) == 0