def test_Aggregate_not_initial_const_SUM():
r = atlas_xaod_dataset() \
.Select("lambda e: e.Jets('AntiKt4EMTopoJets').Select(lambda j: j.pt()/1000).Sum()") \
.value()
lines = get_lines_of_code(r)
print_lines(lines)
l_sets = find_line_numbers_with("/1000", lines)
assert 1 == len(l_sets)
def test_tree_name():
r = atlas_xaod_dataset() \
.Select("lambda e: e.Jets('AntiKt4EMTopoJets').Select(lambda j: j.pt()/1000).Sum()") \
.AsROOTTTree('junk.root', 'analysis', ['fork']) \
.value()
lines = get_lines_of_code(r)
print_lines(lines)
l_sets = find_line_numbers_with('tree("analysis")', lines)
assert 1 == len(l_sets)
def test_First_selects_collection_count():
# Make sure that we have the "First" predicate after if Where's if statement.
r = atlas_xaod_dataset() \
.Select('lambda e: e.Jets("AntiKt4EMTopoJets").Select(lambda j: e.Tracks("InDetTrackParticles")).First().Count()') \
.value()
lines = get_lines_of_code(r)
print_lines(lines)
ln = find_line_numbers_with("for", lines)
assert 2 == len(ln)
def test_metadata_returned_collection_double_ptr():
# The following statement should be a straight sequence, not an array.
r = (atlas_xaod_dataset().MetaData({
'metadata_type': 'add_method_type_info',
'type_string': 'xAOD::Jet',
'method_name': 'pt',
'return_type_element': 'myobj**',
}).MetaData({
'metadata_type': 'add_method_type_info',
'type_string': 'myobj',
'method_name': 'value',
'return_type': 'double',
}).SelectMany(lambda e: e.Jets("AntiKt4EMTopoJets")).SelectMany(
lambda j: j.pt()).Select(lambda pt: pt.value() * 2).value())
# Check to see if there mention of push_back anywhere.
lines = get_lines_of_code(r)
print_lines(lines)
value_ref = find_line_numbers_with(')->value()*2', lines)
assert len(value_ref) == 1
deref = find_line_numbers_with('(*', lines)
assert len(deref) == 1
def test_Range_good_call():
# The following statement should be a straight sequence, not an array.
r = (atlas_xaod_dataset().SelectMany(
lambda e: e.Jets("AntiKt4EMTopoJets")).Select(lambda j: Range(
0, 10).Select(lambda index: j.pt() * index)).value())
# Check to see if there mention of push_back anywhere.
lines = get_lines_of_code(r)
print_lines(lines)
for_loops = find_line_numbers_with('for (', lines)
assert len(for_loops) == 2
find_line_with('(0)', lines)
find_line_with('(10)', lines)
def test_per_jet_with_delta():
# Trying to repro a bug we saw in the wild
r = atlas_xaod_dataset() \
.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_sequence_with_where_first():
r = atlas_xaod_dataset() \
.Select('lambda e: e.Jets("AntiKt4EMTopoJets").Select(lambda j: e.Tracks("InDetTrackParticles").Where(lambda t: t.pt() > 1000.0)).First().Count()') \
.value()
lines = get_lines_of_code(r)
print_lines(lines)
l_first = find_line_numbers_with("if (is_first", lines)
assert 1 == len(l_first)
active_blocks = find_open_blocks(lines[:l_first[0]])
assert 1 == ["for" in a for a in active_blocks].count(True)
l_agg = find_line_with("+1", lines)
active_blocks = find_open_blocks(lines[:l_agg])
assert 1 == [">1000" 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 = atlas_xaod_dataset() \
.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)
ln = find_line_numbers_with("if (0)", lines)
assert len(ln) == 0
def test_Select_Multiple_arrays_2_step():
# The following statement should be a straight sequence, not an array.
r = atlas_xaod_dataset() \
.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([ln for ln in find_open_blocks(lines[:pb]) if "for" in ln]) == 1
for pb in l_push_back
])
assert 2 == ["push_back" in ln for ln 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_Select_of_2D_array():
# This should generate a 2D array.
r = atlas_xaod_dataset() \
.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