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_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_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_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_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_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_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_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_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_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_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
def test_First_Of_Select_is_not_array():
# The following statement should be a straight sequence, not an array.
r = EventDataset("file://root.root") \
.Select('lambda e: e.Jets("AntiKt4EMTopoJets").Select(lambda j: j.pt()/1000.0).Where(lambda jpt: jpt > 10.0).First()') \
.AsPandasDF('FirstJetPt') \
.value(executor=exe_for_test)
# Check to see if there mention of push_back anywhere.
lines = get_lines_of_code(r)
print_lines(lines)
assert all("push_back" not in l for l in lines)
l_fill = find_line_with("Fill()", lines)
active_blocks = find_open_blocks(lines[:l_fill])
assert 3==[(("for" in a) or ("if" in a)) for a in active_blocks].count(True)
l_set = find_line_with("_FirstJetPt", lines)
active_blocks = find_open_blocks(lines[:l_set])
assert 3==[(("for" in a) or ("if" in a)) for a in active_blocks].count(True)
l_true = find_line_with("(true)", lines)
active_blocks = find_open_blocks(lines[:l_true])
assert 0==[(("for" in a) or ("if" in a)) for a in active_blocks].count(True)
def test_Select_of_2D_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.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_numbers_with("push_back", lines)
assert len(l_first_push) == 2
l_first_push_active = len(find_open_blocks(lines[:l_first_push[0]]))
assert (l_vector_active + 1) == l_first_push_active
# Now, make sure the second push_back is at the right level.
l_second_push_active = len(find_open_blocks(lines[:l_first_push[1]]))
assert (l_second_push_active + 1) == l_first_push_active
