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