def test_3v_z():
df = DataFrame()
v1 = a_3v(df)
df1 = v1.z
assert df1 is not None
assert df1.child_expr is not None
assert ast.dump(df1.child_expr) == "Attribute(value=ast_DataFrame(), " \
"attr='z', ctx=Load())"
async def test_run_mapseq(good_xaod, hep_tables_make_local_call): # NOQA
from hl_tables.atlas import a_3v
truth = good_xaod.TruthParticles('TruthParticles')
llp_truth = truth[truth.pdgId == 35]
llp_good_truth = llp_truth[llp_truth.hasProdVtx & llp_truth.hasDecayVtx]
l_prod = a_3v(llp_good_truth.prodVtx)
l_decay = a_3v(llp_good_truth.decayVtx)
lxy = (l_decay - l_prod).xy
lxy_2 = lxy[lxy.Count() == 2]
has_1muon = lxy_2[lxy_2.mapseq(lambda s: s[0] > 1 | s[1] < 2)].Count()
x = xaod_runner()
r = await x.process(has_1muon)
assert isinstance(r, DataFrame)
all_args = hep_tables_make_local_call.call_args_list
for a in all_args:
d = a[0][0]
from dataframe_expressions import render
r, _ = render(d)
print(ast.dump(r))
assert len(all_args) == 4
def test_3v_add_xy():
df = DataFrame()
v1 = a_3v(df)
v2 = a_3v(df)
r = v1 + v2
df1 = r.xy
assert df1.child_expr is not None
assert ast.dump(df1.child_expr) == "Call(func=Attribute(value=ast_DataFrame()," \
" attr='sqrt', ctx=Load()), args=[], keywords=[])"
assert isinstance(df1.child_expr, ast.Call)
assert isinstance(df1.child_expr.func, ast.Attribute)
df1_parent = cast(ast_DataFrame, df1.child_expr.func.value).dataframe
assert df1_parent.child_expr is not None
assert ast.dump(df1_parent.child_expr) == "BinOp(left=ast_DataFrame(), " \
"op=Add(), right=ast_DataFrame())"
assert isinstance(df1_parent.child_expr, ast.BinOp)
df1_parent_parent = cast(ast_DataFrame,
df1_parent.child_expr.left).dataframe
assert df1_parent_parent.child_expr is not None
assert ast.dump(df1_parent_parent.child_expr) == "BinOp(left=ast_DataFrame(), " \
"op=Mult(), right=ast_DataFrame())"
assert isinstance(df1_parent_parent.child_expr, ast.BinOp)
df1_parent_parent_parent = cast(
ast_DataFrame, df1_parent_parent.child_expr.left).dataframe
assert df1_parent_parent_parent.child_expr is not None
assert ast.dump(df1_parent_parent_parent.child_expr) == \
"BinOp(left=ast_DataFrame(), op=Add(), right=ast_DataFrame())"
assert isinstance(df1_parent_parent_parent.child_expr, ast.BinOp)
assert isinstance(df1_parent_parent_parent.child_expr.left, ast_DataFrame)
df1_parent_parent_parent_parent = \
cast(ast_DataFrame, df1_parent_parent_parent.child_expr.left).dataframe
assert df1_parent_parent_parent_parent.child_expr is not None
assert ast.dump(df1_parent_parent_parent_parent.child_expr) == \
"Attribute(value=ast_DataFrame(), attr='x', ctx=Load())"
def test_3v_create():
df = DataFrame()
v = a_3v(df)
assert v is not None